Browse Source

Remove deterministic client
- Use two RNG's, one for public uniform data and one for secret noise
- Remove thread_local! from discrete_gaussian
- Make client methods thread safe (again)

Samir Menon 1 year ago
parent
commit
176e396d7f

+ 3 - 2
spiral-rs/benches/server.rs

@@ -71,7 +71,8 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     let params = get_expansion_testing_params();
     let v_neg1 = params.get_v_neg1();
-    let mut chacha_rng = get_chacha_rng();
+    let mut rng = get_chacha_rng();
+    let mut rng_pub = get_chacha_rng();
     let mut client = Client::init(&params);
     let public_params = client.generate_keys();
 
@@ -82,7 +83,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     let scale_k = params.modulus / params.pt_modulus;
     let mut sigma = PolyMatrixRaw::zero(&params, 1, 1);
     sigma.data[7] = scale_k;
-    v[0] = client.encrypt_matrix_reg(&sigma.ntt(), &mut chacha_rng);
+    v[0] = client.encrypt_matrix_reg(&sigma.ntt(), &mut rng, &mut rng_pub);
 
     let v_w_left = public_params.v_expansion_left.unwrap();
     let v_w_right = public_params.v_expansion_right.unwrap();

+ 0 - 1
spiral-rs/src/bin/client.rs

@@ -1,4 +1,3 @@
-use rand::thread_rng;
 use serde_json::Value;
 use spiral_rs::client::*;
 use spiral_rs::util::*;

+ 1 - 1
spiral-rs/src/bin/e2e.rs

@@ -47,7 +47,7 @@ fn main() {
         params.num_items()
     );
     println!("initializing client");
-    let mut client = Client::init_with_seed(&params, get_chacha_seed());
+    let mut client = Client::init(&params);
     println!("generating public parameters");
     let pub_params = client.generate_keys();
     let pub_params_buf = pub_params.serialize();

+ 75 - 82
spiral-rs/src/client.rs

@@ -1,12 +1,12 @@
 use crate::{
     arith::*, discrete_gaussian::*, gadget::*, number_theory::*, params::*, poly::*, util::*,
 };
-use rand::{thread_rng, Rng, RngCore, SeedableRng};
+use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 use std::{iter::once, mem::size_of};
 
 pub type Seed = <ChaCha20Rng as SeedableRng>::Seed;
-const SEED_LENGTH: usize = 32;
+pub const SEED_LENGTH: usize = 32;
 
 fn new_vec_raw<'a>(
     params: &'a Params,
@@ -171,7 +171,7 @@ impl<'a> PublicParameters<'a> {
     }
 
     pub fn deserialize(params: &'a Params, data: &[u8]) -> Self {
-        assert_eq!(SEED_LENGTH + params.setup_bytes(), data.len());
+        assert_eq!(params.setup_bytes(), data.len());
 
         let mut idx = 0;
 
@@ -254,7 +254,7 @@ impl<'a> Query<'a> {
     }
 
     pub fn deserialize(params: &'a Params, mut data: &[u8]) -> Self {
-        assert_eq!(SEED_LENGTH + params.query_bytes(), data.len());
+        assert_eq!(params.query_bytes(), data.len());
 
         let mut out = Query::empty();
         let seed = data[0..SEED_LENGTH].try_into().unwrap();
@@ -317,22 +317,10 @@ pub struct Client<'a> {
     sk_gsw_full: PolyMatrixRaw<'a>,
     sk_reg_full: PolyMatrixRaw<'a>,
     dg: DiscreteGaussian,
-    pp_seed: Seed,
-    query_seed: Seed,
 }
 
 impl<'a> Client<'a> {
     pub fn init(params: &'a Params) -> Self {
-        let mut root_seed = [0u8; 32];
-        thread_rng().fill_bytes(&mut root_seed);
-        Self::init_with_seed_impl(params, root_seed, false)
-    }
-
-    pub fn init_with_seed(params: &'a Params, root_seed: Seed) -> Self {
-        Self::init_with_seed_impl(params, root_seed, true)
-    }
-
-    fn init_with_seed_impl(params: &'a Params, root_seed: Seed, deterministic: bool) -> Self {
         let sk_gsw_dims = params.get_sk_gsw();
         let sk_reg_dims = params.get_sk_reg();
         let sk_gsw = PolyMatrixRaw::zero(params, sk_gsw_dims.0, sk_gsw_dims.1);
@@ -340,16 +328,8 @@ impl<'a> Client<'a> {
         let sk_gsw_full = matrix_with_identity(&sk_gsw);
         let sk_reg_full = matrix_with_identity(&sk_reg);
 
-        let mut rng = ChaCha20Rng::from_seed(root_seed);
-        let mut pp_seed = [0u8; 32];
-        let mut query_seed = [0u8; 32];
-        rng.fill_bytes(&mut query_seed);
-        rng.fill_bytes(&mut pp_seed);
-        let dg = if deterministic {
-            DiscreteGaussian::init_with_seed(params, root_seed)
-        } else {
-            DiscreteGaussian::init(params)
-        };
+        let dg = DiscreteGaussian::init(params);
+
         Self {
             params,
             sk_gsw,
@@ -357,8 +337,6 @@ impl<'a> Client<'a> {
             sk_gsw_full,
             sk_reg_full,
             dg,
-            pp_seed,
-            query_seed,
         }
     }
 
@@ -367,12 +345,17 @@ impl<'a> Client<'a> {
         &self.sk_reg
     }
 
-    fn get_fresh_gsw_public_key(&self, m: usize, rng: &mut ChaCha20Rng) -> PolyMatrixRaw<'a> {
+    fn get_fresh_gsw_public_key(
+        &self,
+        m: usize,
+        rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
+    ) -> PolyMatrixRaw<'a> {
         let params = self.params;
         let n = params.n;
 
-        let a = PolyMatrixRaw::random_rng(params, 1, m, rng);
-        let e = PolyMatrixRaw::noise(params, n, m, &self.dg);
+        let a = PolyMatrixRaw::random_rng(params, 1, m, rng_pub);
+        let e = PolyMatrixRaw::noise(params, n, m, &self.dg, rng);
         let a_inv = -&a;
         let b_p = &self.sk_gsw.ntt() * &a.ntt();
         let b = &e.ntt() + &b_p;
@@ -380,10 +363,14 @@ impl<'a> Client<'a> {
         p
     }
 
-    fn get_regev_sample(&self, rng: &mut ChaCha20Rng) -> PolyMatrixNTT<'a> {
+    fn get_regev_sample(
+        &self,
+        rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
+    ) -> PolyMatrixNTT<'a> {
         let params = self.params;
-        let a = PolyMatrixRaw::random_rng(params, 1, 1, rng);
-        let e = PolyMatrixRaw::noise(params, 1, 1, &self.dg);
+        let a = PolyMatrixRaw::random_rng(params, 1, 1, rng_pub);
+        let e = PolyMatrixRaw::noise(params, 1, 1, &self.dg, rng);
         let b_p = &self.sk_reg.ntt() * &a.ntt();
         let b = &e.ntt() + &b_p;
         let mut p = PolyMatrixNTT::zero(params, 2, 1);
@@ -392,13 +379,18 @@ impl<'a> Client<'a> {
         p
     }
 
-    fn get_fresh_reg_public_key(&self, m: usize, rng: &mut ChaCha20Rng) -> PolyMatrixNTT<'a> {
+    fn get_fresh_reg_public_key(
+        &self,
+        m: usize,
+        rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
+    ) -> PolyMatrixNTT<'a> {
         let params = self.params;
 
         let mut p = PolyMatrixNTT::zero(params, 2, m);
 
         for i in 0..m {
-            p.copy_into(&self.get_regev_sample(rng), 0, i);
+            p.copy_into(&self.get_regev_sample(rng, rng_pub), 0, i);
         }
         p
     }
@@ -407,9 +399,10 @@ impl<'a> Client<'a> {
         &self,
         ag: &PolyMatrixNTT<'a>,
         rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
     ) -> PolyMatrixNTT<'a> {
         let mx = ag.cols;
-        let p = self.get_fresh_gsw_public_key(mx, rng);
+        let p = self.get_fresh_gsw_public_key(mx, rng, rng_pub);
         let res = &(p.ntt()) + &(ag.pad_top(1));
         res
     }
@@ -418,9 +411,10 @@ impl<'a> Client<'a> {
         &self,
         a: &PolyMatrixNTT<'a>,
         rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
     ) -> PolyMatrixNTT<'a> {
         let m = a.cols;
-        let p = self.get_fresh_reg_public_key(m, rng);
+        let p = self.get_fresh_reg_public_key(m, rng, rng_pub);
         &p + &a.pad_top(1)
     }
 
@@ -437,6 +431,7 @@ impl<'a> Client<'a> {
         num_exp: usize,
         m_exp: usize,
         rng: &mut ChaCha20Rng,
+        rng_pub: &mut ChaCha20Rng,
     ) -> Vec<PolyMatrixNTT<'a>> {
         let params = self.params;
         let g_exp = build_gadget(params, 1, m_exp);
@@ -447,23 +442,26 @@ impl<'a> Client<'a> {
             let t = (params.poly_len / (1 << i)) + 1;
             let tau_sk_reg = automorph_alloc(&self.sk_reg, t);
             let prod = &tau_sk_reg.ntt() * &g_exp_ntt;
-            let w_exp_i = self.encrypt_matrix_reg(&prod, rng);
+            let w_exp_i = self.encrypt_matrix_reg(&prod, rng, rng_pub);
             res.push(w_exp_i);
         }
         res
     }
 
     pub fn generate_keys(&mut self) -> PublicParameters<'a> {
+        let mut rng = ChaCha20Rng::from_entropy();
         let params = self.params;
-        self.dg.sample_matrix(&mut self.sk_gsw);
-        self.dg.sample_matrix(&mut self.sk_reg);
+
+        self.dg.sample_matrix(&mut self.sk_gsw, &mut rng);
+        self.dg.sample_matrix(&mut self.sk_reg, &mut rng);
         self.sk_gsw_full = matrix_with_identity(&self.sk_gsw);
         self.sk_reg_full = matrix_with_identity(&self.sk_reg);
         let sk_reg_ntt = to_ntt_alloc(&self.sk_reg);
 
         let mut pp = PublicParameters::init(params);
-        pp.seed = Some(self.pp_seed);
-        let mut rng = ChaCha20Rng::from_seed(self.pp_seed);
+        let pp_seed = rng.gen();
+        pp.seed = Some(pp_seed);
+        let mut rng_pub = ChaCha20Rng::from_seed(pp_seed);
 
         // Params for packing
         let gadget_conv = build_gadget(params, 1, params.t_conv);
@@ -472,18 +470,23 @@ impl<'a> Client<'a> {
             let scaled = scalar_multiply_alloc(&sk_reg_ntt, &gadget_conv_ntt);
             let mut ag = PolyMatrixNTT::zero(params, params.n, params.t_conv);
             ag.copy_into(&scaled, i, 0);
-            let w = self.encrypt_matrix_gsw(&ag, &mut rng);
+            let w = self.encrypt_matrix_gsw(&ag, &mut rng, &mut rng_pub);
             pp.v_packing.push(w);
         }
 
         if params.expand_queries {
             // Params for expansion
-            pp.v_expansion_left =
-                Some(self.generate_expansion_params(params.g(), params.t_exp_left, &mut rng));
+            pp.v_expansion_left = Some(self.generate_expansion_params(
+                params.g(),
+                params.t_exp_left,
+                &mut rng,
+                &mut rng_pub,
+            ));
             pp.v_expansion_right = Some(self.generate_expansion_params(
                 params.stop_round() + 1,
                 params.t_exp_right,
                 &mut rng,
+                &mut rng_pub,
             ));
 
             // Params for converison
@@ -504,7 +507,7 @@ impl<'a> Client<'a> {
                     let val = g_conv.get_poly(1, i)[0];
                     sigma = &sk_reg_ntt * &single_poly(params, val).ntt();
                 }
-                let ct = self.encrypt_matrix_reg(&sigma, &mut rng);
+                let ct = self.encrypt_matrix_reg(&sigma, &mut rng, &mut rng_pub);
                 pp.v_conversion.as_mut().unwrap()[0].copy_into(&ct, 0, i);
             }
         }
@@ -520,9 +523,12 @@ impl<'a> Client<'a> {
         let scale_k = params.modulus / params.pt_modulus;
         let bits_per = get_bits_per(params, params.t_gsw);
 
+        let mut rng = ChaCha20Rng::from_entropy();
+
         let mut query = Query::empty();
-        query.seed = Some(self.query_seed);
-        let mut rng = ChaCha20Rng::from_seed(self.query_seed);
+        let query_seed = ChaCha20Rng::from_entropy().gen();
+        query.seed = Some(query_seed);
+        let mut rng_pub = ChaCha20Rng::from_seed(query_seed);
         if params.expand_queries {
             // pack query into single ciphertext
             let mut sigma = PolyMatrixRaw::zero(params, 1, 1);
@@ -555,9 +561,11 @@ impl<'a> Client<'a> {
                 }
             }
 
-            query.ct = Some(from_ntt_alloc(
-                &self.encrypt_matrix_reg(&to_ntt_alloc(&sigma), &mut rng),
-            ));
+            query.ct = Some(from_ntt_alloc(&self.encrypt_matrix_reg(
+                &to_ntt_alloc(&sigma),
+                &mut rng,
+                &mut rng_pub,
+            )));
         } else {
             let num_expanded = 1 << params.db_dim_1;
             let mut sigma_v = Vec::<PolyMatrixNTT>::new();
@@ -569,7 +577,11 @@ impl<'a> Client<'a> {
             for i in 0..num_expanded {
                 let value = ((i == idx_dim0) as u64) * scale_k;
                 let sigma = PolyMatrixRaw::single_value(&params, value);
-                reg_cts.push(self.encrypt_matrix_reg(&to_ntt_alloc(&sigma), &mut rng));
+                reg_cts.push(self.encrypt_matrix_reg(
+                    &to_ntt_alloc(&sigma),
+                    &mut rng,
+                    &mut rng_pub,
+                ));
             }
             // reorient into server's preferred indexing
             reorient_reg_ciphertexts(self.params, reg_cts_buf.as_mut_slice(), &reg_cts);
@@ -586,10 +598,10 @@ impl<'a> Client<'a> {
 
                     // important to rng in the right order here
                     let prod = &to_ntt_alloc(&self.sk_reg) * &sigma_ntt;
-                    let ct = &self.encrypt_matrix_reg(&prod, &mut rng);
+                    let ct = &self.encrypt_matrix_reg(&prod, &mut rng, &mut rng_pub);
                     ct_gsw.copy_into(ct, 0, 2 * j);
 
-                    let ct = &self.encrypt_matrix_reg(&sigma_ntt, &mut rng);
+                    let ct = &self.encrypt_matrix_reg(&sigma_ntt, &mut rng, &mut rng_pub);
                     ct_gsw.copy_into(ct, 0, 2 * j + 1);
                 }
                 sigma_v.push(ct_gsw);
@@ -686,11 +698,6 @@ impl<'a> Client<'a> {
 mod test {
     use super::*;
 
-    fn assert_first8(m: &[u64], gold: [u64; 8]) {
-        let got: [u64; 8] = m[0..8].try_into().unwrap();
-        assert_eq!(got, gold);
-    }
-
     fn get_params() -> Params {
         get_short_keygen_params()
     }
@@ -706,30 +713,16 @@ mod test {
     #[test]
     fn keygen_is_correct() {
         let params = get_params();
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
 
-        let pub_params = client.generate_keys();
+        _ = client.generate_keys();
 
-        assert_first8(
-            pub_params.v_conversion.unwrap()[0].data.as_slice(),
-            [
-                160400272, 10738392, 27480222, 201012452, 42036824, 3955189, 201319389, 181880730,
-            ],
-        );
+        let threshold = (10.0 * params.noise_width) as u64;
 
-        assert_first8(
-            client.sk_gsw.data.as_slice(),
-            [
-                1,
-                66974689739603968,
-                4,
-                3,
-                66974689739603965,
-                66974689739603967,
-                1,
-                0,
-            ],
-        );
+        for i in 0..client.sk_gsw.data.len() {
+            let val = client.sk_gsw.data[i];
+            assert!((val < threshold) || ((params.modulus - val) < threshold));
+        }
     }
 
     fn get_vec(v: &Vec<PolyMatrixNTT>) -> Vec<u64> {
@@ -737,7 +730,7 @@ mod test {
     }
 
     fn public_parameters_serialization_is_correct_for_params(params: Params) {
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
         let pub_params = client.generate_keys();
 
         let serialized1 = pub_params.serialize();

+ 7 - 31
spiral-rs/src/discrete_gaussian.rs

@@ -1,20 +1,13 @@
 use rand::distributions::WeightedIndex;
 use rand::prelude::Distribution;
-use rand::thread_rng;
-use rand::Rng;
-use rand::SeedableRng;
 use rand_chacha::ChaCha20Rng;
-use std::cell::*;
 
-use crate::client::*;
 use crate::params::*;
 use crate::poly::*;
 use std::f64::consts::PI;
 
 pub const NUM_WIDTHS: usize = 8;
 
-thread_local!(static RNGS: RefCell<Option<ChaCha20Rng>> = RefCell::new(None));
-
 pub struct DiscreteGaussian {
     choices: Vec<i64>,
     dist: WeightedIndex<f64>,
@@ -22,14 +15,6 @@ pub struct DiscreteGaussian {
 
 impl DiscreteGaussian {
     pub fn init(params: &Params) -> Self {
-        Self::init_impl(params, None)
-    }
-
-    pub fn init_with_seed(params: &Params, seed: Seed) -> Self {
-        Self::init_impl(params, Some(seed))
-    }
-
-    fn init_impl(params: &Params, seed: Option<Seed>) -> Self {
         let max_val = (params.noise_width * (NUM_WIDTHS as f64)).ceil() as i64;
         let mut choices = Vec::new();
         let mut table = vec![0f64; 0];
@@ -40,31 +25,21 @@ impl DiscreteGaussian {
         }
         let dist = WeightedIndex::new(&table).unwrap();
 
-        if seed.is_some() {
-            RNGS.with(|f| *f.borrow_mut() = Some(ChaCha20Rng::from_seed(seed.unwrap())));
-        } else {
-            RNGS.with(|f| *f.borrow_mut() = Some(ChaCha20Rng::from_seed(thread_rng().gen())));
-        }
-
         Self { choices, dist }
     }
 
     // FIXME: not constant-time
-    pub fn sample(&self) -> i64 {
-        RNGS.with(|f| {
-            let mut rng = f.borrow_mut();
-            let rng_mut = rng.as_mut().unwrap();
-            self.choices[self.dist.sample(rng_mut)]
-        })
+    pub fn sample(&self, rng: &mut ChaCha20Rng) -> i64 {
+        self.choices[self.dist.sample(rng)]
     }
 
-    pub fn sample_matrix(&self, p: &mut PolyMatrixRaw) {
+    pub fn sample_matrix(&self, p: &mut PolyMatrixRaw, rng: &mut ChaCha20Rng) {
         let modulus = p.get_params().modulus;
         for r in 0..p.rows {
             for c in 0..p.cols {
                 let poly = p.get_poly_mut(r, c);
                 for z in 0..poly.len() {
-                    let mut s = self.sample();
+                    let mut s = self.sample(rng);
                     s += modulus as i64;
                     s %= modulus as i64; // FIXME: not constant time
                     poly[z] = s as u64;
@@ -82,12 +57,13 @@ mod test {
     #[test]
     fn dg_seems_okay() {
         let params = get_test_params();
-        let dg = DiscreteGaussian::init_with_seed(&params, get_chacha_seed());
+        let dg = DiscreteGaussian::init(&params);
+        let mut rng = get_chacha_rng();
         let mut v = Vec::new();
         let trials = 10000;
         let mut sum = 0;
         for _ in 0..trials {
-            let val = dg.sample();
+            let val = dg.sample(&mut rng);
             v.push(val);
             sum += val;
         }

+ 3 - 3
spiral-rs/src/params.rs

@@ -1,6 +1,6 @@
 use std::mem::size_of;
 
-use crate::{arith::*, ntt::*, number_theory::*, poly::*};
+use crate::{arith::*, client::SEED_LENGTH, ntt::*, number_theory::*, poly::*};
 
 pub const MAX_MODULI: usize = 4;
 
@@ -156,7 +156,7 @@ impl Params {
         }
 
         let sz_bytes = sz_polys * self.poly_len * size_of::<u64>();
-        sz_bytes
+        SEED_LENGTH + sz_bytes
     }
 
     pub fn query_bytes(&self) -> usize {
@@ -171,7 +171,7 @@ impl Params {
         }
 
         let sz_bytes = sz_polys * self.poly_len * size_of::<u64>();
-        sz_bytes
+        SEED_LENGTH + sz_bytes
     }
 
     pub fn query_v_buf_bytes(&self) -> usize {

+ 9 - 2
spiral-rs/src/poly.rs

@@ -3,6 +3,7 @@ use std::arch::x86_64::*;
 
 use rand::distributions::Standard;
 use rand::Rng;
+use rand_chacha::ChaCha20Rng;
 use std::cell::RefCell;
 use std::ops::{Add, Mul, Neg};
 
@@ -172,9 +173,15 @@ impl<'a> PolyMatrixRaw<'a> {
         }
     }
 
-    pub fn noise(params: &'a Params, rows: usize, cols: usize, dg: &DiscreteGaussian) -> Self {
+    pub fn noise(
+        params: &'a Params,
+        rows: usize,
+        cols: usize,
+        dg: &DiscreteGaussian,
+        rng: &mut ChaCha20Rng,
+    ) -> Self {
         let mut out = PolyMatrixRaw::zero(params, rows, cols);
-        dg.sample_matrix(&mut out);
+        dg.sample_matrix(&mut out, rng);
         out
     }
 

+ 23 - 18
spiral-rs/src/server.rs

@@ -789,7 +789,8 @@ pub fn process_query(
 mod test {
     use super::*;
     use crate::client::*;
-    use rand::Rng;
+    use rand::{Rng, SeedableRng};
+    use rand_chacha::ChaCha20Rng;
 
     const TEST_PREPROCESSED_DB_PATH: &'static str = "/home/samir/wiki/enwiki-20220320.dbp";
 
@@ -834,8 +835,9 @@ mod test {
     fn coefficient_expansion_is_correct() {
         let params = get_params();
         let v_neg1 = params.get_v_neg1();
-        let mut chacha_rng = get_chacha_rng();
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut rng = ChaCha20Rng::from_entropy();
+        let mut rng_pub = ChaCha20Rng::from_entropy();
+        let mut client = Client::init(&params);
         let public_params = client.generate_keys();
 
         let mut v = Vec::new();
@@ -847,8 +849,8 @@ mod test {
         let scale_k = params.modulus / params.pt_modulus;
         let mut sigma = PolyMatrixRaw::zero(&params, 1, 1);
         sigma.data[target] = scale_k;
-        v[0] = client.encrypt_matrix_reg(&sigma.ntt(), &mut chacha_rng);
-        let test_ct = client.encrypt_matrix_reg(&sigma.ntt(), &mut chacha_rng);
+        v[0] = client.encrypt_matrix_reg(&sigma.ntt(), &mut rng, &mut rng_pub);
+        let test_ct = client.encrypt_matrix_reg(&sigma.ntt(), &mut rng, &mut rng_pub);
 
         let v_w_left = public_params.v_expansion_left.unwrap();
         let v_w_right = public_params.v_expansion_right.unwrap();
@@ -878,14 +880,15 @@ mod test {
     fn regev_to_gsw_is_correct() {
         let mut params = get_params();
         params.db_dim_2 = 1;
-        let mut chacha_rng = get_chacha_rng();
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut rng = ChaCha20Rng::from_entropy();
+        let mut rng_pub = ChaCha20Rng::from_entropy();
+        let mut client = Client::init(&params);
         let public_params = client.generate_keys();
 
         let mut enc_constant = |val| {
             let mut sigma = PolyMatrixRaw::zero(&params, 1, 1);
             sigma.data[0] = val;
-            client.encrypt_matrix_reg(&sigma.ntt(), &mut chacha_rng)
+            client.encrypt_matrix_reg(&sigma.ntt(), &mut rng, &mut rng_pub)
         };
 
         let v = &public_params.v_conversion.unwrap()[0];
@@ -915,7 +918,8 @@ mod test {
     fn multiply_reg_by_database_is_correct() {
         let params = get_params();
         let mut seeded_rng = get_seeded_rng();
-        let mut chacha_rng = get_chacha_rng();
+        let mut rng = ChaCha20Rng::from_entropy();
+        let mut rng_pub = ChaCha20Rng::from_entropy();
 
         let dim0 = 1 << params.db_dim_1;
         let num_per = 1 << params.db_dim_2;
@@ -925,7 +929,7 @@ mod test {
         let target_idx_dim0 = target_idx / num_per;
         let target_idx_num_per = target_idx % num_per;
 
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
         _ = client.generate_keys();
 
         let (corr_item, db) = generate_random_db_and_get_item(&params, target_idx);
@@ -934,7 +938,7 @@ mod test {
         for i in 0..dim0 {
             let val = if i == target_idx_dim0 { scale_k } else { 0 };
             let sigma = PolyMatrixRaw::single_value(&params, val).ntt();
-            v_reg.push(client.encrypt_matrix_reg(&sigma, &mut chacha_rng));
+            v_reg.push(client.encrypt_matrix_reg(&sigma, &mut rng, &mut rng_pub));
         }
 
         let v_reg_sz = dim0 * 2 * params.poly_len;
@@ -971,7 +975,8 @@ mod test {
     fn fold_ciphertexts_is_correct() {
         let params = get_params();
         let mut seeded_rng = get_seeded_rng();
-        let mut chacha_rng = get_chacha_rng();
+        let mut rng = ChaCha20Rng::from_entropy();
+        let mut rng_pub = ChaCha20Rng::from_entropy();
 
         let dim0 = 1 << params.db_dim_1;
         let num_per = 1 << params.db_dim_2;
@@ -980,14 +985,14 @@ mod test {
         let target_idx = seeded_rng.gen::<usize>() % (dim0 * num_per);
         let target_idx_num_per = target_idx % num_per;
 
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
         _ = client.generate_keys();
 
         let mut v_reg = Vec::new();
         for i in 0..num_per {
             let val = if i == target_idx_num_per { scale_k } else { 0 };
             let sigma = PolyMatrixRaw::single_value(&params, val).ntt();
-            v_reg.push(client.encrypt_matrix_reg(&sigma, &mut chacha_rng));
+            v_reg.push(client.encrypt_matrix_reg(&sigma, &mut rng, &mut rng_pub));
         }
 
         let mut v_reg_raw = Vec::new();
@@ -1005,10 +1010,10 @@ mod test {
                 let value = (1u64 << (bits_per * j)) * bit;
                 let sigma = PolyMatrixRaw::single_value(&params, value);
                 let sigma_ntt = to_ntt_alloc(&sigma);
-                let ct = client.encrypt_matrix_reg(&sigma_ntt, &mut chacha_rng);
+                let ct = client.encrypt_matrix_reg(&sigma_ntt, &mut rng, &mut rng_pub);
                 ct_gsw.copy_into(&ct, 0, 2 * j + 1);
                 let prod = &to_ntt_alloc(client.get_sk_reg()) * &sigma_ntt;
-                let ct = &client.encrypt_matrix_reg(&prod, &mut chacha_rng);
+                let ct = &client.encrypt_matrix_reg(&prod, &mut rng, &mut rng_pub);
                 ct_gsw.copy_into(&ct, 0, 2 * j);
             }
 
@@ -1039,7 +1044,7 @@ mod test {
 
         let target_idx = seeded_rng.gen::<usize>() % (params.db_dim_1 + params.db_dim_2);
 
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
 
         let public_params = client.generate_keys();
         let query = client.generate_query(target_idx);
@@ -1065,7 +1070,7 @@ mod test {
 
         let target_idx = seeded_rng.gen::<usize>() % (params.db_dim_1 + params.db_dim_2);
 
-        let mut client = Client::init_with_seed(&params, get_chacha_static_seed());
+        let mut client = Client::init(&params);
 
         let public_params = client.generate_keys();
         let query = client.generate_query(target_idx);