Add traits and messages for multi-decryptor secure aggregation.

willow/benches/shell_benchmarks.rs

@@ -324,8 +324,10 @@ fn setup_server_recover_aggregation_result(args: &Args) -> ServerRecoverInputs {
     let pd_ct = inputs.verifier.create_partial_decryption_request(inputs.verifier_state).unwrap();
 
     // Decryptor creates partial decryption.
-    let pd =
-        inputs.decryptor.handle_partial_decryption_request(pd_ct, &inputs.decryptor_state).unwrap();
+    let pd = inputs
+        .decryptor
+        .handle_partial_decryption_request(pd_ct, &mut inputs.decryptor_state)
+        .unwrap();
 
     // Server handles the partial decryption.
     inputs.server.handle_partial_decryption(pd, &mut inputs.server_state).unwrap();
@@ -384,7 +386,7 @@ fn run_decryptor_partial_decryption(inputs: &mut DecryptorInputs) {
         .decryptor
         .handle_partial_decryption_request(
             black_box(inputs.partial_decryption_request.clone()),
-            black_box(&inputs.decryptor_state),
+            black_box(&mut inputs.decryptor_state),
         )
         .unwrap();
     let _ = black_box(res); // Prevent optimization.

willow/proto/willow/BUILD

@@ -109,6 +109,7 @@ proto_library(
     name = "messages_proto",
     srcs = ["messages.proto"],
     deps = [
+        ":aggregation_config_proto",
         "//willow/proto/shell:shell_ciphertexts_proto",
         "//willow/proto/zk:proofs_proto",
     ],

willow/proto/willow/messages.proto

@@ -17,6 +17,7 @@ syntax = "proto3";
 package secure_aggregation.willow;
 
 import "willow/proto/shell/ciphertexts.proto";
+import "willow/proto/willow/aggregation_config.proto";
 import "willow/proto/zk/proofs.proto";
 
 option java_multiple_files = true;
@@ -32,10 +33,13 @@ message ClientMessage {
 
 message PartialDecryptionRequest {
   ShellAhePartialDecCiphertext partial_dec_ciphertext = 1;
+  AggregationConfigProto aggregation_config = 2;
 }
 
 message PartialDecryptionResponse {
   ShellAhePartialDecryption partial_decryption = 1;
+  // Noise contribution to the final aggregated result.
+  CiphertextContribution dp_ciphertext_contribution = 2;
 }
 
 message CiphertextContribution {
@@ -51,6 +55,7 @@ message DecryptionRequestContribution {
 
 message DecryptorStateProto {
   ShellAheSecretKeyShare sk_share = 1;
+  AggregationConfigProto aggregation_config = 2;
 }
 
 message ServerStateProto {

willow/protocol/BUILD

@@ -29,11 +29,13 @@ rust_library(
     deps = [
         "@protobuf//rust:protobuf",
         "//ffi_utils:status",
+        "//willow/api:aggregation_config",
         "//willow/api:proto_serialization_traits",
         "//willow/crypto:ahe_traits",
         "//willow/crypto:kahe_traits",
         "//willow/crypto:vahe_traits",
         "//willow/proto/shell:shell_ciphertexts_rust_proto",
+        "//willow/proto/willow:aggregation_config_rust_proto",
         "//willow/proto/willow:messages_rust_proto",
         "//willow/proto/zk:proofs_rust_proto",
     ],
@@ -74,6 +76,8 @@ rust_library(
     deps = [
         ":messages",
         "//ffi_utils:status",
+        "//willow/crypto:ahe_traits",
+        "//willow/crypto:kahe_traits",
         "//willow/crypto:vahe_traits",
     ],
 )
@@ -130,8 +134,10 @@ rust_test(
         "@crate_index//:googletest",
         "//willow/api:proto_serialization_traits",
         "//willow/crypto:ahe_traits",
+        "//willow/crypto:shell_kahe",
         "//willow/crypto:shell_parameters",
         "//willow/crypto:shell_vahe",
+        "//willow/testing_utils",
     ],
 )
 
@@ -145,9 +151,13 @@ rust_library(
         ":messages",
         "@protobuf//rust:protobuf",
         "//ffi_utils:status",
+        "//willow/api:aggregation_config",
         "//willow/api:proto_serialization_traits",
         "//willow/crypto:ahe_traits",
+        "//willow/crypto:kahe_traits",
         "//willow/crypto:prng_traits",
+        "//willow/crypto:shell_kahe",
+        "//willow/crypto:shell_parameters",
         "//willow/crypto:vahe_traits",
         "//willow/proto/shell:shell_ciphertexts_rust_proto",
         "//willow/proto/willow:messages_rust_proto",

willow/protocol/decryptor_traits.rs

@@ -1,4 +1,4 @@
-// Copyright 2025 Google LLC
+// Copyright 2026 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,7 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use messages::{DecryptorPublicKeyShare, PartialDecryptionRequest, PartialDecryptionResponse};
+use ahe_traits::AheBase;
+use kahe_traits::KaheBase;
+use messages::{
+    DecryptorPublicKey, DecryptorPublicKeyShare, PartialDecryptionRequest,
+    PartialDecryptionResponse, RecoveryRequest, RecoveryResponse, SetupContribution,
+    VerifyKeyContributionsRequest,
+};
 use status::StatusError;
 use vahe_traits::HasVahe;
 
@@ -28,11 +34,142 @@ pub trait SecureAggregationDecryptor: HasVahe {
         decryptor_state: &mut Self::DecryptorState,
     ) -> Result<DecryptorPublicKeyShare<<Self as HasVahe>::Vahe>, StatusError>;
 
+    type Kahe: KaheBase;
+
     /// Handles a partial decryption request received from the Server. Returns a
     /// partial decryption to the Server.
     fn handle_partial_decryption_request(
         &self,
         partial_decryption_request: PartialDecryptionRequest<<Self as HasVahe>::Vahe>,
-        decryptor_state: &Self::DecryptorState,
-    ) -> Result<PartialDecryptionResponse<<Self as HasVahe>::Vahe>, StatusError>;
+        decryptor_state: &mut Self::DecryptorState,
+    ) -> Result<PartialDecryptionResponse<Self::Kahe, <Self as HasVahe>::Vahe>, StatusError>;
+}
+
+/// Trait for reputable/non-recoverable decryptors (e.g. TEEs) in a multi-decryptor committee.
+pub trait SecureAggregationBaseMultiDecryptor: HasVahe {
+    /// The state held by the Decryptor between messages.
+    type DecryptorState: Default;
+
+    /// Creates a public key share, a ZK proof of knowledge of the secret key,
+    /// and encrypted shares of the randomness used for key generation.
+    ///
+    /// The randomness shares are encrypted for other committee members.
+    fn create_setup_contribution(
+        &self,
+        decryptor_state: &mut Self::DecryptorState,
+    ) -> Result<SetupContribution<Self::Vahe>, StatusError>;
+
+    /// Creates a public key share to be sent to the Server, updating the
+    /// decryptor state.
+    fn create_public_key_share(
+        &self,
+        decryptor_state: &mut Self::DecryptorState,
+    ) -> 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<Kahe: KaheBase>(
+        &self,
+        partial_decryption_request: PartialDecryptionRequest<<Self as HasVahe>::Vahe>,
+        kahe: Option<&Kahe>,
+        decryptor_state: &mut Self::DecryptorState,
+    ) -> Result<PartialDecryptionResponse<Kahe, <Self as HasVahe>::Vahe>, StatusError>;
+}
+
+/// Trait for the reputable decryptors in a multi-decryptor committee.
+///
+/// Reputable decryptors are assumed to be stable and do not share their
+/// randomness for recovery.
+pub trait SecureAggregationReputableDecryptor: SecureAggregationBaseMultiDecryptor {
+    /// Verifies the ZK proofs of knowledge of the secret key for all public key
+    /// shares, and returns the aggregated public key. Calling code should sign
+    /// the aggregated public key for the aggregation.
+    fn verify_and_aggregate_key_contributions(
+        &self,
+        request: VerifyKeyContributionsRequest<<Self as HasVahe>::Vahe>,
+    ) -> Result<DecryptorPublicKey<<Self as HasVahe>::Vahe>, StatusError>;
+}
+
+/// Trait for the non-reputable decryptors in a multi-decryptor committee.
+pub trait SecureAggregationNonReputableMultiDecryptor: SecureAggregationBaseMultiDecryptor {
+    /// Handles a request to decrypt shares of dropped decryptors.
+    ///
+    /// The decryptor should verify they are not being asked to decrypt more than
+    /// the allowed threshold of shares.
+    fn handle_recovery_request(
+        &self,
+        recovery_request: RecoveryRequest,
+        decryptor_state: &mut Self::DecryptorState,
+    ) -> Result<RecoveryResponse, StatusError>;
+}
+
+/// Trait for the protocol coordinator managing the multi-decryptor committee.
+///
+/// The coordinator manages protocol flow and aggregates messages from all
+/// decryptors. The coordinator itself does not contribute to the public key.
+pub trait SecureAggregationCoordinator: HasVahe {
+    /// The state held by the Coordinator between protocol rounds.
+    type CoordinatorState: Default;
+
+    /// Verifies and stores setup contributions from all decryptors.
+    fn handle_setup_submissions(
+        &self,
+        non_reputable_contributions: Vec<SetupContribution<Self::Vahe>>,
+        reputable_contributions: Vec<SetupContribution<Self::Vahe>>,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<(), StatusError>;
+
+    /// Aggregates all public key shares into the final multi-decryptor public key.
+    ///
+    /// In the multiple reputable decryptor case, the coordinator also creates a verification
+    /// request to be sent to reputable decryptors to check the proofs of all key contributions.
+    fn aggregate_public_key_and_create_verification_request(
+        &self,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<
+        (VerifyKeyContributionsRequest<Self::Vahe>, DecryptorPublicKey<Self::Vahe>),
+        StatusError,
+    >;
+
+    /// Combines the verifier's ciphertext half with the accumulated AHE components.
+    ///
+    /// The result should be forwarded to decryptors for partial decryption.
+    fn prepare_decryption_request(
+        &self,
+        verifier_ciphertext: &<Self::Vahe as AheBase>::PartialDecCiphertext,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<PartialDecryptionRequest<Self::Vahe>, StatusError>;
+
+    /// Accumulates partial decryptions from responding decryptors.
+    fn aggregate_partial_decryptions<Kahe: KaheBase>(
+        &self,
+        partial_responses: Vec<PartialDecryptionResponse<Kahe, Self::Vahe>>,
+        kahe: Option<&Kahe>,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<(), StatusError>;
+
+    /// Creates recovery requests for surviving decryptors to decrypt shares of dropped client
+    /// decryptors.
+    ///
+    /// If the vector is empty, there are no dropped decryptors to recover and
+    /// recover_dropped_decryptors can be called immediately with an empty vector.
+    fn create_recovery_requests(
+        &self,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<Vec<RecoveryRequest>, StatusError>;
+
+    /// Finalizes the decryption by recovering randomness from dropped decryptors.
+    ///
+    /// Uses decrypted shares from survivors to simulate missing partial decryptions.
+    fn recover_dropped_decryptors(
+        &self,
+        recovery_responses: Vec<RecoveryResponse>,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<(), StatusError>;
+
+    /// Returns the resulting plaintext from the final decryption, if available.
+    fn get_plaintext(
+        &self,
+        coordinator_state: &mut Self::CoordinatorState,
+    ) -> Result<<Self::Vahe as AheBase>::Plaintext, StatusError>;
 }