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