Merge pull request #3 from Decentrailzed-Communication-System/change_from_dalek

Changed curve and added bench

Cargo.toml

@@ -13,8 +13,8 @@ license = "MIT"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-x25519-dalek = "1.1.1"
-rand_core = "0.5"
+p256 = {version = "0.8.1", features = ["zeroize", "ecdh", "arithmetic", "pem"]}
+rand_core = {version = "0.6", features = ["getrandom"]}
 hkdf = "0.11.0"
 hmac = "0.11.0"
 aes-gcm-siv = {version = "0.10.0"}
@@ -25,6 +25,13 @@ serde_bytes = "0.11.5"
 bincode = "1.3.3"
 hashbrown = "0.11.2"
 
+[dev-dependencies]
+criterion = "0.3.4"
+
+[[bench]]
+name = "double_ratchet_bench"
+harness = false
+
 [profile.release]
 lto = true
 opt-level = 3

benches/double_ratchet_bench.rs

@@ -0,0 +1,111 @@
+use double_ratchet_2::ratchet::{Ratchet, RatchetEncHeader};
+use criterion::{Criterion, criterion_main, criterion_group};
+
+fn ratchet_enc_single() {
+    let sk = [1; 32];
+    let (mut bob_ratchet, public_key) = Ratchet::init_bob(sk);
+    let mut alice_ratchet = Ratchet::init_alice(sk, public_key);
+    let data = include_bytes!("../src/header.rs").to_vec();
+    let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce, b"");
+}
+
+fn criterion_benchmark_1(c: &mut Criterion) {
+    c.bench_function("Ratchet Enc Single", |b| b.iter(|| ratchet_enc_single()));
+}
+
+fn ratchet_enc_skip() {
+    let sk = [1; 32];
+    let (mut bob_ratchet, public_key) = Ratchet::init_bob(sk);
+    let mut alice_ratchet = Ratchet::init_alice(sk, public_key);
+    let data = include_bytes!("../src/header.rs").to_vec();
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let (header2, encrypted2, nonce2) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2, b"");
+    let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1, b"");
+}
+
+fn criterion_benchmark_2(c: &mut Criterion) {
+    c.bench_function("Ratchet Enc Skip", |b| b.iter(|| ratchet_enc_skip()));
+}
+
+fn ratchet_encryt_decrypt_four() {
+    let sk = [1; 32];
+    let data = include_bytes!("../src/dh.rs").to_vec();
+    let (mut bob_ratchet, public_key) = Ratchet::init_bob(sk);
+    let mut alice_ratchet = Ratchet::init_alice(sk, public_key);
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1, b"");
+    let (header2, encrypted2, nonce2) = bob_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2, b"");
+}
+
+fn criterion_benchmark_3(c: &mut Criterion) {
+    c.bench_function("Ratchet Dec Four", |b| b.iter(|| ratchet_encryt_decrypt_four()));
+}
+
+fn ratchet_ench_enc_single() {
+    let sk = [1; 32];
+    let shared_hka = [2; 32];
+    let shared_nhkb = [3; 32];
+    let (mut bob_ratchet, public_key) = RatchetEncHeader::init_bob(sk,
+                                                                   shared_hka,
+                                                                   shared_nhkb);
+    let mut alice_ratchet = RatchetEncHeader::init_alice(sk,
+                                                         public_key,
+                                                         shared_hka,
+                                                         shared_nhkb);
+    let data = include_bytes!("../src/header.rs").to_vec();
+    let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce, b"");
+    assert_eq!(data, decrypted)
+}
+
+fn criterion_benchmark_4(c: &mut Criterion) {
+    c.bench_function("Encrypted Header Ratchet Enc Single", |b| b.iter(|| ratchet_ench_enc_single()));
+}
+
+fn ratchet_ench_enc_skip() {
+    let sk = [1; 32];
+    let shared_hka = [2; 32];
+    let shared_nhkb = [3; 32];
+    let (mut bob_ratchet, public_key) = RatchetEncHeader::init_bob(sk,
+                                                                   shared_hka,
+                                                                   shared_nhkb);
+    let mut alice_ratchet = RatchetEncHeader::init_alice(sk,
+                                                         public_key,
+                                                         shared_hka,
+                                                         shared_nhkb);
+    let data = include_bytes!("../src/header.rs").to_vec();
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let (header2, encrypted2, nonce2) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2, b"");
+    let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1, b"");
+}
+
+fn criterion_benchmark_5(c: &mut Criterion) {
+    c.bench_function("Encrypted Header Ratchet Enc Skip", |b| b.iter(|| ratchet_ench_enc_skip()));
+}
+
+fn ratchet_ench_decrypt_four() {
+    let sk = [1; 32];
+    let shared_hka = [2; 32];
+    let shared_nhkb = [3; 32];
+    let (mut bob_ratchet, public_key) = RatchetEncHeader::init_bob(sk,
+                                                                   shared_hka,
+                                                                   shared_nhkb);
+    let mut alice_ratchet = RatchetEncHeader::init_alice(sk, public_key, shared_hka, shared_nhkb);
+    let data = include_bytes!("../src/dh.rs").to_vec();
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1, b"");
+    let (header2, encrypted2, nonce2) = bob_ratchet.ratchet_encrypt(&data, b"");
+    let _decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2, b"");
+}
+
+fn criterion_benchmark_6(c: &mut Criterion) {
+    c.bench_function("Encrypted Header Ratchet Dec Four", |b| b.iter(|| ratchet_ench_decrypt_four()));
+}
+
+criterion_group!(without_enc_headerd, criterion_benchmark_1, criterion_benchmark_2, criterion_benchmark_3);
+criterion_group!(with_enc_headerd, criterion_benchmark_4, criterion_benchmark_5, criterion_benchmark_6);
+criterion_main!(without_enc_headerd, with_enc_headerd);

src/dh.rs

@@ -1,19 +1,30 @@
-use x25519_dalek::{PublicKey, SharedSecret, StaticSecret};
 use rand_core::OsRng;
 use core::fmt::{Debug, Formatter};
 use core::fmt;
+use p256::PublicKey as PublicKey;
+use p256::ecdh::SharedSecret;
+use p256::SecretKey;
+use alloc::vec::Vec;
+use alloc::string::ToString;
+use p256::elliptic_curve::ecdh::diffie_hellman;
 
 pub struct DhKeyPair {
-    pub private_key: StaticSecret,
+    pub private_key: SecretKey,
     pub public_key: PublicKey,
 }
 
+impl DhKeyPair {
+    fn ex_public_key_bytes(&self) -> Vec<u8> {
+        self.public_key.to_string().as_bytes().to_vec()
+    }
+}
+
 impl PartialEq for DhKeyPair {
     fn eq(&self, other: &Self) -> bool {
         if self.private_key.to_bytes() != other.private_key.to_bytes() {
             return false
         }
-        if self.public_key.to_bytes() != other.public_key.to_bytes() {
+        if self.ex_public_key_bytes() != other.ex_public_key_bytes() {
             return false
         }
         true
@@ -24,7 +35,7 @@ impl Debug for DhKeyPair {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("DhKeyPair")
             .field("private_key", &self.private_key.to_bytes())
-            .field("public_key", &self.public_key.to_bytes())
+            .field("public_key", &self.ex_public_key_bytes())
             .finish()
     }
 }
@@ -37,8 +48,8 @@ impl Default for DhKeyPair {
 
 impl DhKeyPair {
     pub fn new() -> Self {
-        let secret = StaticSecret::new(OsRng);
-        let public = PublicKey::from(&secret);
+        let secret = SecretKey::random(&mut OsRng);
+        let public = PublicKey::from_secret_scalar(&secret.secret_scalar());
         DhKeyPair {
             private_key: secret,
             public_key: public,
@@ -46,7 +57,7 @@ impl DhKeyPair {
     }
 
     pub fn key_agreement(&self, public_key: &PublicKey) -> SharedSecret {
-        self.private_key.diffie_hellman(public_key)
+        diffie_hellman(self.private_key.secret_scalar(), public_key.as_affine())
     }
 }
 

src/header.rs

@@ -1,4 +1,4 @@
-use x25519_dalek::PublicKey;
+use p256::PublicKey;
 use crate::dh::DhKeyPair;
 use alloc::vec::Vec;
 use serde::{Serialize, Deserialize};
@@ -8,6 +8,8 @@ use aes_gcm_siv::aead::{NewAead, AeadInPlace};
 
 #[cfg(test)]
 use crate::dh::gen_key_pair;
+use alloc::string::{ToString, String};
+use core::str::FromStr;
 
 #[derive(Debug, Copy, Clone)]
 pub struct Header {
@@ -20,7 +22,7 @@ pub struct Header {
 struct ExHeader {
     #[serde(with = "serde_bytes")]
     ad: Vec<u8>,
-    public_key: [u8; 32],
+    public_key: Vec<u8>,
     pn: usize,
     n: usize
 }
@@ -39,7 +41,7 @@ impl Header {
     pub fn concat(&self, ad: &[u8]) -> Vec<u8> {
         let ex_header = ExHeader {
             ad: ad.to_vec(),
-            public_key: self.public_key.to_bytes(),
+            public_key: self.public_key.to_string().as_bytes().to_vec(),
             pn: self.pn,
             n: self.n
         };
@@ -72,13 +74,17 @@ impl Header {
         };
         Some(Header::from(buffer))
     }
+    pub fn ex_public_key_bytes(&self) -> Vec<u8> {
+        self.public_key.to_string().as_bytes().to_vec()
+    }
 }
 
 impl From<Vec<u8>> for Header {
     fn from(d: Vec<u8>) -> Self {
         let ex_header: ExHeader = bincode::deserialize(&d).unwrap();
+        let public_key_string = String::from_utf8(ex_header.public_key).unwrap();
         Header {
-            public_key: PublicKey::from(ex_header.public_key),
+            public_key: PublicKey::from_str(&public_key_string).unwrap(),
             pn: ex_header.pn,
             n: ex_header.n,
         }
@@ -88,8 +94,9 @@ impl From<Vec<u8>> for Header {
 impl From<&[u8]> for Header {
     fn from(d: &[u8]) -> Self {
         let ex_header: ExHeader = bincode::deserialize(d).unwrap();
+        let public_key_string = String::from_utf8(ex_header.public_key).unwrap();
         Header {
-            public_key: PublicKey::from(ex_header.public_key),
+            public_key: PublicKey::from_str(&public_key_string).unwrap(),
             pn: ex_header.pn,
             n: ex_header.n,
         }

src/kdf_root.rs

@@ -1,4 +1,4 @@
-use x25519_dalek::SharedSecret;
+
 use hkdf::Hkdf;
 
 #[cfg(feature = "ring")]
@@ -11,6 +11,7 @@ use core::convert::TryInto;
 
 #[cfg(test)]
 use crate::dh::gen_shared_secret;
+use p256::ecdh::SharedSecret;
 
 pub fn kdf_rk(rk: &[u8; 32], dh_out: &SharedSecret) -> ([u8; 32], [u8; 32]) {
     let h = Hkdf::<Sha512>::new(Some(rk), dh_out.as_bytes());

src/lib.rs

@@ -130,7 +130,7 @@
 
 extern crate alloc;
 
-pub use x25519_dalek::PublicKey;
+pub use p256::PublicKey;
 
 mod aead;
 mod dh;

src/ratchet.rs

@@ -2,13 +2,14 @@
 //!
 
 use crate::dh::DhKeyPair;
-use x25519_dalek::PublicKey;
+use p256::PublicKey;
 use hashbrown::HashMap;
 use crate::kdf_root::{kdf_rk, kdf_rk_he};
 use crate::header::Header;
 use alloc::vec::Vec;
 use crate::kdf_chain::kdf_ck;
 use crate::aead::{encrypt, decrypt};
+use alloc::string::ToString;
 
 const MAX_SKIP: usize = 100;
 
@@ -24,7 +25,7 @@ pub struct Ratchet {
     ns: usize,
     nr: usize,
     pn: usize,
-    mkskipped: HashMap<(PublicKey, usize), [u8; 32]>,
+    mkskipped: HashMap<(Vec<u8>, usize), [u8; 32]>,
 }
 
 impl Ratchet {
@@ -75,10 +76,10 @@ impl Ratchet {
     }
 
     fn try_skipped_message_keys(&mut self, header: &Header, ciphertext: &[u8], nonce: &[u8; 12], ad: &[u8]) -> Option<Vec<u8>> {
-        if self.mkskipped.contains_key(&(header.public_key, header.n)) {
-            let mk = *self.mkskipped.get(&(header.public_key, header.n))
+        if self.mkskipped.contains_key(&(header.ex_public_key_bytes(), header.n)) {
+            let mk = *self.mkskipped.get(&(header.ex_public_key_bytes(), header.n))
                 .unwrap();
-            self.mkskipped.remove(&(header.public_key, header.n)).unwrap();
+            self.mkskipped.remove(&(header.ex_public_key_bytes(), header.n)).unwrap();
             Some(decrypt(&mk, ciphertext, &header.concat(ad), nonce))
         } else {
             None
@@ -94,7 +95,7 @@ impl Ratchet {
                 while self.nr < until {
                     let (ckr, mk) = kdf_ck(&d);
                     self.ckr = Some(ckr);
-                    self.mkskipped.insert((self.dhr.unwrap(), self.nr), mk);
+                    self.mkskipped.insert((self.dhr.unwrap().to_string().as_bytes().to_vec(), self.nr), mk);
                     self.nr += 1
                 }
                 Ok(())