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