diff --git a/willow/src/traits/ahe.rs b/willow/src/traits/ahe.rs index 7dcf317..5fdb85f 100644 --- a/willow/src/traits/ahe.rs +++ b/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( diff --git a/willow/src/traits/client.rs b/willow/src/traits/client.rs index 16657a9..e4fca66 100644 --- a/willow/src/traits/client.rs +++ b/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 { +pub trait SecureAggregationClient: HasKahe + HasVahe { /// The plaintext to be aggregated. type Plaintext; type PlaintextSlice<'a>; @@ -28,7 +28,7 @@ pub trait SecureAggregationClient { fn create_client_message( &mut self, plaintext: &Self::PlaintextSlice<'_>, - signed_public_key: &DecryptorPublicKey, + signed_public_key: &DecryptorPublicKey<::Vahe>, nonce: &[u8], - ) -> Result, StatusError>; + ) -> Result::Kahe, ::Vahe>, StatusError>; } diff --git a/willow/src/traits/decryptor.rs b/willow/src/traits/decryptor.rs index 1d4853d..31f8934 100644 --- a/willow/src/traits/decryptor.rs +++ b/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 { +pub trait SecureAggregationDecryptor: HasVahe { /// The state held by the Decryptor between messages. type DecryptorState: Default; @@ -26,13 +26,13 @@ pub trait SecureAggregationDecryptor { fn create_public_key_share( &mut self, decryptor_state: &mut Self::DecryptorState, - ) -> Result, StatusError>; + ) -> Result::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, + partial_decryption_request: PartialDecryptionRequest<::Vahe>, decryptor_state: &Self::DecryptorState, - ) -> Result, StatusError>; + ) -> Result::Vahe>, StatusError>; } diff --git a/willow/src/traits/kahe.rs b/willow/src/traits/kahe.rs index 911e541..d48f3eb 100644 --- a/willow/src/traits/kahe.rs +++ b/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. diff --git a/willow/src/traits/server.rs b/willow/src/traits/server.rs index 1081496..ee5049f 100644 --- a/willow/src/traits/server.rs +++ b/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 = ::Kahe; +type Vahe = ::Vahe; /// Base trait for the secure aggregation server. Also includes the Coordinator /// functionality of the threshold AHE scheme. /// -pub trait SecureAggregationServer { +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 { /// server state. fn handle_decryptor_public_key_share( &self, - key_share: DecryptorPublicKeyShare, + key_share: DecryptorPublicKeyShare>, decryptor_id: &str, server_state: &mut Self::ServerState, ) -> Result<(), StatusError>; @@ -43,22 +47,22 @@ pub trait SecureAggregationServer { fn create_decryptor_public_key( &self, server_state: &Self::ServerState, - ) -> Result, StatusError>; + ) -> Result>, StatusError>; /// Splits a client message into the ciphertext contribution and the /// decryption request contribution. fn split_client_message( &self, - client_message: ClientMessage, + client_message: ClientMessage, Vahe>, ) -> Result< - (CiphertextContribution, DecryptionRequestContribution), + (CiphertextContribution, Vahe>, DecryptionRequestContribution>), StatusError, >; /// Handles a single client message, updating the server state. fn handle_ciphertext_contribution( &self, - ciphertext_contribution: CiphertextContribution, + ciphertext_contribution: CiphertextContribution, Vahe>, server_state: &mut Self::ServerState, ) -> Result<(), StatusError>; @@ -66,7 +70,7 @@ pub trait SecureAggregationServer { /// server state. fn handle_partial_decryption( &self, - partial_decryption_response: PartialDecryptionResponse, + partial_decryption_response: PartialDecryptionResponse>, server_state: &mut Self::ServerState, ) -> Result<(), StatusError>; diff --git a/willow/src/traits/vahe.rs b/willow/src/traits/vahe.rs index 329c626..f3881cf 100644 --- a/willow/src/traits/vahe.rs +++ b/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. /// diff --git a/willow/src/traits/verifier.rs b/willow/src/traits/verifier.rs index f6c0478..0cc45ba 100644 --- a/willow/src/traits/verifier.rs +++ b/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 { +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, + contribution: DecryptionRequestContribution<::Vahe>, state: &mut Self::VerifierState, ) -> Result<(), StatusError>; @@ -41,5 +41,5 @@ pub trait SecureAggregationVerifier { fn create_partial_decryption_request( &self, state: Self::VerifierState, - ) -> Result, StatusError>; + ) -> Result::Vahe>, StatusError>; } diff --git a/willow/src/willow_v1/client.rs b/willow/src/willow_v1/client.rs index 0b45306..aa89916 100644 --- a/willow/src/willow_v1/client.rs +++ b/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 { @@ -25,10 +25,24 @@ pub struct WillowV1Client { pub prng: Kahe::Rng, // Using a single PRNG for both VAHE and KAHE. } +impl HasKahe for WillowV1Client { + type Kahe = Kahe; + fn kahe(&self) -> &Self::Kahe { + &self.kahe + } +} + +impl HasVahe for WillowV1Client { + 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 SecureAggregationClient for WillowV1Client +impl SecureAggregationClient for WillowV1Client where Vahe: VaheBase + VerifiableEncrypt, // Reusing the same PRNG for both AHE and KAHE. diff --git a/willow/src/willow_v1/decryptor.rs b/willow/src/willow_v1/decryptor.rs index 15ad2fa..7e76977 100644 --- a/willow/src/willow_v1/decryptor.rs +++ b/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 { pub prng: Vahe::Rng, } +impl HasVahe for WillowV1Decryptor { + type Vahe = Vahe; + fn vahe(&self) -> &Self::Vahe { + &self.vahe + } +} + pub struct DecryptorState { sk_share: Option, } @@ -37,7 +44,7 @@ impl Default for DecryptorState { /// 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 SecureAggregationDecryptor for WillowV1Decryptor +impl SecureAggregationDecryptor for WillowV1Decryptor where Vahe: VaheBase + EncryptVerify + PartialDec + AheKeygen, { diff --git a/willow/src/willow_v1/server.rs b/willow/src/willow_v1/server.rs index c6df8a0..6a0b39b 100644 --- a/willow/src/willow_v1/server.rs +++ b/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 { +pub struct WillowV1Server { pub kahe: Kahe, pub vahe: Vahe, } +impl HasKahe for WillowV1Server { + type Kahe = Kahe; + fn kahe(&self) -> &Self::Kahe { + &self.kahe + } +} + +impl HasVahe for WillowV1Server { + type Vahe = Vahe; + fn vahe(&self) -> &Self::Vahe { + &self.vahe + } +} + /// State for the server. pub struct ServerState { /// The public key shares received from Decryptors. The key is the ID of the Decryptor. @@ -59,7 +73,7 @@ impl Clone for ServerState SecureAggregationServer for WillowV1Server +impl SecureAggregationServer for WillowV1Server where Vahe: EncryptVerify + PartialDec + Recover, Kahe: KaheBase + TrySecretKeyFrom + 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) } diff --git a/willow/src/willow_v1/verifier.rs b/willow/src/willow_v1/verifier.rs index 49c38f3..1107c54 100644 --- a/willow/src/willow_v1/verifier.rs +++ b/willow/src/willow_v1/verifier.rs @@ -14,7 +14,7 @@ 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. @@ -22,6 +22,13 @@ pub struct WillowV1Verifier { pub vahe: Vahe, } +impl HasVahe for WillowV1Verifier { + type Vahe = Vahe; + fn vahe(&self) -> &Self::Vahe { + &self.vahe + } +} + // State for the verifier after the first contribution is received. struct NonemptyVerifierState { partial_dec_ciphertext_sum: Vahe::PartialDecCiphertext, @@ -90,7 +97,7 @@ impl Clone for VerifierState { } } -impl SecureAggregationVerifier for WillowV1Verifier +impl SecureAggregationVerifier for WillowV1Verifier where Vahe: EncryptVerify, {