use std::hash::Hash;
use std::time::Duration;
use crate::Timestamp;
use crate::extractor::FlowExtractor;
use crate::session::SessionParser;
use crate::tracker::FlowTrackerConfig;
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
use crate::session::DatagramParser;
type IdleTimeoutFn<K> =
Box<dyn Fn(&K, Option<crate::L4Proto>) -> Option<Duration> + Send + 'static>;
use crate::session_driver::FlowSessionDriver;
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
use crate::datagram_driver::FlowDatagramDriver;
#[cfg(all(feature = "reassembler", feature = "session"))]
pub struct FlowSessionDriverBuilder<E, P>
where
E: FlowExtractor,
P: SessionParser + Clone,
{
extractor: E,
parser: Option<P>,
config: FlowTrackerConfig,
emit_anomalies: bool,
monotonic_timestamps: bool,
dedup: Option<crate::dedup::Dedup>,
idle_timeout_fn: Option<IdleTimeoutFn<E::Key>>,
}
#[cfg(all(feature = "reassembler", feature = "session"))]
impl<E, P> FlowSessionDriverBuilder<E, P>
where
E: FlowExtractor,
E::Key: Hash + Eq + Clone + Send + 'static,
P: SessionParser + Clone + Send + 'static,
{
pub(crate) fn new(extractor: E) -> Self {
Self {
extractor,
parser: None,
config: FlowTrackerConfig::default(),
emit_anomalies: false,
monotonic_timestamps: false,
dedup: None,
idle_timeout_fn: None,
}
}
pub fn parser(mut self, parser: P) -> Self {
self.parser = Some(parser);
self
}
pub fn config(mut self, config: FlowTrackerConfig) -> Self {
self.config = config;
self
}
pub fn emit_anomalies(mut self, enable: bool) -> Self {
self.emit_anomalies = enable;
self
}
pub fn monotonic_timestamps(mut self, enable: bool) -> Self {
self.monotonic_timestamps = enable;
self
}
pub fn dedup(mut self, dedup: crate::dedup::Dedup) -> Self {
self.dedup = Some(dedup);
self
}
pub fn idle_timeout_fn<F>(mut self, f: F) -> Self
where
F: Fn(&E::Key, Option<crate::L4Proto>) -> Option<Duration> + Send + 'static,
{
self.idle_timeout_fn = Some(Box::new(f));
self
}
pub fn build(self) -> FlowSessionDriver<E, P, ()> {
let parser = self
.parser
.expect("FlowSessionDriverBuilder: parser not set");
let mut driver = FlowSessionDriver::with_config(self.extractor, parser, self.config)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic_timestamps);
if let Some(dedup) = self.dedup {
driver = driver.with_dedup(dedup);
}
if let Some(f) = self.idle_timeout_fn {
driver = driver.with_idle_timeout_fn(move |k, l4| f(k, l4));
}
driver
}
}
#[cfg(all(feature = "reassembler", feature = "session"))]
impl<E, P, S> FlowSessionDriver<E, P, S>
where
E: FlowExtractor,
E::Key: Hash + Eq + Clone + Send + 'static,
P: SessionParser + Clone + Send + 'static,
S: Send + 'static,
{
pub fn builder(extractor: E) -> FlowSessionDriverBuilder<E, P> {
FlowSessionDriverBuilder::new(extractor)
}
}
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
pub struct FlowDatagramDriverBuilder<E, P>
where
E: FlowExtractor,
P: DatagramParser + Clone,
{
extractor: E,
parser: Option<P>,
config: FlowTrackerConfig,
emit_anomalies: bool,
monotonic_timestamps: bool,
dedup: Option<crate::dedup::Dedup>,
}
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
impl<E, P> FlowDatagramDriverBuilder<E, P>
where
E: FlowExtractor,
E::Key: Hash + Eq + Clone + Send + 'static,
P: DatagramParser + Clone + Send + 'static,
{
pub(crate) fn new(extractor: E) -> Self {
Self {
extractor,
parser: None,
config: FlowTrackerConfig::default(),
emit_anomalies: false,
monotonic_timestamps: false,
dedup: None,
}
}
pub fn parser(mut self, parser: P) -> Self {
self.parser = Some(parser);
self
}
pub fn config(mut self, config: FlowTrackerConfig) -> Self {
self.config = config;
self
}
pub fn emit_anomalies(mut self, enable: bool) -> Self {
self.emit_anomalies = enable;
self
}
pub fn monotonic_timestamps(mut self, enable: bool) -> Self {
self.monotonic_timestamps = enable;
self
}
pub fn dedup(mut self, dedup: crate::dedup::Dedup) -> Self {
self.dedup = Some(dedup);
self
}
pub fn build(self) -> FlowDatagramDriver<E, P, ()> {
let parser = self
.parser
.expect("FlowDatagramDriverBuilder: parser not set");
let mut driver = FlowDatagramDriver::with_config(self.extractor, parser, self.config)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic_timestamps);
if let Some(dedup) = self.dedup {
driver = driver.with_dedup(dedup);
}
driver
}
}
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
impl<E, P, S> FlowDatagramDriver<E, P, S>
where
E: FlowExtractor,
E::Key: Hash + Eq + Clone + Send + 'static,
P: DatagramParser + Clone + Send + 'static,
S: Send + 'static,
{
pub fn builder(extractor: E) -> FlowDatagramDriverBuilder<E, P> {
FlowDatagramDriverBuilder::new(extractor)
}
}
#[allow(unused_imports)]
use Timestamp as _UseTimestamp;
#[cfg(test)]
mod tests {
use super::*;
use crate::extract::FiveTuple;
#[derive(Default, Clone)]
struct NoopParser;
impl SessionParser for NoopParser {
type Message = ();
fn feed_initiator(&mut self, _b: &[u8], _ts: Timestamp) -> Vec<()> {
Vec::new()
}
fn feed_responder(&mut self, _b: &[u8], _ts: Timestamp) -> Vec<()> {
Vec::new()
}
}
#[derive(Default, Clone)]
struct NoopDatagram;
impl DatagramParser for NoopDatagram {
type Message = ();
fn parse(&mut self, _b: &[u8], _side: crate::event::FlowSide, _ts: Timestamp) -> Vec<()> {
Vec::new()
}
}
#[test]
#[cfg(all(feature = "reassembler", feature = "session"))]
fn session_builder_yields_working_driver() {
let _d = FlowSessionDriver::<_, NoopParser, ()>::builder(FiveTuple::bidirectional())
.parser(NoopParser)
.emit_anomalies(true)
.monotonic_timestamps(true)
.build();
}
#[test]
#[cfg(all(feature = "extractors", feature = "reassembler", feature = "session"))]
fn datagram_builder_yields_working_driver() {
let _d = FlowDatagramDriver::<_, NoopDatagram, ()>::builder(FiveTuple::bidirectional())
.parser(NoopDatagram)
.emit_anomalies(true)
.build();
}
#[test]
#[cfg(all(feature = "reassembler", feature = "session"))]
#[should_panic(expected = "parser not set")]
fn missing_parser_panics_on_build() {
let _ = FlowSessionDriver::<_, NoopParser, ()>::builder(FiveTuple::bidirectional()).build();
}
}