Skip to content

Commit 8afa063

Browse files
schoppmpcopybara-github
schoppmp
authored and
copybara-github
committed
Add nonce to VAHE and client message
PiperOrigin-RevId: 829563668
1 parent 353b9c9 commit 8afa063

File tree

18 files changed

+222
-340
lines changed

18 files changed

+222
-340
lines changed

‎willow/benches/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,13 @@ rust_library(
3333
"//willow/src/shell:single_thread_hkdf",
3434
"//willow/src/shell:vahe_shell",
3535
"//willow/src/testing_utils",
36+
"//willow/src/traits:ahe_traits",
3637
"//willow/src/traits:client_traits",
3738
"//willow/src/traits:decryptor_traits",
3839
"//willow/src/traits:kahe_traits",
3940
"//willow/src/traits:prng_traits",
4041
"//willow/src/traits:server_traits",
42+
"//willow/src/traits:vahe_traits",
4143
"//willow/src/traits:verifier_traits",
4244
"//willow/src/willow_v1:willow_v1_client",
4345
"//willow/src/willow_v1:willow_v1_common",

‎willow/benches/shell_benchmarks.rs‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,18 @@ use std::collections::HashMap;
1717
use std::hint::black_box;
1818
use std::time::Duration;
1919

20+
use ahe_traits::AheBase;
2021
use client_traits::SecureAggregationClient;
2122
use decryptor_traits::SecureAggregationDecryptor;
2223
use kahe_shell::ShellKahe;
2324
use kahe_traits::KaheBase;
25+
use parameters_shell::create_shell_configs;
2426
use prng_traits::SecurePrng;
2527
use server_traits::SecureAggregationServer;
26-
use parameters_shell::create_shell_configs;
2728
use single_thread_hkdf::SingleThreadHkdfPrng;
2829
use testing_utils::{generate_random_unsigned_vector, ShellClient, ShellClientMessage};
2930
use vahe_shell::ShellVahe;
31+
use vahe_traits::VaheBase;
3032
use verifier_traits::SecureAggregationVerifier;
3133
use willow_api_common::AggregationConfig;
3234
use willow_v1_client::WillowV1Client;
@@ -39,6 +41,7 @@ use willow_v1_server::{ServerState, WillowV1Server};
3941
use willow_v1_verifier::{VerifierState, WillowV1Verifier};
4042

4143
const DEFAULT_ID: &str = "default";
44+
const CONTEXT_STRING: &[u8] = b"benchmark_context_string";
4245

4346
#[derive(Parser, Debug)]
4447
#[command(version, about, long_about = None)]
@@ -126,21 +129,20 @@ fn setup_base(args: &Args) -> BaseInputs {
126129
};
127130
let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config).unwrap();
128131
let public_kahe_seed = SingleThreadHkdfPrng::generate_seed().unwrap();
129-
let public_ahe_seed = SingleThreadHkdfPrng::generate_seed().unwrap();
130132

131133
// Create client.
132134
let common = WillowCommon {
133-
kahe: ShellKahe::new(kahe_config.clone(), &public_kahe_seed).unwrap(),
134-
vahe: ShellVahe::new(ahe_config.clone(), &public_ahe_seed).unwrap(),
135+
kahe: ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap(),
136+
vahe: ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap(),
135137
};
136138
let seed = SingleThreadHkdfPrng::generate_seed().unwrap();
137139
let prng = SingleThreadHkdfPrng::create(&seed).unwrap();
138140
let client = ShellClient { common, prng };
139141

140142
// Create decryptor, which needs its own `common` and `prng`.
141143
let common = WillowCommon {
142-
kahe: ShellKahe::new(kahe_config.clone(), &public_kahe_seed).unwrap(),
143-
vahe: ShellVahe::new(ahe_config.clone(), &public_ahe_seed).unwrap(),
144+
kahe: ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap(),
145+
vahe: ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap(),
144146
};
145147
let seed = SingleThreadHkdfPrng::generate_seed().unwrap();
146148
let prng = SingleThreadHkdfPrng::create(&seed).unwrap();
@@ -149,16 +151,16 @@ fn setup_base(args: &Args) -> BaseInputs {
149151

150152
// Create server.
151153
let common = WillowCommon {
152-
kahe: ShellKahe::new(kahe_config.clone(), &public_kahe_seed).unwrap(),
153-
vahe: ShellVahe::new(ahe_config.clone(), &public_ahe_seed).unwrap(),
154+
kahe: ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap(),
155+
vahe: ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap(),
154156
};
155157
let server = WillowV1Server { common };
156158
let mut server_state = ServerState::new();
157159

158160
// Create verifier.
159161
let common = WillowCommon {
160-
kahe: ShellKahe::new(kahe_config.clone(), &public_kahe_seed).unwrap(),
161-
vahe: ShellVahe::new(ahe_config.clone(), &public_ahe_seed).unwrap(),
162+
kahe: ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap(),
163+
vahe: ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap(),
162164
};
163165
let verifier = WillowV1Verifier { common };
164166
let verifier_state = VerifierState::new();

‎willow/src/shell/ahe.rs‎

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -56,21 +56,6 @@ pub struct ShellAhe {
5656
}
5757

5858
impl ShellAhe {
59-
pub fn new(config: ShellAheConfig, public_seed: &Seed) -> Result<Self, status::StatusError> {
60-
let num_coeffs = 1 << config.log_n;
61-
let public_ahe_parameters = ahe::create_public_parameters(
62-
config.log_n,
63-
config.t,
64-
&config.qs,
65-
/* error_variance= */ ERROR_VARIANCE,
66-
/* s_base_flood= */ S_BASE_FLOOD,
67-
config.s_flood,
68-
&public_seed,
69-
)?;
70-
71-
Ok(Self { public_ahe_parameters, num_coeffs })
72-
}
73-
7459
/// Convenience function.
7560
fn add_vec_rns_polynomial_in_place(
7661
&self,
@@ -400,6 +385,29 @@ impl AheBase for ShellAhe {
400385

401386
type Rng = SingleThreadHkdfPrng;
402387

388+
type Config = ShellAheConfig;
389+
390+
fn new(config: Self::Config, context_string: &[u8]) -> Result<Self, status::StatusError> {
391+
let num_coeffs = 1 << config.log_n;
392+
let public_seed = single_thread_hkdf::compute_hkdf(
393+
context_string,
394+
b"",
395+
b"ShellAhe.public_seed",
396+
single_thread_hkdf::seed_length(),
397+
)?;
398+
let public_ahe_parameters = ahe::create_public_parameters(
399+
config.log_n,
400+
config.t,
401+
&config.qs,
402+
/* error_variance= */ ERROR_VARIANCE,
403+
/* s_base_flood= */ S_BASE_FLOOD,
404+
config.s_flood,
405+
&public_seed,
406+
)?;
407+
408+
Ok(Self { public_ahe_parameters, num_coeffs })
409+
}
410+
403411
fn aggregate_public_key_shares(
404412
&self,
405413
public_key_shares: &[Self::PublicKeyShare],
@@ -651,13 +659,13 @@ mod test {
651659
const NUM_DECRYPTORS: usize = 3;
652660
const NUM_CLIENTS: usize = 1000;
653661
const MAX_ABSOLUTE_VALUE: i64 = 72;
662+
const CONTEXT_STRING: &[u8] = b"test_context_string";
654663

655664
#[gtest]
656665
fn test_encrypt_decrypt_one() -> googletest::Result<()> {
657666
const NUM_VALUES: usize = 100;
658667

659-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
660-
let ahe = ShellAhe::new(make_ahe_config(), &public_seed)?;
668+
let ahe = ShellAhe::new(make_ahe_config(), CONTEXT_STRING)?;
661669

662670
let pt = vec![1, 2, 3, 4, 5, 6, 7, 8];
663671
let seed = SingleThreadHkdfPrng::generate_seed()?;
@@ -682,8 +690,7 @@ mod test {
682690
let config = make_ahe_config();
683691
let t = config.t; // Keep a copy of the plaintext modulus.
684692

685-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
686-
let ahe = ShellAhe::new(config, &public_seed)?;
693+
let ahe = ShellAhe::new(config, CONTEXT_STRING)?;
687694
let seed = SingleThreadHkdfPrng::generate_seed()?;
688695
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
689696

@@ -750,8 +757,7 @@ mod test {
750757

751758
#[gtest]
752759
fn test_errors() -> googletest::Result<()> {
753-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
754-
let ahe = ShellAhe::new(make_ahe_config(), &public_seed)?;
760+
let ahe = ShellAhe::new(make_ahe_config(), CONTEXT_STRING)?;
755761
let seed = SingleThreadHkdfPrng::generate_seed()?;
756762
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
757763

@@ -826,11 +832,10 @@ mod test {
826832

827833
#[gtest]
828834
fn test_manual_encryption() -> googletest::Result<()> {
829-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
830835
let config = make_ahe_config();
831836
let q: i128 = config.qs.iter().map(|x| *x as i128).product();
832837

833-
let ahe = ShellAhe::new(config, &public_seed)?;
838+
let ahe = ShellAhe::new(config, CONTEXT_STRING)?;
834839
let seed = SingleThreadHkdfPrng::generate_seed()?;
835840
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
836841
let (_, pk_share, _) = ahe.key_gen(&mut prng)?;
@@ -871,9 +876,8 @@ mod test {
871876

872877
#[gtest]
873878
fn test_export_ciphertext_has_right_order() -> googletest::Result<()> {
874-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
875879
let config = make_ahe_config();
876-
let ahe = ShellAhe::new(config, &public_seed)?;
880+
let ahe = ShellAhe::new(config, CONTEXT_STRING)?;
877881
let seed = SingleThreadHkdfPrng::generate_seed()?;
878882
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
879883
let (_, pk_share, _) = ahe.key_gen(&mut prng)?;

‎willow/src/shell/kahe.rs‎

Lines changed: 32 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -50,22 +50,6 @@ pub struct ShellKahe {
5050
}
5151

5252
impl ShellKahe {
53-
pub fn new(
54-
shell_kahe_config: ShellKaheConfig,
55-
public_seed: &Seed,
56-
) -> Result<Self, status::StatusError> {
57-
Self::validate_kahe_config(&shell_kahe_config)?;
58-
let num_coeffs = 1 << shell_kahe_config.log_n;
59-
let public_kahe_parameters = kahe::create_public_parameters(
60-
shell_kahe_config.log_n as u64,
61-
shell_kahe_config.log_t as u64,
62-
&shell_kahe_config.moduli,
63-
shell_kahe_config.num_public_polynomials,
64-
&public_seed,
65-
)?;
66-
Ok(Self { config: shell_kahe_config, num_coeffs, public_kahe_parameters })
67-
}
68-
6953
/// Validates KAHE parameters in ShellKaheConfig.
7054
fn validate_kahe_config(config: &ShellKaheConfig) -> Result<(), status::StatusError> {
7155
if config.log_t > BIG_INT_BITS {
@@ -113,6 +97,30 @@ impl KaheBase for ShellKahe {
11397

11498
type Rng = SingleThreadHkdfPrng;
11599

100+
type Config = ShellKaheConfig;
101+
102+
fn new(
103+
shell_kahe_config: Self::Config,
104+
context_string: &[u8],
105+
) -> Result<Self, status::StatusError> {
106+
Self::validate_kahe_config(&shell_kahe_config)?;
107+
let num_coeffs = 1 << shell_kahe_config.log_n;
108+
let public_seed = single_thread_hkdf::compute_hkdf(
109+
context_string,
110+
b"",
111+
b"ShellKahe.public_seed",
112+
single_thread_hkdf::seed_length(),
113+
)?;
114+
let public_kahe_parameters = kahe::create_public_parameters(
115+
shell_kahe_config.log_n as u64,
116+
shell_kahe_config.log_t as u64,
117+
&shell_kahe_config.moduli,
118+
shell_kahe_config.num_public_polynomials,
119+
&public_seed,
120+
)?;
121+
Ok(Self { config: shell_kahe_config, num_coeffs, public_kahe_parameters })
122+
}
123+
116124
fn add_keys_in_place(
117125
&self,
118126
left: &Self::SecretKey,
@@ -299,6 +307,8 @@ mod test {
299307
/// Default ID used in tests.
300308
const DEFAULT_ID: &str = "default";
301309

310+
const CONTEXT_STRING: &[u8] = b"test_context_string";
311+
302312
#[gtest]
303313
fn test_encrypt_decrypt_short() -> googletest::Result<()> {
304314
let plaintext_modulus_bits = 39;
@@ -307,8 +317,7 @@ mod test {
307317
PackedVectorConfig { base: 10, dimension: 2, num_packed_coeffs: 5 },
308318
)]);
309319
let kahe_config = make_kahe_config_for(plaintext_modulus_bits, packed_vector_configs)?;
310-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
311-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
320+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
312321

313322
let pt = HashMap::from([(String::from(DEFAULT_ID), vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])]);
314323
let seed = SingleThreadHkdfPrng::generate_seed()?;
@@ -327,8 +336,7 @@ mod test {
327336
PackedVectorConfig { base: 10, dimension: 2, num_packed_coeffs: 5 },
328337
)]);
329338
let kahe_config = make_kahe_config_for(plaintext_modulus_bits, packed_vector_configs)?;
330-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
331-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
339+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
332340

333341
let pt = HashMap::from([(String::from(DEFAULT_ID), vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])]);
334342
let seed = SingleThreadHkdfPrng::generate_seed()?;
@@ -364,8 +372,7 @@ mod test {
364372
packed_vector_config.num_packed_coeffs = num_messages;
365373
set_kahe_num_public_polynomials(&mut kahe_config);
366374

367-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
368-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
375+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
369376

370377
let seed = SingleThreadHkdfPrng::generate_seed()?;
371378
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
@@ -397,8 +404,7 @@ mod test {
397404
)]);
398405
let kahe_config = make_kahe_config_for(plaintext_modulus_bits, packed_vector_configs)?;
399406

400-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
401-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
407+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
402408
let seed = SingleThreadHkdfPrng::generate_seed()?;
403409
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
404410

@@ -434,8 +440,7 @@ mod test {
434440
let packed_vector_configs = HashMap::from([]);
435441
let kahe_config = make_kahe_config_for(plaintext_modulus_bits, packed_vector_configs)?;
436442

437-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
438-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
443+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
439444
let seed = SingleThreadHkdfPrng::generate_seed()?;
440445
let mut prng = SingleThreadHkdfPrng::create(&seed)?;
441446

@@ -477,8 +482,7 @@ mod test {
477482
let plaintext_modulus_bits = 39;
478483
let packed_vector_configs = HashMap::from([]);
479484
let kahe_config = make_kahe_config_for(plaintext_modulus_bits, packed_vector_configs)?;
480-
let public_seed = SingleThreadHkdfPrng::generate_seed()?;
481-
let kahe = ShellKahe::new(kahe_config, &public_seed)?;
485+
let kahe = ShellKahe::new(kahe_config, CONTEXT_STRING)?;
482486

483487
// The seed used to sample the secret keys.
484488
let seed = SingleThreadHkdfPrng::generate_seed()?;

‎willow/src/shell/single_thread_hkdf.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ mod test {
4545
use prng_traits::SecurePrng;
4646

4747
#[gtest]
48-
/// Two sequences of 8 random bytes should be different (w.h.p).
48+
/// Two sequences of 8 random bits should be different (w.h.p).
4949
fn test_rand8() -> googletest::Result<()> {
5050
let mut equal = true;
5151
let seed = SingleThreadHkdfPrng::generate_seed()?;

0 commit comments

Comments
 (0)