Create Shell decryptor util to test client functionality.

tholop · copybara-github · commit 4145e0de90ac · 2025-12-29T08:58:31.000-08:00
PiperOrigin-RevId: 845392033
diff --git a/willow/src/testing_utils/BUILD b/willow/src/testing_utils/BUILD
@@ -89,3 +89,25 @@ rust_test(
         "@crate_index//:googletest",
     ],
 )
+
+rust_library(
+    name = "shell_testing_decryptor",
+    testonly = 1,
+    srcs = [
+        "shell_testing_decryptor.rs",
+    ],
+    deps = [
+        ":shell_testing_parameters",
+        "//shell_wrapper:status",
+        "//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/traits:ahe_traits",
+        "//willow/src/traits:kahe_traits",
+        "//willow/src/traits:messages",
+        "//willow/src/traits:prng_traits",
+        "//willow/src/traits:vahe_traits",
+    ],
+)
diff --git a/willow/src/testing_utils/shell_testing_decryptor.rs b/willow/src/testing_utils/shell_testing_decryptor.rs
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+use aggregation_config::AggregationConfig;
+use ahe_traits::{AheBase, AheKeygen, PartialDec};
+use kahe_shell::ShellKahe;
+use kahe_traits::{KaheBase, KaheDecrypt, TrySecretKeyFrom};
+use messages::ClientMessage;
+use parameters_shell::create_shell_configs;
+use prng_traits::SecurePrng;
+use single_thread_hkdf::SingleThreadHkdfPrng;
+use status::{StatusError, StatusErrorCode};
+use vahe_shell::ShellVahe;
+use vahe_traits::Recover;
+
+/// Basic implementation of a single decryptor that uses Shell operations directly. Useful for
+/// testing Shell clients, by checking that encrypted messages can be decrypted properly.
+pub struct ShellTestingDecryptor {
+    kahe: ShellKahe,
+    vahe: ShellVahe,
+    prng: SingleThreadHkdfPrng,
+    secret_key: Option<<ShellVahe as AheBase>::SecretKeyShare>,
+}
+
+impl ShellTestingDecryptor {
+    /// Creates a new ShellTestingDecryptor, using the given context string to seed KAHE and AHE
+    /// public parameters.
+    pub fn new(
+        aggregation_config: &AggregationConfig,
+        context_string: &[u8],
+    ) -> Result<ShellTestingDecryptor, StatusError> {
+        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 seed = SingleThreadHkdfPrng::generate_seed()?;
+        let prng = SingleThreadHkdfPrng::create(&seed)?;
+        Ok(ShellTestingDecryptor { kahe, vahe, 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)?;
+        self.secret_key = Some(sk_share);
+        let public_key = self.vahe.aggregate_public_key_shares(&[pk_share])?;
+        Ok(public_key)
+    }
+
+    /// Decrypts a client message using the stored AHE secret key, by recovering the KAHE key from
+    /// the AHE ciphertext and then decrypting the KAHE ciphertext. Does not verify the client proof
+    /// contained in the message.
+    pub fn decrypt(
+        &mut self,
+        client_message: &ClientMessage<ShellKahe, ShellVahe>,
+    ) -> Result<<ShellKahe as KaheBase>::Plaintext, StatusError> {
+        let decryption_request =
+            self.vahe.get_partial_dec_ciphertext(&client_message.ahe_ciphertext)?;
+        let rest_of_ciphertext =
+            self.vahe.get_recover_ciphertext(&client_message.ahe_ciphertext)?;
+        match &self.secret_key {
+            None => Err(StatusError::new_with_current_location(
+                StatusErrorCode::InvalidArgument,
+                "No secret key available",
+            )),
+            Some(sk_share) => {
+                let partial_decryption =
+                    self.vahe.partial_decrypt(&decryption_request, sk_share, &mut self.prng)?;
+                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)?;
+                let decrypted_plaintext =
+                    self.kahe.decrypt(&client_message.kahe_ciphertext, &decrypted_kahe_key)?;
+                Ok(decrypted_plaintext)
+            }
+        }
+    }
+}
diff --git a/willow/src/willow_v1/BUILD b/willow/src/willow_v1/BUILD
@@ -43,9 +43,11 @@ rust_test(
         "@crate_index//:googletest",
         "//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",
     ],
diff --git a/willow/src/willow_v1/client.rs b/willow/src/willow_v1/client.rs
@@ -15,7 +15,6 @@
 use client_traits::SecureAggregationClient;
 use kahe_traits::{HasKahe, KaheBase, KaheEncrypt, KaheKeygen, TrySecretKeyInto};
 use messages::{ClientMessage, DecryptorPublicKey};
-use prng_traits::SecurePrng;
 use vahe_traits::{HasVahe, VaheBase, VerifiableEncrypt};
 
 /// Lightweight client directly exposing KAHE/VAHE types.
@@ -84,18 +83,17 @@ mod test {
     use super::*;
 
     use aggregation_config::AggregationConfig;
-    use ahe_traits::{AheBase, AheKeygen, PartialDec};
+    use ahe_traits::AheBase;
     use googletest::prelude::container_eq;
     use googletest::{gtest, verify_eq, verify_that};
     use kahe_shell::ShellKahe;
-    use kahe_traits::{KaheDecrypt, TrySecretKeyFrom};
+    use parameters_shell::create_shell_configs;
     use prng_traits::SecurePrng;
-    use shell_testing_parameters::{make_ahe_config, make_kahe_config};
+    use shell_testing_decryptor::ShellTestingDecryptor;
     use single_thread_hkdf::SingleThreadHkdfPrng;
     use std::collections::HashMap;
     use testing_utils::generate_random_nonce;
     use vahe_shell::ShellVahe;
-    use vahe_traits::Recover;
 
     const CONTEXT_STRING: &[u8] = b"test_context_string";
 
@@ -111,36 +109,27 @@ mod test {
         };
 
         // Create a client.
-        let kahe = ShellKahe::new(make_kahe_config(&aggregation_config), CONTEXT_STRING).unwrap();
-        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING).unwrap();
+        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 client_seed = SingleThreadHkdfPrng::generate_seed()?;
         let prng = SingleThreadHkdfPrng::create(&client_seed)?;
         let mut client = WillowV1Client { kahe, vahe, prng };
 
         // Generate AHE keys.
-        let kahe = ShellKahe::new(make_kahe_config(&aggregation_config), CONTEXT_STRING).unwrap();
-        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING).unwrap();
-        let seed = SingleThreadHkdfPrng::generate_seed()?;
-        let mut prng = SingleThreadHkdfPrng::create(&seed)?;
-        let (sk_share, pk_share, _) = vahe.key_gen(&mut prng)?;
-        let public_key = vahe.aggregate_public_key_shares(&[pk_share])?;
+        let mut testing_decryptor =
+            ShellTestingDecryptor::new(&aggregation_config, CONTEXT_STRING)?;
+        let public_key = testing_decryptor.generate_public_key()?;
 
         // Create client message.
         let input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
         let client_plaintext = HashMap::from([(default_id.as_str(), input_values.as_slice())]);
         let nonce = generate_random_nonce();
         let client_message =
-            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
+            client.create_client_message(&client_plaintext, &public_key, &nonce)?;
 
         // Decrypt client message.
-        let decryption_request = vahe.get_partial_dec_ciphertext(&client_message.ahe_ciphertext)?;
-        let rest_of_ciphertext = vahe.get_recover_ciphertext(&client_message.ahe_ciphertext)?;
-        let partial_decryption = vahe.partial_decrypt(&decryption_request, &sk_share, &mut prng)?;
-        let decrypted_kahe_key = vahe.recover(&partial_decryption, &rest_of_ciphertext, None)?;
-        let decrypted_kahe_key = kahe.try_secret_key_from(decrypted_kahe_key)?;
-        let decrypted_plaintext =
-            kahe.decrypt(&client_message.kahe_ciphertext, &decrypted_kahe_key)?;
-
+        let decrypted_plaintext = testing_decryptor.decrypt(&client_message)?;
         verify_that!(decrypted_plaintext.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
         let client_plaintext_length = client_plaintext.get(default_id.as_str()).unwrap().len();
         verify_eq!(
@@ -161,26 +150,25 @@ mod test {
         };
 
         // Create a client.
-        let kahe = ShellKahe::new(make_kahe_config(&aggregation_config), CONTEXT_STRING).unwrap();
-        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING).unwrap();
+        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_seed = SingleThreadHkdfPrng::generate_seed()?;
         let prng = SingleThreadHkdfPrng::create(&client1_seed)?;
         let mut client1 = WillowV1Client { kahe, vahe, prng };
 
         // Create a second client.
-        let kahe = ShellKahe::new(make_kahe_config(&aggregation_config), CONTEXT_STRING).unwrap();
-        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING).unwrap();
+        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 client2_seed = SingleThreadHkdfPrng::generate_seed()?;
         let prng = SingleThreadHkdfPrng::create(&client2_seed)?;
         let mut client2 = WillowV1Client { kahe, vahe, prng };
 
         // Generate AHE keys.
-        let kahe = ShellKahe::new(make_kahe_config(&aggregation_config), CONTEXT_STRING).unwrap();
-        let vahe = ShellVahe::new(make_ahe_config(), CONTEXT_STRING).unwrap();
-        let seed = SingleThreadHkdfPrng::generate_seed()?;
-        let mut prng = SingleThreadHkdfPrng::create(&seed)?;
-        let (sk_share, pk_share, _) = vahe.key_gen(&mut prng)?;
-        let public_key = vahe.aggregate_public_key_shares(&[pk_share])?;
+        let mut testing_decryptor =
+            ShellTestingDecryptor::new(&aggregation_config, CONTEXT_STRING)?;
+        let public_key = testing_decryptor.generate_public_key()?;
 
         // Create client messages.
         let input_values1 = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
@@ -190,30 +178,23 @@ mod test {
         let expected_output = vec![2, 3, 5, 7, 10, 14, 10, 9, 11, 11, 14, 8, 6, 9, 1];
         let nonce1 = generate_random_nonce();
         let mut client_message =
-            client1.create_client_message(&client1_plaintext, &public_key, &nonce1).unwrap();
+            client1.create_client_message(&client1_plaintext, &public_key, &nonce1)?;
         let nonce2 = generate_random_nonce();
         let extra_message =
-            client2.create_client_message(&client2_plaintext, &public_key, &nonce2).unwrap();
+            client2.create_client_message(&client2_plaintext, &public_key, &nonce2)?;
 
         // Add extra message to the first client message.
-        kahe.add_ciphertexts_in_place(
+        client2.kahe.add_ciphertexts_in_place(
             &extra_message.kahe_ciphertext,
             &mut client_message.kahe_ciphertext,
         )?;
-        vahe.add_ciphertexts_in_place(
+        client2.vahe.add_ciphertexts_in_place(
             &extra_message.ahe_ciphertext,
             &mut client_message.ahe_ciphertext,
         )?;
 
         // Decrypt client message.
-        let decryption_request = vahe.get_partial_dec_ciphertext(&client_message.ahe_ciphertext)?;
-        let rest_of_ciphertext = vahe.get_recover_ciphertext(&client_message.ahe_ciphertext)?;
-        let partial_decryption = vahe.partial_decrypt(&decryption_request, &sk_share, &mut prng)?;
-        let decrypted_kahe_key = vahe.recover(&partial_decryption, &rest_of_ciphertext, None)?;
-        let decrypted_kahe_key = kahe.try_secret_key_from(decrypted_kahe_key)?;
-        let decrypted_plaintext =
-            kahe.decrypt(&client_message.kahe_ciphertext, &decrypted_kahe_key)?;
-
+        let decrypted_plaintext = testing_decryptor.decrypt(&client_message)?;
         verify_that!(decrypted_plaintext.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
         let client_plaintext_length = client1_plaintext.get(default_id.as_str()).unwrap().len();
         verify_eq!(
