diff --git a/src/rust/lqos_bus/src/ip_stats.rs b/src/rust/lqos_bus/src/ip_stats.rs index 0f2a89c2..9c0e8b46 100644 --- a/src/rust/lqos_bus/src/ip_stats.rs +++ b/src/rust/lqos_bus/src/ip_stats.rs @@ -196,4 +196,10 @@ pub struct FlowbeeData { pub tos: u8, /// Raw TCP flags pub flags: u8, + /// Remote ASN + pub remote_asn: u32, + /// Remote ASN Name + pub remote_asn_name: String, + /// Remote ASN Country + pub remote_asn_country: String, } \ No newline at end of file diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/asn.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/asn.rs index ebaa503f..82b89d9e 100644 --- a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/asn.rs +++ b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/asn.rs @@ -98,4 +98,8 @@ impl AsnTable { } }).map(|idx| self.asn_table[idx].clone()).ok() } + + pub fn find_asn_by_id(&self, asn: u32) -> Option { + self.asn_table.iter().find(|row| row.asn == asn).map(|row| row.clone()) + } } diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/mod.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/mod.rs index fa4eaf67..0497b5be 100644 --- a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/mod.rs +++ b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_analysis/mod.rs @@ -1,4 +1,4 @@ -use std::sync::Mutex; +use std::{net::IpAddr, sync::Mutex}; use once_cell::sync::Lazy; use self::asn::AsnTable; @@ -41,3 +41,37 @@ pub fn setup_flow_analysis() -> anyhow::Result<()> { } Ok(()) } + +pub fn lookup_asn_id(ip: IpAddr) -> Option { + let table_lock = ANALYSIS.asn_table.lock(); + if table_lock.is_err() { + return None; + } + let table = table_lock.unwrap(); + if table.is_none() { + return None; + } + let table = table.as_ref().unwrap(); + if let Some(asn) = table.find_asn(ip) { + Some(asn.asn) + } else { + None + } +} + +pub fn get_asn_name_and_country(asn: u32) -> (String, String) { + let table_lock = ANALYSIS.asn_table.lock(); + if table_lock.is_err() { + return ("".to_string(), "".to_string()); + } + let table = table_lock.unwrap(); + if table.is_none() { + return ("".to_string(), "".to_string()); + } + let table = table.as_ref().unwrap(); + if let Some(row) = table.find_asn_by_id(asn) { + (row.owners.clone(), row.country.clone()) + } else { + ("".to_string(), "".to_string()) + } +} \ No newline at end of file diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_tracker.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_tracker.rs index c81cc912..b3dfaafd 100644 --- a/src/rust/lqosd/src/throughput_tracker/flow_data/flow_tracker.rs +++ b/src/rust/lqosd/src/throughput_tracker/flow_data/flow_tracker.rs @@ -5,5 +5,8 @@ use dashmap::DashMap; use lqos_sys::flowbee_data::{FlowbeeData, FlowbeeKey}; use once_cell::sync::Lazy; -pub static ALL_FLOWS: Lazy> = Lazy::new(|| DashMap::new()); +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct AsnId(pub u32); + +pub static ALL_FLOWS: Lazy> = Lazy::new(|| DashMap::new()); diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/mod.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/mod.rs index 7e8ce22d..4b7df971 100644 --- a/src/rust/lqosd/src/throughput_tracker/flow_data/mod.rs +++ b/src/rust/lqosd/src/throughput_tracker/flow_data/mod.rs @@ -7,13 +7,14 @@ mod netflow9; mod flow_analysis; use crate::throughput_tracker::flow_data::{netflow5::Netflow5, netflow9::Netflow9}; -pub(crate) use flow_tracker::ALL_FLOWS; +pub(crate) use flow_tracker::{ALL_FLOWS, AsnId}; use lqos_sys::flowbee_data::{FlowbeeData, FlowbeeKey}; use std::sync::{ mpsc::{channel, Sender}, Arc, }; -pub(crate) use flow_analysis::setup_flow_analysis; +pub(crate) use flow_analysis::{setup_flow_analysis, lookup_asn_id, get_asn_name_and_country}; + trait FlowbeeRecipient { fn enqueue(&self, key: FlowbeeKey, data: FlowbeeData); diff --git a/src/rust/lqosd/src/throughput_tracker/mod.rs b/src/rust/lqosd/src/throughput_tracker/mod.rs index 995bd7b9..f331e4c0 100644 --- a/src/rust/lqosd/src/throughput_tracker/mod.rs +++ b/src/rust/lqosd/src/throughput_tracker/mod.rs @@ -18,7 +18,7 @@ use tokio::{ time::{Duration, Instant}, }; -use self::flow_data::ALL_FLOWS; +use self::flow_data::{get_asn_name_and_country, AsnId, ALL_FLOWS}; const RETIRE_AFTER_SECONDS: u64 = 30; @@ -486,20 +486,25 @@ pub fn all_unknown_ips() -> BusResponse { /// For debugging: dump all active flows! pub fn dump_active_flows() -> BusResponse { let result: Vec = ALL_FLOWS.iter().map(|row| { + let (remote_asn_name, remote_asn_country) = get_asn_name_and_country(row.value().1.0); + lqos_bus::FlowbeeData { remote_ip: row.key().remote_ip.as_ip().to_string(), local_ip: row.key().local_ip.as_ip().to_string(), src_port: row.key().src_port, dst_port: row.key().dst_port, ip_protocol: FlowbeeProtocol::from(row.key().ip_protocol), - bytes_sent: row.value().bytes_sent, - packets_sent: row.value().packets_sent, - rate_estimate_bps: row.value().rate_estimate_bps, - retries: row.value().retries, - last_rtt: row.value().last_rtt, - end_status: row.value().end_status, - tos: row.value().tos, - flags: row.value().flags, + bytes_sent: row.value().0.bytes_sent, + packets_sent: row.value().0.packets_sent, + rate_estimate_bps: row.value().0.rate_estimate_bps, + retries: row.value().0.retries, + last_rtt: row.value().0.last_rtt, + end_status: row.value().0.end_status, + tos: row.value().0.tos, + flags: row.value().0.flags, + remote_asn: row.value().1.0, + remote_asn_name, + remote_asn_country, } }).collect(); @@ -513,7 +518,7 @@ pub fn all_unknown_ips() -> BusResponse { /// Top Flows Report pub fn top_flows(n: u32, flow_type: TopFlowType) -> BusResponse { - let mut table: Vec<(FlowbeeKey, FlowbeeData)> = ALL_FLOWS + let mut table: Vec<(FlowbeeKey, (FlowbeeData, AsnId))> = ALL_FLOWS .iter() .map(|row| ( row.key().clone(), @@ -524,36 +529,36 @@ pub fn all_unknown_ips() -> BusResponse { match flow_type { TopFlowType::RateEstimate => { table.sort_by(|a, b| { - let a_total = a.1.rate_estimate_bps[0] + a.1.rate_estimate_bps[1]; - let b_total = b.1.rate_estimate_bps[0] + b.1.rate_estimate_bps[1]; + let a_total = a.1.0.rate_estimate_bps[0] + a.1.0.rate_estimate_bps[1]; + let b_total = b.1.0.rate_estimate_bps[0] + b.1.0.rate_estimate_bps[1]; b_total.cmp(&a_total) }); } TopFlowType::Bytes => { table.sort_by(|a, b| { - let a_total = a.1.bytes_sent[0] + a.1.bytes_sent[1]; - let b_total = b.1.bytes_sent[0] + b.1.bytes_sent[1]; + let a_total = a.1.0.bytes_sent[0] + a.1.0.bytes_sent[1]; + let b_total = b.1.0.bytes_sent[0] + b.1.0.bytes_sent[1]; b_total.cmp(&a_total) }); } TopFlowType::Packets => { table.sort_by(|a, b| { - let a_total = a.1.packets_sent[0] + a.1.packets_sent[1]; - let b_total = b.1.packets_sent[0] + b.1.packets_sent[1]; + let a_total = a.1.0.packets_sent[0] + a.1.0.packets_sent[1]; + let b_total = b.1.0.packets_sent[0] + b.1.0.packets_sent[1]; b_total.cmp(&a_total) }); } TopFlowType::Drops => { table.sort_by(|a, b| { - let a_total = a.1.retries[0] + a.1.retries[1]; - let b_total = b.1.retries[0] + b.1.retries[1]; + let a_total = a.1.0.retries[0] + a.1.0.retries[1]; + let b_total = b.1.0.retries[0] + b.1.0.retries[1]; b_total.cmp(&a_total) }); } TopFlowType::RoundTripTime => { table.sort_by(|a, b| { - let a_total = a.1.last_rtt[0] + a.1.last_rtt[1]; - let b_total = b.1.last_rtt[0] + b.1.last_rtt[1]; + let a_total = a.1.0.last_rtt[0] + a.1.0.last_rtt[1]; + let b_total = b.1.0.last_rtt[0] + b.1.0.last_rtt[1]; b_total.cmp(&a_total) }); } @@ -563,20 +568,24 @@ pub fn all_unknown_ips() -> BusResponse { .iter() .take(n as usize) .map(|(ip, flow)| { + let (remote_asn_name, remote_asn_country) = get_asn_name_and_country(flow.1.0); lqos_bus::FlowbeeData { remote_ip: ip.remote_ip.as_ip().to_string(), local_ip: ip.local_ip.as_ip().to_string(), src_port: ip.src_port, dst_port: ip.dst_port, ip_protocol: FlowbeeProtocol::from(ip.ip_protocol), - bytes_sent: flow.bytes_sent, - packets_sent: flow.packets_sent, - rate_estimate_bps: flow.rate_estimate_bps, - retries: flow.retries, - last_rtt: flow.last_rtt, - end_status: flow.end_status, - tos: flow.tos, - flags: flow.flags, + bytes_sent: flow.0.bytes_sent, + packets_sent: flow.0.packets_sent, + rate_estimate_bps: flow.0.rate_estimate_bps, + retries: flow.0.retries, + last_rtt: flow.0.last_rtt, + end_status: flow.0.end_status, + tos: flow.0.tos, + flags: flow.0.flags, + remote_asn: flow.1.0, + remote_asn_name, + remote_asn_country, } }) .collect(); diff --git a/src/rust/lqosd/src/throughput_tracker/tracking_data.rs b/src/rust/lqosd/src/throughput_tracker/tracking_data.rs index bf5818fe..7f0ab16a 100644 --- a/src/rust/lqosd/src/throughput_tracker/tracking_data.rs +++ b/src/rust/lqosd/src/throughput_tracker/tracking_data.rs @@ -1,6 +1,6 @@ use std::{sync::atomic::AtomicU64, time::Duration}; use crate::{shaped_devices_tracker::{SHAPED_DEVICES, NETWORK_JSON}, stats::{HIGH_WATERMARK_DOWN, HIGH_WATERMARK_UP}}; -use super::{flow_data::ALL_FLOWS, throughput_entry::ThroughputEntry, RETIRE_AFTER_SECONDS}; +use super::{flow_data::{lookup_asn_id, AsnId, ALL_FLOWS}, throughput_entry::ThroughputEntry, RETIRE_AFTER_SECONDS}; use dashmap::DashMap; use lqos_bus::TcHandle; use lqos_sys::{flowbee_data::{FlowbeeData, FlowbeeKey}, iterate_flows, throughput_for_each}; @@ -203,18 +203,20 @@ impl ThroughputTracker { } else { // We have a valid flow, so it needs to be tracked if let Some(mut this_flow) = ALL_FLOWS.get_mut(&key) { - this_flow.last_seen = data.last_seen; - this_flow.bytes_sent = data.bytes_sent; - this_flow.packets_sent = data.packets_sent; - this_flow.rate_estimate_bps = data.rate_estimate_bps; - this_flow.retries = data.retries; - this_flow.last_rtt = data.last_rtt; - this_flow.end_status = data.end_status; - this_flow.tos = data.tos; - this_flow.flags = data.flags; + this_flow.0.last_seen = data.last_seen; + this_flow.0.bytes_sent = data.bytes_sent; + this_flow.0.packets_sent = data.packets_sent; + this_flow.0.rate_estimate_bps = data.rate_estimate_bps; + this_flow.0.retries = data.retries; + this_flow.0.last_rtt = data.last_rtt; + this_flow.0.end_status = data.end_status; + this_flow.0.tos = data.tos; + this_flow.0.flags = data.flags; } else { // Insert it into the map - ALL_FLOWS.insert(key.clone(), data.clone()); + let asn = lookup_asn_id(key.remote_ip.as_ip()).unwrap_or(0); + + ALL_FLOWS.insert(key.clone(), (data.clone(), AsnId(asn))); // TODO: Submit it for analysis }