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