#![allow(dead_code)]
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[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 += traffic;
}
pub fn add_downstream_traffic(&mut self, traffic: u64) {
self.downstream_traffic += 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 + 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<String, ClientNode>,
}
impl TrafficAudit {
pub fn new() -> Self {
Self::default()
}
pub fn add_client(&mut self, client_id: &str) {
self.client_map.entry(client_id.to_string()).or_default();
}
pub fn remove_client(&mut self, client_id: &str) {
self.client_map.remove(client_id);
}
pub fn get_client_list(&self) -> Vec<String> {
self.client_map.keys().map(|s| s.to_string()).collect()
}
pub fn contain_client(&self, client_id: &str) -> bool {
self.client_map.contains_key(client_id)
}
pub fn add_upstream_traffic_of(&mut self, client_id: &str, 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: &str, 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: &str) -> 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: &str) -> 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: &str) -> 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: &str, enable: bool) {
self.client_map.entry(client_id.to_string()).or_default().set_enable(enable);
}
pub fn get_enable_of(&self, client_id: &str) -> 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: &str) {
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: &str) {
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: &str) {
if let Some(client_node) = self.client_map.get_mut(client_id) {
client_node.reset_downstream_traffic();
}
}
}