arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/bool/mp_api.rs
src/bool/mp_api.rsRaw
  1use std::{cell::RefCell, sync::OnceLock};
  2
  3use crate::{
  4    backend::{ModularOpsU64, ModulusPowerOf2},
  5    ntt::NttBackendU64,
  6    random::{DefaultSecureRng, NewWithSeed},
  7    utils::{Global, WithLocal},
  8};
  9
 10use super::{evaluator::MultiPartyCrs, keys::*, parameters::*, ClientKey, ParameterSelector};
 11
 12pub(crate) type BoolEvaluator = super::evaluator::BoolEvaluator<
 13    Vec<Vec<u64>>,
 14    NttBackendU64,
 15    ModularOpsU64<CiphertextModulus<u64>>,
 16    ModulusPowerOf2<CiphertextModulus<u64>>,
 17    ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>,
 18>;
 19
 20thread_local! {
 21    static BOOL_EVALUATOR: RefCell<Option<BoolEvaluator>> = RefCell::new(None);
 22
 23}
 24static BOOL_SERVER_KEY: OnceLock<ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>> = OnceLock::new();
 25
 26static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new();
 27
 28pub fn set_parameter_set(select: ParameterSelector) {
 29    match select {
 30        ParameterSelector::MultiPartyLessThanOrEqualTo16 => {
 31            BOOL_EVALUATOR.with_borrow_mut(|v| *v = Some(BoolEvaluator::new(SMALL_MP_BOOL_PARAMS)));
 32        }
 33        _ => {
 34            panic!("Paramerters not supported")
 35        }
 36    }
 37}
 38
 39pub fn set_mp_seed(seed: [u8; 32]) {
 40    assert!(
 41        MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(),
 42        "Attempted to set MP SEED twice."
 43    )
 44}
 45
 46pub fn gen_client_key() -> ClientKey {
 47    BoolEvaluator::with_local(|e| e.client_key())
 48}
 49
 50pub fn gen_mp_keys_phase1(
 51    ck: &ClientKey,
 52) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> {
 53    let seed = MultiPartyCrs::global().public_key_share_seed::<DefaultSecureRng>();
 54    BoolEvaluator::with_local(|e| {
 55        let pk_share = e.multi_party_public_key_share(seed, ck);
 56        pk_share
 57    })
 58}
 59
 60pub fn gen_mp_keys_phase2<R, ModOp>(
 61    ck: &ClientKey,
 62    pk: &PublicKey<Vec<Vec<u64>>, R, ModOp>,
 63) -> CommonReferenceSeededMultiPartyServerKeyShare<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
 64    let seed = MultiPartyCrs::global().server_key_share_seed::<DefaultSecureRng>();
 65    BoolEvaluator::with_local_mut(|e| {
 66        let server_key_share = e.multi_party_server_key_share(seed, pk.key(), ck);
 67        server_key_share
 68    })
 69}
 70
 71pub fn aggregate_public_key_shares(
 72    shares: &[CommonReferenceSeededCollectivePublicKeyShare<
 73        Vec<u64>,
 74        [u8; 32],
 75        BoolParameters<u64>,
 76    >],
 77) -> PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>> {
 78    PublicKey::from(shares)
 79}
 80
 81pub fn aggregate_server_key_shares(
 82    shares: &[CommonReferenceSeededMultiPartyServerKeyShare<
 83        Vec<Vec<u64>>,
 84        BoolParameters<u64>,
 85        [u8; 32],
 86    >],
 87) -> SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>> {
 88    BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares))
 89}
 90
 91impl
 92    SeededMultiPartyServerKey<
 93        Vec<Vec<u64>>,
 94        <DefaultSecureRng as NewWithSeed>::Seed,
 95        BoolParameters<u64>,
 96    >
 97{
 98    pub fn set_server_key(&self) {
 99        assert!(
100            BOOL_SERVER_KEY
101                .set(ShoupServerKeyEvaluationDomain::from(
102                    ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(self),
103                ))
104                .is_ok(),
105            "Attempted to set server key twice."
106        );
107    }
108}
109
110//  MULTIPARTY CRS  //
111impl Global for MultiPartyCrs<[u8; 32]> {
112    fn global() -> &'static Self {
113        MULTI_PARTY_CRS
114            .get()
115            .expect("Multi Party Common Reference String not set")
116    }
117}
118
119// BOOL EVALUATOR //
120impl WithLocal for BoolEvaluator {
121    fn with_local<F, R>(func: F) -> R
122    where
123        F: Fn(&Self) -> R,
124    {
125        BOOL_EVALUATOR.with_borrow(|s| func(s.as_ref().expect("Parameters not set")))
126    }
127
128    fn with_local_mut<F, R>(func: F) -> R
129    where
130        F: Fn(&mut Self) -> R,
131    {
132        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
133    }
134
135    fn with_local_mut_mut<F, R>(func: &mut F) -> R
136    where
137        F: FnMut(&mut Self) -> R,
138    {
139        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
140    }
141}
142
143pub(crate) type RuntimeServerKey = ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>;
144impl Global for RuntimeServerKey {
145    fn global() -> &'static Self {
146        BOOL_SERVER_KEY.get().expect("Server key not set!")
147    }
148}
149
150mod impl_enc_dec {
151    use crate::{
152        bool::evaluator::BoolEncoding,
153        pbs::{sample_extract, PbsInfo},
154        rgsw::public_key_encrypt_rlwe,
155        Encryptor, Matrix, MatrixEntity, MultiPartyDecryptor, RowEntity,
156    };
157    use itertools::Itertools;
158    use num_traits::{ToPrimitive, Zero};
159
160    use super::*;
161
162    type Mat = Vec<Vec<u64>>;
163
164    impl<Rng, ModOp> Encryptor<[bool], Vec<Mat>> for PublicKey<Mat, Rng, ModOp> {
165        fn encrypt(&self, m: &[bool]) -> Vec<Mat> {
166            BoolEvaluator::with_local(|e| {
167                DefaultSecureRng::with_local_mut(|rng| {
168                    let parameters = e.parameters();
169                    let ring_size = parameters.rlwe_n().0;
170
171                    let rlwe_count = ((m.len() as f64 / ring_size as f64).ceil())
172                        .to_usize()
173                        .unwrap();
174
175                    // encrypt `m` into ceil(len(m)/N) RLWE ciphertexts
176                    let rlwes = (0..rlwe_count)
177                        .map(|index| {
178                            let mut message = vec![<Mat as Matrix>::MatElement::zero(); ring_size];
179                            m[(index * ring_size)..std::cmp::min(m.len(), (index + 1) * ring_size)]
180                                .iter()
181                                .enumerate()
182                                .for_each(|(i, v)| {
183                                    if *v {
184                                        message[i] = parameters.rlwe_q().true_el()
185                                    } else {
186                                        message[i] = parameters.rlwe_q().false_el()
187                                    }
188                                });
189
190                            // encrypt message
191                            let mut rlwe_out =
192                                <Mat as MatrixEntity>::zeros(2, parameters.rlwe_n().0);
193
194                            public_key_encrypt_rlwe::<_, _, _, _, i32, _>(
195                                &mut rlwe_out,
196                                self.key(),
197                                &message,
198                                e.pbs_info().modop_rlweq(),
199                                e.pbs_info().nttop_rlweq(),
200                                rng,
201                            );
202
203                            rlwe_out
204                        })
205                        .collect_vec();
206                    rlwes
207                })
208            })
209        }
210    }
211
212    impl<Rng, ModOp> Encryptor<bool, <Mat as Matrix>::R> for PublicKey<Mat, Rng, ModOp> {
213        fn encrypt(&self, m: &bool) -> <Mat as Matrix>::R {
214            let m = vec![*m];
215            let rlwe = &self.encrypt(m.as_slice())[0];
216            BoolEvaluator::with_local(|e| {
217                let mut lwe = <Mat as Matrix>::R::zeros(e.parameters().rlwe_n().0 + 1);
218                sample_extract(&mut lwe, rlwe, e.pbs_info().modop_rlweq(), 0);
219                lwe
220            })
221        }
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use itertools::Itertools;
228    use rand::{thread_rng, RngCore};
229
230    use crate::{bool::evaluator::BooleanGates, Encryptor, MultiPartyDecryptor};
231
232    use super::*;
233
234    #[test]
235    fn multi_party_bool_gates() {
236        set_parameter_set(ParameterSelector::MultiPartyLessThanOrEqualTo16);
237        let mut seed = [0u8; 32];
238        thread_rng().fill_bytes(&mut seed);
239        set_mp_seed(seed);
240
241        let parties = 2;
242        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
243
244        // round 1
245        let pk_shares = cks.iter().map(|k| gen_mp_keys_phase1(k)).collect_vec();
246
247        // collective pk
248        let pk = aggregate_public_key_shares(&pk_shares);
249
250        // round 2
251        let server_key_shares = cks.iter().map(|k| gen_mp_keys_phase2(k, &pk)).collect_vec();
252
253        // server key
254        let server_key = aggregate_server_key_shares(&server_key_shares);
255        server_key.set_server_key();
256
257        let mut m0 = false;
258        let mut m1 = true;
259
260        let mut ct0 = pk.encrypt(&m0);
261        let mut ct1 = pk.encrypt(&m1);
262
263        for _ in 0..500 {
264            let ct_out =
265                BoolEvaluator::with_local_mut(|e| e.nand(&ct0, &ct1, RuntimeServerKey::global()));
266
267            let m_expected = !(m0 && m1);
268
269            let decryption_shares = cks
270                .iter()
271                .map(|k| k.gen_decryption_share(&ct_out))
272                .collect_vec();
273            let m_out = cks[0].aggregate_decryption_shares(&ct_out, &decryption_shares);
274
275            assert!(m_out == m_expected, "Expected {m_expected}, got {m_out}");
276
277            m1 = m0;
278            m0 = m_expected;
279
280            ct1 = ct0;
281            ct0 = ct_out;
282        }
283
284        for _ in 0..500 {
285            let ct_out =
286                BoolEvaluator::with_local_mut(|e| e.xnor(&ct0, &ct1, RuntimeServerKey::global()));
287
288            let m_expected = !(m0 ^ m1);
289
290            let decryption_shares = cks
291                .iter()
292                .map(|k| k.gen_decryption_share(&ct_out))
293                .collect_vec();
294            let m_out = cks[0].aggregate_decryption_shares(&ct_out, &decryption_shares);
295
296            assert!(m_out == m_expected, "Expected {m_expected}, got {m_out}");
297
298            m1 = m0;
299            m0 = m_expected;
300
301            ct1 = ct0;
302            ct0 = ct_out;
303        }
304    }
305
306    mod sp_api {
307        use num_traits::ToPrimitive;
308        use rand::Rng;
309
310        use crate::{
311            backend::ModulusPowerOf2, evaluator::BoolEncoding, pbs::PbsInfo,
312            rgsw::secret_key_encrypt_rlwe, utils::WithLocal, Decryptor, ModularOpsU64,
313            NttBackendU64, SampleExtractor,
314        };
315
316        use super::*;
317
318        pub(crate) fn set_single_party_parameter_sets(parameter: BoolParameters<u64>) {
319            BOOL_EVALUATOR.with_borrow_mut(|e| *e = Some(BoolEvaluator::new(parameter)));
320        }
321
322        // SERVER KEY EVAL (/SHOUP) DOMAIN //
323        impl SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
324            pub fn set_server_key(&self) {
325                assert!(
326                    BOOL_SERVER_KEY
327                        .set(
328                            ShoupServerKeyEvaluationDomain::from(ServerKeyEvaluationDomain::<
329                                _,
330                                _,
331                                DefaultSecureRng,
332                                NttBackendU64,
333                            >::from(
334                                self
335                            ),)
336                        )
337                        .is_ok(),
338                    "Attempted to set server key twice."
339                );
340            }
341        }
342
343        pub(crate) fn gen_keys() -> (
344            ClientKey,
345            SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>,
346        ) {
347            super::BoolEvaluator::with_local_mut(|e| {
348                let ck = e.client_key();
349                let sk = e.single_party_server_key(&ck);
350
351                (ck, sk)
352            })
353        }
354
355        impl<K: SinglePartyClientKey<Element = i32>> Encryptor<bool, Vec<u64>> for K {
356            fn encrypt(&self, m: &bool) -> Vec<u64> {
357                BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
358            }
359        }
360
361        impl<K: SinglePartyClientKey<Element = i32>> Decryptor<bool, Vec<u64>> for K {
362            fn decrypt(&self, c: &Vec<u64>) -> bool {
363                BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
364            }
365        }
366        impl<K> Encryptor<[bool], (Vec<Vec<u64>>, [u8; 32])> for K
367        where
368            K: SinglePartyClientKey<Element = i32>,
369        {
370            fn encrypt(&self, m: &[bool]) -> (Vec<Vec<u64>>, [u8; 32]) {
371                BoolEvaluator::with_local(|e| {
372                    DefaultSecureRng::with_local_mut(|rng| {
373                        let parameters = e.parameters();
374                        let ring_size = parameters.rlwe_n().0;
375
376                        let rlwe_count = ((m.len() as f64 / ring_size as f64).ceil())
377                            .to_usize()
378                            .unwrap();
379
380                        let mut seed = <DefaultSecureRng as NewWithSeed>::Seed::default();
381                        rng.fill_bytes(&mut seed);
382                        let mut prng = DefaultSecureRng::new_seeded(seed);
383
384                        let sk_u = self.sk_rlwe();
385
386                        // encrypt `m` into ceil(len(m)/N) RLWE ciphertexts
387                        let rlwes = (0..rlwe_count)
388                            .map(|index| {
389                                let mut message = vec![0; ring_size];
390                                m[(index * ring_size)
391                                    ..std::cmp::min(m.len(), (index + 1) * ring_size)]
392                                    .iter()
393                                    .enumerate()
394                                    .for_each(|(i, v)| {
395                                        if *v {
396                                            message[i] = parameters.rlwe_q().true_el()
397                                        } else {
398                                            message[i] = parameters.rlwe_q().false_el()
399                                        }
400                                    });
401
402                                // encrypt message
403                                let mut rlwe_out = vec![0u64; parameters.rlwe_n().0];
404
405                                secret_key_encrypt_rlwe(
406                                    &message,
407                                    &mut rlwe_out,
408                                    &sk_u,
409                                    e.pbs_info().modop_rlweq(),
410                                    e.pbs_info().nttop_rlweq(),
411                                    &mut prng,
412                                    rng,
413                                );
414
415                                rlwe_out
416                            })
417                            .collect_vec();
418
419                        (rlwes, seed)
420                    })
421                })
422            }
423        }
424
425        #[test]
426        fn batch_extract_works() {
427            set_single_party_parameter_sets(SP_TEST_BOOL_PARAMS);
428
429            let (ck, sk) = gen_keys();
430            sk.set_server_key();
431
432            let batch_size = (SP_TEST_BOOL_PARAMS.rlwe_n().0 * 3 + 123);
433            let m = (0..batch_size)
434                .map(|_| thread_rng().gen::<u8>())
435                .collect_vec();
436
437            let seeded_ct = ck.encrypt(m.as_slice());
438            let ct = seeded_ct.unseed::<Vec<Vec<u64>>>();
439
440            let m_back = (0..batch_size)
441                .map(|i| ck.decrypt(&ct.extract(i)))
442                .collect_vec();
443
444            assert_eq!(m, m_back);
445        }
446
447        #[test]
448        #[cfg(feature = "interactive_mp")]
449        fn all_uint8_apis() {
450            use num_traits::Euclid;
451
452            set_single_party_parameter_sets(SP_TEST_BOOL_PARAMS);
453
454            let (ck, sk) = gen_keys();
455            sk.set_server_key();
456
457            for i in 144..=255 {
458                for j in 100..=255 {
459                    let m0 = i;
460                    let m1 = j;
461                    let c0 = ck.encrypt(&m0);
462                    let c1 = ck.encrypt(&m1);
463
464                    assert!(ck.decrypt(&c0) == m0);
465                    assert!(ck.decrypt(&c1) == m1);
466
467                    // Arithmetic
468                    {
469                        {
470                            // Add
471                            let c_add = &c0 + &c1;
472                            let m0_plus_m1 = ck.decrypt(&c_add);
473                            assert_eq!(
474                                m0_plus_m1,
475                                m0.wrapping_add(m1),
476                                "Expected {} but got {m0_plus_m1} for {i}+{j}",
477                                m0.wrapping_add(m1)
478                            );
479                        }
480                        {
481                            // Sub
482                            let c_sub = &c0 - &c1;
483                            let m0_sub_m1 = ck.decrypt(&c_sub);
484                            assert_eq!(
485                                m0_sub_m1,
486                                m0.wrapping_sub(m1),
487                                "Expected {} but got {m0_sub_m1} for {i}-{j}",
488                                m0.wrapping_sub(m1)
489                            );
490                        }
491
492                        {
493                            // Mul
494                            let c_m0m1 = &c0 * &c1;
495                            let m0m1 = ck.decrypt(&c_m0m1);
496                            assert_eq!(
497                                m0m1,
498                                m0.wrapping_mul(m1),
499                                "Expected {} but got {m0m1} for {i}x{j}",
500                                m0.wrapping_mul(m1)
501                            );
502                        }
503
504                        // Div & Rem
505                        {
506                            let (c_quotient, c_rem) = c0.div_rem(&c1);
507                            let m_quotient = ck.decrypt(&c_quotient);
508                            let m_remainder = ck.decrypt(&c_rem);
509                            if j != 0 {
510                                let (q, r) = i.div_rem_euclid(&j);
511                                assert_eq!(
512                                    m_quotient, q,
513                                    "Expected {} but got {m_quotient} for {i}/{j}",
514                                    q
515                                );
516                                assert_eq!(
517                                    m_remainder, r,
518                                    "Expected {} but got {m_quotient} for {i}%{j}",
519                                    r
520                                );
521                            } else {
522                                assert_eq!(
523                                    m_quotient, 255,
524                                    "Expected 255 but got {m_quotient}. Case div by zero"
525                                );
526                                assert_eq!(
527                                    m_remainder, i,
528                                    "Expected {i} but got {m_quotient}. Case div by zero"
529                                )
530                            }
531                        }
532                    }
533
534                    // Comparisons
535                    {
536                        {
537                            let c_eq = c0.eq(&c1);
538                            let is_eq = ck.decrypt(&c_eq);
539                            assert_eq!(
540                                is_eq,
541                                i == j,
542                                "Expected {} but got {is_eq} for {i}=={j}",
543                                i == j
544                            );
545                        }
546
547                        {
548                            let c_gt = c0.gt(&c1);
549                            let is_gt = ck.decrypt(&c_gt);
550                            assert_eq!(
551                                is_gt,
552                                i > j,
553                                "Expected {} but got {is_gt} for {i}>{j}",
554                                i > j
555                            );
556                        }
557
558                        {
559                            let c_lt = c0.lt(&c1);
560                            let is_lt = ck.decrypt(&c_lt);
561                            assert_eq!(
562                                is_lt,
563                                i < j,
564                                "Expected {} but got {is_lt} for {i}<{j}",
565                                i < j
566                            );
567                        }
568
569                        {
570                            let c_ge = c0.ge(&c1);
571                            let is_ge = ck.decrypt(&c_ge);
572                            assert_eq!(
573                                is_ge,
574                                i >= j,
575                                "Expected {} but got {is_ge} for {i}>={j}",
576                                i >= j
577                            );
578                        }
579
580                        {
581                            let c_le = c0.le(&c1);
582                            let is_le = ck.decrypt(&c_le);
583                            assert_eq!(
584                                is_le,
585                                i <= j,
586                                "Expected {} but got {is_le} for {i}<={j}",
587                                i <= j
588                            );
589                        }
590                    }
591                }
592            }
593        }
594    }
595}