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