// Copyright 2024 The BoringSSL Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef OPENSSL_HEADER_MLKEM_H
#define OPENSSL_HEADER_MLKEM_H

#include <openssl/base.h>   // IWYU pragma: export

#if defined(__cplusplus)
extern "C" {
#endif


// ML-KEM-768.
//
// This implements the Module-Lattice-Based Key-Encapsulation Mechanism from
// https://csrc.nist.gov/pubs/fips/204/final


// MLKEM768_public_key contains an ML-KEM-768 public key. The contents of this
// object should never leave the address space since the format is unstable.
struct MLKEM768_public_key {
  union {
    uint8_t bytes[512 * (3 + 9) + 32 + 32];
    uint16_t alignment;
  } opaque;
};

// MLKEM768_private_key contains an ML-KEM-768 private key. The contents of this
// object should never leave the address space since the format is unstable.
struct MLKEM768_private_key {
  union {
    uint8_t bytes[512 * (3 + 3 + 9) + 32 + 32 + 32];
    uint16_t alignment;
  } opaque;
};

// MLKEM768_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM-768
// public key.
#define MLKEM768_PUBLIC_KEY_BYTES 1184

// MLKEM_SEED_BYTES is the number of bytes in an ML-KEM seed.
#define MLKEM_SEED_BYTES 64

// MLKEM768_generate_key generates a random public/private key pair, writes the
// encoded public key to |out_encoded_public_key| and sets |out_private_key| to
// the private key. If |optional_out_seed| is not NULL then the seed used to
// generate the private key is written to it.
OPENSSL_EXPORT void MLKEM768_generate_key(
    uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
    struct MLKEM768_private_key *out_private_key);

// MLKEM768_private_key_from_seed derives a private key from a seed that was
// generated by |MLKEM768_generate_key|. It fails and returns 0 if |seed_len| is
// incorrect, otherwise it writes |*out_private_key| and returns 1.
OPENSSL_EXPORT int MLKEM768_private_key_from_seed(
    struct MLKEM768_private_key *out_private_key, const uint8_t *seed,
    size_t seed_len);

// MLKEM768_public_from_private sets |*out_public_key| to the public key that
// corresponds to |private_key|. (This is faster than parsing the output of
// |MLKEM768_generate_key| if, for some reason, you need to encapsulate to a key
// that was just generated.)
OPENSSL_EXPORT void MLKEM768_public_from_private(
    struct MLKEM768_public_key *out_public_key,
    const struct MLKEM768_private_key *private_key);

// MLKEM768_CIPHERTEXT_BYTES is number of bytes in the ML-KEM-768 ciphertext.
#define MLKEM768_CIPHERTEXT_BYTES 1088

// MLKEM_SHARED_SECRET_BYTES is the number of bytes in an ML-KEM shared secret.
#define MLKEM_SHARED_SECRET_BYTES 32

// MLKEM768_encap encrypts a random shared secret for |public_key|, writes the
// ciphertext to |out_ciphertext|, and writes the random shared secret to
// |out_shared_secret|.
OPENSSL_EXPORT void MLKEM768_encap(
    uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
    const struct MLKEM768_public_key *public_key);

// MLKEM768_decap decrypts a shared secret from |ciphertext| using |private_key|
// and writes it to |out_shared_secret|. If |ciphertext_len| is incorrect it
// returns 0, otherwise it returns 1. If |ciphertext| is invalid (but of the
// correct length), |out_shared_secret| is filled with a key that will always be
// the same for the same |ciphertext| and |private_key|, but which appears to be
// random unless one has access to |private_key|. These alternatives occur in
// constant time. Any subsequent symmetric encryption using |out_shared_secret|
// must use an authenticated encryption scheme in order to discover the
// decapsulation failure.
OPENSSL_EXPORT int MLKEM768_decap(
    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
    const uint8_t *ciphertext, size_t ciphertext_len,
    const struct MLKEM768_private_key *private_key);


// Serialisation of keys.

// MLKEM768_marshal_public_key serializes |public_key| to |out| in the standard
// format for ML-KEM-768 public keys. It returns one on success or zero on
// allocation error.
OPENSSL_EXPORT int MLKEM768_marshal_public_key(
    CBB *out, const struct MLKEM768_public_key *public_key);

// MLKEM768_parse_public_key parses a public key, in the format generated by
// |MLKEM768_marshal_public_key|, from |in| and writes the result to
// |out_public_key|. It returns one on success or zero on parse error or if
// there are trailing bytes in |in|.
OPENSSL_EXPORT int MLKEM768_parse_public_key(
    struct MLKEM768_public_key *out_public_key, CBS *in);


// ML-KEM-1024
//
// ML-KEM-1024 also exists. You should prefer ML-KEM-768 where possible.

// MLKEM1024_public_key contains an ML-KEM-1024 public key. The contents of this
// object should never leave the address space since the format is unstable.
struct MLKEM1024_public_key {
  union {
    uint8_t bytes[512 * (4 + 16) + 32 + 32];
    uint16_t alignment;
  } opaque;
};

// MLKEM1024_private_key contains a ML-KEM-1024 private key. The contents of
// this object should never leave the address space since the format is
// unstable.
struct MLKEM1024_private_key {
  union {
    uint8_t bytes[512 * (4 + 4 + 16) + 32 + 32 + 32];
    uint16_t alignment;
  } opaque;
};

// MLKEM1024_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM-1024
// public key.
#define MLKEM1024_PUBLIC_KEY_BYTES 1568

// MLKEM1024_generate_key generates a random public/private key pair, writes the
// encoded public key to |out_encoded_public_key| and sets |out_private_key| to
// the private key. If |optional_out_seed| is not NULL then the seed used to
// generate the private key is written to it.
OPENSSL_EXPORT void MLKEM1024_generate_key(
    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
    struct MLKEM1024_private_key *out_private_key);

// MLKEM1024_private_key_from_seed derives a private key from a seed that was
// generated by |MLKEM1024_generate_key|. It fails and returns 0 if |seed_len|
// is incorrect, otherwise it writes |*out_private_key| and returns 1.
OPENSSL_EXPORT int MLKEM1024_private_key_from_seed(
    struct MLKEM1024_private_key *out_private_key, const uint8_t *seed,
    size_t seed_len);

// MLKEM1024_public_from_private sets |*out_public_key| to the public key that
// corresponds to |private_key|. (This is faster than parsing the output of
// |MLKEM1024_generate_key| if, for some reason, you need to encapsulate to a
// key that was just generated.)
OPENSSL_EXPORT void MLKEM1024_public_from_private(
    struct MLKEM1024_public_key *out_public_key,
    const struct MLKEM1024_private_key *private_key);

// MLKEM1024_CIPHERTEXT_BYTES is number of bytes in the ML-KEM-1024 ciphertext.
#define MLKEM1024_CIPHERTEXT_BYTES 1568

// MLKEM1024_encap encrypts a random shared secret for |public_key|, writes the
// ciphertext to |out_ciphertext|, and writes the random shared secret to
// |out_shared_secret|.
OPENSSL_EXPORT void MLKEM1024_encap(
    uint8_t out_ciphertext[MLKEM1024_CIPHERTEXT_BYTES],
    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
    const struct MLKEM1024_public_key *public_key);

// MLKEM1024_decap decrypts a shared secret from |ciphertext| using
// |private_key| and writes it to |out_shared_secret|. If |ciphertext_len| is
// incorrect it returns 0, otherwise it returns 1. If |ciphertext| is invalid
// (but of the correct length), |out_shared_secret| is filled with a key that
// will always be the same for the same |ciphertext| and |private_key|, but
// which appears to be random unless one has access to |private_key|. These
// alternatives occur in constant time. Any subsequent symmetric encryption
// using |out_shared_secret| must use an authenticated encryption scheme in
// order to discover the decapsulation failure.
OPENSSL_EXPORT int MLKEM1024_decap(
    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
    const uint8_t *ciphertext, size_t ciphertext_len,
    const struct MLKEM1024_private_key *private_key);


// Serialisation of ML-KEM-1024 keys.

// MLKEM1024_marshal_public_key serializes |public_key| to |out| in the standard
// format for ML-KEM-1024 public keys. It returns one on success or zero on
// allocation error.
OPENSSL_EXPORT int MLKEM1024_marshal_public_key(
    CBB *out, const struct MLKEM1024_public_key *public_key);

// MLKEM1024_parse_public_key parses a public key, in the format generated by
// |MLKEM1024_marshal_public_key|, from |in| and writes the result to
// |out_public_key|. It returns one on success or zero on parse error or if
// there are trailing bytes in |in|.
OPENSSL_EXPORT int MLKEM1024_parse_public_key(
    struct MLKEM1024_public_key *out_public_key, CBS *in);


#if defined(__cplusplus)
}  // extern C
#endif

#endif  // OPENSSL_HEADER_MLKEM_H