Add nonce to VAHE and client message

willow/src/shell/single_thread_hkdf.rs

@@ -45,7 +45,7 @@ mod test {
     use prng_traits::SecurePrng;
 
     #[gtest]
-    /// Two sequences of 8 random bytes should be different (w.h.p).
+    /// Two sequences of 8 random bits should be different (w.h.p).
     fn test_rand8() -> googletest::Result<()> {
         let mut equal = true;
         let seed = SingleThreadHkdfPrng::generate_seed()?;

willow/src/shell/vahe.rs

@@ -208,6 +208,7 @@ impl VerifiableEncrypt for ShellVahe {
         &self,
         plaintext: &Self::Plaintext,
         pk: &Self::PublicKey,
+        nonce: &[u8],
         prng: &mut Self::Rng,
     ) -> Result<(Self::Ciphertext, Self::EncryptionProof), StatusError> {
         let (ciphertext, metadata, wraparounds) =
@@ -222,6 +223,7 @@ impl VerifiableEncrypt for ShellVahe {
         }
 
         let (mut transcript, proof_seed) = self.get_transcript_and_proof_seed(b"encryption")?;
+        transcript.append_message(b"nonce:", nonce);
         let prover = RlweRelationProver::new(proof_seed.as_bytes(), self.ahe.num_coeffs());
         let mut proof = vec![];
         for i in 0..num_polynomials {
@@ -251,6 +253,7 @@ impl EncryptVerify for ShellVahe {
         &self,
         proof: &Vec<RlweRelationProof>,
         ciphertext_component_a: &Self::PartialDecCiphertext,
+        nonce: &[u8],
     ) -> Status {
         let num_polynomials = ciphertext_component_a.0.len();
         if proof.len() != num_polynomials {
@@ -260,6 +263,7 @@ impl EncryptVerify for ShellVahe {
         }
 
         let (mut transcript, proof_seed) = self.get_transcript_and_proof_seed(b"encryption")?;
+        transcript.append_message(b"nonce:", nonce);
         let verifier = RlweRelationVerifier::new(proof_seed.as_bytes(), self.ahe.num_coeffs());
         for i in 0..num_polynomials {
             let statement = RlweRelationProofStatement {
@@ -399,6 +403,7 @@ mod test {
     use single_thread_hkdf::SingleThreadHkdfPrng;
 
     const CONTEXT_STRING: &[u8] = b"testing_context_string";
+    const NONCE: &[u8] = b"0123456789ABCDEF";
 
     #[gtest]
     fn test_verifiable_key_gen() -> googletest::Result<()> {
@@ -430,8 +435,9 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (_, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, proof) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
-        vahe.verify_encrypt(&proof, &ciphertext.component_a)?;
+        let (ciphertext, proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
+        vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE)?;
         Ok(())
     }
 
@@ -442,8 +448,25 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (_, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 256];
-        let (ciphertext, proof) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
-        vahe.verify_encrypt(&proof, &ciphertext.component_a)?;
+        let (ciphertext, proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
+        vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE)?;
+        Ok(())
+    }
+
+    #[gtest]
+    fn test_verifiable_encrypt_with_bad_nonce() -> googletest::Result<()> {
+        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING)?;
+        let seed = SingleThreadHkdfPrng::generate_seed()?;
+        let mut prng = SingleThreadHkdfPrng::create(&seed)?;
+        let (_, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
+        let plaintext = vec![47i64; 8];
+        let (ciphertext, proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
+        let bad_nonce = b"BADBADBADBAD";
+        let status = vahe.verify_encrypt(&proof, &ciphertext.component_a, bad_nonce);
+        // bad_nonce doesn't match NONCE, so the proof verification should fail.
+        assert!(status.is_err());
         Ok(())
     }
 
@@ -454,9 +477,10 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (_, pk_share, key_gen_proof) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, mut proof) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, mut proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         proof.push(key_gen_proof);
-        let status = vahe.verify_encrypt(&proof, &ciphertext.component_a);
+        let status = vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE);
         assert!(status.is_err());
         Ok(())
     }
@@ -468,9 +492,10 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (_, pk_share, key_gen_proof) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, mut proof) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, mut proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         proof[0] = key_gen_proof;
-        let status = vahe.verify_encrypt(&proof, &ciphertext.component_a);
+        let status = vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE);
         assert!(status.is_err());
         Ok(())
     }
@@ -482,7 +507,7 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (sk_share, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         let (pd, proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
         vahe.verify_partial_dec(&proof, &ciphertext.component_a, &pd)?;
@@ -496,7 +521,7 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (sk_share, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 256];
-        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         let (pd, proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
         vahe.verify_partial_dec(&proof, &ciphertext.component_a, &pd)?;
@@ -510,7 +535,7 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (sk_share, pk_share, key_gen_proof) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         let (pd, mut proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
         proof.push(key_gen_proof);
@@ -526,7 +551,7 @@ mod test {
         let mut prng = SingleThreadHkdfPrng::create(&seed)?;
         let (sk_share, pk_share, key_gen_proof) = vahe.verifiable_key_gen(&mut prng)?;
         let plaintext = vec![47i64; 8];
-        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, &mut prng)?;
+        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk_share, NONCE, &mut prng)?;
         let (pd, mut proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
         proof[0] = key_gen_proof;

willow/src/traits/ahe.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use prng_traits::SecurePrng;
 use status::StatusError;
 
 /// Base trait for (Asymmetric) multiparty Additive Homomorphic Encryption (AHE)
@@ -94,7 +95,7 @@ pub trait AheBase: Sized {
     ) -> Result<(), StatusError>;
 
     /// Randomness source, typically a SecurePrng.
-    type Rng;
+    type Rng: SecurePrng;
 }
 
 pub trait AheKeygen: AheBase {

willow/src/traits/kahe.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 //! Traits for Key Additive Homomorphic Encryption (KAHE) schemes.
+use prng_traits::SecurePrng;
 use status::StatusError;
 
 /// Base trait for KAHE primitives, containing types that are shared across all
@@ -54,7 +55,7 @@ pub trait KaheBase: Sized {
     ) -> Result<(), StatusError>;
 
     /// Randomness source, typically a SecurePrng.
-    type Rng;
+    type Rng: SecurePrng;
 }
 
 /// Key generation

willow/src/traits/vahe.rs

@@ -22,35 +22,51 @@ pub trait VaheBase: AheBase + Sized {
 }
 
 pub trait VerifiableKeyGen: VaheBase {
+    /// Generate a secret key and a public key.
+    ///
+    /// Returns the secret key share, the public key share, and a proof that the
+    /// key generation was correct.
     fn verifiable_key_gen(
         &self,
         prng: &mut Self::Rng,
     ) -> Result<(Self::SecretKeyShare, Self::PublicKeyShare, Self::KeyGenProof), StatusError>;
 }
 
 pub trait KeyGenVerify: VaheBase {
+    /// Verify that the key generation proof is valid.
     fn verify_key_gen(&self, proof: &Self::KeyGenProof, key_share: &Self::PublicKeyShare)
         -> Status;
 }
 
 pub trait VerifiableEncrypt: VaheBase {
+    /// Encrypt a plaintext with a given public key.
+    ///
+    /// `nonce` is a unique identifier for this encryption. It must be used
+    /// when verifying the encryption proof.
     fn verifiable_encrypt(
         &self,
         plaintext: &Self::Plaintext,
         pk: &Self::PublicKey,
+        nonce: &[u8],
         prng: &mut Self::Rng,
     ) -> Result<(Self::Ciphertext, Self::EncryptionProof), StatusError>;
 }
 
 pub trait EncryptVerify: VaheBase {
+    /// Verify that the encryption proof is valid.
+    ///
+    /// `nonce` must match the nonce passed to `verifiable_encrypt`.
     fn verify_encrypt(
         &self,
         proof: &Self::EncryptionProof,
         ciphertext: &Self::PartialDecCiphertext,
+        nonce: &[u8],
     ) -> Status;
 }
 
 pub trait VerifiablePartialDec: VaheBase {
+    /// Decrypt a ciphertext with a given secret key. Returns the partial
+    /// decryption and a proof that the decryption was correct.
     fn verifiable_partial_dec(
         &self,
         ct_1: &Self::PartialDecCiphertext,
@@ -60,6 +76,7 @@ pub trait VerifiablePartialDec: VaheBase {
 }
 
 pub trait PartialDecVerify: VaheBase {
+    /// Verify that the partial decryption proof is valid.
     fn verify_partial_dec(
         &self,
         proof: &Self::PartialDecProof,

willow/src/willow_v1/BUILD

@@ -45,6 +45,7 @@ rust_library(
         "//willow/src/traits:client_traits",
         "//willow/src/traits:common_traits",
         "//willow/src/traits:kahe_traits",
+        "//willow/src/traits:prng_traits",
         "//willow/src/traits:vahe_traits",
     ],
 )

willow/src/willow_v1/client.rs

@@ -15,6 +15,7 @@
 use ahe_traits::AheBase;
 use client_traits::SecureAggregationClient;
 use kahe_traits::{KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
+use prng_traits::SecurePrng;
 use vahe_traits::{VaheBase, VerifiableEncrypt};
 use willow_v1_common::{WillowClientMessage, WillowCommon};
 
@@ -51,13 +52,18 @@ where
         let ahe_plaintext: Vahe::Plaintext =
             self.common.kahe.try_secret_key_into(kahe_secret_key)?;
 
+        // Generate a nonce for the VAHE encryption.
+        let nonce =
+            (0..16).map(|_| self.prng.rand8()).collect::<Result<Vec<u8>, status::StatusError>>()?;
+
         // Encrypt AHE plaintext with public key.
         let (ahe_ciphertext, proof) = self.common.vahe.verifiable_encrypt(
             &ahe_plaintext,
             signed_public_key,
+            &nonce,
             &mut self.prng,
         )?;
-        Ok(WillowClientMessage { kahe_ciphertext, ahe_ciphertext, proof })
+        Ok(WillowClientMessage { kahe_ciphertext, ahe_ciphertext, proof, nonce })
     }
 }
 

willow/src/willow_v1/common.rs

@@ -35,6 +35,7 @@ pub struct WillowClientMessage<Kahe: KaheBase, Vahe: VaheBase> {
     pub kahe_ciphertext: Kahe::Ciphertext,
     pub ahe_ciphertext: Vahe::Ciphertext,
     pub proof: Vahe::EncryptionProof,
+    pub nonce: Vec<u8>,
 }
 
 // Partial decryption request is an aggregated AHE ciphertext.
@@ -65,7 +66,7 @@ pub struct CiphertextContribution<Kahe: KaheBase, Vahe: VaheBase> {
 pub struct DecryptionRequestContribution<Vahe: VaheBase> {
     pub partial_dec_ciphertext: Vahe::PartialDecCiphertext,
     pub proof: Vahe::EncryptionProof,
-    // pub client_id: u128,
+    pub nonce: Vec<u8>,
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase> SecureAggregationCommon for WillowCommon<Kahe, Vahe> {

willow/src/willow_v1/server.rs

@@ -91,7 +91,11 @@ where
                 kahe_ciphertext: client_message.kahe_ciphertext,
                 ahe_recover_ciphertext,
             },
-            DecryptionRequestContribution { partial_dec_ciphertext, proof: client_message.proof },
+            DecryptionRequestContribution {
+                partial_dec_ciphertext,
+                proof: client_message.proof,
+                nonce: client_message.nonce,
+            },
         ))
     }
 

willow/src/willow_v1/verifier.rs

@@ -48,9 +48,11 @@ where
         contribution: DecryptionRequestContribution<Vahe>,
         state: &mut Self::VerifierState,
     ) -> Result<(), status::StatusError> {
-        self.common
-            .vahe
-            .verify_encrypt(&contribution.proof, &contribution.partial_dec_ciphertext)?;
+        self.common.vahe.verify_encrypt(
+            &contribution.proof,
+            &contribution.partial_dec_ciphertext,
+            &contribution.nonce,
+        )?;
         if let Some(ref mut sum) = state.partial_dec_ciphertext_sum {
             self.common
                 .vahe