Add a reasonable ML-KEM API for public use.

Adapt the tests to use this API.

This does not yet make the symbols public in Symbols.list
which will happen shortly with a bump.

This includes some partial rototilling of the non-public
interfaces which will be shortly continued when the internal
code is deduplicated to not have multiple copies for ML-KEM
768 and ML-KEM 1024 (which is just an artifact of unravelling
the boring C++ code).

ok jsing@, tb@

5 files changed, 396 insertions, 593 deletions

diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_iteration_tests.c
index b93243023c..10c4b1f4e0 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.5 2025/05/20 00:33:41 beck Exp $ */
+/*      $OpenBSD: mlkem_iteration_tests.c,v 1.6 2025/08/14 15:48:48 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -21,6 +21,7 @@
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 #include "mlkem.h"
 
@@ -63,46 +64,49 @@ const uint8_t kExpectedAdam1024[32] = {
         0x04, 0xab, 0xdb, 0x94, 0x8b, 0x90, 0x8b, 0x75, 0xba, 0xd5
 };
 
-struct iteration_ctx {
-        uint8_t *encoded_public_key;
-        size_t encoded_public_key_len;
-        uint8_t *ciphertext;
-        size_t ciphertext_len;
-        uint8_t *invalid_ciphertext;
-        size_t invalid_ciphertext_len;
-        void *priv;
-        void *pub;
-
-        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;
-        mlkem_decap_fn decap;
-
-        const uint8_t *start;
-        size_t start_len;
-
-        const uint8_t *expected;
-        size_t expected_len;
-};
-
 static int
-MlkemIterativeTest(struct iteration_ctx *ctx)
+MlkemIterativeTest(int rank)
 {
-        uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
+        const uint8_t *start, *expected;
+        size_t start_len;
         uint8_t encap_entropy[MLKEM_ENCAP_ENTROPY];
-        uint8_t seed[MLKEM_SEED_BYTES] = {0};
+        uint8_t seed[MLKEM_SEED_LENGTH] = {0};
+        uint8_t *shared_secret = NULL;
         sha3_ctx drng, results;
         uint8_t out[32];
         int i;
 
+        start = kExpectedSeedStart;
+        start_len = sizeof(kExpectedSeedStart);
+        switch(rank){
+        case RANK768:
+                expected = kExpectedAdam768;
+                break;
+        case RANK1024:
+                expected = kExpectedAdam1024;
+                break;
+        default:
+                errx(1, "invalid rank %d", rank);
+        }
+
         shake128_init(&drng);
         shake128_init(&results);
 
         shake_xof(&drng);
         for (i = 0; i < 10000; i++) {
-                uint8_t *encoded_private_key = NULL;
-                size_t encoded_private_key_len;
+                uint8_t *encoded_public_key = NULL, *ciphertext = NULL,
+                    *encoded_private_key = NULL, *invalid_ciphertext = NULL;
+                size_t encoded_public_key_len, ciphertext_len,
+                    encoded_private_key_len, invalid_ciphertext_len;
+                MLKEM_private_key *priv;
+                MLKEM_public_key *pub;
+                size_t s_len = 0;
+
+                /* allocate keys for this iteration */
+                if ((priv = MLKEM_private_key_new(rank)) == NULL)
+                        errx(1, "malloc");
+                if ((pub = MLKEM_public_key_new(rank)) == NULL)
+                        errx(1, "malloc");
 
                 /*
                  * This should draw both d and z from DRNG concatenating in
@@ -110,120 +114,91 @@ MlkemIterativeTest(struct iteration_ctx *ctx)
                  */
                 shake_out(&drng, seed, sizeof(seed));
                 if (i == 0) {
-                        if (compare_data(seed, ctx->start, ctx->start_len,
+                        if (compare_data(seed, start, start_len,
                             "seed start") != 0)
                                 errx(1, "compare_data");
                 }
 
                 /* generate ek as encoded_public_key */
-                if (!ctx->generate_key_external_entropy(ctx->encoded_public_key,
-                    ctx->priv, seed)) {
+                if (!MLKEM_generate_key_external_entropy(priv,
+                    &encoded_public_key, &encoded_public_key_len,
+                    seed))
                         errx(1, "generate_key_external_entropy");
-                }
-                ctx->public_from_private(ctx->pub, ctx->priv);
+
+                if (!MLKEM_public_from_private(priv, pub))
+                        errx(1, "public_from_private");
 
                 /* hash in ek */
-                shake_update(&results, ctx->encoded_public_key,
-                    ctx->encoded_public_key_len);
+                shake_update(&results, encoded_public_key,
+                    encoded_public_key_len);
 
                 /* marshal priv to dk as encoded_private_key */
-                if (!ctx->marshal_private_key(ctx->priv, &encoded_private_key,
+                if (!MLKEM_marshal_private_key(priv, &encoded_private_key,
                     &encoded_private_key_len))
-                        errx(1, "encode private key");
+                        errx(1, "marshal private key");
 
                 /* hash in dk */
                 shake_update(&results, encoded_private_key,
                     encoded_private_key_len);
 
-                free(encoded_private_key);
+                freezero(encoded_private_key, encoded_private_key_len);
 
                 /* draw m as encap entropy from DRNG */
                 shake_out(&drng, encap_entropy, sizeof(encap_entropy));
 
                 /* generate ct as ciphertext, k as shared_secret */
-                ctx->encap_external_entropy(ctx->ciphertext, shared_secret,
-                    ctx->pub, encap_entropy);
+                if (!MLKEM_encap_external_entropy(pub, encap_entropy,
+                    &ciphertext, &ciphertext_len, &shared_secret, &s_len))
+                        errx(1, "encap_external_entropy");
 
                 /* hash in ct */
-                shake_update(&results, ctx->ciphertext, ctx->ciphertext_len);
+                shake_update(&results, ciphertext, ciphertext_len);
                 /* hash in k */
-                shake_update(&results, shared_secret, sizeof(shared_secret));
+                shake_update(&results, shared_secret, s_len);
+
+                freezero(shared_secret, s_len);
+                shared_secret = NULL;
+
+                invalid_ciphertext_len = ciphertext_len;
+                if ((invalid_ciphertext = calloc(1, invalid_ciphertext_len))
+                    == NULL)
+                        errx(1, "malloc");
 
                 /* draw ct as invalid_ciphertxt from DRNG */
-                shake_out(&drng, ctx->invalid_ciphertext,
-                    ctx->invalid_ciphertext_len);
+                shake_out(&drng, invalid_ciphertext, invalid_ciphertext_len);
 
                 /* generate k as shared secret from invalid ciphertext */
-                if (!ctx->decap(shared_secret, ctx->invalid_ciphertext,
-                    ctx->invalid_ciphertext_len, ctx->priv))
-                        errx(1, "decap failed");
+                if (!MLKEM_decap(priv, invalid_ciphertext,
+                    invalid_ciphertext_len, &shared_secret, &s_len))
+                        errx(1, "decap failed, iteration %d", i);
 
                 /* hash in k */
-                shake_update(&results, shared_secret, sizeof(shared_secret));
+                shake_update(&results, shared_secret, s_len);
+
+                freezero(shared_secret, s_len);
+                shared_secret = NULL;
+                freezero(invalid_ciphertext, invalid_ciphertext_len);
+                invalid_ciphertext = NULL;
+
+                /* free keys and intermediate products for this iteration */
+                MLKEM_private_key_free(priv);
+                MLKEM_public_key_free(pub);
+                freezero(encoded_public_key, encoded_public_key_len);
+                freezero(ciphertext, ciphertext_len);
         }
         shake_xof(&results);
         shake_out(&results, out, sizeof(out));
 
-        return compare_data(ctx->expected, out, sizeof(out), "final result hash");
+        return compare_data(expected, out, sizeof(out), "final result hash");
 }
 
 int
 main(void)
 {
-        uint8_t encoded_public_key768[MLKEM768_PUBLIC_KEY_BYTES];
-        uint8_t ciphertext768[MLKEM768_CIPHERTEXT_BYTES];
-        uint8_t invalid_ciphertext768[MLKEM768_CIPHERTEXT_BYTES];
-        struct MLKEM768_private_key priv768;
-        struct MLKEM768_public_key pub768;
-        struct iteration_ctx iteration768 = {
-                .encoded_public_key = encoded_public_key768,
-                .encoded_public_key_len = sizeof(encoded_public_key768),
-                .ciphertext = ciphertext768,
-                .ciphertext_len = sizeof(ciphertext768),
-                .invalid_ciphertext = invalid_ciphertext768,
-                .invalid_ciphertext_len = sizeof(invalid_ciphertext768),
-                .priv = &priv768,
-                .pub = &pub768,
-                .encap_external_entropy = mlkem768_encap_external_entropy,
-                .marshal_private_key = mlkem768_marshal_private_key,
-                .generate_key_external_entropy =
-                    mlkem768_generate_key_external_entropy,
-                .public_from_private = mlkem768_public_from_private,
-                .decap = mlkem768_decap,
-                .start = kExpectedSeedStart,
-                .start_len = sizeof(kExpectedSeedStart),
-                .expected = kExpectedAdam768,
-                .expected_len = sizeof(kExpectedAdam768),
-        };
-        uint8_t encoded_public_key1024[MLKEM1024_PUBLIC_KEY_BYTES];
-        uint8_t ciphertext1024[MLKEM1024_CIPHERTEXT_BYTES];
-        uint8_t invalid_ciphertext1024[MLKEM1024_CIPHERTEXT_BYTES];
-        struct MLKEM1024_private_key priv1024;
-        struct MLKEM1024_public_key pub1024;
-        struct iteration_ctx iteration1024 = {
-                .encoded_public_key = encoded_public_key1024,
-                .encoded_public_key_len = sizeof(encoded_public_key1024),
-                .ciphertext = ciphertext1024,
-                .ciphertext_len = sizeof(ciphertext1024),
-                .invalid_ciphertext = invalid_ciphertext1024,
-                .invalid_ciphertext_len = sizeof(invalid_ciphertext1024),
-                .priv = &priv1024,
-                .pub = &pub1024,
-                .encap_external_entropy = mlkem1024_encap_external_entropy,
-                .marshal_private_key = mlkem1024_marshal_private_key,
-                .generate_key_external_entropy =
-                    mlkem1024_generate_key_external_entropy,
-                .public_from_private = mlkem1024_public_from_private,
-                .decap = mlkem1024_decap,
-                .start = kExpectedSeedStart,
-                .start_len = sizeof(kExpectedSeedStart),
-                .expected = kExpectedAdam1024,
-                .expected_len = sizeof(kExpectedAdam1024),
-        };
         int failed = 0;
 
-        failed |= MlkemIterativeTest(&iteration768);
-        failed |= MlkemIterativeTest(&iteration1024);
+        failed |= MlkemIterativeTest(RANK768);
+        failed |= MlkemIterativeTest(RANK1024);
 
         return failed;
 }
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests.c
index 8e04dc6ad2..3269ba951f 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.6 2025/05/20 00:33:41 beck Exp $ */
+/*      $OpenBSD: mlkem_tests.c,v 1.7 2025/08/14 15:48:48 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -38,12 +38,8 @@ enum test_type {
 
 struct decap_ctx {
         struct parse *parse_ctx;
-
-        void *private_key;
-        size_t private_key_len;
-
-        mlkem_parse_private_key_fn parse_private_key;
-        mlkem_decap_fn decap;
+        int rank;
+        MLKEM_private_key *private_key;
 };
 
 enum decap_states {
@@ -89,21 +85,26 @@ decap_init(void *ctx, void *parse_ctx)
 
         decap->parse_ctx = parse_ctx;
 
-        return 1;
+        return (decap->private_key = MLKEM_private_key_new(decap->rank))
+            != NULL;
 }
 
 static void
 decap_finish(void *ctx)
 {
-        (void)ctx;
+        struct decap_ctx *decap = ctx;
+
+        MLKEM_private_key_free(decap->private_key);
+        decap->private_key = NULL;
 }
 
 static int
 MlkemDecapFileTest(struct decap_ctx *decap)
 {
         struct parse *p = decap->parse_ctx;
-        uint8_t shared_secret_buf[MLKEM_SHARED_SECRET_BYTES];
+        uint8_t *shared_secret_buf = NULL;
         CBS ciphertext, shared_secret, private_key;
+        size_t s_len = 0;
         int should_fail;
         int failed = 1;
 
@@ -112,21 +113,28 @@ 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,
+        if (!MLKEM_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;
         }
-        if (!decap->decap(shared_secret_buf,
-            CBS_data(&ciphertext), CBS_len(&ciphertext), decap->private_key)) {
+        if (!MLKEM_decap(decap->private_key, CBS_data(&ciphertext),
+            CBS_len(&ciphertext), &shared_secret_buf, &s_len)) {
                 if ((failed = !should_fail))
                         parse_info(p, "decap");
                 goto err;
         }
 
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                if ((failed = !should_fail))
+                        parse_info(p, "shared secret length %zu != %d", s_len,
+                            MLKEM_SHARED_SECRET_LENGTH);
+                goto err;
+        }
+
         failed = !parse_data_equal(p, "shared_secret", &shared_secret,
-            shared_secret_buf, sizeof(shared_secret_buf));
+            shared_secret_buf, s_len);
 
         if (should_fail != failed) {
                 parse_info(p, "FAIL: should_fail %d, failed %d",
@@ -135,6 +143,7 @@ MlkemDecapFileTest(struct decap_ctx *decap)
         }
 
  err:
+        freezero(shared_secret_buf, s_len);
         return failed;
 }
 
@@ -193,8 +202,9 @@ static int
 MlkemNistDecapFileTest(struct decap_ctx *decap)
 {
         struct parse *p = decap->parse_ctx;
-        uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
+        uint8_t *shared_secret = NULL;
         CBS dk, c, k;
+        size_t s_len = 0;
         int failed = 1;
 
         parse_instruction_get_cbs(p, NIST_DECAP_DK, &dk);
@@ -202,27 +212,34 @@ MlkemNistDecapFileTest(struct decap_ctx *decap)
         parse_get_cbs(p, NIST_DECAP_K, &k);
 
         if (!parse_length_equal(p, "private key",
-            decap->private_key_len, CBS_len(&dk)))
+            MLKEM_private_key_encoded_length(decap->private_key), CBS_len(&dk)))
                 goto err;
         if (!parse_length_equal(p, "shared secret",
             MLKEM_SHARED_SECRET_BYTES, CBS_len(&k)))
                 goto err;
 
-        if (!decap->parse_private_key(decap->private_key, CBS_data(&dk),
+        if (!MLKEM_parse_private_key(decap->private_key, CBS_data(&dk),
             CBS_len(&dk))) {
                 parse_info(p, "parse private key");
                 goto err;
         }
-        if (!decap->decap(shared_secret, CBS_data(&c), CBS_len(&c),
-            decap->private_key)) {
+        if (!MLKEM_decap(decap->private_key, CBS_data(&c), CBS_len(&c),
+            &shared_secret, &s_len)) {
                 parse_info(p, "decap");
                 goto err;
         }
 
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                parse_info(p, "shared secret length %zu != %d", s_len,
+                    MLKEM_SHARED_SECRET_LENGTH);
+                goto err;
+        }
+
         failed = !parse_data_equal(p, "shared secret", &k,
-            shared_secret, MLKEM_SHARED_SECRET_BYTES);
+            shared_secret, s_len);
 
  err:
+        free(shared_secret);
         return failed;
 }
 
@@ -246,46 +263,31 @@ static const struct test_parse nist_decap_parse = {
 };
 
 static int
-mlkem_decap_tests(const char *fn, size_t size, enum test_type test_type)
+mlkem_decap_tests(const char *fn, int rank, enum test_type test_type)
 {
-        struct MLKEM768_private_key private_key768;
-        struct decap_ctx decap768 = {
-                .private_key = &private_key768,
-                .private_key_len = MLKEM768_PRIVATE_KEY_BYTES,
-
-                .parse_private_key = mlkem768_parse_private_key,
-                .decap = mlkem768_decap,
-        };
-        struct MLKEM1024_private_key private_key1024;
-        struct decap_ctx decap1024 = {
-                .private_key = &private_key1024,
-                .private_key_len = MLKEM1024_PRIVATE_KEY_BYTES,
-
-                .parse_private_key = mlkem1024_parse_private_key,
-                .decap = mlkem1024_decap,
+        struct decap_ctx decap = {
+                .rank = rank,
         };
+        int ret = 0;
 
-        if (size == 768 && test_type == TEST_TYPE_NORMAL)
-                return parse_test_file(fn, &decap_parse, &decap768);
-        if (size == 768 && test_type == TEST_TYPE_NIST)
-                return parse_test_file(fn, &nist_decap_parse, &decap768);
-        if (size == 1024 && test_type == TEST_TYPE_NORMAL)
-                return parse_test_file(fn, &decap_parse, &decap1024);
-        if (size == 1024 && test_type == TEST_TYPE_NIST)
-                return parse_test_file(fn, &nist_decap_parse, &decap1024);
+        if (test_type == TEST_TYPE_NORMAL)
+                ret = parse_test_file(fn, &decap_parse, &decap);
+        else if (test_type == TEST_TYPE_NIST)
+                ret = parse_test_file(fn, &nist_decap_parse, &decap);
+        else
+                errx(1, "unknown decap test: rank %d, type %d", rank,
+                    test_type);
 
-        errx(1, "unknown decap test: size %zu, type %d", size, test_type);
+        return ret;
 }
 
 struct encap_ctx {
         struct parse *parse_ctx;
 
-        void *public_key;
+        int rank;
+        MLKEM_public_key *public_key;
         uint8_t *ciphertext;
         size_t ciphertext_len;
-
-        mlkem_parse_public_key_fn parse_public_key;
-        mlkem_encap_external_entropy_fn encap_external_entropy;
 };
 
 enum encap_states {
@@ -338,21 +340,30 @@ encap_init(void *ctx, void *parse_ctx)
 
         encap->parse_ctx = parse_ctx;
 
-        return 1;
+        encap->ciphertext = NULL;
+        return (encap->public_key = MLKEM_public_key_new(encap->rank))
+            != NULL;
 }
 
 static void
 encap_finish(void *ctx)
 {
-        (void)ctx;
+        struct encap_ctx *encap = ctx;
+
+        freezero(encap->ciphertext, encap->ciphertext_len);
+        encap->ciphertext = NULL;
+        MLKEM_public_key_free(encap->public_key);
+        encap->public_key = NULL;
 }
 
 static int
 MlkemEncapFileTest(struct encap_ctx *encap)
 {
-        struct parse *p = encap->parse_ctx;
-        uint8_t shared_secret_buf[MLKEM_SHARED_SECRET_BYTES];
         CBS entropy, public_key, ciphertext, shared_secret;
+        struct parse *p = encap->parse_ctx;
+        uint8_t *shared_secret_buf = NULL;
+        size_t s_len = 0;
+
         int should_fail;
         int failed = 1;
 
@@ -362,20 +373,33 @@ 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, CBS_data(&public_key),
+        if (!MLKEM_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;
         }
-        encap->encap_external_entropy(encap->ciphertext, shared_secret_buf,
-            encap->public_key, CBS_data(&entropy));
+        if (!MLKEM_encap_external_entropy(encap->public_key, CBS_data(&entropy),
+            &encap->ciphertext, &encap->ciphertext_len, &shared_secret_buf,
+            &s_len)) {
+                if ((failed = !should_fail))
+                        parse_info(p, "encap_external_entropy");
+                goto err;
+        }
+
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                if ((failed = !should_fail))
+                        parse_info(p, "shared secret length %zu != %d", s_len,
+                            MLKEM_SHARED_SECRET_LENGTH);
+                goto err;
+        }
 
         failed = !parse_data_equal(p, "shared_secret", &shared_secret,
-            shared_secret_buf, sizeof(shared_secret_buf));
+            shared_secret_buf, s_len);
         failed |= !parse_data_equal(p, "ciphertext", &ciphertext,
             encap->ciphertext, encap->ciphertext_len);
 
+
         if (should_fail != failed) {
                 parse_info(p, "FAIL: should_fail %d, failed %d",
                     should_fail, failed);
@@ -383,6 +407,7 @@ MlkemEncapFileTest(struct encap_ctx *encap)
         }
 
  err:
+        freezero(shared_secret_buf, s_len);
         return failed;
 }
 
@@ -403,48 +428,22 @@ static const struct test_parse encap_parse = {
 };
 
 static int
-mlkem_encap_tests(const char *fn, size_t size)
+mlkem_encap_tests(const char *fn, int rank)
 {
-        struct MLKEM768_public_key public_key768;
-        uint8_t ciphertext768[MLKEM768_CIPHERTEXT_BYTES];
-        struct encap_ctx encap768 = {
-                .public_key = &public_key768,
-                .ciphertext = ciphertext768,
-                .ciphertext_len = sizeof(ciphertext768),
-
-                .parse_public_key = mlkem768_parse_public_key,
-                .encap_external_entropy = mlkem768_encap_external_entropy,
-        };
-        struct MLKEM1024_public_key public_key1024;
-        uint8_t ciphertext1024[MLKEM1024_CIPHERTEXT_BYTES];
-        struct encap_ctx encap1024 = {
-                .public_key = &public_key1024,
-                .ciphertext = ciphertext1024,
-                .ciphertext_len = sizeof(ciphertext1024),
-
-                .parse_public_key = mlkem1024_parse_public_key,
-                .encap_external_entropy = mlkem1024_encap_external_entropy,
+        struct encap_ctx encap = {
+                .rank = rank,
         };
 
-        if (size == 768)
-                return parse_test_file(fn, &encap_parse, &encap768);
-        if (size == 1024)
-                return parse_test_file(fn, &encap_parse, &encap1024);
-
-        errx(1, "unknown encap test: size %zu", size);
+        return parse_test_file(fn, &encap_parse, &encap);
 }
 
 struct keygen_ctx {
         struct parse *parse_ctx;
 
-        void *private_key;
-        void *encoded_public_key;
+        int rank;
+        MLKEM_private_key *private_key;
+        uint8_t *encoded_public_key;
         size_t encoded_public_key_len;
-        size_t private_key_len;
-        size_t public_key_len;
-
-        mlkem_generate_key_external_entropy_fn generate_key_external_entropy;
-        mlkem_marshal_private_key_fn marshal_private_key;
 };
 
 enum keygen_states {
@@ -482,13 +481,19 @@ keygen_init(void *ctx, void *parse_ctx)
 
         keygen->parse_ctx = parse_ctx;
 
-        return 1;
+        return (keygen->private_key = MLKEM_private_key_new(keygen->rank))
+            != NULL;
 }
 
 static void
 keygen_finish(void *ctx)
 {
-        (void)ctx;
+        struct keygen_ctx *keygen = ctx;
+
+        freezero(keygen->encoded_public_key, keygen->encoded_public_key_len);
+        keygen->encoded_public_key = NULL;
+        MLKEM_private_key_free(keygen->private_key);
+        keygen->private_key = NULL;
 }
 
 static int
@@ -504,18 +509,25 @@ MlkemKeygenFileTest(struct keygen_ctx *keygen)
         parse_get_cbs(p, KEYGEN_PUBLIC_KEY, &public_key);
         parse_get_cbs(p, KEYGEN_PRIVATE_KEY, &private_key);
 
-        if (!parse_length_equal(p, "seed", MLKEM_SEED_BYTES, CBS_len(&seed)))
+        if (!parse_length_equal(p, "seed", MLKEM_SEED_LENGTH, CBS_len(&seed)))
                 goto err;
+
+        if (!MLKEM_generate_key_external_entropy(keygen->private_key,
+            &keygen->encoded_public_key, &keygen->encoded_public_key_len,
+            CBS_data(&seed))) {
+                parse_info(p, "generate_key_external_entropy");
+                goto err;
+        }
+
         if (!parse_length_equal(p, "public key",
-            keygen->public_key_len, CBS_len(&public_key)))
+            keygen->encoded_public_key_len, CBS_len(&public_key)))
                 goto err;
         if (!parse_length_equal(p, "private key",
-            keygen->private_key_len, CBS_len(&private_key)))
+            MLKEM_private_key_encoded_length(keygen->private_key),
+            CBS_len(&private_key)))
                 goto err;
 
-        keygen->generate_key_external_entropy(keygen->encoded_public_key,
-            keygen->private_key, CBS_data(&seed));
-        if (!keygen->marshal_private_key(keygen->private_key,
+        if (!MLKEM_marshal_private_key(keygen->private_key,
             &encoded_private_key, &encoded_private_key_len)) {
                 parse_info(p, "encode private key");
                 goto err;
@@ -589,7 +601,7 @@ MlkemNistKeygenFileTest(struct keygen_ctx *keygen)
         struct parse *p = keygen->parse_ctx;
         CBB seed_cbb;
         CBS z, d, ek, dk;
-        uint8_t seed[MLKEM_SEED_BYTES];
+        uint8_t seed[MLKEM_SEED_LENGTH];
         size_t seed_len;
         uint8_t *encoded_private_key = NULL;
         size_t encoded_private_key_len = 0;
@@ -609,12 +621,17 @@ MlkemNistKeygenFileTest(struct keygen_ctx *keygen)
         if (!CBB_finish(&seed_cbb, NULL, &seed_len))
                 parse_errx(p, "CBB_finish");
 
-        if (!parse_length_equal(p, "bogus z or d", MLKEM_SEED_BYTES, seed_len))
+        if (!parse_length_equal(p, "bogus z or d", MLKEM_SEED_LENGTH, seed_len))
                 goto err;
 
-        keygen->generate_key_external_entropy(keygen->encoded_public_key,
-            keygen->private_key, seed);
-        if (!keygen->marshal_private_key(keygen->private_key,
+        if (!MLKEM_generate_key_external_entropy(keygen->private_key,
+            &keygen->encoded_public_key, &keygen->encoded_public_key_len,
+            seed)) {
+                parse_info(p, "generate_key_external_entropy");
+                goto err;
+        }
+
+        if (!MLKEM_marshal_private_key(keygen->private_key,
             &encoded_private_key, &encoded_private_key_len)) {
                 parse_info(p, "encode private key");
                 goto err;
@@ -648,74 +665,49 @@ static const struct test_parse nist_keygen_parse = {
 };
 
 static int
-mlkem_keygen_tests(const char *fn, size_t size, enum test_type test_type)
+mlkem_keygen_tests(const char *fn, int rank, enum test_type test_type)
 {
-        struct MLKEM768_private_key private_key768;
-        uint8_t encoded_public_key768[MLKEM768_PUBLIC_KEY_BYTES];
-        struct keygen_ctx keygen768 = {
-                .private_key = &private_key768,
-                .encoded_public_key = encoded_public_key768,
-                .encoded_public_key_len = sizeof(encoded_public_key768),
-                .private_key_len = MLKEM768_PRIVATE_KEY_BYTES,
-                .public_key_len = MLKEM768_PUBLIC_KEY_BYTES,
-
-                .generate_key_external_entropy =
-                    mlkem768_generate_key_external_entropy,
-                .marshal_private_key =
-                    mlkem768_marshal_private_key,
-        };
-        struct MLKEM1024_private_key private_key1024;
-        uint8_t encoded_public_key1024[MLKEM1024_PUBLIC_KEY_BYTES];
-        struct keygen_ctx keygen1024 = {
-                .private_key = &private_key1024,
-                .encoded_public_key = encoded_public_key1024,
-                .encoded_public_key_len = sizeof(encoded_public_key1024),
-                .private_key_len = MLKEM1024_PRIVATE_KEY_BYTES,
-                .public_key_len = MLKEM1024_PUBLIC_KEY_BYTES,
-
-                .generate_key_external_entropy =
-                    mlkem1024_generate_key_external_entropy,
-                .marshal_private_key =
-                    mlkem1024_marshal_private_key,
+        struct keygen_ctx keygen = {
+                .rank = rank,
         };
+        int ret = 0;
 
-        if (size == 768 && test_type == TEST_TYPE_NORMAL)
-                return parse_test_file(fn, &keygen_parse, &keygen768);
-        if (size == 768 && test_type == TEST_TYPE_NIST)
-                return parse_test_file(fn, &nist_keygen_parse, &keygen768);
-        if (size == 1024 && test_type == TEST_TYPE_NORMAL)
-                return parse_test_file(fn, &keygen_parse, &keygen1024);
-        if (size == 1024 && test_type == TEST_TYPE_NIST)
-                return parse_test_file(fn, &nist_keygen_parse, &keygen1024);
+        if (test_type == TEST_TYPE_NORMAL)
+                ret = parse_test_file(fn, &keygen_parse, &keygen);
+        else if (test_type == TEST_TYPE_NIST)
+                ret = parse_test_file(fn, &nist_keygen_parse, &keygen);
+        else
+                errx(1, "unknown keygen test: rank %d, type %d", rank,
+                    test_type);
 
-        errx(1, "unknown keygen test: size %zu, type %d", size, test_type);
+        return ret;
 }
 
 static int
 run_mlkem_test(const char *test, const char *fn)
 {
         if (strcmp(test, "mlkem768_decap_tests") == 0)
-                return mlkem_decap_tests(fn, 768, TEST_TYPE_NORMAL);
+                return mlkem_decap_tests(fn, RANK768, TEST_TYPE_NORMAL);
         if (strcmp(test, "mlkem768_nist_decap_tests") == 0)
-                return mlkem_decap_tests(fn, 768, TEST_TYPE_NIST);
+                return mlkem_decap_tests(fn, RANK768, TEST_TYPE_NIST);
         if (strcmp(test, "mlkem1024_decap_tests") == 0)
-                return mlkem_decap_tests(fn, 1024, TEST_TYPE_NORMAL);
+                return mlkem_decap_tests(fn, RANK1024, TEST_TYPE_NORMAL);
         if (strcmp(test, "mlkem1024_nist_decap_tests") == 0)
-                return mlkem_decap_tests(fn, 1024, TEST_TYPE_NIST);
+                return mlkem_decap_tests(fn, RANK1024, TEST_TYPE_NIST);
 
         if (strcmp(test, "mlkem768_encap_tests") == 0)
-                return mlkem_encap_tests(fn, 768);
+                return mlkem_encap_tests(fn, RANK768);
         if (strcmp(test, "mlkem1024_encap_tests") == 0)
-                return mlkem_encap_tests(fn, 1024);
+                return mlkem_encap_tests(fn, RANK1024);
 
         if (strcmp(test, "mlkem768_keygen_tests") == 0)
-                return mlkem_keygen_tests(fn, 768, TEST_TYPE_NORMAL);
+                return mlkem_keygen_tests(fn, RANK768, TEST_TYPE_NORMAL);
         if (strcmp(test, "mlkem768_nist_keygen_tests") == 0)
-                return mlkem_keygen_tests(fn, 768, TEST_TYPE_NIST);
+                return mlkem_keygen_tests(fn, RANK768, TEST_TYPE_NIST);
         if (strcmp(test, "mlkem1024_keygen_tests") == 0)
-                return mlkem_keygen_tests(fn, 1024, TEST_TYPE_NORMAL);
+                return mlkem_keygen_tests(fn, RANK1024, TEST_TYPE_NORMAL);
         if (strcmp(test, "mlkem1024_nist_keygen_tests") == 0)
-                return mlkem_keygen_tests(fn, 1024, TEST_TYPE_NIST);
+                return mlkem_keygen_tests(fn, RANK1024, TEST_TYPE_NIST);
 
         errx(1, "unknown test %s (test file %s)", test, fn);
 }
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c b/src/regress/lib/libcrypto/mlkem/mlkem_tests_util.c
index 68bd5d4871..9d6e604386 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.8 2025/05/20 00:33:41 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.c,v 1.9 2025/08/14 15:48:48 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -23,7 +23,7 @@
 #include <string.h>
 
 #include "bytestring.h"
-#include "mlkem.h"
+#include <openssl/mlkem.h>
 
 #include "mlkem_internal.h"
 
@@ -59,157 +59,3 @@ compare_data(const uint8_t *want, const uint8_t *got, size_t len, const char *ms
 
         return 1;
 }
-
-int
-mlkem768_marshal_private_key(const void *private_key, uint8_t **out_buf,
-    size_t *out_len)
-{
-        return MLKEM768_marshal_private_key(private_key, out_buf, out_len);
-}
-
-int
-mlkem768_marshal_public_key(const void *public_key, uint8_t **out_buf,
-    size_t *out_len)
-{
-        return MLKEM768_marshal_public_key(out_buf, out_len, public_key);
-}
-
-int
-mlkem1024_marshal_private_key(const void *private_key, uint8_t **out_buf,
-    size_t *out_len)
-{
-        return MLKEM1024_marshal_private_key(private_key, out_buf, out_len);
-}
-
-int
-mlkem1024_marshal_public_key(const void *public_key, uint8_t **out_buf,
-    size_t *out_len)
-{
-        return MLKEM1024_marshal_public_key(out_buf, out_len, public_key);
-}
-
-int
-mlkem768_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const uint8_t *ciphertext, size_t ciphertext_len, const void *private_key)
-{
-        return MLKEM768_decap(out_shared_secret, ciphertext, ciphertext_len,
-            private_key);
-}
-
-void
-mlkem768_encap(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const void *public_key)
-{
-        MLKEM768_encap(out_ciphertext, out_shared_secret, public_key);
-}
-
-void
-mlkem768_encap_external_entropy(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const void *public_key, const uint8_t entropy[MLKEM_ENCAP_ENTROPY])
-{
-        MLKEM768_encap_external_entropy(out_ciphertext, out_shared_secret,
-            public_key, entropy);
-}
-
-int
-mlkem768_generate_key(uint8_t *out_encoded_public_key,
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key)
-{
-        return MLKEM768_generate_key(out_encoded_public_key, optional_out_seed,
-            out_private_key);
-}
-
-int
-mlkem768_generate_key_external_entropy(uint8_t *out_encoded_public_key,
-    void *out_private_key, const uint8_t entropy[MLKEM_SEED_BYTES])
-{
-        return MLKEM768_generate_key_external_entropy(out_encoded_public_key,
-            out_private_key, entropy);
-}
-
-int
-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,
-            private_key_len);
-}
-
-int
-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,
-            public_key_len);
-}
-
-void
-mlkem768_public_from_private(void *out_public_key, const void *private_key)
-{
-        MLKEM768_public_from_private(out_public_key, private_key);
-}
-
-int
-mlkem1024_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const uint8_t *ciphertext, size_t ciphertext_len, const void *private_key)
-{
-        return MLKEM1024_decap(out_shared_secret, ciphertext, ciphertext_len,
-            private_key);
-}
-
-void
-mlkem1024_encap(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const void *public_key)
-{
-        MLKEM1024_encap(out_ciphertext, out_shared_secret, public_key);
-}
-
-void
-mlkem1024_encap_external_entropy(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const void *public_key, const uint8_t entropy[MLKEM_ENCAP_ENTROPY])
-{
-        MLKEM1024_encap_external_entropy(out_ciphertext, out_shared_secret,
-            public_key, entropy);
-}
-
-int
-mlkem1024_generate_key(uint8_t *out_encoded_public_key,
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_key)
-{
-        return MLKEM1024_generate_key(out_encoded_public_key, optional_out_seed,
-            out_private_key);
-}
-
-int
-mlkem1024_generate_key_external_entropy(uint8_t *out_encoded_public_key,
-    void *out_private_key, const uint8_t entropy[MLKEM_SEED_BYTES])
-{
-        return MLKEM1024_generate_key_external_entropy(out_encoded_public_key,
-            out_private_key, entropy);
-}
-
-int
-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,
-            private_key_len);
-}
-
-void
-mlkem1024_public_from_private(void *out_public_key, const void *private_key)
-{
-        MLKEM1024_public_from_private(out_public_key, private_key);
-}
-
-int
-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,
-            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 1235309f60..a2348c75f3 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.7 2025/05/20 00:33:41 beck Exp $ */
+/*      $OpenBSD: mlkem_tests_util.h,v 1.8 2025/08/14 15:48:48 beck Exp $ */
 /*
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
  * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
@@ -30,60 +30,4 @@
 int compare_data(const uint8_t *want, const uint8_t *got, size_t len,
     const char *msg);
 
-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_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);
-
-int mlkem768_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
-    const uint8_t *ciphertext, size_t ciphertext_len, const void *priv);
-void mlkem768_encap(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES], const void *pub);
-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]);
-int mlkem768_generate_key(uint8_t *out_encoded_public_key,
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_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, 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],
-    const uint8_t *ciphertext, size_t ciphertext_len, const void *priv);
-void mlkem1024_encap(uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES], const void *pub);
-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]);
-int mlkem1024_generate_key(uint8_t *out_encoded_public_key,
-    uint8_t optional_out_seed[MLKEM_SEED_BYTES], void *out_private_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, 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_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 *);
-typedef void (*mlkem_encap_fn)(uint8_t *, uint8_t [MLKEM_SHARED_SECRET_BYTES],
-    const void *);
-typedef void (*mlkem_encap_external_entropy_fn)(uint8_t *,
-    uint8_t [MLKEM_SHARED_SECRET_BYTES], const void *,
-    const uint8_t [MLKEM_ENCAP_ENTROPY]);
-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 *, 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);
-
 #endif /* MLKEM_TEST_UTIL_H */
diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
index adb1c47d8e..417d40555f 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.11 2025/05/21 03:46:20 tb Exp $ */
+/*      $OpenBSD: mlkem_unittest.c,v 1.12 2025/08/14 15:48:48 beck Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -27,95 +27,161 @@
 
 #include "mlkem_tests_util.h"
 
-struct unittest_ctx {
-        void *priv;
-        void *pub;
-        void *priv2;
-        void *pub2;
-        uint8_t *encoded_public_key;
-        size_t encoded_public_key_len;
-        uint8_t *ciphertext;
-        size_t ciphertext_len;
-        mlkem_decap_fn decap;
-        mlkem_encap_fn encap;
-        mlkem_generate_key_fn generate_key;
-        mlkem_parse_private_key_fn parse_private_key;
-        mlkem_parse_public_key_fn parse_public_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;
-};
-
 static int
-MlKemUnitTest(struct unittest_ctx *ctx)
+MlKemUnitTest(int rank)
 {
-        uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES];
-        uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES];
+        MLKEM_private_key *priv = NULL, *priv2 = NULL, *priv3 = NULL;
+        MLKEM_public_key *pub = NULL, *pub2 = NULL, *pub3 = NULL;
+        uint8_t *encoded_public_key = NULL, *ciphertext = NULL,
+            *shared_secret2 = NULL, *shared_secret1 = NULL,
+            *encoded_private_key = NULL, *tmp_buf = NULL, *seed_buf = NULL;
+        size_t encoded_public_key_len, ciphertext_len,
+            encoded_private_key_len, tmp_buf_len;
         uint8_t first_two_bytes[2];
-        uint8_t *encoded_private_key = NULL, *tmp_buf = NULL;
-        size_t encoded_private_key_len, tmp_buf_len;
+        size_t s_len = 0;
         int failed = 0;
 
-        if (!ctx->generate_key(ctx->encoded_public_key, NULL, ctx->priv)) {
+        if ((pub = MLKEM_public_key_new(rank)) == NULL) {
+                warnx("public_key_new");
+                failed |= 1;
+        }
+
+        if ((pub2 = MLKEM_public_key_new(rank)) == NULL) {
+                warnx("public_key_new");
+                failed |= 1;
+        }
+
+        if ((priv = MLKEM_private_key_new(rank)) == NULL) {
+                warnx("private_key_new");
+                failed |= 1;
+        }
+
+        if ((priv2 = MLKEM_private_key_new(rank)) == NULL) {
+                warnx("private_key_new");
+                failed |= 1;
+        }
+
+        if (!MLKEM_generate_key(priv, &encoded_public_key,
+            &encoded_public_key_len, &seed_buf, &s_len)) {
                 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));
+        if (s_len != MLKEM_SEED_LENGTH) {
+                warnx("seed length %zu != %d", s_len, MLKEM_SEED_LENGTH);
+                failed |= 1;
+        }
+
+        if ((priv3 = MLKEM_private_key_new(rank)) == NULL) {
+                warnx("private_key_new");
+                failed |= 1;
+        }
+
+        if ((pub3 = MLKEM_public_key_new(rank)) == NULL) {
+                warnx("public_key_new");
+                failed |= 1;
+        }
+
+        if (!MLKEM_private_key_from_seed(priv3, seed_buf, s_len)) {
+                warnx("private_key_from_seed failed");
+                failed |= 1;
+        }
+
+        free(seed_buf);
+        seed_buf = NULL;
+
+        if (!MLKEM_public_from_private(priv3, pub3)) {
+                warnx("public_from_private");
+                failed |= 1;
+        }
+
+        memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes));
+        memset(encoded_public_key, 0xff, sizeof(first_two_bytes));
 
         /* Parsing should fail because the first coefficient is >= kPrime. */
-        if (ctx->parse_public_key(ctx->pub, ctx->encoded_public_key,
-            ctx->encoded_public_key_len)) {
+        if (MLKEM_parse_public_key(pub, encoded_public_key,
+            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));
-        if (!ctx->parse_public_key(ctx->pub, ctx->encoded_public_key,
-            ctx->encoded_public_key_len)) {
-                warnx("MLKEM768_parse_public_key");
+        memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes));
+
+        MLKEM_public_key_free(pub);
+        if ((pub = MLKEM_public_key_new(rank)) == NULL) {
+                warnx("public_key_new");
+                failed |= 1;
+        }
+        if (!MLKEM_parse_public_key(pub, encoded_public_key,
+            encoded_public_key_len)) {
+                warnx("MLKEM_parse_public_key");
+                failed |= 1;
+        }
+
+        if (!MLKEM_marshal_public_key(pub, &tmp_buf, &tmp_buf_len)) {
+                warnx("marshal_public_key");
+                failed |= 1;
+        }
+        if (encoded_public_key_len != tmp_buf_len) {
+                warnx("encoded public key lengths differ %d != %d",
+                    (int) encoded_public_key_len, (int) tmp_buf_len);
+                failed |= 1;
+        }
+
+        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len,
+            "encoded public keys") != 0) {
+                warnx("compare_data");
                 failed |= 1;
         }
+        free(tmp_buf);
+        tmp_buf = NULL;
+        tmp_buf_len = 0;
 
-        if (!ctx->marshal_public_key(ctx->pub, &tmp_buf, &tmp_buf_len)) {
+        if (!MLKEM_marshal_public_key(pub3, &tmp_buf, &tmp_buf_len)) {
                 warnx("marshal_public_key");
                 failed |= 1;
         }
-        if (ctx->encoded_public_key_len != tmp_buf_len) {
-                warnx("encoded public key lengths differ");
+        if (encoded_public_key_len != tmp_buf_len) {
+                warnx("encoded public key lengths differ %d != %d",
+                    (int) encoded_public_key_len, (int) tmp_buf_len);
                 failed |= 1;
         }
 
-        if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len,
+        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len,
             "encoded public keys") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
         free(tmp_buf);
         tmp_buf = NULL;
+        tmp_buf_len = 0;
 
-        ctx->public_from_private(ctx->pub2, ctx->priv);
-        if (!ctx->marshal_public_key(ctx->pub2, &tmp_buf, &tmp_buf_len)) {
+        if (!MLKEM_public_from_private(priv, pub2)) {
+                warnx("public_from_private");
+                failed |= 1;
+        }
+        if (!MLKEM_marshal_public_key(pub2, &tmp_buf, &tmp_buf_len)) {
                 warnx("marshal_public_key");
                 failed |= 1;
         }
-        if (ctx->encoded_public_key_len != tmp_buf_len) {
-                warnx("encoded public key lengths differ");
+        if (encoded_public_key_len != tmp_buf_len) {
+                warnx("encoded public key lengths differ %d %d",
+                    (int) encoded_public_key_len, (int) tmp_buf_len);
                 failed |= 1;
         }
 
-        if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len,
+        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len,
             "encoded public keys") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
         free(tmp_buf);
         tmp_buf = NULL;
+        tmp_buf_len = 0;
 
-        if (!ctx->marshal_private_key(ctx->priv, &encoded_private_key,
+        if (!MLKEM_marshal_private_key(priv, &encoded_private_key,
             &encoded_private_key_len)) {
-                warnx("mlkem768_encode_private_key");
+                warnx("marshal_private_key");
                 failed |= 1;
         }
 
@@ -123,27 +189,34 @@ MlKemUnitTest(struct unittest_ctx *ctx)
         memset(encoded_private_key, 0xff, sizeof(first_two_bytes));
 
         /*  Parsing should fail because the first coefficient is >= kPrime. */
-        if (ctx->parse_private_key(ctx->priv2, encoded_private_key,
+        if (MLKEM_parse_private_key(priv2, encoded_private_key,
             encoded_private_key_len)) {
-                warnx("MLKEM768_parse_private_key should have failed");
+                warnx("parse_private_key should have failed");
                 failed |= 1;
         }
 
         memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes));
 
-        if (!ctx->parse_private_key(ctx->priv2, encoded_private_key,
+        MLKEM_private_key_free(priv2);
+        priv2 = NULL;
+
+        if ((priv2 = MLKEM_private_key_new(rank)) == NULL) {
+                warnx("private_key_new");
+                failed |= 1;
+        }
+        if (!MLKEM_parse_private_key(priv2, encoded_private_key,
             encoded_private_key_len)) {
-                warnx("MLKEM768_parse_private_key");
+                warnx("parse_private_key");
                 failed |= 1;
         }
 
-        if (!ctx->marshal_private_key(ctx->priv2, &tmp_buf, &tmp_buf_len)) {
-                warnx("encode_private_key");
+        if (!MLKEM_marshal_private_key(priv2, &tmp_buf, &tmp_buf_len)) {
+                warnx("marshal_private_key");
                 failed |= 1;
         }
 
         if (encoded_private_key_len != tmp_buf_len) {
-                warnx("encode private key lengths differ");
+                warnx("encoded private key lengths differ");
                 failed |= 1;
         }
 
@@ -156,106 +229,79 @@ MlKemUnitTest(struct unittest_ctx *ctx)
         free(tmp_buf);
         tmp_buf = NULL;
 
-        ctx->encap(ctx->ciphertext, shared_secret1, ctx->pub);
-        if (!ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len,
-            ctx->priv)) {
+        if (!MLKEM_encap(pub, &ciphertext, &ciphertext_len, &shared_secret1,
+            &s_len)) {
+                warnx("encap failed using pub");
+                failed |= 1;
+        }
+
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                warnx("seed length %zu != %d", s_len,
+                    MLKEM_SHARED_SECRET_LENGTH);
+                failed |= 1;
+        }
+
+        if (!MLKEM_decap(priv, ciphertext, ciphertext_len,
+            &shared_secret2, &s_len)) {
                 warnx("decap() failed using priv");
                 failed |= 1;
         }
-        if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
+
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                warnx("seed length %zu != %d", s_len,
+                    MLKEM_SHARED_SECRET_LENGTH);
+                failed |= 1;
+        }
+
+        if (compare_data(shared_secret1, shared_secret2, s_len,
             "shared secrets with priv") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
 
-        if (!ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len,
-            ctx->priv2)) {
+        free(shared_secret2);
+        shared_secret2 = NULL;
+
+        if (!MLKEM_decap(priv2, ciphertext, ciphertext_len,
+            &shared_secret2, &s_len)){
                 warnx("decap() failed using priv2");
                 failed |= 1;
         }
-        if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
+
+        if (s_len != MLKEM_SHARED_SECRET_LENGTH) {
+                warnx("seed length %zu != %d", s_len,
+                    MLKEM_SHARED_SECRET_LENGTH);
+                failed |= 1;
+        }
+
+        if (compare_data(shared_secret1, shared_secret2, s_len,
             "shared secrets with priv2") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
 
+        MLKEM_public_key_free(pub);
+        MLKEM_public_key_free(pub2);
+        MLKEM_public_key_free(pub3);
+        MLKEM_private_key_free(priv);
+        MLKEM_private_key_free(priv2);
+        MLKEM_private_key_free(priv3);
+        free(encoded_public_key);
+        free(ciphertext);
         free(encoded_private_key);
+        free(shared_secret1);
+        free(shared_secret2);
 
         return failed;
 }
 
-static int
-mlkem768_unittest(void)
-{
-        struct MLKEM768_private_key mlkem768_priv, mlkem768_priv2;
-        struct MLKEM768_public_key mlkem768_pub, mlkem768_pub2;
-        uint8_t mlkem768_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES];
-        uint8_t mlkem768_ciphertext[MLKEM768_CIPHERTEXT_BYTES];
-        struct unittest_ctx mlkem768_test = {
-                .priv = &mlkem768_priv,
-                .pub = &mlkem768_pub,
-                .priv2 = &mlkem768_priv2,
-                .pub2 = &mlkem768_pub2,
-                .encoded_public_key = mlkem768_encoded_public_key,
-                .encoded_public_key_len = sizeof(mlkem768_encoded_public_key),
-                .ciphertext = mlkem768_ciphertext,
-                .ciphertext_len = sizeof(mlkem768_ciphertext),
-                .decap = mlkem768_decap,
-                .encap = mlkem768_encap,
-                .generate_key = mlkem768_generate_key,
-                .parse_private_key = mlkem768_parse_private_key,
-                .parse_public_key = mlkem768_parse_public_key,
-                .marshal_private_key = mlkem768_marshal_private_key,
-                .marshal_public_key = mlkem768_marshal_public_key,
-                .public_from_private = mlkem768_public_from_private,
-        };
-
-        return MlKemUnitTest(&mlkem768_test);
-}
-
-static int
-mlkem1024_unittest(void)
-{
-        struct MLKEM1024_private_key mlkem1024_priv, mlkem1024_priv2;
-        struct MLKEM1024_public_key mlkem1024_pub, mlkem1024_pub2;
-        uint8_t mlkem1024_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES];
-        uint8_t mlkem1024_ciphertext[MLKEM1024_CIPHERTEXT_BYTES];
-        struct unittest_ctx mlkem1024_test = {
-                .priv = &mlkem1024_priv,
-                .pub = &mlkem1024_pub,
-                .priv2 = &mlkem1024_priv2,
-                .pub2 = &mlkem1024_pub2,
-                .encoded_public_key = mlkem1024_encoded_public_key,
-                .encoded_public_key_len = sizeof(mlkem1024_encoded_public_key),
-                .ciphertext = mlkem1024_ciphertext,
-                .ciphertext_len = sizeof(mlkem1024_ciphertext),
-                .decap = mlkem1024_decap,
-                .encap = mlkem1024_encap,
-                .generate_key = mlkem1024_generate_key,
-                .parse_private_key = mlkem1024_parse_private_key,
-                .parse_public_key = mlkem1024_parse_public_key,
-                .marshal_private_key = mlkem1024_marshal_private_key,
-                .marshal_public_key = mlkem1024_marshal_public_key,
-                .public_from_private = mlkem1024_public_from_private,
-        };
-
-        return MlKemUnitTest(&mlkem1024_test);
-}
-
 int
 main(void)
 {
-        int failed = 0;
+        int ret = 0;
 
-        /*
-         * XXX - this is split into two helper functions since having a few
-         * ML-KEM key blobs on the stack makes Emscripten's stack explode,
-         * leading to inscrutable silent failures unless ASAN is enabled.
-         * Go figure.
-         */
+        ret |= MlKemUnitTest(RANK768);
+        ret |= MlKemUnitTest(RANK1024);
 
-        failed |= mlkem768_unittest();
-        failed |= mlkem1024_unittest();
-
-        return failed;
+        return ret;
 }