FINALLY WORKS (yay), integration test

.idea/diffie-hellman.iml

@@ -3,6 +3,7 @@
   <component name="NewModuleRootManager">
     <content url="file://$MODULE_DIR$">
       <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/tests" isTestSource="true" />
       <excludeFolder url="file://$MODULE_DIR$/target" />
     </content>
     <orderEntry type="inheritedJdk" />

Cargo.toml

@@ -4,4 +4,7 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-crypto-bigint = { version = "0.6.1", features = ["rand"] }
+crypto-bigint = { version = "0.6.1" }
+
+[features]
+rand = ["crypto-bigint/rand"]

src/lib.rs

@@ -1,7 +1,8 @@
-use std::fmt::{Display, Formatter, Write};
+use crypto_bigint::{ConstChoice, Encoding, NonZero, U256};
-use crypto_bigint::{ConstChoice, Encoding, MulMod, NonZero, Odd, Random, U256};
+use std::fmt::{Display, Formatter};
-use std::ops::Mul;
+use std::ops::{Div, Mul};
-use crypto_bigint::rand_core::OsRng;
+
+#[cfg(feature = "rand")] use crypto_bigint::{Random, rand_core::OsRng};
 
 
 /// By^2 = x^3 + Ax^2 + x
@@ -31,15 +32,14 @@ impl MontgomeryCurve {
 }
 
 fn clamp_u256(x: U256) -> U256 {
-    let mut words: [u8; 32] = x.to_le_bytes();
+    let mut bytes: [u8; 32] = x.to_le_bytes();
 
     // clamp value
-    words[0] &= 0b1111_1000;
+    bytes[0] &= 248;
-    words[31] &= 0b0111_1111;
+    bytes[31] &= 127;
-    words[31] |= 0b0100_0000;
+    bytes[31] |= 64;
 
-
+    U256::from_le_bytes(bytes)
-    U256::from_le_bytes(words)
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -52,7 +52,7 @@ pub struct MontgomeryPoint {
 
 impl MontgomeryPoint {
     pub fn add(&self, rhs: MontgomeryPoint, neg: MontgomeryPoint) -> MontgomeryPoint {
-        let p = &NonZero::new(CURVE_25519.p).unwrap();
+        let p = &NonZero::new(self.curve.p).unwrap();
 
         let mut v_0 = self.x.add_mod(&self.z, p); // 1: V_0 = X_P + Z_P
         let mut v_1 = rhs.x.sub_mod(&rhs.z, p);   // 2: V_1 = X_Q - Z_Q
@@ -79,7 +79,7 @@ impl MontgomeryPoint {
     }
 
     pub fn double(&self) -> MontgomeryPoint {
-        let p = &NonZero::new(CURVE_25519.p).unwrap();
+        let p = &NonZero::new(self.curve.p).unwrap();
 
         let mut v_1 = self.x.add_mod(&self.z, p);
         v_1 = v_1.mul_mod(&v_1, p);
@@ -90,7 +90,7 @@ impl MontgomeryPoint {
         let x_2p = v_1.mul_mod(&v_2, p);
 
         v_1 = v_1.sub_mod(&v_2, p);
-        let mut v_3 = ((self.curve.a + U256::from_u8(2))/U256::from_u8(4)).mul_mod(&v_1, p);
+        let mut v_3 = self.curve.a.add_mod(&U256::from_u8(2), p).div(U256::from_u8(4)).mul_mod(&v_1, p);
 
         v_3 = v_3.add_mod(&v_2, p);
 
@@ -105,10 +105,65 @@ impl MontgomeryPoint {
 
     pub fn get_x(&self) -> U256 {
         self.x.mul_mod(
-            &self.z.inv_mod(&CURVE_25519.p).unwrap(),
+            &self.z.inv_mod(&self.curve.p).unwrap(),
             &NonZero::new(self.curve.p).unwrap()
         )
     }
+
+    pub fn ladder(self, rhs: U256) -> MontgomeryPoint {
+        let mut x_0 = self;
+        let mut x_1 = x_0.double();
+
+        for i in (0..rhs.bits() - 1).rev() {
+            if rhs.bit(i).eq(&ConstChoice::FALSE) {
+                x_1 = x_0.add(x_1, self);
+                x_0 = x_0.double();
+            } else {
+                x_0 = x_0.add(x_1, self);
+                x_1 = x_1.double();
+            }
+        }
+
+        x_0
+    }
+
+    pub fn ct_ladder(self, rhs: U256) -> MontgomeryPoint {
+        let mut x_0 = self;
+        let mut x_1 = x_0.double();
+
+        let mut prev: u8 = 0;
+        for i in (0..rhs.bits() - 1).rev() {
+            let curr: u8 = bool::from(rhs.bit(i)).into();
+            swap(&mut x_0, &mut x_1, curr ^ prev);
+            prev = curr;
+
+            x_1 = x_0.add(x_1, self);
+            x_0 = x_0.double();
+        }
+        swap(&mut x_0, &mut x_1, prev);
+
+        x_0
+    }
+}
+
+const SWAP_MASKS: [U256; 2] = [U256::ZERO, U256::from_be_hex("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF")];
+fn swap(a: &mut MontgomeryPoint, b: &mut MontgomeryPoint, c: u8) {
+    let m = SWAP_MASKS[c as usize];
+
+    let x = m & (a.x ^ b.x);
+    let z = m & (a.z ^ b.z);
+
+    *a = MontgomeryPoint{
+        x: a.x ^ x,
+        z: a.z ^ z,
+        curve: a.curve,
+    };
+
+    *b = MontgomeryPoint{
+        x: b.x ^ x,
+        z: b.z ^ z,
+        curve: b.curve,
+    }
 }
 
 impl Display for MontgomeryPoint {
@@ -122,20 +177,7 @@ impl Mul<U256> for MontgomeryPoint {
 
     /// (not)constant-time montgomery ladder
     fn mul(self, rhs: U256) -> Self::Output {
-        let mut x_0 = self;
+        self.ct_ladder(rhs)
-        let mut x_1 = x_0.double();
-
-        for i in (0..rhs.bits() - 1).rev() {
-            if rhs.bit(i).eq(&ConstChoice::FALSE) {
-                x_0 = x_0.double();
-                x_1 = x_0.add(x_1, self);
-            } else {
-                x_0 = x_0.add(x_1, self);
-                x_1 = x_1.double();
-            }
-        }
-
-        x_0
     }
 }
 
@@ -196,24 +238,37 @@ impl From<&Secret> for Public {
     }
 }
 
+impl From<[u8; 32]> for Public {
+    fn from(bytes: [u8; 32]) -> Public {
+        Public(U256::from_le_bytes(bytes))
+    }
+}
+
+impl Public {
+    pub fn as_bytes(&self) -> [u8; 32] {
+        self.0.to_le_bytes()
+    }
+}
+
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 pub struct Secret(U256);
 
 impl Secret {
+    #[cfg(feature = "rand")]
     pub fn random() -> Secret {
         Secret::from(U256::random(&mut OsRng))
     }
 
-    pub fn diffie_hellman(&self, public: Public) -> Secret {
+    pub fn diffie_hellman(&self, public: &Public) -> Secret {
-        let p = public.0; // public
+        let point = CURVE_25519.point(public.0);
-        let s = self.0; // secret
+        let p = point * self.0;
-        println!("p: {}\ns: {}", p, s);
+        
-
-        let point = CURVE_25519.point(p);
-        let p = point * s;
-
         Secret(p.get_x())
     }
+
+    pub fn to_bytes(&self) -> [u8; 32] {
+        self.0.to_le_bytes()
+    }
 }
 
 impl Display for Secret {
@@ -232,6 +287,8 @@ impl From<U256> for Secret {
 mod tests {
     use super::*;
 
+    const BIG_VALID_VALUE: U256 = U256::ONE.shl(251).add_mod(&U256::from_u64(1232039487203*8), &CURVE_25519.p);
+
     #[test]
     fn public_key_derivation() {
         let a = Secret(U256::from_u8(3));
@@ -246,7 +303,7 @@ mod tests {
         let s = Secret(U256::from_u8(2));
         let p = Public(U256::from_u8(3));
 
-        let x = s.diffie_hellman(p);
+        let x = s.diffie_hellman(&p);
 
         let x_k = U256::from_be_hex("0933dc6ed7122bdf2a5f1b0516e218b743868769778787cfa1ac0ae2dfa89891");
         assert_eq!(x.0, x_k);
@@ -254,16 +311,16 @@ mod tests {
 
     #[test]
     fn diffie_hellman() {
-        let a = Secret::random();
+        let a = Secret::from(U256::from_u32(123457689));
-        let b = Secret::random();
+        let b = Secret::from(U256::from_u64(987654321));
 
         // Public keys
         let pa = Public::from(&a);
         let pb = Public::from(&b);
 
         // Shared secret
-        let sa = a.diffie_hellman(pb);
+        let sa = a.diffie_hellman(&pb);
-        let sb = b.diffie_hellman(pa);
+        let sb = b.diffie_hellman(&pa);
 
         assert_eq!(sa, sb);
     }
@@ -273,25 +330,23 @@ mod tests {
         let o = CURVE_25519.point(U256::from_u8(9));
 
         // multiples of 8
-        let a = U256::from_u8(2 << 3);
+        let a = clamp_u256(U256::from_u32(1 << 9));
-        let b = U256::from_u8(3 << 3);
+        let b = clamp_u256(U256::from_u32(1 << 8));
 
-        let pa = (o * a).get_x();
+        let pa = o.ladder(a).get_x();
-        let pb = (o * b).get_x();
+        let pb = o.ladder(b).get_x();
 
         let pap = CURVE_25519.point(pa);
-        let sa = (pap * b).get_x();
+        let sa = pap.ladder(b).get_x();
 
         let pbp = CURVE_25519.point(pb);
-        let sb = (pbp * a).get_x();
+        let sb = pbp.ladder(a).get_x();
 
         assert_eq!(sa, sb);
     }
 
     #[test]
     fn montgomery_point_double() {
-        let p = &NonZero::new(CURVE_25519.p).unwrap();
-
         let point = CURVE_25519.point(U256::from_u8(2));
 
         let dbl = point.double();
@@ -301,6 +356,17 @@ mod tests {
         assert_eq!(dbl.get_x(), x_k);
     }
 
+    #[test]
+    fn montgomery_point_big_double() {
+        let point = CURVE_25519.point(BIG_VALID_VALUE);
+
+        let dbl = point.double();
+
+        let x_k = U256::from_be_hex("3f4d5612f4df4042680d41c7f5d9673fef3590c0c40bc04959848a88a117f311");
+
+        assert_eq!(dbl.get_x(), x_k);
+    }
+
     #[test]
     fn montgomery_point_get_x() {
         let point = MontgomeryPoint{
@@ -315,7 +381,6 @@ mod tests {
 
     #[test]
     fn montgomery_point_add() {
-        let p = &NonZero::new(CURVE_25519.p).unwrap();
         let point_p = CURVE_25519.point(U256::from_u8(2));
         let point_q = point_p.double();
 
@@ -325,6 +390,17 @@ mod tests {
         assert_eq!(sum.get_x(), x_k);
     }
 
+    #[test]
+    fn montgomery_point_big_add() {
+        let point_p = CURVE_25519.point(BIG_VALID_VALUE);
+        let point_q = point_p.double();
+
+        let sum = point_p.add(point_q, point_p);
+
+        let x_k = U256::from_be_hex("1431b13075c5e0c75f2aad0f604f7286fbeac04cdc6099a63c54cdd96cbd3333");
+        assert_eq!(sum.get_x(), x_k);
+    }
+
     #[test]
     fn montgomery_ladder_double() {
         let p = MontgomeryPoint{
@@ -344,14 +420,38 @@ mod tests {
     fn montgomery_ladder() {
         let p = CURVE_25519.point(U256::from_u8(2));
 
-        let n = U256::from_u8(3);
+        let n = U256::from_u8(4);
 
-        let a = p * n;
+        let a = p.ladder(n);
 
-        let x_k = U256::from_be_hex("56f0b3e7e53fed658bf39e9f8691055ba8a58a935482b00b21b5678c16261fec");
+        let x_k = U256::from_be_hex("2eceef1936e6df00c49e7aedac94446cc3b156165b50f247a15fdcee5e065582");
         assert_eq!(a.get_x(), x_k);
     }
 
+    #[test]
+    fn large_montgomery_ladder() {
+        let p = CURVE_25519.point(U256::from_u8(9));
+
+        let n = BIG_VALID_VALUE;
+
+        let a = p.ladder(n);
+
+        let x_k = U256::from_be_hex("0b379a815f36aca005dd1b19b1d483cd73fad06225c6ad927fde9316214a68c1");
+        assert_eq!(a.get_x(), x_k);
+    }
+
+    #[test]
+    fn constant_time_montgomery_ladder() {
+        let p = CURVE_25519.point(U256::from_u8(9));
+
+        let n = BIG_VALID_VALUE;
+
+        let a = p.ct_ladder(n);
+        let b = p.ladder(n);
+
+        assert_eq!(b.get_x(), a.get_x());
+    }
+
     #[test]
     fn montgomery_ladder_commutative() {
         let p = CURVE_25519.point(U256::from_u8(2));

tests/random.rs

@@ -1,14 +1,17 @@
+#![cfg(feature = "rand")]
+
 use diffie_hellman::{Public, Secret};
 
-fn main() {
+#[test]
+fn random_dh_key_exchange() {
     let a = Secret::random();
     let b = Secret::random();
 
     let pa = Public::from(&a);
     let pb = Public::from(&b);
 
-    let sa = a.diffie_hellman(pb);
+    let sa = a.diffie_hellman(&pb);
-    let sb = b.diffie_hellman(pa);
+    let sb = b.diffie_hellman(&pa);
 
     println!("{} == {}", sa, sb)
 }