hacker-challenge-sgx-general/rewrite/src/datastore.rs

#![allow(dead_code)]
use crate::grpc::challenge::NodeUpdate;
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use dashmap::DashSet;
use solana_sdk::signature::keypair::Keypair;
use std::time::SystemTime;
use std::time::{Duration, UNIX_EPOCH};

type IP = String;
pub const LOCALHOST: &str = "localhost";

#[derive(Clone, PartialEq, Debug)]
pub struct NodeInfo {
    pub started_at: SystemTime,
    pub keepalive: SystemTime,
    pub mint_requests: u64,
    pub mints: u64,
    pub ratls_conns: u64,
    pub ratls_attacks: u64,
    pub public: bool,
}

impl NodeInfo {
    pub fn newer_than(&self, other_node: &Self) -> bool {
        if self.keepalive > other_node.keepalive
            || self.mint_requests > other_node.mint_requests
            || self.ratls_attacks > other_node.ratls_attacks
            || self.ratls_conns > other_node.ratls_conns
            || self.mints > other_node.mints
            || (self.public && !other_node.public)
        {
            return true;
        }
        false
    }

    pub fn to_node_update(self, ip: &str) -> NodeUpdate {
        NodeUpdate {
            ip: ip.to_string(),
            started_at: Some(prost_types::Timestamp::from(self.started_at)),
            keepalive: Some(prost_types::Timestamp::from(self.keepalive)),
            mint_requests: self.mint_requests,
            mints: self.mints,
            ratls_conns: self.ratls_conns,
            ratls_attacks: self.ratls_attacks,
            public: false,
        }
    }
}

/// Keypair must already be known when creating a Store
/// This means the first node of the network creates the key
/// Second node will grab the key from the first node
pub struct Store {
    key: Keypair,
    nodes: DashMap<IP, NodeInfo>,
    conns: DashSet<IP>,
    // TODO: write persistence
    // persistence: FileManager,
}

impl Store {
    pub fn init(key: Keypair) -> Self {
        let store = Self {
            key,
            nodes: DashMap::new(),
            conns: DashSet::new(),
        };
        store.nodes.insert(
            LOCALHOST.to_string(),
            NodeInfo {
                started_at: SystemTime::now(),
                keepalive: SystemTime::now(),
                mint_requests: 0,
                mints: 0,
                ratls_conns: 0,
                ratls_attacks: 0,
                public: false,
            },
        );
        store
    }

    pub fn get_keypair_base58(&self) -> String {
        self.key.to_base58_string()
    }

    pub fn add_conn(&self, ip: &str) {
        self.conns.insert(ip.to_string());
    }

    pub fn delete_conn(&self, ip: &str) {
        self.conns.remove(ip);
    }

    pub fn increase_mint_requests(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.mint_requests += 1;
        }
    }

    pub fn increase_mints(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.mints += 1;
        }
    }

    pub fn get_localhost(&self) -> NodeInfo {
        let mut localhost = self.nodes.get_mut(LOCALHOST).expect("no localhost node");
        localhost.keepalive = SystemTime::now();
        localhost.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 ip = node.ip;
        let started_at: SystemTime = match node.started_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(),
        };
        let keepalive: SystemTime = match node.keepalive {
            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(),
        };
        let mut node_info = NodeInfo {
            started_at,
            keepalive,
            mint_requests: node.mint_requests,
            mints: node.mints,
            ratls_conns: node.ratls_conns,
            ratls_attacks: node.ratls_attacks,
            public: node.public,
        };
        if let Some(old_node) = self.nodes.get(&ip) {
            if !node_info.newer_than(&old_node) {
                return false;
            }
            node_info.public = node_info.public || old_node.public;
        }
        self.nodes.insert(ip, node_info);
        true
    }

    pub fn get_http_node_list(&self) -> Vec<crate::http_server::NodesResp> {
        self.nodes
            .iter()
            .map(|node| {
                let joined_at: DateTime<Utc> = node.value().started_at.into();
                let last_keepalive: DateTime<Utc> = node.value().keepalive.into();
                let joined_at = joined_at.format("%Y-%m-%d %H:%M:%S").to_string();
                let last_keepalive = last_keepalive.format("%Y-%m-%d %H:%M:%S").to_string();
                crate::http_server::NodesResp {
                    ip: node.key().to_string(),
                    joined_at,
                    last_keepalive,
                    mints: node.value().mints,
                    ratls_connections: node.value().ratls_conns,
                    ratls_attacks: node.value().ratls_attacks,
                    public: node.value().public,
                    mint_requests: node.value().mint_requests,
                }
            })
            .collect()
    }

    pub fn get_grpc_node_list(&self) -> Vec<NodeUpdate> {
        self.nodes
            .iter()
            .map(|node| NodeUpdate {
                ip: node.key().to_string(),
                started_at: Some(prost_types::Timestamp::from(node.value().started_at)),
                keepalive: Some(prost_types::Timestamp::from(node.value().keepalive)),
                mint_requests: node.value().mint_requests,
                mints: node.value().mints,
                ratls_conns: node.value().ratls_conns,
                ratls_attacks: node.value().ratls_attacks,
                public: node.value().public,
            })
            .collect()
    }

    // returns a random node that does not have an active connection
    pub fn get_random_node(&self) -> Option<String> {
        use rand::{rngs::OsRng, 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))
    }

    pub fn remove_inactive_nodes(&self) {
        self.nodes.retain(|_, v| {
            let age = SystemTime::now()
                .duration_since(v.keepalive)
                .unwrap_or(Duration::ZERO)
                .as_secs();
            if age > 600 {
                false
            } else {
                true
            }
        });
    }
}