From 406afe08d370a8a288e6c982dfc02f3418bc5b27 Mon Sep 17 00:00:00 2001
From: Hannes <hannes.furmans@web.de>
Date: Fri, 18 Jun 2021 15:15:55 +0200
Subject: [PATCH] Improved coverage

---
 src/dh.rs     | 33 ++++++++++++++++++++++++++++++++-
 src/header.rs | 37 +++++++++++++++++++++++++++++++++++--
 tests/mod.rs  | 39 ++++++++++++++++++++++++++++++++++++++-
 3 files changed, 105 insertions(+), 4 deletions(-)

diff --git a/src/dh.rs b/src/dh.rs
index 17863c4..aaac3f0 100644
--- a/src/dh.rs
+++ b/src/dh.rs
@@ -10,6 +10,7 @@ use p256::elliptic_curve::ecdh::diffie_hellman;
 
 use zeroize::Zeroize;
 
+#[derive(Clone)]
 pub struct DhKeyPair {
     pub private_key: SecretKey,
     pub public_key: PublicKey,
@@ -65,7 +66,7 @@ impl Default for DhKeyPair {
 impl DhKeyPair {
     pub fn new() -> Self {
         let secret = SecretKey::random(&mut OsRng);
-        let public = PublicKey::from_secret_scalar(&secret.to_secret_scalar());
+        let public = secret.public_key();
         DhKeyPair {
             private_key: secret,
             public_key: public,
@@ -92,6 +93,7 @@ pub fn gen_key_pair() -> DhKeyPair {
 #[cfg(test)]
 mod tests {
     use crate::dh::DhKeyPair;
+    use alloc::string::ToString;
 
     #[test]
     fn key_generation() {
@@ -108,4 +110,33 @@ mod tests {
         let bob_shared_secret = bob_pair.key_agreement(&alice_pair.public_key);
         assert_eq!(alice_shared_secret.as_bytes(), bob_shared_secret.as_bytes())
     }
+
+    #[test]
+    fn ex_public_key() {
+        let key_pair = DhKeyPair::new();
+        let public_key_bytes = key_pair.ex_public_key_bytes();
+        let extracted_pk = key_pair.public_key.to_string().as_bytes().to_vec();
+        assert_eq!(extracted_pk, public_key_bytes)
+    }
+
+    #[test]
+    fn nq_key_pair() {
+        let key_pair1 = DhKeyPair::new();
+        let mut key_pair2 = DhKeyPair::new();
+        key_pair2.private_key = key_pair1.private_key.clone();
+        assert_ne!(key_pair1, key_pair2)
+    }
+
+    #[test]
+    fn eq_key_pair() {
+        let key_pair1 = DhKeyPair::new();
+        let key_pair2 = key_pair1.clone();
+        assert_eq!(key_pair1, key_pair2)
+    }
+
+    #[test]
+    fn test_format() {
+        let key_pair = DhKeyPair::default();
+        let _str = alloc::format!("{:?}", key_pair);
+    }
 }
\ No newline at end of file
diff --git a/src/header.rs b/src/header.rs
index 6d13f99..7402040 100644
--- a/src/header.rs
+++ b/src/header.rs
@@ -130,9 +130,10 @@ pub fn gen_header() -> Header {
     Header::new(&dh_pair, pn, n)
 }
 
+
 #[cfg(test)]
 mod tests {
-    use crate::header::{gen_header, Header};
+    use crate::header::{gen_header, Header, ExHeader};
     use crate::kdf_chain::gen_mk;
     use crate::aead::{encrypt, decrypt};
 
@@ -141,7 +142,7 @@ mod tests {
         let ad = b"";
         let header = gen_header();
         let serialized = header.concat(ad);
-        let created = Header::from(serialized);
+        let created = Header::from(serialized.as_slice());
         assert_eq!(header, created)
     }
 
@@ -155,4 +156,36 @@ mod tests {
         let decrypted = decrypt(&mk, &encrypted, &header_data, &nonce);
         assert_eq!(decrypted, data.to_vec())
     }
+
+    #[test]
+    fn test_eq_header() {
+        let header1 = gen_header();
+        let header2 = gen_header();
+        assert_ne!(header1, header2)
+    }
+
+    #[test]
+    fn debug_header() {
+        let header = gen_header();
+        let _string = alloc::format!("{:?}", header);
+    }
+
+    #[test]
+    fn gen_ex_header() {
+        let ex_header = ExHeader {
+            ad: alloc::vec![0],
+            public_key: alloc::vec![1],
+            pn: 0,
+            n: 0
+        };
+        let _string = alloc::format!("{:?}", ex_header);
+    }
+
+    #[test]
+    fn dec_header() {
+        let header = gen_header();
+        let (encrypted, nonce) = header.encrypt(&[0; 32], &[0]);
+        let decrypted = Header::decrypt(&Some([1_u8; 32]), &encrypted, &nonce);
+        assert_eq!(None, decrypted)
+    }
 }
\ No newline at end of file
diff --git a/tests/mod.rs b/tests/mod.rs
index 3a1073c..0b2e8ec 100644
--- a/tests/mod.rs
+++ b/tests/mod.rs
@@ -1,4 +1,5 @@
 use double_ratchet_2::ratchet::{Ratchet, RatchetEncHeader};
+extern crate alloc;
 
 #[test]
 fn ratchet_init() {
@@ -135,4 +136,40 @@ fn ratchet_ench_decrypt_four() {
     let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2, b"");
     let comp_res = decrypted1 == data && decrypted2 == data;
     assert!(comp_res)
-}
\ No newline at end of file
+}
+
+#[test]
+#[should_panic]
+fn ratchet_ench_enc_skip_panic() {
+    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 mut headers = alloc::vec![];
+    let mut encrypteds = alloc::vec![];
+    let mut nonces = alloc::vec![];
+    let mut decrypteds = alloc::vec![];
+    for _ in 0..200 {
+        let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data, b"");
+        headers.push(header);
+        encrypteds.push(encrypted);
+        nonces.push(nonce);
+    }
+    headers.reverse();
+    encrypteds.reverse();
+    nonces.reverse();
+    for idx in 0..200 {
+        let header = headers.get(idx).unwrap();
+        let encrypted = encrypteds.get(idx).unwrap();
+        let nonce = nonces.get(idx).unwrap();
+        let decrypted = bob_ratchet.ratchet_decrypt(header, encrypted, nonce, b"");
+        decrypteds.push(decrypted);
+    }
+}