borer-core 0.5.9

use std::{
    collections::{HashMap, HashSet},
    ops::{Deref, DerefMut},
    sync::Arc,
    time::{SystemTime, UNIX_EPOCH},
};

use dashmap::DashMap;
use once_cell::sync::Lazy;
use serde::Serialize;

static STORE: Lazy<Arc<Store>> = Lazy::new(Store::arc);

/// Return the process-wide shared traffic store.
pub fn store() -> Arc<Store> {
    STORE.clone()
}

/// Shared in-memory state for active connections and aggregated traffic.
#[derive(Debug)]
pub struct Store {
    user_map: DashMap<String /*user*/, HashSet<String /*hash*/>>,
    hash_traffic_map: DashMap<String /*hash*/, Traffic>,
    conns_map: DashMap<String /*hash*/, HashSet<String /*conn*/>>,
    conn_map: DashMap<String /*conn*/, Connection>,
}

/// Traffic counters measured in bytes.
#[derive(Clone, Debug, Default, Serialize)]
pub struct Traffic {
    pub up: usize,
    pub down: usize,
}

/// Aggregated per-user traffic and active connection count.
#[derive(Clone, Debug, Default, Serialize)]
pub struct UserStats {
    pub up: usize,
    pub down: usize,
    pub conn_count: usize,
}

/// Snapshot of a tracked proxied connection.
#[derive(Clone, Debug, Serialize)]
pub struct Connection {
    pub conn_id: String,
    pub user: String,
    pub hash: String,
    pub peer_addr: String,
    pub req_addr: String,
    pub traffic: Traffic,
    pub padding: bool,
    pub create_at: u64,
}

impl Store {
    pub fn arc() -> Arc<Self> {
        Arc::new(Store {
            user_map: DashMap::new(),
            hash_traffic_map: DashMap::new(),
            conns_map: DashMap::new(),
            conn_map: DashMap::new(),
        })
    }

    pub fn get_traffic_all(&self) -> HashMap<String, UserStats> {
        self.user_map
            .iter()
            .filter_map(|m| {
                let user = m.key().clone();
                let conns = self.get_conns_by_user(&user);
                let conn_count = conns.map(|v| v.len()).unwrap_or(0);
                self.get_traffic_by_user(&user).map(|t| {
                    (
                        user,
                        UserStats {
                            up: t.up,
                            down: t.down,
                            conn_count,
                        },
                    )
                })
            })
            .collect()
    }

    pub fn get_conns_all(&self) -> HashMap<String, Vec<Connection>> {
        self.user_map
            .iter()
            .filter_map(|m| {
                let user = m.key().clone();
                self.get_conns_by_user(&user).map(|conns| (user, conns))
            })
            .collect()
    }

    pub fn get_traffic_by_user(&self, user: &str) -> Option<Traffic> {
        if let Some(hashs) = self.user_map.get(user) {
            let (up, down) = hashs
                .iter()
                .map(|h| self.get_traffic_by_hash(h))
                .filter(|t| t.is_some())
                .map(|t| t.map(|t| (t.get_up(), t.get_down())).unwrap_or((0, 0)))
                .reduce(|(up1, down1), (up2, down2)| (up1 + up2, down1 + down2))
                .unwrap_or((0, 0));
            Some(Traffic { up, down })
        } else {
            None
        }
    }

    pub fn get_traffic_by_hash(&self, hash: &str) -> Option<Traffic> {
        self.hash_traffic_map.get(hash).map(|t| t.deref().clone())
    }

    pub fn get_conns_by_user(&self, user: &str) -> Option<Vec<Connection>> {
        if let Some(hashs) = self.user_map.get(user) {
            let conns = hashs
                .iter()
                .filter_map(|h| self.get_conns_by_hash(h))
                .flatten()
                .collect();
            Some(conns)
        } else {
            None
        }
    }

    pub fn get_conns_by_hash(&self, hash: &str) -> Option<Vec<Connection>> {
        if let Some(conn_ids) = self.conns_map.get(hash) {
            let conns = conn_ids
                .iter()
                .filter_map(|conn_id| self.conn_map.get(conn_id))
                .map(|c| c.clone())
                .collect();
            Some(conns)
        } else {
            None
        }
    }

    pub fn insert_conn<S>(
        &self,
        user: S,
        hash: S,
        conn_id: S,
        peer_addr: S,
        req_addr: S,
        padding: bool,
    ) where
        S: Into<String>,
    {
        let user = user.into();
        let hash = hash.into();
        let conn_id = conn_id.into();
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system time failed!")
            .as_secs();
        let conn = Connection {
            user: user.to_string(),
            hash: hash.to_string(),
            conn_id: conn_id.to_string(),
            peer_addr: peer_addr.into(),
            req_addr: req_addr.into(),
            traffic: Default::default(),
            padding,
            create_at: timestamp,
        };
        self.insert_hash(&user, &hash);
        self.insert_conns(&hash, &conn_id);
        self.conn_map.insert(conn_id, conn);
    }

    fn insert_hash(&self, user: &str, hash: &str) {
        if let Some(mut v) = self.user_map.get_mut(user) {
            v.insert(hash.to_string());
        } else {
            self.user_map
                .insert(user.to_string(), HashSet::from_iter([hash.to_string()]));
        }
    }

    fn insert_conns(&self, hash: &str, conn_id: &str) {
        if let Some(mut conns) = self.conns_map.get_mut(hash) {
            conns.insert(conn_id.to_string());
        } else {
            self.conns_map
                .insert(hash.to_string(), HashSet::from_iter([conn_id.to_string()]));
        }
    }

    pub fn delete_conn(&self, conn_id: &str) {
        if let Some((_, conn)) = self.conn_map.remove(conn_id) {
            if let Some(mut conns) = self.conns_map.get_mut(&conn.hash) {
                conns.retain(|v| v != conn_id);
            }
        }
    }

    pub fn add_up(&self, conn_id: &str, v: usize) {
        if let Some(mut conn) = self.conn_map.get_mut(conn_id) {
            let hash = &conn.hash;
            self.add_up_by_hash(hash, v);
            conn.deref_mut().traffic.add_up(v);
        }
    }

    pub fn add_down(&self, conn_id: &str, v: usize) {
        if let Some(mut conn) = self.conn_map.get_mut(conn_id) {
            let hash = &conn.hash;
            self.add_donw_by_hash(hash, v);
            conn.deref_mut().traffic.add_down(v);
        }
    }

    fn add_up_by_hash(&self, hash: &str, v: usize) {
        if let Some(mut t) = self.hash_traffic_map.get_mut(hash) {
            t.add_up(v);
        } else {
            self.hash_traffic_map
                .insert(hash.to_string(), Traffic { up: v, down: 0 });
        }
    }

    fn add_donw_by_hash(&self, hash: &str, v: usize) {
        if let Some(mut t) = self.hash_traffic_map.get_mut(hash) {
            t.add_down(v);
        } else {
            self.hash_traffic_map
                .insert(hash.to_string(), Traffic { up: 0, down: v });
        }
    }
}

impl Traffic {
    pub fn get_up(&self) -> usize {
        self.up
    }
    pub fn get_down(&self) -> usize {
        self.down
    }
    pub fn add_up(&mut self, v: usize) {
        self.up += v
    }
    pub fn add_down(&mut self, v: usize) {
        self.down += v
    }
}

#[cfg(test)]
mod tests {
    use super::Store;

    #[test]
    fn insert_and_aggregate_traffic_by_user() {
        let store = Store::arc();
        store.insert_conn(
            "alice",
            "hash-a",
            "conn-1",
            "127.0.0.1:1",
            "example.com:443",
            false,
        );
        store.insert_conn(
            "alice",
            "hash-b",
            "conn-2",
            "127.0.0.1:2",
            "example.com:80",
            true,
        );
        store.add_up("conn-1", 5);
        store.add_down("conn-1", 7);
        store.add_up("conn-2", 11);

        let traffic = store.get_traffic_by_user("alice").unwrap();
        let all = store.get_traffic_all();

        assert_eq!(traffic.up, 16);
        assert_eq!(traffic.down, 7);
        assert_eq!(all.get("alice").unwrap().conn_count, 2);
    }

    #[test]
    fn get_conns_by_hash_and_delete_conn_updates_views() {
        let store = Store::arc();
        store.insert_conn(
            "alice",
            "hash-a",
            "conn-1",
            "127.0.0.1:1",
            "example.com:443",
            false,
        );
        store.insert_conn(
            "alice",
            "hash-a",
            "conn-2",
            "127.0.0.1:2",
            "example.com:80",
            true,
        );

        let before = store.get_conns_by_hash("hash-a").unwrap();
        store.delete_conn("conn-1");
        let after = store.get_conns_by_hash("hash-a").unwrap();

        assert_eq!(before.len(), 2);
        assert_eq!(after.len(), 1);
        assert_eq!(after[0].conn_id, "conn-2");
    }

    #[test]
    fn unknown_user_returns_none() {
        let store = Store::arc();

        assert!(store.get_traffic_by_user("missing").is_none());
        assert!(store.get_conns_by_user("missing").is_none());
    }
}