use crate::grpc::challenge::NodeUpdate; use dashmap::{DashMap, DashSet}; use ed25519_dalek::{Signer, SigningKey, VerifyingKey, PUBLIC_KEY_LENGTH}; use rand::rngs::OsRng; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; use tabled::{Table, Tabled}; #[derive(Clone, PartialEq, Debug)] pub struct NodeInfo { pub pubkey: VerifyingKey, pub updated_at: SystemTime, pub public: bool, } /// Needs to be surrounded in an Arc. pub struct Store { nodes: DashMap, conns: DashSet, keys: DashMap, } pub enum SigningError { CorruptedKey, KeyNotFound, } impl From for SigningError { fn from(_: hex::FromHexError) -> Self { Self::CorruptedKey } } impl From for SigningError { fn from(_: ed25519_dalek::ed25519::Error) -> Self { Self::CorruptedKey } } impl From for SigningError { fn from(_: std::array::TryFromSliceError) -> Self { Self::CorruptedKey } } type IP = String; impl std::fmt::Display for SigningError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let error_message = match self { SigningError::CorruptedKey => "The public key is corrupted", SigningError::KeyNotFound => "Did not find the public key", }; write!(f, "{}", error_message) } } impl Store { // app should exit if any error happens here so unwrap() is good pub fn init() -> Self { Self { nodes: DashMap::new(), keys: DashMap::new(), conns: DashSet::new(), } } pub async fn add_conn(&self, ip: &str) { self.conns.insert(ip.to_string()); } pub async fn delete_conn(&self, ip: &str) { self.conns.remove(ip); } pub async fn tabled_node_list(&self) -> String { #[derive(Tabled)] struct OutputRow { ip: String, pubkey: String, age: u64, public: bool, } let mut output = vec![]; for (ip, node_info) in self .nodes .iter() .map(|n| (n.key().clone(), n.value().clone())) { let pubkey = hex::encode(node_info.pubkey.as_bytes()); let age = SystemTime::now() .duration_since(node_info.updated_at) .unwrap_or(Duration::ZERO) .as_secs(); let public = node_info.public; output.push(OutputRow { ip, pubkey, age, public, }); } Table::new(output).to_string() } pub async fn sign_message_with_key( &self, message: &str, pubkey: &str, ) -> Result { let mut pubkey_bytes = [0u8; PUBLIC_KEY_LENGTH]; hex::decode_to_slice(pubkey, &mut pubkey_bytes)?; let pubkey = VerifyingKey::from_bytes(&pubkey_bytes)?; let signing_key = match self.get_privkey(&pubkey).await { Some(k) => k, None => return Err(SigningError::KeyNotFound), }; // let signature = format!("{:?}", signing_key.sign(message.as_bytes())); let signature = hex::encode(signing_key.sign(message.as_bytes()).to_bytes()); Ok(signature) } pub async fn add_key(&self, pubkey: VerifyingKey, privkey: SigningKey) { self.keys.insert(pubkey, privkey); } pub async fn remove_key(&self, pubkey: &VerifyingKey) { self.keys.remove(pubkey); } async fn get_privkey(&self, pubkey: &VerifyingKey) -> Option { self.keys.get(pubkey).map(|k| k.value().clone()) } /// This returns true if NodeInfo got modified. /// /// On a side note, there are two types of people in this world: /// 1. Those that can extrapolate... WAT? pub async fn process_node_update(&self, node: NodeUpdate) -> bool { let key_bytes = match hex::decode(node.keypair.clone()) { Ok(k) => k, Err(_) => return false, }; let privkey = SigningKey::from_bytes(match &key_bytes.as_slice().try_into() { Ok(p) => p, Err(_) => return false, }); let pubkey = privkey.verifying_key(); // TODO: check this suggestion // let updated_at_std = node // .updated_at // .map(SystemTime::try_from) // .unwrap_or(Ok(SystemTime::now())) // .unwrap_or(SystemTime::now()); let updated_at_std: SystemTime = match node.updated_at { Some(ts) => { let duration = Duration::new(ts.seconds as u64, ts.nanos as u32); UNIX_EPOCH .checked_add(duration) .unwrap_or(SystemTime::now()) } None => SystemTime::now(), }; self.add_key(pubkey, privkey).await; let node_info = NodeInfo { pubkey, updated_at: updated_at_std, public: node.public, }; if let Some(mut old_node_info) = self.update_node(node.ip, node_info.clone()).await { if !node_info.public { old_node_info.public = false; } match old_node_info.ne(&node_info) { true => { self.remove_key(&old_node_info.pubkey).await; true } false => false, } } else { true } } /// returns old pubkey if node got updated async fn update_node(&self, ip: String, info: NodeInfo) -> Option { if let Some(old_node) = self.nodes.get(&ip) { if old_node.updated_at > info.updated_at { return None; } } self.nodes.insert(ip, info.clone()) } // pub async fn remove_node(&self, ip: &str) { // let mut nodes = self.nodes.lock().await; // nodes.remove(ip); // } pub async fn get_localhost(&self) -> NodeUpdate { // TODO trigger reset_localhost_keys on error instead of expects let node = self.nodes.get("localhost").expect("no localhost node"); let key = self.keys.get(&node.pubkey).expect("no localhost key"); NodeUpdate { ip: "localhost".to_string(), keypair: hex::encode(key.value().as_bytes()), updated_at: Some(prost_types::Timestamp::from(node.value().updated_at)), public: false, } } /// refreshes the keys of the node and returns a protobuf for the network pub async fn reset_localhost_keys(&self) -> NodeUpdate { let mut csprng = OsRng; let keypair_raw = SigningKey::generate(&mut csprng); let keypair = hex::encode(keypair_raw.as_bytes()); let pubkey = keypair_raw.verifying_key(); let ip = "localhost".to_string(); let updated_at = SystemTime::now(); let public = false; self.add_key(pubkey, keypair_raw.clone()).await; if let Some(old_data) = self .update_node( ip.clone(), NodeInfo { pubkey, updated_at, public, }, ) .await { self.remove_key(&old_data.pubkey).await; }; let updated_at = Some(prost_types::Timestamp::from(updated_at)); NodeUpdate { ip, keypair, updated_at, public, } } pub async fn get_full_node_list(&self) -> Vec { self.nodes .iter() .filter_map(|node| { self.keys .get(&node.value().pubkey) .map(|signing_key| NodeUpdate { ip: node.key().to_string(), keypair: hex::encode(signing_key.as_bytes()), updated_at: Some(prost_types::Timestamp::from(node.value().updated_at)), public: node.value().public, }) }) .collect() } // returns a random node that does not have an active connection pub async fn get_random_node(&self) -> Option { use rand::rngs::OsRng; use rand::RngCore; let len = self.nodes.len(); if len == 0 { return None; } let skip = OsRng.next_u64().try_into().unwrap_or(0) % len; self.nodes .iter() .map(|n| n.key().clone()) .cycle() .skip(skip) .find(|k| !self.conns.contains(k)) } }