API changes for ML-KEM

- Get rid of CBB/CBS usage in public api
- Make void functions return int that can fail if malloc fails.

Along with some fallout and resulting bikeshedding in the regress tests.

ok jsing@, tb@

9 files changed, 203 insertions, 176 deletions

diff --git a/src/lib/libcrypto/mlkem/mlkem.h b/src/lib/libcrypto/mlkem/mlkem.h
index 055d92290e..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.5 2025/03/28 12:17:16 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
  *
@@ -81,7 +77,7 @@ struct MLKEM768_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 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 ce6f26e66c..04e106299a 100644
--- a/src/lib/libcrypto/mlkem/mlkem1024.c
+++ b/src/lib/libcrypto/mlkem/mlkem1024.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem1024.c,v 1.7 2025/05/03 08:39:33 tb Exp $ */
+/* $OpenBSD: mlkem1024.c,v 1.8 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
@@ -819,7 +819,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)
@@ -829,7 +829,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);
@@ -843,10 +843,8 @@ 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);
 
@@ -865,7 +863,7 @@ mlkem_marshal_public_key(CBB *out, const struct public_key *pub)
         return 1;
 }
 
-void
+int
 MLKEM1024_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
     struct MLKEM1024_private_key *out_private_key,
@@ -879,7 +877,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);
@@ -894,16 +894,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
@@ -1049,11 +1056,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, MLKEM768_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);
 
@@ -1078,16 +1100,19 @@ 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) ||
+            CBS_len(&cbs) != 0) {
                 return 0;
         }
-        hash_h(pub->public_key_hash, CBS_data(&orig_in), CBS_len(&orig_in));
+        hash_h(pub->public_key_hash, input, input_len);
+
         return 1;
 }
 LCRYPTO_ALIAS(MLKEM1024_parse_public_key);
@@ -1116,26 +1141,28 @@ MLKEM1024_marshal_private_key(CBB *out,
 
 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;
+
+        CBS_init(&cbs, input, input_len);
 
-        if (!CBS_get_bytes(in, &s_bytes, kEncodedVectorSize) ||
+        if (!CBS_get_bytes(&cbs, &s_bytes, kEncodedVectorSize) ||
             !vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime) ||
-            !mlkem_parse_public_key_no_hash(&priv->pub, in)) {
+            !mlkem_parse_public_key_no_hash(&priv->pub, &cbs)) {
                 return 0;
         }
-        memcpy(priv->pub.public_key_hash, CBS_data(in),
+        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 73e293d542..a76971778c 100644
--- a/src/lib/libcrypto/mlkem/mlkem768.c
+++ b/src/lib/libcrypto/mlkem/mlkem768.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem768.c,v 1.8 2025/05/03 08:39:33 tb Exp $ */
+/* $OpenBSD: mlkem768.c,v 1.9 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
@@ -818,7 +818,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)
@@ -828,7 +828,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);
@@ -842,10 +842,8 @@ 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);
 
@@ -864,7 +862,7 @@ mlkem_marshal_public_key(CBB *out, const struct public_key *pub)
         return 1;
 }
 
-void
+int
 MLKEM768_generate_key_external_entropy(
     uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
     struct MLKEM768_private_key *out_private_key,
@@ -878,7 +876,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);
@@ -893,16 +893,22 @@ 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
@@ -965,8 +971,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);
 
@@ -1048,11 +1054,25 @@ 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);
 
@@ -1077,16 +1097,19 @@ 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) ||
+            CBS_len(&cbs) != 0) {
                 return 0;
         }
-        hash_h(pub->public_key_hash, CBS_data(&orig_in), CBS_len(&orig_in));
+        hash_h(pub->public_key_hash, input, input_len);
+
         return 1;
 }
 LCRYPTO_ALIAS(MLKEM768_parse_public_key);
@@ -1115,26 +1138,28 @@ MLKEM768_marshal_private_key(CBB *out,
 
 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;
+
+        CBS_init(&cbs, input, input_len);
 
-        if (!CBS_get_bytes(in, &s_bytes, kEncodedVectorSize) ||
+        if (!CBS_get_bytes(&cbs, &s_bytes, kEncodedVectorSize) ||
             !vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime) ||
-            !mlkem_parse_public_key_no_hash(&priv->pub, in)) {
+            !mlkem_parse_public_key_no_hash(&priv->pub, &cbs)) {
                 return 0;
         }
-        memcpy(priv->pub.public_key_hash, CBS_data(in),
+        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..7a51197c36 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.5 2025/05/19 06:47: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]);
@@ -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]);
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
index 5a61248090..a8495f55e3 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_iteration_tests.c,v 1.2 2024/12/26 07:26:45 tb Exp $ */
+/*      $OpenBSD: mlkem_iteration_tests.c,v 1.3 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -116,8 +116,10 @@ MlkemIterativeTest(struct iteration_ctx *ctx)
                 }
 
                 /* generate ek as encoded_public_key */
-                ctx->generate_key_external_entropy(ctx->encoded_public_key,
-                    ctx->priv, seed);
+                if (!ctx->generate_key_external_entropy(ctx->encoded_public_key,
+                    ctx->priv, seed)) {
+                        errx(1, "generate_key_external_entropy");
+                }
                 ctx->public_from_private(ctx->pub, ctx->priv);
 
                 /* hash in ek */
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
index e9ae417887..a4e7208c76 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_tests.c,v 1.3 2025/05/03 08:34:07 tb Exp $ */
+/*      $OpenBSD: mlkem_tests.c,v 1.4 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -112,7 +112,8 @@ MlkemDecapFileTest(struct decap_ctx *decap)
         parse_get_cbs(p, DECAP_PRIVATE_KEY, &private_key);
         parse_get_int(p, DECAP_RESULT, &should_fail);
 
-        if (!decap->parse_private_key(decap->private_key, &private_key)) {
+        if (!decap->parse_private_key(decap->private_key,
+            CBS_data(&private_key), CBS_len(&private_key))) {
                 if ((failed = !should_fail))
                         parse_info(p, "parse private key");
                 goto err;
@@ -207,7 +208,8 @@ MlkemNistDecapFileTest(struct decap_ctx *decap)
             MLKEM_SHARED_SECRET_BYTES, CBS_len(&k)))
                 goto err;
 
-        if (!decap->parse_private_key(decap->private_key, &dk)) {
+        if (!decap->parse_private_key(decap->private_key, CBS_data(&dk),
+            CBS_len(&dk))) {
                 parse_info(p, "parse private key");
                 goto err;
         }
@@ -360,7 +362,8 @@ MlkemEncapFileTest(struct encap_ctx *encap)
         parse_get_cbs(p, ENCAP_SHARED_SECRET, &shared_secret);
         parse_get_int(p, ENCAP_RESULT, &should_fail);
 
-        if (!encap->parse_public_key(encap->public_key, &public_key)) {
+        if (!encap->parse_public_key(encap->public_key, CBS_data(&public_key),
+            CBS_len(&public_key))) {
                 if ((failed = !should_fail))
                         parse_info(p, "parse public key");
                 goto err;
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
index 1bb2ed3a8b..8677713c8e 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_tests_util.c,v 1.5 2024/12/26 00:04:24 tb Exp $ */
+/*      $OpenBSD: mlkem_tests_util.c,v 1.6 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -83,25 +83,10 @@ mlkem768_encode_private_key(const void *private_key, uint8_t **out_buf,
 }
 
 int
-mlkem768_encode_public_key(const void *public_key, uint8_t **out_buf,
+mlkem768_marshal_public_key(const void *public_key, uint8_t **out_buf,
     size_t *out_len)
 {
-        CBB cbb;
-        int ret = 0;
-
-        if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES))
-                goto err;
-        if (!MLKEM768_marshal_public_key(&cbb, public_key))
-                goto err;
-        if (!CBB_finish(&cbb, out_buf, out_len))
-                goto err;
-
-        ret = 1;
-
- err:
-        CBB_cleanup(&cbb);
-
-        return ret;
+        return MLKEM768_marshal_public_key(out_buf, out_len, public_key);
 }
 
 int
@@ -127,25 +112,10 @@ mlkem1024_encode_private_key(const void *private_key, uint8_t **out_buf,
 }
 
 int
-mlkem1024_encode_public_key(const void *public_key, uint8_t **out_buf,
+mlkem1024_marshal_public_key(const void *public_key, uint8_t **out_buf,
     size_t *out_len)
 {
-        CBB cbb;
-        int ret = 0;
-
-        if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES))
-                goto err;
-        if (!MLKEM1024_marshal_public_key(&cbb, public_key))
-                goto err;
-        if (!CBB_finish(&cbb, out_buf, out_len))
-                goto err;
-
-        ret = 1;
-
- err:
-        CBB_cleanup(&cbb);
-
-        return ret;
+        return MLKEM1024_marshal_public_key(out_buf, out_len, public_key);
 }
 
 int
@@ -173,32 +143,36 @@ mlkem768_encap_external_entropy(uint8_t *out_ciphertext,
             public_key, entropy);
 }
 
-void
+int
 mlkem768_generate_key(uint8_t *out_encoded_public_key,
     uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key)
 {
-        MLKEM768_generate_key(out_encoded_public_key, optional_out_seed,
+        return MLKEM768_generate_key(out_encoded_public_key, optional_out_seed,
             out_private_key);
 }
 
-void
+int
 mlkem768_generate_key_external_entropy(uint8_t *out_encoded_public_key,
     void *out_private_key, const uint8_t 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);
 }
 
 int
-mlkem768_parse_private_key(void *out_private_key, CBS *private_key_cbs)
+mlkem768_parse_private_key(void *out_private_key, const uint8_t *private_key,
+    size_t private_key_len)
 {
-        return MLKEM768_parse_private_key(out_private_key, private_key_cbs);
+        return MLKEM768_parse_private_key(out_private_key, private_key,
+            private_key_len);
 }
 
 int
-mlkem768_parse_public_key(void *out_public_key, CBS *public_key_cbs)
+mlkem768_parse_public_key(void *out_public_key, const uint8_t *public_key,
+    size_t public_key_len)
 {
-        return MLKEM768_parse_public_key(out_public_key, public_key_cbs);
+        return MLKEM768_parse_public_key(out_public_key, public_key,
+            public_key_len);
 }
 
 void
@@ -232,26 +206,28 @@ mlkem1024_encap_external_entropy(uint8_t *out_ciphertext,
             public_key, entropy);
 }
 
-void
+int
 mlkem1024_generate_key(uint8_t *out_encoded_public_key,
     uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key)
 {
-        MLKEM1024_generate_key(out_encoded_public_key, optional_out_seed,
+        return MLKEM1024_generate_key(out_encoded_public_key, optional_out_seed,
             out_private_key);
 }
 
-void
+int
 mlkem1024_generate_key_external_entropy(uint8_t *out_encoded_public_key,
     void *out_private_key, const uint8_t 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);
 }
 
 int
-mlkem1024_parse_private_key(void *out_private_key, CBS *private_key_cbs)
+mlkem1024_parse_private_key(void *out_private_key, const uint8_t *private_key,
+    size_t private_key_len)
 {
-        return MLKEM1024_parse_private_key(out_private_key, private_key_cbs);
+        return MLKEM1024_parse_private_key(out_private_key, private_key,
+            private_key_len);
 }
 
 void
@@ -261,7 +237,9 @@ mlkem1024_public_from_private(void *out_public_key, const void *private_key)
 }
 
 int
-mlkem1024_parse_public_key(void *out_public_key, CBS *public_key_cbs)
+mlkem1024_parse_public_key(void *out_public_key, const uint8_t *public_key,
+    size_t public_key_len)
 {
-        return MLKEM1024_parse_public_key(out_public_key, public_key_cbs);
+        return MLKEM1024_parse_public_key(out_public_key, public_key,
+            public_key_len);
 }
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h
index 7fbe6f76a9..a3b255082f 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h
+++ b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_tests_util.h,v 1.4 2024/12/26 00:04:24 tb Exp $ */
+/*      $OpenBSD: mlkem_tests_util.h,v 1.5 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -32,11 +32,11 @@ int compare_data(const uint8_t *want, const uint8_t *got, size_t len,
 
 int mlkem768_encode_private_key(const void *priv, uint8_t **out_buf,
     size_t *out_len);
-int mlkem768_encode_public_key(const void *pub, uint8_t **out_buf,
+int mlkem768_marshal_public_key(const void *pub, uint8_t **out_buf,
     size_t *out_len);
 int mlkem1024_encode_private_key(const void *priv, uint8_t **out_buf,
     size_t *out_len);
-int mlkem1024_encode_public_key(const void *pub, uint8_t **out_buf,
+int mlkem1024_marshal_public_key(const void *pub, uint8_t **out_buf,
     size_t *out_len);
 
 int mlkem768_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
@@ -46,12 +46,12 @@ void mlkem768_encap(uint8_t *out_ciphertext,
 void mlkem768_encap_external_entropy(uint8_t *out_ciphertext,
     uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES], const void *pub,
     const uint8_t entropy[MLKEM_ENCAP_ENTROPY]);
-void mlkem768_generate_key(uint8_t *out_encoded_public_key,
+int mlkem768_generate_key(uint8_t *out_encoded_public_key,
     uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key);
-void mlkem768_generate_key_external_entropy(uint8_t *out_encoded_public_key,
+int mlkem768_generate_key_external_entropy(uint8_t *out_encoded_public_key,
     void *out_private_key, const uint8_t entropy[MLKEM_SEED_BYTES]);
-int mlkem768_parse_private_key(void *priv, CBS *private_key_cbs);
-int mlkem768_parse_public_key(void *pub, CBS *in);
+int mlkem768_parse_private_key(void *priv, const uint8_t *in, size_t in_len);
+int mlkem768_parse_public_key(void *pub, const uint8_t *in, size_t in_len);
 void mlkem768_public_from_private(void *out_public_key, const void *private_key);
 
 int mlkem1024_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
@@ -61,16 +61,16 @@ void mlkem1024_encap(uint8_t *out_ciphertext,
 void mlkem1024_encap_external_entropy(uint8_t *out_ciphertext,
     uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES], const void *pub,
     const uint8_t entropy[MLKEM_ENCAP_ENTROPY]);
-void mlkem1024_generate_key(uint8_t *out_encoded_public_key,
+int mlkem1024_generate_key(uint8_t *out_encoded_public_key,
     uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key);
-void mlkem1024_generate_key_external_entropy(uint8_t *out_encoded_public_key,
+int mlkem1024_generate_key_external_entropy(uint8_t *out_encoded_public_key,
     void *out_private_key, const uint8_t entropy[MLKEM_SEED_BYTES]);
-int mlkem1024_parse_private_key(void *priv, CBS *private_key_cbs);
-int mlkem1024_parse_public_key(void *pub, CBS *in);
+int mlkem1024_parse_private_key(void *priv, const uint8_t *in, size_t in_len);
+int mlkem1024_parse_public_key(void *pub, const uint8_t *in, size_t in_len);
 void mlkem1024_public_from_private(void *out_public_key, const void *private_key);
 
 typedef int (*mlkem_encode_private_key_fn)(const void *, uint8_t **, size_t *);
-typedef int (*mlkem_encode_public_key_fn)(const void *, uint8_t **, size_t *);
+typedef int (*mlkem_marshal_public_key_fn)(const void *, uint8_t **, size_t *);
 typedef int (*mlkem_decap_fn)(uint8_t [MLKEM_SHARED_SECRET_BYTES],
     const uint8_t *, size_t, const void *);
 typedef void (*mlkem_encap_fn)(uint8_t *, uint8_t [MLKEM_SHARED_SECRET_BYTES],
@@ -78,11 +78,11 @@ typedef void (*mlkem_encap_fn)(uint8_t *, uint8_t [MLKEM_SHARED_SECRET_BYTES],
 typedef void (*mlkem_encap_external_entropy_fn)(uint8_t *,
     uint8_t [MLKEM_SHARED_SECRET_BYTES], const void *,
     const uint8_t [MLKEM_ENCAP_ENTROPY]);
-typedef void (*mlkem_generate_key_fn)(uint8_t *, uint8_t *, void *);
-typedef void (*mlkem_generate_key_external_entropy_fn)(uint8_t *, void *,
+typedef int (*mlkem_generate_key_fn)(uint8_t *, uint8_t *, void *);
+typedef int (*mlkem_generate_key_external_entropy_fn)(uint8_t *, void *,
     const uint8_t [MLKEM_SEED_BYTES]);
-typedef int (*mlkem_parse_private_key_fn)(void *, CBS *);
-typedef int (*mlkem_parse_public_key_fn)(void *, CBS *);
+typedef int (*mlkem_parse_private_key_fn)(void *, const uint8_t *, size_t);
+typedef int (*mlkem_parse_public_key_fn)(void *, const uint8_t *, size_t);
 typedef void (*mlkem_public_from_private_fn)(void *out_public_key,
     const void *private_key);
 
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
index 597297b8cc..a1adc88569 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
@@ -1,4 +1,4 @@
-/*      $OpenBSD: mlkem_unittest.c,v 1.7 2025/05/03 08:34:55 tb Exp $ */
+/*      $OpenBSD: mlkem_unittest.c,v 1.8 2025/05/19 06:47:40 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -42,7 +42,7 @@ struct unittest_ctx {
         mlkem_parse_private_key_fn parse_private_key;
         mlkem_parse_public_key_fn parse_public_key;
         mlkem_encode_private_key_fn encode_private_key;
-        mlkem_encode_public_key_fn encode_public_key;
+        mlkem_marshal_public_key_fn marshal_public_key;
         mlkem_public_from_private_fn public_from_private;
 };
 
@@ -54,36 +54,32 @@ MlKemUnitTest(struct unittest_ctx *ctx)
         uint8_t first_two_bytes[2];
         uint8_t *encoded_private_key = NULL, *tmp_buf = NULL;
         size_t encoded_private_key_len, tmp_buf_len;
-        CBS cbs;
         int failed = 0;
 
-        ctx->generate_key(ctx->encoded_public_key, NULL, ctx->priv);
+        if (!ctx->generate_key(ctx->encoded_public_key, NULL, ctx->priv)) {
+                warnx("generate_key failed");
+                failed |= 1;
+        }
 
         memcpy(first_two_bytes, ctx->encoded_public_key, sizeof(first_two_bytes));
         memset(ctx->encoded_public_key, 0xff, sizeof(first_two_bytes));
 
-        CBS_init(&cbs, ctx->encoded_public_key, ctx->encoded_public_key_len);
-
         /* Parsing should fail because the first coefficient is >= kPrime. */
-        if (ctx->parse_public_key(ctx->pub, &cbs)) {
+        if (ctx->parse_public_key(ctx->pub, ctx->encoded_public_key,
+            ctx->encoded_public_key_len)) {
                 warnx("parse_public_key should have failed");
                 failed |= 1;
         }
 
         memcpy(ctx->encoded_public_key, first_two_bytes, sizeof(first_two_bytes));
-        CBS_init(&cbs, ctx->encoded_public_key, ctx->encoded_public_key_len);
-        if (!ctx->parse_public_key(ctx->pub, &cbs)) {
+        if (!ctx->parse_public_key(ctx->pub, ctx->encoded_public_key,
+            ctx->encoded_public_key_len)) {
                 warnx("MLKEM768_parse_public_key");
                 failed |= 1;
         }
 
-        if (CBS_len(&cbs) != 0u) {
-                warnx("CBS_len must be 0");
-                failed |= 1;
-        }
-
-        if (!ctx->encode_public_key(ctx->pub, &tmp_buf, &tmp_buf_len)) {
-                warnx("encode_public_key");
+        if (!ctx->marshal_public_key(ctx->pub, &tmp_buf, &tmp_buf_len)) {
+                warnx("marshal_public_key");
                 failed |= 1;
         }
         if (ctx->encoded_public_key_len != tmp_buf_len) {
@@ -100,8 +96,8 @@ MlKemUnitTest(struct unittest_ctx *ctx)
         tmp_buf = NULL;
 
         ctx->public_from_private(ctx->pub2, ctx->priv);
-        if (!ctx->encode_public_key(ctx->pub2, &tmp_buf, &tmp_buf_len)) {
-                warnx("encode_public_key");
+        if (!ctx->marshal_public_key(ctx->pub2, &tmp_buf, &tmp_buf_len)) {
+                warnx("marshal_public_key");
                 failed |= 1;
         }
         if (ctx->encoded_public_key_len != tmp_buf_len) {
@@ -125,18 +121,18 @@ MlKemUnitTest(struct unittest_ctx *ctx)
 
         memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes));
         memset(encoded_private_key, 0xff, sizeof(first_two_bytes));
-        CBS_init(&cbs, encoded_private_key, encoded_private_key_len);
 
         /*  Parsing should fail because the first coefficient is >= kPrime. */
-        if (ctx->parse_private_key(ctx->priv2, &cbs)) {
+        if (ctx->parse_private_key(ctx->priv2, encoded_private_key,
+            encoded_private_key_len)) {
                 warnx("MLKEM768_parse_private_key should have failed");
                 failed |= 1;
         }
 
         memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes));
-        CBS_init(&cbs, encoded_private_key, encoded_private_key_len);
 
-        if (!ctx->parse_private_key(ctx->priv2, &cbs)) {
+        if (!ctx->parse_private_key(ctx->priv2, encoded_private_key,
+            encoded_private_key_len)) {
                 warnx("MLKEM768_parse_private_key");
                 failed |= 1;
         }
@@ -210,7 +206,7 @@ mlkem768_unittest(void)
                 .parse_private_key = mlkem768_parse_private_key,
                 .parse_public_key = mlkem768_parse_public_key,
                 .encode_private_key = mlkem768_encode_private_key,
-                .encode_public_key = mlkem768_encode_public_key,
+                .marshal_public_key = mlkem768_marshal_public_key,
                 .public_from_private = mlkem768_public_from_private,
         };
 
@@ -239,7 +235,7 @@ mlkem1024_unittest(void)
                 .parse_private_key = mlkem1024_parse_private_key,
                 .parse_public_key = mlkem1024_parse_public_key,
                 .encode_private_key = mlkem1024_encode_private_key,
-                .encode_public_key = mlkem1024_encode_public_key,
+                .marshal_public_key = mlkem1024_marshal_public_key,
                 .public_from_private = mlkem1024_public_from_private,
         };