Skip to content

Commit 684456b

Browse files
schoppmpcopybara-github
schoppmp
authored and
copybara-github
committed
Wrap Kahe and Vahe struct members with Rc
PiperOrigin-RevId: 859733028
1 parent 1c0b146 commit 684456b

File tree

8 files changed

+193
-264
lines changed

8 files changed

+193
-264
lines changed

‎willow/benches/shell_benchmarks.rs‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
use clap::Parser;
1616
use std::collections::HashMap;
1717
use std::hint::black_box;
18+
use std::rc::Rc;
1819
use std::time::Duration;
1920

2021
use aggregation_config::AggregationConfig;
@@ -126,25 +127,25 @@ fn setup_base(args: &Args) -> BaseInputs {
126127
let ahe_config = create_shell_ahe_config(aggregation_config.max_number_of_decryptors).unwrap();
127128
let kahe_config = create_shell_kahe_config(&aggregation_config).unwrap();
128129

130+
// Create common KAHE/VAHE instances.
131+
let kahe = Rc::new(ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap());
132+
let vahe = Rc::new(ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap());
133+
129134
// Create client.
130-
let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();
131-
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
132-
let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe).unwrap();
135+
let client =
136+
WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
137+
.unwrap();
133138

134139
// Create decryptor.
135-
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
136140
let mut decryptor_state = DecryptorState::default();
137-
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe).unwrap();
141+
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();
138142

139143
// Create server.
140-
let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();
141-
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
142-
let server = WillowV1Server { kahe, vahe };
144+
let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
143145
let mut server_state = ServerState::default();
144146

145147
// Create verifier.
146-
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
147-
let verifier = WillowV1Verifier { vahe };
148+
let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
148149
let verifier_state = VerifierState::default();
149150

150151
// Decryptor generates public key share.

‎willow/src/api/client.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use shell_ciphertexts_rust_proto::ShellAhePublicKey;
2828
use status::ffi::FfiStatus;
2929
use status::StatusError;
3030
use std::collections::HashMap;
31+
use std::rc::Rc;
3132
use vahe_shell::ShellVahe;
3233
use willow_v1_client::WillowV1Client;
3334

@@ -82,8 +83,8 @@ impl WillowShellClient {
8283
let aggregation_config = AggregationConfig::from_proto(aggregation_config_proto, ())?;
8384
let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
8485
let context_bytes = &aggregation_config.key_id;
85-
let kahe = ShellKahe::new(kahe_config, &context_bytes)?;
86-
let vahe = ShellVahe::new(ahe_config, &context_bytes)?;
86+
let kahe = Rc::new(ShellKahe::new(kahe_config, &context_bytes)?);
87+
let vahe = Rc::new(ShellVahe::new(ahe_config, &context_bytes)?);
8788
let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
8889
Ok(WillowShellClient(client))
8990
}
@@ -102,7 +103,7 @@ impl WillowShellClient {
102103
}
103104
let public_key_proto = ShellAhePublicKey::parse(public_key.as_bytes())
104105
.map_err(|e| status::internal(format!("Failed to parse ShellAhePublicKey: {}", e)))?;
105-
let public_key_rust = PublicKey::from_proto(public_key_proto, &self.0.vahe)?;
106+
let public_key_rust = PublicKey::from_proto(public_key_proto, self.0.vahe.as_ref())?;
106107
let message = self.0.create_client_message(&plaintext_slice, &public_key_rust, nonce)?;
107108
Ok(message
108109
.to_proto(&self.0)?

‎willow/src/api/server_accumulator.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use server_traits::SecureAggregationServer;
3131
use status::StatusError;
3232
use std::collections::BTreeMap;
3333
use std::ops::Range;
34+
use std::rc::Rc;
3435
use vahe_shell::ShellVahe;
3536
use verifier_traits::SecureAggregationVerifier;
3637
use willow_v1_server::{ServerState, WillowV1Server};
@@ -147,11 +148,10 @@ impl ServerAccumulator {
147148
fn new(aggregation_config: AggregationConfig) -> Result<Self, StatusError> {
148149
let (kahe_config, vahe_config) = create_shell_configs(&aggregation_config)?;
149150
let context_bytes = &aggregation_config.key_id;
150-
let server_kahe = ShellKahe::new(kahe_config, context_bytes)?;
151-
let server_vahe = ShellVahe::new(vahe_config.clone(), context_bytes)?;
152-
let verifier_vahe = ShellVahe::new(vahe_config, context_bytes)?;
153-
let server = WillowV1Server { kahe: server_kahe, vahe: server_vahe };
154-
let verifier = WillowV1Verifier { vahe: verifier_vahe };
151+
let kahe = Rc::new(ShellKahe::new(kahe_config, &context_bytes)?);
152+
let vahe = Rc::new(ShellVahe::new(vahe_config, &context_bytes)?);
153+
let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
154+
let verifier = WillowV1Verifier { vahe };
155155
Ok(Self {
156156
server: server,
157157
server_state: Default::default(),
@@ -661,8 +661,8 @@ impl FinalResultDecryptor {
661661
let aggregation_config = AggregationConfig::from_proto(aggregation_config_proto, ())?;
662662
let (kahe_config, vahe_config) = create_shell_configs(&aggregation_config)?;
663663
let context_bytes = &aggregation_config.key_id;
664-
let kahe = ShellKahe::new(kahe_config, context_bytes)?;
665-
let vahe = ShellVahe::new(vahe_config, context_bytes)?;
664+
let kahe = Rc::new(ShellKahe::new(kahe_config, context_bytes)?);
665+
let vahe = Rc::new(ShellVahe::new(vahe_config, context_bytes)?);
666666
let server = WillowV1Server { kahe, vahe };
667667
let server_state = ServerState::from_proto(server_state_proto, &server)?;
668668

‎willow/src/willow_v1/client.rs‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,33 +17,34 @@ use kahe_traits::{HasKahe, KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
1717
use messages::{ClientMessage, DecryptorPublicKey};
1818
use prng_traits::SecurePrng;
1919
use std::cell::RefCell;
20+
use std::rc::Rc;
2021
use vahe_traits::{HasVahe, VaheBase, VerifiableEncrypt};
2122

2223
/// Lightweight client directly exposing KAHE/VAHE types.
2324
pub struct WillowV1Client<Kahe: KaheBase, Vahe: VaheBase> {
24-
pub kahe: Kahe,
25-
pub vahe: Vahe,
25+
pub kahe: Rc<Kahe>,
26+
pub vahe: Rc<Vahe>,
2627
pub prng: RefCell<Kahe::Rng>, // Using a single PRNG for both VAHE and KAHE.
2728
}
2829

2930
impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Client<Kahe, Vahe> {
3031
type Kahe = Kahe;
3132
fn kahe(&self) -> &Self::Kahe {
32-
&self.kahe
33+
self.kahe.as_ref()
3334
}
3435
}
3536

3637
impl<Kahe: KaheBase, Vahe: VaheBase> HasVahe for WillowV1Client<Kahe, Vahe> {
3738
type Vahe = Vahe;
3839
fn vahe(&self) -> &Self::Vahe {
39-
&self.vahe
40+
self.vahe.as_ref()
4041
}
4142
}
4243

4344
impl<Kahe: KaheBase, Vahe: VaheBase> WillowV1Client<Kahe, Vahe> {
4445
pub fn new_with_randomly_generated_seed(
45-
kahe: Kahe,
46-
vahe: Vahe,
46+
kahe: Rc<Kahe>,
47+
vahe: Rc<Vahe>,
4748
) -> Result<Self, status::StatusError> {
4849
let seed = Kahe::Rng::generate_seed()?;
4950
let prng = RefCell::new(Kahe::Rng::create(&seed)?);
@@ -122,8 +123,8 @@ mod test {
122123

123124
// Create a client.
124125
let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
125-
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
126-
let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
126+
let kahe = Rc::new(ShellKahe::new(kahe_config, CONTEXT_STRING)?);
127+
let vahe = Rc::new(ShellVahe::new(ahe_config, CONTEXT_STRING)?);
127128
let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
128129

129130
// Generate AHE keys.
@@ -159,16 +160,14 @@ mod test {
159160
key_id: b"test".to_vec(),
160161
};
161162

162-
// Create a client.
163+
// Create common KAHE/VAHE instances.
163164
let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
164-
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
165-
let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
166-
let client1 = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
165+
let kahe = Rc::new(ShellKahe::new(kahe_config, CONTEXT_STRING)?);
166+
let vahe = Rc::new(ShellVahe::new(ahe_config, CONTEXT_STRING)?);
167167

168-
// Create a second client.
169-
let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
170-
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
171-
let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
168+
// Create clients.
169+
let client1 =
170+
WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))?;
172171
let client2 = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
173172

174173
// Generate AHE keys.

‎willow/src/willow_v1/decryptor.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,13 @@ use protobuf::AsView;
2222
use shell_ciphertexts_rust_proto::ShellAheSecretKeyShare;
2323
use status::StatusError;
2424
use std::cell::RefCell;
25+
use std::rc::Rc;
2526
use vahe_traits::{EncryptVerify, HasVahe, VaheBase};
2627

2728
/// Lightweight decryptor directly exposing KAHE/VAHE types. It verifies only the client proofs,
2829
/// does not provide verifiable partial decryptions.
2930
pub struct WillowV1Decryptor<Vahe: VaheBase> {
30-
pub vahe: Vahe,
31+
pub vahe: Rc<Vahe>,
3132
pub prng: RefCell<Vahe::Rng>,
3233
}
3334

@@ -39,7 +40,7 @@ impl<Vahe: VaheBase> HasVahe for WillowV1Decryptor<Vahe> {
3940
}
4041

4142
impl<Vahe: VaheBase> WillowV1Decryptor<Vahe> {
42-
pub fn new_with_randomly_generated_seed(vahe: Vahe) -> Result<Self, status::StatusError> {
43+
pub fn new_with_randomly_generated_seed(vahe: Rc<Vahe>) -> Result<Self, status::StatusError> {
4344
let seed = Vahe::Rng::generate_seed()?;
4445
let prng = RefCell::new(Vahe::Rng::create(&seed)?);
4546
Ok(Self { vahe, prng })
@@ -145,13 +146,15 @@ mod tests {
145146
use googletest::{gtest, verify_true};
146147
use parameters_shell::create_shell_ahe_config;
147148
use proto_serialization_traits::{FromProto, ToProto};
149+
use std::rc::Rc;
148150
use vahe_shell::ShellVahe;
149151

150152
const CONTEXT_STRING: &[u8] = b"testing_context_string";
151153

152154
#[gtest]
153155
fn decryptor_state_serialization_roundtrip() -> googletest::Result<()> {
154-
let vahe = ShellVahe::new(create_shell_ahe_config(1).unwrap(), CONTEXT_STRING).unwrap();
156+
let vahe =
157+
Rc::new(ShellVahe::new(create_shell_ahe_config(1).unwrap(), CONTEXT_STRING).unwrap());
155158
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
156159
let mut decryptor_state = DecryptorState::default();
157160

‎willow/src/willow_v1/server.rs‎

Lines changed: 21 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,15 @@ use shell_ciphertexts_rust_proto::{
2828
};
2929
use status::StatusError;
3030
use std::collections::HashMap;
31+
use std::rc::Rc;
3132
use vahe_traits::{EncryptVerify, HasVahe, Recover, VaheBase};
3233

3334
/// Implements the `server` role in the Willow protocol. This includes aggregating public key shares
3435
/// from the decryptors, aggregating client ciphertexts, and recovering the aggregation result after
3536
/// receiving partial decryption responses from the decryptors.
3637
pub struct WillowV1Server<Kahe: KaheBase, Vahe: VaheBase> {
37-
pub kahe: Kahe,
38-
pub vahe: Vahe,
38+
pub kahe: Rc<Kahe>,
39+
pub vahe: Rc<Vahe>,
3940
}
4041

4142
impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Server<Kahe, Vahe> {
@@ -381,45 +382,33 @@ mod tests {
381382
generate_aggregation_config(DEFAULT_VECTOR_ID.to_string(), 16, 10, 1, 1);
382383
let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;
383384

384-
// Create client.
385-
let kahe =
385+
// Create common KAHE/VAHE instances.
386+
let kahe = Rc::new(
386387
ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
387-
.unwrap();
388-
let vahe = ShellVahe::new(
389-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
390-
CONTEXT_STRING,
391-
)
392-
.unwrap();
393-
let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
388+
.unwrap(),
389+
);
390+
let vahe = Rc::new(
391+
ShellVahe::new(
392+
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
393+
CONTEXT_STRING,
394+
)
395+
.unwrap(),
396+
);
397+
398+
// Create client.
399+
let client = WillowV1Client::new_with_randomly_generated_seed(kahe.clone(), vahe.clone())?;
394400

395401
// Create decryptor.
396-
let vahe = ShellVahe::new(
397-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
398-
CONTEXT_STRING,
399-
)
400-
.unwrap();
402+
401403
let mut decryptor_state = DecryptorState::default();
402-
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
404+
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe))?;
403405

404406
// Create server.
405-
let kahe =
406-
ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
407-
.unwrap();
408-
let vahe = ShellVahe::new(
409-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
410-
CONTEXT_STRING,
411-
)
412-
.unwrap();
413-
let server = WillowV1Server { kahe, vahe };
407+
let server = WillowV1Server { kahe: kahe.clone(), vahe: Rc::clone(&vahe) };
414408
let mut server_state = ServerState::default();
415409

416410
// Create verifier.
417-
let vahe = ShellVahe::new(
418-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
419-
CONTEXT_STRING,
420-
)
421-
.unwrap();
422-
let verifier = WillowV1Verifier { vahe };
411+
let verifier = WillowV1Verifier { vahe: vahe.clone() };
423412
let mut verifier_state = VerifierState::default();
424413

425414
// Check empty state serialization

‎willow/src/willow_v1/verifier.rs‎

Lines changed: 22 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,13 @@ use protobuf::{proto, AsView};
1919
use shell_ciphertexts_rust_proto::ShellAhePartialDecCiphertext;
2020
use status::StatusError;
2121
use std::fmt::Debug;
22+
use std::rc::Rc;
2223
use vahe_traits::{EncryptVerify, HasVahe, VaheBase};
2324
use verifier_traits::SecureAggregationVerifier;
2425

2526
/// The verifier struct, containing a WillowCommon instance.
2627
pub struct WillowV1Verifier<Vahe: VaheBase> {
27-
pub vahe: Vahe,
28+
pub vahe: Rc<Vahe>,
2829
}
2930

3031
impl<Vahe: VaheBase> HasVahe for WillowV1Verifier<Vahe> {
@@ -294,46 +295,34 @@ mod tests {
294295
generate_aggregation_config(DEFAULT_VECTOR_ID.to_string(), 16, 10, 1, 1);
295296
let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;
296297

297-
// Create client.
298-
let kahe =
298+
// Create common KAHE/VAHE instances.
299+
let kahe = Rc::new(
299300
ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
301+
.unwrap(),
302+
);
303+
let vahe = Rc::new(
304+
ShellVahe::new(
305+
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
306+
CONTEXT_STRING,
307+
)
308+
.unwrap(),
309+
);
310+
311+
// Create client.
312+
let client =
313+
WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
300314
.unwrap();
301-
let vahe = ShellVahe::new(
302-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
303-
CONTEXT_STRING,
304-
)
305-
.unwrap();
306-
let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
307-
308-
// Create decryptor, which needs its own `vahe` (with same public polynomials
309-
// generated from the seeds) and `prng`.
310-
let vahe = ShellVahe::new(
311-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
312-
CONTEXT_STRING,
313-
)
314-
.unwrap();
315+
316+
// Create decryptor.
315317
let mut decryptor_state = DecryptorState::default();
316-
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
318+
let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe))?;
317319

318320
// Create server.
319-
let kahe =
320-
ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
321-
.unwrap();
322-
let vahe = ShellVahe::new(
323-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
324-
CONTEXT_STRING,
325-
)
326-
.unwrap();
327-
let server = WillowV1Server { kahe, vahe };
321+
let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
328322
let mut server_state = ServerState::default();
329323

330324
// Create verifier.
331-
let vahe = ShellVahe::new(
332-
create_shell_ahe_config(max_number_of_decryptors).unwrap(),
333-
CONTEXT_STRING,
334-
)
335-
.unwrap();
336-
let verifier = WillowV1Verifier { vahe };
325+
let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
337326

338327
// Decryptor generates public key share.
339328
let public_key_share = decryptor.create_public_key_share(&mut decryptor_state)?;

0 commit comments

Comments
 (0)