Added tests and fixed some issues.

I also added documentation.

Cargo.toml

@@ -9,7 +9,6 @@ keywords = ["double-ratchet", "crypto", "cryptography", "signal"]
 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

@@ -3,8 +3,6 @@ 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]) {
     let key = Key::from_slice(mk);
     let cipher = Aes256GcmSiv::new(key);

src/header.rs

@@ -4,7 +4,7 @@ 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, Error};
+use aes_gcm_siv::aead::{NewAead, AeadInPlace};
 
 #[cfg(test)]
 use crate::dh::gen_key_pair;
@@ -93,9 +93,9 @@ impl From<&[u8]> for Header {
     }
 }
 
-impl Into<Vec<u8>> for Header {
-    fn into(self) -> Vec<u8> {
-        self.concat()
+impl From<Header> for Vec<u8> {
+    fn from(s: Header) -> Self {
+        s.concat()
     }
 }
 

src/lib.rs

@@ -89,6 +89,23 @@
 //! 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

src/ratchet.rs

@@ -12,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,
@@ -200,7 +202,7 @@ impl RatchetEncHeader {
         (ratchet, public_key)
     }
 
-    pub fn ratchet_encrypt(&mut self, plaintext: &[u8]) -> ((Vec<u8>, [u8; 12]), Vec<u8>, [u8; 12]) {
+    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);
@@ -212,28 +214,28 @@ impl RatchetEncHeader {
 
     fn try_skipped_message_keys(&mut self, enc_header: &(Vec<u8>, [u8; 12]),
                                 ciphertext: &[u8], nonce: &[u8; 12]) -> Option<Vec<u8>> {
-        for ((hk, n), mk) in self.mkskipped.iter() {
-            let header = Header::decrypt(hk, &enc_header.0, &enc_header.1);
+
+        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 => { continue },
-                Some(h) => {
-                    if h.n == *n {
-                        let mk = mk.clone();
-                        self.mkskipped.remove(&(*hk, *n));
-                        return Some(decrypt(&mk, ciphertext, &h.concat(), nonce))
-                    }
-                }
+                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))
             }
         }
-        None
     }
 
     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);
-        match header {
-            Some(h) => { return Ok((h, false)) },
-            None => {},
-        }
+        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)),
@@ -245,16 +247,13 @@ impl RatchetEncHeader {
         if self.nr + MAX_SKIP < until {
             return Err("Skipping went wrong")
         }
-        match self.ckr {
-            Some(d) => {
-                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
-                }
+        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
             }
-            None => {}
         }
         Ok(())
     }
@@ -281,10 +280,7 @@ impl RatchetEncHeader {
 
     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);
-        match plaintext {
-            Some(d) => { return d },
-            None => {}
-        };
+        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();

tests/mod.rs

@@ -112,7 +112,7 @@ fn ratchet_ench_panic_bob() {
     let sk = [1; 32];
     let shared_hka = [2; 32];
     let shared_nhkb = [3; 32];
-    let (mut bob_ratchet, public_key) = RatchetEncHeader::init_bob(sk,
+    let (mut bob_ratchet, _) = RatchetEncHeader::init_bob(sk,
                                                                    shared_hka,
                                                                    shared_nhkb);
     let data = include_bytes!("../src/header.rs").to_vec();