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, 129 insertions, 69 deletions

diff --git a/src/regress/lib/libcrypto/mlkem/mlkem768_keygen_tests.c b/src/regress/lib/libcrypto/mlkem/mlkem768_keygen_tests.c
index 7214dd8a80..be1aff3d04 100644
--- a/src/regress/lib/libcrypto/mlkem/mlkem768_keygen_tests.c
+++ b/src/regress/lib/libcrypto/mlkem/mlkem768_keygen_tests.c
@@ -1,7 +1,8 @@
-/*      $OpenBSD: mlkem768_keygen_tests.c,v 1.4 2024/12/17 07:20:10 tb Exp $ */
+/*      $OpenBSD: mlkem768_keygen_tests.c,v 1.5 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,115 +17,174 @@
  * 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 int
-encode_private_key(const struct MLKEM768_private_key *priv, uint8_t **out_buf,
-    size_t *out_len)
-{
-        CBB cbb;
-        if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES))
-                return 0;
-        if (!MLKEM768_marshal_private_key(&cbb, priv))
-                return 0;
-        if (!CBB_finish(&cbb, out_buf, out_len))
-                return 0;
-        CBB_cleanup(&cbb);
-        return 1;
-}
-
-static void
-MlkemKeygenFileTest(CBS *seed, CBS *public_key, CBS *private_key)
+MlkemKeygenFileTest(CBB *seed_cbb, CBB *public_key_cbb, CBB *private_key_cbb,
+    size_t line)
 {
         struct MLKEM768_private_key priv;
+        uint8_t *seed = NULL, *public_key = NULL, *private_key = NULL;
+        size_t seed_len = 0, public_key_len = 0, private_key_len = 0;
         uint8_t *encoded_private_key = NULL;
         uint8_t encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES];
         size_t len;
+        int failed = 1;
+
+        if (!CBB_finish(seed_cbb, &seed, &seed_len))
+                goto err;
+        if (!compare_length(MLKEM_SEED_BYTES, seed_len, line, "seed length"))
+                goto err;
+        if (!CBB_finish(public_key_cbb, &public_key, &public_key_len))
+                goto err;
+        if (!compare_length(MLKEM768_PUBLIC_KEY_BYTES, public_key_len, line,
+            "public key length"))
+                goto err;
+        if (!CBB_finish(private_key_cbb, &private_key, &private_key_len))
+                goto err;
+        if (!compare_length(MLKEM768_PUBLIC_KEY_BYTES, public_key_len, line,
+            "public key length"))
+                goto err;
 
-        TEST(CBS_len(seed) != MLKEM_SEED_BYTES, "seed len bogus");
-        TEST(CBS_len(private_key) != MLKEM768_PRIVATE_KEY_BYTES,
-            "expected private key len bogus");
-        TEST(CBS_len(public_key) != MLKEM768_PUBLIC_KEY_BYTES,
-            "expected public key len bogus");
         MLKEM768_generate_key_external_entropy(encoded_public_key, &priv,
-            CBS_data(seed));
-        TEST(!encode_private_key(&priv, &encoded_private_key,
-            &len), "encode_private_key");
-        TEST(len != MLKEM768_PRIVATE_KEY_BYTES, "private key len bogus");
-        TEST_DATAEQ(encoded_public_key, CBS_data(public_key),
-            MLKEM768_PUBLIC_KEY_BYTES, "public key");
-        TEST_DATAEQ(encoded_private_key, CBS_data(private_key),
-            MLKEM768_PRIVATE_KEY_BYTES, "private key");
+            seed);
+        if (!mlkem768_encode_private_key(&priv, &encoded_private_key, &len)) {
+                warnx("#%zu: encoded_private_key", line);
+                goto err;
+        }
+
+        if (!compare_length(MLKEM768_PRIVATE_KEY_BYTES, len, line,
+            "private key length"))
+                goto err;
+
+        failed = compare_data(private_key, encoded_private_key,
+            MLKEM768_PRIVATE_KEY_BYTES, line, "private key");
+        failed |= compare_data(public_key, encoded_public_key,
+            MLKEM768_PUBLIC_KEY_BYTES, line, "public key");
+
+ err:
+        CBB_cleanup(seed_cbb);
+        CBB_cleanup(public_key_cbb);
+        CBB_cleanup(private_key_cbb);
+        freezero(seed, seed_len);
+        freezero(public_key, public_key_len);
+        freezero(private_key, private_key_len);
         free(encoded_private_key);
+
+        return failed;
 }
 
-#define S_START 0
-#define S_SEED 1
-#define S_PUBLIC_KEY  2
-#define S_PRIVATE_KEY   3
+#define S_START         0
+#define S_COMMENT       1
+#define S_SEED          2
+#define S_PUBLIC_KEY    3
+#define S_PRIVATE_KEY   4
+
+#define S2S(x) case x: return #x
+
+static const char *
+state2str(int state)
+{
+        switch (state) {
+        S2S(S_START);
+        S2S(S_COMMENT);
+        S2S(S_SEED);
+        S2S(S_PUBLIC_KEY);
+        S2S(S_PRIVATE_KEY);
+        default:
+                errx(1, "unknown state %d", state);
+        }
+}
 
 int
 main(int argc, char **argv)
 {
-        CBS seed, public_key, private_key;
-        char *buf;
+        CBB seed = { 0 }, public_key = { 0 }, private_key = { 0 };
+        const char *test;
+        size_t line = 0;
+        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];
+
+        if ((fp = fopen(test, "r")) == NULL)
+                err(1, "cant't open test file");
+
+        state = S_COMMENT;
+        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');
 
-        fprintf(stderr, "Testing keygen 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_SEED;
-        test_number = 1;
-        while (fgets(buf, 16*1024, fp) != NULL) {
                 switch (state) {
                 case S_START:
-                        if (strcmp(buf, "\n") != 0)
-                                break;
+                        state = S_COMMENT;
+                        break;
+                case S_COMMENT:
+                        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_SEED;
                         break;
                 case S_SEED:
-                        if (strncmp(buf, "seed: ", strlen("seed: ")) != 0) {
-                                break;
-                        }
-                        grab_data(&seed, buf, strlen("seed: "));
+                        if (!get_string_cbs(&cbs, "seed: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &seed, 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: "));
+                        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_PRIVATE_KEY;
                         break;
                 case S_PRIVATE_KEY:
-                        if (strncmp(buf, "private_key: ",
-                            strlen("private_key: ")) != 0)
-                                break;
-                        grab_data(&private_key, buf, strlen("private_key: "));
-                        state = S_START;
+                        if (!get_string_cbs(&cbs, "private_key: ", line, msg))
+                                errx(1, "#%zu %s: get_string_cbs", line, msg);
+                        hex_decode_cbs(&cbs, &private_key, line, msg);
 
-                        MlkemKeygenFileTest(&seed, &public_key, &private_key);
-                        free((void *)CBS_data(&seed));
-                        free((void *)CBS_data(&public_key));
-                        free((void *)CBS_data(&private_key));
+                        failed |= MlkemKeygenFileTest(&seed, &public_key,
+                            &private_key, line);
 
-                        test_number++;
                         state = S_START;
                         break;
                 }
+                if (CBS_len(&cbs) > 0)
+                        errx(1, "#%zu %s: CBS_len", line, msg);
         }
-
         free(buf);
+
+        if (ferror(fp))
+                err(1, NULL);
         fclose(fp);
-        exit(failure);
+
+        return failed;
 }