ML-KEM decapsulation key hash check (aws#1873)

ML-KEM decapsulation key check as specified in Section 7.3 of FIPS 203: https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf
justsmth · Sep 25, 2024 · c5d3f3d · c5d3f3d
1 parent 8669544
commit c5d3f3d
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 3 deletions.
diff --git a/crypto/evp_extra/evp_extra_test.cc b/crypto/evp_extra/evp_extra_test.cc
@@ -2906,15 +2906,33 @@ TEST_P(PerMLKEMTest, InputValidation) {
   std::vector<uint8_t> ss(ss_len);
 
   // Encapsulate.
-  ASSERT_TRUE(
-      EVP_PKEY_encapsulate(ctx.get(), ct.data(), &ct_len, ss.data(), &ss_len));
+  ASSERT_TRUE(EVP_PKEY_encapsulate(ctx.get(), ct.data(), &ct_len, ss.data(), &ss_len));
 
   // ---- 3. Test invalid public key ----
   // FIPS 203 Section 7.2 Encapsulation key check (Modulus check).
-  // Invalidate the key by forcing a coefficient out of range.
+  // Invalidate the key by forcing a coefficient out of range
+  // (save the original values to reset later).
+  uint8_t tmp0 = ctx->pkey->pkey.kem_key->public_key[0];
+  uint8_t tmp1 = ctx->pkey->pkey.kem_key->public_key[1];
   ctx->pkey->pkey.kem_key->public_key[0] = 0xff;
   ctx->pkey->pkey.kem_key->public_key[1] = 0xff;
 
   ASSERT_FALSE(
       EVP_PKEY_encapsulate(ctx.get(), ct.data(), &ct_len, ss.data(), &ss_len));
+
+  // Reset the public key and make sure encapsulation/decapsulation succeeds.
+  ctx->pkey->pkey.kem_key->public_key[0] = tmp0;
+  ctx->pkey->pkey.kem_key->public_key[1] = tmp1;
+
+  std::vector<uint8_t> ss_expected(ss_len); // The shared secret.
+  ASSERT_TRUE(EVP_PKEY_encapsulate(ctx.get(), ct.data(), &ct_len, ss.data(), &ss_len));
+  ASSERT_TRUE(EVP_PKEY_decapsulate(ctx.get(), ss_expected.data(), &ss_len, ct.data(), ct_len));
+  EXPECT_EQ(Bytes(ss_expected), Bytes(ss));
+
+  // ---- 4. Test invalid secret key ----
+  // FIPS 203 Section 7.3 Decapsulation key check (Hash check).
+  // Invalidate the key by changing the hash of the public key within the secret key.
+  // The 32-byte hash is stored right before the last 32 bytes of the secret key.
+  ctx->pkey->pkey.kem_key->secret_key[GetParam().secret_key_len - 64] ^= 1;
+  ASSERT_FALSE(EVP_PKEY_decapsulate(ctx.get(), ss_expected.data(), &ss_len, ct.data(), ct_len));
 }
diff --git a/crypto/fipsmodule/ml_kem/ml_kem_ref/kem.c b/crypto/fipsmodule/ml_kem/ml_kem_ref/kem.c
@@ -156,6 +156,22 @@ static int encapsulation_key_modulus_check(ml_kem_params *params, const uint8_t
   return verify(ek_recoded, ek, params->k * BYTE_ENCODE_12_OUT_SIZE);
 }
 
+// FIPS 203. Section 7.3 Decapsulation key hash check
+// The spec defines the decapsulation key as following:
+//   dk <-- (dk_pke || ek || H(ek) || z).
+// This check takes |ek| out of |dk|, computes H(ek), and verifies that it is
+// the same as the H(ek) portion stored in |dk|.
+static int decapsulation_key_hash_check(ml_kem_params *params, const uint8_t *dk) {
+  uint8_t dk_pke_hash_computed[KYBER_SYMBYTES] = {0};
+
+  hash_h(dk_pke_hash_computed, &dk[params->indcpa_secret_key_bytes],
+                               params->indcpa_public_key_bytes);
+  const uint8_t *dk_pke_hash_expected = &dk[params->indcpa_secret_key_bytes +
+                                            params->indcpa_public_key_bytes];
+
+  return verify(dk_pke_hash_computed, dk_pke_hash_expected, KYBER_SYMBYTES);
+}
+
 /*************************************************
 * Name:        crypto_kem_enc_derand
 *
@@ -248,6 +264,10 @@ int crypto_kem_dec(ml_kem_params *params,
                    const uint8_t *ct,
                    const uint8_t *sk)
 {
+  if (decapsulation_key_hash_check(params, sk) != 0) {
+    return 1;
+  }
+
   int fail;
   uint8_t buf[2*KYBER_SYMBYTES];
   /* Will contain key, coins */