1 /* $OpenBSD: mlkem_unittest.c,v 1.4 2024/12/20 00:07:12 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 static int 31 MlKem768UnitTest(void) 32 { 33 struct MLKEM768_private_key priv = { 0 }, priv2 = { 0 }; 34 struct MLKEM768_public_key pub = { 0 }, pub2 = { 0 }; 35 uint8_t encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES]; 36 uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES]; 37 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 38 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 39 uint8_t first_two_bytes[2]; 40 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 41 size_t encoded_private_key_len, tmp_buf_len; 42 CBS cbs; 43 int failed = 0; 44 45 MLKEM768_generate_key(encoded_public_key, NULL, &priv); 46 47 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 48 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 49 50 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 51 52 /* Parsing should fail because the first coefficient is >= kPrime. */ 53 if (MLKEM768_parse_public_key(&pub, &cbs)) { 54 warnx("MLKEM768_parse_public_key should have failed"); 55 failed |= 1; 56 } 57 58 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 59 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 60 if (!MLKEM768_parse_public_key(&pub, &cbs)) { 61 warnx("MLKEM768_parse_public_key"); 62 failed |= 1; 63 } 64 65 if (CBS_len(&cbs) != 0u) { 66 warnx("CBS_len must be 0"); 67 failed |= 1; 68 } 69 70 if (!mlkem768_encode_public_key(&pub, &tmp_buf, &tmp_buf_len)) { 71 warnx("encode_public_key"); 72 failed |= 1; 73 } 74 if (sizeof(encoded_public_key) != tmp_buf_len) { 75 warnx("mlkem768 encoded public key lengths differ"); 76 failed |= 1; 77 } 78 79 if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 768, 80 "encoded public keys") != 0) { 81 warnx("compare_data"); 82 failed |= 1; 83 } 84 free(tmp_buf); 85 tmp_buf = NULL; 86 87 MLKEM768_public_from_private(&pub2, &priv); 88 if (!mlkem768_encode_public_key(&pub2, &tmp_buf, &tmp_buf_len)) { 89 warnx("mlkem768_encode_public_key"); 90 failed |= 1; 91 } 92 if (sizeof(encoded_public_key) != tmp_buf_len) { 93 warnx("mlkem768 encoded public key lengths differ"); 94 failed |= 1; 95 } 96 97 if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 768, 98 "encoded public keys") != 0) { 99 warnx("compare_data"); 100 failed |= 1; 101 } 102 free(tmp_buf); 103 tmp_buf = NULL; 104 105 if (!mlkem768_encode_private_key(&priv, &encoded_private_key, 106 &encoded_private_key_len)) { 107 warnx("mlkem768_encode_private_key"); 108 failed |= 1; 109 } 110 111 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 112 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 113 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 114 115 /* Parsing should fail because the first coefficient is >= kPrime. */ 116 if (MLKEM768_parse_private_key(&priv2, &cbs)) { 117 warnx("MLKEM768_parse_private_key should have failed"); 118 failed |= 1; 119 } 120 121 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 122 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 123 124 if (!MLKEM768_parse_private_key(&priv2, &cbs)) { 125 warnx("MLKEM768_parse_private_key"); 126 failed |= 1; 127 } 128 129 if (!mlkem768_encode_private_key(&priv2, &tmp_buf, &tmp_buf_len)) { 130 warnx("mlkem768_encode_private_key"); 131 failed |= 1; 132 } 133 134 if (encoded_private_key_len != tmp_buf_len) { 135 warnx("mlkem768 encode private key lengths differ"); 136 failed |= 1; 137 } 138 139 if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len, 768, 140 "encoded private key") != 0) { 141 warnx("compare_data"); 142 failed |= 1; 143 } 144 145 free(tmp_buf); 146 tmp_buf = NULL; 147 148 MLKEM768_encap(ciphertext, shared_secret1, &pub); 149 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 150 &priv); 151 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 152 768, "shared secrets with priv") != 0) { 153 warnx("compare_data"); 154 failed |= 1; 155 } 156 157 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 158 &priv2); 159 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 160 768, "shared secrets with priv2") != 0) { 161 warnx("compare_data"); 162 failed |= 1; 163 } 164 165 free(encoded_private_key); 166 167 return failed; 168 } 169 170 static int 171 MlKem1024UnitTest(void) 172 { 173 struct MLKEM1024_private_key priv = { 0 }, priv2 = { 0 }; 174 struct MLKEM1024_public_key pub = { 0 }, pub2 = { 0 }; 175 uint8_t encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES]; 176 uint8_t ciphertext[MLKEM1024_CIPHERTEXT_BYTES]; 177 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 178 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 179 uint8_t first_two_bytes[2]; 180 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 181 size_t encoded_private_key_len, tmp_buf_len; 182 CBS cbs; 183 int failed = 0; 184 185 MLKEM1024_generate_key(encoded_public_key, NULL, &priv); 186 187 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 188 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 189 190 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 191 192 /* Parsing should fail because the first coefficient is >= kPrime. */ 193 if (MLKEM1024_parse_public_key(&pub, &cbs)) { 194 warnx("MLKEM1024_parse_public_key should have failed"); 195 failed |= 1; 196 } 197 198 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 199 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 200 if (!MLKEM1024_parse_public_key(&pub, &cbs)) { 201 warnx("MLKEM1024_parse_public_key"); 202 failed |= 1; 203 } 204 205 if (CBS_len(&cbs) != 0u) { 206 warnx("CBS_len must be 0"); 207 failed |= 1; 208 } 209 210 if (!mlkem1024_encode_public_key(&pub, &tmp_buf, &tmp_buf_len)) { 211 warnx("encode_public_key"); 212 failed |= 1; 213 } 214 if (sizeof(encoded_public_key) != tmp_buf_len) { 215 warnx("mlkem1024 encoded public key lengths differ"); 216 failed |= 1; 217 } 218 219 if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 1024, 220 "encoded public keys") != 0) { 221 warnx("compare_data"); 222 failed |= 1; 223 } 224 free(tmp_buf); 225 tmp_buf = NULL; 226 227 MLKEM1024_public_from_private(&pub2, &priv); 228 if (!mlkem1024_encode_public_key(&pub2, &tmp_buf, &tmp_buf_len)) { 229 warnx("mlkem1024_encode_public_key"); 230 failed |= 1; 231 } 232 if (sizeof(encoded_public_key) != tmp_buf_len) { 233 warnx("mlkem1024 encoded public key lengths differ"); 234 failed |= 1; 235 } 236 237 if (compare_data(encoded_public_key, tmp_buf, tmp_buf_len, 1024, 238 "encoded public keys") != 0) { 239 warnx("compare_data"); 240 failed |= 1; 241 } 242 free(tmp_buf); 243 tmp_buf = NULL; 244 245 if (!mlkem1024_encode_private_key(&priv, &encoded_private_key, 246 &encoded_private_key_len)) { 247 warnx("mlkem1024_encode_private_key"); 248 failed |= 1; 249 } 250 251 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 252 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 253 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 254 255 /* Parsing should fail because the first coefficient is >= kPrime. */ 256 if (MLKEM1024_parse_private_key(&priv2, &cbs)) { 257 warnx("MLKEM1024_parse_private_key should have failed"); 258 failed |= 1; 259 } 260 261 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 262 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 263 264 if (!MLKEM1024_parse_private_key(&priv2, &cbs)) { 265 warnx("MLKEM1024_parse_private_key"); 266 failed |= 1; 267 } 268 269 if (!mlkem1024_encode_private_key(&priv2, &tmp_buf, &tmp_buf_len)) { 270 warnx("mlkem1024_encode_private_key"); 271 failed |= 1; 272 } 273 274 if (encoded_private_key_len != tmp_buf_len) { 275 warnx("mlkem1024 encode private key lengths differ"); 276 failed |= 1; 277 } 278 279 if (compare_data(encoded_private_key, tmp_buf, tmp_buf_len, 1024, 280 "encoded private key") != 0) { 281 warnx("compare_data"); 282 failed |= 1; 283 } 284 285 free(tmp_buf); 286 tmp_buf = NULL; 287 288 MLKEM1024_encap(ciphertext, shared_secret1, &pub); 289 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 290 &priv); 291 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 292 1024, "shared secrets with priv") != 0) { 293 warnx("compare_data"); 294 failed |= 1; 295 } 296 297 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 298 &priv2); 299 if (compare_data(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 300 1024, "shared secrets with priv2") != 0) { 301 warnx("compare_data"); 302 failed |= 1; 303 } 304 305 free(encoded_private_key); 306 307 return failed; 308 } 309 310 int 311 main(int argc, char **argv) 312 { 313 int failed = 0; 314 315 failed |= MlKem768UnitTest(); 316 failed |= MlKem1024UnitTest(); 317 318 return failed; 319 } 320