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::BoolEncoding,
269 keys::tests::{ideal_sk_rlwe, measure_noise_lwe},
270 BooleanGates,
271 },
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 bool::impl_bool_frontend::FheBool, pbs::PbsInfo, rgsw::seeded_secret_key_encrypt_rlwe,
372 Decryptor, SampleExtractor,
373 };
374
375 use super::*;
376
377 pub(crate) fn set_single_party_parameter_sets(parameter: BoolParameters<u64>) {
378 BOOL_EVALUATOR.with_borrow_mut(|e| *e = Some(BoolEvaluator::new(parameter)));
379 }
380
381 // SERVER KEY EVAL (/SHOUP) DOMAIN //
382 impl SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
383 pub fn set_server_key(&self) {
384 assert!(
385 BOOL_SERVER_KEY
386 .set(
387 ShoupServerKeyEvaluationDomain::from(ServerKeyEvaluationDomain::<
388 _,
389 _,
390 DefaultSecureRng,
391 NttBackendU64,
392 >::from(
393 self
394 ),)
395 )
396 .is_ok(),
397 "Attempted to set server key twice."
398 );
399 }
400 }
401
402 pub(crate) fn gen_keys() -> (
403 ClientKey,
404 SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>,
405 ) {
406 super::BoolEvaluator::with_local_mut(|e| {
407 let ck = e.client_key();
408 let sk = e.single_party_server_key(&ck);
409
410 (ck, sk)
411 })
412 }
413
414 impl<K: SinglePartyClientKey<Element = i32>> Encryptor<bool, Vec<u64>> for K {
415 fn encrypt(&self, m: &bool) -> Vec<u64> {
416 BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
417 }
418 }
419
420 impl<K: SinglePartyClientKey<Element = i32>> Decryptor<bool, Vec<u64>> for K {
421 fn decrypt(&self, c: &Vec<u64>) -> bool {
422 BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
423 }
424 }
425
426 impl<K: SinglePartyClientKey<Element = i32>, C> Encryptor<bool, FheBool<C>> for K
427 where
428 K: Encryptor<bool, C>,
429 {
430 fn encrypt(&self, m: &bool) -> FheBool<C> {
431 FheBool {
432 data: self.encrypt(m),
433 }
434 }
435 }
436
437 impl<K: SinglePartyClientKey<Element = i32>, C> Decryptor<bool, FheBool<C>> for K
438 where
439 K: Decryptor<bool, C>,
440 {
441 fn decrypt(&self, c: &FheBool<C>) -> bool {
442 self.decrypt(c.data())
443 }
444 }
445
446 impl<K> Encryptor<[bool], (Vec<Vec<u64>>, [u8; 32])> for K
447 where
448 K: SinglePartyClientKey<Element = i32>,
449 {
450 fn encrypt(&self, m: &[bool]) -> (Vec<Vec<u64>>, [u8; 32]) {
451 BoolEvaluator::with_local(|e| {
452 DefaultSecureRng::with_local_mut(|rng| {
453 let parameters = e.parameters();
454 let ring_size = parameters.rlwe_n().0;
455
456 let rlwe_count = ((m.len() as f64 / ring_size as f64).ceil())
457 .to_usize()
458 .unwrap();
459
460 let mut seed = <DefaultSecureRng as NewWithSeed>::Seed::default();
461 rng.fill_bytes(&mut seed);
462 let mut prng = DefaultSecureRng::new_seeded(seed);
463
464 let sk_u = self.sk_rlwe();
465
466 // encrypt `m` into ceil(len(m)/N) RLWE ciphertexts
467 let rlwes = (0..rlwe_count)
468 .map(|index| {
469 let mut message = vec![0; ring_size];
470 m[(index * ring_size)
471 ..std::cmp::min(m.len(), (index + 1) * ring_size)]
472 .iter()
473 .enumerate()
474 .for_each(|(i, v)| {
475 if *v {
476 message[i] = parameters.rlwe_q().true_el()
477 } else {
478 message[i] = parameters.rlwe_q().false_el()
479 }
480 });
481
482 // encrypt message
483 let mut rlwe_out = vec![0u64; parameters.rlwe_n().0];
484 seeded_secret_key_encrypt_rlwe(
485 &message,
486 &mut rlwe_out,
487 &sk_u,
488 e.pbs_info().modop_rlweq(),
489 e.pbs_info().nttop_rlweq(),
490 &mut prng,
491 rng,
492 );
493
494 rlwe_out
495 })
496 .collect_vec();
497
498 (rlwes, seed)
499 })
500 })
501 }
502 }
503
504 #[test]
505 fn batch_extract_works() {
506 set_single_party_parameter_sets(SP_TEST_BOOL_PARAMS);
507
508 let (ck, sk) = gen_keys();
509 sk.set_server_key();
510
511 let batch_size = (SP_TEST_BOOL_PARAMS.rlwe_n().0 * 3 + 123);
512 let m = (0..batch_size)
513 .map(|_| thread_rng().gen::<u8>())
514 .collect_vec();
515
516 let seeded_ct = ck.encrypt(m.as_slice());
517 let ct = seeded_ct.unseed::<Vec<Vec<u64>>>();
518
519 let m_back = (0..batch_size)
520 .map(|i| ck.decrypt(&ct.extract(i)))
521 .collect_vec();
522
523 assert_eq!(m, m_back);
524 }
525
526 #[test]
527 #[cfg(feature = "interactive_mp")]
528 fn all_uint8_apis() {
529 use num_traits::Euclid;
530
531 use crate::div_zero_error_flag;
532
533 set_single_party_parameter_sets(SP_TEST_BOOL_PARAMS);
534
535 let (ck, sk) = gen_keys();
536 sk.set_server_key();
537
538 for i in 0..=255 {
539 for j in 0..=255 {
540 let m0 = i;
541 let m1 = j;
542 let c0 = ck.encrypt(&m0);
543 let c1 = ck.encrypt(&m1);
544
545 assert!(ck.decrypt(&c0) == m0);
546 assert!(ck.decrypt(&c1) == m1);
547
548 // Arithmetic
549 {
550 {
551 // Add
552 let c_add = &c0 + &c1;
553 let m0_plus_m1 = ck.decrypt(&c_add);
554 assert_eq!(
555 m0_plus_m1,
556 m0.wrapping_add(m1),
557 "Expected {} but got {m0_plus_m1} for
558 {i}+{j}",
559 m0.wrapping_add(m1)
560 );
561 }
562 {
563 // Sub
564 let c_sub = &c0 - &c1;
565 let m0_sub_m1 = ck.decrypt(&c_sub);
566 assert_eq!(
567 m0_sub_m1,
568 m0.wrapping_sub(m1),
569 "Expected {} but got {m0_sub_m1} for
570 {i}-{j}",
571 m0.wrapping_sub(m1)
572 );
573 }
574
575 {
576 // Mul
577 let c_m0m1 = &c0 * &c1;
578 let m0m1 = ck.decrypt(&c_m0m1);
579 assert_eq!(
580 m0m1,
581 m0.wrapping_mul(m1),
582 "Expected {} but got {m0m1} for {i}x{j}",
583 m0.wrapping_mul(m1)
584 );
585 }
586
587 // Div & Rem
588 {
589 let (c_quotient, c_rem) = c0.div_rem(&c1);
590 let m_quotient = ck.decrypt(&c_quotient);
591 let m_remainder = ck.decrypt(&c_rem);
592 if j != 0 {
593 let (q, r) = i.div_rem_euclid(&j);
594 assert_eq!(
595 m_quotient, q,
596 "Expected {} but got {m_quotient} for
597 {i}/{j}",
598 q
599 );
600 assert_eq!(
601 m_remainder, r,
602 "Expected {} but got {m_remainder} for
603 {i}%{j}",
604 r
605 );
606 } else {
607 assert_eq!(
608 m_quotient, 255,
609 "Expected 255 but got {m_quotient}. Case
610 div by zero"
611 );
612 assert_eq!(
613 m_remainder, i,
614 "Expected {i} but got {m_remainder}. Case
615 div by zero"
616 );
617
618 let div_by_zero = ck.decrypt(&div_zero_error_flag().unwrap());
619 assert_eq!(
620 div_by_zero, true,
621 "Expected true but got {div_by_zero}"
622 );
623 }
624 }
625 }
626
627 // Comparisons
628 {
629 {
630 let c_eq = c0.eq(&c1);
631 let is_eq = ck.decrypt(&c_eq);
632 assert_eq!(
633 is_eq,
634 i == j,
635 "Expected {} but got {is_eq} for {i}=={j}",
636 i == j
637 );
638 }
639
640 {
641 let c_gt = c0.gt(&c1);
642 let is_gt = ck.decrypt(&c_gt);
643 assert_eq!(
644 is_gt,
645 i > j,
646 "Expected {} but got {is_gt} for {i}>{j}",
647 i > j
648 );
649 }
650
651 {
652 let c_lt = c0.lt(&c1);
653 let is_lt = ck.decrypt(&c_lt);
654 assert_eq!(
655 is_lt,
656 i < j,
657 "Expected {} but got {is_lt} for {i}<{j}",
658 i < j
659 );
660 }
661
662 {
663 let c_ge = c0.ge(&c1);
664 let is_ge = ck.decrypt(&c_ge);
665 assert_eq!(
666 is_ge,
667 i >= j,
668 "Expected {} but got {is_ge} for {i}>={j}",
669 i >= j
670 );
671 }
672
673 {
674 let c_le = c0.le(&c1);
675 let is_le = ck.decrypt(&c_le);
676 assert_eq!(
677 is_le,
678 i <= j,
679 "Expected {} but got {is_le} for {i}<={j}",
680 i <= j
681 );
682 }
683 }
684 }
685 }
686 }
687
688 #[test]
689 #[cfg(feature = "interactive_mp")]
690 fn all_bool_apis() {
691 use crate::FheBool;
692
693 set_single_party_parameter_sets(SP_TEST_BOOL_PARAMS);
694
695 let (ck, sk) = gen_keys();
696 sk.set_server_key();
697
698 for _ in 0..100 {
699 let a = thread_rng().gen_bool(0.5);
700 let b = thread_rng().gen_bool(0.5);
701
702 let c_a: FheBool = ck.encrypt(&a);
703 let c_b: FheBool = ck.encrypt(&b);
704
705 let c_out = &c_a & &c_b;
706 let out = ck.decrypt(&c_out);
707 assert_eq!(out, a & b, "Expected {} but got {out}", a & b);
708
709 let c_out = &c_a | &c_b;
710 let out = ck.decrypt(&c_out);
711 assert_eq!(out, a | b, "Expected {} but got {out}", a | b);
712
713 let c_out = &c_a ^ &c_b;
714 let out = ck.decrypt(&c_out);
715 assert_eq!(out, a ^ b, "Expected {} but got {out}", a ^ b);
716
717 let c_out = !(&c_a);
718 let out = ck.decrypt(&c_out);
719 assert_eq!(out, !a, "Expected {} but got {out}", !a);
720 }
721 }
722 }
723}