1 files changed, 131 insertions, 72 deletions

diff --git a/regress/lib/libcrypto/mlkem/mlkem1024_nist_keygen_tests.c b/regress/lib/libcrypto/mlkem/mlkem1024_nist_keygen_tests.c
index d293d121d46..e84f7dafc6c 100644
--- a/regress/lib/libcrypto/mlkem/mlkem1024_nist_keygen_tests.c
+++ b/regress/lib/libcrypto/mlkem/mlkem1024_nist_keygen_tests.c
@@ -1,7 +1,8 @@
-/*	$OpenBSD: mlkem1024_nist_keygen_tests.c,v 1.3 2024/12/17 07:20:10 tb Exp $ */
+/*	$OpenBSD: mlkem1024_nist_keygen_tests.c,v 1.4 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,123 +17,181 @@
  * 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 MLKEM1024_private_key *priv, uint8_t **out_buf,
-    size_t *out_len)
-{
-	CBB cbb;
-	if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES))
-		return 0;
-	if (!MLKEM1024_marshal_private_key(&cbb, priv))
-		return 0;
-	if (!CBB_finish(&cbb, out_buf, out_len))
-		return 0;
-	CBB_cleanup(&cbb);
-	return 1;
-}
-
-static void
-MlkemNistKeygenFileTest(CBS *z, CBS *d, CBS *ek, CBS *dk)
+MlkemNistKeygenFileTest(CBB *z_cbb, CBB *d_cbb, CBB *ek_cbb, CBB *dk_cbb,
+    size_t line)
 {
+	CBB seed_cbb;
+	uint8_t *z = NULL, *d = NULL, *ek = NULL, *dk = NULL;
+	size_t z_len = 0, d_len = 0, ek_len = 0, dk_len = 0;
 	uint8_t seed[MLKEM_SEED_BYTES];
 	struct MLKEM1024_private_key priv;
 	uint8_t *encoded_private_key = NULL;
 	uint8_t encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES];
 	size_t len;
+	int failed = 1;
+
+	if (!CBB_init_fixed(&seed_cbb, seed, sizeof(seed)))
+		goto err;
+
+	if (!CBB_finish(z_cbb, &z, &z_len))
+		goto err;
+	if (!CBB_finish(d_cbb, &d, &d_len))
+		goto err;
+	if (!CBB_finish(ek_cbb, &ek, &ek_len))
+		goto err;
+	if (!CBB_finish(dk_cbb, &dk, &dk_len))
+		goto err;
+
+	if (!CBB_add_bytes(&seed_cbb, d, d_len))
+		goto err;
+	if (!CBB_add_bytes(&seed_cbb, z, z_len))
+		goto err;
+	if (!CBB_finish(&seed_cbb, NULL, &len))
+		goto err;
+
+	if (!compare_length(MLKEM_SEED_BYTES, len, line, "z or d length bogus"))
+		goto err;
 
-	TEST(CBS_len(d) != (MLKEM_SEED_BYTES / 2), "d len bogus");
-	TEST(CBS_len(z) != (MLKEM_SEED_BYTES / 2), "z len bogus");
-	TEST(CBS_len(dk) != MLKEM1024_PRIVATE_KEY_BYTES,
-	    "expected private key len bogus");
-	TEST(CBS_len(ek) != MLKEM1024_PUBLIC_KEY_BYTES,
-	    "expected public key len bogus");
-	memcpy(&seed[0], CBS_data(d), CBS_len(d));
-	memcpy(&seed[MLKEM_SEED_BYTES / 2], CBS_data(z), CBS_len(z));
 	MLKEM1024_generate_key_external_entropy(encoded_public_key, &priv, seed);
-	TEST(!encode_private_key(&priv, &encoded_private_key,
-	    &len), "encode_private_key");
-	TEST(len != MLKEM1024_PRIVATE_KEY_BYTES, "private key len bogus");
-	TEST_DATAEQ(encoded_public_key, CBS_data(ek),
-	    MLKEM1024_PUBLIC_KEY_BYTES, "public key");
-	TEST_DATAEQ(encoded_private_key, CBS_data(dk),
-	    MLKEM1024_PRIVATE_KEY_BYTES, "private key");
+
+	if (!mlkem1024_encode_private_key(&priv, &encoded_private_key, &len)) {
+		warnx("#%zu mlkem1024_encode_private_key", line);
+		goto err;
+	}
+
+	if (!compare_length(MLKEM1024_PRIVATE_KEY_BYTES, len, line,
+	    "private key length"))
+		goto err;
+
+	failed = compare_data(ek, encoded_public_key, MLKEM1024_PUBLIC_KEY_BYTES,
+	    line, "public key");
+	failed |= compare_data(dk, encoded_private_key, MLKEM1024_PRIVATE_KEY_BYTES,
+	    line, "private key");
+
+ err:
+	CBB_cleanup(&seed_cbb);
+	CBB_cleanup(z_cbb);
+	CBB_cleanup(d_cbb);
+	CBB_cleanup(ek_cbb);
+	CBB_cleanup(dk_cbb);
+	freezero(z, z_len);
+	freezero(d, d_len);
+	freezero(ek, ek_len);
+	freezero(dk, dk_len);
 	free(encoded_private_key);
+
+	return failed;
 }
 
-#define S_START 0
-#define S_Z 1
-#define S_D  2
-#define S_EK	3
-#define S_DK	4
+#define S_START		0
+#define S_Z		1
+#define S_D		2
+#define S_EK		3
+#define S_DK		4
+
+#define S2S(x) case x: return #x
+
+static const char *
+state2str(int state)
+{
+	switch (state) {
+	S2S(S_START);
+	S2S(S_Z);
+	S2S(S_D);
+	S2S(S_EK);
+	S2S(S_DK);
+	default:
+		errx(1, "unknown state %d", state);
+	}
+}
 
 int
 main(int argc, char **argv)
 {
-	CBS z, d, ek, dk;
-	char *buf;
+	CBB z = { 0 }, d = { 0 }, ek = { 0 }, dk = { 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");
 
-	fprintf(stderr, "Testing NIST 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_Z;
-	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_Z;
 			break;
 		case S_Z:
-			if (strncmp(buf, "z: ", strlen("z: ")) != 0) {
-				break;
-			}
-			grab_data(&z, buf, strlen("z: "));
+			if (!get_string_cbs(&cbs, "z: ", line, msg))
+				errx(1, "#%zu %s: get_string_cbs", line, msg);
+			hex_decode_cbs(&cbs, &z, line, msg);
 			state = S_D;
 			break;
 		case S_D:
-			if (strncmp(buf, "d: ", strlen("d: ")) != 0)
-				break;
-			grab_data(&d, buf, strlen("d: "));
+			if (!get_string_cbs(&cbs, "d: ", line, msg))
+				errx(1, "#%zu %s: get_string_cbs", line, msg);
+			hex_decode_cbs(&cbs, &d, line, msg);
 			state = S_EK;
 			break;
 		case S_EK:
-			if (strncmp(buf, "ek: ", strlen("ek: ")) != 0)
-				break;
-			grab_data(&ek, buf, strlen("ek: "));
+			if (!get_string_cbs(&cbs, "ek: ", line, msg))
+				errx(1, "#%zu %s: get_string_cbs", line, msg);
+			hex_decode_cbs(&cbs, &ek, line, msg);
 			state = S_DK;
 			break;
 		case S_DK:
-			if (strncmp(buf, "dk: ", strlen("dk: ")) != 0)
-				break;
-			grab_data(&dk, buf, strlen("dk: "));
+			if (!get_string_cbs(&cbs, "dk: ", line, msg))
+				errx(1, "#%zu %s: get_string_cbs", line, msg);
+			hex_decode_cbs(&cbs, &dk, line, msg);
 
-			MlkemNistKeygenFileTest(&z, &d, &ek, &dk);
-			free((void *)CBS_data(&z));
-			free((void *)CBS_data(&d));
-			free((void *)CBS_data(&ek));
-			free((void *)CBS_data(&dk));
+			failed |= MlkemNistKeygenFileTest(&z, &d, &ek, &dk, line);
 
-			test_number++;
 			state = S_START;
 			break;
 		}
 	}
-
 	free(buf);
+
+	if (ferror(fp))
+		err(1, NULL);
 	fclose(fp);
-	exit(failure);
+
+	return failed;
 }