105 lines
2.8 KiB
Rust
105 lines
2.8 KiB
Rust
use crate::lease::PeerId;
|
|
use std::collections::HashMap;
|
|
use std::net::IpAddr;
|
|
use std::time::{Duration, Instant};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
|
pub enum TransportProtocol {
|
|
Tcp,
|
|
Udp,
|
|
Icmp,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
|
pub struct FlowKey {
|
|
pub protocol: TransportProtocol,
|
|
pub source: IpAddr,
|
|
pub source_port: u16,
|
|
pub destination: IpAddr,
|
|
pub destination_port: u16,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct FlowEntry {
|
|
pub key: FlowKey,
|
|
pub peer_id: PeerId,
|
|
pub created_at: Instant,
|
|
pub last_seen_at: Instant,
|
|
pub bytes_from_peer: u64,
|
|
pub bytes_to_peer: u64,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct FlowTable {
|
|
entries: HashMap<FlowKey, FlowEntry>,
|
|
}
|
|
|
|
impl FlowTable {
|
|
pub fn upsert(&mut self, key: FlowKey, peer_id: PeerId, now: Instant, bytes_from_peer: u64) {
|
|
self.entries
|
|
.entry(key.clone())
|
|
.and_modify(|entry| {
|
|
entry.last_seen_at = now;
|
|
entry.bytes_from_peer = entry.bytes_from_peer.saturating_add(bytes_from_peer);
|
|
})
|
|
.or_insert(FlowEntry {
|
|
key,
|
|
peer_id,
|
|
created_at: now,
|
|
last_seen_at: now,
|
|
bytes_from_peer,
|
|
bytes_to_peer: 0,
|
|
});
|
|
}
|
|
|
|
pub fn record_return_bytes(&mut self, key: &FlowKey, bytes_to_peer: u64) {
|
|
if let Some(entry) = self.entries.get_mut(key) {
|
|
entry.bytes_to_peer = entry.bytes_to_peer.saturating_add(bytes_to_peer);
|
|
}
|
|
}
|
|
|
|
pub fn expire_idle(&mut self, now: Instant, idle_timeout: Duration) -> usize {
|
|
let before = self.entries.len();
|
|
self.entries
|
|
.retain(|_, entry| now.duration_since(entry.last_seen_at) <= idle_timeout);
|
|
before - self.entries.len()
|
|
}
|
|
|
|
pub fn clear(&mut self) {
|
|
self.entries.clear();
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.entries.len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.entries.is_empty()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::net::{IpAddr, Ipv4Addr};
|
|
|
|
#[test]
|
|
fn expires_idle_flows() {
|
|
let now = Instant::now();
|
|
let mut table = FlowTable::default();
|
|
let key = FlowKey {
|
|
protocol: TransportProtocol::Tcp,
|
|
source: IpAddr::V4(Ipv4Addr::new(10, 241, 0, 2)),
|
|
source_port: 1234,
|
|
destination: IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)),
|
|
destination_port: 443,
|
|
};
|
|
table.upsert(key, PeerId::from_u128(1), now, 10);
|
|
|
|
assert_eq!(table.len(), 1);
|
|
let expired = table.expire_idle(now + Duration::from_secs(61), Duration::from_secs(60));
|
|
assert_eq!(expired, 1);
|
|
assert!(table.is_empty());
|
|
}
|
|
}
|