hacker-challenge/src/datastore.rs

use crate::persistence::{Logfile, SealError, SealedFile};
use dashmap::{DashMap, DashSet};
use serde::{Deserialize, Serialize};
use serde_with::{serde_as, TimestampSeconds};
use std::time::{Duration, SystemTime};

type IP = String;
pub const LOCALHOST: &str = "localhost";
const LOG_PATH: &str = "/host/main/logs";
const LOCAL_INFO_FILE: &str = "/host/main/node_info";

#[serde_as]
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Ord, PartialOrd)]
pub struct NodeInfo {
    #[serde_as(as = "TimestampSeconds")]
    pub started_at: SystemTime,
    #[serde_as(as = "TimestampSeconds")]
    pub keepalive: SystemTime,
    pub mint_requests: u64,
    pub mints: u64,
    pub mratls_conns: u64,
    pub net_attacks: u64,
    pub public: bool,
    pub restarts: u64,
    pub disk_attacks: u64,
}

impl NodeInfo {
    pub fn is_newer_than(&self, older_self: &Self) -> bool {
        self > older_self
    }

    pub fn to_json(&self) -> String {
        serde_json::to_string(self).unwrap() // can fail only if time goes backwards :D
    }

    pub fn log(&self, ip: &IP) {
        let json = format!("{{\"ip\":\"{}\",", ip) + &self.to_json()[1..];
        if let Err(e) = Logfile::append(LOG_PATH, &format!("{}: {}", ip, &json)) {
            println!("Could not log node info: {:?}", e);
        }
    }

    pub fn load() -> Self {
        match Self::read(LOCAL_INFO_FILE) {
            Ok(mut info) => {
                info.mratls_conns = 0;
                info.restarts += 1;
                info
            }
            Err(SealError::Attack(e)) => {
                println!("The local node file is corrupted: {}", e);
                NodeInfo {
                    started_at: SystemTime::now(),
                    keepalive: SystemTime::now(),
                    mint_requests: 0,
                    mints: 0,
                    mratls_conns: 0,
                    net_attacks: 0,
                    public: false,
                    restarts: 0,
                    disk_attacks: 1, // add very first disk attack
                }
            }
            Err(_) => NodeInfo {
                started_at: SystemTime::now(),
                keepalive: SystemTime::now(),
                mint_requests: 0,
                mints: 0,
                mratls_conns: 0,
                net_attacks: 0,
                public: false,
                restarts: 0,
                disk_attacks: 0,
            },
        }
    }

    pub fn save(&self) {
        if let Err(e) = self.write(LOCAL_INFO_FILE) {
            println!("Could not save node info: {}", e);
        }
    }
}

/// Multithreaded state, designed to be
/// shared everywhere in the code
pub struct State {
    nodes: DashMap<IP, NodeInfo>,
    conns: DashSet<IP>,
}

impl State {
    pub fn new_with_localhost() -> Self {
        let localhost_info = NodeInfo::load();
        let state = Self { nodes: DashMap::new(), conns: DashSet::new() };
        state.nodes.insert(LOCALHOST.to_string(), localhost_info);
        state
    }

    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;
            localhost_info.log(localhost_info.key());
            localhost_info.save();
        }
    }

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

    pub fn increase_mratls_conns(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.mratls_conns += 1;
            localhost_info.log(localhost_info.key());
            localhost_info.save();
        }
    }

    pub fn increase_disk_attacks(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.disk_attacks += 1;
            localhost_info.log(localhost_info.key());
            localhost_info.save();
        }
    }

    pub fn increase_net_attacks(&self) {
        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
            localhost_info.net_attacks += 1;
            localhost_info.log(localhost_info.key());
            localhost_info.save();
        }
    }

    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.
    pub fn process_node_update(&self, (ip, mut node_info): (String, NodeInfo)) -> bool {
        if let Some(old_node) = self.nodes.get(&ip) {
            if !node_info.is_newer_than(&old_node) {
                return false;
            }
            node_info.public = node_info.public || old_node.public;
        }
        println!("Inserting: {}, {}", ip, node_info.to_json());
        node_info.log(&ip);
        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| {
            // HACK: Check if it is possible to abuse network by corrupting system time
            let age =
                SystemTime::now().duration_since(v.keepalive).unwrap_or(Duration::ZERO).as_secs();
            age <= 600
        });
    }
}