1 files changed, 177 insertions, 283 deletions

diff --git a/src/lib/libcrypto/mlkem/mlkem_internal.h b/src/lib/libcrypto/mlkem/mlkem_internal.h
index 7e6c313aa9..2b3157256e 100644
--- a/src/lib/libcrypto/mlkem/mlkem_internal.h
+++ b/src/lib/libcrypto/mlkem/mlkem_internal.h
@@ -1,6 +1,7 @@
-/*      $OpenBSD: mlkem_internal.h,v 1.9 2025/08/19 21:37:08 tb Exp $ */
+/*      $OpenBSD: mlkem_internal.h,v 1.10 2025/09/05 23:30:12 beck Exp $ */
 /*
  * Copyright (c) 2023, Google Inc.
+ * Copyright (c) 2025, Bob Beck <beck@obtuse.com>
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -26,402 +27,295 @@ extern "C" {
 #endif
 
 __BEGIN_HIDDEN_DECLS
-/*
- * MLKEM_SEED_LENGTH is the number of bytes in an ML-KEM seed. An ML-KEM
- * seed is normally used to represent a private key.
- */
-#define MLKEM_SEED_LENGTH 64
 
-/*
- * MLKEM_SHARED_SECRET_LENGTH is the number of bytes in an ML-KEM shared
- * secret.
- */
-#define MLKEM_SHARED_SECRET_LENGTH 32
+/* Public opaque ML-KEM key structures. */
 
-/*
- * |MLKEM_encap_external_entropy| behaves exactly like the public |MLKEM_encap|
- * with the entropy provided by the caller. It is directly called internally
- * and by tests.
- */
-int
-MLKEM_encap_external_entropy(const MLKEM_public_key *public_key,
-    const uint8_t *entropy, uint8_t **out_ciphertext,
-    size_t *out_ciphertext_len, uint8_t **out_shared_secret,
-    size_t *out_shared_secret_len);
+#define MLKEM_PUBLIC_KEY_UNINITIALIZED 1
+#define MLKEM_PUBLIC_KEY_INITIALIZED 2
+#define MLKEM_PRIVATE_KEY_UNINITIALIZED 3
+#define MLKEM_PRIVATE_KEY_INITIALIZED 4
 
-/*
- * |MLKEM_generate_key_external_entropy| behaves exactly like the public
- * |MLKEM_generate_key| with the entropy provided by the caller.
- * It is directly called internally and by tests.
- */
-int
-MLKEM_generate_key_external_entropy(MLKEM_private_key *private_key,
-    uint8_t **out_encoded_public_key, size_t *out_encoded_public_key_len,
-    const uint8_t *entropy);
+struct MLKEM_public_key_st {
+        uint16_t rank;
+        int state;
+        struct MLKEM768_public_key *key_768;
+        struct MLKEM1024_public_key *key_1024;
+};
+
+struct MLKEM_private_key_st {
+        uint16_t rank;
+        int state;
+        struct MLKEM768_private_key *key_768;
+        struct MLKEM1024_private_key *key_1024;
+};
 
 /*
- * ML-KEM-768
+ * ML-KEM-768 and ML-KEM-1024
  *
  * This implements the Module-Lattice-Based Key-Encapsulation Mechanism from
  * https://csrc.nist.gov/pubs/fips/204/final
+ *
+ * You should prefer ML-KEM-768 where possible. ML-KEM-1024 is larger and exists
+ * for people who are obsessed with more 'bits of crypto', and who are also
+ * lacking the knowledge to realize that anything that can count to 256 bits
+ * must likely use an equivalent amount of energy to that of an entire star to
+ * do so.
+ *
+ * ML-KEM-768 is adequate to protect against a future cryptographically relevant
+ * quantum computer, VIC 20, abacus, or carefully calibrated reference dog. I
+ * for one plan on welcoming our new Kardashev-II civilization overlords with
+ * open arms. In the meantime will not waste bytes on the wire by to adding
+ * the fear of the possible future existence of a cryptographically relevant
+ * Dyson sphere to the aforementioned list of fear-inducing future
+ * cryptographically relevant hypotheticals.
+ *
+ * If your carefully calibrated reference dog notices the sun starting to dim,
+ * you might need ML-KEM-1024, but you probably have bigger concerns than
+ * the decryption of your stored past TLS sessions at that point.
  */
 
 /*
- * MLKEM768_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM768 public
- * key.
+ * 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.
  */
-#define MLKEM768_PUBLIC_KEY_BYTES 1184
-
-/* MLKEM_SEED_BYTES is the number of bytes in an ML-KEM seed. */
-#define MLKEM_SEED_BYTES 64
+struct MLKEM1024_public_key {
+        uint8_t bytes[512 * (4 + 16) + 32 + 32];
+        uint16_t alignment;
+};
 
 /*
- *  MLKEM_SHARED_SECRET_BYTES is the number of bytes in the ML-KEM768 shared
- * secret. Although the round-3 specification has a variable-length output, the
- * final ML-KEM construction is expected to use a fixed 32-byte output. To
- * simplify the future transition, we apply the same restriction.
+ * 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.
  */
-#define MLKEM_SHARED_SECRET_BYTES 32
+struct MLKEM1024_private_key {
+        uint8_t bytes[512 * (4 + 4 + 16) + 32 + 32 + 32];
+        uint16_t alignment;
+};
 
 /*
- * MLKEM_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.
+ * MLKEM768_public_key contains a ML-KEM-768 public key. The contents of this
+ * object should never leave the address space since the format is unstable.
  */
-int MLKEM768_generate_key(
-    uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
-    MLKEM_private_key *out_private_key);
+struct MLKEM768_public_key {
+        uint8_t bytes[512 * (3 + 9) + 32 + 32];
+        uint16_t alignment;
+};
 
 /*
- * 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.
+ * MLKEM768_private_key contains a ML-KEM-768 private key. The contents of this
+ * object should never leave the address space since the format is unstable.
  */
-int MLKEM768_private_key_from_seed(const uint8_t *seed, size_t seed_len,
-    MLKEM_private_key *out_private_key);
+struct MLKEM768_private_key {
+        uint8_t bytes[512 * (3 + 3 + 9) + 32 + 32 + 32];
+        uint16_t alignment;
+};
 
 /*
- * MLKEM_public_from_private sets |*out_public_key| to the public key that
- * corresponds to |private_key|. (This is faster than parsing the output of
- * |MLKEM_generate_key| if, for some reason, you need to encapsulate to a key
- * that was just generated.)
+ * MLKEM_SEED_LENGTH is the number of bytes in an ML-KEM seed. An ML-KEM
+ * seed is normally used to represent a private key.
  */
-void MLKEM768_public_from_private(const MLKEM_private_key *private_key,
-    MLKEM_public_key *out_public_key);
-
-/* MLKEM768_CIPHERTEXT_BYTES is number of bytes in the ML-KEM768 ciphertext. */
-#define MLKEM768_CIPHERTEXT_BYTES 1088
+#define MLKEM_SEED_LENGTH 64
 
 /*
- * 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|.
+ * MLKEM_SHARED_SECRET_LENGTH is the number of bytes in an ML-KEM shared
+ * secret.
  */
-void MLKEM768_encap(const MLKEM_public_key *public_key,
-    uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES]);
+#define MLKEM_SHARED_SECRET_LENGTH 32
 
 /*
- * 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 rreturns 1. If |ciphertext| is invalid,
- * |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.
+ * MLKEM_ENCAP_ENTROPY is the number of bytes of uniformly random entropy
+ * necessary to encapsulate a secret. The entropy will be leaked to the
+ * decapsulating party.
  */
-int MLKEM768_decap(const MLKEM_private_key *private_key,
-    const uint8_t *ciphertext, size_t ciphertext_len,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES]);
+#define MLKEM_ENCAP_ENTROPY 32
 
-/* Serialisation of keys. */
+/* MLKEM1024_CIPHERTEXT_BYTES is number of bytes in the ML-KEM-1024 ciphertext. */
+#define MLKEM1024_CIPHERTEXT_BYTES 1568
+
+/* MLKEM768_CIPHERTEXT_BYTES is number of bytes in the ML-KEM768 ciphertext. */
+#define MLKEM768_CIPHERTEXT_BYTES 1088
 
 /*
- * MLKEM768_marshal_public_key serializes |public_key| to |out| in the standard
- * format for ML-KEM public keys. It returns one on success or zero on allocation
- * error.
+ * MLKEM768_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM768 public
+ * key.
  */
-int MLKEM768_marshal_public_key(const MLKEM_public_key *public_key,
-    uint8_t **output, size_t *output_len);
+#define MLKEM768_PUBLIC_KEY_BYTES 1184
 
 /*
- * MLKEM768_parse_public_key parses a public key, in the format generated by
- * |MLKEM_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|.
+ * MLKEM1024_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM-1024
+ * public key.
  */
-int MLKEM768_parse_public_key(const uint8_t *input, size_t input_len,
-    MLKEM_public_key *out_public_key);
+#define MLKEM1024_PUBLIC_KEY_BYTES 1568
 
 /*
- * MLKEM_parse_private_key parses a private key, in the format generated by
- * |MLKEM_marshal_private_key|, from |in| and writes the result to
- * |out_private_key|. It returns one on success or zero on parse error or if
- * there are trailing bytes in |in|. This formate is verbose and should be avoided.
- * Private keys should be stored as seeds and parsed using |MLKEM768_private_key_from_seed|.
+ * MLKEM768_PRIVATE_KEY_BYTES is the length of the data produced by
+ * |marshal_private_key| for a RANK768 MLKEM_private_key.
  */
-int MLKEM768_parse_private_key(const uint8_t *input, size_t input_len,
-    MLKEM_private_key *out_private_key);
+#define MLKEM768_PRIVATE_KEY_BYTES 2400
 
 /*
- * ML-KEM-1024
- *
- * ML-KEM-1024 also exists. You should prefer ML-KEM-768 where possible.
+ * MLKEM1024_PRIVATE_KEY_BYTES is the length of the data produced by
+ * |marshal_private_key| for a RANK1024 MLKEM_private_key.
  */
+#define MLKEM1024_PRIVATE_KEY_BYTES 3168
 
 /*
- * MLKEM1024_PUBLIC_KEY_BYTES is the number of bytes in an encoded ML-KEM-1024
- * public key.
+ * Internal MLKEM 768 and MLKEM 1024 functions come largely from BoringSSL, but
+ * converted to C from templated C++. Due to this history, most internal
+ * functions do not allocate, and are expected to be handed memory allocated by
+ * the caller. The caller is generally expected to know what sizes to allocate
+ * based upon the rank of the key (either public or private) that they are
+ * starting with. This avoids the need to handle memory allocation failures
+ * (which boring in C++ just crashes by choice) deep in the implementation, as
+ * what is needed is allocated up front in the public facing functions, and
+ * failure is handled there.
  */
-#define MLKEM1024_PUBLIC_KEY_BYTES 1568
+
+/* Key generation. */
 
 /*
- * MLKEM1024_generate_key generates a random public/private key pair, writes the
+ * mlkem_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.
+ * generate the private key is written to it. The caller is responsible for
+ * ensuring that |out_encoded_public_key| and |out_optonal_seed| point to
+ * enough memory to contain a key and seed for the rank of |out_private_key|.
  */
-int MLKEM1024_generate_key(
-    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
-    MLKEM_private_key *out_private_key);
+int mlkem_generate_key(uint8_t *out_encoded_public_key,
+    uint8_t *optional_out_seed, MLKEM_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.
+ * mlkem_private_key_from_seed modifies |out_private_key| to generate a key of
+ * the rank of |*out_private_key| from a seed that was generated by
+ * |mlkem_generate_key|. It fails and returns 0 if |seed_len| is incorrect, or
+ * if |*out_private_key| has not been initialized. otherwise it writes to
+ * |*out_private_key| and returns 1.
  */
-int MLKEM1024_private_key_from_seed(
-    MLKEM_private_key *out_private_key, const uint8_t *seed,
-    size_t seed_len);
+int mlkem_private_key_from_seed(const uint8_t *seed, size_t seed_len,
+    MLKEM_private_key *out_private_key);
 
 /*
- * 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.)
+ * mlkem_public_from_private sets |*out_public_key| to the public key that
+ * corresponds to |*private_key|. (This is faster than parsing the output of
+ * |MLKEM_generate_key| if, for some reason, you need to encapsulate to a key
+ * that was just generated.)
  */
-void MLKEM1024_public_from_private(const MLKEM_private_key *private_key,
+void mlkem_public_from_private(const MLKEM_private_key *private_key,
     MLKEM_public_key *out_public_key);
 
-/* MLKEM1024_CIPHERTEXT_BYTES is number of bytes in the ML-KEM-1024 ciphertext. */
-#define MLKEM1024_CIPHERTEXT_BYTES 1568
+
+/* Encapsulation and decapsulation of secrets. */
 
 /*
- * MLKEM1024_encap encrypts a random shared secret for |public_key|, writes the
+ * mlkem_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|.
  */
-void MLKEM1024_encap(const MLKEM_public_key *public_key,
-    uint8_t out_ciphertext[MLKEM1024_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES]);
-
+void mlkem_encap(const MLKEM_public_key *public_key,
+    uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_LENGTH]);
 
 /*
- * 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.
+ * mlkem_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,
+ * |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.
  */
-int MLKEM1024_decap(const MLKEM_private_key *private_key,
+int mlkem_decap(const MLKEM_private_key *private_key,
     const uint8_t *ciphertext, size_t ciphertext_len,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES]);
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_LENGTH]);
+
+
+/* Serialisation of keys. */
 
 /*
- * 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.
+ * mlkem_marshal_public_key serializes |public_key| to |output| in the standard
+ * format for ML-KEM public keys. It returns one on success or zero on allocation
+ * error.
  */
-int MLKEM1024_marshal_public_key(const MLKEM_public_key *public_key,
+int mlkem_marshal_public_key(const MLKEM_public_key *public_key,
     uint8_t **output, size_t *output_len);
 
-
 /*
- * MLKEM1024_parse_public_key parses a public key, in the format generated by
- * |MLKEM1024_marshal_public_key|, from |in| and writes the result to
+ * mlkem_parse_public_key parses a public key, in the format generated by
+ * |MLKEM_marshal_public_key|, from |input| 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|.
+ * there are trailing bytes in |input|.
  */
-int MLKEM1024_parse_public_key(const uint8_t *input, size_t input_len,
+int mlkem_parse_public_key(const uint8_t *input, size_t input_len,
     MLKEM_public_key *out_public_key);
 
-
 /*
- * MLKEM1024_parse_private_key parses a private key, in NIST's format for
- * private keys, from |in| and writes the result to |out_private_key|. It
- * returns one on success or zero on parse error or if there are trailing bytes
- * in |in|. This format is verbose and should be avoided. Private keys should be
- * stored as seeds and parsed using |MLKEM1024_private_key_from_seed|.
+ * mlkem_parse_private_key parses a private key, in the format generated by
+ * |MLKEM_marshal_private_key|, from |input| and writes the result to
+ * |out_private_key|. It returns one on success or zero on parse error or if
+ * there are trailing bytes in |input|. This formate is verbose and should be avoided.
+ * Private keys should be stored as seeds and parsed using |mlkem_private_key_from_seed|.
  */
-int MLKEM1024_parse_private_key(const uint8_t *input, size_t input_len,
- MLKEM_private_key *out_private_key);
-
-
-/* XXXX Truly internal stuff below, also in need of de-duping */
+int mlkem_parse_private_key(const uint8_t *input, size_t input_len,
+    MLKEM_private_key *out_private_key);
 
-/*
- * MLKEM_ENCAP_ENTROPY is the number of bytes of uniformly random entropy
- * necessary to encapsulate a secret. The entropy will be leaked to the
- * decapsulating party.
- */
-#define MLKEM_ENCAP_ENTROPY 32
 
-/*
- * MLKEM768_public_key contains a 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;
-};
+/* Functions that are only used for test purposes. */
 
 /*
- * MLKEM768_private_key contains a 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;
-};
-
-/* Public opaque ML-KEM key structures. */
-
-#define MLKEM_PUBLIC_KEY_UNINITIALIZED 1
-#define MLKEM_PUBLIC_KEY_INITIALIZED 2
-#define MLKEM_PRIVATE_KEY_UNINITIALIZED 3
-#define MLKEM_PRIVATE_KEY_INITIALIZED 4
-
-struct MLKEM_public_key_st {
-        uint16_t rank;
-        int state;
-        struct MLKEM768_public_key *key_768;
-        struct MLKEM1024_public_key *key_1024;
-};
-
-struct MLKEM_private_key_st {
-        uint16_t rank;
-        int state;
-        struct MLKEM768_private_key *key_768;
-        struct MLKEM1024_private_key *key_1024;
-};
-
-/*
- * MLKEM768_generate_key_external_entropy is a deterministic function to create a
+ * mlkem_generate_key_external_entropy is a deterministic function to create a
  * pair of ML-KEM 768 keys, using the supplied entropy. The entropy needs to be
  * uniformly random generated. This function is should only be used for tests,
  * regular callers should use the non-deterministic |MLKEM_generate_key|
  * directly.
  */
-int MLKEM768_generate_key_external_entropy(
+int mlkem_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
     MLKEM_private_key *out_private_key,
-    const uint8_t entropy[MLKEM_SEED_BYTES]);
-
-/*
- * MLKEM768_PRIVATE_KEY_BYTES is the length of the data produced by
- * |MLKEM768_marshal_private_key|.
- */
-#define MLKEM768_PRIVATE_KEY_BYTES 2400
+    const uint8_t entropy[MLKEM_SEED_LENGTH]);
 
 /*
- * MLKEM768_marshal_private_key serializes |private_key| to |out| in the standard
+ * mlkem_marshal_private_key serializes |private_key| to |out_private_key| in the standard
  * format for ML-KEM private keys. It returns one on success or zero on
  * allocation error.
  */
-int MLKEM768_marshal_private_key(const MLKEM_private_key *private_key,
+int mlkem_marshal_private_key(const MLKEM_private_key *private_key,
     uint8_t **out_private_key, size_t *out_private_key_len);
 
 /*
- * MLKEM768_encap_external_entropy behaves like |MLKEM768_encap|, but uses
+ * mlkem_encap_external_entropy behaves like |mlkem_encap|, but uses
  * |MLKEM_ENCAP_ENTROPY| bytes of |entropy| for randomization. The decapsulating
  * side will be able to recover |entropy| in full. This function should only be
  * used for tests, regular callers should use the non-deterministic
  * |MLKEM_encap| directly.
  */
-void MLKEM768_encap_external_entropy(
+void mlkem_encap_external_entropy(
     uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_LENGTH],
     const MLKEM_public_key *public_key,
     const uint8_t entropy[MLKEM_ENCAP_ENTROPY]);
 
-
 /*
- * 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_generate_key_external_entropy is a deterministic function to create a
- * pair of ML-KEM 1024 keys, using the supplied entropy. The entropy needs to be
- * uniformly random generated. This function is should only be used for tests,
- * regular callers should use the non-deterministic |MLKEM_generate_key|
- * directly.
- */
-int MLKEM1024_generate_key_external_entropy(
-    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
-    MLKEM_private_key *out_private_key,
-    const uint8_t entropy[MLKEM_SEED_BYTES]);
-
-/*
- * MLKEM1024_PRIVATE_KEY_BYTES is the length of the data produced by
- * |MLKEM1024_marshal_private_key|.
+ * |MLKEM_encap_external_entropy| behaves exactly like the public |MLKEM_encap|
+ * with the entropy provided by the caller. It is directly called internally
+ * and by tests.
  */
-#define MLKEM1024_PRIVATE_KEY_BYTES 3168
+int MLKEM_encap_external_entropy(const MLKEM_public_key *public_key,
+    const uint8_t *entropy, uint8_t **out_ciphertext,
+    size_t *out_ciphertext_len, uint8_t **out_shared_secret,
+    size_t *out_shared_secret_len);
 
 /*
- * MLKEM1024_marshal_private_key serializes |private_key| to |out| in the
- * standard format for ML-KEM private keys. It returns one on success or zero on
- * allocation error.
+ * |MLKEM_generate_key_external_entropy| behaves exactly like the public
+ * |MLKEM_generate_key| with the entropy provided by the caller.
+ * It is directly called internally and by tests.
  */
-int MLKEM1024_marshal_private_key(
-    const MLKEM_private_key *private_key, uint8_t **out_private_key,
-    size_t *out_private_key_len);
+int MLKEM_generate_key_external_entropy(MLKEM_private_key *private_key,
+    uint8_t **out_encoded_public_key, size_t *out_encoded_public_key_len,
+    const uint8_t *entropy);
 
-/*
- * MLKEM_encap_external_entropy behaves like |MLKEM_encap|, but uses
- * |MLKEM_ENCAP_ENTROPY| bytes of |entropy| for randomization. The decapsulating
- * side will be able to recover |entropy| in full. This function should only be
- * used for tests, regular callers should use the non-deterministic
- * |MLKEM_encap| directly.
- */
-void MLKEM1024_encap_external_entropy(
-    uint8_t out_ciphertext[MLKEM1024_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const MLKEM_public_key *public_key,
-    const uint8_t entropy[MLKEM_ENCAP_ENTROPY]);
 
 __END_HIDDEN_DECLS