1 /* Copyright (c) 2024, Google Inc. 2 * Copyright (c) 2024, Bob Beck <beck@obtuse.com> 3 * 4 * Permission to use, copy, modify, and/or distribute this software for any 5 * purpose with or without fee is hereby granted, provided that the above 6 * copyright notice and this permission notice appear in all copies. 7 * 8 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 9 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 10 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY 11 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 12 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION 13 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN 14 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. */ 15 16 #include <stdlib.h> 17 #include <stdio.h> 18 #include <string.h> 19 20 #include "bytestring.h" 21 #include "mlkem.h" 22 #include "mlkem_internal.h" 23 #include "mlkem_tests_util.h" 24 25 static int 26 encode_public_key(const struct MLKEM768_public_key *pub, uint8_t **out_buf, 27 size_t *out_len) 28 { 29 CBB cbb; 30 if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES)) 31 return 0; 32 if (!MLKEM768_marshal_public_key(&cbb, pub)) 33 return 0; 34 if (!CBB_finish(&cbb, out_buf, out_len)) 35 return 0; 36 CBB_cleanup(&cbb); 37 return 1; 38 } 39 40 static int 41 encode_private_key(const struct MLKEM768_private_key *priv, uint8_t **out_buf, 42 size_t *out_len) 43 { 44 CBB cbb; 45 if (!CBB_init(&cbb, MLKEM768_PUBLIC_KEY_BYTES)) 46 return 0; 47 if (!MLKEM768_marshal_private_key(&cbb, priv)) 48 return 0; 49 if (!CBB_finish(&cbb, out_buf, out_len)) 50 return 0; 51 CBB_cleanup(&cbb); 52 return 1; 53 } 54 55 static void 56 MlKem768UnitTest() 57 { 58 struct MLKEM768_private_key *priv, *priv2; 59 struct MLKEM768_public_key *pub, *pub2; 60 uint8_t encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES]; 61 uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES]; 62 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 63 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 64 uint8_t first_two_bytes[2]; 65 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 66 size_t encoded_private_key_len, tmp_buf_len; 67 CBS cbs; 68 69 fprintf(stderr, "ML-KEM 768...\n"); 70 71 MALLOC(priv, sizeof(struct MLKEM768_private_key)); 72 MLKEM768_generate_key(encoded_public_key, NULL, priv); 73 74 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 75 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 76 CBS_init(&cbs, encoded_public_key, 77 sizeof(encoded_public_key)); 78 MALLOC(pub, sizeof(struct MLKEM768_public_key)); 79 /* Parsing should fail because the first coefficient is >= kPrime; */ 80 TEST(MLKEM768_parse_public_key(pub, &cbs), 81 "Kyber_parse_public_key should have failed"); 82 83 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 84 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 85 TEST(!MLKEM768_parse_public_key(pub, &cbs), 86 "MLKEM768_parse_public_key"); 87 TEST(CBS_len(&cbs) != 0u, "CBS_len must be 0"); 88 89 TEST(!encode_public_key(pub, &tmp_buf, &tmp_buf_len), 90 "encode_public_key"); 91 TEST(sizeof(encoded_public_key) != tmp_buf_len, 92 "encoded public key lengths differ"); 93 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 94 "encoded public keys"); 95 free(tmp_buf); 96 tmp_buf = NULL; 97 98 MALLOC(pub2, sizeof(struct MLKEM768_public_key)); 99 MLKEM768_public_from_private(pub2, priv); 100 TEST(!encode_public_key(pub2, &tmp_buf, &tmp_buf_len), 101 "encode_public_key"); 102 TEST(sizeof(encoded_public_key) != tmp_buf_len, 103 "encoded public key lengths differ"); 104 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 105 "encoded pubic keys"); 106 free(tmp_buf); 107 tmp_buf = NULL; 108 109 TEST(!encode_private_key(priv, &encoded_private_key, 110 &encoded_private_key_len), "encode_private_key"); 111 112 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 113 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 114 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 115 MALLOC(priv2, sizeof(struct MLKEM768_private_key)); 116 /* Parsing should fail because the first coefficient is >= kPrime. */ 117 TEST(MLKEM768_parse_private_key(priv2, &cbs), "Should not have parsed"); 118 119 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 120 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 121 TEST(!MLKEM768_parse_private_key(priv2, &cbs), 122 "MLKEM768_parse_private_key"); 123 TEST(!encode_private_key(priv2, &tmp_buf, &tmp_buf_len), 124 "encode_private_key"); 125 TEST(encoded_private_key_len != tmp_buf_len, 126 "encoded private key lengths differ"); 127 TEST_DATAEQ(tmp_buf, encoded_private_key, encoded_private_key_len, 128 "encoded private keys"); 129 free(tmp_buf); 130 tmp_buf = NULL; 131 132 MLKEM768_encap(ciphertext, shared_secret1, pub); 133 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 134 priv); 135 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 136 "shared secrets with priv"); 137 MLKEM768_decap(shared_secret2, ciphertext, MLKEM768_CIPHERTEXT_BYTES, 138 priv2); 139 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 140 "shared secrets with priv2"); 141 142 free(encoded_private_key); 143 free(pub); 144 free(pub2); 145 free(priv); 146 free(priv2); 147 148 } 149 150 static int 151 encode_1024public_key(const struct MLKEM1024_public_key *pub, uint8_t **out_buf, 152 size_t *out_len) 153 { 154 CBB cbb; 155 if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES)) 156 return 0; 157 if (!MLKEM1024_marshal_public_key(&cbb, pub)) 158 return 0; 159 if (!CBB_finish(&cbb, out_buf, out_len)) 160 return 0; 161 CBB_cleanup(&cbb); 162 return 1; 163 } 164 165 static int 166 encode_1024private_key(const struct MLKEM1024_private_key *priv, uint8_t **out_buf, 167 size_t *out_len) 168 { 169 CBB cbb; 170 if (!CBB_init(&cbb, MLKEM1024_PUBLIC_KEY_BYTES)) 171 return 0; 172 if (!MLKEM1024_marshal_private_key(&cbb, priv)) 173 return 0; 174 if (!CBB_finish(&cbb, out_buf, out_len)) 175 return 0; 176 CBB_cleanup(&cbb); 177 return 1; 178 } 179 180 static void 181 MlKem1024UnitTest() 182 { 183 struct MLKEM1024_private_key *priv, *priv2; 184 struct MLKEM1024_public_key *pub, *pub2; 185 uint8_t encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES]; 186 uint8_t ciphertext[MLKEM1024_CIPHERTEXT_BYTES]; 187 uint8_t shared_secret1[MLKEM_SHARED_SECRET_BYTES]; 188 uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES]; 189 uint8_t first_two_bytes[2]; 190 uint8_t *encoded_private_key = NULL, *tmp_buf = NULL; 191 size_t encoded_private_key_len, tmp_buf_len; 192 CBS cbs; 193 194 fprintf(stderr, "ML-KEM 1024...\n"); 195 196 MALLOC(priv, sizeof(struct MLKEM1024_private_key)); 197 MLKEM1024_generate_key(encoded_public_key, NULL, priv); 198 199 memcpy(first_two_bytes, encoded_public_key, sizeof(first_two_bytes)); 200 memset(encoded_public_key, 0xff, sizeof(first_two_bytes)); 201 CBS_init(&cbs, encoded_public_key, 202 sizeof(encoded_public_key)); 203 MALLOC(pub, sizeof(struct MLKEM1024_public_key)); 204 /* Parsing should fail because the first coefficient is >= kPrime; */ 205 TEST(MLKEM1024_parse_public_key(pub, &cbs), 206 "Kyber_parse_public_key should have failed"); 207 208 memcpy(encoded_public_key, first_two_bytes, sizeof(first_two_bytes)); 209 CBS_init(&cbs, encoded_public_key, sizeof(encoded_public_key)); 210 TEST(!MLKEM1024_parse_public_key(pub, &cbs), 211 "MLKEM1024_parse_public_key"); 212 TEST(CBS_len(&cbs) != 0u, "CBS_len must be 0"); 213 214 TEST(!encode_1024public_key(pub, &tmp_buf, &tmp_buf_len), 215 "encode_1024public_key"); 216 TEST(sizeof(encoded_public_key) != tmp_buf_len, 217 "encoded public key lengths differ"); 218 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 219 "encoded public keys"); 220 free(tmp_buf); 221 tmp_buf = NULL; 222 223 MALLOC(pub2, sizeof(struct MLKEM1024_public_key)); 224 MLKEM1024_public_from_private(pub2, priv); 225 TEST(!encode_1024public_key(pub2, &tmp_buf, &tmp_buf_len), 226 "encode_public_key"); 227 TEST(sizeof(encoded_public_key) != tmp_buf_len, 228 "encoded public key lengths differ"); 229 TEST_DATAEQ(tmp_buf, encoded_public_key, tmp_buf_len, 230 "encoded pubic keys"); 231 free(tmp_buf); 232 tmp_buf = NULL; 233 234 TEST(!encode_1024private_key(priv, &encoded_private_key, 235 &encoded_private_key_len), "encode_1024private_key"); 236 237 memcpy(first_two_bytes, encoded_private_key, sizeof(first_two_bytes)); 238 memset(encoded_private_key, 0xff, sizeof(first_two_bytes)); 239 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 240 MALLOC(priv2, sizeof(struct MLKEM1024_private_key)); 241 /* Parsing should fail because the first coefficient is >= kPrime. */ 242 TEST(MLKEM1024_parse_private_key(priv2, &cbs), "Should not have parsed"); 243 244 memcpy(encoded_private_key, first_two_bytes, sizeof(first_two_bytes)); 245 CBS_init(&cbs, encoded_private_key, encoded_private_key_len); 246 TEST(!MLKEM1024_parse_private_key(priv2, &cbs), 247 "MLKEM1024_parse_private_key"); 248 TEST(!encode_1024private_key(priv2, &tmp_buf, &tmp_buf_len), 249 "encode_private_key"); 250 TEST(encoded_private_key_len != tmp_buf_len, 251 "encoded private key lengths differ"); 252 TEST_DATAEQ(tmp_buf, encoded_private_key, encoded_private_key_len, 253 "encoded private keys"); 254 free(tmp_buf); 255 tmp_buf = NULL; 256 257 MLKEM1024_encap(ciphertext, shared_secret1, pub); 258 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 259 priv); 260 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 261 "shared secrets with priv"); 262 MLKEM1024_decap(shared_secret2, ciphertext, MLKEM1024_CIPHERTEXT_BYTES, 263 priv2); 264 TEST_DATAEQ(shared_secret1, shared_secret2, MLKEM_SHARED_SECRET_BYTES, 265 "shared secrets with priv2"); 266 267 free(encoded_private_key); 268 free(pub); 269 free(pub2); 270 free(priv); 271 free(priv2); 272 } 273 274 int 275 main(int argc, char **argv) 276 { 277 MlKem768UnitTest(); 278 MlKem1024UnitTest(); 279 280 exit(failure); 281 } 282