diff --git a/src/dh_node.rs b/src/dh_node.rs index 962c604..ab441e4 100644 --- a/src/dh_node.rs +++ b/src/dh_node.rs @@ -95,6 +95,12 @@ impl DhNode { #[cfg(test)] mod tests { + const PRIV_A: &str = "bde93e8e88997c7b130a827b86b5e3a33522c478a3a725b104be938d339de05d"; + const PRIV_B: &str = "80ce9bab99bd10ac0a6762e3c20a70d386eab16918741fa77d567dcfbe7ef2d5"; + + const PUB_A: &str = "1607af185dcf7c552cb8d37af5cdef29df26dde738a9da2d9f878e3c0d85e442"; + const PUB_B: &str = "636d16b0ab3f2c9d25a0aaf191415ab0817b82e51b3bcadeb325eaeac4b41802"; + use super::*; #[test] @@ -125,4 +131,30 @@ mod tests { pair12b.secret_key.unwrap().to_bytes() ); } + + #[test] + fn can_create_a_valid_shared_secret() { + let mut node_a = DhNode::import(None, Some(PRIV_A), PUB_A); + let node_b = DhNode::import(None, Some(PRIV_B), PUB_B); + + let shared_secret = node_a + .secret_key + .as_ref() + .unwrap() + .diffie_hellman(&node_b.pub_key); + assert_eq!( + "57d7c061f48818e811bd07662b263b45d081b4919e7ff54dceb05ca62f3dd47e", + hex::encode(shared_secret.as_bytes()) + ); + + let res = node_a.derive_key(&node_b, None).unwrap(); + assert_eq!( + "5e490105efc4afea52a2904f376ee8d5c0b79a703c1094d41325d247eaad53e9", + hex::encode(res) + ); + + node_a.secret_key = None; + let res2 = node_a.derive_key(&node_b, None).unwrap(); + assert_eq!(res2, res); + } } diff --git a/src/dh_tree.rs b/src/dh_tree.rs index fd2ed1a..d00f234 100644 --- a/src/dh_tree.rs +++ b/src/dh_tree.rs @@ -50,14 +50,16 @@ impl RatchetKeyTree { simulation .update(new_location) - .expect("Failed to create update on simulation"); + .expect("Failed to create update"); - let new_keys: VecDeque<[u8; 32]> = simulation + let mut new_keys: VecDeque<[u8; 32]> = simulation .tree .iter(TreeIteratorMode::Parent(new_location)) .map(|n| n.value.pub_key.as_bytes().clone()) .collect(); + new_keys.truncate(new_keys.len() - 1); + ( RatchetTreeKeyUpdate { new_keys, @@ -125,6 +127,9 @@ impl RatchetKeyTree { }; next_index = self.tree.parent_index(index); + if next_index == Some(0) { + break; + } } Ok(()) @@ -141,8 +146,7 @@ impl RatchetKeyTree { .value .secret_key .clone() - .map(|s| s.to_bytes()) - .expect("No secret key present"); + .map(|s| s.to_bytes()); let mut self_idx = index; @@ -153,27 +157,14 @@ impl RatchetKeyTree { .get(self_idx) .ok_or(RatchetTreeError::NodeNotFound(self_idx))?; - let sibling_idx = self - .tree - .sibling_index(self_idx) - .ok_or(RatchetTreeError::SiblingNodeNotFound(self_idx))?; - let sibling = self .tree - .nodes - .get(sibling_idx) + .sibling_index(self_idx) + .and_then(|idx| self.tree.nodes.get(idx)) .ok_or(RatchetTreeError::SiblingNodeNotFound(self_idx))?; - let (left, right) = if me.value.secret_key.is_some() { - (me, sibling) - } else if sibling.value.secret_key.is_some() { - (sibling, me) - } else { - return Err(RatchetTreeError::MissingPrivateKey); - }; - - let replace_node = left.value.derive( - &right.value, + let replace_node = me.value.derive( + &sibling.value, if parent_idx == 0 { self.prev_tree_secret.as_ref() } else { @@ -198,25 +189,24 @@ impl RatchetKeyTree { .map(|s| s.to_bytes()) .expect("No secret key present"); - if !prev_secret.eq(&new_secret) { - self.prev_tree_secret = Some(prev_secret); + self.prev_tree_secret = prev_secret; - println!( - "Updated secret: {:x?}, Previous: {:x?}", - new_secret, prev_secret - ); - } + println!( + "Updated secret: {:x?}, Previous: {:x?}", + new_secret, prev_secret + ); Ok(()) } - // body_encryption_key = hkdf(ikm: dh(tree.priv_key, node2.pub), salt: ratchet_id, info: tree.pub_key) + // body_encryption_key = hkdf(ikm: dh(tree.priv_key, usernode.pub), salt: ratchet_id, info: tree.pub_key) fn get_body_encryption_key(&self, leaf_index: usize) -> Result<[u8; 32], RatchetTreeError> { let child = self .tree .leaves() .get(leaf_index) .ok_or(RatchetTreeError::NodeNotFound(self.tree.index(leaf_index)))?; + let root = self .tree .nodes @@ -273,13 +263,35 @@ impl RatchetKeyTree { Ok((key, salt)) } + + fn clone_public(&self) -> Self { + let mut res = self.clone(); + for node in res.tree.nodes.iter_mut() { + node.value.secret_key = None; + } + res + } +} + +impl std::fmt::Debug for RatchetKeyTree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let nodes: Vec = self.tree.nodes.iter().map(|n|format!( + "{:?}, {}", + n.value.user_id, + if n.value.secret_key.is_some() { + "P" + } else { + "X" + } + )).collect(); + + writeln!(f, "Tree: {:?}", nodes).unwrap(); + Ok(()) + } } #[cfg(test)] mod tests { - use super::RatchetKeyTree; - use crate::dh_node::DhNode; - const PRIV_A: &str = "bde93e8e88997c7b130a827b86b5e3a33522c478a3a725b104be938d339de05d"; const PRIV_B: &str = "80ce9bab99bd10ac0a6762e3c20a70d386eab16918741fa77d567dcfbe7ef2d5"; const PRIV_C: &str = "ce6dbb15ddca4c5c3bf66436ecf1eef7e34716e651786adcc74ce4cf5d1f9fd2"; @@ -292,17 +304,53 @@ mod tests { const PUB_D: &str = "fa6e6eea0dd830a407aa184d5a56736965b19abcee7ccfec55b4099f8478860c"; const PUB_E: &str = "1fe0fafc670b454663ebef01d2e4821960a3cf12591e726d4e44e6fbdbfc076c"; + use super::*; + #[test] fn can_create_a_tree() { let mut rt = RatchetKeyTree::new(); + assert_eq!(None, rt.prev_tree_secret); rt.add(DhNode::new(Some(123))) .expect("First node to be able to add"); - assert_eq!(None, rt.prev_tree_secret); assert_eq!(1, rt.tree.leaf_count()); assert_eq!(1, rt.tree.nodes.len()); assert_eq!(Some(123), rt.tree.nodes[0].value.user_id); } + #[test] + fn can_create_shared_secret_with_two_nodes() { + let mut rt = RatchetKeyTree::new(); + let node1 = DhNode::import(Some(1), Some(PRIV_A), PUB_A); + let secret1 = node1.secret_key.clone(); + let node2 = DhNode::import(Some(2), None, PUB_B); + rt.add(node1).unwrap(); + rt.add(node2).unwrap(); + println!("{:?}", rt); + + assert_eq!( + &secret1.map(|s| s.to_bytes()).unwrap(), + &rt.prev_tree_secret.unwrap() + ); + + let mut rt2 = rt.clone_public(); + rt2.tree.nodes[1].value = DhNode::import(Some(2), Some(PRIV_B), PUB_B); + rt2.update(1).unwrap(); + println!("{:?}", rt2); + + assert_eq!( + rt.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()), + rt2.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()) + ); + } + #[test] fn can_add_a_user() { let mut rt = RatchetKeyTree::new(); @@ -313,14 +361,66 @@ mod tests { } #[test] - fn can_create_a_valid_shared_secret() { - let node_a = DhNode::import(None, Some(PRIV_A), PUB_A); - let node_b = DhNode::import(None, Some(PRIV_B), PUB_B); + fn can_create_and_apply_a_tree_update() { + let mut rt = RatchetKeyTree::new(); + rt.add(DhNode::new(Some(123))).unwrap(); + rt.add(DhNode::new(Some(234))).unwrap(); + rt.add(DhNode::new(Some(567))).unwrap(); + rt.add(DhNode::new(Some(890))).unwrap(); + println!("{:?}", rt); - let shared_secret = node_a.secret_key.unwrap().diffie_hellman(&node_b.pub_key); - assert_eq!( - "57d7c061f48818e811bd07662b263b45d081b4919e7ff54dceb05ca62f3dd47e", - hex::encode(shared_secret.as_bytes()) + let mut rt2 = rt.clone_public(); + + let index = rt.tree.index(3); + let secret_4 = rt.tree.nodes[index - 3].value.secret_key.clone(); + + let current = rt.tree.nodes[index].value.clone(); + println!("\nCurrent: {:?}", current.secret_key.map(|s| s.to_bytes())); + println!("Current id: {:?}", current.user_id); + + let update = rt.create_key_update(index); + rt.tree.nodes[index].value.pub_key = PublicKey::from(&update.1); + rt.tree.nodes[index].value.secret_key = Some(update.1); + rt.update(index).unwrap(); + + println!( + "New secret: {:x?}", + rt.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()) ); + let new_secret = rt.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()) + .clone() + .unwrap(); + + rt2.tree.nodes[index - 3].value.secret_key = secret_4; + println!("{:?}", rt2); + rt2.apply_key_update(update.0) + .expect("Failed to apply update"); + rt2.update(index - 3).expect("Failed to update"); + + let other_secret = rt2.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()) + .clone() + .unwrap(); + println!( + "Other secret: {:x?}", + rt2.tree.nodes[0] + .value + .secret_key + .as_ref() + .map(|s| s.to_bytes()) + ); + + assert_eq!(new_secret, other_secret); } }