Add traits for accessing VAHE/KAHE from role traits

willow/src/traits/ahe.rs

@@ -100,6 +100,12 @@ pub trait AheBase: Sized {
     type Rng: SecurePrng;
 }
 
+/// Accessor trait for composition.
+pub trait HasAhe {
+    type Ahe: AheBase;
+    fn ahe(&self) -> &Self::Ahe;
+}
+
 pub trait AheKeygen: AheBase {
     /// Sample a new secret key and public key share.
     fn key_gen(

willow/src/traits/client.rs

@@ -12,13 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use kahe_traits::KaheBase;
+use kahe_traits::HasKahe;
 use messages::{ClientMessage, DecryptorPublicKey};
 use status::StatusError;
-use vahe_traits::VaheBase;
+use vahe_traits::HasVahe;
 
 /// Base trait for the secure aggregation Client.
-pub trait SecureAggregationClient<Kahe: KaheBase, Vahe: VaheBase> {
+pub trait SecureAggregationClient: HasKahe + HasVahe {
     /// The plaintext to be aggregated.
     type Plaintext;
     type PlaintextSlice<'a>;
@@ -28,7 +28,7 @@ pub trait SecureAggregationClient<Kahe: KaheBase, Vahe: VaheBase> {
     fn create_client_message(
         &mut self,
         plaintext: &Self::PlaintextSlice<'_>,
-        signed_public_key: &DecryptorPublicKey<Vahe>,
+        signed_public_key: &DecryptorPublicKey<<Self as HasVahe>::Vahe>,
         nonce: &[u8],
-    ) -> Result<ClientMessage<Kahe, Vahe>, StatusError>;
+    ) -> Result<ClientMessage<<Self as HasKahe>::Kahe, <Self as HasVahe>::Vahe>, StatusError>;
 }

willow/src/traits/decryptor.rs

@@ -14,10 +14,10 @@
 
 use messages::{DecryptorPublicKeyShare, PartialDecryptionRequest, PartialDecryptionResponse};
 use status::StatusError;
-use vahe_traits::VaheBase;
+use vahe_traits::HasVahe;
 
 /// Base trait for the Decryptor.
-pub trait SecureAggregationDecryptor<Vahe: VaheBase> {
+pub trait SecureAggregationDecryptor: HasVahe {
     /// The state held by the Decryptor between messages.
     type DecryptorState: Default;
 
@@ -26,13 +26,13 @@ pub trait SecureAggregationDecryptor<Vahe: VaheBase> {
     fn create_public_key_share(
         &mut self,
         decryptor_state: &mut Self::DecryptorState,
-    ) -> Result<DecryptorPublicKeyShare<Vahe>, StatusError>;
+    ) -> Result<DecryptorPublicKeyShare<<Self as HasVahe>::Vahe>, StatusError>;
 
     /// Handles a partial decryption request received from the Server. Returns a
     /// partial decryption to the Server.
     fn handle_partial_decryption_request(
         &mut self,
-        partial_decryption_request: PartialDecryptionRequest<Vahe>,
+        partial_decryption_request: PartialDecryptionRequest<<Self as HasVahe>::Vahe>,
         decryptor_state: &Self::DecryptorState,
-    ) -> Result<PartialDecryptionResponse<Vahe>, StatusError>;
+    ) -> Result<PartialDecryptionResponse<<Self as HasVahe>::Vahe>, StatusError>;
 }

willow/src/traits/kahe.rs

@@ -60,6 +60,12 @@ pub trait KaheBase: Sized {
     type Rng: SecurePrng;
 }
 
+/// Accessor trait for composition.
+pub trait HasKahe {
+    type Kahe: KaheBase;
+    fn kahe(&self) -> &Self::Kahe;
+}
+
 /// Key generation
 pub trait KaheKeygen: KaheBase {
     /// Sample a new secret key.

willow/src/traits/server.rs

@@ -12,18 +12,22 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use kahe_traits::KaheBase;
+use kahe_traits::{HasKahe, KaheBase};
 use messages::{
     CiphertextContribution, ClientMessage, DecryptionRequestContribution, DecryptorPublicKey,
     DecryptorPublicKeyShare, PartialDecryptionResponse,
 };
 use status::StatusError;
-use vahe_traits::VaheBase;
+use vahe_traits::{HasVahe, VaheBase};
+
+// Helper aliases for the generic types.
+type Kahe<T> = <T as HasKahe>::Kahe;
+type Vahe<T> = <T as HasVahe>::Vahe;
 
 /// Base trait for the secure aggregation server. Also includes the Coordinator
 /// functionality of the threshold AHE scheme.
 ///
-pub trait SecureAggregationServer<Kahe: KaheBase, Vahe: VaheBase> {
+pub trait SecureAggregationServer: HasKahe + HasVahe {
     /// The state held by the server between messages.
     type ServerState: Default + Clone;
     /// The result of the aggregation.
@@ -33,7 +37,7 @@ pub trait SecureAggregationServer<Kahe: KaheBase, Vahe: VaheBase> {
     /// server state.
     fn handle_decryptor_public_key_share(
         &self,
-        key_share: DecryptorPublicKeyShare<Vahe>,
+        key_share: DecryptorPublicKeyShare<Vahe<Self>>,
         decryptor_id: &str,
         server_state: &mut Self::ServerState,
     ) -> Result<(), StatusError>;
@@ -43,30 +47,30 @@ pub trait SecureAggregationServer<Kahe: KaheBase, Vahe: VaheBase> {
     fn create_decryptor_public_key(
         &self,
         server_state: &Self::ServerState,
-    ) -> Result<DecryptorPublicKey<Vahe>, StatusError>;
+    ) -> Result<DecryptorPublicKey<Vahe<Self>>, StatusError>;
 
     /// Splits a client message into the ciphertext contribution and the
     /// decryption request contribution.
     fn split_client_message(
         &self,
-        client_message: ClientMessage<Kahe, Vahe>,
+        client_message: ClientMessage<Kahe<Self>, Vahe<Self>>,
     ) -> Result<
-        (CiphertextContribution<Kahe, Vahe>, DecryptionRequestContribution<Vahe>),
+        (CiphertextContribution<Kahe<Self>, Vahe<Self>>, DecryptionRequestContribution<Vahe<Self>>),
         StatusError,
     >;
 
     /// Handles a single client message, updating the server state.
     fn handle_ciphertext_contribution(
         &self,
-        ciphertext_contribution: CiphertextContribution<Kahe, Vahe>,
+        ciphertext_contribution: CiphertextContribution<Kahe<Self>, Vahe<Self>>,
         server_state: &mut Self::ServerState,
     ) -> Result<(), StatusError>;
 
     /// Handles a partial decryption received from a Decryptor, updating the
     /// server state.
     fn handle_partial_decryption(
         &self,
-        partial_decryption_response: PartialDecryptionResponse<Vahe>,
+        partial_decryption_response: PartialDecryptionResponse<Vahe<Self>>,
         server_state: &mut Self::ServerState,
     ) -> Result<(), StatusError>;
 

willow/src/traits/vahe.rs

@@ -21,6 +21,12 @@ pub trait VaheBase: AheBase + Sized {
     type PartialDecProof;
 }
 
+/// Accessor trait for composition.
+pub trait HasVahe {
+    type Vahe: VaheBase;
+    fn vahe(&self) -> &Self::Vahe;
+}
+
 pub trait VerifiableKeyGen: VaheBase {
     /// Generate a secret key and a public key.
     ///

willow/src/traits/verifier.rs

@@ -14,18 +14,18 @@
 
 use messages::{DecryptionRequestContribution, PartialDecryptionRequest};
 use status::StatusError;
-use vahe_traits::VaheBase;
+use vahe_traits::HasVahe;
 
 /// Base trait for the secure aggregation verifier.
 ///
-pub trait SecureAggregationVerifier<Vahe: VaheBase> {
+pub trait SecureAggregationVerifier: HasVahe {
     /// The state held by the verifier between messages.
     type VerifierState: Default;
 
     /// Verifies a clients decryption request contribution.
     fn verify_and_include(
         &self,
-        contribution: DecryptionRequestContribution<Vahe>,
+        contribution: DecryptionRequestContribution<<Self as HasVahe>::Vahe>,
         state: &mut Self::VerifierState,
     ) -> Result<(), StatusError>;
 
@@ -41,5 +41,5 @@ pub trait SecureAggregationVerifier<Vahe: VaheBase> {
     fn create_partial_decryption_request(
         &self,
         state: Self::VerifierState,
-    ) -> Result<PartialDecryptionRequest<Vahe>, StatusError>;
+    ) -> Result<PartialDecryptionRequest<<Self as HasVahe>::Vahe>, StatusError>;
 }

willow/src/willow_v1/client.rs

@@ -13,10 +13,10 @@
 // limitations under the License.
 
 use client_traits::SecureAggregationClient;
-use kahe_traits::{KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
+use kahe_traits::{HasKahe, KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
 use messages::{ClientMessage, DecryptorPublicKey};
 use prng_traits::SecurePrng;
-use vahe_traits::{VaheBase, VerifiableEncrypt};
+use vahe_traits::{HasVahe, VaheBase, VerifiableEncrypt};
 
 /// Lightweight client directly exposing KAHE/VAHE types.
 pub struct WillowV1Client<Kahe: KaheBase, Vahe: VaheBase> {
@@ -25,10 +25,24 @@ pub struct WillowV1Client<Kahe: KaheBase, Vahe: VaheBase> {
     pub prng: Kahe::Rng, // Using a single PRNG for both VAHE and KAHE.
 }
 
+impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Client<Kahe, Vahe> {
+    type Kahe = Kahe;
+    fn kahe(&self) -> &Self::Kahe {
+        &self.kahe
+    }
+}
+
+impl<Kahe: KaheBase, Vahe: VaheBase> HasVahe for WillowV1Client<Kahe, Vahe> {
+    type Vahe = Vahe;
+    fn vahe(&self) -> &Self::Vahe {
+        &self.vahe
+    }
+}
+
 /// Implementation of the `SecureAggregationClient` trait for the generic
 /// KAHE/VAHE client, using WillowCommon as the common types (e.g. protocol
 /// messages are directly the AHE public key and ciphertexts).
-impl<Kahe, Vahe> SecureAggregationClient<Kahe, Vahe> for WillowV1Client<Kahe, Vahe>
+impl<Kahe, Vahe> SecureAggregationClient for WillowV1Client<Kahe, Vahe>
 where
     Vahe: VaheBase + VerifiableEncrypt,
     // Reusing the same PRNG for both AHE and KAHE.

willow/src/willow_v1/decryptor.rs

@@ -15,7 +15,7 @@
 use ahe_traits::{AheKeygen, PartialDec};
 use decryptor_traits::SecureAggregationDecryptor;
 use messages::{DecryptorPublicKeyShare, PartialDecryptionRequest, PartialDecryptionResponse};
-use vahe_traits::{EncryptVerify, VaheBase};
+use vahe_traits::{EncryptVerify, HasVahe, VaheBase};
 
 /// Lightweight decryptor directly exposing KAHE/VAHE types. It verifies only the client proofs,
 /// does not provide verifiable partial decryptions.
@@ -24,6 +24,13 @@ pub struct WillowV1Decryptor<Vahe: VaheBase> {
     pub prng: Vahe::Rng,
 }
 
+impl<Vahe: VaheBase> HasVahe for WillowV1Decryptor<Vahe> {
+    type Vahe = Vahe;
+    fn vahe(&self) -> &Self::Vahe {
+        &self.vahe
+    }
+}
+
 pub struct DecryptorState<Vahe: VaheBase> {
     sk_share: Option<Vahe::SecretKeyShare>,
 }
@@ -37,7 +44,7 @@ impl<Vahe: VaheBase> Default for DecryptorState<Vahe> {
 /// Implementation of the `SecureAggregationDecryptor` trait for the generic
 /// KAHE/AHE decryptor, using WillowCommon as the common types (e.g. protocol
 /// messages are directly the AHE public key and ciphertexts).
-impl<Vahe> SecureAggregationDecryptor<Vahe> for WillowV1Decryptor<Vahe>
+impl<Vahe> SecureAggregationDecryptor for WillowV1Decryptor<Vahe>
 where
     Vahe: VaheBase + EncryptVerify + PartialDec + AheKeygen,
 {

willow/src/willow_v1/server.rs

@@ -13,22 +13,36 @@
 // limitations under the License.
 
 use ahe_traits::PartialDec;
-use kahe_traits::{KaheBase, KaheDecrypt, TrySecretKeyFrom};
+use kahe_traits::{HasKahe, KaheBase, KaheDecrypt, TrySecretKeyFrom};
 use messages::{
     CiphertextContribution, ClientMessage, DecryptionRequestContribution, DecryptorPublicKey,
     DecryptorPublicKeyShare, PartialDecryptionResponse,
 };
 use server_traits::SecureAggregationServer;
 use std::collections::HashMap;
-use vahe_traits::{EncryptVerify, Recover, VaheBase};
+use vahe_traits::{EncryptVerify, HasVahe, Recover, VaheBase};
 
 /// The server struct, containing a WillowCommon instance. Only the clients messages are verified,
 /// not the key generation or partial decryptions.
-pub struct WillowV1Server<Kahe, Vahe: VaheBase> {
+pub struct WillowV1Server<Kahe: KaheBase, Vahe: VaheBase> {
     pub kahe: Kahe,
     pub vahe: Vahe,
 }
 
+impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Server<Kahe, Vahe> {
+    type Kahe = Kahe;
+    fn kahe(&self) -> &Self::Kahe {
+        &self.kahe
+    }
+}
+
+impl<Kahe: KaheBase, Vahe: VaheBase> HasVahe for WillowV1Server<Kahe, Vahe> {
+    type Vahe = Vahe;
+    fn vahe(&self) -> &Self::Vahe {
+        &self.vahe
+    }
+}
+
 /// State for the server.
 pub struct ServerState<Kahe: KaheBase, Vahe: VaheBase + PartialDec> {
     /// The public key shares received from Decryptors. The key is the ID of the Decryptor.
@@ -59,7 +73,7 @@ impl<Kahe: KaheBase, Vahe: VaheBase + PartialDec> Clone for ServerState<Kahe, Va
     }
 }
 
-impl<Kahe, Vahe> SecureAggregationServer<Kahe, Vahe> for WillowV1Server<Kahe, Vahe>
+impl<Kahe, Vahe> SecureAggregationServer for WillowV1Server<Kahe, Vahe>
 where
     Vahe: EncryptVerify + PartialDec + Recover,
     Kahe: KaheBase + TrySecretKeyFrom<Vahe::Plaintext> + KaheDecrypt,
@@ -228,17 +242,18 @@ where
                 (None, None) => None,
             };
 
-        merged_server_state.partial_decryption_sum =
-            match (&server_state_1.partial_decryption_sum, &server_state_2.partial_decryption_sum)
-            {
-                (Some(sum1), Some(sum2)) => {
-                    let mut merged_sum = sum1.clone();
-                    self.vahe.add_partial_decryptions_in_place(sum2, &mut merged_sum)?;
-                    Some(merged_sum)
-                }
-                (Some(s), None) | (None, Some(s)) => Some(s.clone()),
-                (None, None) => None,
-            };
+        merged_server_state.partial_decryption_sum = match (
+            &server_state_1.partial_decryption_sum,
+            &server_state_2.partial_decryption_sum,
+        ) {
+            (Some(sum1), Some(sum2)) => {
+                let mut merged_sum = sum1.clone();
+                self.vahe.add_partial_decryptions_in_place(sum2, &mut merged_sum)?;
+                Some(merged_sum)
+            }
+            (Some(s), None) | (None, Some(s)) => Some(s.clone()),
+            (None, None) => None,
+        };
 
         Ok(merged_server_state)
     }

willow/src/willow_v1/verifier.rs

@@ -14,14 +14,21 @@
 
 use messages::{DecryptionRequestContribution, PartialDecryptionRequest};
 use std::fmt::Debug;
-use vahe_traits::{EncryptVerify, VaheBase};
+use vahe_traits::{EncryptVerify, HasVahe, VaheBase};
 use verifier_traits::SecureAggregationVerifier;
 
 /// The verifier struct, containing a WillowCommon instance.
 pub struct WillowV1Verifier<Vahe: VaheBase> {
     pub vahe: Vahe,
 }
 
+impl<Vahe: VaheBase> HasVahe for WillowV1Verifier<Vahe> {
+    type Vahe = Vahe;
+    fn vahe(&self) -> &Self::Vahe {
+        &self.vahe
+    }
+}
+
 // State for the verifier after the first contribution is received.
 struct NonemptyVerifierState<Vahe: VaheBase> {
     partial_dec_ciphertext_sum: Vahe::PartialDecCiphertext,
@@ -90,7 +97,7 @@ impl<Vahe: VaheBase> Clone for VerifierState<Vahe> {
     }
 }
 
-impl<Vahe> SecureAggregationVerifier<Vahe> for WillowV1Verifier<Vahe>
+impl<Vahe> SecureAggregationVerifier for WillowV1Verifier<Vahe>
 where
     Vahe: EncryptVerify,
 {