use std::sync::LazyLock;
use flowscope::Timestamp;
use flowscope::well_known::LabelTable;
mod from_ctx;
mod split;
pub use from_ctx::{CounterRegistry, FlowStateRegistry, StateMap};
use crate::correlate::TimeBucketedCounter;
use crate::protocol::FlowKey;
static DEFAULT_LABEL_TABLE: LazyLock<LabelTable> = LazyLock::new(LabelTable::new);
pub(crate) fn default_label_table() -> &'static LabelTable {
&DEFAULT_LABEL_TABLE
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceIdx(pub u8);
pub struct Ctx<'a> {
pub flow: Option<FlowKey>,
pub ts: Timestamp,
pub source: SourceIdx,
pub monitor_name: Option<&'a str>,
pub(crate) state_map: &'a mut StateMap,
pub(crate) sink: &'a mut dyn crate::anomaly::sink::AnomalySink,
pub(crate) counters: &'a mut CounterRegistry,
pub(crate) flow_states: &'a mut FlowStateRegistry,
pub(crate) label_table: &'a LabelTable,
pub(crate) tracker: Option<&'a flowscope::FlowTracker<flowscope::extract::FiveTuple, ()>>,
#[cfg_attr(not(feature = "arp"), allow(dead_code))]
pub(crate) arp_table:
Option<&'a flowscope::correlate::NeighborTable<std::net::Ipv4Addr, flowscope::MacAddr>>,
}
impl<'a> Ctx<'a> {
#[cfg(feature = "bench-zero-alloc")]
pub fn new_for_bench(
ts: Timestamp,
state_map: &'a mut StateMap,
sink: &'a mut dyn crate::anomaly::sink::AnomalySink,
counters: &'a mut CounterRegistry,
flow_states: &'a mut FlowStateRegistry,
) -> Self {
Self {
flow: None,
ts,
source: SourceIdx(0),
monitor_name: None,
state_map,
sink,
counters,
flow_states,
label_table: default_label_table(),
tracker: None,
arp_table: None,
}
}
#[doc(hidden)]
pub fn new(
flow: Option<FlowKey>,
ts: Timestamp,
source: SourceIdx,
state_map: &'a mut StateMap,
sink: &'a mut dyn crate::anomaly::sink::AnomalySink,
counters: &'a mut CounterRegistry,
flow_states: &'a mut FlowStateRegistry,
) -> Self {
Self {
flow,
ts,
source,
monitor_name: None,
state_map,
sink,
counters,
flow_states,
label_table: default_label_table(),
tracker: None,
arp_table: None,
}
}
#[inline]
pub fn community_id(&self) -> Option<String> {
self.flow
.as_ref()
.and_then(flowscope::KeyFields::community_id)
}
#[cfg(feature = "dns")]
pub fn names(&self, ip: std::net::IpAddr) -> Vec<String> {
match self.state::<crate::monitor::dns::NameMapState>() {
Some(st) => st
.map
.peek_names(ip, self.ts)
.into_iter()
.map(|c| c.name.clone())
.collect(),
None => Vec::new(),
}
}
#[inline]
pub fn state_mut<T: Default + Send + 'static>(&mut self) -> &mut T {
self.state_map.get_or_init_mut::<T>()
}
#[inline]
pub fn state<T: 'static>(&self) -> Option<&T> {
self.state_map.get::<T>()
}
#[inline]
pub fn label_table(&self) -> &LabelTable {
self.label_table
}
#[inline]
pub fn counter<K>(&self) -> Option<&TimeBucketedCounter<K>>
where
K: std::hash::Hash + Eq + Clone + Send + 'static,
{
self.counters.get::<K>()
}
#[inline]
pub fn bandwidth(&self) -> Option<crate::monitor::BandwidthReport<'_>> {
let state = self.state::<crate::monitor::bandwidth::BandwidthState>()?;
Some(crate::monitor::BandwidthReport {
rate: &state.0,
now: self.ts,
})
}
#[cfg(feature = "icmp")]
#[inline]
pub fn lookup_icmp_flow(
&self,
inner: &flowscope::icmp::IcmpInner,
) -> Option<(FlowKey, flowscope::FlowStats)> {
self.tracker?.stats_for_inner(inner)
}
#[cfg(feature = "arp")]
#[inline]
pub fn arp_table(&self) -> Option<&flowscope::correlate::ArpTable> {
self.arp_table
}
#[inline]
pub fn counter_mut<K>(&mut self) -> &mut TimeBucketedCounter<K>
where
K: std::hash::Hash + Eq + Clone + Send + 'static,
{
self.counters.get_mut::<K>()
}
#[inline]
pub fn sink_mut(&mut self) -> &mut dyn crate::anomaly::sink::AnomalySink {
self.sink
}
#[inline]
pub fn flow_state_mut<T>(&mut self) -> Option<&mut T>
where
T: Default + Send + 'static,
{
let key = self.flow?;
let ts = self.ts;
let map = self.flow_states.get_mut::<T>()?;
Some(map.get_or_default(&key, ts))
}
#[inline]
pub fn emit(
&mut self,
kind: &'static str,
severity: crate::anomaly::Severity,
) -> crate::anomaly::sink::AnomalyWriter<'_> {
let ts = self.ts;
self.sink.begin(kind, severity, ts)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::anomaly::sink::NoopSink;
#[derive(Default)]
struct DemoState {
n: u64,
}
fn make_ctx<'a>(
state: &'a mut StateMap,
sink: &'a mut NoopSink,
counters: &'a mut CounterRegistry,
flow_states: &'a mut FlowStateRegistry,
) -> Ctx<'a> {
Ctx {
flow: None,
ts: Timestamp::new(0, 0),
source: SourceIdx(0),
monitor_name: None,
state_map: state,
sink,
counters,
flow_states,
label_table: default_label_table(),
tracker: None,
arp_table: None,
}
}
#[test]
fn ctx_constructs_from_borrowed_fields() {
let mut state = StateMap::default();
let mut counters = CounterRegistry::default();
let mut sink = NoopSink;
let mut flow_states = FlowStateRegistry::default();
let _ctx = make_ctx(&mut state, &mut sink, &mut counters, &mut flow_states);
}
#[test]
fn state_mut_lazy_creates_then_returns_same() {
let mut state = StateMap::default();
let mut counters = CounterRegistry::default();
let mut sink = NoopSink;
let mut flow_states = FlowStateRegistry::default();
let mut ctx = make_ctx(&mut state, &mut sink, &mut counters, &mut flow_states);
ctx.state_mut::<DemoState>().n = 7;
assert_eq!(ctx.state_mut::<DemoState>().n, 7);
}
#[test]
fn sink_mut_returns_dyn_anomalysink() {
let mut state = StateMap::default();
let mut counters = CounterRegistry::default();
let mut sink = NoopSink;
let mut flow_states = FlowStateRegistry::default();
let mut ctx = make_ctx(&mut state, &mut sink, &mut counters, &mut flow_states);
let _: &mut dyn crate::anomaly::sink::AnomalySink = ctx.sink_mut();
}
#[test]
fn counter_mut_returns_registered_counter() {
use std::time::Duration;
let mut state = StateMap::default();
let mut counters = CounterRegistry::default();
counters.register::<u16>(TimeBucketedCounter::<u16>::new_unbounded(
Duration::from_secs(60),
Duration::from_secs(1),
));
let mut sink = NoopSink;
let mut flow_states = FlowStateRegistry::default();
let mut ctx = make_ctx(&mut state, &mut sink, &mut counters, &mut flow_states);
ctx.counter_mut::<u16>().bump(42u16, Timestamp::new(0, 0));
}
#[test]
fn source_idx_roundtrip() {
assert_eq!(SourceIdx(3), SourceIdx(3));
assert_ne!(SourceIdx(1), SourceIdx(2));
}
}