github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-09 14:47:56 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add re randomization key to XofKeySet

Browse Source
This commit is contained in:
Thomas Montaigu
2025-09-17 12:14:16 +02:00
committed by tmontaigu
parent d1c190fac6
commit 6e7aaac90f
5 changed files with 172 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tfhe/src/high_level_api/config.rs
View File

@@ -91,10 +91,12 @@ impl ConfigBuilder {
/// Enable the re-randomization of ciphertexts after compression using a
/// [crate::CompactPublicKey].
///
/// # Note
/// # Panics
///
/// This requires dedicated [crate::CompactPublicKey] parameters to be enabled via
/// [Self::use_dedicated_compact_public_key_parameters].
///
/// The given parameters must target the `EncryptionKeyChoice::Big`
pub fn enable_ciphertext_re_randomization(
mut self,
cpk_re_randomization_ksk_params: ShortintKeySwitchingParameters,

12
tfhe/src/high_level_api/keys/cpk_re_randomization.rs
View File

@@ -1,8 +1,20 @@
use crate::high_level_api::backward_compatibility::keys::{
CompressedReRandomizationKeySwitchingKeyVersions, ReRandomizationKeySwitchingKeyVersions,
};
use crate::shortint::parameters::ShortintKeySwitchingParameters;
use tfhe_versionable::Versionize;
#[derive(Debug, Clone)]
pub(crate) enum ReRandomizationKeyGenerationInfo<'a> {
UseCPKEncryptionKSK,
DedicatedKSK(
(
&'a crate::integer::CompactPrivateKey<Vec<u64>>,
ShortintKeySwitchingParameters,
),
),
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Versionize)]
#[versionize(ReRandomizationKeySwitchingKeyVersions)]
pub enum ReRandomizationKeySwitchingKey {

110
tfhe/src/high_level_api/keys/inner.rs
View File

@@ -2,9 +2,11 @@ use crate::conformance::ParameterSetConformant;
use crate::core_crypto::commons::generators::DeterministicSeeder;
use crate::core_crypto::prelude::{DefaultRandomGenerator, LweKeyswitchKeyConformanceParams};
use crate::high_level_api::backward_compatibility::keys::*;
use crate::high_level_api::errors::UnwrapResultExt;
use crate::high_level_api::keys::cpk_re_randomization::{
CompressedReRandomizationKeySwitchingKey, ReRandomizationKeySwitchingKey,
};
use crate::high_level_api::keys::ReRandomizationKeyGenerationInfo;
use crate::integer::ciphertext::{
CompressedNoiseSquashingCompressionKey, NoiseSquashingCompressionKey,
NoiseSquashingCompressionPrivateKey,
@@ -267,6 +269,36 @@ impl IntegerClientKey {
pub(crate) fn block_parameters(&self) -> crate::shortint::parameters::AtomicPatternParameters {
self.key.parameters()
}
pub(crate) fn re_randomization_ksk_gen_info(
&self,
) -> Result<Option<ReRandomizationKeyGenerationInfo<'_>>, crate::Error> {
let maybe_cpk = self.dedicated_compact_private_key.as_ref();
let maybe_re_rand_ksk_params = self.cpk_re_randomization_ksk_params.as_ref();
match (maybe_cpk, maybe_re_rand_ksk_params) {
(Some((cpk, cpke_ksk_params)), Some(re_rand_ksk_params)) => {
if re_rand_ksk_params.destination_key != EncryptionKeyChoice::Big {
return Err(crate::error!(
"CompactPublicKey re-randomization can only be enabled \
targeting the large secret key."
));
}
if cpke_ksk_params == re_rand_ksk_params {
Ok(Some(ReRandomizationKeyGenerationInfo::UseCPKEncryptionKSK))
} else {
Ok(Some(ReRandomizationKeyGenerationInfo::DedicatedKSK((
cpk,
*re_rand_ksk_params,
))))
}
}
(_, None) => Ok(None),
_ => Err(crate::error!(
"Inconsistent ClientKey set-up for CompactPublicKey re-randomization."
)),
}
}
}
impl From<IntegerConfig> for IntegerClientKey {
@@ -372,38 +404,27 @@ impl IntegerServerKey {
},
);
let cpk_re_randomization_key_switching_key_material = match (
&client_key.dedicated_compact_private_key,
client_key.cpk_re_randomization_ksk_params,
) {
(Some(compact_private_key), Some(cpk_re_randomization_ksk_params)) => {
assert!(
matches!(
cpk_re_randomization_ksk_params.destination_key,
EncryptionKeyChoice::Big
),
"CompactPublicKey re-randomization can only be enabled \
targeting the large secret key."
);
if cpk_re_randomization_ksk_params == compact_private_key.1 {
Some(ReRandomizationKeySwitchingKey::UseCPKEncryptionKSK)
} else {
let cpk_re_randomization_key_switching_key_material = client_key
.re_randomization_ksk_gen_info()
.unwrap_display()
.map(|key_gen_info| match key_gen_info {
ReRandomizationKeyGenerationInfo::UseCPKEncryptionKSK => {
ReRandomizationKeySwitchingKey::UseCPKEncryptionKSK
}
ReRandomizationKeyGenerationInfo::DedicatedKSK((
input_cpk,
cpk_re_randomization_ksk_params,
)) => {
let build_helper =
crate::integer::key_switching_key::KeySwitchingKeyBuildHelper::new(
(&compact_private_key.0, None),
(input_cpk, None),
(cks, &base_integer_key),
cpk_re_randomization_ksk_params,
);
Some(ReRandomizationKeySwitchingKey::DedicatedKSK(
build_helper.into(),
))
ReRandomizationKeySwitchingKey::DedicatedKSK(build_helper.into())
}
}
(_, None) => None,
_ => panic!("Inconsistent ClientKey set-up for CompactPublicKey re-randomization."),
};
});
Self {
key: base_integer_key,
@@ -534,38 +555,27 @@ impl IntegerCompressedServerKey {
(Some(noise_squashing_key), noise_squashing_compression_key)
});
let cpk_re_randomization_key_switching_key_material = match (
&client_key.dedicated_compact_private_key,
client_key.cpk_re_randomization_ksk_params,
) {
(Some(compact_private_key), Some(cpk_re_randomization_ksk_params)) => {
assert!(
matches!(
cpk_re_randomization_ksk_params.destination_key,
EncryptionKeyChoice::Big
),
"CompactPublicKey re-randomization can only be enabled \
targeting the large secret key."
);
if cpk_re_randomization_ksk_params == compact_private_key.1 {
Some(CompressedReRandomizationKeySwitchingKey::UseCPKEncryptionKSK)
} else {
let cpk_re_randomization_key_switching_key_material = client_key
.re_randomization_ksk_gen_info()
.unwrap_display()
.map(|key_gen_info| match key_gen_info {
ReRandomizationKeyGenerationInfo::UseCPKEncryptionKSK => {
CompressedReRandomizationKeySwitchingKey::UseCPKEncryptionKSK
}
ReRandomizationKeyGenerationInfo::DedicatedKSK((
input_cpk,
cpk_re_randomization_ksk_params,
)) => {
let build_helper =
crate::integer::key_switching_key::CompressedKeySwitchingKeyBuildHelper::new(
(&compact_private_key.0, None),
(input_cpk, None),
(cks, &key),
cpk_re_randomization_ksk_params,
);
Some(CompressedReRandomizationKeySwitchingKey::DedicatedKSK(
build_helper.into(),
))
CompressedReRandomizationKeySwitchingKey::DedicatedKSK(build_helper.into())
}
}
(_, None) => None,
_ => panic!("Inconsistent ClientKey set-up for CompactPublicKey re-randomization."),
};
});
Self {
key,

1
tfhe/src/high_level_api/keys/mod.rs
View File

@@ -8,6 +8,7 @@ mod key_switching_key;
use crate::high_level_api::config::Config;
pub use client::ClientKey;
pub(crate) use cpk_re_randomization::ReRandomizationKeyGenerationInfo;
pub use cpk_re_randomization::{
CompressedReRandomizationKeySwitchingKey, ReRandomizationKeySwitchingKey,
};

102
tfhe/src/high_level_api/xof_key_set.rs
View File

@@ -36,7 +36,7 @@ use crate::shortint::{PBSParameters, ShortintParameterSet};
use crate::ClientKey;
use crate::{
integer, shortint, CompactPublicKey, CompressedCompactPublicKey, CompressedServerKey,
ServerKey, Tag,
ReRandomizationKeySwitchingKey, ServerKey, Tag,
};
use aligned_vec::ABox;
use serde::{Deserialize, Serialize};
@@ -78,6 +78,8 @@ pub struct CompressedXofKeySet {
impl CompressedXofKeySet {
#[cfg(test)]
fn with_seed(pub_seed: XofSeed, priv_seed: XofSeed, ck: &ClientKey) -> crate::Result<Self> {
use crate::high_level_api::keys::ReRandomizationKeyGenerationInfo;
let Some(dedicated_pk_key) = ck.key.dedicated_compact_private_key.as_ref() else {
return Err(crate::error!("Dedicated compact private key is required"));
};
@@ -228,7 +230,7 @@ impl CompressedXofKeySet {
)
});
// Lastly, generate the key switching material that will allow going from
// Generate the key switching material that will allow going from
// the public key's dedicated parameters to the computation parameters
let pk_to_sk_ksk_params = dedicated_pk_key.1;
let (target_private_key, noise_distrib) = match pk_to_sk_ksk_params.destination_key {
@@ -263,6 +265,52 @@ impl CompressedXofKeySet {
}
};
// Generate the key switching material that will allow going from
// the public key's dedicated parameters to the re-rand
let cpk_re_randomization_key_switching_key_material = ck
.key
.re_randomization_ksk_gen_info()?
.map(|key_gen_info| match key_gen_info {
ReRandomizationKeyGenerationInfo::UseCPKEncryptionKSK => {
use crate::CompressedReRandomizationKeySwitchingKey;
CompressedReRandomizationKeySwitchingKey::UseCPKEncryptionKSK
}
ReRandomizationKeyGenerationInfo::DedicatedKSK((
input_cpk,
cpk_re_randomization_ksk_params,
)) => {
use crate::CompressedReRandomizationKeySwitchingKey;
let (target_private_key, noise_distrib) = (
glwe_secret_key.as_lwe_secret_key(),
computation_parameters.glwe_noise_distribution(),
);
let key_switching_key =
allocate_and_generate_lwe_key_switching_key_with_pre_seeded_generator(
&input_cpk.key.key(),
&target_private_key,
cpk_re_randomization_ksk_params.ks_base_log,
cpk_re_randomization_ksk_params.ks_level,
noise_distrib,
computation_parameters.ciphertext_modulus(),
&mut encryption_rand_gen,
);
let key = CompressedKeySwitchingKeyMaterial {
material: shortint::key_switching_key::CompressedKeySwitchingKeyMaterial {
key_switching_key,
cast_rshift: 0,
destination_key: EncryptionKeyChoice::Big,
},
};
CompressedReRandomizationKeySwitchingKey::DedicatedKSK(key)
}
});
if let PBSParameters::PBS(pbs_params) = computation_parameters.pbs_parameters().unwrap() {
let mod_switch_noise_key =
CompressedModulusSwitchConfiguration::generate_from_existing_generator(
@@ -295,6 +343,7 @@ impl CompressedXofKeySet {
}
}
// TODO
let noise_squashing_compression_key = None;
let compressed_server_key = CompressedServerKey::from_raw_parts(
integer_compressed_server_key,
@@ -303,7 +352,7 @@ impl CompressedXofKeySet {
decompression_key,
noise_squashing_bs_key,
noise_squashing_compression_key,
None,
cpk_re_randomization_key_switching_key_material,
Tag::default(),
);
@@ -438,6 +487,21 @@ impl CompressedXofKeySet {
.cpk_key_switching_key_material
.map(|k| k.decompress_with_pre_seeded_generator(&mut mask_generator));
let integer_cpk_re_rand_ksk = self
.compressed_server_key
.integer_key
.cpk_re_randomization_key_switching_key_material
.map(|k| match k {
super::CompressedReRandomizationKeySwitchingKey::UseCPKEncryptionKSK => {
ReRandomizationKeySwitchingKey::UseCPKEncryptionKSK
}
super::CompressedReRandomizationKeySwitchingKey::DedicatedKSK(key) => {
ReRandomizationKeySwitchingKey::DedicatedKSK(
key.decompress_with_pre_seeded_generator(&mut mask_generator),
)
}
});
match &mut integer_sk.key.atomic_pattern {
AtomicPatternServerKey::Standard(ap) => {
match &mut ap.bootstrapping_key {
@@ -488,7 +552,7 @@ impl CompressedXofKeySet {
decompression_key,
noise_squashing_key,
noise_squashing_compression_key,
None,
integer_cpk_re_rand_ksk,
Tag::default(),
)
};
@@ -1114,11 +1178,13 @@ mod test {
let noise_squashing_params = shortint::parameters::test_params::TEST_PARAM_NOISE_SQUASHING_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
let compression_params =
shortint::parameters::test_params::TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
let re_rand_ksk_params = shortint::parameters::test_params::TEST_PARAM_KEYSWITCH_PKE_TO_BIG_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
let config = ConfigBuilder::with_custom_parameters(params)
.use_dedicated_compact_public_key_parameters((cpk_params, casting_params))
.enable_noise_squashing(noise_squashing_params)
.enable_compression(compression_params)
.enable_ciphertext_re_randomization(re_rand_ksk_params)
.build();
let cks = ClientKey::generate(config);
@@ -1149,8 +1215,32 @@ mod test {
.push(clear_b)
.build();
let expander = list.expand().unwrap();
let a = expander.get::<FheUint32>(0).unwrap().unwrap();
let b = expander.get::<FheUint32>(1).unwrap().unwrap();
let mut a = expander.get::<FheUint32>(0).unwrap().unwrap();
let mut b = expander.get::<FheUint32>(1).unwrap().unwrap();
// Test re-randomization
{
// Simulate a 256 bits nonce
let nonce: [u8; 256 / 8] = core::array::from_fn(|_| rand::random());
let compact_public_encryption_domain_separator = *b"TFHE.Enc";
let rerand_domain_separator = *b"TFHE.Rrd";
let mut re_rand_context = ReRandomizationContext::new(
rerand_domain_separator,
// First is the function description, second is a nonce
[b"FheUint32 bin ops".as_slice(), nonce.as_slice()],
compact_public_encryption_domain_separator,
);
re_rand_context.add_ciphertext(&a);
re_rand_context.add_ciphertext(&b);
let mut seed_gen = re_rand_context.finalize();
a.re_randomize(&pk, seed_gen.next_seed().unwrap()).unwrap();
b.re_randomize(&pk, seed_gen.next_seed().unwrap()).unwrap();
}
let c = &a * &b;
let d = &a & &b;