client.rs 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711
  1. use crate::{
  2. arith::*, discrete_gaussian::*, gadget::*, number_theory::*, params::*, poly::*, util::*,
  3. };
  4. use rand::{Rng, SeedableRng};
  5. use std::{iter::once, mem::size_of};
  6. use rand_chacha::ChaCha20Rng;
  7. fn new_vec_raw<'a>(
  8. params: &'a Params,
  9. num: usize,
  10. rows: usize,
  11. cols: usize,
  12. ) -> Vec<PolyMatrixRaw<'a>> {
  13. let mut v = Vec::with_capacity(num);
  14. for _ in 0..num {
  15. v.push(PolyMatrixRaw::zero(params, rows, cols));
  16. }
  17. v
  18. }
  19. fn serialize_polymatrix(vec: &mut Vec<u8>, a: &PolyMatrixRaw) {
  20. for i in 0..a.rows * a.cols * a.params.poly_len {
  21. vec.extend_from_slice(&u64::to_ne_bytes(a.data[i]));
  22. }
  23. }
  24. fn serialize_vec_polymatrix(vec: &mut Vec<u8>, a: &Vec<PolyMatrixRaw>) {
  25. for i in 0..a.len() {
  26. serialize_polymatrix(vec, &a[i]);
  27. }
  28. }
  29. fn mat_sz_bytes(a: &PolyMatrixRaw) -> usize {
  30. a.rows * a.cols * a.params.poly_len * size_of::<u64>()
  31. }
  32. fn deserialize_polymatrix(a: &mut PolyMatrixRaw, data: &[u8]) -> usize {
  33. for (i, chunk) in data.chunks(size_of::<u64>()).enumerate() {
  34. a.data[i] = u64::from_ne_bytes(chunk.try_into().unwrap());
  35. }
  36. mat_sz_bytes(a)
  37. }
  38. fn deserialize_vec_polymatrix(a: &mut Vec<PolyMatrixRaw>, data: &[u8]) -> usize {
  39. let mut chunks = data.chunks(mat_sz_bytes(&a[0]));
  40. let mut bytes_read = 0;
  41. for i in 0..a.len() {
  42. bytes_read += deserialize_polymatrix(&mut a[i], chunks.next().unwrap());
  43. }
  44. bytes_read
  45. }
  46. pub struct PublicParameters<'a> {
  47. pub v_packing: Vec<PolyMatrixNTT<'a>>, // Ws
  48. pub v_expansion_left: Option<Vec<PolyMatrixNTT<'a>>>,
  49. pub v_expansion_right: Option<Vec<PolyMatrixNTT<'a>>>,
  50. pub v_conversion: Option<Vec<PolyMatrixNTT<'a>>>, // V
  51. }
  52. impl<'a> PublicParameters<'a> {
  53. pub fn init(params: &'a Params) -> Self {
  54. if params.expand_queries {
  55. PublicParameters {
  56. v_packing: Vec::new(),
  57. v_expansion_left: Some(Vec::new()),
  58. v_expansion_right: Some(Vec::new()),
  59. v_conversion: Some(Vec::new()),
  60. }
  61. } else {
  62. PublicParameters {
  63. v_packing: Vec::new(),
  64. v_expansion_left: None,
  65. v_expansion_right: None,
  66. v_conversion: None,
  67. }
  68. }
  69. }
  70. fn from_ntt_alloc_vec(v: &Vec<PolyMatrixNTT<'a>>) -> Option<Vec<PolyMatrixRaw<'a>>> {
  71. Some(v.iter().map(from_ntt_alloc).collect())
  72. }
  73. fn from_ntt_alloc_opt_vec(
  74. v: &Option<Vec<PolyMatrixNTT<'a>>>,
  75. ) -> Option<Vec<PolyMatrixRaw<'a>>> {
  76. Some(v.as_ref()?.iter().map(from_ntt_alloc).collect())
  77. }
  78. fn to_ntt_alloc_vec(v: &Vec<PolyMatrixRaw<'a>>) -> Option<Vec<PolyMatrixNTT<'a>>> {
  79. Some(v.iter().map(to_ntt_alloc).collect())
  80. }
  81. pub fn to_raw(&self) -> Vec<Option<Vec<PolyMatrixRaw>>> {
  82. vec![
  83. Self::from_ntt_alloc_vec(&self.v_packing),
  84. Self::from_ntt_alloc_opt_vec(&self.v_expansion_left),
  85. Self::from_ntt_alloc_opt_vec(&self.v_expansion_right),
  86. Self::from_ntt_alloc_opt_vec(&self.v_conversion),
  87. ]
  88. }
  89. pub fn serialize(&self) -> Vec<u8> {
  90. let mut data = Vec::new();
  91. for v in self.to_raw().iter() {
  92. if v.is_some() {
  93. serialize_vec_polymatrix(&mut data, v.as_ref().unwrap());
  94. }
  95. }
  96. data
  97. }
  98. pub fn deserialize(params: &'a Params, data: &[u8]) -> Self {
  99. assert_eq!(params.setup_bytes(), data.len());
  100. let mut idx = 0;
  101. let mut v_packing = new_vec_raw(params, params.n, params.n + 1, params.t_conv);
  102. idx += deserialize_vec_polymatrix(&mut v_packing, &data[idx..]);
  103. if params.expand_queries {
  104. let mut v_expansion_left = new_vec_raw(params, params.g(), 2, params.t_exp_left);
  105. idx += deserialize_vec_polymatrix(&mut v_expansion_left, &data[idx..]);
  106. let mut v_expansion_right =
  107. new_vec_raw(params, params.stop_round() + 1, 2, params.t_exp_right);
  108. idx += deserialize_vec_polymatrix(&mut v_expansion_right, &data[idx..]);
  109. let mut v_conversion = new_vec_raw(params, 1, 2, 2 * params.t_conv);
  110. _ = deserialize_vec_polymatrix(&mut v_conversion, &data[idx..]);
  111. Self {
  112. v_packing: Self::to_ntt_alloc_vec(&v_packing).unwrap(),
  113. v_expansion_left: Self::to_ntt_alloc_vec(&v_expansion_left),
  114. v_expansion_right: Self::to_ntt_alloc_vec(&v_expansion_right),
  115. v_conversion: Self::to_ntt_alloc_vec(&v_conversion),
  116. }
  117. } else {
  118. Self {
  119. v_packing: Self::to_ntt_alloc_vec(&v_packing).unwrap(),
  120. v_expansion_left: None,
  121. v_expansion_right: None,
  122. v_conversion: None,
  123. }
  124. }
  125. }
  126. }
  127. pub struct Query<'a> {
  128. pub ct: Option<PolyMatrixRaw<'a>>,
  129. pub v_buf: Option<Vec<u64>>,
  130. pub v_ct: Option<Vec<PolyMatrixRaw<'a>>>,
  131. }
  132. impl<'a> Query<'a> {
  133. pub fn empty() -> Self {
  134. Query {
  135. ct: None,
  136. v_ct: None,
  137. v_buf: None,
  138. }
  139. }
  140. pub fn serialize(&self) -> Vec<u8> {
  141. let mut data = Vec::new();
  142. if self.ct.is_some() {
  143. let ct = self.ct.as_ref().unwrap();
  144. serialize_polymatrix(&mut data, &ct);
  145. }
  146. if self.v_buf.is_some() {
  147. let v_buf = self.v_buf.as_ref().unwrap();
  148. data.extend(v_buf.iter().map(|x| x.to_ne_bytes()).flatten());
  149. }
  150. if self.v_ct.is_some() {
  151. let v_ct = self.v_ct.as_ref().unwrap();
  152. for x in v_ct {
  153. serialize_polymatrix(&mut data, x);
  154. }
  155. }
  156. data
  157. }
  158. pub fn deserialize(params: &'a Params, data: &[u8]) -> Self {
  159. assert_eq!(params.query_bytes(), data.len());
  160. let mut out = Query::empty();
  161. if params.expand_queries {
  162. let mut ct = PolyMatrixRaw::zero(params, 2, 1);
  163. deserialize_polymatrix(&mut ct, data);
  164. out.ct = Some(ct);
  165. } else {
  166. let v_buf_bytes = params.query_v_buf_bytes();
  167. let v_buf = (&data[..v_buf_bytes])
  168. .chunks(size_of::<u64>())
  169. .map(|x| u64::from_ne_bytes(x.try_into().unwrap()))
  170. .collect();
  171. out.v_buf = Some(v_buf);
  172. let mut v_ct = new_vec_raw(params, params.db_dim_2, 2, 2 * params.t_gsw);
  173. deserialize_vec_polymatrix(&mut v_ct, &data[v_buf_bytes..]);
  174. out.v_ct = Some(v_ct);
  175. }
  176. out
  177. }
  178. }
  179. pub struct Client<'a, T: Rng> {
  180. params: &'a Params,
  181. sk_gsw: PolyMatrixRaw<'a>,
  182. sk_reg: PolyMatrixRaw<'a>,
  183. sk_gsw_full: PolyMatrixRaw<'a>,
  184. sk_reg_full: PolyMatrixRaw<'a>,
  185. dg: DiscreteGaussian<'a, T>,
  186. public_rng: ChaCha20Rng,
  187. public_seed: <ChaCha20Rng as SeedableRng>::Seed,
  188. }
  189. fn matrix_with_identity<'a>(p: &PolyMatrixRaw<'a>) -> PolyMatrixRaw<'a> {
  190. assert_eq!(p.cols, 1);
  191. let mut r = PolyMatrixRaw::zero(p.params, p.rows, p.rows + 1);
  192. r.copy_into(p, 0, 0);
  193. r.copy_into(&PolyMatrixRaw::identity(p.params, p.rows, p.rows), 0, 1);
  194. r
  195. }
  196. fn params_with_moduli(params: &Params, moduli: &Vec<u64>) -> Params {
  197. Params::init(
  198. params.poly_len,
  199. moduli,
  200. params.noise_width,
  201. params.n,
  202. params.pt_modulus,
  203. params.q2_bits,
  204. params.t_conv,
  205. params.t_exp_left,
  206. params.t_exp_right,
  207. params.t_gsw,
  208. params.expand_queries,
  209. params.db_dim_1,
  210. params.db_dim_2,
  211. params.instances,
  212. params.db_item_size,
  213. )
  214. }
  215. impl<'a, T: Rng> Client<'a, T> {
  216. pub fn init(params: &'a Params, rng: &'a mut T) -> Self {
  217. let sk_gsw_dims = params.get_sk_gsw();
  218. let sk_reg_dims = params.get_sk_reg();
  219. let sk_gsw = PolyMatrixRaw::zero(params, sk_gsw_dims.0, sk_gsw_dims.1);
  220. let sk_reg = PolyMatrixRaw::zero(params, sk_reg_dims.0, sk_reg_dims.1);
  221. let sk_gsw_full = matrix_with_identity(&sk_gsw);
  222. let sk_reg_full = matrix_with_identity(&sk_reg);
  223. let mut public_seed = [0u8; 32];
  224. rng.fill_bytes(&mut public_seed);
  225. let public_rng = ChaCha20Rng::from_seed(public_seed);
  226. let dg = DiscreteGaussian::init(params, rng);
  227. Self {
  228. params,
  229. sk_gsw,
  230. sk_reg,
  231. sk_gsw_full,
  232. sk_reg_full,
  233. dg,
  234. public_rng,
  235. public_seed
  236. }
  237. }
  238. #[allow(dead_code)]
  239. pub(crate) fn get_sk_reg(&self) -> &PolyMatrixRaw<'a> {
  240. &self.sk_reg
  241. }
  242. pub fn get_rng(&mut self) -> &mut T {
  243. &mut self.dg.rng
  244. }
  245. pub fn get_public_rng(&mut self) -> &mut ChaCha20Rng {
  246. &mut self.public_rng
  247. }
  248. pub fn get_public_seed(&mut self) -> <ChaCha20Rng as SeedableRng>::Seed {
  249. self.public_seed
  250. }
  251. fn get_fresh_gsw_public_key(&mut self, m: usize) -> PolyMatrixRaw<'a> {
  252. let params = self.params;
  253. let n = params.n;
  254. let a = PolyMatrixRaw::random_rng(params, 1, m, self.get_rng());
  255. let e = PolyMatrixRaw::noise(params, n, m, &mut self.dg);
  256. let a_inv = -&a;
  257. let b_p = &self.sk_gsw.ntt() * &a.ntt();
  258. let b = &e.ntt() + &b_p;
  259. let p = stack(&a_inv, &b.raw());
  260. p
  261. }
  262. fn get_regev_sample(&mut self) -> PolyMatrixNTT<'a> {
  263. let params = self.params;
  264. let a = PolyMatrixRaw::random_rng(params, 1, 1, self.get_public_rng());
  265. let e = PolyMatrixRaw::noise(params, 1, 1, &mut self.dg);
  266. let b_p = &self.sk_reg.ntt() * &a.ntt();
  267. let b = &e.ntt() + &b_p;
  268. let mut p = PolyMatrixNTT::zero(params, 2, 1);
  269. p.copy_into(&(-&a).ntt(), 0, 0);
  270. p.copy_into(&b, 1, 0);
  271. p
  272. }
  273. fn get_fresh_reg_public_key(&mut self, m: usize) -> PolyMatrixNTT<'a> {
  274. let params = self.params;
  275. let mut p = PolyMatrixNTT::zero(params, 2, m);
  276. for i in 0..m {
  277. p.copy_into(&self.get_regev_sample(), 0, i);
  278. }
  279. p
  280. }
  281. fn encrypt_matrix_gsw(&mut self, ag: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
  282. let mx = ag.cols;
  283. let p = self.get_fresh_gsw_public_key(mx);
  284. let res = &(p.ntt()) + &(ag.pad_top(1));
  285. res
  286. }
  287. pub fn encrypt_matrix_reg(&mut self, a: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
  288. let m = a.cols;
  289. let p = self.get_fresh_reg_public_key(m);
  290. &p + &a.pad_top(1)
  291. }
  292. pub fn decrypt_matrix_reg(&mut self, a: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
  293. &self.sk_reg_full.ntt() * a
  294. }
  295. pub fn decrypt_matrix_gsw(&mut self, a: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
  296. &self.sk_gsw_full.ntt() * a
  297. }
  298. fn generate_expansion_params(
  299. &mut self,
  300. num_exp: usize,
  301. m_exp: usize,
  302. ) -> Vec<PolyMatrixNTT<'a>> {
  303. let params = self.params;
  304. let g_exp = build_gadget(params, 1, m_exp);
  305. let g_exp_ntt = g_exp.ntt();
  306. let mut res = Vec::new();
  307. for i in 0..num_exp {
  308. let t = (params.poly_len / (1 << i)) + 1;
  309. let tau_sk_reg = automorph_alloc(&self.sk_reg, t);
  310. let prod = &tau_sk_reg.ntt() * &g_exp_ntt;
  311. let w_exp_i = self.encrypt_matrix_reg(&prod);
  312. res.push(w_exp_i);
  313. }
  314. res
  315. }
  316. pub fn generate_keys(&mut self) -> PublicParameters<'a> {
  317. let params = self.params;
  318. self.dg.sample_matrix(&mut self.sk_gsw);
  319. self.dg.sample_matrix(&mut self.sk_reg);
  320. self.sk_gsw_full = matrix_with_identity(&self.sk_gsw);
  321. self.sk_reg_full = matrix_with_identity(&self.sk_reg);
  322. let sk_reg_ntt = to_ntt_alloc(&self.sk_reg);
  323. let mut pp = PublicParameters::init(params);
  324. // Params for packing
  325. let gadget_conv = build_gadget(params, 1, params.t_conv);
  326. let gadget_conv_ntt = to_ntt_alloc(&gadget_conv);
  327. for i in 0..params.n {
  328. let scaled = scalar_multiply_alloc(&sk_reg_ntt, &gadget_conv_ntt);
  329. let mut ag = PolyMatrixNTT::zero(params, params.n, params.t_conv);
  330. ag.copy_into(&scaled, i, 0);
  331. let w = self.encrypt_matrix_gsw(&ag);
  332. pp.v_packing.push(w);
  333. }
  334. if params.expand_queries {
  335. // Params for expansion
  336. pp.v_expansion_left = Some(self.generate_expansion_params(params.g(), params.t_exp_left));
  337. pp.v_expansion_right =
  338. Some(self.generate_expansion_params(params.stop_round() + 1, params.t_exp_right));
  339. // Params for converison
  340. let g_conv = build_gadget(params, 2, 2 * params.t_conv);
  341. let sk_reg_ntt = self.sk_reg.ntt();
  342. let sk_reg_squared_ntt = &sk_reg_ntt * &sk_reg_ntt;
  343. pp.v_conversion = Some(Vec::from_iter(once(PolyMatrixNTT::zero(
  344. params,
  345. 2,
  346. 2 * params.t_conv,
  347. ))));
  348. for i in 0..2 * params.t_conv {
  349. let sigma;
  350. if i % 2 == 0 {
  351. let val = g_conv.get_poly(0, i)[0];
  352. sigma = &sk_reg_squared_ntt * &single_poly(params, val).ntt();
  353. } else {
  354. let val = g_conv.get_poly(1, i)[0];
  355. sigma = &sk_reg_ntt * &single_poly(params, val).ntt();
  356. }
  357. let ct = self.encrypt_matrix_reg(&sigma);
  358. pp.v_conversion.as_mut().unwrap()[0].copy_into(&ct, 0, i);
  359. }
  360. }
  361. pp
  362. }
  363. pub fn generate_query(&mut self, idx_target: usize) -> Query<'a> {
  364. let params = self.params;
  365. let further_dims = params.db_dim_2;
  366. let idx_dim0 = idx_target / (1 << further_dims);
  367. let idx_further = idx_target % (1 << further_dims);
  368. let scale_k = params.modulus / params.pt_modulus;
  369. let bits_per = get_bits_per(params, params.t_gsw);
  370. let mut query = Query::empty();
  371. if params.expand_queries {
  372. // pack query into single ciphertext
  373. let mut sigma = PolyMatrixRaw::zero(params, 1, 1);
  374. sigma.data[2 * idx_dim0] = scale_k;
  375. for i in 0..further_dims as u64 {
  376. let bit: u64 = ((idx_further as u64) & (1 << i)) >> i;
  377. for j in 0..params.t_gsw {
  378. let val = (1u64 << (bits_per * j)) * bit;
  379. let idx = (i as usize) * params.t_gsw + (j as usize);
  380. sigma.data[2 * idx + 1] = val;
  381. }
  382. }
  383. let inv_2_g_first = invert_uint_mod(1 << params.g(), params.modulus).unwrap();
  384. let inv_2_g_rest = invert_uint_mod(1 << (params.stop_round() + 1), params.modulus).unwrap();
  385. for i in 0..params.poly_len / 2 {
  386. sigma.data[2 * i] =
  387. multiply_uint_mod(sigma.data[2 * i], inv_2_g_first, params.modulus);
  388. sigma.data[2 * i + 1] =
  389. multiply_uint_mod(sigma.data[2 * i + 1], inv_2_g_rest, params.modulus);
  390. }
  391. query.ct = Some(from_ntt_alloc(
  392. &self.encrypt_matrix_reg(&to_ntt_alloc(&sigma)),
  393. ));
  394. } else {
  395. let num_expanded = 1 << params.db_dim_1;
  396. let mut sigma_v = Vec::<PolyMatrixNTT>::new();
  397. // generate regev ciphertexts
  398. let reg_cts_buf_words = num_expanded * 2 * params.poly_len;
  399. let mut reg_cts_buf = vec![0u64; reg_cts_buf_words];
  400. let mut reg_cts = Vec::<PolyMatrixNTT>::new();
  401. for i in 0..num_expanded {
  402. let value = ((i == idx_dim0) as u64) * scale_k;
  403. let sigma = PolyMatrixRaw::single_value(&params, value);
  404. reg_cts.push(self.encrypt_matrix_reg(&to_ntt_alloc(&sigma)));
  405. }
  406. // reorient into server's preferred indexing
  407. reorient_reg_ciphertexts(self.params, reg_cts_buf.as_mut_slice(), &reg_cts);
  408. // generate GSW ciphertexts
  409. for i in 0..further_dims {
  410. let bit = ((idx_further as u64) & (1 << (i as u64))) >> (i as u64);
  411. let mut ct_gsw = PolyMatrixNTT::zero(&params, 2, 2 * params.t_gsw);
  412. for j in 0..params.t_gsw {
  413. let value = (1u64 << (bits_per * j)) * bit;
  414. let sigma = PolyMatrixRaw::single_value(&params, value);
  415. let sigma_ntt = to_ntt_alloc(&sigma);
  416. let ct = &self.encrypt_matrix_reg(&sigma_ntt);
  417. ct_gsw.copy_into(ct, 0, 2 * j + 1);
  418. let prod = &to_ntt_alloc(&self.sk_reg) * &sigma_ntt;
  419. let ct = &self.encrypt_matrix_reg(&prod);
  420. ct_gsw.copy_into(ct, 0, 2 * j);
  421. }
  422. sigma_v.push(ct_gsw);
  423. }
  424. query.v_buf = Some(reg_cts_buf);
  425. query.v_ct = Some(sigma_v.iter().map(|x| from_ntt_alloc(x)).collect());
  426. }
  427. query
  428. }
  429. pub fn decode_response(&self, data: &[u8]) -> Vec<u8> {
  430. /*
  431. 0. NTT over q2 the secret key
  432. 1. read first row in q2_bit chunks
  433. 2. read rest in q1_bit chunks
  434. 3. NTT over q2 the first row
  435. 4. Multiply the results of (0) and (3)
  436. 5. Divide and round correctly
  437. */
  438. let params = self.params;
  439. let p = params.pt_modulus;
  440. let p_bits = log2_ceil(params.pt_modulus);
  441. let q1 = 4 * params.pt_modulus;
  442. let q1_bits = log2_ceil(q1) as usize;
  443. let q2 = Q2_VALUES[params.q2_bits as usize];
  444. let q2_bits = params.q2_bits as usize;
  445. let q2_params = params_with_moduli(params, &vec![q2]);
  446. // this only needs to be done during keygen
  447. let mut sk_gsw_q2 = PolyMatrixRaw::zero(&q2_params, params.n, 1);
  448. for i in 0..params.poly_len * params.n {
  449. sk_gsw_q2.data[i] = recenter(self.sk_gsw.data[i], params.modulus, q2);
  450. }
  451. let mut sk_gsw_q2_ntt = PolyMatrixNTT::zero(&q2_params, params.n, 1);
  452. to_ntt(&mut sk_gsw_q2_ntt, &sk_gsw_q2);
  453. let mut result = PolyMatrixRaw::zero(&params, params.instances * params.n, params.n);
  454. let mut bit_offs = 0;
  455. for instance in 0..params.instances {
  456. // this must be done during decoding
  457. let mut first_row = PolyMatrixRaw::zero(&q2_params, 1, params.n);
  458. let mut rest_rows = PolyMatrixRaw::zero(&params, params.n, params.n);
  459. for i in 0..params.n * params.poly_len {
  460. first_row.data[i] = read_arbitrary_bits(data, bit_offs, q2_bits);
  461. bit_offs += q2_bits;
  462. }
  463. for i in 0..params.n * params.n * params.poly_len {
  464. rest_rows.data[i] = read_arbitrary_bits(data, bit_offs, q1_bits);
  465. bit_offs += q1_bits;
  466. }
  467. let mut first_row_q2 = PolyMatrixNTT::zero(&q2_params, 1, params.n);
  468. to_ntt(&mut first_row_q2, &first_row);
  469. let sk_prod = (&sk_gsw_q2_ntt * &first_row_q2).raw();
  470. let q1_i64 = q1 as i64;
  471. let q2_i64 = q2 as i64;
  472. let p_i128 = p as i128;
  473. for i in 0..params.n * params.n * params.poly_len {
  474. let mut val_first = sk_prod.data[i] as i64;
  475. if val_first >= q2_i64 / 2 {
  476. val_first -= q2_i64;
  477. }
  478. let mut val_rest = rest_rows.data[i] as i64;
  479. if val_rest >= q1_i64 / 2 {
  480. val_rest -= q1_i64;
  481. }
  482. let denom = (q2 * (q1 / p)) as i64;
  483. let mut r = val_first * q1_i64;
  484. r += val_rest * q2_i64;
  485. // divide r by q2, rounding
  486. let sign: i64 = if r >= 0 { 1 } else { -1 };
  487. let mut res = ((r + sign * (denom / 2)) as i128) / (denom as i128);
  488. res = (res + (denom as i128 / p_i128) * (p_i128) + 2 * (p_i128)) % (p_i128);
  489. let idx = instance * params.n * params.n * params.poly_len + i;
  490. result.data[idx] = res as u64;
  491. }
  492. }
  493. // println!("{:?}", result.data.as_slice().to_vec());
  494. result.to_vec(p_bits as usize, params.modp_words_per_chunk())
  495. }
  496. }
  497. #[cfg(test)]
  498. mod test {
  499. use rand::thread_rng;
  500. use super::*;
  501. fn assert_first8(m: &[u64], gold: [u64; 8]) {
  502. let got: [u64; 8] = m[0..8].try_into().unwrap();
  503. assert_eq!(got, gold);
  504. }
  505. fn get_params() -> Params {
  506. get_short_keygen_params()
  507. }
  508. #[test]
  509. fn init_is_correct() {
  510. let params = get_params();
  511. let mut rng = thread_rng();
  512. let client = Client::init(&params, &mut rng);
  513. assert_eq!(*client.params, params);
  514. }
  515. #[test]
  516. fn keygen_is_correct() {
  517. let params = get_params();
  518. let mut seeded_rng = get_static_seeded_rng();
  519. let mut client = Client::init(&params, &mut seeded_rng);
  520. let pub_params = client.generate_keys();
  521. assert_first8(
  522. pub_params.v_conversion.unwrap()[0].data.as_slice(),
  523. [
  524. 122680182, 165987256, 137892309, 95732358, 221787731, 13233184, 156136764,
  525. 259944211,
  526. ],
  527. );
  528. assert_first8(
  529. client.sk_gsw.data.as_slice(),
  530. [
  531. 66974689739603965,
  532. 66974689739603965,
  533. 0,
  534. 1,
  535. 0,
  536. 5,
  537. 66974689739603967,
  538. 2,
  539. ],
  540. );
  541. }
  542. fn get_vec(v: &Vec<PolyMatrixNTT>) -> Vec<u64> {
  543. v.iter().map(|d| d.as_slice().to_vec()).flatten().collect()
  544. }
  545. fn public_parameters_serialization_is_correct_for_params(params: Params) {
  546. let mut seeded_rng = get_static_seeded_rng();
  547. let mut client = Client::init(&params, &mut seeded_rng);
  548. let pub_params = client.generate_keys();
  549. let serialized1 = pub_params.serialize();
  550. let deserialized1 = PublicParameters::deserialize(&params, &serialized1);
  551. let serialized2 = deserialized1.serialize();
  552. assert_eq!(serialized1, serialized2);
  553. assert_eq!(
  554. get_vec(&pub_params.v_packing),
  555. get_vec(&deserialized1.v_packing)
  556. );
  557. assert_eq!(
  558. get_vec(&pub_params.v_conversion.unwrap()),
  559. get_vec(&deserialized1.v_conversion.unwrap())
  560. );
  561. assert_eq!(
  562. get_vec(&pub_params.v_expansion_left.unwrap()),
  563. get_vec(&deserialized1.v_expansion_left.unwrap())
  564. );
  565. assert_eq!(
  566. get_vec(&pub_params.v_expansion_right.unwrap()),
  567. get_vec(&deserialized1.v_expansion_right.unwrap())
  568. );
  569. }
  570. #[test]
  571. fn public_parameters_serialization_is_correct() {
  572. public_parameters_serialization_is_correct_for_params(get_params())
  573. }
  574. #[test]
  575. fn real_public_parameters_serialization_is_correct() {
  576. let cfg_expand = r#"
  577. {'n': 2,
  578. 'nu_1': 10,
  579. 'nu_2': 6,
  580. 'p': 512,
  581. 'q2_bits': 21,
  582. 's_e': 85.83255142749422,
  583. 't_gsw': 10,
  584. 't_conv': 4,
  585. 't_exp_left': 16,
  586. 't_exp_right': 56,
  587. 'instances': 11,
  588. 'db_item_size': 100000 }
  589. "#;
  590. let cfg = cfg_expand.replace("'", "\"");
  591. let params = params_from_json(&cfg);
  592. public_parameters_serialization_is_correct_for_params(params)
  593. }
  594. #[test]
  595. fn no_expansion_public_parameters_serialization_is_correct() {
  596. public_parameters_serialization_is_correct_for_params(get_no_expansion_testing_params())
  597. }
  598. fn query_serialization_is_correct_for_params(params: Params) {
  599. let mut seeded_rng = get_static_seeded_rng();
  600. let mut client = Client::init(&params, &mut seeded_rng);
  601. _ = client.generate_keys();
  602. let query = client.generate_query(1);
  603. let serialized1 = query.serialize();
  604. let deserialized1 = Query::deserialize(&params, &serialized1);
  605. let serialized2 = deserialized1.serialize();
  606. assert_eq!(serialized1, serialized2);
  607. }
  608. #[test]
  609. fn query_serialization_is_correct() {
  610. query_serialization_is_correct_for_params(get_params())
  611. }
  612. #[test]
  613. fn no_expansion_query_serialization_is_correct() {
  614. query_serialization_is_correct_for_params(get_no_expansion_testing_params())
  615. }
  616. }