From 0a66954a57450910742eb80a887c59532ac9d894 Mon Sep 17 00:00:00 2001
From: Aaron Feickert <66188213+AaronFeickert@users.noreply.github.com>
Date: Mon, 15 Jan 2024 16:39:45 -0600
Subject: [PATCH] Add constant-time trait bounds

---
 Cargo.toml                      |  4 +++-
 src/dhke.rs                     | 15 ++++++++++++---
 src/keys.rs                     |  7 +++++--
 src/ristretto/ristretto_keys.rs | 23 ++++++++++++++++++-----
 4 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 8be02791..0f16442e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,6 +24,7 @@ rand_core = { version = "0.6" , default-features = false}
 serde = { version = "1.0", optional = true }
 sha3 = { version = "0.10", default-features = false  }
 snafu = { version = "0.7", default-features = false}
+subtle = { version = "2.5.0", default-features = false }
 zeroize = {version = "1" , default-features = false}
 
 [dev-dependencies]
@@ -48,6 +49,7 @@ std = [
     "serde?/std",
     "sha3/std",
     "snafu/std",
+    "subtle/std",
     "tari_utilities/std",
     "zeroize/std",
 ]
@@ -61,4 +63,4 @@ crate-type = ["lib", "cdylib"]
 [[bench]]
 name = "benches"
 path = "benches/mod.rs"
-harness = false
\ No newline at end of file
+harness = false
diff --git a/src/dhke.rs b/src/dhke.rs
index 41c08c93..2282ab89 100644
--- a/src/dhke.rs
+++ b/src/dhke.rs
@@ -11,18 +11,19 @@
 
 use core::ops::Mul;
 
+use subtle::{Choice, ConstantTimeEq};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
 use crate::keys::PublicKey;
 
 /// The result of a Diffie-Hellman key exchange
-#[derive(Zeroize, ZeroizeOnDrop)]
+#[derive(PartialEq, Eq, Zeroize, ZeroizeOnDrop)]
 pub struct DiffieHellmanSharedSecret<P>(P)
-where P: Zeroize;
+where P: PublicKey;
 
 impl<P> DiffieHellmanSharedSecret<P>
 where
-    P: PublicKey + Zeroize,
+    P: PublicKey,
     for<'a> &'a <P as PublicKey>::K: Mul<&'a P, Output = P>,
 {
     /// Perform a Diffie-Hellman key exchange
@@ -36,6 +37,14 @@ where
     }
 }
 
+impl<P> ConstantTimeEq for DiffieHellmanSharedSecret<P>
+where P: PublicKey
+{
+    fn ct_eq(&self, other: &DiffieHellmanSharedSecret<P>) -> Choice {
+        self.0.ct_eq(&other.0)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use rand_core::OsRng;
diff --git a/src/keys.rs b/src/keys.rs
index b1692f26..5032ed55 100644
--- a/src/keys.rs
+++ b/src/keys.rs
@@ -9,6 +9,7 @@
 use core::ops::Add;
 
 use rand_core::{CryptoRng, RngCore};
+use subtle::ConstantTimeEq;
 use tari_utilities::{ByteArray, ByteArrayError};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
@@ -27,7 +28,7 @@ use zeroize::{Zeroize, ZeroizeOnDrop};
 /// let p = RistrettoPublicKey::from_secret_key(&k);
 /// ```
 pub trait SecretKey:
-    ByteArray + Clone + PartialEq + Eq + Add<Output = Self> + Default + Zeroize + ZeroizeOnDrop
+    ByteArray + Clone + ConstantTimeEq + PartialEq + Eq + Add<Output = Self> + Default + Zeroize + ZeroizeOnDrop
 {
     /// The length of the byte encoding of a key, in bytes
     const KEY_LEN: usize;
@@ -54,7 +55,9 @@ pub trait SecretKey:
 /// implementations need to implement this trait for them to be used in Tari.
 ///
 /// See [SecretKey](trait.SecretKey.html) for an example.
-pub trait PublicKey: ByteArray + Add<Output = Self> + Clone + PartialOrd + Ord + Default + Zeroize {
+pub trait PublicKey:
+    ByteArray + ConstantTimeEq + PartialEq + Eq + Add<Output = Self> + Clone + PartialOrd + Ord + Default + Zeroize
+{
     /// The length of the byte encoding of a key, in bytes
     const KEY_LEN: usize;
 
diff --git a/src/ristretto/ristretto_keys.rs b/src/ristretto/ristretto_keys.rs
index fc5d1efc..c18fb2d4 100644
--- a/src/ristretto/ristretto_keys.rs
+++ b/src/ristretto/ristretto_keys.rs
@@ -21,6 +21,7 @@ use curve25519_dalek::{
 use digest::{consts::U64, Digest};
 use once_cell::sync::OnceCell;
 use rand_core::{CryptoRng, RngCore};
+use subtle::ConstantTimeEq;
 use tari_utilities::{hex::Hex, ByteArray, ByteArrayError, Hashable};
 use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};
 
@@ -51,7 +52,7 @@ use crate::{
 /// let _k2 = RistrettoSecretKey::from_hex(&"100000002000000030000000040000000");
 /// let _k3 = RistrettoSecretKey::random(&mut rng);
 /// ```
-#[derive(Eq, Clone, Default, Zeroize, ZeroizeOnDrop)]
+#[derive(Clone, Default, Zeroize, ZeroizeOnDrop)]
 pub struct RistrettoSecretKey(pub(crate) Scalar);
 
 #[cfg(feature = "borsh")]
@@ -132,12 +133,20 @@ impl Hash for RistrettoSecretKey {
     }
 }
 
+impl ConstantTimeEq for RistrettoSecretKey {
+    fn ct_eq(&self, other: &Self) -> subtle::Choice {
+        self.0.ct_eq(&other.0)
+    }
+}
+
 impl PartialEq for RistrettoSecretKey {
     fn eq(&self, other: &Self) -> bool {
-        self.0.eq(&other.0)
+        self.ct_eq(other).into()
     }
 }
 
+impl Eq for RistrettoSecretKey {}
+
 //----------------------------------   RistrettoSecretKey Debug --------------------------------------------//
 impl fmt::Debug for RistrettoSecretKey {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -423,6 +432,12 @@ impl fmt::Display for RistrettoPublicKey {
     }
 }
 
+impl ConstantTimeEq for RistrettoPublicKey {
+    fn ct_eq(&self, other: &Self) -> subtle::Choice {
+        self.point.ct_eq(&other.point)
+    }
+}
+
 impl RistrettoPublicKey {
     // Formats a 64 char hex string to a given width.
     // If w >= 64, we pad the result.
@@ -471,9 +486,7 @@ impl fmt::Debug for RistrettoPublicKey {
 
 impl PartialEq for RistrettoPublicKey {
     fn eq(&self, other: &RistrettoPublicKey) -> bool {
-        // Although this is slower than `self.compressed == other.compressed`, expanded point comparison is an equal
-        // time comparison
-        self.point == other.point
+        self.ct_eq(other).into()
     }
 }