From 9fa9f4a8bc7fe016ca1840ef82c4eb7960b2302b Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@inbox.lv>
Date: Sat, 19 Apr 2025 18:52:57 +0200
Subject: [PATCH] fix random number generation + asymmetric key generation

---
 src/dghv.rs      | 36 ++++++++++++--------------------
 src/dghv_asym.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++++++
 src/utils.rs     | 20 ++++++++++++++++++
 3 files changed, 87 insertions(+), 23 deletions(-)
 create mode 100644 src/dghv_asym.rs
 create mode 100644 src/utils.rs

diff --git a/src/dghv.rs b/src/dghv.rs
index c30db91..cb73c79 100644
--- a/src/dghv.rs
+++ b/src/dghv.rs
@@ -1,18 +1,14 @@
-use rand::rngs::StdRng;
-use rand::RngCore;
-use rand::SeedableRng;
+use crate::utils::{generate_random_integer, generate_random_odd_integer};
 use rug::{Assign, Integer};
 
-pub fn encrypt_bit(bit: u8, key: &Integer) -> Integer {
+pub fn encrypt_bit(bit: u8, key: &Integer, gamma: u32, rho: u32) -> Integer {
     assert!(bit == 0 || bit == 1, "Only bits (0 or 1) are allowed");
 
-    let mut rng = StdRng::from_os_rng();
-
-    let q: Integer = Integer::from(rng.next_u64()) % key;
+    let q: Integer = generate_random_integer(gamma);
 
     // noise
     let p_div_4: Integer = key.clone() / 4;
-    let r = std::iter::repeat_with(|| Integer::from(rng.next_u64()))
+    let r = std::iter::repeat_with(|| generate_random_integer(rho))
         .find(|r| r.clone() * 2 < p_div_4)
         .unwrap();
 
@@ -29,13 +25,7 @@ pub fn decrypt_bit(encrypted: &Integer, key: &Integer) -> u8 {
 }
 
 pub fn generate_secret_key(eta: u32) -> Integer {
-    let mut rng = StdRng::from_os_rng();
-    let lower = Integer::from(1) << (eta - 1);
-    let upper = Integer::from(1) << eta;
-    let range: Integer = upper - &lower;
-    std::iter::repeat_with(|| Integer::from(rng.next_u64()) % &range + &lower)
-        .find(|p: &Integer| p.is_odd())
-        .unwrap()
+    generate_random_odd_integer(eta)
 }
 
 #[cfg(test)]
@@ -44,7 +34,7 @@ mod tests {
 
     #[test]
     fn test_secret_key_generation_properties() {
-        let eta = 64;
+        let eta = 1024;
         let key = generate_secret_key(eta);
         assert!(key.significant_bits() <= eta);
         assert!(key.significant_bits() >= eta - 1);
@@ -57,12 +47,12 @@ mod tests {
         let key = generate_secret_key(eta);
 
         let bit = 0;
-        let ciphertext = encrypt_bit(bit, &key);
+        let ciphertext = encrypt_bit(bit, &key, 128, 61);
         let decrypted = decrypt_bit(&ciphertext, &key);
         assert_eq!(decrypted, bit);
 
         let bit = 1;
-        let ciphertext = encrypt_bit(bit, &key);
+        let ciphertext = encrypt_bit(bit, &key, 128, 61);
         let decrypted = decrypt_bit(&ciphertext, &key);
         assert_eq!(decrypted, bit);
     }
@@ -73,8 +63,8 @@ mod tests {
         let key = generate_secret_key(eta);
 
         let bit = 1;
-        let c1 = encrypt_bit(bit, &key);
-        let c2 = encrypt_bit(bit, &key);
+        let c1 = encrypt_bit(bit, &key, 128, 61);
+        let c2 = encrypt_bit(bit, &key, 128, 61);
 
         assert_ne!(c1, c2);
         assert_eq!(decrypt_bit(&c1, &key), bit);
@@ -83,12 +73,12 @@ mod tests {
 
     #[test]
     fn test_is_homomorphic() {
-        let eta = 1024;
+        let eta = 2048;
         let key = generate_secret_key(eta);
 
         for &(b1, b2) in &[(0, 0), (0, 1), (1, 0), (1, 1)] {
-            let c1 = encrypt_bit(b1, &key);
-            let c2 = encrypt_bit(b2, &key);
+            let c1 = encrypt_bit(b1, &key, 4096, 512);
+            let c2 = encrypt_bit(b2, &key, 4096, 512);
 
             let sum: Integer = c1.clone() + &c2;
             let product: Integer = c1 * c2;
diff --git a/src/dghv_asym.rs b/src/dghv_asym.rs
new file mode 100644
index 0000000..a4a2a9a
--- /dev/null
+++ b/src/dghv_asym.rs
@@ -0,0 +1,54 @@
+use crate::dghv::{encrypt_bit, generate_secret_key};
+use rug::Integer;
+
+pub struct PublicKey {
+    pub xs: Vec<Integer>,
+}
+
+pub struct PrivateKey {
+    pub p: Integer,
+}
+
+pub fn generate_keys(gamma: u32, eta: u32, rho: u32, theta: usize) -> (PrivateKey, PublicKey) {
+    let p: Integer = generate_secret_key(eta);
+
+    let mut xs: Vec<Integer> = Vec::with_capacity(theta);
+    for _ in 0..theta {
+        let xi = encrypt_bit(0u8, &p, gamma, rho);
+        xs.push(xi);
+    }
+    xs.sort_by(|a, b| b.cmp(a));
+
+    while xs[0].is_even() {
+        xs.remove(0);
+        let xi = encrypt_bit(0u8, &p, gamma, rho);
+        xs.push(xi);
+        xs.sort_by(|a, b| b.cmp(a));
+    }
+
+    (PrivateKey { p }, PublicKey { xs })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_key_generation_properties() {
+        let eta: u32 = 64;
+        let gamma: u32 = 80;
+        let theta: usize = 10;
+        let rho: u32 = 16;
+        let (sk, pk) = generate_keys(gamma, eta, rho, theta);
+
+        assert_eq!(sk.p.significant_bits(), eta);
+        assert!(sk.p.is_odd());
+
+        assert_eq!(pk.xs.len(), theta);
+        assert!(pk.xs[0].is_odd());
+
+        for xi in pk.xs.iter().skip(1) {
+            assert!(pk.xs[0] >= *xi);
+        }
+    }
+}
diff --git a/src/utils.rs b/src/utils.rs
new file mode 100644
index 0000000..5280bf7
--- /dev/null
+++ b/src/utils.rs
@@ -0,0 +1,20 @@
+use rand::rngs::StdRng;
+use rand::Rng;
+use rand::SeedableRng;
+use rug::Integer;
+
+pub fn generate_random_odd_integer(num_bits: u32) -> Integer {
+    std::iter::repeat_with(|| generate_random_integer(num_bits))
+        .find(|p: &Integer| p.is_odd())
+        .unwrap()
+}
+
+pub fn generate_random_integer(num_bits: u32) -> Integer {
+    let mut rng = StdRng::from_os_rng();
+    let mut x = Integer::from(1);
+    for _ in 1..num_bits {
+        x <<= 1;
+        x += rng.random_range(0..=1);
+    }
+    x
+}