client.rs 25 KB

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