Implement VAHE serialization

willow/proto/zk/proofs.proto

@@ -37,3 +37,7 @@ message RlweRelationProofProto {
   repeated bytes z_vw = 8;
   LinearInnerProductProofProto lip_proof = 9;
 }
+
+message RlweRelationProofListProto {
+  repeated RlweRelationProofProto proofs = 1;
+}

willow/src/shell/BUILD

@@ -183,15 +183,19 @@ rust_library(
     deps = [
         ":ahe_shell",
         ":single_thread_hkdf",
+        "@protobuf//rust:protobuf",
         "@crate_index//:merlin",
         "//shell_wrapper:ahe",
         "//shell_wrapper:ahe_cxx",  # fixdeps: keep
         "//shell_wrapper:shell_types",
         "//shell_wrapper:status",
+        "//willow/proto/zk:proofs_rust_proto",
         "//willow/src/traits:ahe_traits",
+        "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:vahe_traits",
         "//willow/src/traits:zk_traits",
         "//willow/src/zk:rlwe_relation",
+        "//willow/src/zk:rlwe_relation_serialization",
     ],
 )
 
@@ -207,6 +211,7 @@ rust_test(
         "//willow/src/testing_utils:shell_testing_parameters",
         "//willow/src/traits:ahe_traits",
         "//willow/src/traits:prng_traits",
+        "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:vahe_traits",
         "//willow/src/traits:zk_traits",
         "//willow/src/zk:rlwe_relation",

willow/src/shell/vahe.rs

@@ -16,7 +16,11 @@ use ahe_shell::{ShellAhe, ShellAheConfig};
 use ahe_traits::Recover as AheRecover;
 use ahe_traits::{AheBase, AheKeygen, PartialDec};
 use merlin::Transcript as MerlinTranscript;
+use proofs_rust_proto::{RlweRelationProofListProto, RlweRelationProofProto};
+use proto_serialization_traits::{FromProto, ToProto};
+use protobuf::{proto, AsView};
 use rlwe_relation::{RlweRelationProof, RlweRelationProver, RlweRelationVerifier};
+use rlwe_relation_serialization::{rlwe_relation_proof_from_proto, rlwe_relation_proof_to_proto};
 use single_thread_hkdf::{compute_hkdf, Seed};
 use status::Status;
 use status::StatusError;
@@ -29,6 +33,68 @@ use zk_traits::{
     ZeroKnowledgeVerifier,
 };
 
+/// Wrapper for key generation proof to support serialization.
+#[derive(Clone)]
+pub struct ShellKeyGenProof(pub RlweRelationProof);
+
+impl ToProto for ShellKeyGenProof {
+    type Proto = RlweRelationProofProto;
+
+    fn to_proto(&self, _ctx: ()) -> Result<Self::Proto, StatusError> {
+        Ok(rlwe_relation_proof_to_proto(&self.0))
+    }
+}
+
+impl FromProto for ShellKeyGenProof {
+    type Proto = RlweRelationProofProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        _ctx: (),
+    ) -> Result<Self, StatusError> {
+        Ok(ShellKeyGenProof(rlwe_relation_proof_from_proto(proto)?))
+    }
+}
+
+/// Wrapper for encryption proof to support serialization.
+#[derive(Clone)]
+pub struct ShellEncryptionProof(pub Vec<RlweRelationProof>);
+
+/// Wrapper for partial decryption proof to support serialization.
+#[derive(Clone)]
+pub struct ShellPartialDecProof(pub Vec<RlweRelationProof>);
+
+macro_rules! impl_proof_list_serialization {
+    ($wrapper_ty:ident) => {
+        impl ToProto for $wrapper_ty {
+            type Proto = RlweRelationProofListProto;
+
+            fn to_proto(&self, _ctx: ()) -> Result<Self::Proto, StatusError> {
+                Ok(proto!(RlweRelationProofListProto {
+                    proofs: self.0.iter().map(|p| rlwe_relation_proof_to_proto(p))
+                }))
+            }
+        }
+
+        impl FromProto for $wrapper_ty {
+            type Proto = RlweRelationProofListProto;
+
+            fn from_proto(
+                proto: impl AsView<Proxied = Self::Proto>,
+                _ctx: (),
+            ) -> Result<Self, StatusError> {
+                let proto = proto.as_view();
+                let proofs: Result<Vec<_>, _> =
+                    proto.proofs().iter().map(|p| rlwe_relation_proof_from_proto(p)).collect();
+                Ok($wrapper_ty(proofs?))
+            }
+        }
+    };
+}
+
+impl_proof_list_serialization!(ShellEncryptionProof);
+impl_proof_list_serialization!(ShellPartialDecProof);
+
 /// Base type holding public VAHE configuration and C++ parameters.
 pub struct ShellVahe {
     ahe: ShellAhe,
@@ -53,6 +119,20 @@ impl ShellVahe {
     }
 }
 
+impl AsRef<ShellVahe> for ShellVahe {
+    fn as_ref(&self) -> &ShellVahe {
+        self
+    }
+}
+
+// Allows calling to_proto and from_proto on the underlying AHE keys and ciphertexts with ShellVahe
+// as the context.
+impl AsRef<ShellAhe> for ShellVahe {
+    fn as_ref(&self) -> &ShellAhe {
+        &self.ahe
+    }
+}
+
 impl AheBase for ShellVahe {
     // This entire implementation is just a simulation of inheritance from ShellAhe.
     type KeyGenMetadata = <ShellAhe as AheBase>::KeyGenMetadata;
@@ -148,9 +228,9 @@ impl AheBase for ShellVahe {
 }
 
 impl VaheBase for ShellVahe {
-    type KeyGenProof = RlweRelationProof;
-    type EncryptionProof = Vec<RlweRelationProof>;
-    type PartialDecProof = Vec<RlweRelationProof>;
+    type KeyGenProof = ShellKeyGenProof;
+    type EncryptionProof = ShellEncryptionProof;
+    type PartialDecProof = ShellPartialDecProof;
 }
 
 impl VerifiableKeyGen for ShellVahe {
@@ -176,16 +256,12 @@ impl VerifiableKeyGen for ShellVahe {
         let (mut transcript, proof_seed) = self.get_transcript_and_proof_seed(b"key_gen")?;
         let prover = RlweRelationProver::new(proof_seed.as_bytes(), self.ahe.num_coeffs());
         let key_gen_proof = prover.prove(&rlwe_statement, &rlwe_witness, &mut transcript)?;
-        Ok((sk_share, pk_share_b, key_gen_proof))
+        Ok((sk_share, pk_share_b, ShellKeyGenProof(key_gen_proof)))
     }
 }
 
 impl KeyGenVerify for ShellVahe {
-    fn verify_key_gen(
-        &self,
-        proof: &RlweRelationProof,
-        key_share: &Self::PublicKeyShare,
-    ) -> Status {
+    fn verify_key_gen(&self, proof: &ShellKeyGenProof, key_share: &Self::PublicKeyShare) -> Status {
         let statement = RlweRelationProofStatement {
             n: self.ahe.num_coeffs(),
             context: self.ahe.rns_context(),
@@ -199,7 +275,7 @@ impl KeyGenVerify for ShellVahe {
 
         let (mut transcript, proof_seed) = self.get_transcript_and_proof_seed(b"key_gen")?;
         let verifier = RlweRelationVerifier::new(proof_seed.as_bytes(), self.ahe.num_coeffs());
-        verifier.verify(&statement, &proof, &mut transcript)
+        verifier.verify(&statement, &proof.0, &mut transcript)
     }
 }
 
@@ -244,19 +320,19 @@ impl VerifiableEncrypt for ShellVahe {
             };
             proof.push(prover.prove(&rlwe_statement, &rlwe_witness, &mut transcript)?);
         }
-        Ok((ciphertext, proof))
+        Ok((ciphertext, ShellEncryptionProof(proof)))
     }
 }
 
 impl EncryptVerify for ShellVahe {
     fn verify_encrypt(
         &self,
-        proof: &Vec<RlweRelationProof>,
+        proof: &ShellEncryptionProof,
         ciphertext_component_a: &Self::PartialDecCiphertext,
         nonce: &[u8],
     ) -> Status {
         let num_polynomials = ciphertext_component_a.0.len();
-        if proof.len() != num_polynomials {
+        if proof.0.len() != num_polynomials {
             return Err(status::permission_denied(
                 "Invalid proof. Proof length does not match number of polynomials in ciphertext.",
             ));
@@ -276,7 +352,7 @@ impl EncryptVerify for ShellVahe {
                 bound_r: 1,
                 bound_e: 16,
             };
-            verifier.verify(&statement, &proof[i], &mut transcript)?;
+            verifier.verify(&statement, &proof.0[i], &mut transcript)?;
         }
         Ok(())
     }
@@ -318,19 +394,19 @@ impl VerifiablePartialDec for ShellVahe {
                 RlweRelationProofWitness { r: &sk.0, e: &errors[i], v: &wraparounds[i] };
             proof.push(prover.prove(&rlwe_statement, &rlwe_witness, &mut transcript)?);
         }
-        Ok((pd, proof))
+        Ok((pd, ShellPartialDecProof(proof)))
     }
 }
 
 impl PartialDecVerify for ShellVahe {
     fn verify_partial_dec(
         &self,
-        proof: &Vec<RlweRelationProof>,
+        proof: &ShellPartialDecProof,
         ct_1: &Self::PartialDecCiphertext,
         pd: &Self::PartialDecryption,
     ) -> Status {
         let num_polynomials = pd.0.len();
-        if proof.len() != num_polynomials {
+        if proof.0.len() != num_polynomials {
             return Err(status::permission_denied(
                 "Invalid proof. Proof length does not match number of polynomials in decryption.",
             ));
@@ -350,7 +426,7 @@ impl PartialDecVerify for ShellVahe {
                 bound_r: 1,
                 bound_e: self.ahe.flood_bound()?,
             };
-            verifier.verify(&statement, &proof[i], &mut transcript)?;
+            verifier.verify(&statement, &proof.0[i], &mut transcript)?;
         }
         Ok(())
     }
@@ -399,6 +475,7 @@ mod test {
 
     use googletest::gtest;
     use prng_traits::SecurePrng;
+    use proto_serialization_traits::{FromProto, ToProto};
     use shell_testing_parameters::make_ahe_config;
     use single_thread_hkdf::SingleThreadHkdfPrng;
 
@@ -414,6 +491,26 @@ mod test {
         vahe.verify_key_gen(&key_gen_proof, &pk_share)?;
         Ok(())
     }
+    #[gtest]
+    fn test_verifiable_key_gen_with_serialization() -> 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, key_gen_proof) = vahe.verifiable_key_gen(&mut prng)?;
+
+        // Serialize and deserialize the public key share.
+        let pk_share_proto = pk_share.to_proto(&vahe)?;
+        let pk_share_deserialized =
+            <ShellVahe as AheBase>::PublicKeyShare::from_proto(pk_share_proto, &vahe)?;
+
+        // Serialize and deserialize the proof.
+        let key_gen_proof_proto = key_gen_proof.to_proto(())?;
+        let key_gen_proof_deserialized =
+            <ShellVahe as VaheBase>::KeyGenProof::from_proto(key_gen_proof_proto, ())?;
+
+        vahe.verify_key_gen(&key_gen_proof_deserialized, &pk_share_deserialized)?;
+        Ok(())
+    }
 
     #[gtest]
     fn test_verifiable_key_gen_with_bad_proof() -> googletest::Result<()> {
@@ -441,6 +538,36 @@ mod test {
         Ok(())
     }
 
+    #[gtest]
+    fn test_verifiable_encrypt_with_serialization() -> 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 pk = vahe.aggregate_public_key_shares(&[pk_share])?;
+
+        // Serialize and deserialize the public key.
+        let pk_proto = pk.to_proto(&vahe)?;
+        let pk_deserialized = <ShellVahe as AheBase>::PublicKey::from_proto(pk_proto, &vahe)?;
+
+        let plaintext = vec![47i64; 8];
+        let (ciphertext, proof) =
+            vahe.verifiable_encrypt(&plaintext, &pk_deserialized, NONCE, &mut prng)?;
+
+        // Serialize and deserialize the ciphertext.
+        let ciphertext_proto = ciphertext.to_proto(&vahe)?;
+        let ciphertext_deserialized =
+            <ShellVahe as AheBase>::Ciphertext::from_proto(ciphertext_proto, &vahe)?;
+
+        // Serialize and deserialize the proof.
+        let proof_proto = proof.to_proto(())?;
+        let proof_deserialized =
+            <ShellVahe as VaheBase>::EncryptionProof::from_proto(proof_proto, ())?;
+
+        vahe.verify_encrypt(&proof_deserialized, &ciphertext_deserialized.component_a, NONCE)?;
+        Ok(())
+    }
+
     #[gtest]
     fn test_verifiable_encrypt_long_plaintext() -> googletest::Result<()> {
         let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING)?;
@@ -479,7 +606,7 @@ mod test {
         let pk = vahe.aggregate_public_key_shares(&[pk_share])?;
         let plaintext = vec![47i64; 8];
         let (ciphertext, mut proof) = vahe.verifiable_encrypt(&plaintext, &pk, NONCE, &mut prng)?;
-        proof.push(key_gen_proof);
+        proof.0.push(key_gen_proof.0);
         let status = vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE);
         assert!(status.is_err());
         Ok(())
@@ -494,7 +621,7 @@ mod test {
         let pk = vahe.aggregate_public_key_shares(&[pk_share])?;
         let plaintext = vec![47i64; 8];
         let (ciphertext, mut proof) = vahe.verifiable_encrypt(&plaintext, &pk, NONCE, &mut prng)?;
-        proof[0] = key_gen_proof;
+        proof.0[0] = key_gen_proof.0;
         let status = vahe.verify_encrypt(&proof, &ciphertext.component_a, NONCE);
         assert!(status.is_err());
         Ok(())
@@ -530,6 +657,30 @@ mod test {
         Ok(())
     }
 
+    #[gtest]
+    fn test_verifiable_partial_dec_with_serialization() -> googletest::Result<()> {
+        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING)?;
+        let seed = SingleThreadHkdfPrng::generate_seed()?;
+        let mut prng = SingleThreadHkdfPrng::create(&seed)?;
+        let (sk_share, pk_share, _) = vahe.verifiable_key_gen(&mut prng)?;
+        let pk = vahe.aggregate_public_key_shares(&[pk_share])?;
+        let plaintext = vec![47i64; 8];
+        let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk, NONCE, &mut prng)?;
+        let (pd, proof) =
+            vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
+
+        // Serialize and deserialize the partial decryption and the proof.
+        let proof_proto = proof.to_proto(())?;
+        let pd_proto = pd.to_proto(&vahe)?;
+        let proof_deserialized =
+            <ShellVahe as VaheBase>::PartialDecProof::from_proto(proof_proto, ())?;
+        let pd_deserialized =
+            <ShellVahe as AheBase>::PartialDecryption::from_proto(pd_proto, &vahe)?;
+
+        vahe.verify_partial_dec(&proof_deserialized, &ciphertext.component_a, &pd_deserialized)?;
+        Ok(())
+    }
+
     #[gtest]
     fn test_verifiable_partial_dec_with_bad_length_proof() -> googletest::Result<()> {
         let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING)?;
@@ -541,7 +692,7 @@ mod test {
         let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk, NONCE, &mut prng)?;
         let (pd, mut proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
-        proof.push(key_gen_proof);
+        proof.0.push(key_gen_proof.0);
         let status = vahe.verify_partial_dec(&proof, &ciphertext.component_a, &pd);
         assert!(status.is_err());
         Ok(())
@@ -558,7 +709,7 @@ mod test {
         let (ciphertext, _) = vahe.verifiable_encrypt(&plaintext, &pk, NONCE, &mut prng)?;
         let (pd, mut proof) =
             vahe.verifiable_partial_dec(&ciphertext.component_a, &sk_share, &mut prng)?;
-        proof[0] = key_gen_proof;
+        proof.0[0] = key_gen_proof.0;
         let status = vahe.verify_partial_dec(&proof, &ciphertext.component_a, &pd);
         assert!(status.is_err());
         Ok(())