fhe-study/gfhe/src/glwe.rs
gfhe/src/glwe.rsRaw
1use anyhow::Result;
2use rand::Rng;
3use rand_distr::{Normal, Uniform};
4use std::ops::{Add, Mul};
5
6use arith::{Ring, Rq, TR};
7
8const ERR_SIGMA: f64 = 3.2;
9
10pub struct GLWE<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>, Rq<Q, N>);
11
12#[derive(Clone, Debug)]
13pub struct SecretKey<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>);
14#[derive(Clone, Debug)]
15pub struct PublicKey<const Q: u64, const N: usize, const K: usize>(Rq<Q, N>, TR<Rq<Q, N>, K>);
16
17impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
18 pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<Q, N, K>, PublicKey<Q, N, K>)> {
19 let Xi_key = Uniform::new(0_f64, 2_f64);
20 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
21
22 let s: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
23 let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Uniform::new(0_f64, Q as f64));
24 let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
25
26 let pk: PublicKey<Q, N, K> = PublicKey((&a * &s) + e, a);
27 Ok((SecretKey(s), pk))
28 }
29
30 // TODO delta not as input
31 pub fn encrypt_s<const T: u64>(
32 mut rng: impl Rng,
33 sk: &SecretKey<Q, N, K>,
34 m: &Rq<T, N>,
35 delta: u64,
36 ) -> Result<Self> {
37 let m: Rq<Q, N> = m.remodule::<Q>();
38
39 let Xi_key = Uniform::new(0_f64, 2_f64);
40 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
41
42 let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
43 let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
44
45 let b: Rq<Q, N> = (&a * &sk.0) + m * delta + e;
46 Ok(Self(a, b))
47 }
48 pub fn encrypt<const T: u64>(
49 mut rng: impl Rng,
50 pk: &PublicKey<Q, N, K>,
51 m: &Rq<T, N>,
52 delta: u64,
53 ) -> Result<Self> {
54 let m: Rq<Q, N> = m.remodule::<Q>();
55
56 let Xi_key = Uniform::new(0_f64, 2_f64);
57 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
58
59 let u: Rq<Q, N> = Rq::rand(&mut rng, Xi_key);
60
61 let e0 = Rq::<Q, N>::rand(&mut rng, Xi_err);
62 let e1 = TR::<Rq<Q, N>, K>::rand(&mut rng, Xi_err);
63
64 let b: Rq<Q, N> = pk.0 * u + m * delta + e0;
65 let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1;
66
67 Ok(Self(d, b))
68 }
69 pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<T, N> {
70 let (d, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
71 let r: Rq<Q, N> = b - &d * &sk.0;
72 let r = r.mul_div_round(T, Q);
73 // let r_scaled: Vec<f64> = r
74 // .coeffs()
75 // .iter()
76 // // .map(|e| (e.0 as f64 / delta as f64).round())
77 // .map(|e| e.mul_div_round(T, Q))
78 // .collect();
79 // let r = Rq::<Q, N>::from_vec_f64(r_scaled);
80 r.remodule::<T>()
81 }
82
83 pub fn mod_switch<const P: u64>(&self) -> GLWE<P, N, K> {
84 let a: TR<Rq<P, N>, K> = TR(self.0 .0.iter().map(|r| r.mod_switch::<P>()).collect());
85 let b: Rq<P, N> = self.1.mod_switch::<P>();
86 GLWE(a, b)
87 }
88}
89
90impl<const Q: u64, const N: usize, const K: usize> Add<GLWE<Q, N, K>> for GLWE<Q, N, K> {
91 type Output = Self;
92 fn add(self, other: Self) -> Self {
93 let a: TR<Rq<Q, N>, K> = self.0 + other.0;
94 let b: Rq<Q, N> = self.1 + other.1;
95 Self(a, b)
96 }
97}
98
99impl<const Q: u64, const N: usize, const K: usize> Add<Rq<Q, N>> for GLWE<Q, N, K> {
100 type Output = Self;
101 fn add(self, plaintext: Rq<Q, N>) -> Self {
102 let a: TR<Rq<Q, N>, K> = self.0;
103 let b: Rq<Q, N> = self.1 + plaintext;
104 Self(a, b)
105 }
106}
107impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N, K> {
108 type Output = Self;
109 fn mul(self, plaintext: Rq<Q, N>) -> Self {
110 // first compute the NTT for plaintext, to avoid computing it at each
111 // iteration, speeding up the multiplications
112 let mut plaintext = plaintext.clone();
113 plaintext.compute_evals();
114
115 let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
116 let b: Rq<Q, N> = self.1 * plaintext;
117 Self(a, b)
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use anyhow::Result;
124 use rand::distributions::Uniform;
125
126 use super::*;
127
128 #[test]
129 fn test_encrypt_decrypt() -> Result<()> {
130 const Q: u64 = 2u64.pow(16) + 1;
131 const N: usize = 128;
132 const T: u64 = 32; // plaintext modulus
133 const K: usize = 16;
134 type S = GLWE<Q, N, K>;
135
136 let delta: u64 = Q / T; // floored
137 let mut rng = rand::thread_rng();
138
139 for _ in 0..200 {
140 let (sk, pk) = S::new_key(&mut rng)?;
141
142 let msg_dist = Uniform::new(0_u64, T);
143 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
144
145 let c = S::encrypt(&mut rng, &pk, &m, delta)?;
146 let m_recovered = c.decrypt(&sk, delta);
147
148 assert_eq!(m, m_recovered);
149 }
150
151 Ok(())
152 }
153
154 #[test]
155 fn test_addition() -> Result<()> {
156 const Q: u64 = 2u64.pow(16) + 1;
157 const N: usize = 128;
158 const T: u64 = 20;
159 const K: usize = 16;
160 type S = GLWE<Q, N, K>;
161
162 let delta: u64 = Q / T; // floored
163 let mut rng = rand::thread_rng();
164
165 for _ in 0..200 {
166 let (sk, pk) = S::new_key(&mut rng)?;
167
168 let msg_dist = Uniform::new(0_u64, T);
169 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
170 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
171
172 let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
173 let c2 = S::encrypt(&mut rng, &pk, &m2, delta)?;
174
175 let c3 = c1 + c2;
176
177 let m3_recovered = c3.decrypt(&sk, delta);
178
179 assert_eq!(m1 + m2, m3_recovered);
180 }
181
182 Ok(())
183 }
184
185 #[test]
186 fn test_add_plaintext() -> Result<()> {
187 const Q: u64 = 2u64.pow(16) + 1;
188 const N: usize = 128;
189 const T: u64 = 32;
190 const K: usize = 16;
191 type S = GLWE<Q, N, K>;
192
193 let delta: u64 = Q / T; // floored
194 let mut rng = rand::thread_rng();
195
196 for _ in 0..200 {
197 let (sk, pk) = S::new_key(&mut rng)?;
198
199 let msg_dist = Uniform::new(0_u64, T);
200 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
201 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
202 let m2_scaled: Rq<Q, N> = m2.remodule::<Q>() * delta;
203
204 let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
205
206 let c3 = c1 + m2_scaled;
207
208 let m3_recovered = c3.decrypt(&sk, delta);
209
210 assert_eq!(m1 + m2, m3_recovered);
211 }
212
213 Ok(())
214 }
215
216 #[test]
217 fn test_mul_plaintext() -> Result<()> {
218 const Q: u64 = 2u64.pow(16) + 1;
219 const N: usize = 16;
220 const T: u64 = 4;
221 const K: usize = 16;
222 type S = GLWE<Q, N, K>;
223
224 let delta: u64 = Q / T; // floored
225 let mut rng = rand::thread_rng();
226
227 for _ in 0..200 {
228 let (sk, pk) = S::new_key(&mut rng)?;
229
230 let msg_dist = Uniform::new(0_u64, T);
231 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
232 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
233 let m2: Rq<Q, N> = m2.remodule::<Q>();
234 let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?;
235
236 let c3 = c1 * m2;
237
238 let m3_recovered: Rq<T, N> = c3.decrypt(&sk, delta);
239 assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
240 }
241
242 Ok(())
243 }
244
245 #[test]
246 fn test_mod_switch() -> Result<()> {
247 const Q: u64 = 2u64.pow(16) + 1;
248 const P: u64 = 2u64.pow(8) + 1;
249 // note: wip, Q and P chosen so that P/Q is an integer
250 const N: usize = 8;
251 const T: u64 = 8; // plaintext modulus, must be a prime or power of a prime
252 const K: usize = 16;
253 type S = GLWE<Q, N, K>;
254
255 let delta: u64 = Q / T; // floored
256 let mut rng = rand::thread_rng();
257
258 dbg!(P as f64 / Q as f64);
259 dbg!(delta);
260 dbg!(delta as f64 * P as f64 / Q as f64);
261 dbg!(delta as f64 * (P as f64 / Q as f64));
262
263 for _ in 0..200 {
264 let (sk, pk) = S::new_key(&mut rng)?;
265
266 let msg_dist = Uniform::new(0_u64, T);
267 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
268
269 let c = S::encrypt(&mut rng, &pk, &m, delta)?;
270
271 let c2 = c.mod_switch::<P>();
272 let sk2: SecretKey<P, N, K> =
273 SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect()));
274 let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64;
275
276 let m_recovered = c2.decrypt(&sk2, delta2);
277
278 assert_eq!(m, m_recovered);
279 }
280
281 Ok(())
282 }
283}