hacker-challenge/src/datastore.rs

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<IP, NodeInfo>,
    conns: DashSet<IP>,
    keys: DashMap<VerifyingKey, SigningKey>,
}
pub enum SigningError {
    CorruptedKey,
    KeyNotFound,
}

impl From<hex::FromHexError> for SigningError {
    fn from(_: hex::FromHexError) -> Self {
        Self::CorruptedKey
    }
}

impl From<ed25519_dalek::ed25519::Error> for SigningError {
    fn from(_: ed25519_dalek::ed25519::Error) -> Self {
        Self::CorruptedKey
    }
}

impl From<std::array::TryFromSliceError> 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<String, SigningError> {
        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<SigningKey> {
        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<NodeInfo> {
        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<NodeUpdate> {
        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<String> {
        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))
    }
}