arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/bool/ni_mp_api.rs
src/bool/ni_mp_api.rsRaw
  1use std::{cell::RefCell, sync::OnceLock};
  2
  3use crate::{
  4    backend::ModulusPowerOf2,
  5    bool::parameters::ParameterVariant,
  6    parameters::NI_4P,
  7    random::DefaultSecureRng,
  8    utils::{Global, WithLocal},
  9    ModularOpsU64, NttBackendU64,
 10};
 11
 12use super::{
 13    evaluator::NonInteractiveMultiPartyCrs,
 14    keys::{
 15        CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare,
 16        NonInteractiveServerKeyEvaluationDomain, SeededNonInteractiveMultiPartyServerKey,
 17        ShoupNonInteractiveServerKeyEvaluationDomain,
 18    },
 19    parameters::{BoolParameters, CiphertextModulus, NI_2P},
 20    ClientKey,
 21};
 22
 23pub(crate) type BoolEvaluator = super::evaluator::BoolEvaluator<
 24    Vec<Vec<u64>>,
 25    NttBackendU64,
 26    ModularOpsU64<CiphertextModulus<u64>>,
 27    ModulusPowerOf2<CiphertextModulus<u64>>,
 28    ShoupNonInteractiveServerKeyEvaluationDomain<Vec<Vec<u64>>>,
 29>;
 30
 31thread_local! {
 32    static BOOL_EVALUATOR: RefCell<Option<BoolEvaluator>> = RefCell::new(None);
 33
 34}
 35static BOOL_SERVER_KEY: OnceLock<ShoupNonInteractiveServerKeyEvaluationDomain<Vec<Vec<u64>>>> =
 36    OnceLock::new();
 37
 38static MULTI_PARTY_CRS: OnceLock<NonInteractiveMultiPartyCrs<[u8; 32]>> = OnceLock::new();
 39
 40pub enum ParameterSelector {
 41    NonInteractiveLTE2Party,
 42    NonInteractiveLTE4Party,
 43}
 44
 45pub fn set_parameter_set(select: ParameterSelector) {
 46    match select {
 47        ParameterSelector::NonInteractiveLTE2Party => {
 48            BOOL_EVALUATOR.with_borrow_mut(|v| *v = Some(BoolEvaluator::new(NI_2P)));
 49        }
 50        ParameterSelector::NonInteractiveLTE4Party => {
 51            BOOL_EVALUATOR.with_borrow_mut(|v| *v = Some(BoolEvaluator::new(NI_4P)));
 52        }
 53
 54        _ => {
 55            panic!("Paramerters not supported")
 56        }
 57    }
 58}
 59
 60pub fn set_common_reference_seed(seed: [u8; 32]) {
 61    BoolEvaluator::with_local(|e| {
 62        assert_eq!(
 63            e.parameters().variant(),
 64            &ParameterVariant::NonInteractiveMultiParty,
 65            "Set parameters do not support Non interactive multi-party"
 66        );
 67    });
 68
 69    assert!(
 70        MULTI_PARTY_CRS
 71            .set(NonInteractiveMultiPartyCrs { seed: seed })
 72            .is_ok(),
 73        "Attempted to set MP SEED twice."
 74    )
 75}
 76
 77pub fn gen_client_key() -> ClientKey {
 78    BoolEvaluator::with_local(|e| e.client_key())
 79}
 80
 81pub fn gen_server_key_share(
 82    user_id: usize,
 83    total_users: usize,
 84    client_key: &ClientKey,
 85) -> CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare<
 86    Vec<Vec<u64>>,
 87    BoolParameters<u64>,
 88    NonInteractiveMultiPartyCrs<[u8; 32]>,
 89> {
 90    BoolEvaluator::with_local(|e| {
 91        let cr_seed = NonInteractiveMultiPartyCrs::global();
 92        e.gen_non_interactive_multi_party_key_share(cr_seed, user_id, total_users, client_key)
 93    })
 94}
 95
 96pub fn aggregate_server_key_shares(
 97    shares: &[CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare<
 98        Vec<Vec<u64>>,
 99        BoolParameters<u64>,
100        NonInteractiveMultiPartyCrs<[u8; 32]>,
101    >],
102) -> SeededNonInteractiveMultiPartyServerKey<
103    Vec<Vec<u64>>,
104    NonInteractiveMultiPartyCrs<[u8; 32]>,
105    BoolParameters<u64>,
106> {
107    BoolEvaluator::with_local(|e| {
108        let cr_seed = NonInteractiveMultiPartyCrs::global();
109        e.aggregate_non_interactive_multi_party_server_key_shares(cr_seed, shares)
110    })
111}
112
113impl
114    SeededNonInteractiveMultiPartyServerKey<
115        Vec<Vec<u64>>,
116        NonInteractiveMultiPartyCrs<[u8; 32]>,
117        BoolParameters<u64>,
118    >
119{
120    pub fn set_server_key(&self) {
121        let eval_key = NonInteractiveServerKeyEvaluationDomain::<
122            _,
123            BoolParameters<u64>,
124            DefaultSecureRng,
125            NttBackendU64,
126        >::from(self);
127        assert!(
128            BOOL_SERVER_KEY
129                .set(ShoupNonInteractiveServerKeyEvaluationDomain::from(eval_key))
130                .is_ok(),
131            "Attempted to set server key twice!"
132        );
133    }
134}
135
136impl Global for NonInteractiveMultiPartyCrs<[u8; 32]> {
137    fn global() -> &'static Self {
138        MULTI_PARTY_CRS
139            .get()
140            .expect("Non-interactive multi-party common reference string not set")
141    }
142}
143
144// BOOL EVALUATOR //
145impl WithLocal for BoolEvaluator {
146    fn with_local<F, R>(func: F) -> R
147    where
148        F: Fn(&Self) -> R,
149    {
150        BOOL_EVALUATOR.with_borrow(|s| func(s.as_ref().expect("Parameters not set")))
151    }
152
153    fn with_local_mut<F, R>(func: F) -> R
154    where
155        F: Fn(&mut Self) -> R,
156    {
157        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
158    }
159
160    fn with_local_mut_mut<F, R>(func: &mut F) -> R
161    where
162        F: FnMut(&mut Self) -> R,
163    {
164        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
165    }
166}
167
168pub(crate) type RuntimeServerKey = ShoupNonInteractiveServerKeyEvaluationDomain<Vec<Vec<u64>>>;
169impl Global for RuntimeServerKey {
170    fn global() -> &'static Self {
171        BOOL_SERVER_KEY.get().expect("Server key not set!")
172    }
173}
174
175/// Batch of bool ciphertexts stored as vector of RLWE ciphertext under user j's
176/// RLWE secret `u_j`
177///
178/// To use the bool ciphertexts in multi-party protocol first key switch the
179/// ciphertexts from u_j to ideal RLWE secret `s` with
180/// `self.key_switch(user_id)` where `user_id` is user j's id. Key switch
181/// returns `BatchedFheBools` that stored key vector of key switched RLWE
182/// ciphertext.
183pub(super) struct NonInteractiveBatchedFheBools<C> {
184    data: Vec<C>,
185}
186
187/// Batch of Bool cipphertexts stored as vector of RLWE ciphertexts under the
188/// ideal RLWE secret key `s` of the protocol
189///
190/// Bool ciphertext at `index` can be extracted from the coefficient at `index %
191/// N` of `index / N`th RLWE ciphertext.
192///
193/// To extract bool ciphertext at `index` as LWE ciphertext use
194/// `self.extract(index)`
195pub(super) struct BatchedFheBools<C> {
196    pub(in super::super) data: Vec<C>,
197}
198
199/// Non interactive multi-party specfic encryptor decryptor routines
200mod impl_enc_dec {
201    use crate::{
202        bool::{evaluator::BoolEncoding, keys::NonInteractiveMultiPartyClientKey},
203        pbs::{sample_extract, PbsInfo, WithShoupRepr},
204        random::{NewWithSeed, RandomFillUniformInModulus},
205        rgsw::{rlwe_key_switch, seeded_secret_key_encrypt_rlwe},
206        utils::TryConvertFrom1,
207        Encryptor, KeySwitchWithId, Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut,
208    };
209    use itertools::Itertools;
210    use num_traits::{ToPrimitive, Zero};
211
212    use super::*;
213
214    type Mat = Vec<Vec<u64>>;
215
216    // Implement `extract` to extract Bool LWE ciphertext at `index` from
217    // `BatchedFheBools`
218    impl<C: MatrixMut<MatElement = u64>> BatchedFheBools<C>
219    where
220        C::R: RowEntity + RowMut,
221    {
222        pub(super) fn extract(&self, index: usize) -> C::R {
223            BoolEvaluator::with_local(|e| {
224                let ring_size = e.parameters().rlwe_n().0;
225                let ct_index = index / ring_size;
226                let coeff_index = index % ring_size;
227                let mut lwe_out = C::R::zeros(e.parameters().rlwe_n().0 + 1);
228                sample_extract(
229                    &mut lwe_out,
230                    &self.data[ct_index],
231                    e.pbs_info().modop_rlweq(),
232                    coeff_index,
233                );
234                lwe_out
235            })
236        }
237    }
238
239    impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&(Vec<M::R>, [u8; 32])>
240        for NonInteractiveBatchedFheBools<M>
241    where
242        <M as Matrix>::R: RowMut,
243    {
244        /// Derive `NonInteractiveBatchedFheBools` from a vector seeded RLWE
245        /// ciphertexts (Vec<RLWE>, Seed)
246        ///
247        /// Unseed the RLWE ciphertexts and store them as vector RLWE
248        /// ciphertexts in `NonInteractiveBatchedFheBools`
249        fn from(value: &(Vec<M::R>, [u8; 32])) -> Self {
250            BoolEvaluator::with_local(|e| {
251                let parameters = e.parameters();
252                let ring_size = parameters.rlwe_n().0;
253                let rlwe_q = parameters.rlwe_q();
254
255                let mut prng = DefaultSecureRng::new_seeded(value.1);
256                let rlwes = value
257                    .0
258                    .iter()
259                    .map(|partb| {
260                        let mut rlwe = M::zeros(2, ring_size);
261
262                        // sample A
263                        RandomFillUniformInModulus::random_fill(
264                            &mut prng,
265                            rlwe_q,
266                            rlwe.get_row_mut(0),
267                        );
268
269                        // Copy over B
270                        rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
271
272                        rlwe
273                    })
274                    .collect_vec();
275                Self { data: rlwes }
276            })
277        }
278    }
279
280    impl<K> Encryptor<[bool], NonInteractiveBatchedFheBools<Mat>> for K
281    where
282        K: Encryptor<[bool], (Mat, [u8; 32])>,
283    {
284        /// Encrypt a vector bool of arbitrary length as vector of unseeded RLWE
285        /// ciphertexts in `NonInteractiveBatchedFheBools`
286        fn encrypt(&self, m: &[bool]) -> NonInteractiveBatchedFheBools<Mat> {
287            NonInteractiveBatchedFheBools::from(&K::encrypt(&self, m))
288        }
289    }
290
291    impl<K> Encryptor<[bool], (Vec<<Mat as Matrix>::R>, [u8; 32])> for K
292    where
293        K: NonInteractiveMultiPartyClientKey,
294        <Mat as Matrix>::R:
295            TryConvertFrom1<[K::Element], CiphertextModulus<<Mat as Matrix>::MatElement>>,
296    {
297        /// Encrypt a vector of bool of arbitrary length as vector of seeded
298        /// RLWE ciphertexts and returns (Vec<RLWE>, Seed)
299        fn encrypt(&self, m: &[bool]) -> (Mat, [u8; 32]) {
300            BoolEvaluator::with_local(|e| {
301                DefaultSecureRng::with_local_mut(|rng| {
302                    let parameters = e.parameters();
303                    let ring_size = parameters.rlwe_n().0;
304
305                    let rlwe_count = ((m.len() as f64 / ring_size as f64).ceil())
306                        .to_usize()
307                        .unwrap();
308
309                    let mut seed = <DefaultSecureRng as NewWithSeed>::Seed::default();
310                    rng.fill_bytes(&mut seed);
311                    let mut prng = DefaultSecureRng::new_seeded(seed);
312
313                    let sk_u = self.sk_u_rlwe();
314
315                    // encrypt `m` into ceil(len(m)/N) RLWE ciphertexts
316                    let rlwes = (0..rlwe_count)
317                        .map(|index| {
318                            let mut message = vec![<Mat as Matrix>::MatElement::zero(); ring_size];
319                            m[(index * ring_size)..std::cmp::min(m.len(), (index + 1) * ring_size)]
320                                .iter()
321                                .enumerate()
322                                .for_each(|(i, v)| {
323                                    if *v {
324                                        message[i] = parameters.rlwe_q().true_el()
325                                    } else {
326                                        message[i] = parameters.rlwe_q().false_el()
327                                    }
328                                });
329
330                            // encrypt message
331                            let mut rlwe_out =
332                                <<Mat as Matrix>::R as RowEntity>::zeros(parameters.rlwe_n().0);
333
334                            seeded_secret_key_encrypt_rlwe(
335                                &message,
336                                &mut rlwe_out,
337                                &sk_u,
338                                e.pbs_info().modop_rlweq(),
339                                e.pbs_info().nttop_rlweq(),
340                                &mut prng,
341                                rng,
342                            );
343
344                            rlwe_out
345                        })
346                        .collect_vec();
347
348                    (rlwes, seed)
349                })
350            })
351        }
352    }
353
354    impl KeySwitchWithId<Mat> for Mat {
355        /// Key switch RLWE ciphertext `Self` from user j's RLWE secret u_j
356        /// to ideal RLWE secret `s` of non-interactive multi-party protocol.
357        ///
358        /// - user_id: user j's user_id in the protocol
359        fn key_switch(&self, user_id: usize) -> Mat {
360            BoolEvaluator::with_local(|e| {
361                assert!(self.dimension() == (2, e.parameters().rlwe_n().0));
362                let server_key = BOOL_SERVER_KEY.get().unwrap();
363                let ksk = server_key.ui_to_s_ksk(user_id);
364                let decomposer = e.ni_ui_to_s_ks_decomposer().as_ref().unwrap();
365
366                // perform key switch
367                rlwe_key_switch(
368                    self,
369                    ksk.as_ref(),
370                    ksk.shoup_repr(),
371                    decomposer,
372                    e.pbs_info().nttop_rlweq(),
373                    e.pbs_info().modop_rlweq(),
374                )
375            })
376        }
377    }
378
379    impl<C> KeySwitchWithId<BatchedFheBools<C>> for NonInteractiveBatchedFheBools<C>
380    where
381        C: KeySwitchWithId<C>,
382    {
383        /// Key switch `Self`'s vector of RLWE ciphertexts from user j's RLWE
384        /// secret u_j to ideal RLWE secret `s` of non-interactive
385        /// multi-party protocol.
386        ///
387        /// Returns vector of key switched RLWE ciphertext as `BatchedFheBools`
388        /// which can then be used to extract individual Bool LWE ciphertexts.
389        ///
390        /// - user_id: user j's user_id in the protocol
391        fn key_switch(&self, user_id: usize) -> BatchedFheBools<C> {
392            let data = self
393                .data
394                .iter()
395                .map(|c| c.key_switch(user_id))
396                .collect_vec();
397            BatchedFheBools { data }
398        }
399    }
400}
401
402#[cfg(test)]
403mod tests {
404    use itertools::Itertools;
405    use rand::{thread_rng, RngCore};
406
407    use crate::{
408        bool::{
409            keys::tests::{ideal_sk_rlwe, measure_noise_lwe},
410            BooleanGates,
411        },
412        Encoder, Encryptor, KeySwitchWithId, MultiPartyDecryptor,
413    };
414
415    use super::*;
416
417    #[test]
418    fn non_interactive_mp_bool_nand() {
419        set_parameter_set(ParameterSelector::NonInteractiveLTE4Party);
420        let mut seed = [0u8; 32];
421        thread_rng().fill_bytes(&mut seed);
422        set_common_reference_seed(seed);
423
424        let parties = 3;
425
426        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
427
428        let key_shares = cks
429            .iter()
430            .enumerate()
431            .map(|(user_index, ck)| gen_server_key_share(user_index, parties, ck))
432            .collect_vec();
433
434        let seeded_server_key = aggregate_server_key_shares(&key_shares);
435        seeded_server_key.set_server_key();
436
437        let parameters = BoolEvaluator::with_local(|e| e.parameters().clone());
438        let rlwe_modop = parameters.default_rlwe_modop();
439
440        let ideal_sk_rlwe = ideal_sk_rlwe(&cks);
441
442        let mut m0 = false;
443        let mut m1 = true;
444
445        let mut ct0 = {
446            let ct: NonInteractiveBatchedFheBools<_> = cks[0].encrypt(vec![m0].as_slice());
447            let ct = ct.key_switch(0);
448            ct.extract(0)
449        };
450        let mut ct1 = {
451            let ct: NonInteractiveBatchedFheBools<_> = cks[1].encrypt(vec![m1].as_slice());
452            let ct = ct.key_switch(1);
453            ct.extract(0)
454        };
455
456        for _ in 0..1000 {
457            // let now = std::time::Instant::now();
458            let ct_out =
459                BoolEvaluator::with_local_mut(|e| e.xor(&ct0, &ct1, RuntimeServerKey::global()));
460            // println!("Time: {:?}", now.elapsed());
461
462            let decryption_shares = cks
463                .iter()
464                .map(|k| k.gen_decryption_share(&ct_out))
465                .collect_vec();
466            let m_out = cks[0].aggregate_decryption_shares(&ct_out, &decryption_shares);
467
468            let m_expected = m0 ^ m1;
469
470            {
471                let noise = measure_noise_lwe(
472                    &ct_out,
473                    parameters.rlwe_q().encode(m_expected),
474                    &ideal_sk_rlwe,
475                    &rlwe_modop,
476                );
477                println!("Noise: {noise}");
478            }
479
480            assert!(m_out == m_expected, "Expected {m_expected} but got {m_out}");
481
482            m1 = m0;
483            m0 = m_out;
484
485            ct1 = ct0;
486            ct0 = ct_out;
487        }
488    }
489}