diff --git a/key_protocol/src/key_management/key_tree.rs b/key_protocol/src/key_management/key_tree.rs new file mode 100644 index 0000000..8f3772f --- /dev/null +++ b/key_protocol/src/key_management/key_tree.rs @@ -0,0 +1,205 @@ +use std::{ + collections::{BTreeMap, HashMap}, + str::FromStr, + u32, +}; + +use crate::key_management::secret_holders::SeedHolder; + +#[derive(Debug)] +pub struct ChildKeysPublic { + pub csk: nssa::PrivateKey, + pub cpk: nssa::PublicKey, + pub ccc: [u8; 32], + ///Can be None if root + pub cci: Option, +} + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)] +pub struct ChainIndex(Vec); + +impl FromStr for ChainIndex { + type Err = hex::FromHexError; + + fn from_str(s: &str) -> Result { + if s == "" { + return Ok(Self(vec![])); + } + + let hex_decoded = hex::decode(s)?; + + if !hex_decoded.len().is_multiple_of(4) { + Err(hex::FromHexError::InvalidStringLength) + } else { + let mut res_vec = vec![]; + + for i in 0..(hex_decoded.len() / 4) { + res_vec.push(u32::from_le_bytes([ + hex_decoded[4 * i], + hex_decoded[4 * i + 1], + hex_decoded[4 * i + 2], + hex_decoded[4 * i + 3], + ])); + } + + Ok(Self(res_vec)) + } + } +} + +impl ToString for ChainIndex { + fn to_string(&self) -> String { + if self.0.is_empty() { + return "".to_string(); + } + + let mut res_vec = vec![]; + + for index in &self.0 { + res_vec.extend_from_slice(&index.to_le_bytes()); + } + + hex::encode(res_vec) + } +} + +impl ChainIndex { + pub fn next_in_line(&self) -> ChainIndex { + let mut chain = self.0.clone(); + //ToDo: Add overflow check + chain.last_mut().map(|last_p| *last_p += 1); + + ChainIndex(chain) + } + + pub fn n_th_son(&self, son_id: u32) -> ChainIndex { + let mut chain = self.0.clone(); + chain.push(son_id); + + ChainIndex(chain) + } +} + +#[derive(Debug)] +pub struct KeyTreePublic { + pub key_map: BTreeMap, + pub addr_map: HashMap, +} + +impl KeyTreePublic { + pub fn new(seed: &SeedHolder) -> Self { + let seed_fit: [u8; 64] = seed.seed.clone().try_into().unwrap(); + let hash_value = hmac_sha512::HMAC::mac(&seed_fit, "NSSA_master_pub"); + + let csk = nssa::PrivateKey::try_new(*hash_value.first_chunk::<32>().unwrap()).unwrap(); + let ccc = *hash_value.last_chunk::<32>().unwrap(); + let cpk = nssa::PublicKey::new_from_private_key(&csk); + let address = nssa::Address::from(&cpk); + + let root_keys = ChildKeysPublic { + csk, + cpk, + ccc, + cci: None, + }; + + let mut key_map = BTreeMap::new(); + let mut addr_map = HashMap::new(); + + key_map.insert(ChainIndex::from_str("").unwrap(), root_keys); + addr_map.insert(address, ChainIndex::from_str("").unwrap()); + + Self { key_map, addr_map } + } + + pub fn find_last_son_of_id(&self, father_id: &ChainIndex) -> Option { + if !self.key_map.contains_key(father_id) { + return None; + } + + let leftmost_son = father_id.n_th_son(u32::MIN); + + if !self.key_map.contains_key(&leftmost_son) { + Some(0) + } else { + let mut right = u32::MAX - 1; + let mut left_border = u32::MIN; + let mut right_border = u32::MAX; + + loop { + let rightmost_son = father_id.n_th_son(right); + + let rightmost_ref = self.key_map.get(&rightmost_son); + let rightmost_ref_next = self.key_map.get(&rightmost_son.next_in_line()); + + match (&rightmost_ref, &rightmost_ref_next) { + (Some(_), Some(_)) => { + left_border = right; + right = (right + right_border) / 2; + } + (Some(_), None) => { + break Some(right); + } + (None, None) => { + right_border = right; + right = (left_border + right) / 2; + } + (None, Some(_)) => { + unreachable!(); + } + } + } + } + } + + pub fn generate_new_pub_keys(&mut self, father_cci: ChainIndex) -> Option { + if !self.key_map.contains_key(&father_cci) { + return None; + } + + let father_keys = self.key_map.get(&father_cci).unwrap(); + let next_son_id = self.find_last_son_of_id(&father_cci).unwrap(); + let next_son_cci = father_cci.n_th_son(next_son_id); + + let mut hash_input = vec![]; + hash_input.extend_from_slice(father_keys.csk.value()); + hash_input.extend_from_slice(&next_son_id.to_le_bytes()); + + let hash_value = hmac_sha512::HMAC::mac(&hash_input, father_keys.ccc); + + let csk = nssa::PrivateKey::try_new(*hash_value.first_chunk::<32>().unwrap()).unwrap(); + let ccc = *hash_value.last_chunk::<32>().unwrap(); + let cpk = nssa::PublicKey::new_from_private_key(&csk); + let address = nssa::Address::from(&cpk); + + let child_keys = ChildKeysPublic { + csk, + cpk, + ccc, + cci: None, + }; + + self.key_map.insert(next_son_cci.clone(), child_keys); + self.addr_map.insert(address, next_son_cci); + + Some(address) + } + + pub fn get_pub_keys(&self, addr: nssa::Address) -> Option<&ChildKeysPublic> { + self.addr_map + .get(&addr) + .map(|chain_id| self.key_map.get(chain_id)) + .flatten() + } + + pub fn topology_hexdump(&self) -> String { + let mut hex_dump = String::new(); + + //Very inefficient + for chain_id in self.key_map.keys() { + hex_dump = format!("{hex_dump}{}", chain_id.to_string()); + } + + hex_dump + } +} diff --git a/key_protocol/src/key_management/mod.rs b/key_protocol/src/key_management/mod.rs index b642e13..4c3d5e0 100644 --- a/key_protocol/src/key_management/mod.rs +++ b/key_protocol/src/key_management/mod.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; pub type PublicAccountSigningKey = [u8; 32]; pub mod ephemeral_key_holder; +pub mod key_tree; pub mod secret_holders; #[derive(Serialize, Deserialize, Clone, Debug)]