#![allow(dead_code)]
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub(crate) struct ClientNode {
enable: bool,
upstream_traffic: u64,
downstream_traffic: u64,
}
impl ClientNode {
pub fn new() -> Self {
Self::default()
}
pub fn add_upstream_traffic(&mut self, traffic: u64) {
self.upstream_traffic = self.upstream_traffic.saturating_add(traffic);
}
pub fn add_downstream_traffic(&mut self, traffic: u64) {
self.downstream_traffic = self.downstream_traffic.saturating_add(traffic);
}
pub fn get_upstream_traffic(&self) -> u64 {
self.upstream_traffic
}
pub fn get_downstream_traffic(&self) -> u64 {
self.downstream_traffic
}
pub fn get_total_traffic(&self) -> u64 {
self.upstream_traffic.saturating_add(self.downstream_traffic)
}
pub fn set_enable(&mut self, enable: bool) {
self.enable = enable;
}
pub fn get_enable(&self) -> bool {
self.enable
}
pub fn reset(&mut self) {
self.upstream_traffic = 0;
self.downstream_traffic = 0;
}
pub fn reset_upstream_traffic(&mut self) {
self.upstream_traffic = 0;
}
pub fn reset_downstream_traffic(&mut self) {
self.downstream_traffic = 0;
}
}
pub(crate) type TrafficAuditPtr = std::sync::Arc<tokio::sync::Mutex<TrafficAudit>>;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub(crate) struct TrafficAudit {
client_map: HashMap<Uuid, ClientNode>,
}
impl TrafficAudit {
pub fn new() -> Self {
Self::default()
}
pub fn add_client(&mut self, client_id: &Uuid) {
self.client_map.entry(*client_id).or_default();
}
pub fn remove_client(&mut self, client_id: &Uuid) {
self.client_map.remove(client_id);
}
pub fn get_client_list(&self) -> Vec<Uuid> {
self.client_map.keys().copied().collect()
}
pub fn contain_client(&self, client_id: &Uuid) -> bool {
self.client_map.contains_key(client_id)
}
pub fn add_upstream_traffic_of(&mut self, client_id: &Uuid, traffic: u64) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.add_upstream_traffic(traffic);
}
}
pub fn add_downstream_traffic_of(&mut self, client_id: &Uuid, traffic: u64) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.add_downstream_traffic(traffic);
}
}
pub fn get_upstream_traffic_of(&self, client_id: &Uuid) -> u64 {
let f = |client_node: &ClientNode| client_node.get_upstream_traffic();
self.client_map.get(client_id).map_or(0, f)
}
pub fn get_downstream_traffic_of(&self, client_id: &Uuid) -> u64 {
let f = |client_node: &ClientNode| client_node.get_downstream_traffic();
self.client_map.get(client_id).map_or(0, f)
}
pub fn get_traffic_of(&self, client_id: &Uuid) -> u64 {
let f = |client_node: &ClientNode| client_node.get_total_traffic();
self.client_map.get(client_id).map_or(0, f)
}
pub fn set_enable_of(&mut self, client_id: &Uuid, enable: bool) {
self.client_map.entry(*client_id).or_default().set_enable(enable);
}
pub fn get_enable_of(&self, client_id: &Uuid) -> bool {
let f = |client_node: &ClientNode| client_node.get_enable();
self.client_map.get(client_id).map(f).unwrap_or(false)
}
pub fn reset(&mut self) {
self.client_map.iter_mut().for_each(|(_, client_node)| client_node.reset());
}
pub fn reset_of(&mut self, client_id: &Uuid) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.reset();
}
}
pub fn reset_upstream_traffic(&mut self) {
self.client_map
.iter_mut()
.for_each(|(_, client_node)| client_node.reset_upstream_traffic());
}
pub fn reset_downstream_traffic(&mut self) {
self.client_map
.iter_mut()
.for_each(|(_, client_node)| client_node.reset_downstream_traffic());
}
pub fn reset_upstream_traffic_of(&mut self, client_id: &Uuid) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.reset_upstream_traffic();
}
}
pub fn reset_downstream_traffic_of(&mut self, client_id: &Uuid) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.reset_downstream_traffic();
}
}
}