Use interior mutability for Decryptor PRNG too; add constructor; clean up imports.

PiperOrigin-RevId: 860138561

Author: tholop

willow/benches/shell_benchmarks.rs

@@ -28,9 +28,7 @@ use messages::{
     PartialDecryptionRequest,
 };
 use parameters_shell::{create_shell_ahe_config, create_shell_kahe_config};
-use prng_traits::SecurePrng;
 use server_traits::SecureAggregationServer;
-use single_thread_hkdf::SingleThreadHkdfPrng;
 use testing_utils::{generate_random_nonce, generate_random_unsigned_vector};
 use vahe_shell::ShellVahe;
 use verifier_traits::SecureAggregationVerifier;
@@ -135,10 +133,8 @@ fn setup_base(args: &Args) -> BaseInputs {
 
     // Create decryptor.
     let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
-    let seed = SingleThreadHkdfPrng::generate_seed().unwrap();
-    let prng = SingleThreadHkdfPrng::create(&seed).unwrap();
     let mut decryptor_state = DecryptorState::default();
-    let mut decryptor = WillowV1Decryptor { vahe, prng };
+    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe).unwrap();
 
     // Create server.
     let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();

willow/src/api/client.rs

@@ -22,11 +22,9 @@ use client_traits::SecureAggregationClient;
 use kahe_shell::ShellKahe;
 use kahe_traits::KaheBase;
 use parameters_shell::create_shell_configs;
-use prng_traits::SecurePrng;
 use proto_serialization_traits::{FromProto, ToProto};
 use protobuf::prelude::*;
 use shell_ciphertexts_rust_proto::ShellAhePublicKey;
-use single_thread_hkdf::SingleThreadHkdfPrng;
 use status::ffi::FfiStatus;
 use status::StatusError;
 use std::collections::HashMap;

willow/src/testing_utils/shell_testing_decryptor.rs

@@ -31,6 +31,7 @@ use protobuf::prelude::*;
 use single_thread_hkdf::SingleThreadHkdfPrng;
 use status::ffi::FfiStatus;
 use status::{StatusError, StatusErrorCode};
+use std::cell::RefCell;
 use vahe_shell::ShellVahe;
 use vahe_traits::Recover;
 use vahe_traits::{HasVahe, VaheBase};
@@ -41,7 +42,7 @@ use vahe_traits::{HasVahe, VaheBase};
 pub struct ShellTestingDecryptor {
     kahe: ShellKahe,
     vahe: ShellVahe,
-    prng: SingleThreadHkdfPrng,
+    prng: RefCell<SingleThreadHkdfPrng>,
     secret_key: Option<<ShellVahe as AheBase>::SecretKeyShare>,
 }
 
@@ -64,14 +65,14 @@ impl ShellTestingDecryptor {
         let vahe = ShellVahe::new(ahe_config, context_bytes)?;
         let seed = SingleThreadHkdfPrng::generate_seed()?;
         let prng = SingleThreadHkdfPrng::create(&seed)?;
-        Ok(ShellTestingDecryptor { kahe, vahe, prng, secret_key: None })
+        Ok(ShellTestingDecryptor { kahe, vahe, prng: RefCell::new(prng), secret_key: None })
     }
 
     /// Generates a new AHE public key, and stores the corresponding secret key.
     pub fn generate_public_key(
         &mut self,
     ) -> Result<<ShellVahe as AheBase>::PublicKey, StatusError> {
-        let (sk_share, pk_share, _) = self.vahe.key_gen(&mut self.prng)?;
+        let (sk_share, pk_share, _) = self.vahe.key_gen(&mut self.prng.borrow_mut())?;
         self.secret_key = Some(sk_share);
         let public_key = self.vahe.aggregate_public_key_shares(&[pk_share])?;
         Ok(public_key)
@@ -81,7 +82,7 @@ impl ShellTestingDecryptor {
     /// the AHE ciphertext and then decrypting the KAHE ciphertext. Does not verify the client proof
     /// contained in the message.
     pub fn decrypt(
-        &mut self,
+        &self,
         client_message: &ClientMessage<ShellKahe, ShellVahe>,
     ) -> Result<<ShellKahe as KaheBase>::Plaintext, StatusError> {
         let partial_dec_ciphertext =
@@ -94,8 +95,11 @@ impl ShellTestingDecryptor {
                 "No secret key available",
             )),
             Some(sk_share) => {
-                let partial_decryption =
-                    self.vahe.partial_decrypt(&partial_dec_ciphertext, sk_share, &mut self.prng)?;
+                let partial_decryption = self.vahe.partial_decrypt(
+                    &partial_dec_ciphertext,
+                    sk_share,
+                    &mut self.prng.borrow_mut(),
+                )?;
                 let decrypted_kahe_key =
                     self.vahe.recover(&partial_decryption, &rest_of_ciphertext, None)?;
                 let decrypted_kahe_key = self.kahe.try_secret_key_from(decrypted_kahe_key)?;
@@ -134,7 +138,7 @@ impl ShellTestingDecryptor {
     }
 
     fn decrypt_serialized(
-        &mut self,
+        &self,
         contribution: &[u8],
     ) -> Result<Vec<ffi::EncodedDataEntry>, StatusError> {
         let client_message_proto = ClientMessageProto::parse(contribution)
@@ -192,7 +196,7 @@ impl ShellTestingDecryptor {
                 let partial_decryption = self.vahe.partial_decrypt(
                     &request.partial_dec_ciphertext,
                     sk_share,
-                    &mut self.prng,
+                    &mut self.prng.borrow_mut(),
                 )?;
                 Ok(PartialDecryptionResponse { partial_decryption })
             }

willow/src/traits/decryptor.rs

@@ -24,14 +24,14 @@ pub trait SecureAggregationDecryptor: HasVahe {
     /// Creates a public key share to be sent to the Server, updating the
     /// decryptor state.
     fn create_public_key_share(
-        &mut self,
+        &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(
-        &mut self,
+        &self,
         partial_decryption_request: PartialDecryptionRequest<<Self as HasVahe>::Vahe>,
         decryptor_state: &Self::DecryptorState,
     ) -> Result<PartialDecryptionResponse<<Self as HasVahe>::Vahe>, StatusError>;

willow/src/willow_v1/BUILD

@@ -44,12 +44,10 @@ rust_test(
         "//willow/src/api:aggregation_config",
         "//willow/src/shell:kahe_shell",
         "//willow/src/shell:parameters_shell",
-        "//willow/src/shell:single_thread_hkdf",
         "//willow/src/shell:vahe_shell",
         "//willow/src/testing_utils",
         "//willow/src/testing_utils:shell_testing_decryptor",
         "//willow/src/testing_utils:shell_testing_parameters",
-        "//willow/src/traits:prng_traits",
     ],
 )
 
@@ -59,11 +57,9 @@ rust_test(
     deps = [
         "@crate_index//:googletest",
         "//willow/src/shell:parameters_shell",
-        "//willow/src/shell:single_thread_hkdf",
         "//willow/src/shell:vahe_shell",
         "//willow/src/traits:ahe_traits",
         "//willow/src/traits:decryptor_traits",
-        "//willow/src/traits:prng_traits",
         "//willow/src/traits:proto_serialization_traits",
     ],
 )
@@ -81,6 +77,7 @@ rust_library(
         "//willow/src/traits:ahe_traits",
         "//willow/src/traits:decryptor_traits",
         "//willow/src/traits:messages",
+        "//willow/src/traits:prng_traits",
         "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:vahe_traits",
     ],
@@ -96,13 +93,11 @@ rust_test(
         "@crate_index//:googletest",
         "//willow/src/shell:kahe_shell",
         "//willow/src/shell:parameters_shell",
-        "//willow/src/shell:single_thread_hkdf",
         "//willow/src/shell:vahe_shell",
         "//willow/src/testing_utils",
         "//willow/src/traits:ahe_traits",
         "//willow/src/traits:client_traits",
         "//willow/src/traits:decryptor_traits",
-        "//willow/src/traits:prng_traits",
         "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:server_traits",
         "//willow/src/traits:verifier_traits",
@@ -158,15 +153,13 @@ rust_test(
         "//shell_wrapper:status_matchers_rs",
         "//willow/src/shell:kahe_shell",
         "//willow/src/shell:parameters_shell",
-        "//willow/src/shell:single_thread_hkdf",
         "//willow/src/shell:vahe_shell",
         "//willow/src/testing_utils",
         "//willow/src/testing_utils:shell_testing_parameters",
         "//willow/src/traits:ahe_traits",
         "//willow/src/traits:client_traits",
         "//willow/src/traits:decryptor_traits",
         "//willow/src/traits:kahe_traits",
-        "//willow/src/traits:prng_traits",
         "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:server_traits",
         "//willow/src/traits:vahe_traits",

willow/src/willow_v1/client.rs

@@ -102,9 +102,7 @@ mod test {
     use googletest::{gtest, verify_eq, verify_that};
     use kahe_shell::ShellKahe;
     use parameters_shell::create_shell_configs;
-    use prng_traits::SecurePrng;
     use shell_testing_decryptor::ShellTestingDecryptor;
-    use single_thread_hkdf::SingleThreadHkdfPrng;
     use std::collections::HashMap;
     use testing_utils::generate_random_nonce;
     use vahe_shell::ShellVahe;

willow/src/willow_v1/decryptor.rs

@@ -16,17 +16,19 @@ use ahe_traits::{AheKeygen, PartialDec};
 use decryptor_traits::SecureAggregationDecryptor;
 use messages::{DecryptorPublicKeyShare, PartialDecryptionRequest, PartialDecryptionResponse};
 use messages_rust_proto::DecryptorStateProto;
+use prng_traits::SecurePrng;
 use proto_serialization_traits::{FromProto, ToProto};
 use protobuf::AsView;
 use shell_ciphertexts_rust_proto::ShellAheSecretKeyShare;
 use status::StatusError;
+use std::cell::RefCell;
 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 prng: Vahe::Rng,
+    pub prng: RefCell<Vahe::Rng>,
 }
 
 impl<Vahe: VaheBase> HasVahe for WillowV1Decryptor<Vahe> {
@@ -36,6 +38,14 @@ 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> {
+        let seed = Vahe::Rng::generate_seed()?;
+        let prng = RefCell::new(Vahe::Rng::create(&seed)?);
+        Ok(Self { vahe, prng })
+    }
+}
+
 pub struct DecryptorState<Vahe: VaheBase> {
     sk_share: Option<Vahe::SecretKeyShare>,
 }
@@ -97,18 +107,18 @@ where
     /// Creates a public key share to be sent to the Server, updating the
     /// decryptor state.
     fn create_public_key_share(
-        &mut self,
+        &self,
         decryptor_state: &mut Self::DecryptorState,
     ) -> Result<DecryptorPublicKeyShare<Vahe>, status::StatusError> {
-        let (sk_share, pk_share, _) = self.vahe.key_gen(&mut self.prng)?;
+        let (sk_share, pk_share, _) = self.vahe.key_gen(&mut self.prng.borrow_mut())?;
         decryptor_state.sk_share = Some(sk_share);
         Ok(pk_share)
     }
 
     /// Handles a partial decryption request received from the Server. Returns a
     /// partial decryption to the Server.
     fn handle_partial_decryption_request(
-        &mut self,
+        &self,
         partial_decryption_request: PartialDecryptionRequest<Vahe>,
         decryptor_state: &Self::DecryptorState,
     ) -> Result<PartialDecryptionResponse<Vahe>, status::StatusError> {
@@ -121,7 +131,7 @@ where
         let pd = self.vahe.partial_decrypt(
             &partial_decryption_request.partial_dec_ciphertext,
             sk_share,
-            &mut self.prng,
+            &mut self.prng.borrow_mut(),
         )?;
         Ok(PartialDecryptionResponse { partial_decryption: pd })
     }
@@ -134,19 +144,15 @@ mod tests {
     use decryptor_traits::SecureAggregationDecryptor;
     use googletest::{gtest, verify_true};
     use parameters_shell::create_shell_ahe_config;
-    use prng_traits::SecurePrng;
     use proto_serialization_traits::{FromProto, ToProto};
-    use single_thread_hkdf::SingleThreadHkdfPrng;
     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 seed = SingleThreadHkdfPrng::generate_seed()?;
-        let prng = SingleThreadHkdfPrng::create(&seed)?;
-        let mut decryptor = WillowV1Decryptor { vahe, prng };
+        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
         let mut decryptor_state = DecryptorState::default();
 
         // Check empty state serialization.

willow/src/willow_v1/server.rs

@@ -362,10 +362,8 @@ mod tests {
     use googletest::{gtest, verify_true};
     use kahe_shell::ShellKahe;
     use parameters_shell::{create_shell_ahe_config, create_shell_kahe_config};
-    use prng_traits::SecurePrng;
     use proto_serialization_traits::{FromProto, ToProto};
     use server_traits::SecureAggregationServer;
-    use single_thread_hkdf::SingleThreadHkdfPrng;
     use std::collections::HashMap;
     use testing_utils::{generate_aggregation_config, generate_random_nonce};
     use vahe_shell::ShellVahe;
@@ -400,10 +398,8 @@ mod tests {
             CONTEXT_STRING,
         )
         .unwrap();
-        let seed = SingleThreadHkdfPrng::generate_seed()?;
-        let prng = SingleThreadHkdfPrng::create(&seed)?;
         let mut decryptor_state = DecryptorState::default();
-        let mut decryptor = WillowV1Decryptor { vahe, prng };
+        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
 
         // Create server.
         let kahe =

willow/src/willow_v1/verifier.rs

@@ -271,10 +271,8 @@ mod tests {
     use kahe_shell::ShellKahe;
     use kahe_traits::KaheBase;
     use parameters_shell::{create_shell_ahe_config, create_shell_kahe_config};
-    use prng_traits::SecurePrng;
     use proto_serialization_traits::{FromProto, ToProto};
     use server_traits::SecureAggregationServer;
-    use single_thread_hkdf::SingleThreadHkdfPrng;
     use status_matchers_rs::status_is;
     use std::collections::HashMap;
     use testing_utils::{generate_aggregation_config, generate_random_nonce};
@@ -314,10 +312,8 @@ mod tests {
             CONTEXT_STRING,
         )
         .unwrap();
-        let seed = SingleThreadHkdfPrng::generate_seed()?;
-        let prng = SingleThreadHkdfPrng::create(&seed)?;
         let mut decryptor_state = DecryptorState::default();
-        let mut decryptor = WillowV1Decryptor { vahe, prng };
+        let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(vahe)?;
 
         // Create server.
         let kahe =