fhe-study/ckks/src/lib.rs
ckks/src/lib.rsRaw
1//! Implementation of CKKS https://eprint.iacr.org/2016/421.pdf
2#![allow(non_snake_case)]
3#![allow(non_upper_case_globals)]
4#![allow(non_camel_case_types)]
5#![allow(clippy::upper_case_acronyms)]
6#![allow(dead_code)] // TMP
7
8use arith::{Rq, C, R};
9
10use anyhow::Result;
11use rand::Rng;
12use rand_distr::{Normal, Uniform};
13
14pub mod encoder;
15pub use encoder::Encoder;
16
17// error deviation for the Gaussian(Normal) distribution
18// sigma=3.2 from: https://eprint.iacr.org/2016/421.pdf page 17
19const ERR_SIGMA: f64 = 3.2;
20
21#[derive(Clone, Copy, Debug)]
22pub struct Params {
23 q: u64,
24 n: usize,
25 t: u64,
26}
27
28#[derive(Debug)]
29pub struct PublicKey(Rq, Rq);
30
31pub struct SecretKey(Rq);
32
33pub struct CKKS {
34 params: Params,
35 encoder: Encoder,
36}
37
38impl CKKS {
39 pub fn new(params: &Params, delta: C<f64>) -> Self {
40 let encoder = Encoder::new(params.n, delta);
41 Self {
42 params: params.clone(),
43 encoder,
44 }
45 }
46 /// generate a new key pair (privK, pubK)
47 pub fn new_key(&self, mut rng: impl Rng) -> Result<(SecretKey, PublicKey)> {
48 let params = &self.params;
49
50 let Xi_key = Uniform::new(-1_f64, 1_f64);
51 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
52
53 let e = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
54
55 let mut s = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
56 // since s is going to be multiplied by other Rq elements, already
57 // compute its NTT
58 s.compute_evals();
59
60 let a = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
61
62 let pk: PublicKey = PublicKey((&(-a.clone()) * &s) + e, a.clone()); // TODO rm clones
63 Ok((SecretKey(s), pk))
64 }
65
66 // encrypts a plaintext \in R=Z_Q[X]/(X^N+1)
67 fn encrypt(
68 &self, // TODO maybe rm?
69 mut rng: impl Rng,
70 pk: &PublicKey,
71 m: &R,
72 ) -> Result<(Rq, Rq)> {
73 let params = self.params;
74 let Xi_key = Uniform::new(-1_f64, 1_f64);
75 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
76
77 let e_0 = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
78 let e_1 = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
79
80 let v = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
81
82 // let m: Rq = Rq::from(*m);
83 let m: Rq = m.clone().to_rq(params.q); // TODO rm clone
84
85 Ok((m + e_0 + &v * &pk.0.clone(), &v * &pk.1 + e_1))
86 }
87
88 fn decrypt(
89 &self, // TODO maybe rm?
90 sk: &SecretKey,
91 c: (Rq, Rq),
92 ) -> Result<R> {
93 let m = c.0.clone() + &c.1 * &sk.0;
94 Ok(m.mod_centered_q())
95 }
96
97 pub fn encode_and_encrypt(
98 &self,
99 mut rng: impl Rng,
100 pk: &PublicKey,
101 z: &[C<f64>],
102 ) -> Result<(Rq, Rq)> {
103 let m: R = self.encoder.encode(&z)?; // polynomial (encoded vec) \in R
104
105 self.encrypt(&mut rng, pk, &m)
106 }
107
108 pub fn decrypt_and_decode(&self, sk: SecretKey, c: (Rq, Rq)) -> Result<Vec<C<f64>>> {
109 let d = self.decrypt(&sk, c)?;
110
111 self.encoder.decode(&d)
112 }
113
114 pub fn add(&self, c0: &(Rq, Rq), c1: &(Rq, Rq)) -> Result<(Rq, Rq)> {
115 Ok((&c0.0 + &c1.0, &c0.1 + &c1.1))
116 }
117 pub fn sub(&self, c0: &(Rq, Rq), c1: &(Rq, Rq)) -> Result<(Rq, Rq)> {
118 Ok((&c0.0 - &c1.0, &c0.1 + &c1.1))
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125
126 #[test]
127 fn test_encrypt_decrypt() -> Result<()> {
128 let q: u64 = 2u64.pow(16) + 1;
129 let n: usize = 32;
130 let t: u64 = 50;
131 let params = Params { q, n, t };
132 let scale_factor_u64 = 512_u64; // delta
133 let scale_factor = C::<f64>::new(scale_factor_u64 as f64, 0.0); // delta
134
135 let mut rng = rand::thread_rng();
136
137 for _ in 0..1000 {
138 let ckks = CKKS::new(¶ms, scale_factor);
139
140 let (sk, pk) = ckks.new_key(&mut rng)?;
141
142 let m_raw: R = Rq::rand_f64(&mut rng, Uniform::new(0_f64, t as f64), q, n)?.to_r();
143 let m = &m_raw * &scale_factor_u64;
144
145 let ct = ckks.encrypt(&mut rng, &pk, &m)?;
146 let m_decrypted = ckks.decrypt(&sk, ct)?;
147
148 let m_decrypted: Vec<u64> = m_decrypted
149 .coeffs()
150 .iter()
151 .map(|e| (*e as f64 / (scale_factor_u64 as f64)).round() as u64)
152 .collect();
153 let m_decrypted = Rq::from_vec_u64(q, n, m_decrypted);
154 // assert_eq!(m_decrypted, Rq::from(m_raw));
155 assert_eq!(m_decrypted, m_raw.to_rq(q));
156 }
157
158 Ok(())
159 }
160
161 #[test]
162 fn test_encode_encrypt_decrypt_decode() -> Result<()> {
163 let q: u64 = 2u64.pow(16) + 1;
164 let n: usize = 16;
165 let t: u64 = 8;
166 let params = Params { q, n, t };
167 let scale_factor = C::<f64>::new(512.0, 0.0); // delta
168
169 let mut rng = rand::thread_rng();
170
171 for _ in 0..1000 {
172 let ckks = CKKS::new(¶ms, scale_factor);
173 let (sk, pk) = ckks.new_key(&mut rng)?;
174
175 let z: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
176 .take(n / 2)
177 .collect();
178 let m: R = ckks.encoder.encode(&z)?;
179 println!("{}", m);
180
181 // sanity check
182 {
183 let z_decoded = ckks.encoder.decode(&m)?;
184 let rounded_z_decoded: Vec<C<f64>> = z_decoded
185 .iter()
186 .map(|c| C::<f64>::new(c.re.round(), c.im.round()))
187 .collect();
188 assert_eq!(rounded_z_decoded, z);
189 }
190
191 let ct = ckks.encrypt(&mut rng, &pk, &m)?;
192 let m_decrypted = ckks.decrypt(&sk, ct)?;
193 println!("{}", m_decrypted);
194
195 let z_decrypted = ckks.encoder.decode(&m_decrypted)?;
196
197 let rounded_z_decrypted: Vec<C<f64>> = z_decrypted
198 .iter()
199 .map(|&c| C::<f64>::new(c.re.round(), c.im.round()))
200 .collect();
201 assert_eq!(rounded_z_decrypted, z);
202 }
203
204 Ok(())
205 }
206
207 #[test]
208 fn test_add() -> Result<()> {
209 let q: u64 = 2u64.pow(16) + 1;
210 let n: usize = 16;
211 let t: u64 = 8;
212 let params = Params { q, n, t };
213 let scale_factor = C::<f64>::new(1024.0, 0.0); // delta
214
215 let mut rng = rand::thread_rng();
216
217 for _ in 0..1000 {
218 let ckks = CKKS::new(¶ms, scale_factor);
219
220 let (sk, pk) = ckks.new_key(&mut rng)?;
221
222 let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
223 .take(n / 2)
224 .collect();
225 let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
226 .take(n / 2)
227 .collect();
228 let m0: R = ckks.encoder.encode(&z0)?;
229 let m1: R = ckks.encoder.encode(&z1)?;
230
231 let ct0 = ckks.encrypt(&mut rng, &pk, &m0)?;
232 let ct1 = ckks.encrypt(&mut rng, &pk, &m1)?;
233
234 let ct2 = ckks.add(&ct0, &ct1)?;
235
236 let m2_decrypted = ckks.decrypt(&sk, ct2)?;
237
238 let z_decrypted = ckks.encoder.decode(&m2_decrypted)?;
239 let rounded_z_decrypted: Vec<C<f64>> = z_decrypted
240 .iter()
241 .map(|&c| C::<f64>::new(c.re.round(), c.im.round()))
242 .collect();
243
244 let expected_z2: Vec<C<f64>> = itertools::zip_eq(z0, z1).map(|(a, b)| a + b).collect();
245 assert_eq!(rounded_z_decrypted, expected_z2);
246 }
247
248 Ok(())
249 }
250
251 #[test]
252 fn test_sub() -> Result<()> {
253 let q: u64 = 2u64.pow(16) + 1;
254 let n: usize = 16;
255 let t: u64 = 8;
256 let params = Params { q, n, t };
257 let scale_factor = C::<f64>::new(1024.0, 0.0); // delta
258
259 let mut rng = rand::thread_rng();
260
261 for _ in 0..1000 {
262 let ckks = CKKS::new(¶ms, scale_factor);
263
264 let (sk, pk) = ckks.new_key(&mut rng)?;
265
266 let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
267 .take(n / 2)
268 .collect();
269 let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
270 .take(n / 2)
271 .collect();
272 let m0: R = ckks.encoder.encode(&z0)?;
273 let m1: R = ckks.encoder.encode(&z1)?;
274
275 let ct0 = ckks.encrypt(&mut rng, &pk, &m0)?;
276 let ct1 = ckks.encrypt(&mut rng, &pk, &m1)?;
277
278 let ct2 = ckks.sub(&ct0, &ct1)?;
279
280 let m2_decrypted = ckks.decrypt(&sk, ct2)?;
281
282 let z_decrypted = ckks.encoder.decode(&m2_decrypted)?;
283 let rounded_z_decrypted: Vec<C<f64>> = z_decrypted
284 .iter()
285 .map(|&c| C::<f64>::new(c.re.round(), c.im.round()))
286 .collect();
287
288 let expected_z2: Vec<C<f64>> = itertools::zip_eq(z0, z1).map(|(a, b)| a - b).collect();
289 assert_eq!(rounded_z_decrypted, expected_z2);
290 }
291
292 Ok(())
293 }
294}