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