Fix up MLKEM768_marshal_private_key to not use a passed in CBB

Even though this should remain internal, make it the same
as the public key marshal function, and make the needed
fallout changes in regress.

This does not yet do the bikeshed of renaming the structure
field in the regress ctx, that will wait until a follow on
to convert 1024 in a similar manner

ok tb@

7 files changed, 54 insertions, 53 deletions

diff --git a/src/lib/libcrypto/mlkem/mlkem768.c b/src/lib/libcrypto/mlkem/mlkem768.c
index b20545defc..82adea7b42 100644
--- a/src/lib/libcrypto/mlkem/mlkem768.c
+++ b/src/lib/libcrypto/mlkem/mlkem768.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: mlkem768.c,v 1.10 2025/05/19 07:40:17 beck Exp $ */
+/* $OpenBSD: mlkem768.c,v 1.11 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024, Google Inc.
  * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
@@ -611,6 +611,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 >=
@@ -850,16 +863,9 @@ 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)) {
+        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
@@ -1117,27 +1123,37 @@ MLKEM768_parse_public_key(struct MLKEM768_public_key *public_key,
 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))
-                return 0;
-        if (!CBB_add_bytes(out, priv->pub.public_key_hash,
+        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)))
-                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 7a51197c36..aed051e980 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.5 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_internal.h,v 1.6 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2023, Google Inc.
  *
@@ -57,8 +57,8 @@ int 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
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
index a8495f55e3..e0fd9ca241 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.3 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_iteration_tests.c,v 1.4 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -185,7 +185,7 @@ main(void)
                 .priv = &priv768,
                 .pub = &pub768,
                 .encap_external_entropy = mlkem768_encap_external_entropy,
-                .encode_private_key = mlkem768_encode_private_key,
+                .encode_private_key = mlkem768_marshal_private_key,
                 .generate_key_external_entropy =
                     mlkem768_generate_key_external_entropy,
                 .public_from_private = mlkem768_public_from_private,
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
index a4e7208c76..84b71aebb0 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.4 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_tests.c,v 1.5 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -662,7 +662,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 =
-                    mlkem768_encode_private_key,
+                    mlkem768_marshal_private_key,
         };
         struct MLKEM1024_private_key private_key1024;
         uint8_t encoded_public_key1024[MLKEM1024_PUBLIC_KEY_BYTES];
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
index 8677713c8e..5ec8c08585 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.6 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.c,v 1.7 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -61,25 +61,10 @@ compare_data(const uint8_t *want, const uint8_t *got, size_t len, const char *ms
 }
 
 int
-mlkem768_encode_private_key(const void *private_key, uint8_t **out_buf,
+mlkem768_marshal_private_key(const void *private_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_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 MLKEM768_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 a3b255082f..5c2c400ea5 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.5 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.h,v 1.6 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -30,7 +30,7 @@
 int compare_data(const uint8_t *want, const uint8_t *got, size_t len,
     const char *msg);
 
-int mlkem768_encode_private_key(const void *priv, uint8_t **out_buf,
+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);
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
index a1adc88569..ce1e797904 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.8 2025/05/19 06:47:40 beck Exp $ */
+/*      $OpenBSD: mlkem_unittest.c,v 1.9 2025/05/19 07:53:00 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -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_encode_private_key,
+                .encode_private_key = mlkem768_marshal_private_key,
                 .marshal_public_key = mlkem768_marshal_public_key,
                 .public_from_private = mlkem768_public_from_private,
         };