1 files changed, 148 insertions, 73 deletions

diff --git a/src/regress/lib/libcrypto/mlkem/mlkem768_encap_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem768_encap_tests.c
index be6c6149da..55e3fe66bb 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem768_encap_tests.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem768_encap_tests.c
@@ -1,7 +1,8 @@
-/*      $OpenBSD: mlkem768_encap_tests.c,v 1.2 2024/12/14 19:16:24 tb Exp $ */
+/*      $OpenBSD: mlkem768_encap_tests.c,v 1.3 2024/12/20 00:07:12 tb Exp $ */
 /*
- * Copyright (c) 2024, Google Inc.
- * Copyright (c) 2024, Bob Beck <beck@obtuse.com>
+ * Copyright (c) 2024 Google Inc.
+ * Copyright (c) 2024 Bob Beck <beck@obtuse.com>
+ * Copyright (c) 2024 Theo Buehler <tb@openbsd.org>
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -16,120 +17,194 @@
  * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  */
 
+#include <assert.h>
+#include <err.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <string.h>
 
-#include <openssl/bytestring.h>
-#include <openssl/mlkem.h>
+#include "bytestring.h"
+#include "mlkem.h"
 
+#include "mlkem_internal.h"
 #include "mlkem_tests_util.h"
 
-static void
-MlkemEncapFileTest(CBS *entropy, CBS *public_key, CBS *expected_ciphertext,
-    CBS *expected_shared_secret, int should_fail)
+static int
+MlkemEncapFileTest(CBB *entropy_cbb, CBB *pubkey_cbb, CBB *ciphertext_cbb,
+    CBB *shared_secret_cbb, int should_fail, size_t line)
 {
-        uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
-        uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES];
         struct MLKEM768_public_key pub;
-        int parse_ok;
+        uint8_t *entropy = NULL, *public_key = NULL, *ciphertext = NULL;
+        uint8_t *shared_secret = NULL;
+        size_t entropy_len = 0, public_key_len = 0, ciphertext_len = 0;
+        size_t shared_secret_len = 0;
+        uint8_t shared_secret_buf[MLKEM_SHARED_SECRET_BYTES];
+        uint8_t ciphertext_buf[MLKEM768_CIPHERTEXT_BYTES];
+        CBS public_key_cbs;
+        int failed = 1;
 
-        parse_ok = MLKEM768_parse_public_key(&pub, public_key);
-        if (!parse_ok) {
-                TEST(!should_fail, "parse_public_key");
-                return;
+        if (!CBB_finish(entropy_cbb, &entropy, &entropy_len))
+                goto err;
+        if (!CBB_finish(pubkey_cbb, &public_key, &public_key_len))
+                goto err;
+        if (!CBB_finish(ciphertext_cbb, &ciphertext, &ciphertext_len))
+                goto err;
+        if (!CBB_finish(shared_secret_cbb, &shared_secret, &shared_secret_len))
+                goto err;
+
+        CBS_init(&public_key_cbs, public_key, public_key_len);
+
+        if (!MLKEM768_parse_public_key(&pub, &public_key_cbs)) {
+                if ((failed = !should_fail))
+                        warnx("#%zu: parse_public_key", line);
+                goto err;
         }
-        MLKEM768_encap(ciphertext, shared_secret, &pub);
-        TEST_DATAEQ(shared_secret, CBS_data(expected_shared_secret),
-            MLKEM_SHARED_SECRET_BYTES, "shared_secret");
-        TEST_DATAEQ(ciphertext, CBS_data(expected_ciphertext),
-            MLKEM768_CIPHERTEXT_BYTES, "shared_secret");
+        MLKEM768_encap_external_entropy(ciphertext_buf, shared_secret_buf,
+            &pub, entropy);
+
+        failed = compare_data(shared_secret, shared_secret_buf,
+            MLKEM_SHARED_SECRET_BYTES, line, "shared_secret");
+        failed |= compare_data(ciphertext, ciphertext_buf,
+            MLKEM768_CIPHERTEXT_BYTES, line, "ciphertext");
+
+        if (should_fail != failed) {
+                warnx("FAIL: #%zu: should_fail %d, failed %d",
+                    line, should_fail, failed);
+                failed = 1;
+        }
+
+ err:
+        CBB_cleanup(entropy_cbb);
+        CBB_cleanup(pubkey_cbb);
+        CBB_cleanup(ciphertext_cbb);
+        CBB_cleanup(shared_secret_cbb);
+        freezero(entropy, entropy_len);
+        freezero(public_key, public_key_len);
+        freezero(ciphertext, ciphertext_len);
+        freezero(shared_secret, shared_secret_len);
+
+        return failed;
 }
 
-#define S_START 0
-#define S_COMMENT 1
-#define S_ENTROPY 2
-#define S_PUBLIC_KEY 3
-#define S_RESULT 4
-#define S_CIPHERTEXT 5
-#define S_SHARED_SECRET  6
+#define S_START                 0
+#define S_COMMENT               1
+#define S_ENTROPY               2
+#define S_PUBLIC_KEY            3
+#define S_RESULT                4
+#define S_CIPHERTEXT            5
+#define S_SHARED_SECRET         6
+
+#define S2S(x) case x: return #x
+
+static const char *
+state2str(int state)
+{
+        switch (state) {
+        S2S(S_START);
+        S2S(S_COMMENT);
+        S2S(S_ENTROPY);
+        S2S(S_PUBLIC_KEY);
+        S2S(S_RESULT);
+        S2S(S_CIPHERTEXT);
+        S2S(S_SHARED_SECRET);
+        default:
+                errx(1, "unknown state %d", state);
+        }
+}
 
 int
 main(int argc, char **argv)
 {
-        CBS entropy, public_key, ciphertext, shared_secret;
-        const uint8_t *p = NULL;
+        CBB entropy = { 0 }, public_key = { 0 }, ciphertext = { 0 }, shared_secret = { 0 };
         int should_fail = 0;
-        char *buf;
+        const char *test;
+        size_t line;
+        char *buf = NULL;
+        size_t buflen = 0;
+        ssize_t len;
         FILE *fp;
         int state;
+        int failed = 0;
+
+        if (argc < 2)
+                errx(1, "%s: missing test file", argv[0]);
+
+        test = argv[1];
+        line = 0;
+
+        if ((fp = fopen(test, "r")) == NULL)
+                err(1, "cant't open test file");
 
-        fprintf(stderr, "Testing encap test vectors in %s\n", argv[1]);
-        TEST((fp = fopen(argv[1], "r")) == NULL, "can't open test file");
-        MALLOC(buf, 16*1024);
         state = S_COMMENT;
-        test_number = 1;
-        while (fgets(buf, 16*1024, fp) != NULL) {
+        line = 0;
+
+        while ((len = getline(&buf, &buflen, fp)) != -1) {
+                const char *msg = state2str(state);
+                CBS cbs;
+                uint8_t u8;
+
+                line++;
+                CBS_init(&cbs, buf, len);
+
+                if (!CBS_get_last_u8(&cbs, &u8))
+                        errx(1, "#%zu %s: CBB_get_last_u8", line, msg);
+                assert(u8 == '\n');
+
                 switch (state) {
                 case S_START:
-                        if (strcmp(buf, "\n") != 0)
-                                break;
                         state = S_COMMENT;
                         break;
                 case S_COMMENT:
-                        if (strncmp(buf, "#", 1) != 0)
-                                break;
+                        if (!CBS_get_u8(&cbs, &u8))
+                                errx(1, "#%zu %s: CBB_get_u8", line, msg);
+                        assert(u8 == '#');
+                        if (!CBS_skip(&cbs, CBS_len(&cbs)))
+                                errx(1, "#%zu %s: CBB_skip", line, msg);
                         state = S_ENTROPY;
                         break;
                 case S_ENTROPY:
-                        if (strncmp(buf, "entropy: ", strlen("entropy: ")) != 0)
-                                break;
-                        grab_data(&entropy, buf, strlen("entropy: "));
-                        p = CBS_data(&entropy);
+                        if (!get_string_cbs(&cbs, "entropy: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &entropy, line, msg);
                         state = S_PUBLIC_KEY;
                         break;
                 case S_PUBLIC_KEY:
-                        if (strncmp(buf, "public_key: ",
-                            strlen("public_key: ")) != 0)
-                                break;
-                        grab_data(&public_key, buf, strlen("public_key: "));
-                        p = CBS_data(&public_key);
+                        if (!get_string_cbs(&cbs, "public_key = ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &public_key, line, msg);
                         state = S_RESULT;
                         break;
                 case S_RESULT:
-                        if (strncmp(buf, "result: pass",
-                            strlen("result: pass")) != 0)
-                                should_fail = 1;
-                        else
-                                should_fail = 0;
+                        if (!get_string_cbs(&cbs, "result: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        should_fail = get_string_cbs(&cbs, "fail", line, msg);
                         state = S_CIPHERTEXT;
                         break;
                 case S_CIPHERTEXT:
-                        if (strncmp(buf, "ciphertext: ",
-                            strlen("ciphertext: ")) != 0)
-                                break;
-                        grab_data(&ciphertext, buf, strlen("ciphertext: "));
-                        state = S_RESULT;
+                        if (!get_string_cbs(&cbs, "ciphertext: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &ciphertext, line, msg);
+                        state = S_SHARED_SECRET;
                         break;
                 case S_SHARED_SECRET:
-                        if (strncmp(buf, "shared_secret: ",
-                            strlen("shared_secret: ")) != 0)
-                                break;
-                        grab_data(&shared_secret, buf,
-                            strlen("shared_secret: "));
-                        MlkemEncapFileTest(&entropy, &public_key, &ciphertext,
-                            &shared_secret, should_fail);
-                        free((void *)CBS_data(&ciphertext));
-                        free((void *)CBS_data(&shared_secret));
-                        free((void *)p);
-
-                        test_number++;
+                        if (!get_string_cbs(&cbs, "shared_secret: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &shared_secret, line, msg);
+
+                        failed |= MlkemEncapFileTest(&entropy, &public_key,
+                            &ciphertext, &shared_secret, should_fail, line);
+
                         state = S_START;
                         break;
                 }
+                if (CBS_len(&cbs) > 0)
+                        errx(1, "#%zu %s: CBS_len", line, msg);
         }
-
         free(buf);
-        exit(failure);
+
+        if (ferror(fp))
+                err(1, NULL);
+        fclose(fp);
+
+        return failed;
 }