1 files changed, 101 insertions, 182 deletions

diff --git a/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c b/src/regress/lib/libcrypto/mlkem/mlkem_unittest.c
index 18bf128bea..3f7e2886b9 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.4 2024/12/20 00:07:12 tb Exp $ */
+/*      $OpenBSD: mlkem_unittest.c,v 1.5 2024/12/26 00:04:24 tb Exp $ */
 /*
  * Copyright (c) 2024 Google Inc.
  * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
@@ -27,13 +27,28 @@
 
 #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_encode_private_key_fn encode_private_key;
+        mlkem_encode_public_key_fn encode_public_key;
+        mlkem_public_from_private_fn public_from_private;
+};
+
 static int
-MlKem768UnitTest(void)
+MlKemUnitTest(struct unittest_ctx *ctx)
 {
-        struct MLKEM768_private_key priv = { 0 }, priv2 = { 0 };
-        struct MLKEM768_public_key pub = { 0 }, pub2 = { 0 };
-        uint8_t encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES];
-        uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES];
         uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES];
         uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES];
         uint8_t first_two_bytes[2];
@@ -42,22 +57,22 @@ MlKem768UnitTest(void)
         CBS cbs;
         int failed = 0;
 
-        MLKEM768_generate_key(encoded_public_key, NULL, &priv);
+        ctx->generate_key(ctx->encoded_public_key, NULL, ctx->priv);
 
-        memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes));
-        memset(encoded_public_key, 0xff, sizeof(first_two_bytes));
+        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, encoded_public_key, sizeof(encoded_public_key));
+        CBS_init(&cbs, ctx->encoded_public_key, ctx->encoded_public_key_len);
 
         /* Parsing should fail because the first coefficient is >= kPrime. */
-        if (MLKEM768_parse_public_key(&pub, &cbs)) {
-                warnx("MLKEM768_parse_public_key should have failed");
+        if (ctx->parse_public_key(ctx->pub, &cbs)) {
+                warnx("parse_public_key should have failed");
                 failed |= 1;
         }
 
-        memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes));
-        CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key));
-        if (!MLKEM768_parse_public_key(&pub, &cbs)) {
+        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)) {
                 warnx("MLKEM768_parse_public_key");
                 failed |= 1;
         }
@@ -67,16 +82,16 @@ MlKem768UnitTest(void)
                 failed |= 1;
         }
 
-        if (!mlkem768_encode_public_key(&pub, &tmp_buf, &tmp_buf_len)) {
+        if (!ctx->encode_public_key(ctx->pub, &tmp_buf, &tmp_buf_len)) {
                 warnx("encode_public_key");
                 failed |= 1;
         }
-        if (sizeof(encoded_public_key) != tmp_buf_len) {
-                warnx("mlkem768 encoded public key lengths differ");
+        if (ctx->encoded_public_key_len != tmp_buf_len) {
+                warnx("encoded public key lengths differ");
                 failed |= 1;
         }
 
-        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 768,
+        if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len,
             "encoded public keys") != 0) {
                 warnx("compare_data");
                 failed |= 1;
@@ -84,17 +99,17 @@ MlKem768UnitTest(void)
         free(tmp_buf);
         tmp_buf = NULL;
 
-        MLKEM768_public_from_private(&pub2, &priv);
-        if (!mlkem768_encode_public_key(&pub2, &tmp_buf, &tmp_buf_len)) {
-                warnx("mlkem768_encode_public_key");
+        ctx->public_from_private(ctx->pub2, ctx->priv);
+        if (!ctx->encode_public_key(ctx->pub2, &tmp_buf, &tmp_buf_len)) {
+                warnx("encode_public_key");
                 failed |= 1;
         }
-        if (sizeof(encoded_public_key) != tmp_buf_len) {
-                warnx("mlkem768 encoded public key lengths differ");
+        if (ctx->encoded_public_key_len != tmp_buf_len) {
+                warnx("encoded public key lengths differ");
                 failed |= 1;
         }
 
-        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 768,
+        if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len,
             "encoded public keys") != 0) {
                 warnx("compare_data");
                 failed |= 1;
@@ -102,7 +117,7 @@ MlKem768UnitTest(void)
         free(tmp_buf);
         tmp_buf = NULL;
 
-        if (!mlkem768_encode_private_key(&priv, &encoded_private_key,
+        if (!ctx->encode_private_key(ctx->priv, &encoded_private_key,
             &encoded_private_key_len)) {
                 warnx("mlkem768_encode_private_key");
                 failed |= 1;
@@ -113,7 +128,7 @@ MlKem768UnitTest(void)
         CBS_init(&cbs, encoded_private_key, encoded_private_key_len);
 
         /*  Parsing should fail because the first coefficient is >= kPrime. */
-        if (MLKEM768_parse_private_key(&priv2, &cbs)) {
+        if (ctx->parse_private_key(ctx->priv2, &cbs)) {
                 warnx("MLKEM768_parse_private_key should have failed");
                 failed |= 1;
         }
@@ -121,162 +136,22 @@ MlKem768UnitTest(void)
         memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes));
         CBS_init(&cbs, encoded_private_key, encoded_private_key_len);
 
-        if (!MLKEM768_parse_private_key(&priv2, &cbs)) {
+        if (!ctx->parse_private_key(ctx->priv2, &cbs)) {
                 warnx("MLKEM768_parse_private_key");
                 failed |= 1;
         }
 
-        if (!mlkem768_encode_private_key(&priv2, &tmp_buf, &tmp_buf_len)) {
-                warnx("mlkem768_encode_private_key");
-                failed |= 1;
-        }
-
-        if (encoded_private_key_len != tmp_buf_len) {
-                warnx("mlkem768 encode private key lengths differ");
-                failed |= 1;
-        }
-
-        if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len, 768,
-            "encoded private key") != 0) {
-                warnx("compare_data");
-                failed |= 1;
-        }
-
-        free(tmp_buf);
-        tmp_buf = NULL;
-
-        MLKEM768_encap(ciphertext, shared_secret1, &pub);
-        MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES,
-            &priv);
-        if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
-            768, "shared secrets with priv") != 0) {
-                warnx("compare_data");
-                failed |= 1;
-        }
-
-        MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES,
-            &priv2);
-        if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
-            768, "shared secrets with priv2") != 0) {
-                warnx("compare_data");
-                failed |= 1;
-        }
-
-        free(encoded_private_key);
-
-        return failed;
-}
-
-static int
-MlKem1024UnitTest(void)
-{
-        struct MLKEM1024_private_key priv = { 0 }, priv2 = { 0 };
-        struct MLKEM1024_public_key pub = { 0 }, pub2 = { 0 };
-        uint8_t encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES];
-        uint8_t ciphertext[MLKEM1024_CIPHERTEXT_BYTES];
-        uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES];
-        uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES];
-        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;
-
-        MLKEM1024_generate_key(encoded_public_key, NULL, &priv);
-
-        memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes));
-        memset(encoded_public_key, 0xff, sizeof(first_two_bytes));
-
-        CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key));
-
-        /* Parsing should fail because the first coefficient is >= kPrime. */
-        if (MLKEM1024_parse_public_key(&pub, &cbs)) {
-                warnx("MLKEM1024_parse_public_key should have failed");
-                failed |= 1;
-        }
-
-        memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes));
-        CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key));
-        if (!MLKEM1024_parse_public_key(&pub, &cbs)) {
-                warnx("MLKEM1024_parse_public_key");
-                failed |= 1;
-        }
-
-        if (CBS_len(&cbs) != 0u) {
-                warnx("CBS_len must be 0");
-                failed |= 1;
-        }
-
-        if (!mlkem1024_encode_public_key(&pub, &tmp_buf, &tmp_buf_len)) {
-                warnx("encode_public_key");
-                failed |= 1;
-        }
-        if (sizeof(encoded_public_key) != tmp_buf_len) {
-                warnx("mlkem1024 encoded public key lengths differ");
-                failed |= 1;
-        }
-
-        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 1024,
-            "encoded public keys") != 0) {
-                warnx("compare_data");
-                failed |= 1;
-        }
-        free(tmp_buf);
-        tmp_buf = NULL;
-
-        MLKEM1024_public_from_private(&pub2, &priv);
-        if (!mlkem1024_encode_public_key(&pub2, &tmp_buf, &tmp_buf_len)) {
-                warnx("mlkem1024_encode_public_key");
-                failed |= 1;
-        }
-        if (sizeof(encoded_public_key) != tmp_buf_len) {
-                warnx("mlkem1024 encoded public key lengths differ");
-                failed |= 1;
-        }
-
-        if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 1024,
-            "encoded public keys") != 0) {
-                warnx("compare_data");
-                failed |= 1;
-        }
-        free(tmp_buf);
-        tmp_buf = NULL;
-
-        if (!mlkem1024_encode_private_key(&priv, &encoded_private_key,
-            &encoded_private_key_len)) {
-                warnx("mlkem1024_encode_private_key");
-                failed |= 1;
-        }
-
-        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 (MLKEM1024_parse_private_key(&priv2, &cbs)) {
-                warnx("MLKEM1024_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 (!MLKEM1024_parse_private_key(&priv2, &cbs)) {
-                warnx("MLKEM1024_parse_private_key");
-                failed |= 1;
-        }
-
-        if (!mlkem1024_encode_private_key(&priv2, &tmp_buf, &tmp_buf_len)) {
-                warnx("mlkem1024_encode_private_key");
+        if (!ctx->encode_private_key(ctx->priv2, &tmp_buf, &tmp_buf_len)) {
+                warnx("encode_private_key");
                 failed |= 1;
         }
 
         if (encoded_private_key_len != tmp_buf_len) {
-                warnx("mlkem1024 encode private key lengths differ");
+                warnx("encode private key lengths differ");
                 failed |= 1;
         }
 
-        if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len, 1024,
+        if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len,
             "encoded private key") != 0) {
                 warnx("compare_data");
                 failed |= 1;
@@ -285,19 +160,19 @@ MlKem1024UnitTest(void)
         free(tmp_buf);
         tmp_buf = NULL;
 
-        MLKEM1024_encap(ciphertext, shared_secret1, &pub);
-        MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES,
-            &priv);
+        ctx->encap(ctx->ciphertext, shared_secret1, ctx->pub);
+        ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len,
+            ctx->priv);
         if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
-            1024, "shared secrets with priv") != 0) {
+            "shared secrets with priv") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
 
-        MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES,
-            &priv2);
+        ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len,
+            ctx->priv2);
         if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES,
-            1024, "shared secrets with priv2") != 0) {
+            "shared secrets with priv2") != 0) {
                 warnx("compare_data");
                 failed |= 1;
         }
@@ -308,12 +183,56 @@ MlKem1024UnitTest(void)
 }
 
 int
-main(int argc, char **argv)
+main(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,
+                .encode_private_key = mlkem768_encode_private_key,
+                .encode_public_key = mlkem768_encode_public_key,
+                .public_from_private = mlkem768_public_from_private,
+        };
+        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,
+                .encode_private_key = mlkem1024_encode_private_key,
+                .encode_public_key = mlkem1024_encode_public_key,
+                .public_from_private = mlkem1024_public_from_private,
+        };
         int failed = 0;
 
-        failed |= MlKem768UnitTest();
-        failed |= MlKem1024UnitTest();
+        failed |= MlKemUnitTest(&mlkem768_test);
+        failed |= MlKemUnitTest(&mlkem1024_test);
 
         return failed;
 }