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

#![allow(dead_code)]
use crate::solana::Client as SolClient;
use dashmap::{DashMap, DashSet};
use std::time::{Duration, SystemTime};

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 let Ok(duration) = self.keepalive.duration_since(other_node.keepalive) {
            if duration > Duration::from_secs(30) {
                return true;
            }
        }
        if 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
    }
}

/// 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 {
    sol_client: SolClient,
    nodes: DashMap<IP, NodeInfo>,
    conns: DashSet<IP>,
}

impl Store {
    pub fn init(sol_client: SolClient) -> Self {
        let store = Self { sol_client, 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_token_address(&self) -> String {
        self.sol_client.token_address()
    }

    pub fn get_pubkey_base58(&self) -> String {
        self.sol_client.wallet_address()
    }

    pub fn get_keypair_bytes(&self) -> Vec<u8> {
        self.sol_client.get_keypair_bytes()
    }

    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 increase_ratls_conns(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.ratls_conns += 1;
        }
    }

    pub fn mint(&self, recipient: &str) -> Result<String, Box<dyn std::error::Error>> {
        use std::str::FromStr;
        let recipient = solana_sdk::pubkey::Pubkey::from_str(recipient)?;
        let sig = self.sol_client.mint(&recipient)?;
        self.increase_mints();
        Ok(sig)
    }

    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, (ip, mut node_info): (String, NodeInfo)) -> bool {
        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;
        }
        println!("Inserting: {}, {:?}", ip, node_info);
        self.nodes.insert(ip, node_info);
        true
    }

    pub fn get_node_list(&self) -> Vec<(String, NodeInfo)> {
        self.nodes.clone().into_iter().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();
            age <= 600
        });
    }
}