1 files changed, 99 insertions, 0 deletions
diff --git a/private_join_and_compute/crypto/dodis_yampolskiy_prf/dy_verifiable_random_function_test.cc b/private_join_and_compute/crypto/dodis_yampolskiy_prf/dy_verifiable_random_function_test.cc
index ae16efa..096c23b 100644
--- a/private_join_and_compute/crypto/dodis_yampolskiy_prf/dy_verifiable_random_function_test.cc
+++ b/private_join_and_compute/crypto/dodis_yampolskiy_prf/dy_verifiable_random_function_test.cc
@@ -163,6 +163,105 @@ TEST_F(DyVerifiableRandomFunctionTest, GenerateKeyPairProducesDifferentValues) {
             std::get<1>(key_pair_2).open_commit_prf_key());
 }
 
+TEST_F(DyVerifiableRandomFunctionTest, GenerateKeyProofSucceeds) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  EXPECT_OK(dy_vrf_->VerifyGenerateKeysProof(public_key_proto,
+                                             generate_keys_proof_proto));
+}
+
+TEST_F(DyVerifiableRandomFunctionTest, EmptyGenerateKeyProofFails) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  // Empty proof should fail.
+  EXPECT_THAT(
+      dy_vrf_->VerifyGenerateKeysProof(public_key_proto,
+                                       proto::DyVrfGenerateKeysProof()),
+      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("Failed")));
+}
+
+TEST_F(DyVerifiableRandomFunctionTest, GenerateKeyProofFailsForDifferentKeys) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  // Using this proof with the keys generated by the test fixture should fail.
+  EXPECT_THAT(
+      dy_vrf_->VerifyGenerateKeysProof(public_key_, generate_keys_proof_proto),
+      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("Failed")));
+}
+
+TEST_F(DyVerifiableRandomFunctionTest,
+       GenerateKeyProofFailsWhenPrfCommitmentIsMissing) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  public_key_proto.clear_commit_prf_key();
+  // Technically this proof fails because the verification method fails to
+  // compute a modular inverse.
+  EXPECT_THAT(
+      dy_vrf_->VerifyGenerateKeysProof(public_key_proto,
+                                       generate_keys_proof_proto),
+      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("Inverse")));
+}
+
+TEST_F(DyVerifiableRandomFunctionTest,
+       GenerateKeyProofFailsWhenMaskedDummyPrfKeyIsTooLarge) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  BigNum too_large =
+      ctx_.CreateBigNum(
+              generate_keys_proof_proto.message_2().masked_dummy_prf_key())
+          .Lshift(20);
+
+  generate_keys_proof_proto.mutable_message_2()->set_masked_dummy_prf_key(
+      too_large.ToBytes());
+
+  EXPECT_THAT(dy_vrf_->VerifyGenerateKeysProof(public_key_proto,
+                                               generate_keys_proof_proto),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("masked_dummy_prf_key")));
+}
+
+TEST_F(DyVerifiableRandomFunctionTest,
+       GenerateKeyProofFailsWhenMaskedDummyPrfKeyOpeningIsMissing) {
+  proto::DyVrfPublicKey public_key_proto;
+  proto::DyVrfPrivateKey private_key_proto;
+  proto::DyVrfGenerateKeysProof generate_keys_proof_proto;
+  ASSERT_OK_AND_ASSIGN(
+      std::tie(public_key_proto, private_key_proto, generate_keys_proof_proto),
+      dy_vrf_->GenerateKeyPair());
+
+  generate_keys_proof_proto.mutable_message_2()->clear_masked_dummy_prf_key();
+
+  EXPECT_THAT(
+      dy_vrf_->VerifyGenerateKeysProof(public_key_proto,
+                                       generate_keys_proof_proto),
+      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("Failed")));
+}
+
 TEST_F(DyVerifiableRandomFunctionTest, ApplySucceeds) {
   std::vector<BigNum> messages = {ctx_.CreateBigNum(0),
                                   ctx_.CreateBigNum(5),