7 files changed, 68 insertions, 66 deletions

diff --git a/src/lib/libcrypto/mlkem/mlkem1024.c b/src/lib/libcrypto/mlkem/mlkem1024.c
index 27e34898d1..59ebaaa3cb 100644
--- a/src/lib/libcrypto/mlkem/mlkem1024.c
+++ b/src/lib/libcrypto/mlkem/mlkem1024.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem1024.c,v 1.9 2025/05/19 07:40:17 beck Exp $ */
+/* $OpenBSD: mlkem1024.c,v 1.10 2025/05/20 00:33:40 beck 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 >=
@@ -851,16 +864,9 @@ 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));
 }
 
 int
@@ -1120,27 +1126,37 @@ MLKEM1024_parse_public_key(struct MLKEM1024_public_key *public_key,
 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))
-                return 0;
-        if (!CBB_add_bytes(out, priv->pub.public_key_hash,
+        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)))
-                return 0;
-        if (!CBB_add_bytes(out, priv->fo_failure_secret,
+                goto err;
+        if (!CBB_add_bytes(&cbb, priv->fo_failure_secret,
             sizeof(priv->fo_failure_secret)))
-                return 0;
+                goto err;
 
-        return 1;
+        if (!CBB_finish(&cbb, out_private_key, out_private_key_len))
+                goto err;
+
+        ret = 1;
+
+ err:
+        CBB_cleanup(&cbb);
+
+        return ret;
 }
 
 int
diff --git a/src/lib/libcrypto/mlkem/mlkem_internal.h b/src/lib/libcrypto/mlkem/mlkem_internal.h
index aed051e980..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.6 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_internal.h,v 1.7 2025/05/20 00:33:40 beck Exp $ */
 /*
  * Copyright (c) 2023, Google Inc.
  *
@@ -96,8 +96,9 @@ int 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
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
index e0fd9ca241..b93243023c 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.4 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_iteration_tests.c,v 1.5 2025/05/20 00:33:41 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -73,7 +73,7 @@ struct iteration_ctx {
         void *priv;
         void *pub;
 
-        mlkem_encode_private_key_fn encode_private_key;
+        mlkem_marshal_private_key_fn marshal_private_key;
         mlkem_encap_external_entropy_fn encap_external_entropy;
         mlkem_generate_key_external_entropy_fn generate_key_external_entropy;
         mlkem_public_from_private_fn public_from_private;
@@ -127,7 +127,7 @@ MlkemIterativeTest(struct iteration_ctx *ctx)
                     ctx->encoded_public_key_len);
 
                 /* marshal priv to dk as encoded_private_key */
-                if (!ctx->encode_private_key(ctx->priv, &encoded_private_key,
+                if (!ctx->marshal_private_key(ctx->priv, &encoded_private_key,
                     &encoded_private_key_len))
                         errx(1, "encode private key");
 
@@ -185,7 +185,7 @@ main(void)
                 .priv = &priv768,
                 .pub = &pub768,
                 .encap_external_entropy = mlkem768_encap_external_entropy,
-                .encode_private_key = mlkem768_marshal_private_key,
+                .marshal_private_key = mlkem768_marshal_private_key,
                 .generate_key_external_entropy =
                     mlkem768_generate_key_external_entropy,
                 .public_from_private = mlkem768_public_from_private,
@@ -210,7 +210,7 @@ main(void)
                 .priv = &priv1024,
                 .pub = &pub1024,
                 .encap_external_entropy = mlkem1024_encap_external_entropy,
-                .encode_private_key = mlkem1024_encode_private_key,
+                .marshal_private_key = mlkem1024_marshal_private_key,
                 .generate_key_external_entropy =
                     mlkem1024_generate_key_external_entropy,
                 .public_from_private = mlkem1024_public_from_private,
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
index 84b71aebb0..8e04dc6ad2 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.5 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_tests.c,v 1.6 2025/05/20 00:33:41 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -444,7 +444,7 @@ struct keygen_ctx {
         size_t public_key_len;
 
         mlkem_generate_key_external_entropy_fn generate_key_external_entropy;
-        mlkem_encode_private_key_fn encode_private_key;
+        mlkem_marshal_private_key_fn marshal_private_key;
 };
 
 enum keygen_states {
@@ -515,7 +515,7 @@ MlkemKeygenFileTest(struct keygen_ctx *keygen)
 
         keygen->generate_key_external_entropy(keygen->encoded_public_key,
             keygen->private_key, CBS_data(&seed));
-        if (!keygen->encode_private_key(keygen->private_key,
+        if (!keygen->marshal_private_key(keygen->private_key,
             &encoded_private_key, &encoded_private_key_len)) {
                 parse_info(p, "encode private key");
                 goto err;
@@ -614,7 +614,7 @@ MlkemNistKeygenFileTest(struct keygen_ctx *keygen)
 
         keygen->generate_key_external_entropy(keygen->encoded_public_key,
             keygen->private_key, seed);
-        if (!keygen->encode_private_key(keygen->private_key,
+        if (!keygen->marshal_private_key(keygen->private_key,
             &encoded_private_key, &encoded_private_key_len)) {
                 parse_info(p, "encode private key");
                 goto err;
@@ -661,7 +661,7 @@ mlkem_keygen_tests(const char *fn, size_t size, enum test_type test_type)
 
                 .generate_key_external_entropy =
                     mlkem768_generate_key_external_entropy,
-                .encode_private_key =
+                .marshal_private_key =
                     mlkem768_marshal_private_key,
         };
         struct MLKEM1024_private_key private_key1024;
@@ -675,8 +675,8 @@ mlkem_keygen_tests(const char *fn, size_t size, enum test_type test_type)
 
                 .generate_key_external_entropy =
                     mlkem1024_generate_key_external_entropy,
-                .encode_private_key =
-                    mlkem1024_encode_private_key,
+                .marshal_private_key =
+                    mlkem1024_marshal_private_key,
         };
 
         if (size == 768 && test_type == TEST_TYPE_NORMAL)
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
index 5ec8c08585..68bd5d4871 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.7 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.c,v 1.8 2025/05/20 00:33:41 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -75,25 +75,10 @@ mlkem768_marshal_public_key(const void *public_key, uint8_t **out_buf,
 }
 
 int
-mlkem1024_encode_private_key(const void *private_key, uint8_t **out_buf,
+mlkem1024_marshal_private_key(const void *private_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_private_key(&cbb, private_key))
-                goto err;
-        if (!CBB_finish(&cbb, out_buf, out_len))
-                goto err;
-
-        ret = 1;
-
- err:
-        CBB_cleanup(&cbb);
-
-        return ret;
+        return MLKEM1024_marshal_private_key(private_key, out_buf, out_len);
 }
 
 int
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.h
index 5c2c400ea5..1235309f60 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.6 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.h,v 1.7 2025/05/20 00:33:41 beck Exp $ */
 /*
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -34,7 +34,7 @@ int mlkem768_marshal_private_key(const void *priv, uint8_t **out_buf,
     size_t *out_len);
 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,
+int mlkem1024_marshal_private_key(const void *priv, uint8_t **out_buf,
     size_t *out_len);
 int mlkem1024_marshal_public_key(const void *pub, uint8_t **out_buf,
     size_t *out_len);
@@ -69,7 +69,7 @@ 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_marshal_private_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 *);
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
index ce1e797904..3c6aab13c7 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.9 2025/05/19 07:53:00 beck Exp $ */
+/*      $OpenBSD: mlkem_unittest.c,v 1.10 2025/05/20 00:33:41 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -41,7 +41,7 @@ struct unittest_ctx {
         mlkem_generate_key_fn generate_key;
         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_marshal_private_key_fn marshal_private_key;
         mlkem_marshal_public_key_fn marshal_public_key;
         mlkem_public_from_private_fn public_from_private;
 };
@@ -113,7 +113,7 @@ MlKemUnitTest(struct unittest_ctx *ctx)
         free(tmp_buf);
         tmp_buf = NULL;
 
-        if (!ctx->encode_private_key(ctx->priv, &encoded_private_key,
+        if (!ctx->marshal_private_key(ctx->priv, &encoded_private_key,
             &encoded_private_key_len)) {
                 warnx("mlkem768_encode_private_key");
                 failed |= 1;
@@ -137,7 +137,7 @@ MlKemUnitTest(struct unittest_ctx *ctx)
                 failed |= 1;
         }
 
-        if (!ctx->encode_private_key(ctx->priv2, &tmp_buf, &tmp_buf_len)) {
+        if (!ctx->marshal_private_key(ctx->priv2, &tmp_buf, &tmp_buf_len)) {
                 warnx("encode_private_key");
                 failed |= 1;
         }
@@ -205,7 +205,7 @@ mlkem768_unittest(void)
                 .generate_key = mlkem768_generate_key,
                 .parse_private_key = mlkem768_parse_private_key,
                 .parse_public_key = mlkem768_parse_public_key,
-                .encode_private_key = mlkem768_marshal_private_key,
+                .marshal_private_key = mlkem768_marshal_private_key,
                 .marshal_public_key = mlkem768_marshal_public_key,
                 .public_from_private = mlkem768_public_from_private,
         };
@@ -234,7 +234,7 @@ mlkem1024_unittest(void)
                 .generate_key = mlkem1024_generate_key,
                 .parse_private_key = mlkem1024_parse_private_key,
                 .parse_public_key = mlkem1024_parse_public_key,
-                .encode_private_key = mlkem1024_encode_private_key,
+                .marshal_private_key = mlkem1024_marshal_private_key,
                 .marshal_public_key = mlkem1024_marshal_public_key,
                 .public_from_private = mlkem1024_public_from_private,
         };