1 /* $OpenBSD: mlkem_unittest.c,v 1.6 2024/12/26 12:35:25 tb Exp $ */ 2 /* 3 * Copyright (c) 2024 Google Inc. 4 * Copyright (c) 2024 Bob Beck <beck@obtuse.com> 5 * 6 * Permission to use, copy, modify, and/or distribute this software for any 7 * purpose with or without fee is hereby granted, provided that the above 8 * copyright notice and this permission notice appear in all copies. 9 * 10 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 11 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 12 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY 13 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 14 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION 15 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN 16 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 17 */ 18 19 #include <err.h> 20 #include <stdint.h> 21 #include <stdio.h> 22 #include <stdlib.h> 23 #include <string.h> 24 25 #include "bytestring.h" 26 #include "mlkem.h" 27 28 #include "mlkem_tests_util.h" 29 30 struct unittest_ctx { 31 void *priv; 32 void *pub; 33 void *priv2; 34 void *pub2; 35 uint8_t *encoded_public_key; 36 size_t encoded_public_key_len; 37 uint8_t *ciphertext; 38 size_t ciphertext_len; 39 mlkem_decap_fn decap; 40 mlkem_encap_fn encap; 41 mlkem_generate_key_fn generate_key; 42 mlkem_parse_private_key_fn parse_private_key; 43 mlkem_parse_public_key_fn parse_public_key; 44 mlkem_encode_private_key_fn encode_private_key; 45 mlkem_encode_public_key_fn encode_public_key; 46 mlkem_public_from_private_fn public_from_private; 47 }; 48 49 static int 50 MlKemUnitTest(struct unittest_ctx *ctx) 51 { 52 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 53 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 54 uint8_t first_two_bytes[2]; 55 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 56 size_t encoded_private_key_len, tmp_buf_len; 57 CBS cbs; 58 int failed = 0; 59 60 ctx->generate_key(ctx->encoded_public_key, NULL, ctx->priv); 61 62 memcpy(first_two_bytes, ctx->encoded_public_key, sizeof(first_two_bytes)); 63 memset(ctx->encoded_public_key, 0xff, sizeof(first_two_bytes)); 64 65 CBS_init(&cbs, ctx->encoded_public_key, ctx->encoded_public_key_len); 66 67 /* Parsing should fail because the first coefficient is >= kPrime. */ 68 if (ctx->parse_public_key(ctx->pub, &cbs)) { 69 warnx("parse_public_key should have failed"); 70 failed |= 1; 71 } 72 73 memcpy(ctx->encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 74 CBS_init(&cbs, ctx->encoded_public_key, ctx->encoded_public_key_len); 75 if (!ctx->parse_public_key(ctx->pub, &cbs)) { 76 warnx("MLKEM768_parse_public_key"); 77 failed |= 1; 78 } 79 80 if (CBS_len(&cbs) != 0u) { 81 warnx("CBS_len must be 0"); 82 failed |= 1; 83 } 84 85 if (!ctx->encode_public_key(ctx->pub, &tmp_buf, &tmp_buf_len)) { 86 warnx("encode_public_key"); 87 failed |= 1; 88 } 89 if (ctx->encoded_public_key_len != tmp_buf_len) { 90 warnx("encoded public key lengths differ"); 91 failed |= 1; 92 } 93 94 if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len, 95 "encoded public keys") != 0) { 96 warnx("compare_data"); 97 failed |= 1; 98 } 99 free(tmp_buf); 100 tmp_buf = NULL; 101 102 ctx->public_from_private(ctx->pub2, ctx->priv); 103 if (!ctx->encode_public_key(ctx->pub2, &tmp_buf, &tmp_buf_len)) { 104 warnx("encode_public_key"); 105 failed |= 1; 106 } 107 if (ctx->encoded_public_key_len != tmp_buf_len) { 108 warnx("encoded public key lengths differ"); 109 failed |= 1; 110 } 111 112 if (compare_data(ctx->encoded_public_key, tmp_buf, tmp_buf_len, 113 "encoded public keys") != 0) { 114 warnx("compare_data"); 115 failed |= 1; 116 } 117 free(tmp_buf); 118 tmp_buf = NULL; 119 120 if (!ctx->encode_private_key(ctx->priv, &encoded_private_key, 121 &encoded_private_key_len)) { 122 warnx("mlkem768_encode_private_key"); 123 failed |= 1; 124 } 125 126 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 127 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 128 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 129 130 /* Parsing should fail because the first coefficient is >= kPrime. */ 131 if (ctx->parse_private_key(ctx->priv2, &cbs)) { 132 warnx("MLKEM768_parse_private_key should have failed"); 133 failed |= 1; 134 } 135 136 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 137 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 138 139 if (!ctx->parse_private_key(ctx->priv2, &cbs)) { 140 warnx("MLKEM768_parse_private_key"); 141 failed |= 1; 142 } 143 144 if (!ctx->encode_private_key(ctx->priv2, &tmp_buf, &tmp_buf_len)) { 145 warnx("encode_private_key"); 146 failed |= 1; 147 } 148 149 if (encoded_private_key_len != tmp_buf_len) { 150 warnx("encode private key lengths differ"); 151 failed |= 1; 152 } 153 154 if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len, 155 "encoded private key") != 0) { 156 warnx("compare_data"); 157 failed |= 1; 158 } 159 160 free(tmp_buf); 161 tmp_buf = NULL; 162 163 ctx->encap(ctx->ciphertext, shared_secret1, ctx->pub); 164 ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len, 165 ctx->priv); 166 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 167 "shared secrets with priv") != 0) { 168 warnx("compare_data"); 169 failed |= 1; 170 } 171 172 ctx->decap(shared_secret2, ctx->ciphertext, ctx->ciphertext_len, 173 ctx->priv2); 174 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 175 "shared secrets with priv2") != 0) { 176 warnx("compare_data"); 177 failed |= 1; 178 } 179 180 free(encoded_private_key); 181 182 return failed; 183 } 184 185 static int 186 mlkem768_unittest(void) 187 { 188 struct MLKEM768_private_key mlkem768_priv, mlkem768_priv2; 189 struct MLKEM768_public_key mlkem768_pub, mlkem768_pub2; 190 uint8_t mlkem768_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES]; 191 uint8_t mlkem768_ciphertext[MLKEM768_CIPHERTEXT_BYTES]; 192 struct unittest_ctx mlkem768_test = { 193 .priv = &mlkem768_priv, 194 .pub = &mlkem768_pub, 195 .priv2 = &mlkem768_priv2, 196 .pub2 = &mlkem768_pub2, 197 .encoded_public_key = mlkem768_encoded_public_key, 198 .encoded_public_key_len = sizeof(mlkem768_encoded_public_key), 199 .ciphertext = mlkem768_ciphertext, 200 .ciphertext_len = sizeof(mlkem768_ciphertext), 201 .decap = mlkem768_decap, 202 .encap = mlkem768_encap, 203 .generate_key = mlkem768_generate_key, 204 .parse_private_key = mlkem768_parse_private_key, 205 .parse_public_key = mlkem768_parse_public_key, 206 .encode_private_key = mlkem768_encode_private_key, 207 .encode_public_key = mlkem768_encode_public_key, 208 .public_from_private = mlkem768_public_from_private, 209 }; 210 211 return MlKemUnitTest(&mlkem768_test); 212 } 213 214 static int 215 mlkem1024_unittest(void) 216 { 217 struct MLKEM1024_private_key mlkem1024_priv, mlkem1024_priv2; 218 struct MLKEM1024_public_key mlkem1024_pub, mlkem1024_pub2; 219 uint8_t mlkem1024_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES]; 220 uint8_t mlkem1024_ciphertext[MLKEM1024_CIPHERTEXT_BYTES]; 221 struct unittest_ctx mlkem1024_test = { 222 .priv = &mlkem1024_priv, 223 .pub = &mlkem1024_pub, 224 .priv2 = &mlkem1024_priv2, 225 .pub2 = &mlkem1024_pub2, 226 .encoded_public_key = mlkem1024_encoded_public_key, 227 .encoded_public_key_len = sizeof(mlkem1024_encoded_public_key), 228 .ciphertext = mlkem1024_ciphertext, 229 .ciphertext_len = sizeof(mlkem1024_ciphertext), 230 .decap = mlkem1024_decap, 231 .encap = mlkem1024_encap, 232 .generate_key = mlkem1024_generate_key, 233 .parse_private_key = mlkem1024_parse_private_key, 234 .parse_public_key = mlkem1024_parse_public_key, 235 .encode_private_key = mlkem1024_encode_private_key, 236 .encode_public_key = mlkem1024_encode_public_key, 237 .public_from_private = mlkem1024_public_from_private, 238 }; 239 240 return MlKemUnitTest(&mlkem1024_test); 241 } 242 243 int 244 main(void) 245 { 246 int failed = 0; 247 248 /* 249 * XXX - this is split into two helper functions since having a few 250 * ML-KEM key blobs on the stack makes Emscripten's stack explode, 251 * leading to inscrutable silent failures unles ASAN is enabled. 252 * Go figure. 253 */ 254 255 failed |= mlkem768_unittest(); 256 failed |= mlkem1024_unittest(); 257 258 return failed; 259 } 260