Rework and fix the mlkem tests

Make proper use of CBB and CBS. If a CBS ever owns data, you're holding
it wrong. Ditch gross macros, sscanf, and globals. The use of fgets is
annoying here, so replace it with getline, which be provided by portable
if needed.

Most importantly, make the tests actually signal failure rather than
only printing an error. Fix the state machines in a few of them. Some
tests didn't parse the .txt file at all. Others mostly did but didn't
actually test what they were supposed to be testing. Such failures
were hidden by the way the tests were written.

This basically needed a complete revamp. It still isn't pretty and much
of it could be deduplicated, but I only have so much time alotted on this
blue planet.

1 files changed, 148 insertions, 73 deletions

diff --git a/src/regress/lib/libcrypto/mlkem/mlkem1024_encap_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem1024_encap_tests.c
index 83924ba6d2..06b00a4b75 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem1024_encap_tests.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem1024_encap_tests.c
@@ -1,7 +1,8 @@
-/*      $OpenBSD: mlkem1024_encap_tests.c,v 1.2 2024/12/14 19:16:24 tb Exp $ */
+/*      $OpenBSD: mlkem1024_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[MLKEM1024_CIPHERTEXT_BYTES];
         struct MLKEM1024_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[MLKEM1024_CIPHERTEXT_BYTES];
+        CBS public_key_cbs;
+        int failed = 1;
 
-        parse_ok = MLKEM1024_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 (!MLKEM1024_parse_public_key(&pub, &public_key_cbs)) {
+                if ((failed = !should_fail))
+                        warnx("#%zu: parse_public_key", line);
+                goto err;
         }
-        MLKEM1024_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),
-            MLKEM1024_CIPHERTEXT_BYTES, "shared_secret");
+        MLKEM1024_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,
+            MLKEM1024_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;
 }