Add multi-party Uint8
@@ -1,5 +1,6 @@
use std::{ cell::{OnceCell, RefCell}, + clone, collections::HashMap, fmt::{Debug, Display}, iter::Once,@@ -20,8 +21,8 @@ lwe::{decrypt_lwe, encrypt_lwe, lwe_key_switch, lwe_ksk_keygen, measure_noise_lwe, LweSecret},
multi_party::public_key_share, ntt::{self, Ntt, NttBackendU64, NttInit}, random::{ - DefaultSecureRng, NewWithSeed, RandomFillGaussianInModulus, RandomFillUniformInModulus, - RandomGaussianElementInModulus, + DefaultSecureRng, NewWithSeed, RandomFill, RandomFillGaussianInModulus, + RandomFillUniformInModulus, RandomGaussianElementInModulus, }, rgsw::{ decrypt_rlwe, galois_auto, galois_key_gen, generate_auto_map, public_key_encrypt_rgsw,@@ -32,10 +33,11 @@ utils::{
fill_random_ternary_secret_with_hamming_weight, generate_prime, mod_exponent, Global, TryConvertFrom1, WithLocal, }, - Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret, + Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, Row, RowEntity, + RowMut, Secret, }; -use super::parameters::{BoolParameters, CiphertextModulus}; +use super::parameters::{self, BoolParameters, CiphertextModulus}; thread_local! { pub(crate) static BOOL_EVALUATOR: RefCell<BoolEvaluator<Vec<Vec<u64>>, NttBackendU64, ModularOpsU64<CiphertextModulus<u64>>, ModularOpsU64<CiphertextModulus<u64>>>> = RefCell::new(BoolEvaluator::new(MP_BOOL_PARAMS));@@ -45,10 +47,19 @@ pub(crate) static BOOL_SERVER_KEY: OnceLock<
ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>, > = OnceLock::new(); +pub(crate) static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new(); + pub fn set_parameter_set(parameter: &BoolParameters<u64>) { BoolEvaluator::with_local_mut(|e| *e = BoolEvaluator::new(parameter.clone())) } +pub fn set_mp_seed(seed: [u8; 32]) { + assert!( + MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(), + "Attempted to set MP SEED twice." + ) +} + fn set_server_key(key: ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>) { assert!( BOOL_SERVER_KEY.set(key).is_ok(),@@ -56,7 +67,7 @@ "Attempted to set server key twice."
); } -pub fn gen_keys() -> ( +pub(crate) fn gen_keys() -> ( ClientKey, SeededServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>, ) {@@ -67,6 +78,93 @@
(ck, sk) }) } + +pub fn gen_client_key() -> ClientKey { + BoolEvaluator::with_local(|e| e.client_key()) +} + +pub fn gen_mp_keys_phase1( + ck: &ClientKey, +) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> { + let seed = MultiPartyCrs::global().public_key_share_seed::<DefaultSecureRng>(); + BoolEvaluator::with_local(|e| { + let pk_share = e.multi_party_public_key_share(seed, &ck); + pk_share + }) +} + +pub fn gen_mp_keys_phase2<R, ModOp>( + ck: &ClientKey, + pk: &PublicKey<Vec<Vec<u64>>, R, ModOp>, +) -> CommonReferenceSeededMultiPartyServerKeyShare<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> { + let seed = MultiPartyCrs::global().server_key_share_seed::<DefaultSecureRng>(); + BoolEvaluator::with_local_mut(|e| { + let server_key_share = e.multi_party_server_key_share(seed, &pk.key, ck); + server_key_share + }) +} + +pub fn aggregate_public_key_shares( + shares: &[CommonReferenceSeededCollectivePublicKeyShare< + Vec<u64>, + [u8; 32], + BoolParameters<u64>, + >], +) -> PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>> { + PublicKey::from(shares) +} + +pub fn aggregate_server_key_shares( + shares: &[CommonReferenceSeededMultiPartyServerKeyShare< + Vec<Vec<u64>>, + BoolParameters<u64>, + [u8; 32], + >], +) -> SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>> { + BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares)) +} + +// GENERIC BELOW + +pub struct MultiPartyCrs<S> { + seed: S, +} + +impl<S: Default + Copy> MultiPartyCrs<S> { + /// Seed to generate public key share using MultiPartyCrs as the main seed. + /// + /// Public key seed equals the 1st seed extracted from PRNG Seeded with + /// MiltiPartyCrs's seed. + fn public_key_share_seed<Rng: NewWithSeed<Seed = S> + RandomFill<S>>(&self) -> S { + let mut prng = Rng::new_with_seed(self.seed); + + let mut seed = S::default(); + RandomFill::<S>::random_fill(&mut prng, &mut seed); + seed + } + + /// Seed to generate server key share using MultiPartyCrs as the main seed. + /// + /// Server key seed equals the 2nd seed extracted from PRNG Seeded with + /// MiltiPartyCrs's seed. + fn server_key_share_seed<Rng: NewWithSeed<Seed = S> + RandomFill<S>>(&self) -> S { + let mut prng = Rng::new_with_seed(self.seed); + + let mut seed = S::default(); + RandomFill::<S>::random_fill(&mut prng, &mut seed); + RandomFill::<S>::random_fill(&mut prng, &mut seed); + seed + } +} + +impl Global for MultiPartyCrs<[u8; 32]> { + fn global() -> &'static Self { + MULTI_PARTY_CRS + .get() + .expect("Multi Party Common Reference String not set") + } +} + pub(crate) trait BooleanGates { type Ciphertext: RowEntity; type Key;@@ -323,19 +421,116 @@ BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
} } -struct MultiPartyDecryptionShare<E> { - share: E, +impl MultiPartyDecryptor<bool, Vec<u64>> for ClientKey { + type DecryptionShare = u64; + + fn gen_decryption_share(&self, c: &Vec<u64>) -> Self::DecryptionShare { + BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, &self)) + } + + fn aggregate_decryption_shares(&self, c: &Vec<u64>, shares: &[Self::DecryptionShare]) -> bool { + BoolEvaluator::with_local(|e| e.multi_party_decrypt(shares, c)) + } } -struct CommonReferenceSeededCollectivePublicKeyShare<R, S, P> { +// struct MultiPartyDecryptionShare<E> { +// share: E, +// } + +pub struct CommonReferenceSeededCollectivePublicKeyShare<R, S, P> { share: R, cr_seed: S, parameters: P, } -struct PublicKey<M, R, O> { +struct SeededPublicKey<R, S, P, ModOp> { + part_b: R, + seed: S, + parameters: P, + _phantom: PhantomData<ModOp>, +} + +impl<R, S, ModOp> + From<&[CommonReferenceSeededCollectivePublicKeyShare<R, S, BoolParameters<R::Element>>]> + for SeededPublicKey<R, S, BoolParameters<R::Element>, ModOp> +where + ModOp: VectorOps<Element = R::Element> + ModInit<M = CiphertextModulus<R::Element>>, + S: PartialEq + Clone, + R: RowMut + RowEntity + Clone, + R::Element: Clone + PartialEq, +{ + fn from( + value: &[CommonReferenceSeededCollectivePublicKeyShare<R, S, BoolParameters<R::Element>>], + ) -> Self { + assert!(value.len() > 0); + + let parameters = &value[0].parameters; + let cr_seed = value[0].cr_seed.clone(); + + // Sum all Bs + let rlweq_modop = ModOp::new(parameters.rlwe_q().clone()); + let mut part_b = value[0].share.clone(); + value.iter().skip(1).for_each(|share_i| { + assert!(&share_i.cr_seed == &cr_seed); + assert!(&share_i.parameters == parameters); + + rlweq_modop.elwise_add_mut(part_b.as_mut(), share_i.share.as_ref()); + }); + + Self { + part_b, + seed: cr_seed, + parameters: parameters.clone(), + _phantom: PhantomData, + } + } +} + +pub struct PublicKey<M, Rng, ModOp> { key: M, - _phantom: PhantomData<(R, O)>, + _phantom: PhantomData<(Rng, ModOp)>, +} + +impl<Rng, ModOp> Encryptor<bool, Vec<u64>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> { + fn encrypt(&self, m: &bool) -> Vec<u64> { + BoolEvaluator::with_local(|e| e.pk_encrypt(&self.key, *m)) + } +} + +impl<Rng, ModOp> Encryptor<[bool], Vec<Vec<u64>>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> { + fn encrypt(&self, m: &[bool]) -> Vec<Vec<u64>> { + BoolEvaluator::with_local(|e| e.pk_encrypt_batched(&self.key, m)) + } +} + +impl< + M: MatrixMut + MatrixEntity, + Rng: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>, + ModOp, + > From<SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>> + for PublicKey<M, Rng, ModOp> +where + <M as Matrix>::R: RowMut, + M::MatElement: Copy, +{ + fn from(value: SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>) -> Self { + let mut prng = Rng::new_with_seed(value.seed); + + let mut key = M::zeros(2, value.parameters.rlwe_n().0); + // sample A + RandomFillUniformInModulus::random_fill( + &mut prng, + value.parameters.rlwe_q(), + key.get_row_mut(0), + ); + // Copy over B + key.get_row_mut(1).copy_from_slice(value.part_b.as_ref()); + + PublicKey { + key, + _phantom: PhantomData, + } + } } impl<@@ -392,7 +587,7 @@ }
} } -struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> { +pub struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> { rgsw_cts: Vec<M>, /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding /// to -g is at 0@@ -402,7 +597,7 @@ /// Common reference seed
cr_seed: S, parameters: P, } -struct SeededMultiPartyServerKey<M: Matrix, S, P> { +pub struct SeededMultiPartyServerKey<M: Matrix, S, P> { rgsw_cts: Vec<M>, /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding /// to -g is at 0@@ -412,6 +607,22 @@ cr_seed: S,
parameters: P, } +impl + SeededMultiPartyServerKey< + Vec<Vec<u64>>, + <DefaultSecureRng as NewWithSeed>::Seed, + BoolParameters<u64>, + > +{ + pub fn set_server_key(&self) { + set_server_key(ServerKeyEvaluationDomain::< + Vec<Vec<u64>>, + DefaultSecureRng, + NttBackendU64, + >::from(self)) + } +} + /// Seeded single party server key pub struct SeededServerKey<M: Matrix, P, S> { /// Rgsw cts of LWE secret elements@@ -708,7 +919,6 @@ _phanton: PhantomData,
} } } - impl<M: Matrix, R, N> PbsKey for ServerKeyEvaluationDomain<M, R, N> { type M = M;@@ -1241,7 +1451,7 @@ fn multi_party_decryption_share(
&self, lwe_ct: &M::R, client_key: &ClientKey, - ) -> MultiPartyDecryptionShare<<M as Matrix>::MatElement> { + ) -> <M as Matrix>::MatElement { assert!(lwe_ct.as_ref().len() == self.pbs_info.parameters.rlwe_n().0 + 1); let modop = &self.pbs_info.rlwe_modop; let mut neg_s = M::R::try_convert_from(@@ -1262,34 +1472,44 @@ e
}); let share = modop.add(&neg_sa, &e); - MultiPartyDecryptionShare { share } + share } - pub(crate) fn multi_party_decrypt( - &self, - shares: &[MultiPartyDecryptionShare<M::MatElement>], - lwe_ct: &M::R, - ) -> bool { + pub(crate) fn multi_party_decrypt(&self, shares: &[M::MatElement], lwe_ct: &M::R) -> bool { let modop = &self.pbs_info.rlwe_modop; let mut sum_a = M::MatElement::zero(); shares .iter() - .for_each(|share_i| sum_a = modop.add(&sum_a, &share_i.share)); + .for_each(|share_i| sum_a = modop.add(&sum_a, &share_i)); let encoded_m = modop.add(&lwe_ct.as_ref()[0], &sum_a); self.pbs_info.parameters.rlwe_q().decode(encoded_m) } - /// First encrypt as RLWE(m) with m as constant polynomial and extract it as - /// LWE ciphertext pub(crate) fn pk_encrypt(&self, pk: &M, m: bool) -> M::R { + self.pk_encrypt_batched(pk, &vec![m]).remove(0) + } + + /// Encrypts a batch booleans as multiple LWE ciphertexts. + /// + /// For public key encryption we first encrypt `m` as a RLWE ciphertext and + /// then sample extract LWE samples at required indices. + /// + /// - TODO(Jay:) Communication can be improved by not sample exctracting and + /// instead just truncate degree 0 values (part Bs) + pub(crate) fn pk_encrypt_batched(&self, pk: &M, m: &[bool]) -> Vec<M::R> { DefaultSecureRng::with_local_mut(|rng| { + let ring_size = self.pbs_info.parameters.rlwe_n().0; + assert!( + m.len() <= ring_size, + "Cannot batch encrypt > ring_size{ring_size} elements at once" + ); + let modop = &self.pbs_info.rlwe_modop; let nttop = &self.pbs_info.rlwe_nttop; // RLWE(0) // sample ephemeral key u - let ring_size = self.pbs_info.parameters.rlwe_n().0; let mut u = vec![0i32; ring_size]; fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng); let mut u = M::R::try_convert_from(&u, &self.pbs_info.parameters.rlwe_q());@@ -1326,22 +1546,31 @@ // b*u + e1
modop.elwise_add_mut(rlwe.get_row_mut(1), ub.as_ref()); //FIXME(Jay): Figure out a way to get Q/8 form modulus - let m = if m { - // Q/8 - self.pbs_info.rlwe_q().true_el() - } else { - // -Q/8 - self.pbs_info.rlwe_q().false_el() - }; + let mut m_vec = M::R::zeros(ring_size); + izip!(m_vec.as_mut().iter_mut(), m.iter()).for_each(|(m_el, m_bool)| { + if *m_bool { + // Q/8 + *m_el = self.pbs_info.rlwe_q().true_el() + } else { + // -Q/8 + *m_el = self.pbs_info.rlwe_q().false_el() + } + }); - // b*u + e1 + m, where m is constant polynomial - rlwe.set(1, 0, modop.add(rlwe.get(1, 0), &m)); + // b*u + e1 + m + modop.elwise_add_mut(rlwe.get_row_mut(1), m_vec.as_ref()); + // rlwe.set(1, 0, modop.add(rlwe.get(1, 0), &m)); - // sample extract index 0 - let mut lwe_out = M::R::zeros(ring_size + 1); - sample_extract(&mut lwe_out, &rlwe, modop, 0); - - lwe_out + // sample extract index required indices + let samples = m.len(); + (0..samples) + .into_iter() + .map(|i| { + let mut lwe_out = M::R::zeros(ring_size + 1); + sample_extract(&mut lwe_out, &rlwe, modop, i); + lwe_out + }) + .collect_vec() }) }@@ -2103,7 +2332,6 @@ pbs_info,
pbs_key, ); - // sample extract sample_extract(lwe_in, &trivial_rlwe_test_poly, pbs_info.modop_rlweq(), 0); }@@ -2731,7 +2959,7 @@ ModularOpsU64<CiphertextModulus<u64>>,
>::new(MP_BOOL_PARAMS); let (parties, collective_pk, _, _, server_key_eval, ideal_client_key) = - _multi_party_all_keygen(&bool_evaluator, 8); + _multi_party_all_keygen(&bool_evaluator, 64); let mut m0 = true; let mut m1 = false;
@@ -307,10 +307,10 @@ rgrg_decomposer_count: (DecompositionCount(4), DecompositionCount(4)),
auto_decomposer_base: DecompostionLogBase(7), auto_decomposer_count: DecompositionCount(4), g: 5, - w: 10, + w: 5, }; -pub(super) const MP_BOOL_PARAMS: BoolParameters<u64> = BoolParameters::<u64> { +pub(crate) const MP_BOOL_PARAMS: BoolParameters<u64> = BoolParameters::<u64> { rlwe_q: CiphertextModulus::new_non_native(1152921504606830593), lwe_q: CiphertextModulus::new_non_native(1 << 20), br_q: 1 << 11,@@ -325,7 +325,7 @@ rgrg_decomposer_count: (DecompositionCount(5), DecompositionCount(4)),
auto_decomposer_base: DecompostionLogBase(12), auto_decomposer_count: DecompositionCount(5), g: 5, - w: 5, + w: 10, }; #[cfg(test)]
@@ -1,3 +1,5 @@
+use std::{iter::Once, sync::OnceLock}; + use itertools::{izip, Itertools}; use num::UnsignedInteger; use num_traits::{abs, Zero};@@ -156,10 +158,17 @@ vec![T::zero(); col]
} } -trait Encryptor<M, C> { +trait Encryptor<M: ?Sized, C> { fn encrypt(&self, m: &M) -> C; } -trait Decryptor<M, C: ?Sized> { +trait Decryptor<M, C> { fn decrypt(&self, c: &C) -> M; } + +trait MultiPartyDecryptor<M, C> { + type DecryptionShare; + + fn gen_decryption_share(&self, c: &C) -> Self::DecryptionShare; + fn aggregate_decryption_shares(&self, c: &C, shares: &[Self::DecryptionShare]) -> M; +}
@@ -138,6 +138,21 @@ });
} } +impl<T> RandomFill<[T; 32]> for DefaultSecureRng +where + T: PrimInt + SampleUniform, +{ + fn random_fill(&mut self, container: &mut [T; 32]) { + izip!( + (&mut self.rng).sample_iter(Uniform::new_inclusive(T::zero(), T::max_value())), + container.iter_mut() + ) + .for_each(|(from, to)| { + *to = from; + }); + } +} + impl<T> RandomElement<T> for DefaultSecureRng where T: PrimInt + SampleUniform,
@@ -1,9 +1,11 @@
use itertools::Itertools; use crate::{ - bool::evaluator::{BoolEvaluator, ClientKey, ServerKeyEvaluationDomain, BOOL_SERVER_KEY}, + bool::evaluator::{ + BoolEvaluator, ClientKey, PublicKey, ServerKeyEvaluationDomain, BOOL_SERVER_KEY, + }, utils::{Global, WithLocal}, - Decryptor, Encryptor, + Decryptor, Encryptor, Matrix, MultiPartyDecryptor, }; mod ops;@@ -26,6 +28,7 @@ }
impl Decryptor<u8, FheUint8> for ClientKey { fn decrypt(&self, c: &FheUint8) -> u8 { + assert!(c.data.len() == 8); let mut out = 0u8; c.data().iter().enumerate().for_each(|(index, bit_c)| { let bool = Decryptor::<bool, Vec<u64>>::decrypt(self, bit_c);@@ -33,6 +36,60 @@ if bool {
out += 1 << index; } }); + out + } +} + +impl<M, R, Mo> Encryptor<u8, FheUint8> for PublicKey<M, R, Mo> +where + PublicKey<M, R, Mo>: Encryptor<bool, Vec<u64>>, +{ + fn encrypt(&self, m: &u8) -> FheUint8 { + let cts = (0..8) + .into_iter() + .map(|i| { + let bit = ((m >> i) & 1) == 1; + Encryptor::<bool, Vec<u64>>::encrypt(self, &bit) + }) + .collect_vec(); + FheUint8 { data: cts } + } +} + +impl MultiPartyDecryptor<u8, FheUint8> for ClientKey +where + ClientKey: MultiPartyDecryptor<bool, Vec<u64>>, +{ + type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, Vec<u64>>>::DecryptionShare>; + fn gen_decryption_share(&self, c: &FheUint8) -> Self::DecryptionShare { + assert!(c.data().len() == 8); + c.data() + .iter() + .map(|bit_c| { + let decryption_share = + MultiPartyDecryptor::<bool, Vec<u64>>::gen_decryption_share(self, bit_c); + decryption_share + }) + .collect_vec() + } + + fn aggregate_decryption_shares(&self, c: &FheUint8, shares: &[Self::DecryptionShare]) -> u8 { + let mut out = 0u8; + + (0..8).into_iter().for_each(|i| { + // Collect bit i^th decryption share of each party + let bit_i_decryption_shares = shares.iter().map(|s| s[i]).collect_vec(); + let bit_i = MultiPartyDecryptor::<bool, Vec<u64>>::aggregate_decryption_shares( + self, + &c.data()[i], + &bit_i_decryption_shares, + ); + + if bit_i { + out += 1 << i; + } + }); + out } }@@ -245,15 +302,20 @@ }
#[cfg(test)] mod tests { + use itertools::Itertools; use num_traits::Euclid; use crate::{ bool::{ - evaluator::{gen_keys, set_parameter_set, BoolEvaluator}, - parameters::SP_BOOL_PARAMS, + evaluator::{ + aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key, gen_keys, + gen_mp_keys_phase1, gen_mp_keys_phase2, set_mp_seed, set_parameter_set, + BoolEvaluator, ClientKey, + }, + parameters::{MP_BOOL_PARAMS, SP_BOOL_PARAMS}, }, shortint::types::FheUint8, - Decryptor, Encryptor, + Decryptor, Encryptor, MultiPartyDecryptor, }; #[test]@@ -402,5 +464,58 @@ }
} } } + } + + #[test] + fn fheuint8_test_multi_party() { + set_parameter_set(&MP_BOOL_PARAMS); + set_mp_seed([0; 32]); + + let parties = 8; + + // client keys and public key share + let cks = (0..parties) + .into_iter() + .map(|i| gen_client_key()) + .collect_vec(); + + // round 1: generate pulic key shares + let pk_shares = cks.iter().map(|key| gen_mp_keys_phase1(key)).collect_vec(); + + let public_key = aggregate_public_key_shares(&pk_shares); + + // round 2: generate server key shares + let server_key_shares = cks + .iter() + .map(|key| gen_mp_keys_phase2(key, &public_key)) + .collect_vec(); + + // server aggregates the server key + let server_key = aggregate_server_key_shares(&server_key_shares); + server_key.set_server_key(); + + // Clients use Pk to encrypt private inputs + let a = 8u8; + let b = 10u8; + let c = 155u8; + let ct_a = public_key.encrypt(&a); + let ct_b = public_key.encrypt(&b); + let ct_c = public_key.encrypt(&c); + + // server computes + // a*b + c + let mut ct_ab = &ct_a * &ct_b; + ct_ab += &ct_c; + + // decrypt ab and check + // generate decryption shares + let decryption_shares = cks + .iter() + .map(|k| k.gen_decryption_share(&ct_ab)) + .collect_vec(); + + // aggregate and decryption ab + let ab_add_c = cks[0].aggregate_decryption_shares(&ct_ab, &decryption_shares); + assert!(ab_add_c == (a.wrapping_mul(b)).wrapping_add(c)); } }