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