From 98a29b3799397c8618f3a37ca53423a23edc08c9 Mon Sep 17 00:00:00 2001
From: Vaibhav Pathak <pathakvaibhav@protonmail.com>
Date: Tue, 31 Dec 2024 01:12:14 +0530
Subject: [PATCH] improve keygen algo

---
 src/dsa/eddsa/mod.rs   | 22 +++++++++++++++++-----
 src/dsa/eddsa/tests.rs | 19 ++++++++++++++-----
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/dsa/eddsa/mod.rs b/src/dsa/eddsa/mod.rs
index fc33ecf6..8c452f09 100644
--- a/src/dsa/eddsa/mod.rs
+++ b/src/dsa/eddsa/mod.rs
@@ -4,6 +4,7 @@
 //!     1. [RFC8032] "Edwards-Curve Digital Signature Algorithm (EdDSA)".
 use crypto_bigint::{Encoding, U256, U512};
 use curve::{Coordinate, ScalarField, ScalarField64, GENERATOR, ORDER};
+use rand::Rng;
 use sha2::{Digest, Sha512};
 
 pub mod curve;
@@ -46,18 +47,29 @@ fn reduce_by_order(x: [u8; 64]) -> [u8; 32] {
 pub struct Ed25519;
 
 impl Ed25519 {
-  /// Generates the `public_key` given the `private_key`.
+  /// Generates the `public_key` using the `private_key`, if given, otherwise randomly generates
+  /// one.
   ///
-  /// The `private_key` should be a randomly generated 32-byte array.
   /// Returns both the keys as a (private_key, public_key) tuple.
-  pub fn keygen(private_key: [u8; 32]) -> ([u8; 32], [u8; 32]) {
-    let keyhash = sha512_hash(&private_key);
+  pub fn keygen(private_key: Option<[u8; 32]>) -> ([u8; 32], [u8; 32]) {
+    let pk = match private_key {
+      Some(pk) => pk,
+      None => {
+        let mut rng = rand::thread_rng();
+        let v: Vec<_> = (0..32).map(|_| rng.gen_range(0..=255)).collect();
+        let mut a = [0u8; 32];
+        a.copy_from_slice(&v);
+        a
+      },
+    };
+
+    let keyhash = sha512_hash(&pk);
     let mut h = [0u8; 32];
     h.copy_from_slice(&keyhash[..32]);
 
     let a = ScalarField::new(&U256::from_le_bytes(clamp(h)));
     let public_key = GENERATOR * a;
-    (private_key, public_key.encode())
+    (pk, public_key.encode())
   }
 
   /// Sign the `message` using the `public_key` and `private_key`.
diff --git a/src/dsa/eddsa/tests.rs b/src/dsa/eddsa/tests.rs
index 5c771078..482abfb4 100644
--- a/src/dsa/eddsa/tests.rs
+++ b/src/dsa/eddsa/tests.rs
@@ -10,6 +10,15 @@ use test::Bencher;
 
 use super::*;
 
+#[test]
+fn test_simple() {
+  let (sk, pk) = Ed25519::keygen(None);
+  let msg = b"Hello World";
+
+  let signature = Ed25519::sign(sk, pk, msg);
+  assert!(Ed25519::verify(pk, msg, signature));
+}
+
 /// Test the `Ed25519` digital signature scheme using the test vectors given in Section 7.1 of RFC
 /// 8032.
 #[rstest]
@@ -49,7 +58,7 @@ fn test_small(
   #[case] msg: &[u8],
   #[case] expected_signature: [u8; 64],
 ) {
-  let (_, public_key) = Ed25519::keygen(secret_key);
+  let (_, public_key) = Ed25519::keygen(Some(secret_key));
   assert_eq!(public_key, expected_public_key);
 
   let signature = Ed25519::sign(secret_key, public_key, msg);
@@ -82,7 +91,7 @@ fn test_large() {
 
     let msg = decode_hex(&v[2]).unwrap();
 
-    let (_, public_key) = Ed25519::keygen(secret_key);
+    let (_, public_key) = Ed25519::keygen(Some(secret_key));
     assert_eq!(public_key, expected_public_key);
 
     let signature = Ed25519::sign(secret_key, public_key, &msg);
@@ -101,7 +110,7 @@ fn bench_keygen(b: &mut Bencher) {
 
   b.iter(|| {
     let sk = test::black_box(sk_b);
-    Ed25519::keygen(sk)
+    Ed25519::keygen(Some(sk))
   });
 }
 
@@ -113,7 +122,7 @@ macro_rules! bench_sign {
       let sk_v: Vec<_> = (0..32).map(|_| rng.gen_range(0..=255)).collect();
       let mut sk_b = [0u8; 32];
       sk_b.copy_from_slice(&sk_v);
-      let (_, pk_b) = Ed25519::keygen(sk_b);
+      let (_, pk_b) = Ed25519::keygen(Some(sk_b));
 
       let msg_v: Vec<_> = (0..$n).map(|_| rng.gen_range(0..=255)).collect();
       let mut msg_b = [0u8; $n];
@@ -147,7 +156,7 @@ macro_rules! bench_verify {
       let sk_v: Vec<_> = (0..32).map(|_| rng.gen_range(0..=255)).collect();
       let mut sk_b = [0u8; 32];
       sk_b.copy_from_slice(&sk_v);
-      let (_, pk_b) = Ed25519::keygen(sk_b);
+      let (_, pk_b) = Ed25519::keygen(Some(sk_b));
 
       let msg_v: Vec<_> = (0..$n).map(|_| rng.gen_range(0..=255)).collect();
       let mut msg_b = [0u8; $n];