Simplify the accumulator trait to directly recover the plaintext from a FinalizedPartialDecryption, and use coordinator/multidecryptor traits for end-to-end tests.

willow/api/server_accumulator.rs

@@ -17,7 +17,7 @@ use aggregation_config::AggregationConfig;
 use aggregation_config_rust_proto::AggregationConfigProto;
 use ahe_traits::AheBase;
 use kahe_traits::KaheBase;
-use messages::{ClientMessage, PartialDecryptionResponse};
+use messages::{ClientMessage, FinalizedPartialDecryption, PartialDecryptionResponse};
 use messages_rust_proto::PartialDecryptionResponse as PartialDecryptionResponseProto;
 use proto_serialization_traits::{FromProto, ToProto};
 use protobuf::prelude::*;
@@ -187,7 +187,7 @@ impl ServerAccumulator {
             self.accumulator.split_client_message(client_message)?;
         self.verifier.verify_and_include(decryption_request_contribution, verifier_state)?;
         self.accumulator
-            .handle_ciphertext_contribution(ciphertext_contribution, accumulator_state)?;
+            .accumulate_ciphertext_contribution(ciphertext_contribution, accumulator_state)?;
         Ok(())
     }
 
@@ -630,9 +630,10 @@ impl FinalResultDecryptor {
 
         // Receives a single partial decryption response and attempts to recover right away.
         // This only works in the single-decryptor case.
-        self.accumulator.handle_partial_decryption(pd, &mut self.accumulator_state)?;
+        let finalized_pd =
+            FinalizedPartialDecryption { partial_decryption_sum: pd.partial_decryption };
         let aggregation_result =
-            self.accumulator.recover_aggregation_result(&self.accumulator_state)?;
+            self.accumulator.recover_aggregation_result(&self.accumulator_state, &finalized_pd)?;
 
         // `aggregation_result` is a Kahe::Plaintext, i.e. HashMap<String, Vec<u64>>
         // Flatten hashmap for FFI like in shell_testing_decryptor.rs

willow/benches/shell_benchmarks.rs

@@ -27,7 +27,7 @@ use decryptor_traits::SecureAggregationDecryptor;
 use kahe_traits::KaheBase;
 use messages::{
     CiphertextContribution, DecryptionRequestContribution, DecryptorPublicKey,
-    PartialDecryptionRequest,
+    FinalizedPartialDecryption, PartialDecryptionRequest,
 };
 use shell_kahe::ShellKahe;
 use shell_parameters::create_shell_configs;
@@ -275,7 +275,7 @@ fn setup_server_handle_client_message(args: &Args) -> ServerInputs {
 fn run_server_handle_client_message(inputs: &mut ServerInputs) {
     inputs
         .accumulator
-        .handle_ciphertext_contribution(
+        .accumulate_ciphertext_contribution(
             black_box(inputs.ciphertext_contributions.pop().unwrap()),
             black_box(&mut inputs.accumulator_state),
         )
@@ -288,6 +288,7 @@ fn run_server_handle_client_message(inputs: &mut ServerInputs) {
 struct ServerRecoverInputs {
     accumulator: WillowV1CiphertextAccumulator<ShellKahe, ShellVahe>,
     accumulator_state: CiphertextAccumulatorState<ShellKahe, ShellVahe>,
+    finalized_partial_decryption: FinalizedPartialDecryption<ShellVahe>,
 }
 
 fn setup_server_recover_aggregation_result(args: &Args) -> ServerRecoverInputs {
@@ -314,7 +315,7 @@ fn setup_server_recover_aggregation_result(args: &Args) -> ServerRecoverInputs {
     // Accumulator handles its part.
     inputs
         .accumulator
-        .handle_ciphertext_contribution(ciphertext_contribution, &mut inputs.accumulator_state)
+        .accumulate_ciphertext_contribution(ciphertext_contribution, &mut inputs.accumulator_state)
         .unwrap();
 
     // Verifier creates the partial decryption request.
@@ -327,18 +328,22 @@ fn setup_server_recover_aggregation_result(args: &Args) -> ServerRecoverInputs {
         .unwrap();
 
     // Accumulator handles the partial decryption.
-    inputs.accumulator.handle_partial_decryption(pd, &mut inputs.accumulator_state).unwrap();
+    let finalized_pd = FinalizedPartialDecryption { partial_decryption_sum: pd.partial_decryption };
 
     ServerRecoverInputs {
         accumulator: inputs.accumulator,
         accumulator_state: inputs.accumulator_state,
+        finalized_partial_decryption: finalized_pd,
     }
 }
 
 fn run_server_recover_aggregation_result(inputs: &mut ServerRecoverInputs) {
     let res = inputs
         .accumulator
-        .recover_aggregation_result(black_box(&inputs.accumulator_state))
+        .recover_aggregation_result(
+            black_box(&inputs.accumulator_state),
+            black_box(&inputs.finalized_partial_decryption),
+        )
         .unwrap();
     let _ = black_box(res); // Prevent optimization.
 }
@@ -368,7 +373,7 @@ fn setup_decryptor_partial_decryption(args: &Args) -> DecryptorInputs {
     // The accumulator and verifier each handle their part of the client message.
     inputs
         .accumulator
-        .handle_ciphertext_contribution(ciphertext_contribution, &mut inputs.accumulator_state)
+        .accumulate_ciphertext_contribution(ciphertext_contribution, &mut inputs.accumulator_state)
         .unwrap();
     inputs
         .verifier

willow/proto/willow/messages.proto

@@ -62,7 +62,7 @@ message ServerState {
   map<string, ShellAhePublicKeyShare> decryptor_public_key_shares = 1;
   ShellKaheCiphertext client_sum_kahe = 2;
   ShellAheRecoverCiphertext client_sum_ahe_recover = 3;
-  ShellAhePartialDecryption partial_decryption_sum = 4;
+  ShellAhePartialDecryption partial_decryption_sum = 4 [deprecated = true];
 }
 
 message VerifierState {

willow/protocol/accumulator_traits.rs

@@ -14,7 +14,8 @@
 
 use kahe_traits::HasKahe;
 use messages::{
-    CiphertextContribution, ClientMessage, DecryptionRequestContribution, PartialDecryptionResponse,
+    CiphertextContribution, ClientMessage, DecryptionRequestContribution,
+    FinalizedPartialDecryption,
 };
 use status::StatusError;
 use vahe_traits::HasVahe;
@@ -25,7 +26,6 @@ type Vahe<T> = <T as HasVahe>::Vahe;
 
 /// Base trait for the secure aggregation accumulator. Also includes the Coordinator
 /// functionality of the threshold AHE scheme.
-///
 pub trait SecureAggregationCiphertextAccumulator: HasKahe + HasVahe {
     /// The state held by the accumulator between messages.
     type CiphertextAccumulatorState: Default + Clone;
@@ -42,26 +42,19 @@ pub trait SecureAggregationCiphertextAccumulator: HasKahe + HasVahe {
         StatusError,
     >;
 
-    /// Handles a single client message, updating the accumulator state.
-    fn handle_ciphertext_contribution(
+    /// Accumulates a single client ciphertext contribution into the accumulator state.
+    fn accumulate_ciphertext_contribution(
         &self,
         ciphertext_contribution: CiphertextContribution<Kahe<Self>, Vahe<Self>>,
         accumulator_state: &mut Self::CiphertextAccumulatorState,
     ) -> Result<(), StatusError>;
 
-    /// Handles a partial decryption received from a Decryptor, updating the
-    /// accumulator state.
-    fn handle_partial_decryption(
-        &self,
-        partial_decryption_response: PartialDecryptionResponse<Kahe<Self>, Vahe<Self>>,
-        accumulator_state: &mut Self::CiphertextAccumulatorState,
-    ) -> Result<(), StatusError>;
-
-    /// Recovers the aggregation result after enough partial decryptions have
-    /// been received from Decryptors.
+    /// Recovers the aggregation result from the accumulated ciphertext contributions
+    /// and the finalized partial decryption.
     fn recover_aggregation_result(
         &self,
         accumulator_state: &Self::CiphertextAccumulatorState,
+        finalized_partial_decryption: &FinalizedPartialDecryption<Vahe<Self>>,
     ) -> Result<Self::AggregationResult, StatusError>;
 
     /// Merges two accumulator states into one.

willow/protocol/willow_v1_accumulator.rs

@@ -16,7 +16,8 @@ use accumulator_traits::SecureAggregationCiphertextAccumulator;
 use ahe_traits::PartialDec;
 use kahe_traits::{HasKahe, KaheBase, KaheDecrypt, TrySecretKeyFrom};
 use messages::{
-    CiphertextContribution, ClientMessage, DecryptionRequestContribution, PartialDecryptionResponse,
+    CiphertextContribution, ClientMessage, DecryptionRequestContribution,
+    FinalizedPartialDecryption,
 };
 use messages_rust_proto::ServerState as ServerStateProto;
 use proto_serialization_traits::{FromProto, ToProto};
@@ -54,24 +55,19 @@ impl<Kahe: KaheBase, Vahe: VaheBase> HasVahe for WillowV1CiphertextAccumulator<K
 pub struct CiphertextAccumulatorState<Kahe: KaheBase, Vahe: VaheBase + PartialDec> {
     /// Running sum of client ciphertexts.
     client_sum: Option<(Kahe::Ciphertext, Vahe::RecoverCiphertext)>,
-    /// Running sum of partial decryptions.
-    partial_decryption_sum: Option<Vahe::PartialDecryption>,
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase + PartialDec> Default
     for CiphertextAccumulatorState<Kahe, Vahe>
 {
     fn default() -> Self {
-        Self { client_sum: None, partial_decryption_sum: None }
+        Self { client_sum: None }
     }
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase + PartialDec> Clone for CiphertextAccumulatorState<Kahe, Vahe> {
     fn clone(&self) -> Self {
-        Self {
-            client_sum: self.client_sum.clone(),
-            partial_decryption_sum: self.partial_decryption_sum.clone(),
-        }
+        Self { client_sum: self.client_sum.clone() }
     }
 }
 
@@ -82,7 +78,6 @@ where
     Vahe: VaheBase + PartialDec + 'a,
     Kahe::Ciphertext: ToProto<&'a Kahe, Proto = ShellKaheCiphertext>, // TODO: Rename protos to be generic once cl/836370582 has landed.
     Vahe::RecoverCiphertext: ToProto<&'a Vahe, Proto = ShellAheRecoverCiphertext>,
-    Vahe::PartialDecryption: ToProto<&'a Vahe, Proto = ShellAhePartialDecryption>,
 {
     type Proto = ServerStateProto;
 
@@ -94,10 +89,6 @@ where
             proto.set_client_sum_ahe_recover(ahe.to_proto(context.vahe())?);
         }
 
-        if let Some(pd) = &self.partial_decryption_sum {
-            proto.set_partial_decryption_sum(pd.to_proto(context.vahe())?);
-        }
-
         Ok(proto)
     }
 }
@@ -109,7 +100,6 @@ where
     Vahe: VaheBase + PartialDec + 'a,
     Kahe::Ciphertext: FromProto<&'a Kahe, Proto = ShellKaheCiphertext>,
     Vahe::RecoverCiphertext: FromProto<&'a Vahe, Proto = ShellAheRecoverCiphertext>,
-    Vahe::PartialDecryption: FromProto<&'a Vahe, Proto = ShellAhePartialDecryption>,
 {
     type Proto = ServerStateProto;
 
@@ -136,16 +126,7 @@ where
             ));
         };
 
-        let partial_decryption_sum = if proto.has_partial_decryption_sum() {
-            Some(Vahe::PartialDecryption::from_proto(
-                proto.partial_decryption_sum(),
-                context.vahe(),
-            )?)
-        } else {
-            None
-        };
-
-        Ok(CiphertextAccumulatorState { client_sum, partial_decryption_sum })
+        Ok(CiphertextAccumulatorState { client_sum })
     }
 }
 
@@ -186,8 +167,8 @@ where
         ))
     }
 
-    /// Handles a single client's ciphertext contribution, updating the accumulator state.
-    fn handle_ciphertext_contribution(
+    /// Accumulates a single client's ciphertext contribution, updating the accumulator state.
+    fn accumulate_ciphertext_contribution(
         &self,
         contribution: CiphertextContribution<Kahe, Vahe>,
         accumulator_state: &mut Self::CiphertextAccumulatorState,
@@ -207,46 +188,20 @@ where
         Ok(())
     }
 
-    /// Handles a partial decryption response received from a Decryptor, updating the
-    /// accumulator state.
-    fn handle_partial_decryption(
-        &self,
-        partial_decryption_response: PartialDecryptionResponse<Kahe, Vahe>,
-        accumulator_state: &mut Self::CiphertextAccumulatorState,
-    ) -> Result<(), status::StatusError> {
-        if partial_decryption_response.dp_ciphertext_contribution.is_some() {
-            return Err(status::failed_precondition(
-                "DP ciphertext contributions are not yet supported.",
-            ));
-        }
-        let partial_decryption = partial_decryption_response.partial_decryption;
-        if let Some(ref mut partial_decryption_sum) = accumulator_state.partial_decryption_sum {
-            self.vahe
-                .add_partial_decryptions_in_place(&partial_decryption, partial_decryption_sum)?;
-        } else {
-            accumulator_state.partial_decryption_sum = Some(partial_decryption);
-        }
-        Ok(())
-    }
-
-    /// Recovers the aggregation result after enough partial decryptions have
-    /// been received from Decryptors.
+    /// Recovers the aggregation result from the accumulated ciphertext contributions
+    /// and the finalized partial decryption.
     fn recover_aggregation_result(
         &self,
         accumulator_state: &Self::CiphertextAccumulatorState,
+        finalized_partial_decryption: &FinalizedPartialDecryption<Vahe>,
     ) -> Result<Self::AggregationResult, status::StatusError> {
         if let Some((ref kahe_ciphertext, ref recover_ciphertext)) = accumulator_state.client_sum {
-            if let Some(ref partial_decryption_sum) = accumulator_state.partial_decryption_sum {
-                let ahe_plaintext =
-                    self.vahe.recover(&partial_decryption_sum, &recover_ciphertext, None)?;
-                let kahe_secret_key = self.kahe.try_secret_key_from(ahe_plaintext)?;
-                let kahe_plaintext = self.kahe.decrypt(kahe_ciphertext, &kahe_secret_key)?;
-                Ok(kahe_plaintext)
-            } else {
-                Err(status::failed_precondition(
-                    "Must handle at least one partial decryption before requesting recovery",
-                ))?
-            }
+            let partial_decryption_sum = &finalized_partial_decryption.partial_decryption_sum;
+            let ahe_plaintext =
+                self.vahe.recover(partial_decryption_sum, &recover_ciphertext, None)?;
+            let kahe_secret_key = self.kahe.try_secret_key_from(ahe_plaintext)?;
+            let kahe_plaintext = self.kahe.decrypt(kahe_ciphertext, &kahe_secret_key)?;
+            Ok(kahe_plaintext)
         } else {
             Err(status::failed_precondition(
                 "Must handle at least one client message before requesting recovery",
@@ -285,19 +240,6 @@ where
                 (None, None) => None,
             };
 
-        merged_accumulator_state.partial_decryption_sum = match (
-            accumulator_state_1.partial_decryption_sum,
-            accumulator_state_2.partial_decryption_sum,
-        ) {
-            (Some(sum1), Some(sum2)) => {
-                let mut merged_sum = sum1;
-                self.vahe.add_partial_decryptions_in_place(&sum2, &mut merged_sum)?;
-                Some(merged_sum)
-            }
-            (Some(s), None) | (None, Some(s)) => Some(s),
-            (None, None) => None,
-        };
-
         Ok(merged_accumulator_state)
     }
 }
@@ -364,15 +306,14 @@ mod tests {
         let accumulator_state_roundtrip =
             CiphertextAccumulatorState::from_proto(accumulator_state_proto, &accumulator)?;
         verify_true!(accumulator_state_roundtrip.client_sum.is_none())?;
-        verify_true!(accumulator_state_roundtrip.partial_decryption_sum.is_none())?;
 
         // Populate accumulator state.
         let public_key_share = decryptor.create_public_key_share(&mut decryptor_state)?;
         // Aggregate public key share directly.
         let public_key = vahe.aggregate_public_key_shares(std::iter::once(&public_key_share))?;
         let client_plaintext = HashMap::from([(
             DEFAULT_VECTOR_ID.to_string(),
-            vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1],
+            vec![1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1],
         )]);
         let nonce = generate_random_nonce();
         let client_message = client.create_client_message(
@@ -384,19 +325,23 @@ mod tests {
             accumulator.split_client_message(client_message)?;
         verifier.verify_and_include(decryption_request_contribution, &mut verifier_state)?;
         accumulator
-            .handle_ciphertext_contribution(ciphertext_contribution, &mut accumulator_state)?;
+            .accumulate_ciphertext_contribution(ciphertext_contribution, &mut accumulator_state)?;
         let pd_ct = verifier.create_partial_decryption_request(verifier_state)?;
         let pd = decryptor.handle_partial_decryption_request(pd_ct, &mut decryptor_state)?;
-        accumulator.handle_partial_decryption(pd, &mut accumulator_state)?;
 
         // Check populated state serialization
         verify_true!(accumulator_state.client_sum.is_some())?;
-        verify_true!(accumulator_state.partial_decryption_sum.is_some())?;
         let accumulator_state_proto = accumulator_state.to_proto(&accumulator)?;
         let accumulator_state_roundtrip =
             CiphertextAccumulatorState::from_proto(accumulator_state_proto, &accumulator)?;
         verify_true!(accumulator_state_roundtrip.client_sum.is_some())?;
-        verify_true!(accumulator_state_roundtrip.partial_decryption_sum.is_some())?;
+
+        // Check recovery
+        let finalized_pd =
+            FinalizedPartialDecryption { partial_decryption_sum: pd.partial_decryption };
+        let recovered =
+            accumulator.recover_aggregation_result(&accumulator_state, &finalized_pd)?;
+        verify_eq!(recovered, client_plaintext)?;
 
         Ok(())
     }

willow/tests/BUILD

@@ -44,6 +44,7 @@ rust_test(
         "//willow/protocol:verifier_traits",
         "//willow/protocol:willow_v1_accumulator",
         "//willow/protocol:willow_v1_client",
+        "//willow/protocol:willow_v1_coordinator",
         "//willow/protocol:willow_v1_decryptor",
         "//willow/protocol:willow_v1_verifier",
         "//willow/testing_utils",