Merge GH #3368 Fix Dilithium verification bug

randombit · Mar 14, 2023 · f471186 · f471186
2 parents 563f186 + 7d220ba
commit f471186
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 9 deletions.
diff --git a/src/lib/pubkey/dilithium/dilithium_common/dilithium.cpp b/src/lib/pubkey/dilithium/dilithium_common/dilithium.cpp
@@ -534,10 +534,14 @@ class Dilithium_Verification_Operation final : public PK_Ops::Verification
       */
       bool is_valid_signature(const  uint8_t* sig, size_t sig_len) override
          {
+         const auto& mode = m_pub_key.m_public->mode();
+
          /* Compute CRH(H(rho, t1), msg) */
          const auto mu = m_shake.final_stdvec();
 
-         const auto& mode = m_pub_key.m_public->mode();
+         // Reset shake context for the next message
+         m_shake.update(mode.H(m_pub_key.m_public->raw_pk(), DilithiumModeConstants::SEEDBYTES));
+
          if(sig_len != mode.crypto_bytes())
             {
             return false;

diff --git a/src/tests/test_dilithium.cpp b/src/tests/test_dilithium.cpp
@@ -68,15 +68,19 @@ class Dilithium_KAT_Tests : public Text_Based_Test
 
 
          Botan::Dilithium_PublicKey pub_key(priv_key.public_key_bits(), DerivedT::mode, Botan::DilithiumKeyEncoding::Raw);
-         auto verificator = Botan::PK_Verifier(pub_key,"");
-         verificator.update(ref_msg.data(), ref_msg.size());
+         auto verifier = Botan::PK_Verifier(pub_key,"");
+         verifier.update(ref_msg.data(), ref_msg.size());
          result.confirm("signature verifies",
-                        verificator.check_signature(signature.data(), signature.size()));
+                        verifier.check_signature(signature.data(), signature.size()));
 
          // test validating incorrect wrong signagture
          auto mutated_signature = Test::mutate_vec(signature);
          result.confirm("invalid signature rejected",
-                        !verificator.check_signature(mutated_signature.data(), mutated_signature.size()));
+                        !verifier.check_signature(mutated_signature.data(), mutated_signature.size()));
+
+         verifier.update(ref_msg.data(), ref_msg.size());
+         result.confirm("signature verifies",
+                        verifier.check_signature(signature.data(), signature.size()));
 
          return result;
          }

diff --git a/src/tests/test_kyber.cpp b/src/tests/test_kyber.cpp
@@ -121,8 +121,7 @@ class KYBER_Tests final : public Test
          Botan::Kyber_PrivateKey alice_priv_key(priv_key_bits, mode, Botan::KyberKeyEncoding::Full);
          auto dec = Botan::PK_KEM_Decryptor(alice_priv_key, Test::rng(), "Raw", "base");
          const auto key_alice = dec.decrypt(cipher_text, 0 /* no KDF */, std::vector<uint8_t>());
-
-         result.confirm("shared secrets are equal", key_alice == key_bob);
+         result.test_eq("shared secrets are equal", key_alice, key_bob);
 
          //
          // negative tests
@@ -139,9 +138,13 @@ class KYBER_Tests final : public Test
          // Invalid cipher_text from Alice
          Botan::secure_vector<uint8_t> reverse_cipher_text;
          std::copy(cipher_text.crbegin(), cipher_text.crend(), std::back_inserter(reverse_cipher_text));
-         const auto key_alice2 =
+         const auto key_alice_rev =
             dec.decrypt(reverse_cipher_text, 0, std::vector<uint8_t>());
-         result.confirm("shared secrets are not equal", key_alice != key_alice2);
+         result.confirm("shared secrets are not equal", key_alice != key_alice_rev);
+
+         // Try to decrypt the valid ciphertext again
+         const auto key_alice_try2 = dec.decrypt(cipher_text, 0 /* no KDF */, std::vector<uint8_t>());
+         result.test_eq("shared secrets are equal", key_alice_try2, key_bob);
 
          //
          // regression tests

diff --git a/src/tests/test_pubkey.cpp b/src/tests/test_pubkey.cpp
@@ -153,6 +153,9 @@ PK_Signature_Generation_Test::run_one_test(const std::string& pad_hdr, const Var
       result.test_eq("KAT signature valid", verifier->verify_message(message, signature), true);
 
       check_invalid_signatures(result, *verifier, message, signature);
+
+      result.test_eq("KAT signature valid (try 2)", verifier->verify_message(message, signature), true);
+
       verifiers.push_back(std::move(verifier));
       }