Merge pull request #2 from Decentrailzed-Communication-System/encrypted_header

Adding encrypted headers

Cargo.toml

@@ -6,10 +6,9 @@ homepage = "https://github.com/Decentrailzed-Communication-System/double-ratchet
 repository = "https://github.com/Decentrailzed-Communication-System/double-ratchet-2"
 readme = "README.md"
 keywords = ["double-ratchet", "crypto", "cryptography", "signal"]
-version = "0.1.3"
+version = "0.2.0"
 edition = "2018"
 license = "MIT"
-license-file = "LICENSE"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

src/aead.rs

@@ -1,30 +1,30 @@
 use aes_gcm_siv::{Key, Aes256GcmSiv, Nonce};
 use aes_gcm_siv::aead::{NewAead, AeadInPlace};
 use alloc::vec::Vec;
+use rand_core::{OsRng, RngCore};
 
-const CONSTANT_NONCE: &[u8] = b"Super Noncel";
+pub fn encrypt(mk: &[u8; 32], plaintext: &[u8], associated_data: &[u8]) -> (Vec<u8>, [u8; 12]) {
-
-pub fn encrypt(mk: &[u8; 32], plaintext: &[u8], associated_data: &[u8]) -> Vec<u8> {
     let key = Key::from_slice(mk);
     let cipher = Aes256GcmSiv::new(key);
-
+    let mut nonce_data = [0_u8; 12];
-    let nonce = Nonce::from_slice(&CONSTANT_NONCE);
+    OsRng::fill_bytes(&mut OsRng, &mut nonce_data);
+    let nonce = Nonce::from_slice(&nonce_data);
     let mut buffer = Vec::new();
     buffer.extend_from_slice(plaintext);
 
     cipher.encrypt_in_place(nonce, associated_data, &mut buffer)
         .expect("Encryption failed");
-    buffer
+    (buffer, nonce_data)
 }
 
-pub fn decrypt(mk: &[u8; 32], ciphertext: &[u8], associated_data: &[u8]) -> Vec<u8> {
+pub fn decrypt(mk: &[u8; 32], ciphertext: &[u8], associated_data: &[u8], nonce: &[u8; 12]) -> Vec<u8> {
     let key = Key::from_slice(mk);
     let cipher = Aes256GcmSiv::new(key);
 
-    let nonce = Nonce::from_slice(&CONSTANT_NONCE);
+    let nonce = Nonce::from_slice(nonce);
     let mut buffer = Vec::new();
     buffer.extend_from_slice(ciphertext);
-    cipher.decrypt_in_place(nonce, associated_data, &mut buffer).expect("Decryption failure");
+    cipher.decrypt_in_place(nonce, associated_data, &mut buffer).expect("Decryption failure {}");
     buffer
 }
 
@@ -38,8 +38,8 @@ mod tests {
         let test_data = include_bytes!("aead.rs").to_vec();
         let associated_data = include_bytes!("lib.rs").to_vec();
         let mk = gen_mk();
-        let ciphertext = encrypt(&mk, &test_data, &associated_data);
+        let (ciphertext, nonce) = encrypt(&mk, &test_data, &associated_data);
-        let plaintext = decrypt(&mk, &ciphertext, &associated_data);
+        let plaintext = decrypt(&mk, &ciphertext, &associated_data, &nonce);
         assert_eq!(test_data, plaintext)
     }
 }

src/header.rs

@@ -2,6 +2,9 @@ use x25519_dalek::PublicKey;
 use crate::dh::DhKeyPair;
 use alloc::vec::Vec;
 use serde::{Serialize, Deserialize};
+use crate::aead::encrypt;
+use aes_gcm_siv::{Key, Nonce, Aes256GcmSiv};
+use aes_gcm_siv::aead::{NewAead, AeadInPlace};
 
 #[cfg(test)]
 use crate::dh::gen_key_pair;
@@ -22,7 +25,7 @@ struct ExHeader {
 
 // Message Header
 impl Header {
-    #[doc(hidden)]
+    // #[doc(hidden)]
     pub fn new(dh_pair: &DhKeyPair, pn: usize, n: usize) -> Self {
         Header {
             public_key: dh_pair.public_key,
@@ -30,7 +33,7 @@ impl Header {
             n,
         }
     }
-    #[doc(hidden)]
+    // #[doc(hidden)]
     pub fn concat(&self) -> Vec<u8> {
         let ex_header = ExHeader{
             public_key: self.public_key.to_bytes(),
@@ -39,6 +42,33 @@ impl Header {
         };
         bincode::serialize(&ex_header).expect("Failed to serialize Header")
     }
+
+    pub fn encrypt(&self, hk: &[u8; 32]) -> (Vec<u8>, [u8; 12]) {
+        let header_data = self.concat();
+        encrypt(hk, &header_data, b"")
+    }
+
+    pub fn decrypt(hk: &Option<[u8; 32]>, ciphertext: &[u8], nonce: &[u8; 12]) -> Option<Self> {
+        let key_d = match hk {
+            None => {
+                return None
+            },
+            Some(d) => d
+        };
+        let key = Key::from_slice(key_d);
+        let cipher = Aes256GcmSiv::new(key);
+
+        let nonce = Nonce::from_slice(nonce);
+        let mut buffer = Vec::new();
+        buffer.extend_from_slice(ciphertext);
+        match cipher.decrypt_in_place(nonce, b"", &mut buffer) {
+            Ok(_) => {}
+            Err(_) => {
+                return None
+            }
+        };
+        Some(Header::from(buffer))
+    }
 }
 
 impl From<Vec<u8>> for Header {
@@ -63,9 +93,9 @@ impl From<&[u8]> for Header {
     }
 }
 
-impl Into<Vec<u8>> for Header {
+impl From<Header> for Vec<u8> {
-    fn into(self) -> Vec<u8> {
+    fn from(s: Header) -> Self {
-        self.concat()
+        s.concat()
     }
 }
 
@@ -108,8 +138,8 @@ mod tests {
         let mk = gen_mk();
         let header_data = header.concat();
         let data = include_bytes!("aead.rs");
-        let encrypted = encrypt(&mk, data, &header_data);
+        let (encrypted, nonce) = encrypt(&mk, data, &header_data);
-        let decrypted = decrypt(&mk, &encrypted, &header_data);
+        let decrypted = decrypt(&mk, &encrypted, &header_data, &nonce);
         assert_eq!(decrypted, data.to_vec())
     }
 }

src/kdf_root.rs

@@ -24,6 +24,20 @@ pub fn kdf_rk(rk: &[u8; 32], dh_out: &SharedSecret) -> ([u8; 32], [u8; 32]) {
          .expect("Incorrect length"))
 }
 
+pub fn kdf_rk_he(rk: &[u8; 32], dh_out: &SharedSecret) -> ([u8; 32], [u8; 32], [u8; 32]) {
+    let h = Hkdf::<Sha512>::new(Some(rk), dh_out.as_bytes());
+    let mut okm = [0u8; 96];
+    let info = b"Root Key Generator";
+    h.expand(info, &mut okm).unwrap();
+    let (rk, a) = okm.split_at(32);
+    let (ck, nhk) = a.split_at(32);
+    (
+        rk.try_into().expect("Wrong length"),
+        ck.try_into().expect("Wrong length"),
+        nhk.try_into().expect("Wrong length")
+    )
+}
+
 #[cfg(test)]
 pub fn gen_ck() -> [u8; 32] {
     let shared_secret = gen_shared_secret();

src/lib.rs

@@ -18,8 +18,8 @@
 //! let mut alice_ratchet = Ratchet::init_alice(sk, public_key);      // Creating Alice Ratchet with Bobs PublicKey
 //! let data = b"Hello World".to_vec();                               // Data to be encrypted
 //!
-//! let (header, encrypted) = alice_ratchet.ratchet_encrypt(&data);   // Encrypting message with Alice Ratchet (Alice always needs to send the first message)
+//! let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data);   // Encrypting message with Alice Ratchet (Alice always needs to send the first message)
-//! let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted); // Decrypt message with Bobs Ratchet
+//! let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce); // Decrypt message with Bobs Ratchet
 //! assert_eq!(data, decrypted)
 //! ```
 //!
@@ -32,11 +32,11 @@
 //! let mut alice_ratchet = Ratchet::init_alice(sk, public_key);      // Creating Alice Ratchet with Bobs PublicKey
 //! let data = b"Hello World".to_vec();                               // Data to be encrypted
 //!
-//! let (header1, encrypted1) = alice_ratchet.ratchet_encrypt(&data); // Lost message
+//! let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data); // Lost message
-//! let (header2, encrypted2) = alice_ratchet.ratchet_encrypt(&data); // Successful message
+//! let (header2, encrypted2, nonce2) = alice_ratchet.ratchet_encrypt(&data); // Successful message
 //!
-//! let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2); // Decrypting second message first
+//! let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2); // Decrypting second message first
-//! let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1); // Decrypting latter message
+//! let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1); // Decrypting latter message
 //!
 //! let comp = decrypted1 == data && decrypted2 == data;
 //! assert!(comp);
@@ -51,7 +51,7 @@
 //! let (mut bob_ratchet, _) = Ratchet::init_bob(sk);
 //! let data = b"Hello World".to_vec();
 //!
-//! let (_, _) = bob_ratchet.ratchet_encrypt(&data);
+//! let (_, _, _) = bob_ratchet.ratchet_encrypt(&data);
 //! ```
 //!
 //! ## Encryption after recieving initial message
@@ -66,11 +66,11 @@
 //!
 //! let data = b"Hello World".to_vec();
 //!
-//! let (header1, encrypted1) = alice_ratchet.ratchet_encrypt(&data);
+//! let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data);
-//! let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1);
+//! let _decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1);
 //!
-//! let (header2, encrypted2) = bob_ratchet.ratchet_encrypt(&data);
+//! let (header2, encrypted2, nonce2) = bob_ratchet.ratchet_encrypt(&data);
-//! let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2);
+//! let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2);
 //!
 //! assert_eq!(data, decrypted2);
 //! ```
@@ -83,12 +83,29 @@
 //! # let (mut bob_ratchet, public_key) = Ratchet::init_bob(sk);
 //! # let mut alice_ratchet = Ratchet::init_alice(sk, public_key);
 //! # let data = b"hello World".to_vec();
-//! # let (header, _) = alice_ratchet.ratchet_encrypt(&data);
+//! # let (header, _, _) = alice_ratchet.ratchet_encrypt(&data);
 //! let header_bytes: Vec<u8> = header.clone().into();
 //! let header_const = Header::from(header_bytes);
 //! assert_eq!(header, header_const);
 //! ```
 //!
+//! # Example Ratchet with encrypted headers
+//!
+//! ```
+//! use double_ratchet_2::ratchet::RatchetEncHeader;
+//! let sk = [0; 32];
+//! let shared_hka = [1; 32];
+//! let shared_nhkb = [2; 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 = b"Hello World".to_vec();
+//!
+//! let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data);
+//! let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce);
+//! assert_eq!(data, decrypted)
+//! ```
+//!
 //! # Features
 //!
 //! Currently the crate only supports one feature: ring. If feature is enabled the crate switches
@@ -97,7 +114,7 @@
 //!
 //! TODO:
 //! - [x] Standard Double Ratchet
-//! - [ ] [Double Ratchet with encrypted headers][3]
+//! - [x] [Double Ratchet with encrypted headers][3]
 //!
 //! [1]: https://signal.org/docs/specifications/doubleratchet/
 //! [2]: https://signal.org/docs/specifications/doubleratchet/#recommended-cryptographic-algorithms
@@ -115,7 +132,6 @@ mod dh;
 mod kdf_root;
 mod kdf_chain;
 
-/// Providing essential functions
 pub mod ratchet;
 
 /// Message Header

src/ratchet.rs

@@ -1,7 +1,10 @@
+//! Encryption with encrypted Headers
+//!
+
 use crate::dh::DhKeyPair;
 use x25519_dalek::PublicKey;
 use hashbrown::HashMap;
-use crate::kdf_root::kdf_rk;
+use crate::kdf_root::{kdf_rk, kdf_rk_he};
 use crate::header::Header;
 use alloc::vec::Vec;
 use crate::kdf_chain::kdf_ck;
@@ -9,6 +12,8 @@ use crate::aead::{encrypt, decrypt};
 
 const MAX_SKIP: usize = 100;
 
+type HeaderNonceCipherNonce = ((Vec<u8>, [u8; 12]), Vec<u8>, [u8; 12]);
+
 /// Object Representing Ratchet
 pub struct Ratchet {
     dhs: DhKeyPair,
@@ -60,21 +65,21 @@ impl Ratchet {
     }
 
     /// Encrypt Plaintext with [Ratchet]. Returns Message [Header] and ciphertext.
-    pub fn ratchet_encrypt(&mut self, plaintext: &[u8]) -> (Header, Vec<u8>) {
+    pub fn ratchet_encrypt(&mut self, plaintext: &[u8]) -> (Header, Vec<u8>, [u8; 12]) {
         let (cks, mk) = kdf_ck(&self.cks.unwrap());
         self.cks = Some(cks);
         let header = Header::new(&self.dhs, self.pn, self.ns);
         self.ns += 1;
-        let encrypted_data = encrypt(&mk, plaintext, &header.concat());
+        let (encrypted_data, nonce) = encrypt(&mk, plaintext, &header.concat());
-        (header, encrypted_data)
+        (header, encrypted_data, nonce)
     }
 
-    fn try_skipped_message_keys(&mut self, header: &Header, ciphertext: &[u8]) -> Option<Vec<u8>> {
+    fn try_skipped_message_keys(&mut self, header: &Header, ciphertext: &[u8], nonce: &[u8; 12]) -> Option<Vec<u8>> {
         if self.mkskipped.contains_key(&(header.public_key, header.n)) {
             let mk = *self.mkskipped.get(&(header.public_key, header.n))
                 .unwrap();
             self.mkskipped.remove(&(header.public_key, header.n)).unwrap();
-            Some(decrypt(&mk, ciphertext, &header.concat()))
+            Some(decrypt(&mk, ciphertext, &header.concat(), nonce))
         } else {
             None
         }
@@ -99,8 +104,8 @@ impl Ratchet {
     }
 
     /// Decrypt ciphertext with ratchet. Requires Header. Returns plaintext.
-    pub fn ratchet_decrypt(&mut self, header: &Header, ciphertext: &[u8]) -> Vec<u8> {
+    pub fn ratchet_decrypt(&mut self, header: &Header, ciphertext: &[u8], nonce: &[u8; 12]) -> Vec<u8> {
-        let plaintext = self.try_skipped_message_keys(header, ciphertext);
+        let plaintext = self.try_skipped_message_keys(header, ciphertext, nonce);
         match plaintext {
             Some(d) => d,
             None => {
@@ -114,7 +119,7 @@ impl Ratchet {
                 let (ckr, mk) = kdf_ck(&self.ckr.unwrap());
                 self.ckr = Some(ckr);
                 self.nr += 1;
-                decrypt(&mk, ciphertext, &header.concat())
+                decrypt(&mk, ciphertext, &header.concat(), nonce)
             }
         }
     }
@@ -136,3 +141,155 @@ impl Ratchet {
     }
 }
 
+pub struct RatchetEncHeader {
+    dhs: DhKeyPair,
+    dhr: Option<PublicKey>,
+    rk: [u8; 32],
+    cks: Option<[u8; 32]>,
+    ckr: Option<[u8; 32]>,
+    ns: usize,
+    nr: usize,
+    pn: usize,
+    hks: Option<[u8; 32]>,
+    hkr: Option<[u8; 32]>,
+    nhks: Option<[u8; 32]>,
+    nhkr: Option<[u8; 32]>,
+    mkskipped: HashMap<(Option<[u8; 32]>, usize), [u8; 32]>
+}
+
+impl RatchetEncHeader {
+    pub fn init_alice(sk: [u8; 32],
+                      bob_dh_public_key: PublicKey,
+                      shared_hka: [u8; 32],
+                      shared_nhkb: [u8; 32]) -> Self {
+        let dhs = DhKeyPair::new();
+        let (rk, cks, nhks) = kdf_rk_he(&sk, &dhs.key_agreement(&bob_dh_public_key));
+        RatchetEncHeader {
+            dhs,
+            dhr: Some(bob_dh_public_key),
+            rk,
+            cks: Some(cks),
+            ckr: None,
+            ns: 0,
+            nr: 0,
+            pn: 0,
+            mkskipped: HashMap::new(),
+            hks: Some(shared_hka),
+            hkr: None,
+            nhkr: Some(shared_nhkb),
+            nhks: Some(nhks),
+        }
+    }
+
+    pub fn init_bob(sk: [u8; 32], shared_hka: [u8; 32], shared_nhkb: [u8; 32]) -> (Self, PublicKey) {
+        let dhs = DhKeyPair::new();
+        let public_key = dhs.public_key;
+        let ratchet = Self {
+            dhs,
+            dhr: None,
+            rk: sk,
+            cks: None,
+            ckr: None,
+            ns: 0,
+            nr: 0,
+            pn: 0,
+            mkskipped: HashMap::new(),
+            hks: None,
+            nhks: Some(shared_nhkb),
+            hkr: None,
+            nhkr: Some(shared_hka),
+        };
+        (ratchet, public_key)
+    }
+
+    pub fn ratchet_encrypt(&mut self, plaintext: &[u8]) -> HeaderNonceCipherNonce {
+        let (cks, mk) = kdf_ck(&self.cks.unwrap());
+        self.cks = Some(cks);
+        let header = Header::new(&self.dhs, self.pn, self.ns);
+        let enc_header = header.encrypt(&self.hks.unwrap());
+        self.ns += 1;
+        let encrypted = encrypt(&mk, plaintext, &header.concat());
+        (enc_header, encrypted.0, encrypted.1)
+    }
+
+    fn try_skipped_message_keys(&mut self, enc_header: &(Vec<u8>, [u8; 12]),
+                                ciphertext: &[u8], nonce: &[u8; 12]) -> Option<Vec<u8>> {
+
+        let ret_data = self.mkskipped.clone().into_iter().find(|e| {
+            let header = Header::decrypt(&e.0.0, &enc_header.0, &enc_header.1);
+            match header {
+                None => false,
+                Some(h) => h.n == e.0.1
+            }
+        });
+        match ret_data {
+            None => { None },
+            Some(data) => {
+                let header = Header::decrypt(&data.0.0, &enc_header.0, &enc_header.1);
+                let mk = data.1;
+                self.mkskipped.remove(&(data.0.0, data.0.1));
+                Some(decrypt(&mk, ciphertext, &header.unwrap().concat(), nonce))
+            }
+        }
+    }
+
+    fn decrypt_header(&mut self, enc_header: &(Vec<u8>, [u8; 12])) -> Result<(Header, bool), &str> {
+        let header = Header::decrypt(&self.hkr, &enc_header.0, &enc_header.1);
+        if let Some(h) = header { return Ok((h, false)) };
+        let header = Header::decrypt(&self.nhkr, &enc_header.0, &enc_header.1);
+        match header {
+            Some(h) => Ok((h, true)),
+            None => Err("Header is unencryptable!")
+        }
+    }
+
+    fn skip_message_keys(&mut self, until: usize) -> Result<(), &str> {
+        if self.nr + MAX_SKIP < until {
+            return Err("Skipping went wrong")
+        }
+        if let Some(d) = self.ckr {
+            while self.nr < until {
+                let (ckr, mk) = kdf_ck(&d);
+                self.ckr = Some(ckr);
+                self.mkskipped.insert((self.hkr, self.nr), mk);
+                self.nr += 1
+            }
+        }
+        Ok(())
+    }
+
+    fn dhratchet(&mut self, header: &Header) {
+        self.pn = self.ns;
+        self.ns = 0;
+        self.nr = 0;
+        self.hks = self.nhks;
+        self.hkr = self.nhkr;
+        self.dhr = Some(header.public_key);
+        let (rk, ckr, nhkr) = kdf_rk_he(&self.rk,
+                                        &self.dhs.key_agreement(&self.dhr.unwrap()));
+        self.rk = rk;
+        self.ckr = Some(ckr);
+        self.nhkr = Some(nhkr);
+        self.dhs = DhKeyPair::new();
+        let (rk, cks, nhks) = kdf_rk_he(&self.rk,
+                                        &self.dhs.key_agreement(&self.dhr.unwrap()));
+        self.rk = rk;
+        self.cks = Some(cks);
+        self.nhks = Some(nhks);
+    }
+
+    pub fn ratchet_decrypt(&mut self, enc_header: &(Vec<u8>, [u8; 12]), ciphertext: &[u8], nonce: &[u8; 12]) -> Vec<u8> {
+        let plaintext = self.try_skipped_message_keys(enc_header, ciphertext, nonce);
+        if let Some(d) = plaintext { return d };
+        let (header, dh_ratchet) = self.decrypt_header(enc_header).unwrap();
+        if dh_ratchet {
+            self.skip_message_keys(header.pn).unwrap();
+            self.dhratchet(&header);
+        }
+        self.skip_message_keys(header.n).unwrap();
+        let (ckr, mk) = kdf_ck(&self.ckr.unwrap());
+        self.ckr = Some(ckr);
+        self.nr += 1;
+        decrypt(&mk, ciphertext, &header.concat(), nonce)
+    }
+}

tests/mod.rs

@@ -1,4 +1,4 @@
-use double_ratchet_2::ratchet::Ratchet;
+use double_ratchet_2::ratchet::{Ratchet, RatchetEncHeader};
 
 #[test]
 fn ratchet_init() {
@@ -13,8 +13,8 @@ fn ratchet_enc_single() {
     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) = alice_ratchet.ratchet_encrypt(&data);
+    let (header, encrypted, nonce) = alice_ratchet.ratchet_encrypt(&data);
-    let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted);
+    let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce);
     assert_eq!(data, decrypted)
 }
 
@@ -24,10 +24,10 @@ fn ratchet_enc_skip() {
     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) = alice_ratchet.ratchet_encrypt(&data);
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data);
-    let (header2, encrypted2) = alice_ratchet.ratchet_encrypt(&data);
+    let (header2, encrypted2, nonce2) = alice_ratchet.ratchet_encrypt(&data);
-    let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2);
+    let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2);
-    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1);
+    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1);
     let comp_res = decrypted1 == data && decrypted2 == data;
     assert!(comp_res)
 }
@@ -38,7 +38,7 @@ fn ratchet_panic_bob() {
     let sk = [1; 32];
     let (mut bob_ratchet, _) = Ratchet::init_bob(sk);
     let data = include_bytes!("../src/header.rs").to_vec();
-    let (_, _) = bob_ratchet.ratchet_encrypt(&data);
+    let (_, _, _) = bob_ratchet.ratchet_encrypt(&data);
 }
 
 #[test]
@@ -47,10 +47,92 @@ fn ratchet_encryt_decrypt_four() {
     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) = alice_ratchet.ratchet_encrypt(&data);
+    let (header1, encrypted1, nonce1) = alice_ratchet.ratchet_encrypt(&data);
-    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1);
+    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1);
-    let (header2, encrypted2) = bob_ratchet.ratchet_encrypt(&data);
+    let (header2, encrypted2, nonce2) = bob_ratchet.ratchet_encrypt(&data);
-    let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2);
+    let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2);
+    let comp_res = decrypted1 == data && decrypted2 == data;
+    assert!(comp_res)
+}
+
+#[test]
+fn ratchet_ench_init() {
+    let sk = [1; 32];
+    let shared_hka = [2; 32];
+    let shared_nhkb = [3; 32];
+    let (_bob_ratchet, public_key) = RatchetEncHeader::init_bob(sk,
+                                                                   shared_hka,
+                                                                   shared_nhkb);
+    let _alice_ratchet = RatchetEncHeader::init_alice(sk, public_key,
+                                                      shared_hka, shared_nhkb);
+}
+
+#[test]
+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);
+    let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted, &nonce);
+    assert_eq!(data, decrypted)
+}
+
+#[test]
+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);
+    let (header2, encrypted2, nonce2) = alice_ratchet.ratchet_encrypt(&data);
+    let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2);
+    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1);
+    let comp_res = decrypted1 == data && decrypted2 == data;
+    assert!(comp_res)
+}
+
+#[test]
+#[should_panic]
+fn ratchet_ench_panic_bob() {
+    let sk = [1; 32];
+    let shared_hka = [2; 32];
+    let shared_nhkb = [3; 32];
+    let (mut bob_ratchet, _) = RatchetEncHeader::init_bob(sk,
+                                                                   shared_hka,
+                                                                   shared_nhkb);
+    let data = include_bytes!("../src/header.rs").to_vec();
+    let (_, _, _) = bob_ratchet.ratchet_encrypt(&data);
+}
+
+#[test]
+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);
+    let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1, &nonce1);
+    let (header2, encrypted2, nonce2) = bob_ratchet.ratchet_encrypt(&data);
+    let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2, &nonce2);
     let comp_res = decrypted1 == data && decrypted2 == data;
     assert!(comp_res)
 }