1- // Copyright 2025 Google LLC
1+ // Copyright 2026 Google LLC
22//
33// Licensed under the Apache License, Version 2.0 (the "License");
44// you may not use this file except in compliance with the License.
1212// See the License for the specific language governing permissions and
1313// limitations under the License.
1414
15- use messages:: { DecryptorPublicKeyShare , PartialDecryptionRequest , PartialDecryptionResponse } ;
15+ use ahe_traits:: AheBase ;
16+ use kahe_traits:: KaheBase ;
17+ use messages:: {
18+ DecryptorPublicKey , DecryptorPublicKeyShare , PartialDecryptionRequest ,
19+ PartialDecryptionResponse , RecoveryRequest , RecoveryResponse , SetupContribution ,
20+ VerifyKeyContributionsRequest ,
21+ } ;
1622use status:: StatusError ;
1723use vahe_traits:: HasVahe ;
1824
@@ -28,11 +34,142 @@ pub trait SecureAggregationDecryptor: HasVahe {
2834 decryptor_state : & mut Self :: DecryptorState ,
2935 ) -> Result < DecryptorPublicKeyShare < <Self as HasVahe >:: Vahe > , StatusError > ;
3036
37+ type Kahe : KaheBase ;
38+
3139 /// Handles a partial decryption request received from the Server. Returns a
3240/// partial decryption to the Server.
3341fn handle_partial_decryption_request (
3442 & self ,
3543 partial_decryption_request : PartialDecryptionRequest < <Self as HasVahe >:: Vahe > ,
36- decryptor_state : & Self :: DecryptorState ,
37- ) -> Result < PartialDecryptionResponse < <Self as HasVahe >:: Vahe > , StatusError > ;
44+ decryptor_state : & mut Self :: DecryptorState ,
45+ ) -> Result < PartialDecryptionResponse < Self :: Kahe , <Self as HasVahe >:: Vahe > , StatusError > ;
46+ }
47+
48+ /// Trait for reputable/non-recoverable decryptors (e.g. TEEs) in a multi-decryptor committee.
49+ pub trait SecureAggregationBaseMultiDecryptor : HasVahe {
50+ /// The state held by the Decryptor between messages.
51+ type DecryptorState : Default ;
52+
53+ /// Creates a public key share, a ZK proof of knowledge of the secret key,
54+ /// and encrypted shares of the randomness used for key generation.
55+ ///
56+ /// The randomness shares are encrypted for other committee members.
57+ fn create_setup_contribution (
58+ & self ,
59+ decryptor_state : & mut Self :: DecryptorState ,
60+ ) -> Result < SetupContribution < Self :: Vahe > , StatusError > ;
61+
62+ /// Creates a public key share to be sent to the Server, updating the
63+ /// decryptor state.
64+ fn create_public_key_share (
65+ & self ,
66+ decryptor_state : & mut Self :: DecryptorState ,
67+ ) -> Result < DecryptorPublicKeyShare < <Self as HasVahe >:: Vahe > , StatusError > ;
68+
69+ /// Handles a partial decryption request received from the Server. Returns a
70+ /// partial decryption to the Server.
71+ fn handle_partial_decryption_request < Kahe : KaheBase > (
72+ & self ,
73+ partial_decryption_request : PartialDecryptionRequest < <Self as HasVahe >:: Vahe > ,
74+ kahe : Option < & Kahe > ,
75+ decryptor_state : & mut Self :: DecryptorState ,
76+ ) -> Result < PartialDecryptionResponse < Kahe , <Self as HasVahe >:: Vahe > , StatusError > ;
77+ }
78+
79+ /// Trait for the reputable decryptors in a multi-decryptor committee.
80+ ///
81+ /// Reputable decryptors are assumed to be stable and do not share their
82+ /// randomness for recovery.
83+ pub trait SecureAggregationReputableDecryptor : SecureAggregationBaseMultiDecryptor {
84+ /// Verifies the ZK proofs of knowledge of the secret key for all public key
85+ /// shares, and returns the aggregated public key. Calling code should sign
86+ /// the aggregated public key for the aggregation.
87+ fn verify_and_aggregate_key_contributions (
88+ & self ,
89+ request : VerifyKeyContributionsRequest < <Self as HasVahe >:: Vahe > ,
90+ ) -> Result < DecryptorPublicKey < <Self as HasVahe >:: Vahe > , StatusError > ;
91+ }
92+
93+ /// Trait for the non-reputable decryptors in a multi-decryptor committee.
94+ pub trait SecureAggregationNonReputableMultiDecryptor : SecureAggregationBaseMultiDecryptor {
95+ /// Handles a request to decrypt shares of dropped decryptors.
96+ ///
97+ /// The decryptor should verify they are not being asked to decrypt more than
98+ /// the allowed threshold of shares.
99+ fn handle_recovery_request (
100+ & self ,
101+ recovery_request : RecoveryRequest ,
102+ decryptor_state : & mut Self :: DecryptorState ,
103+ ) -> Result < RecoveryResponse , StatusError > ;
104+ }
105+
106+ /// Trait for the protocol coordinator managing the multi-decryptor committee.
107+ ///
108+ /// The coordinator manages protocol flow and aggregates messages from all
109+ /// decryptors. The coordinator itself does not contribute to the public key.
110+ pub trait SecureAggregationCoordinator : HasVahe {
111+ /// The state held by the Coordinator between protocol rounds.
112+ type CoordinatorState : Default ;
113+
114+ /// Verifies and stores setup contributions from all decryptors.
115+ fn handle_setup_submissions (
116+ & self ,
117+ non_reputable_contributions : Vec < SetupContribution < Self :: Vahe > > ,
118+ reputable_contributions : Vec < SetupContribution < Self :: Vahe > > ,
119+ coordinator_state : & mut Self :: CoordinatorState ,
120+ ) -> Result < ( ) , StatusError > ;
121+
122+ /// Aggregates all public key shares into the final multi-decryptor public key.
123+ ///
124+ /// In the multiple reputable decryptor case, the coordinator also creates a verification
125+ /// request to be sent to reputable decryptors to check the proofs of all key contributions.
126+ fn aggregate_public_key_and_create_verification_request (
127+ & self ,
128+ coordinator_state : & mut Self :: CoordinatorState ,
129+ ) -> Result <
130+ ( VerifyKeyContributionsRequest < Self :: Vahe > , DecryptorPublicKey < Self :: Vahe > ) ,
131+ StatusError ,
132+ > ;
133+
134+ /// Combines the verifier's ciphertext half with the accumulated AHE components.
135+ ///
136+ /// The result should be forwarded to decryptors for partial decryption.
137+ fn prepare_decryption_request (
138+ & self ,
139+ verifier_ciphertext : & <Self :: Vahe as AheBase >:: PartialDecCiphertext ,
140+ coordinator_state : & mut Self :: CoordinatorState ,
141+ ) -> Result < PartialDecryptionRequest < Self :: Vahe > , StatusError > ;
142+
143+ /// Accumulates partial decryptions from responding decryptors.
144+ fn aggregate_partial_decryptions < Kahe : KaheBase > (
145+ & self ,
146+ partial_responses : Vec < PartialDecryptionResponse < Kahe , Self :: Vahe > > ,
147+ kahe : Option < & Kahe > ,
148+ coordinator_state : & mut Self :: CoordinatorState ,
149+ ) -> Result < ( ) , StatusError > ;
150+
151+ /// Creates recovery requests for surviving decryptors to decrypt shares of dropped client
152+ /// decryptors.
153+ ///
154+ /// If the vector is empty, there are no dropped decryptors to recover and
155+ /// recover_dropped_decryptors can be called immediately with an empty vector.
156+ fn create_recovery_requests (
157+ & self ,
158+ coordinator_state : & mut Self :: CoordinatorState ,
159+ ) -> Result < Vec < RecoveryRequest > , StatusError > ;
160+
161+ /// Finalizes the decryption by recovering randomness from dropped decryptors.
162+ ///
163+ /// Uses decrypted shares from survivors to simulate missing partial decryptions.
164+ fn recover_dropped_decryptors (
165+ & self ,
166+ recovery_responses : Vec < RecoveryResponse > ,
167+ coordinator_state : & mut Self :: CoordinatorState ,
168+ ) -> Result < ( ) , StatusError > ;
169+
170+ /// Returns the resulting plaintext from the final decryption, if available.
171+ fn get_plaintext (
172+ & self ,
173+ coordinator_state : & mut Self :: CoordinatorState ,
174+ ) -> Result < <Self :: Vahe as AheBase >:: Plaintext , StatusError > ;
38175}
0 commit comments