From e43ccfa7d44126c66b185fcc2999121f4036d5be Mon Sep 17 00:00:00 2001 From: Hannes Date: Thu, 13 May 2021 13:24:52 +0200 Subject: [PATCH] Initial Commit --- Cargo.toml | 22 ++++++++ src/aead.rs | 45 ++++++++++++++++ src/dh.rs | 84 ++++++++++++++++++++++++++++++ src/header.rs | 107 ++++++++++++++++++++++++++++++++++++++ src/kdf_chain.rs | 39 ++++++++++++++ src/kdf_root.rs | 42 +++++++++++++++ src/lib.rs | 12 +++++ src/ratchet.rs | 133 +++++++++++++++++++++++++++++++++++++++++++++++ tests/mod.rs | 56 ++++++++++++++++++++ 9 files changed, 540 insertions(+) create mode 100644 Cargo.toml create mode 100644 src/aead.rs create mode 100644 src/dh.rs create mode 100644 src/header.rs create mode 100644 src/kdf_chain.rs create mode 100644 src/kdf_root.rs create mode 100644 src/lib.rs create mode 100644 src/ratchet.rs create mode 100644 tests/mod.rs diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..4ac4d61 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "double-ratchet-2" +version = "0.1.0" +edition = "2018" + +# 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" +hkdf = "0.11.0" +hmac = "0.11.0" +aes-gcm-siv = {version = "0.10.0"} +ring-compat = {version = "0.2.1", features = ["digest"]} +serde = {version = "1.0.125", default-features = false, features = ["derive"]} +serde_bytes = "0.11.5" +bincode = "1.3.3" +hashbrown = "0.11.2" + +[profile.release] +lto = true +opt-level = 3 \ No newline at end of file diff --git a/src/aead.rs b/src/aead.rs new file mode 100644 index 0000000..8abdac9 --- /dev/null +++ b/src/aead.rs @@ -0,0 +1,45 @@ +use aes_gcm_siv::{Key, Aes256GcmSiv, Nonce}; +use aes_gcm_siv::aead::{NewAead, AeadInPlace}; +use alloc::vec::Vec; + +const CONSTANT_NONCE: &[u8] = b"Super Noncel"; + +pub fn encrypt(mk: &[u8; 32], plaintext: &[u8], associated_data: &[u8]) -> Vec { + let key = Key::from_slice(mk); + let cipher = Aes256GcmSiv::new(key); + + let nonce = Nonce::from_slice(&CONSTANT_NONCE); + let mut buffer = Vec::new(); + buffer.extend_from_slice(plaintext); + + cipher.encrypt_in_place(nonce, associated_data, &mut buffer) + .expect("Encryption failed"); + buffer +} + +pub fn decrypt(mk: &[u8; 32], ciphertext: &[u8], associated_data: &[u8]) -> Vec { + let key = Key::from_slice(mk); + let cipher = Aes256GcmSiv::new(key); + + let nonce = Nonce::from_slice(&CONSTANT_NONCE); + let mut buffer = Vec::new(); + buffer.extend_from_slice(ciphertext); + cipher.decrypt_in_place(nonce, associated_data, &mut buffer).expect("Decryption failure"); + buffer +} + +#[cfg(test)] +mod tests { + use crate::kdf_chain::gen_mk; + use crate::aead::{encrypt, decrypt}; + + #[test] + fn enc_a_dec() { + 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 plaintext = decrypt(&mk, &ciphertext, &associated_data); + assert_eq!(test_data, plaintext) + } +} \ No newline at end of file diff --git a/src/dh.rs b/src/dh.rs new file mode 100644 index 0000000..09911b8 --- /dev/null +++ b/src/dh.rs @@ -0,0 +1,84 @@ +use x25519_dalek::{PublicKey, SharedSecret, StaticSecret}; +use rand_core::OsRng; +use core::fmt::{Debug, Formatter}; +use core::fmt; + +pub struct DhKeyPair { + pub private_key: StaticSecret, + pub public_key: PublicKey, +} + +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() { + return false + } + true + } +} + +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()) + .finish() + } +} + +impl Default for DhKeyPair { + fn default() -> Self { + Self::new() + } +} + +impl DhKeyPair { + pub fn new() -> Self { + let secret = StaticSecret::new(OsRng); + let public = PublicKey::from(&secret); + DhKeyPair { + private_key: secret, + public_key: public, + } + } + + pub fn key_agreement(&self, public_key: &PublicKey) -> SharedSecret { + self.private_key.diffie_hellman(public_key) + } +} + +#[cfg(test)] +pub fn gen_shared_secret() -> SharedSecret { + let alice_pair = DhKeyPair::new(); + let bob_pair = DhKeyPair::new(); + alice_pair.key_agreement(&bob_pair.public_key) +} + +#[cfg(test)] +pub fn gen_key_pair() -> DhKeyPair { + DhKeyPair::new() +} + +#[cfg(test)] +mod tests { + use crate::dh::DhKeyPair; + + #[test] + fn key_generation() { + let pair_1 = DhKeyPair::new(); + let pair_2 = DhKeyPair::new(); + assert_ne!(pair_1, pair_2) + } + + #[test] + fn key_agreement() { + let alice_pair = DhKeyPair::new(); + let bob_pair = DhKeyPair::new(); + let alice_shared_secret = alice_pair.key_agreement(&bob_pair.public_key); + let bob_shared_secret = bob_pair.key_agreement(&alice_pair.public_key); + assert_eq!(alice_shared_secret.as_bytes(), bob_shared_secret.as_bytes()) + } +} \ No newline at end of file diff --git a/src/header.rs b/src/header.rs new file mode 100644 index 0000000..28fdf43 --- /dev/null +++ b/src/header.rs @@ -0,0 +1,107 @@ +use x25519_dalek::PublicKey; +use crate::dh::DhKeyPair; +use alloc::vec::Vec; +use serde::{Serialize, Deserialize}; + +#[cfg(test)] +use crate::dh::gen_key_pair; + +#[derive(Debug)] +pub struct Header { + pub public_key: PublicKey, + pub pn: usize, // Previous Chain Length + pub n: usize, // Message Number +} + +#[derive(Serialize, Deserialize, Debug)] +struct ExHeader { + public_key: [u8; 32], + pn: usize, + n: usize +} + +impl Header { + pub fn new(dh_pair: &DhKeyPair, pn: usize, n: usize) -> Self { + Header { + public_key: dh_pair.public_key, + pn, + n, + } + } + + pub fn concat(&self) -> Vec { + let ex_header = ExHeader{ + public_key: self.public_key.to_bytes(), + pn: self.pn, + n: self.n + }; + bincode::serialize(&ex_header).expect("Failed to serialize Header") + } +} + +impl From> for Header { + fn from(d: Vec) -> Self { + let ex_header: ExHeader = bincode::deserialize(&d).unwrap(); + Header { + public_key: PublicKey::from(ex_header.public_key), + pn: ex_header.pn, + n: ex_header.n, + } + } +} + +impl From<&[u8]> for Header { + fn from(d: &[u8]) -> Self { + let ex_header: ExHeader = bincode::deserialize(d).unwrap(); + Header { + public_key: PublicKey::from(ex_header.public_key), + pn: ex_header.pn, + n: ex_header.n, + } + } +} + +impl PartialEq for Header { + fn eq(&self, other: &Self) -> bool { + if self.public_key == other.public_key + && self.pn == other.pn + && self.n == other.n { + return true + } + false + } +} + +#[cfg(test)] +pub fn gen_header() -> Header { + let dh_pair = gen_key_pair(); + let pn = 10; + let n = 50; + Header::new(&dh_pair, pn, n) +} + +#[cfg(test)] +mod tests { + use crate::header::{gen_header, Header}; + use crate::kdf_chain::gen_mk; + use crate::aead::{encrypt, decrypt}; + + #[test] + fn ser_des() { + let header = gen_header(); + let serialized = header.concat(); + let created = Header::from(serialized); + assert_eq!(header, created) + } + + #[test] + fn enc_header() { + let header = gen_header(); + let mk = gen_mk(); + let header_data = header.concat(); + let data = include_bytes!("aead.rs"); + let encrypted = encrypt(&mk, data, &header_data); + let decrypted = decrypt(&mk, &encrypted, &header_data); + assert_eq!(decrypted, data.to_vec()) + } +} \ No newline at end of file diff --git a/src/kdf_chain.rs b/src/kdf_chain.rs new file mode 100644 index 0000000..181ef94 --- /dev/null +++ b/src/kdf_chain.rs @@ -0,0 +1,39 @@ +use hmac::{Hmac, Mac, NewMac}; +use ring_compat::digest::Sha512; +use core::convert::TryInto; + +#[cfg(test)] +use crate::kdf_root::gen_ck; + +type HmacSha512 = Hmac; + +pub fn kdf_ck(ck: &[u8; 32]) -> ([u8; 32], [u8; 32]) { + let mac = HmacSha512::new_from_slice(ck) + .expect("Invalid Key Length"); + let result = mac.finalize().into_bytes(); + let (a, b) = result.split_at(32); + (a.try_into() + .expect("Incorrect Length"), + b.try_into() + .expect("Incorrect Length")) +} + +#[cfg(test)] +pub fn gen_mk() -> [u8; 32] { + let ck = gen_ck(); + let (_, mk) = kdf_ck(&ck); + mk +} + +#[cfg(test)] +mod tests { + use crate::kdf_root::gen_ck; + use crate::kdf_chain::kdf_ck; + #[test] + fn kdf_chain_ratchet() { + let ck = gen_ck(); + let (ck, mk1) = kdf_ck(&ck); + let (_, mk2) = kdf_ck(&ck); + assert_ne!(mk1, mk2) + } +} \ No newline at end of file diff --git a/src/kdf_root.rs b/src/kdf_root.rs new file mode 100644 index 0000000..b08b8f4 --- /dev/null +++ b/src/kdf_root.rs @@ -0,0 +1,42 @@ +use x25519_dalek::SharedSecret; +use hkdf::Hkdf; +use ring_compat::digest::Sha512; +use core::convert::TryInto; + +#[cfg(test)] +use crate::dh::gen_shared_secret; + +pub fn kdf_rk(rk: &[u8; 32], dh_out: &SharedSecret) -> ([u8; 32], [u8; 32]) { + let h = Hkdf::::new(Some(rk), dh_out.as_bytes()); + let mut okm = [0u8; 64]; + let info = b"Root Key Info"; + h.expand(info, &mut okm).unwrap(); + let (a, b) = okm.split_at(32); + (a.try_into() + .expect("Incorrect length"), + b.try_into() + .expect("Incorrect length")) +} + +#[cfg(test)] +pub fn gen_ck() -> [u8; 32] { + let shared_secret = gen_shared_secret(); + let rk = [0; 32]; + let (_, ck) = kdf_rk(&rk, &shared_secret); + ck +} + +#[cfg(test)] +mod tests { + use crate::dh::gen_shared_secret; + use crate::kdf_root::kdf_rk; + + #[test] + fn kdf_root_ratchet() { + let rk = [0; 32]; + let shared_secret = gen_shared_secret(); + let (rk1, _) = kdf_rk(&rk, &shared_secret); + let (rk2, _) = kdf_rk(&rk1, &shared_secret); + assert_ne!(rk1, rk2) + } +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..ddacdac --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,12 @@ +#![no_std] +#![allow(stable_features)] +#![feature(alloc)] + +extern crate alloc; + +mod aead; +mod dh; +mod kdf_root; +mod kdf_chain; +mod header; +pub mod ratchet; diff --git a/src/ratchet.rs b/src/ratchet.rs new file mode 100644 index 0000000..50e0ae4 --- /dev/null +++ b/src/ratchet.rs @@ -0,0 +1,133 @@ +use crate::dh::DhKeyPair; +use x25519_dalek::PublicKey; +use hashbrown::HashMap; +use crate::kdf_root::kdf_rk; +use crate::header::Header; +use alloc::vec::Vec; +use crate::kdf_chain::kdf_ck; +use crate::aead::{encrypt, decrypt}; + +const MAX_SKIP: usize = 100; + +pub struct Ratchet { + dhs: DhKeyPair, + dhr: Option, + rk: [u8; 32], + ckr: Option<[u8; 32]>, + cks: Option<[u8; 32]>, + ns: usize, + nr: usize, + pn: usize, + mkskipped: HashMap<(PublicKey, usize), [u8; 32]>, +} + +impl Ratchet { + pub fn init_alice(sk: [u8; 32], bob_dh_public_key: PublicKey) -> Self { + let dhs = DhKeyPair::new(); + let (rk, cks) = kdf_rk(&sk, + &dhs.key_agreement(&bob_dh_public_key)); + Ratchet { + dhs, + dhr: Some(bob_dh_public_key), + rk, + cks: Some(cks), + ckr: None, + ns: 0, + nr: 0, + pn: 0, + mkskipped: HashMap::new(), + } + } + + pub fn init_bob(sk: [u8; 32]) -> (Self, PublicKey) { + let dhs = DhKeyPair::new(); + let public_key = dhs.public_key; + let ratchet = Ratchet { + dhs, + dhr: None, + rk: sk, + cks: None, + ckr: None, + ns: 0, + nr: 0, + pn: 0, + mkskipped: HashMap::new(), + }; + (ratchet, public_key) + } + + pub fn ratchet_encrypt(&mut self, plaintext: &[u8]) -> (Header, Vec) { + 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()); + (header, encrypted_data) + } + + fn try_skipped_message_keys(&mut self, header: &Header, ciphertext: &[u8]) -> Option> { + 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())) + } else { + None + } + } + + fn skip_message_keys(&mut self, until: usize) -> Result<(), &str> { + if self.nr + MAX_SKIP < until { + return Err("Skipped to many keys"); + } + match self.ckr { + Some(d) => { + while self.nr < until { + let (ckr, mk) = kdf_ck(&d); + self.ckr = Some(ckr); + self.mkskipped.insert((self.dhr.unwrap(), self.nr), mk); + self.nr += 1 + } + Ok(()) + }, + None => { Err("No Ckr set") } + } + } + + pub fn ratchet_decrypt(&mut self, header: &Header, ciphertext: &[u8]) -> Vec { + let plaintext = self.try_skipped_message_keys(header, ciphertext); + match plaintext { + Some(d) => d, + None => { + if Some(header.public_key) != self.dhr { + if self.ckr != None { + 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()) + } + } + } + + fn dhratchet(&mut self, header: &Header) { + self.pn = self.ns; + self.ns = 0; + self.nr = 0; + self.dhr = Some(header.public_key); + let (rk, ckr) = kdf_rk(&self.rk, + &self.dhs.key_agreement(&self.dhr.unwrap())); + self.rk = rk; + self.ckr = Some(ckr); + self.dhs = DhKeyPair::new(); + let (rk, cks) = kdf_rk(&self.rk, + &self.dhs.key_agreement(&self.dhr.unwrap())); + self.rk = rk; + self.cks = Some(cks); + } +} + diff --git a/tests/mod.rs b/tests/mod.rs new file mode 100644 index 0000000..26f54a4 --- /dev/null +++ b/tests/mod.rs @@ -0,0 +1,56 @@ +use double_ratchet_2::ratchet::Ratchet; + +#[test] +fn ratchet_init() { + let sk = [1; 32]; + let (_bob_ratchet, public_key) = Ratchet::init_bob(sk); + let _alice_ratchet = Ratchet::init_alice(sk, public_key); +} + +#[test] +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) = alice_ratchet.ratchet_encrypt(&data); + let decrypted = bob_ratchet.ratchet_decrypt(&header, &encrypted); + assert_eq!(data, decrypted) +} + +#[test] +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) = alice_ratchet.ratchet_encrypt(&data); + let (header2, encrypted2) = alice_ratchet.ratchet_encrypt(&data); + let decrypted2 = bob_ratchet.ratchet_decrypt(&header2, &encrypted2); + let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1); + let comp_res = decrypted1 == data && decrypted2 == data; + assert!(comp_res) +} + +#[test] +#[should_panic] +fn ratchet_panic_bob() { + let sk = [1; 32]; + let (mut bob_ratchet, public_key) = Ratchet::init_bob(sk); + let data = include_bytes!("../src/header.rs").to_vec(); + let (header, encrypted) = bob_ratchet.ratchet_encrypt(&data); +} + +#[test] +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) = alice_ratchet.ratchet_encrypt(&data); + let decrypted1 = bob_ratchet.ratchet_decrypt(&header1, &encrypted1); + let (header2, encrypted2) = bob_ratchet.ratchet_encrypt(&data); + let decrypted2 = alice_ratchet.ratchet_decrypt(&header2, &encrypted2); + let comp_res = decrypted1 == data && decrypted2 == data; + assert!(comp_res) +} \ No newline at end of file