Wrap Kahe and Vahe struct members with Rc

willow/benches/shell_benchmarks.rs

@@ -15,6 +15,7 @@
 use clap::Parser;
 use std::collections::HashMap;
 use std::hint::black_box;
+use std::rc::Rc;
 use std::time::Duration;
 
 use aggregation_config::AggregationConfig;
@@ -126,25 +127,25 @@ fn setup_base(args: &Args) -> BaseInputs {
     let ahe_config = create_shell_ahe_config(aggregation_config.max_number_of_decryptors).unwrap();
     let kahe_config = create_shell_kahe_config(&aggregation_config).unwrap();
 
+    // Create common KAHE/VAHE instances.
+    let kahe = Rc::new(ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap());
+    let vahe = Rc::new(ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap());
+
     // Create client.
-    let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();
-    let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
-    let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe).unwrap();
+    let client =
+        WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
+            .unwrap();
 
     // Create decryptor.
-    let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
     let mut decryptor_state = DecryptorState::default();
-    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe).unwrap();
+    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();
 
     // Create server.
-    let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();
-    let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
-    let server = WillowV1Server { kahe, vahe };
+    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
     let mut server_state = ServerState::default();
 
     // Create verifier.
-    let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
-    let verifier = WillowV1Verifier { vahe };
+    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
     let verifier_state = VerifierState::default();
 
     // Decryptor generates public key share.

willow/src/api/client.rs

@@ -28,6 +28,7 @@ use shell_ciphertexts_rust_proto::ShellAhePublicKey;
 use status::ffi::FfiStatus;
 use status::StatusError;
 use std::collections::HashMap;
+use std::rc::Rc;
 use vahe_shell::ShellVahe;
 use willow_v1_client::WillowV1Client;
 
@@ -82,8 +83,8 @@ impl WillowShellClient {
         let aggregation_config = AggregationConfig::from_proto(aggregation_config_proto, ())?;
         let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
         let context_bytes = &aggregation_config.key_id;
-        let kahe = ShellKahe::new(kahe_config, &context_bytes)?;
-        let vahe = ShellVahe::new(ahe_config, &context_bytes)?;
+        let kahe = Rc::new(ShellKahe::new(kahe_config, &context_bytes)?);
+        let vahe = Rc::new(ShellVahe::new(ahe_config, &context_bytes)?);
         let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
         Ok(WillowShellClient(client))
     }
@@ -102,7 +103,7 @@ impl WillowShellClient {
         }
         let public_key_proto = ShellAhePublicKey::parse(public_key.as_bytes())
             .map_err(|e| status::internal(format!("Failed to parse ShellAhePublicKey: {}", e)))?;
-        let public_key_rust = PublicKey::from_proto(public_key_proto, &self.0.vahe)?;
+        let public_key_rust = PublicKey::from_proto(public_key_proto, self.0.vahe.as_ref())?;
         let message = self.0.create_client_message(&plaintext_slice, &public_key_rust, nonce)?;
         Ok(message
             .to_proto(&self.0)?

willow/src/api/server_accumulator.rs

@@ -31,6 +31,7 @@ use server_traits::SecureAggregationServer;
 use status::StatusError;
 use std::collections::BTreeMap;
 use std::ops::Range;
+use std::rc::Rc;
 use vahe_shell::ShellVahe;
 use verifier_traits::SecureAggregationVerifier;
 use willow_v1_server::{ServerState, WillowV1Server};
@@ -147,11 +148,10 @@ impl ServerAccumulator {
     fn new(aggregation_config: AggregationConfig) -> Result<Self, StatusError> {
         let (kahe_config, vahe_config) = create_shell_configs(&aggregation_config)?;
         let context_bytes = &aggregation_config.key_id;
-        let server_kahe = ShellKahe::new(kahe_config, context_bytes)?;
-        let server_vahe = ShellVahe::new(vahe_config.clone(), context_bytes)?;
-        let verifier_vahe = ShellVahe::new(vahe_config, context_bytes)?;
-        let server = WillowV1Server { kahe: server_kahe, vahe: server_vahe };
-        let verifier = WillowV1Verifier { vahe: verifier_vahe };
+        let kahe = Rc::new(ShellKahe::new(kahe_config, &context_bytes)?);
+        let vahe = Rc::new(ShellVahe::new(vahe_config, &context_bytes)?);
+        let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
+        let verifier = WillowV1Verifier { vahe };
         Ok(Self {
             server: server,
             server_state: Default::default(),
@@ -661,8 +661,8 @@ impl FinalResultDecryptor {
         let aggregation_config = AggregationConfig::from_proto(aggregation_config_proto, ())?;
         let (kahe_config, vahe_config) = create_shell_configs(&aggregation_config)?;
         let context_bytes = &aggregation_config.key_id;
-        let kahe = ShellKahe::new(kahe_config, context_bytes)?;
-        let vahe = ShellVahe::new(vahe_config, context_bytes)?;
+        let kahe = Rc::new(ShellKahe::new(kahe_config, context_bytes)?);
+        let vahe = Rc::new(ShellVahe::new(vahe_config, context_bytes)?);
         let server = WillowV1Server { kahe, vahe };
         let server_state = ServerState::from_proto(server_state_proto, &server)?;
 

willow/src/willow_v1/client.rs

@@ -17,33 +17,34 @@ use kahe_traits::{HasKahe, KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
 use messages::{ClientMessage, DecryptorPublicKey};
 use prng_traits::SecurePrng;
 use std::cell::RefCell;
+use std::rc::Rc;
 use vahe_traits::{HasVahe, VaheBase, VerifiableEncrypt};
 
 /// Lightweight client directly exposing KAHE/VAHE types.
 pub struct WillowV1Client<Kahe: KaheBase, Vahe: VaheBase> {
-    pub kahe: Kahe,
-    pub vahe: Vahe,
+    pub kahe: Rc<Kahe>,
+    pub vahe: Rc<Vahe>,
     pub prng: RefCell<Kahe::Rng>, // Using a single PRNG for both VAHE and KAHE.
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Client<Kahe, Vahe> {
     type Kahe = Kahe;
     fn kahe(&self) -> &Self::Kahe {
-        &self.kahe
+        self.kahe.as_ref()
     }
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase> HasVahe for WillowV1Client<Kahe, Vahe> {
     type Vahe = Vahe;
     fn vahe(&self) -> &Self::Vahe {
-        &self.vahe
+        self.vahe.as_ref()
     }
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase> WillowV1Client<Kahe, Vahe> {
     pub fn new_with_randomly_generated_seed(
-        kahe: Kahe,
-        vahe: Vahe,
+        kahe: Rc<Kahe>,
+        vahe: Rc<Vahe>,
     ) -> Result<Self, status::StatusError> {
         let seed = Kahe::Rng::generate_seed()?;
         let prng = RefCell::new(Kahe::Rng::create(&seed)?);
@@ -122,8 +123,8 @@ mod test {
 
         // Create a client.
         let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
-        let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
-        let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
+        let kahe = Rc::new(ShellKahe::new(kahe_config, CONTEXT_STRING)?);
+        let vahe = Rc::new(ShellVahe::new(ahe_config, CONTEXT_STRING)?);
         let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
 
         // Generate AHE keys.
@@ -159,16 +160,14 @@ mod test {
             key_id: b"test".to_vec(),
         };
 
-        // Create a client.
+        // Create common KAHE/VAHE instances.
         let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
-        let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
-        let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
-        let client1 = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
+        let kahe = Rc::new(ShellKahe::new(kahe_config, CONTEXT_STRING)?);
+        let vahe = Rc::new(ShellVahe::new(ahe_config, CONTEXT_STRING)?);
 
-        // Create a second client.
-        let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?;
-        let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
-        let vahe = ShellVahe::new(ahe_config, CONTEXT_STRING)?;
+        // Create clients.
+        let client1 =
+            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))?;
         let client2 = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
 
         // Generate AHE keys.

willow/src/willow_v1/decryptor.rs

@@ -22,12 +22,13 @@ use protobuf::AsView;
 use shell_ciphertexts_rust_proto::ShellAheSecretKeyShare;
 use status::StatusError;
 use std::cell::RefCell;
+use std::rc::Rc;
 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.
 pub struct WillowV1Decryptor<Vahe: VaheBase> {
-    pub vahe: Vahe,
+    pub vahe: Rc<Vahe>,
     pub prng: RefCell<Vahe::Rng>,
 }
 
@@ -39,7 +40,7 @@ impl<Vahe: VaheBase> HasVahe for WillowV1Decryptor<Vahe> {
 }
 
 impl<Vahe: VaheBase> WillowV1Decryptor<Vahe> {
-    pub fn new_with_randomly_generated_seed(vahe: Vahe) -> Result<Self, status::StatusError> {
+    pub fn new_with_randomly_generated_seed(vahe: Rc<Vahe>) -> Result<Self, status::StatusError> {
         let seed = Vahe::Rng::generate_seed()?;
         let prng = RefCell::new(Vahe::Rng::create(&seed)?);
         Ok(Self { vahe, prng })
@@ -145,13 +146,15 @@ mod tests {
     use googletest::{gtest, verify_true};
     use parameters_shell::create_shell_ahe_config;
     use proto_serialization_traits::{FromProto, ToProto};
+    use std::rc::Rc;
     use vahe_shell::ShellVahe;
 
     const CONTEXT_STRING: &[u8] = b"testing_context_string";
 
     #[gtest]
     fn decryptor_state_serialization_roundtrip() -> googletest::Result<()> {
-        let vahe = ShellVahe::new(create_shell_ahe_config(1).unwrap(), CONTEXT_STRING).unwrap();
+        let vahe =
+            Rc::new(ShellVahe::new(create_shell_ahe_config(1).unwrap(), CONTEXT_STRING).unwrap());
         let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
         let mut decryptor_state = DecryptorState::default();
 

willow/src/willow_v1/server.rs

@@ -28,14 +28,15 @@ use shell_ciphertexts_rust_proto::{
 };
 use status::StatusError;
 use std::collections::HashMap;
+use std::rc::Rc;
 use vahe_traits::{EncryptVerify, HasVahe, Recover, VaheBase};
 
 /// Implements the `server` role in the Willow protocol. This includes aggregating public key shares
 /// from the decryptors, aggregating client ciphertexts, and recovering the aggregation result after
 /// receiving partial decryption responses from the decryptors.
 pub struct WillowV1Server<Kahe: KaheBase, Vahe: VaheBase> {
-    pub kahe: Kahe,
-    pub vahe: Vahe,
+    pub kahe: Rc<Kahe>,
+    pub vahe: Rc<Vahe>,
 }
 
 impl<Kahe: KaheBase, Vahe: VaheBase> HasKahe for WillowV1Server<Kahe, Vahe> {
@@ -381,45 +382,33 @@ mod tests {
             generate_aggregation_config(DEFAULT_VECTOR_ID.to_string(), 16, 10, 1, 1);
         let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;
 
-        // Create client.
-        let kahe =
+        // Create common KAHE/VAHE instances.
+        let kahe = Rc::new(
             ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
-                .unwrap();
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
+                .unwrap(),
+        );
+        let vahe = Rc::new(
+            ShellVahe::new(
+                create_shell_ahe_config(max_number_of_decryptors).unwrap(),
+                CONTEXT_STRING,
+            )
+            .unwrap(),
+        );
+
+        // Create client.
+        let client = WillowV1Client::new_with_randomly_generated_seed(kahe.clone(), vahe.clone())?;
 
         // Create decryptor.
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
+
         let mut decryptor_state = DecryptorState::default();
-        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
+        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe))?;
 
         // Create server.
-        let kahe =
-            ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
-                .unwrap();
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let server = WillowV1Server { kahe, vahe };
+        let server = WillowV1Server { kahe: kahe.clone(), vahe: Rc::clone(&vahe) };
         let mut server_state = ServerState::default();
 
         // Create verifier.
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let verifier = WillowV1Verifier { vahe };
+        let verifier = WillowV1Verifier { vahe: vahe.clone() };
         let mut verifier_state = VerifierState::default();
 
         // Check empty state serialization

willow/src/willow_v1/verifier.rs

@@ -19,12 +19,13 @@ use protobuf::{proto, AsView};
 use shell_ciphertexts_rust_proto::ShellAhePartialDecCiphertext;
 use status::StatusError;
 use std::fmt::Debug;
+use std::rc::Rc;
 use vahe_traits::{EncryptVerify, HasVahe, VaheBase};
 use verifier_traits::SecureAggregationVerifier;
 
 /// The verifier struct, containing a WillowCommon instance.
 pub struct WillowV1Verifier<Vahe: VaheBase> {
-    pub vahe: Vahe,
+    pub vahe: Rc<Vahe>,
 }
 
 impl<Vahe: VaheBase> HasVahe for WillowV1Verifier<Vahe> {
@@ -294,46 +295,34 @@ mod tests {
             generate_aggregation_config(DEFAULT_VECTOR_ID.to_string(), 16, 10, 1, 1);
         let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;
 
-        // Create client.
-        let kahe =
+        // Create common KAHE/VAHE instances.
+        let kahe = Rc::new(
             ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
+                .unwrap(),
+        );
+        let vahe = Rc::new(
+            ShellVahe::new(
+                create_shell_ahe_config(max_number_of_decryptors).unwrap(),
+                CONTEXT_STRING,
+            )
+            .unwrap(),
+        );
+
+        // Create client.
+        let client =
+            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                 .unwrap();
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let client = WillowV1Client::new_with_randomly_generated_seed(kahe, vahe)?;
-
-        // Create decryptor, which needs its own `vahe` (with same public polynomials
-        // generated from the seeds) and `prng`.
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
+
+        // Create decryptor.
         let mut decryptor_state = DecryptorState::default();
-        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
+        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe))?;
 
         // Create server.
-        let kahe =
-            ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
-                .unwrap();
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let server = WillowV1Server { kahe, vahe };
+        let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
         let mut server_state = ServerState::default();
 
         // Create verifier.
-        let vahe = ShellVahe::new(
-            create_shell_ahe_config(max_number_of_decryptors).unwrap(),
-            CONTEXT_STRING,
-        )
-        .unwrap();
-        let verifier = WillowV1Verifier { vahe };
+        let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
 
         // Decryptor generates public key share.
         let public_key_share = decryptor.create_public_key_share(&mut decryptor_state)?;