4 files changed, 248 insertions, 144 deletions

diff --git a/src/lib/libcrypto/mlkem/mlkem.h b/src/lib/libcrypto/mlkem/mlkem.h
index a5645facc6..a2c5d7fed0 100644
--- a/src/lib/libcrypto/mlkem/mlkem.h
+++ b/src/lib/libcrypto/mlkem/mlkem.h
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem.h,v 1.4 2024/12/19 23:52:26 tb Exp $ */
+/*      $OpenBSD: mlkem.h,v 1.6 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  *
@@ -25,10 +25,6 @@
 extern "C" {
 #endif
 
-/* Hack for now */
-struct cbs_st;
-struct cbb_st;
-
 /*
  * ML-KEM-768
  *
@@ -78,10 +74,10 @@ struct MLKEM768_private_key {
 /*
  * 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| us not NULL then te seed used to
- * generate te private key is written to it.
+ * the private key. If |optional_out_seed| is not NULL then the seed used to
+ * generate the private key is written to it.
  */
-void MLKEM768_generate_key(
+int 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);
@@ -137,7 +133,7 @@ int MLKEM768_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
  * format for ML-KEM public keys. It returns one on success or zero on allocation
  * error.
  */
-int MLKEM768_marshal_public_key(struct cbb_st *out,
+int MLKEM768_marshal_public_key(uint8_t **output, size_t *output_len,
     const struct MLKEM768_public_key *public_key);
 
 /*
@@ -147,7 +143,7 @@ int MLKEM768_marshal_public_key(struct cbb_st *out,
  * there are trailing bytes in |in|.
  */
 int MLKEM768_parse_public_key(struct MLKEM768_public_key *out_public_key,
-    struct cbs_st *in);
+    const uint8_t *input, size_t input_len);
 
 /*
  * MLKEM_parse_private_key parses a private key, in the format generated by
@@ -157,7 +153,7 @@ int MLKEM768_parse_public_key(struct MLKEM768_public_key *out_public_key,
  * Private keys should be stored as seeds and parsed using |MLKEM768_private_key_from_seed|.
  */
 int MLKEM768_parse_private_key(struct MLKEM768_private_key *out_private_key,
-    struct cbs_st *in);
+    const uint8_t *input, size_t input_len);
 
 /*
  * ML-KEM-1024
@@ -200,7 +196,7 @@ struct MLKEM1024_private_key {
  * the private key. If |optional_out_seed| is not NULL then the seed used to
  * generate the private key is written to it.
  */
-void MLKEM1024_generate_key(
+int 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);
@@ -256,7 +252,7 @@ int MLKEM1024_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
  * format for ML-KEM-1024 public keys. It returns one on success or zero on
  * allocation error.
  */
-int MLKEM1024_marshal_public_key(struct cbb_st *out,
+int MLKEM1024_marshal_public_key(uint8_t **output, size_t *output_len,
     const struct MLKEM1024_public_key *public_key);
 
 /*
@@ -266,7 +262,7 @@ int MLKEM1024_marshal_public_key(struct cbb_st *out,
  * there are trailing bytes in |in|.
  */
 int MLKEM1024_parse_public_key(struct MLKEM1024_public_key *out_public_key,
-    struct cbs_st *in);
+    const uint8_t *input, size_t input_len);
 
 /*
  * MLKEM1024_parse_private_key parses a private key, in NIST's format for
@@ -276,7 +272,7 @@ int MLKEM1024_parse_public_key(struct MLKEM1024_public_key *out_public_key,
  * stored as seeds and parsed using |MLKEM1024_private_key_from_seed|.
  */
 int MLKEM1024_parse_private_key(struct MLKEM1024_private_key *out_private_key,
-    struct cbs_st *in);
+    const uint8_t *input, size_t input_len);
 
 #if defined(__cplusplus)
 }
diff --git a/src/lib/libcrypto/mlkem/mlkem1024.c b/src/lib/libcrypto/mlkem/mlkem1024.c
index f6fccdf6a8..26c4716539 100644
--- a/src/lib/libcrypto/mlkem/mlkem1024.c
+++ b/src/lib/libcrypto/mlkem/mlkem1024.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem1024.c,v 1.6 2025/01/03 08:19:24 tb Exp $ */
+/* $OpenBSD: mlkem1024.c,v 1.11 2025/05/21 02:18:11 kenjiro Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
@@ -612,6 +612,19 @@ vector_encode(uint8_t *out, const vector *a, int bits)
         }
 }
 
+/* Encodes an entire vector as above, but adding it to a CBB */
+static int
+vector_encode_cbb(CBB *cbb, const vector *a, int bits)
+{
+        uint8_t *encoded_vector;
+
+        if (!CBB_add_space(cbb, &encoded_vector, kEncodedVectorSize))
+                return 0;
+        vector_encode(encoded_vector, a, bits);
+
+        return 1;
+}
+
 /*
  * scalar_decode parses |DEGREE * bits| bits from |in| into |DEGREE| values in
  * |out|. It returns one on success and zero if any parsed value is >=
@@ -793,6 +806,8 @@ struct public_key {
         matrix m;
 };
 
+CTASSERT(sizeof(struct MLKEM1024_public_key) == sizeof(struct public_key));
+
 static struct public_key *
 public_key_1024_from_external(const struct MLKEM1024_public_key *external)
 {
@@ -805,6 +820,8 @@ struct private_key {
         uint8_t fo_failure_secret[32];
 };
 
+CTASSERT(sizeof(struct MLKEM1024_private_key) == sizeof(struct private_key));
+
 static struct private_key *
 private_key_1024_from_external(const struct MLKEM1024_private_key *external)
 {
@@ -815,7 +832,7 @@ private_key_1024_from_external(const struct MLKEM1024_private_key *external)
  * Calls |MLKEM1024_generate_key_external_entropy| with random bytes from
  * |RAND_bytes|.
  */
-void
+int
 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)
@@ -825,7 +842,7 @@ MLKEM1024_generate_key(uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES
             entropy_buf;
 
         arc4random_buf(entropy, MLKEM_SEED_BYTES);
-        MLKEM1024_generate_key_external_entropy(out_encoded_public_key,
+        return MLKEM1024_generate_key_external_entropy(out_encoded_public_key,
             out_private_key, entropy);
 }
 LCRYPTO_ALIAS(MLKEM1024_generate_key);
@@ -839,29 +856,20 @@ MLKEM1024_private_key_from_seed(struct MLKEM1024_private_key *out_private_key,
         if (seed_len != MLKEM_SEED_BYTES) {
                 return 0;
         }
-        MLKEM1024_generate_key_external_entropy(public_key_bytes,
+        return MLKEM1024_generate_key_external_entropy(public_key_bytes,
             out_private_key, seed);
-
-        return 1;
 }
 LCRYPTO_ALIAS(MLKEM1024_private_key_from_seed);
 
 static int
 mlkem_marshal_public_key(CBB *out, const struct public_key *pub)
 {
-        uint8_t *vector_output;
-
-        if (!CBB_add_space(out, &vector_output, kEncodedVectorSize)) {
+        if (!vector_encode_cbb(out, &pub->t, kLog2Prime))
                 return 0;
-        }
-        vector_encode(vector_output, &pub->t, kLog2Prime);
-        if (!CBB_add_bytes(out, pub->rho, sizeof(pub->rho))) {
-                return 0;
-        }
-        return 1;
+        return CBB_add_bytes(out, pub->rho, sizeof(pub->rho));
 }
 
-void
+int
 MLKEM1024_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
     struct MLKEM1024_private_key *out_private_key,
@@ -875,7 +883,9 @@ MLKEM1024_generate_key_external_entropy(
         uint8_t hashed[64];
         vector error;
         CBB cbb;
+        int ret = 0;
 
+        memset(&cbb, 0, sizeof(CBB));
         memcpy(augmented_seed, entropy, 32);
         augmented_seed[32] = RANK1024;
         hash_g(hashed, augmented_seed, 33);
@@ -890,16 +900,23 @@ MLKEM1024_generate_key_external_entropy(
         matrix_mult_transpose(&priv->pub.t, &priv->pub.m, &priv->s);
         vector_add(&priv->pub.t, &error);
 
-        /* XXX - error checking. */
-        CBB_init_fixed(&cbb, out_encoded_public_key, MLKEM1024_PUBLIC_KEY_BYTES);
-        if (!mlkem_marshal_public_key(&cbb, &priv->pub)) {
-                abort();
-        }
-        CBB_cleanup(&cbb);
+        if (!CBB_init_fixed(&cbb, out_encoded_public_key,
+            MLKEM1024_PUBLIC_KEY_BYTES))
+                goto err;
+
+        if (!mlkem_marshal_public_key(&cbb, &priv->pub))
+                goto err;
 
         hash_h(priv->pub.public_key_hash, out_encoded_public_key,
             MLKEM1024_PUBLIC_KEY_BYTES);
         memcpy(priv->fo_failure_secret, entropy + 32, 32);
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 
 void
@@ -1045,11 +1062,26 @@ MLKEM1024_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
 LCRYPTO_ALIAS(MLKEM1024_decap);
 
 int
-MLKEM1024_marshal_public_key(CBB *out,
+MLKEM1024_marshal_public_key(uint8_t **output, size_t *output_len,
     const struct MLKEM1024_public_key *public_key)
 {
-        return mlkem_marshal_public_key(out,
-            public_key_1024_from_external(public_key));
+        int ret = 0;
+        CBB cbb;
+
+        if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES))
+                goto err;
+        if (!mlkem_marshal_public_key(&cbb,
+            public_key_1024_from_external(public_key)))
+                goto err;
+        if (!CBB_finish(&cbb, output, output_len))
+                goto err;
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 LCRYPTO_ALIAS(MLKEM1024_marshal_public_key);
 
@@ -1062,10 +1094,11 @@ mlkem_parse_public_key_no_hash(struct public_key *pub, CBS *in)
 {
         CBS t_bytes;
 
-        if (!CBS_get_bytes(in, &t_bytes, kEncodedVectorSize) ||
-            !vector_decode(&pub->t, CBS_data(&t_bytes), kLog2Prime)) {
+        if (!CBS_get_bytes(in, &t_bytes, kEncodedVectorSize))
                 return 0;
-        }
+        if (!vector_decode(&pub->t, CBS_data(&t_bytes), kLog2Prime))
+                return 0;
+
         memcpy(pub->rho, CBS_data(in), sizeof(pub->rho));
         if (!CBS_skip(in, sizeof(pub->rho)))
                 return 0;
@@ -1074,64 +1107,84 @@ mlkem_parse_public_key_no_hash(struct public_key *pub, CBS *in)
 }
 
 int
-MLKEM1024_parse_public_key(struct MLKEM1024_public_key *public_key, CBS *in)
+MLKEM1024_parse_public_key(struct MLKEM1024_public_key *public_key,
+    const uint8_t *input, size_t input_len)
 {
         struct public_key *pub = public_key_1024_from_external(public_key);
-        CBS orig_in = *in;
+        CBS cbs;
 
-        if (!mlkem_parse_public_key_no_hash(pub, in) ||
-            CBS_len(in) != 0) {
+        CBS_init(&cbs, input, input_len);
+        if (!mlkem_parse_public_key_no_hash(pub, &cbs))
                 return 0;
-        }
-        hash_h(pub->public_key_hash, CBS_data(&orig_in), CBS_len(&orig_in));
+        if (CBS_len(&cbs) != 0)
+                return 0;
+
+        hash_h(pub->public_key_hash, input, input_len);
+
         return 1;
 }
 LCRYPTO_ALIAS(MLKEM1024_parse_public_key);
 
 int
-MLKEM1024_marshal_private_key(CBB *out,
-    const struct MLKEM1024_private_key *private_key)
+MLKEM1024_marshal_private_key(const struct MLKEM1024_private_key *private_key,
+    uint8_t **out_private_key, size_t *out_private_key_len)
 {
         const struct private_key *const priv = private_key_1024_from_external(
             private_key);
-        uint8_t *s_output;
+        CBB cbb;
+        int ret = 0;
 
-        if (!CBB_add_space(out, &s_output, kEncodedVectorSize)) {
-                return 0;
-        }
-        vector_encode(s_output, &priv->s, kLog2Prime);
-        if (!mlkem_marshal_public_key(out, &priv->pub) ||
-            !CBB_add_bytes(out, priv->pub.public_key_hash,
-            sizeof(priv->pub.public_key_hash)) ||
-            !CBB_add_bytes(out, priv->fo_failure_secret,
-            sizeof(priv->fo_failure_secret))) {
-                return 0;
-        }
-        return 1;
+        if (!CBB_init(&cbb, MLKEM1024_PRIVATE_KEY_BYTES))
+                goto err;
+
+        if (!vector_encode_cbb(&cbb, &priv->s, kLog2Prime))
+                goto err;
+        if (!mlkem_marshal_public_key(&cbb, &priv->pub))
+                goto err;
+        if (!CBB_add_bytes(&cbb, priv->pub.public_key_hash,
+            sizeof(priv->pub.public_key_hash)))
+                goto err;
+        if (!CBB_add_bytes(&cbb, priv->fo_failure_secret,
+            sizeof(priv->fo_failure_secret)))
+                goto err;
+
+        if (!CBB_finish(&cbb, out_private_key, out_private_key_len))
+                goto err;
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 
 int
 MLKEM1024_parse_private_key(struct MLKEM1024_private_key *out_private_key,
-    CBS *in)
+    const uint8_t *input, size_t input_len)
 {
         struct private_key *const priv = private_key_1024_from_external(
             out_private_key);
-        CBS s_bytes;
+        CBS cbs, s_bytes;
 
-        if (!CBS_get_bytes(in, &s_bytes, kEncodedVectorSize) ||
-            !vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime) ||
-            !mlkem_parse_public_key_no_hash(&priv->pub, in)) {
+        CBS_init(&cbs, input, input_len);
+
+        if (!CBS_get_bytes(&cbs, &s_bytes, kEncodedVectorSize))
                 return 0;
-        }
-        memcpy(priv->pub.public_key_hash, CBS_data(in),
+        if (!vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime))
+                return 0;
+        if (!mlkem_parse_public_key_no_hash(&priv->pub, &cbs))
+                return 0;
+
+        memcpy(priv->pub.public_key_hash, CBS_data(&cbs),
             sizeof(priv->pub.public_key_hash));
-        if (!CBS_skip(in, sizeof(priv->pub.public_key_hash)))
+        if (!CBS_skip(&cbs, sizeof(priv->pub.public_key_hash)))
                 return 0;
-        memcpy(priv->fo_failure_secret, CBS_data(in),
+        memcpy(priv->fo_failure_secret, CBS_data(&cbs),
             sizeof(priv->fo_failure_secret));
-        if (!CBS_skip(in, sizeof(priv->fo_failure_secret)))
+        if (!CBS_skip(&cbs, sizeof(priv->fo_failure_secret)))
                 return 0;
-        if (CBS_len(in) != 0)
+        if (CBS_len(&cbs) != 0)
                 return 0;
 
         return 1;
diff --git a/src/lib/libcrypto/mlkem/mlkem768.c b/src/lib/libcrypto/mlkem/mlkem768.c
index bacde0c0b7..653b92d8d8 100644
--- a/src/lib/libcrypto/mlkem/mlkem768.c
+++ b/src/lib/libcrypto/mlkem/mlkem768.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem768.c,v 1.7 2025/01/03 08:19:24 tb Exp $ */
+/* $OpenBSD: mlkem768.c,v 1.12 2025/05/20 00:30:38 beck Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
@@ -89,6 +89,7 @@ static const int kLog2Prime = 12;
 static const uint16_t kHalfPrime = (/*kPrime=*/3329 - 1) / 2;
 static const int kDU768 = 10;
 static const int kDV768 = 4;
+
 /*
  * kInverseDegree is 128^-1 mod 3329; 128 because kPrime does not have a 512th
  * root of unity.
@@ -611,6 +612,19 @@ vector_encode(uint8_t *out, const vector *a, int bits)
         }
 }
 
+/* Encodes an entire vector as above, but adding it to a CBB */
+static int
+vector_encode_cbb(CBB *cbb, const vector *a, int bits)
+{
+        uint8_t *encoded_vector;
+
+        if (!CBB_add_space(cbb, &encoded_vector, kEncodedVectorSize))
+                return 0;
+        vector_encode(encoded_vector, a, bits);
+
+        return 1;
+}
+
 /*
  * scalar_decode parses |DEGREE * bits| bits from |in| into |DEGREE| values in
  * |out|. It returns one on success and zero if any parsed value is >=
@@ -792,6 +806,8 @@ struct public_key {
         matrix m;
 };
 
+CTASSERT(sizeof(struct MLKEM768_public_key) == sizeof(struct public_key));
+
 static struct public_key *
 public_key_768_from_external(const struct MLKEM768_public_key *external)
 {
@@ -804,6 +820,8 @@ struct private_key {
         uint8_t fo_failure_secret[32];
 };
 
+CTASSERT(sizeof(struct MLKEM768_private_key) == sizeof(struct private_key));
+
 static struct private_key *
 private_key_768_from_external(const struct MLKEM768_private_key *external)
 {
@@ -814,7 +832,7 @@ private_key_768_from_external(const struct MLKEM768_private_key *external)
  * Calls |MLKEM768_generate_key_external_entropy| with random bytes from
  * |RAND_bytes|.
  */
-void
+int
 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)
@@ -824,7 +842,7 @@ MLKEM768_generate_key(uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
             entropy_buf;
 
         arc4random_buf(entropy, MLKEM_SEED_BYTES);
-        MLKEM768_generate_key_external_entropy(out_encoded_public_key,
+        return MLKEM768_generate_key_external_entropy(out_encoded_public_key,
             out_private_key, entropy);
 }
 LCRYPTO_ALIAS(MLKEM768_generate_key);
@@ -838,29 +856,20 @@ MLKEM768_private_key_from_seed(struct MLKEM768_private_key *out_private_key,
         if (seed_len != MLKEM_SEED_BYTES) {
                 return 0;
         }
-        MLKEM768_generate_key_external_entropy(public_key_bytes,
+        return MLKEM768_generate_key_external_entropy(public_key_bytes,
             out_private_key, seed);
-
-        return 1;
 }
 LCRYPTO_ALIAS(MLKEM768_private_key_from_seed);
 
 static int
 mlkem_marshal_public_key(CBB *out, const struct public_key *pub)
 {
-        uint8_t *vector_output;
-
-        if (!CBB_add_space(out, &vector_output, kEncodedVectorSize)) {
-                return 0;
-        }
-        vector_encode(vector_output, &pub->t, kLog2Prime);
-        if (!CBB_add_bytes(out, pub->rho, sizeof(pub->rho))) {
+        if (!vector_encode_cbb(out, &pub->t, kLog2Prime))
                 return 0;
-        }
-        return 1;
+        return CBB_add_bytes(out, pub->rho, sizeof(pub->rho));
 }
 
-void
+int
 MLKEM768_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
     struct MLKEM768_private_key *out_private_key,
@@ -874,7 +883,9 @@ MLKEM768_generate_key_external_entropy(
         uint8_t hashed[64];
         vector error;
         CBB cbb;
+        int ret = 0;
 
+        memset(&cbb, 0, sizeof(CBB));
         memcpy(augmented_seed, entropy, 32);
         augmented_seed[32] = RANK768;
         hash_g(hashed, augmented_seed, 33);
@@ -889,16 +900,23 @@ MLKEM768_generate_key_external_entropy(
         matrix_mult_transpose(&priv->pub.t, &priv->pub.m, &priv->s);
         vector_add(&priv->pub.t, &error);
 
-        /* XXX - error checking */
-        CBB_init_fixed(&cbb, out_encoded_public_key, MLKEM768_PUBLIC_KEY_BYTES);
-        if (!mlkem_marshal_public_key(&cbb, &priv->pub)) {
-                abort();
-        }
-        CBB_cleanup(&cbb);
+        if (!CBB_init_fixed(&cbb, out_encoded_public_key,
+            MLKEM768_PUBLIC_KEY_BYTES))
+                goto err;
+
+        if (!mlkem_marshal_public_key(&cbb, &priv->pub))
+                goto err;
 
         hash_h(priv->pub.public_key_hash, out_encoded_public_key,
             MLKEM768_PUBLIC_KEY_BYTES);
         memcpy(priv->fo_failure_secret, entropy + 32, 32);
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 
 void
@@ -961,8 +979,8 @@ MLKEM768_encap(uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
         uint8_t entropy[MLKEM_ENCAP_ENTROPY];
 
         arc4random_buf(entropy, MLKEM_ENCAP_ENTROPY);
-        MLKEM768_encap_external_entropy(out_ciphertext, out_shared_secret,
-            public_key, entropy);
+        MLKEM768_encap_external_entropy(out_ciphertext,
+            out_shared_secret, public_key, entropy);
 }
 LCRYPTO_ALIAS(MLKEM768_encap);
 
@@ -1044,11 +1062,26 @@ MLKEM768_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
 LCRYPTO_ALIAS(MLKEM768_decap);
 
 int
-MLKEM768_marshal_public_key(CBB *out,
+MLKEM768_marshal_public_key(uint8_t **output, size_t *output_len,
     const struct MLKEM768_public_key *public_key)
 {
-        return mlkem_marshal_public_key(out,
-            public_key_768_from_external(public_key));
+        int ret = 0;
+        CBB cbb;
+
+        if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES))
+                goto err;
+        if (!mlkem_marshal_public_key(&cbb,
+            public_key_768_from_external(public_key)))
+                goto err;
+        if (!CBB_finish(&cbb, output, output_len))
+                goto err;
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 LCRYPTO_ALIAS(MLKEM768_marshal_public_key);
 
@@ -1061,10 +1094,11 @@ mlkem_parse_public_key_no_hash(struct public_key *pub, CBS *in)
 {
         CBS t_bytes;
 
-        if (!CBS_get_bytes(in, &t_bytes, kEncodedVectorSize) ||
-            !vector_decode(&pub->t, CBS_data(&t_bytes), kLog2Prime)) {
+        if (!CBS_get_bytes(in, &t_bytes, kEncodedVectorSize))
                 return 0;
-        }
+        if (!vector_decode(&pub->t, CBS_data(&t_bytes), kLog2Prime))
+                return 0;
+
         memcpy(pub->rho, CBS_data(in), sizeof(pub->rho));
         if (!CBS_skip(in, sizeof(pub->rho)))
                 return 0;
@@ -1073,64 +1107,84 @@ mlkem_parse_public_key_no_hash(struct public_key *pub, CBS *in)
 }
 
 int
-MLKEM768_parse_public_key(struct MLKEM768_public_key *public_key, CBS *in)
+MLKEM768_parse_public_key(struct MLKEM768_public_key *public_key,
+    const uint8_t *input, size_t input_len)
 {
         struct public_key *pub = public_key_768_from_external(public_key);
-        CBS orig_in = *in;
+        CBS cbs;
 
-        if (!mlkem_parse_public_key_no_hash(pub, in) ||
-            CBS_len(in) != 0) {
+        CBS_init(&cbs, input, input_len);
+        if (!mlkem_parse_public_key_no_hash(pub, &cbs))
                 return 0;
-        }
-        hash_h(pub->public_key_hash, CBS_data(&orig_in), CBS_len(&orig_in));
+        if (CBS_len(&cbs) != 0)
+                return 0;
+
+        hash_h(pub->public_key_hash, input, input_len);
+
         return 1;
 }
 LCRYPTO_ALIAS(MLKEM768_parse_public_key);
 
 int
-MLKEM768_marshal_private_key(CBB *out,
-    const struct MLKEM768_private_key *private_key)
+MLKEM768_marshal_private_key(const struct MLKEM768_private_key *private_key,
+    uint8_t **out_private_key, size_t *out_private_key_len)
 {
         const struct private_key *const priv = private_key_768_from_external(
             private_key);
-        uint8_t *s_output;
+        CBB cbb;
+        int ret = 0;
 
-        if (!CBB_add_space(out, &s_output, kEncodedVectorSize)) {
-                return 0;
-        }
-        vector_encode(s_output, &priv->s, kLog2Prime);
-        if (!mlkem_marshal_public_key(out, &priv->pub) ||
-            !CBB_add_bytes(out, priv->pub.public_key_hash,
-            sizeof(priv->pub.public_key_hash)) ||
-            !CBB_add_bytes(out, priv->fo_failure_secret,
-            sizeof(priv->fo_failure_secret))) {
-                return 0;
-        }
-        return 1;
+        if (!CBB_init(&cbb, MLKEM768_PRIVATE_KEY_BYTES))
+                goto err;
+
+        if (!vector_encode_cbb(&cbb, &priv->s, kLog2Prime))
+                goto err;
+        if (!mlkem_marshal_public_key(&cbb, &priv->pub))
+                goto err;
+        if (!CBB_add_bytes(&cbb, priv->pub.public_key_hash,
+            sizeof(priv->pub.public_key_hash)))
+                goto err;
+        if (!CBB_add_bytes(&cbb, priv->fo_failure_secret,
+            sizeof(priv->fo_failure_secret)))
+                goto err;
+
+        if (!CBB_finish(&cbb, out_private_key, out_private_key_len))
+                goto err;
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 
 int
 MLKEM768_parse_private_key(struct MLKEM768_private_key *out_private_key,
-    CBS *in)
+    const uint8_t *input, size_t input_len)
 {
         struct private_key *const priv = private_key_768_from_external(
             out_private_key);
-        CBS s_bytes;
+        CBS cbs, s_bytes;
 
-        if (!CBS_get_bytes(in, &s_bytes, kEncodedVectorSize) ||
-            !vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime) ||
-            !mlkem_parse_public_key_no_hash(&priv->pub, in)) {
+        CBS_init(&cbs, input, input_len);
+
+        if (!CBS_get_bytes(&cbs, &s_bytes, kEncodedVectorSize))
                 return 0;
-        }
-        memcpy(priv->pub.public_key_hash, CBS_data(in),
+        if (!vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime))
+                return 0;
+        if (!mlkem_parse_public_key_no_hash(&priv->pub, &cbs))
+                return 0;
+
+        memcpy(priv->pub.public_key_hash, CBS_data(&cbs),
             sizeof(priv->pub.public_key_hash));
-        if (!CBS_skip(in, sizeof(priv->pub.public_key_hash)))
+        if (!CBS_skip(&cbs, sizeof(priv->pub.public_key_hash)))
                 return 0;
-        memcpy(priv->fo_failure_secret, CBS_data(in),
+        memcpy(priv->fo_failure_secret, CBS_data(&cbs),
             sizeof(priv->fo_failure_secret));
-        if (!CBS_skip(in, sizeof(priv->fo_failure_secret)))
+        if (!CBS_skip(&cbs, sizeof(priv->fo_failure_secret)))
                 return 0;
-        if (CBS_len(in) != 0)
+        if (CBS_len(&cbs) != 0)
                 return 0;
 
         return 1;
diff --git a/src/lib/libcrypto/mlkem/mlkem_internal.h b/src/lib/libcrypto/mlkem/mlkem_internal.h
index d3f325932f..1e051970a8 100644
--- a/src/lib/libcrypto/mlkem/mlkem_internal.h
+++ b/src/lib/libcrypto/mlkem/mlkem_internal.h
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_internal.h,v 1.4 2024/12/19 23:52:26 tb Exp $ */
+/*      $OpenBSD: mlkem_internal.h,v 1.7 2025/05/20 00:33:40 beck Exp $ */
 /*
  * Copyright (c) 2023, Google Inc.
  *
@@ -41,7 +41,7 @@ __BEGIN_HIDDEN_DECLS
  * regular callers should use the non-deterministic |MLKEM_generate_key|
  * directly.
  */
-void MLKEM768_generate_key_external_entropy(
+int MLKEM768_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
     struct MLKEM768_private_key *out_private_key,
     const uint8_t entropy[MLKEM_SEED_BYTES]);
@@ -57,8 +57,8 @@ void MLKEM768_generate_key_external_entropy(
  * format for ML-KEM private keys. It returns one on success or zero on
  * allocation error.
  */
-int MLKEM768_marshal_private_key(CBB *out,
-    const struct MLKEM768_private_key *private_key);
+int MLKEM768_marshal_private_key(const struct MLKEM768_private_key *private_key,
+    uint8_t **out_private_key, size_t *out_private_key_len);
 
 /*
  * MLKEM_encap_external_entropy behaves like |MLKEM_encap|, but uses
@@ -80,7 +80,7 @@ void MLKEM768_encap_external_entropy(
  * regular callers should use the non-deterministic |MLKEM_generate_key|
  * directly.
  */
-void MLKEM1024_generate_key_external_entropy(
+int MLKEM1024_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
     struct MLKEM1024_private_key *out_private_key,
     const uint8_t entropy[MLKEM_SEED_BYTES]);
@@ -96,8 +96,9 @@ void MLKEM1024_generate_key_external_entropy(
  * standard format for ML-KEM private keys. It returns one on success or zero on
  * allocation error.
  */
-int MLKEM1024_marshal_private_key(CBB *out,
-    const struct MLKEM1024_private_key *private_key);
+int MLKEM1024_marshal_private_key(
+    const struct MLKEM1024_private_key *private_key, uint8_t **out_private_key,
+    size_t *out_private_key_len);
 
 /*
  * MLKEM_encap_external_entropy behaves like |MLKEM_encap|, but uses