Update derived traits, add ability to zeroize on drop, and refactor

Cargo.lock

@@ -30,6 +30,7 @@ name = "diffie-hellman"
 version = "0.1.0"
 dependencies = [
  "crypto-bigint",
+ "zeroize",
 ]
 
 [[package]]
@@ -78,3 +79,9 @@ name = "wasi"
 version = "0.11.1+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+
+[[package]]
+name = "zeroize"
+version = "1.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0"

Cargo.toml

@@ -5,10 +5,17 @@ edition = "2024"
 
 [dependencies]
 crypto-bigint = { version = "0.6.1" }
+zeroize = { version = "1.8.2", optional = true }
 
 [features]
 rand = ["crypto-bigint/rand"]
 edwards = []
+zeroize = ["dep:zeroize"]
+
+[[test]]
+name = "random"
+path = "tests/random.rs"
+required-features = ["rand"]
 
 [[bench]]
 name = "montgomery"

src/lib.rs

@@ -6,11 +6,11 @@ pub mod edwards;
 pub mod montgomery;
 pub mod traits;
 
+use crypto_bigint::subtle::{Choice, ConstantTimeEq};
 pub use crypto_bigint::{ConstChoice, Encoding, NonZero, U256};
-use std::fmt::{Display, Formatter};
 
 #[cfg(feature = "rand")]
-use crypto_bigint::{Random, rand_core::OsRng};
+use crypto_bigint::{rand_core::OsRng, Random};
 
 use crate::traits::{Curve, Point};
 pub(crate) use montgomery::*;
@@ -40,6 +40,7 @@ impl Public {
 }
 
 #[derive(Debug, Eq, PartialEq)]
+#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
 pub struct Secret(U256);
 
 impl Secret {
@@ -54,16 +55,6 @@ impl Secret {
 
         SharedSecret(p.get_x().to_le_bytes())
     }
-
-    pub fn to_bytes(&self) -> [u8; 32] {
-        self.0.to_le_bytes()
-    }
-}
-
-impl Display for Secret {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        f.write_fmt(format_args!("{}", self.0))
-    }
 }
 
 impl From<U256> for Secret {
@@ -72,7 +63,7 @@ impl From<U256> for Secret {
     }
 }
 
-#[derive(Debug, PartialEq)]
+#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
 pub struct SharedSecret([u8; 32]);
 
 impl SharedSecret {
@@ -81,6 +72,12 @@ impl SharedSecret {
     }
 }
 
+impl ConstantTimeEq for SharedSecret {
+    fn ct_eq(&self, other: &Self) -> Choice {
+        self.0.ct_eq(&other.0)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -90,7 +87,7 @@ mod tests {
     fn random_secret() {
         let secret = Secret::random();
 
-        assert_ne!(secret.0, U256::ZERO);
+        assert!(secret.0 < CURVE_25519.p);
     }
 
     #[test]
@@ -128,6 +125,6 @@ mod tests {
         let sa = a.diffie_hellman(&pb);
         let sb = b.diffie_hellman(&pa);
 
-        assert_eq!(sa, sb);
+        assert!(bool::from(sa.ct_eq(&sb)));
     }
 }

tests/random.rs

@@ -1,5 +1,4 @@
-#![cfg(feature = "rand")]
-
+use crypto_bigint::subtle::ConstantTimeEq;
 use diffie_hellman::{Public, Secret};
 
 #[test]
@@ -13,5 +12,5 @@ fn random_dh_key_exchange() {
     let sa = a.diffie_hellman(&pb);
     let sb = b.diffie_hellman(&pa);
 
-    println!("{} == {}", sa, sb)
+    assert!(bool::from(sa.ct_eq(&sb)));
 }