fhe-study/gfhe/src/glwe.rs
gfhe/src/glwe.rsRaw
1//! Generalized LWE.
2//!
3
4use anyhow::Result;
5use itertools::zip_eq;
6use rand::Rng;
7use rand_distr::{Normal, Uniform};
8use std::iter::Sum;
9use std::ops::{Add, AddAssign, Mul, Sub};
10
11use arith::{Ring, Rq, Zq, TR};
12
13use crate::glev::GLev;
14
15// const ERR_SIGMA: f64 = 3.2;
16const ERR_SIGMA: f64 = 0.0; // TODO WIP
17
18/// GLWE implemented over the `Ring` trait, so that it can be also instantiated
19/// over the Torus polynomials 𝕋_<N,q>[X] = 𝕋_q[X]/ (X^N+1).
20#[derive(Clone, Debug)]
21pub struct GLWE<R: Ring, const K: usize>(pub TR<R, K>, pub R);
22
23#[derive(Clone, Debug)]
24pub struct SecretKey<R: Ring, const K: usize>(pub TR<R, K>);
25#[derive(Clone, Debug)]
26pub struct PublicKey<R: Ring, const K: usize>(pub R, pub TR<R, K>);
27
28// K GLevs, each KSK_i=l GLWEs
29#[derive(Clone, Debug)]
30pub struct KSK<R: Ring, const K: usize>(Vec<GLev<R, K>>);
31
32impl<R: Ring, const K: usize> GLWE<R, K> {
33 pub fn zero() -> Self {
34 Self(TR::zero(), R::zero())
35 }
36
37 pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> {
38 let Xi_key = Uniform::new(0_f64, 2_f64);
39 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
40
41 let s: TR<R, K> = TR::rand(&mut rng, Xi_key);
42 let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
43 let e = R::rand(&mut rng, Xi_err);
44
45 let pk: PublicKey<R, K> = PublicKey((&a * &s) + e, a);
46 Ok((SecretKey(s), pk))
47 }
48 pub fn pk_from_sk(mut rng: impl Rng, sk: SecretKey<R, K>) -> Result<PublicKey<R, K>> {
49 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
50
51 let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
52 let e = R::rand(&mut rng, Xi_err);
53
54 let pk: PublicKey<R, K> = PublicKey((&a * &sk.0) + e, a);
55 Ok(pk)
56 }
57
58 pub fn new_ksk(
59 mut rng: impl Rng,
60 beta: u32,
61 l: u32,
62 sk: &SecretKey<R, K>,
63 new_sk: &SecretKey<R, K>,
64 ) -> Result<KSK<R, K>> {
65 let r: Vec<GLev<R, K>> = (0..K)
66 .into_iter()
67 .map(|i|
68 // treat sk_i as the msg being encrypted
69 GLev::<R, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
70 .collect::<Result<Vec<_>>>()?;
71
72 Ok(KSK(r))
73 }
74 pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<R, K>) -> Self {
75 let (a, b): (TR<R, K>, R) = (self.0.clone(), self.1);
76
77 let lhs: GLWE<R, K> = GLWE(TR::zero(), b);
78
79 // K iterations, ksk.0 contains K times GLev
80 let rhs: GLWE<R, K> = zip_eq(a.0, ksk.0.clone())
81 .map(|(a_i, ksk_i)| ksk_i * a_i.decompose(beta, l)) // dot_product
82 .sum();
83
84 lhs - rhs
85 }
86
87 // encrypts with the given SecretKey (instead of PublicKey)
88 pub fn encrypt_s(
89 mut rng: impl Rng,
90 sk: &SecretKey<R, K>,
91 m: &R, // already scaled
92 ) -> Result<Self> {
93 let Xi_key = Uniform::new(0_f64, 2_f64);
94 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
95
96 let a: TR<R, K> = TR::rand(&mut rng, Xi_key);
97 let e = R::rand(&mut rng, Xi_err);
98
99 let b: R = (&a * &sk.0) + *m + e;
100 Ok(Self(a, b))
101 }
102 pub fn encrypt(
103 mut rng: impl Rng,
104 pk: &PublicKey<R, K>,
105 m: &R, // already scaled
106 ) -> Result<Self> {
107 let Xi_key = Uniform::new(0_f64, 2_f64);
108 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
109
110 let u: R = R::rand(&mut rng, Xi_key);
111
112 let e0 = R::rand(&mut rng, Xi_err);
113 let e1 = TR::<R, K>::rand(&mut rng, Xi_err);
114
115 let b: R = pk.0.clone() * u.clone() + *m + e0;
116 let d: TR<R, K> = &pk.1 * &u + e1;
117
118 Ok(Self(d, b))
119 }
120 // returns m' not downscaled
121 pub fn decrypt(&self, sk: &SecretKey<R, K>) -> R {
122 let (d, b): (TR<R, K>, R) = (self.0.clone(), self.1);
123 let p: R = b - &d * &sk.0;
124 p
125 }
126}
127
128// Methods for when Ring=Rq<Q,N>
129impl<const Q: u64, const N: usize, const K: usize> GLWE<Rq<Q, N>, K> {
130 // scale up
131 pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
132 let m = m.remodule::<Q>();
133 let delta = Q / T; // floored
134 m * delta
135 }
136 // scale down
137 pub fn decode<const T: u64>(m: &Rq<Q, N>) -> Rq<T, N> {
138 let r = m.mul_div_round(T, Q);
139 let r: Rq<T, N> = r.remodule::<T>();
140 r
141 }
142 pub fn mod_switch<const P: u64>(&self) -> GLWE<Rq<P, N>, K> {
143 let a: TR<Rq<P, N>, K> = TR(self
144 .0
145 .0
146 .iter()
147 .map(|r| r.mod_switch::<P>())
148 .collect::<Vec<_>>());
149 let b: Rq<P, N> = self.1.mod_switch::<P>();
150 GLWE(a, b)
151 }
152}
153
154impl<R: Ring, const K: usize> Add<GLWE<R, K>> for GLWE<R, K> {
155 type Output = Self;
156 fn add(self, other: Self) -> Self {
157 let a: TR<R, K> = self.0 + other.0;
158 let b: R = self.1 + other.1;
159 Self(a, b)
160 }
161}
162
163impl<R: Ring, const K: usize> Add<R> for GLWE<R, K> {
164 type Output = Self;
165 fn add(self, plaintext: R) -> Self {
166 let a: TR<R, K> = self.0;
167 let b: R = self.1 + plaintext;
168 Self(a, b)
169 }
170}
171impl<R: Ring, const K: usize> AddAssign for GLWE<R, K> {
172 fn add_assign(&mut self, rhs: Self) {
173 for i in 0..K {
174 self.0 .0[i] = self.0 .0[i].clone() + rhs.0 .0[i].clone();
175 }
176 self.1 = self.1.clone() + rhs.1.clone();
177 }
178}
179impl<R: Ring, const K: usize> Sum<GLWE<R, K>> for GLWE<R, K> {
180 fn sum<I>(iter: I) -> Self
181 where
182 I: Iterator<Item = Self>,
183 {
184 let mut acc = GLWE::<R, K>::zero();
185 for e in iter {
186 acc += e;
187 }
188 acc
189 }
190}
191
192impl<R: Ring, const K: usize> Sub<GLWE<R, K>> for GLWE<R, K> {
193 type Output = Self;
194 fn sub(self, other: Self) -> Self {
195 let a: TR<R, K> = self.0 - other.0;
196 let b: R = self.1 - other.1;
197 Self(a, b)
198 }
199}
200
201impl<R: Ring, const K: usize> Mul<R> for GLWE<R, K> {
202 type Output = Self;
203 fn mul(self, plaintext: R) -> Self {
204 let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
205 let b: R = self.1 * plaintext;
206 Self(a, b)
207 }
208}
209// for when R = Rq<Q,N>
210// impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Rq<Q, N>, K> {
211// type Output = Self;
212// fn mul(self, plaintext: Rq<Q, N>) -> Self {
213// // first compute the NTT for plaintext, to avoid computing it at each
214// // iteration, speeding up the multiplications
215// let mut plaintext = plaintext.clone();
216// plaintext.compute_evals();
217//
218// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
219// let b: Rq<Q, N> = self.1 * plaintext;
220// Self(a, b)
221// }
222// }
223
224// impl<R: Ring, const K: usize> Mul<R::C> for GLWE<R, K>
225// // where
226// // // R: std::ops::Mul<<R as arith::Ring>::C>,
227// // // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
228// // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
229// {
230// type Output = Self;
231// fn mul(self, e: R::C) -> Self {
232// let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e.clone()).collect());
233// let b: R = self.1 * e.clone();
234// Self(a, b)
235// }
236// }
237
238// impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
239// type Output = Self;
240// fn mul(self, e: Zq<Q>) -> Self {
241// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
242// let b: Rq<Q, N> = self.1 * e;
243// Self(a, b)
244// }
245// }
246
247#[cfg(test)]
248mod tests {
249 use anyhow::Result;
250 use rand::distributions::Uniform;
251
252 use super::*;
253
254 #[test]
255 fn test_encrypt_decrypt() -> Result<()> {
256 const Q: u64 = 2u64.pow(16) + 1;
257 const N: usize = 128;
258 const T: u64 = 32; // plaintext modulus
259 const K: usize = 16;
260 type S = GLWE<Rq<Q, N>, K>;
261
262 let mut rng = rand::thread_rng();
263 let msg_dist = Uniform::new(0_u64, T);
264
265 for _ in 0..200 {
266 let (sk, pk) = S::new_key(&mut rng)?;
267
268 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; // msg
269 // let m: Rq<Q, N> = m.remodule::<Q>();
270
271 let p = S::encode::<T>(&m); // plaintext
272 let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
273 let p_recovered = c.decrypt(&sk);
274 let m_recovered = S::decode::<T>(&p_recovered);
275
276 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
277
278 // same but using encrypt_s (with sk instead of pk))
279 let c = S::encrypt_s(&mut rng, &sk, &p)?;
280 let p_recovered = c.decrypt(&sk);
281 let m_recovered = S::decode::<T>(&p_recovered);
282
283 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
284 }
285
286 Ok(())
287 }
288
289 use arith::{Tn, T64};
290 use std::array;
291 pub fn t_encode<const P: u64>(m: &Rq<P, 4>) -> Tn<4> {
292 let delta = u64::MAX / P; // floored
293 let coeffs = m.coeffs();
294 Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
295 }
296 pub fn t_decode<const P: u64>(p: &Tn<4>) -> Rq<P, 4> {
297 let p = p.mul_div_round(P, u64::MAX);
298 Rq::<P, 4>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
299 }
300 #[test]
301 fn test_encrypt_decrypt_torus() -> Result<()> {
302 const N: usize = 128;
303 const T: u64 = 32; // plaintext modulus
304 const K: usize = 16;
305 type S = GLWE<Tn<4>, K>;
306
307 let mut rng = rand::thread_rng();
308 let msg_dist = Uniform::new(0_f64, T as f64);
309
310 for _ in 0..200 {
311 let (sk, pk) = S::new_key(&mut rng)?;
312
313 let m = Rq::<T, 4>::rand(&mut rng, msg_dist); // msg
314
315 let p = t_encode::<T>(&m); // plaintext
316 let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
317 let p_recovered = c.decrypt(&sk);
318 let m_recovered = t_decode::<T>(&p_recovered);
319
320 assert_eq!(m, m_recovered);
321
322 // same but using encrypt_s (with sk instead of pk))
323 let c = S::encrypt_s(&mut rng, &sk, &p)?;
324 let p_recovered = c.decrypt(&sk);
325 let m_recovered = t_decode::<T>(&p_recovered);
326
327 assert_eq!(m, m_recovered);
328 }
329
330 Ok(())
331 }
332
333 #[test]
334 fn test_addition() -> Result<()> {
335 const Q: u64 = 2u64.pow(16) + 1;
336 const N: usize = 128;
337 const T: u64 = 20;
338 const K: usize = 16;
339 type S = GLWE<Rq<Q, N>, K>;
340
341 let mut rng = rand::thread_rng();
342 let msg_dist = Uniform::new(0_u64, T);
343
344 for _ in 0..200 {
345 let (sk, pk) = S::new_key(&mut rng)?;
346
347 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
348 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
349 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
350 let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext
351
352 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
353 let c2 = S::encrypt(&mut rng, &pk, &p2)?;
354
355 let c3 = c1 + c2;
356
357 let p3_recovered = c3.decrypt(&sk);
358 let m3_recovered = S::decode::<T>(&p3_recovered);
359
360 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
361 }
362
363 Ok(())
364 }
365
366 #[test]
367 fn test_add_plaintext() -> Result<()> {
368 const Q: u64 = 2u64.pow(16) + 1;
369 const N: usize = 128;
370 const T: u64 = 32;
371 const K: usize = 16;
372 type S = GLWE<Rq<Q, N>, K>;
373
374 let mut rng = rand::thread_rng();
375 let msg_dist = Uniform::new(0_u64, T);
376
377 for _ in 0..200 {
378 let (sk, pk) = S::new_key(&mut rng)?;
379
380 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
381 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
382 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
383 let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext
384
385 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
386
387 let c3 = c1 + p2;
388
389 let p3_recovered = c3.decrypt(&sk);
390 let m3_recovered = S::decode::<T>(&p3_recovered);
391
392 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
393 }
394
395 Ok(())
396 }
397
398 #[test]
399 fn test_mul_plaintext() -> Result<()> {
400 const Q: u64 = 2u64.pow(16) + 1;
401 const N: usize = 16;
402 const T: u64 = 4;
403 const K: usize = 16;
404 type S = GLWE<Rq<Q, N>, K>;
405
406 let mut rng = rand::thread_rng();
407 let msg_dist = Uniform::new(0_u64, T);
408
409 for _ in 0..200 {
410 let (sk, pk) = S::new_key(&mut rng)?;
411
412 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
413 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
414 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
415 let p2 = m2.remodule::<Q>(); // notice we don't encode (scale by delta)
416
417 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
418
419 let c3 = c1 * p2;
420
421 let p3_recovered: Rq<Q, N> = c3.decrypt(&sk);
422 let m3_recovered: Rq<T, N> = S::decode::<T>(&p3_recovered);
423 assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
424 }
425
426 Ok(())
427 }
428
429 #[test]
430 fn test_mod_switch() -> Result<()> {
431 const Q: u64 = 2u64.pow(16) + 1;
432 const P: u64 = 2u64.pow(8) + 1;
433 // note: wip, Q and P chosen so that P/Q is an integer
434 const N: usize = 8;
435 const T: u64 = 4; // plaintext modulus, must be a prime or power of a prime
436 const K: usize = 16;
437 type S = GLWE<Rq<Q, N>, K>;
438
439 let mut rng = rand::thread_rng();
440 let msg_dist = Uniform::new(0_u64, T);
441
442 for _ in 0..200 {
443 let (sk, pk) = S::new_key(&mut rng)?;
444
445 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
446
447 let p = S::encode::<T>(&m);
448 let c = S::encrypt(&mut rng, &pk, &p)?;
449
450 let c2: GLWE<Rq<P, N>, K> = c.mod_switch::<P>();
451 let sk2: SecretKey<Rq<P, N>, K> =
452 SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect()));
453
454 let p_recovered = c2.decrypt(&sk2);
455 let m_recovered = GLWE::<Rq<P, N>, K>::decode::<T>(&p_recovered);
456
457 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
458 }
459
460 Ok(())
461 }
462
463 #[test]
464 fn test_key_switch() -> Result<()> {
465 const Q: u64 = 2u64.pow(16) + 1;
466 const N: usize = 128;
467 const T: u64 = 2; // plaintext modulus
468 const K: usize = 16;
469 type S = GLWE<Rq<Q, N>, K>;
470
471 let beta: u32 = 2;
472 let l: u32 = 16;
473
474 let mut rng = rand::thread_rng();
475
476 let (sk, pk) = S::new_key(&mut rng)?;
477 let (sk2, _) = S::new_key(&mut rng)?;
478 // ksk to switch from sk to sk2
479 let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?;
480
481 let msg_dist = Uniform::new(0_u64, T);
482 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
483 let p = S::encode::<T>(&m); // plaintext
484 //
485 let c = S::encrypt_s(&mut rng, &sk, &p)?;
486
487 let c2 = c.key_switch(beta, l, &ksk);
488
489 // decrypt with the 2nd secret key
490 let p_recovered = c2.decrypt(&sk2);
491 let m_recovered = S::decode::<T>(&p_recovered);
492 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
493
494 // do the same but now encrypting with pk
495 let c = S::encrypt(&mut rng, &pk, &p)?;
496 let c2 = c.key_switch(beta, l, &ksk);
497 let p_recovered = c2.decrypt(&sk2);
498 let m_recovered = S::decode::<T>(&p_recovered);
499 assert_eq!(m, m_recovered);
500
501 Ok(())
502 }
503}