1 /* $OpenBSD: mlkem_unittest.c,v 1.3 2024/12/14 19:16:24 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 <stdint.h> 20 #include <stdio.h> 21 #include <stdlib.h> 22 #include <string.h> 23 24 #include <openssl/bytestring.h> 25 #include <openssl/mlkem.h> 26 27 #include "mlkem_internal.h" 28 #include "mlkem_tests_util.h" 29 30 static int 31 encode_public_key(const struct MLKEM768_public_key *pub, uint8_t **out_buf, 32 size_t *out_len) 33 { 34 CBB cbb; 35 if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES)) 36 return 0; 37 if (!MLKEM768_marshal_public_key(&cbb, pub)) 38 return 0; 39 if (!CBB_finish(&cbb, out_buf, out_len)) 40 return 0; 41 CBB_cleanup(&cbb); 42 return 1; 43 } 44 45 static int 46 encode_private_key(const struct MLKEM768_private_key *priv, uint8_t **out_buf, 47 size_t *out_len) 48 { 49 CBB cbb; 50 if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES)) 51 return 0; 52 if (!MLKEM768_marshal_private_key(&cbb, priv)) 53 return 0; 54 if (!CBB_finish(&cbb, out_buf, out_len)) 55 return 0; 56 CBB_cleanup(&cbb); 57 return 1; 58 } 59 60 static void 61 MlKem768UnitTest() 62 { 63 struct MLKEM768_private_key *priv, *priv2; 64 struct MLKEM768_public_key *pub, *pub2; 65 uint8_t encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES]; 66 uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES]; 67 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 68 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 69 uint8_t first_two_bytes[2]; 70 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 71 size_t encoded_private_key_len, tmp_buf_len; 72 CBS cbs; 73 74 fprintf(stderr, "ML-KEM 768...\n"); 75 76 MALLOC(priv, sizeof(struct MLKEM768_private_key)); 77 MLKEM768_generate_key(encoded_public_key, NULL, priv); 78 79 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 80 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 81 CBS_init(&cbs, encoded_public_key, 82 sizeof(encoded_public_key)); 83 MALLOC(pub, sizeof(struct MLKEM768_public_key)); 84 /* Parsing should fail because the first coefficient is >= kPrime; */ 85 TEST(MLKEM768_parse_public_key(pub, &cbs), 86 "Kyber_parse_public_key should have failed"); 87 88 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 89 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 90 TEST(!MLKEM768_parse_public_key(pub, &cbs), 91 "MLKEM768_parse_public_key"); 92 TEST(CBS_len(&cbs) != 0u, "CBS_len must be 0"); 93 94 TEST(!encode_public_key(pub, &tmp_buf, &tmp_buf_len), 95 "encode_public_key"); 96 TEST(sizeof(encoded_public_key) != tmp_buf_len, 97 "encoded public key lengths differ"); 98 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 99 "encoded public keys"); 100 free(tmp_buf); 101 tmp_buf = NULL; 102 103 MALLOC(pub2, sizeof(struct MLKEM768_public_key)); 104 MLKEM768_public_from_private(pub2, priv); 105 TEST(!encode_public_key(pub2, &tmp_buf, &tmp_buf_len), 106 "encode_public_key"); 107 TEST(sizeof(encoded_public_key) != tmp_buf_len, 108 "encoded public key lengths differ"); 109 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 110 "encoded pubic keys"); 111 free(tmp_buf); 112 tmp_buf = NULL; 113 114 TEST(!encode_private_key(priv, &encoded_private_key, 115 &encoded_private_key_len), "encode_private_key"); 116 117 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 118 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 119 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 120 MALLOC(priv2, sizeof(struct MLKEM768_private_key)); 121 /* Parsing should fail because the first coefficient is >= kPrime. */ 122 TEST(MLKEM768_parse_private_key(priv2, &cbs), "Should not have parsed"); 123 124 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 125 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 126 TEST(!MLKEM768_parse_private_key(priv2, &cbs), 127 "MLKEM768_parse_private_key"); 128 TEST(!encode_private_key(priv2, &tmp_buf, &tmp_buf_len), 129 "encode_private_key"); 130 TEST(encoded_private_key_len != tmp_buf_len, 131 "encoded private key lengths differ"); 132 TEST_DATAEQ(tmp_buf, encoded_private_key, encoded_private_key_len, 133 "encoded private keys"); 134 free(tmp_buf); 135 tmp_buf = NULL; 136 137 MLKEM768_encap(ciphertext, shared_secret1, pub); 138 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 139 priv); 140 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 141 "shared secrets with priv"); 142 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 143 priv2); 144 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 145 "shared secrets with priv2"); 146 147 free(encoded_private_key); 148 free(pub); 149 free(pub2); 150 free(priv); 151 free(priv2); 152 153 } 154 155 static int 156 encode_1024public_key(const struct MLKEM1024_public_key *pub, uint8_t **out_buf, 157 size_t *out_len) 158 { 159 CBB cbb; 160 if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES)) 161 return 0; 162 if (!MLKEM1024_marshal_public_key(&cbb, pub)) 163 return 0; 164 if (!CBB_finish(&cbb, out_buf, out_len)) 165 return 0; 166 CBB_cleanup(&cbb); 167 return 1; 168 } 169 170 static int 171 encode_1024private_key(const struct MLKEM1024_private_key *priv, uint8_t **out_buf, 172 size_t *out_len) 173 { 174 CBB cbb; 175 if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES)) 176 return 0; 177 if (!MLKEM1024_marshal_private_key(&cbb, priv)) 178 return 0; 179 if (!CBB_finish(&cbb, out_buf, out_len)) 180 return 0; 181 CBB_cleanup(&cbb); 182 return 1; 183 } 184 185 static void 186 MlKem1024UnitTest() 187 { 188 struct MLKEM1024_private_key *priv, *priv2; 189 struct MLKEM1024_public_key *pub, *pub2; 190 uint8_t encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES]; 191 uint8_t ciphertext[MLKEM1024_CIPHERTEXT_BYTES]; 192 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 193 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 194 uint8_t first_two_bytes[2]; 195 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 196 size_t encoded_private_key_len, tmp_buf_len; 197 CBS cbs; 198 199 fprintf(stderr, "ML-KEM 1024...\n"); 200 201 MALLOC(priv, sizeof(struct MLKEM1024_private_key)); 202 MLKEM1024_generate_key(encoded_public_key, NULL, priv); 203 204 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 205 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 206 CBS_init(&cbs, encoded_public_key, 207 sizeof(encoded_public_key)); 208 MALLOC(pub, sizeof(struct MLKEM1024_public_key)); 209 /* Parsing should fail because the first coefficient is >= kPrime; */ 210 TEST(MLKEM1024_parse_public_key(pub, &cbs), 211 "Kyber_parse_public_key should have failed"); 212 213 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 214 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 215 TEST(!MLKEM1024_parse_public_key(pub, &cbs), 216 "MLKEM1024_parse_public_key"); 217 TEST(CBS_len(&cbs) != 0u, "CBS_len must be 0"); 218 219 TEST(!encode_1024public_key(pub, &tmp_buf, &tmp_buf_len), 220 "encode_1024public_key"); 221 TEST(sizeof(encoded_public_key) != tmp_buf_len, 222 "encoded public key lengths differ"); 223 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 224 "encoded public keys"); 225 free(tmp_buf); 226 tmp_buf = NULL; 227 228 MALLOC(pub2, sizeof(struct MLKEM1024_public_key)); 229 MLKEM1024_public_from_private(pub2, priv); 230 TEST(!encode_1024public_key(pub2, &tmp_buf, &tmp_buf_len), 231 "encode_public_key"); 232 TEST(sizeof(encoded_public_key) != tmp_buf_len, 233 "encoded public key lengths differ"); 234 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 235 "encoded pubic keys"); 236 free(tmp_buf); 237 tmp_buf = NULL; 238 239 TEST(!encode_1024private_key(priv, &encoded_private_key, 240 &encoded_private_key_len), "encode_1024private_key"); 241 242 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 243 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 244 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 245 MALLOC(priv2, sizeof(struct MLKEM1024_private_key)); 246 /* Parsing should fail because the first coefficient is >= kPrime. */ 247 TEST(MLKEM1024_parse_private_key(priv2, &cbs), "Should not have parsed"); 248 249 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 250 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 251 TEST(!MLKEM1024_parse_private_key(priv2, &cbs), 252 "MLKEM1024_parse_private_key"); 253 TEST(!encode_1024private_key(priv2, &tmp_buf, &tmp_buf_len), 254 "encode_private_key"); 255 TEST(encoded_private_key_len != tmp_buf_len, 256 "encoded private key lengths differ"); 257 TEST_DATAEQ(tmp_buf, encoded_private_key, encoded_private_key_len, 258 "encoded private keys"); 259 free(tmp_buf); 260 tmp_buf = NULL; 261 262 MLKEM1024_encap(ciphertext, shared_secret1, pub); 263 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 264 priv); 265 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 266 "shared secrets with priv"); 267 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 268 priv2); 269 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 270 "shared secrets with priv2"); 271 272 free(encoded_private_key); 273 free(pub); 274 free(pub2); 275 free(priv); 276 free(priv2); 277 } 278 279 int 280 main(int argc, char **argv) 281 { 282 MlKem768UnitTest(); 283 MlKem1024UnitTest(); 284 285 exit(failure); 286 } 287