flowscope 0.22.0

//! Crate-internal sync UDP-datagram engine. Wraps a [`FlowDriver`]
//! (with a no-op reassembler factory) and adds per-flow
//! [`DatagramParser`] dispatch, yielding [`SessionEvent`]s.
//!
//! This was the public `FlowDatagramDriver` through 0.19; it was
//! demoted to a private engine in 0.20 (#99) — the typed
//! [`crate::driver::Driver`] plus one datagram slot is the supported
//! single-parser surface now. The engine survives because the typed
//! slots ([`crate::driver`]) and the offline [`crate::pcap`] source
//! both need parser dispatch.

use std::{collections::HashMap, hash::Hash};

use ahash::RandomState;

use crate::{
    Timestamp,
    event::{AnomalyKind, EndReason, FlowEvent, FlowSide},
    extractor::FlowExtractor,
    flow_driver::FlowDriver,
    reassembler::{Reassembler, ReassemblerFactory},
    session::{DatagramParser, SessionEvent},
    tracker::FlowTrackerConfig,
    view::PacketView,
};

/// Boxed per-flow parser factory closure. Each new flow gets its
/// parser by calling this on the flow's key.
type ParserFactory<K, P> = Box<dyn FnMut(&K) -> P + Send + Sync>;

/// Cap on the size of `poison_reason()` strings carried through
/// [`AnomalyKind::SessionParseError`]. Matches the
/// [`crate::session_driver`] cap.
const POISON_REASON_MAX_BYTES: usize = 256;

fn truncate_reason(s: &str) -> String {
    let mut owned = String::from(s);
    if owned.len() > POISON_REASON_MAX_BYTES {
        let cap = (0..=POISON_REASON_MAX_BYTES)
            .rev()
            .find(|i| owned.is_char_boundary(*i))
            .unwrap_or(0);
        owned.truncate(cap);
    }
    owned
}

/// A no-op reassembler. UDP traffic doesn't reassemble, but the
/// inner `FlowDriver` requires a `ReassemblerFactory`. This stub
/// is created on TCP-payload callbacks (which UDP traffic never
/// triggers); the slot stays empty in practice.
#[derive(Debug, Default)]
struct NoopReassembler;

impl Reassembler for NoopReassembler {
    fn segment(&mut self, _seq: u32, _payload: &[u8], _ts: crate::Timestamp) {}
}

#[derive(Debug, Default)]
struct NoopReassemblerFactory;

impl<K: Send + 'static> ReassemblerFactory<K> for NoopReassemblerFactory {
    type Reassembler = NoopReassembler;
    fn new_reassembler(&mut self, _key: &K, _side: FlowSide) -> NoopReassembler {
        NoopReassembler
    }
}

/// Crate-internal sync UDP-datagram engine. Owns a flow tracker +
/// per-flow [`DatagramParser`] instances, yielding [`SessionEvent`]s.
/// `S` is fixed at `()` since 0.20 (the stateful constructors were
/// dropped with the public surface).
pub(crate) struct FlowDatagramDriver<E, P, S = ()>
where
    E: FlowExtractor,
    E::Key: Hash + Eq + Clone + Send + Sync + 'static,
    P: DatagramParser + Send + Sync + 'static,
    S: Send + 'static,
{
    driver: FlowDriver<E, NoopReassemblerFactory, S>,
    parser_factory: ParserFactory<E::Key, P>,
    parsers: HashMap<E::Key, P, RandomState>,
}

// Template-parser path — `S = ()`, `P: Clone`.
impl<E, P> FlowDatagramDriver<E, P, ()>
where
    E: FlowExtractor,
    E::Key: Hash + Eq + Clone + Send + Sync + 'static,
    P: DatagramParser + Clone + Send + Sync + 'static,
{
    /// Construct with default tracker config and `S = ()`. `parser`
    /// is cloned once per flow to give each flow a fresh instance.
    ///
    /// Convenience entry used by the offline [`crate::pcap`] source
    /// and the unit tests; the typed-driver slots build via
    /// [`Self::with_config`].
    #[cfg(any(feature = "pcap", test))]
    pub fn new(extractor: E, parser: P) -> Self {
        Self::with_config(extractor, parser, FlowTrackerConfig::default())
    }

    /// Construct with explicit tracker config and `S = ()`.
    pub fn with_config(extractor: E, parser: P, config: FlowTrackerConfig) -> Self {
        Self {
            driver: FlowDriver::with_config(extractor, NoopReassemblerFactory, config),
            parser_factory: Box::new(move |_key| parser.clone()),
            parsers: HashMap::with_hasher(RandomState::new()),
        }
    }
}

// All non-construction methods — apply to every `S` and every
// parser-source variant.
impl<E, P, S> FlowDatagramDriver<E, P, S>
where
    E: FlowExtractor,
    E::Key: Hash + Eq + Clone + Send + Sync + 'static,
    P: DatagramParser + Send + Sync + 'static,
    S: Send + 'static,
{
    /// Opt in to monotonic timestamps.
    pub fn with_monotonic_timestamps(mut self, enable: bool) -> Self {
        self.driver = self.driver.with_monotonic_timestamps(enable);
        self
    }

    /// Drive one packet. Returns zero or more [`SessionEvent`]s.
    ///
    /// Allocating convenience used by the offline [`crate::pcap`]
    /// source and the unit tests; the typed-driver slots use the
    /// zero-alloc [`Self::track_into`].
    #[cfg(any(feature = "pcap", test))]
    pub fn track<'v>(
        &mut self,
        view: impl Into<PacketView<'v>>,
    ) -> Vec<SessionEvent<E::Key, P::Message>> {
        let mut out = Vec::new();
        self.track_into(view, &mut out);
        out
    }

    /// Append-only variant of [`Self::track`]. Reuses the caller's
    /// capacity — zero allocation at this surface in steady state.
    /// Plan 121 zero-alloc threading.
    pub fn track_into<'v>(
        &mut self,
        view: impl Into<PacketView<'v>>,
        out: &mut Vec<SessionEvent<E::Key, P::Message>>,
    ) {
        let view: PacketView<'v> = view.into();
        let mut flow_events = self.driver.track_pending(view);
        // `PacketView` is `Copy` and `frame: &[u8]` is immutable
        // after `track_pending`; dedup / monotonic_timestamps only
        // modify timestamps, not frame bytes. Re-extract the UDP
        // payload here as a zero-copy &[u8] instead of allocating
        // a per-packet `to_vec()` clone before the call.
        let udp_payload: Option<&[u8]> = extract_udp_payload(view);
        self.translate_events_into(&flow_events, udp_payload, out);
        self.driver.finalize(flow_events.as_mut_slice());
    }

    /// Run the idle-timeout sweep. Also drives each still-live
    /// parser's [`DatagramParser::on_tick`] hook with `now`, emitting
    /// any time-driven messages as initiator-side `Application`
    /// events.
    pub fn sweep(&mut self, now: Timestamp) -> Vec<SessionEvent<E::Key, P::Message>> {
        let mut flow_events = self.driver.sweep_pending(now);
        // Fire `on_tick` on every live parser *before* translating
        // the swept events, so a flow this sweep closes still gets
        // its final tick ahead of its `Closed` event.
        let mut out: Vec<SessionEvent<E::Key, P::Message>> = Vec::new();
        // Plan 80: defer tear-down of any parser that finished
        // during on_tick to after the iteration so we don't mutate
        // self.parsers while borrowing it.
        enum TickClose<R> {
            Poison(FlowSide, Option<R>),
            Done,
        }
        let mut closes: Vec<(E::Key, TickClose<String>)> = Vec::new();
        let mut tick_scratch: Vec<P::Message> = Vec::new();
        for (key, parser) in self.parsers.iter_mut() {
            let kind = parser.parser_kind();
            tick_scratch.clear();
            parser.on_tick(now, &mut tick_scratch);
            for m in tick_scratch.drain(..) {
                crate::obs::trace_session_message(FlowSide::Initiator, &m);
                out.push(SessionEvent::Application {
                    key: key.clone(),
                    side: FlowSide::Initiator,
                    message: m,
                    ts: now,
                    parser_kind: kind,
                });
            }
            if parser.is_poisoned() {
                let reason = parser.poison_reason().map(truncate_reason);
                closes.push((key.clone(), TickClose::Poison(FlowSide::Initiator, reason)));
            } else if parser.is_done() {
                closes.push((key.clone(), TickClose::Done));
            }
        }
        let closed_keys: std::collections::HashSet<E::Key> =
            closes.iter().map(|(k, _)| k.clone()).collect();
        for (key, status) in closes {
            match status {
                TickClose::Poison(side, reason) => {
                    self.synthesise_parser_poison(&key, side, reason, now, &mut out);
                }
                TickClose::Done => {
                    self.synthesise_parser_done(&key, now, &mut out);
                }
            }
        }
        if !closed_keys.is_empty() {
            flow_events.retain(|ev| match ev {
                FlowEvent::Ended { key, .. } => !closed_keys.contains(key),
                _ => true,
            });
        }
        out.extend(self.translate_events(&flow_events, None));
        self.driver.finalize(flow_events.as_mut_slice());
        out
    }

    /// Sweep every remaining flow, emitting `Closed` events. Call
    /// once at end of input — equivalent to `sweep(Timestamp::MAX)`.
    pub fn finish(&mut self) -> Vec<SessionEvent<E::Key, P::Message>> {
        self.sweep(Timestamp::MAX)
    }

    /// Force-end the UDP flow with this key. New in 0.8.0. Mirror of
    /// the session-driver / flow-driver counterparts. No reassembler
    /// to tear down; just removes the parser slot, calls
    /// [`FlowDriver::force_close`], and translates the driver's
    /// terminal event into `SessionEvent::Closed { reason:
    /// ForceClosed, .. }`.
    pub fn force_close(
        &mut self,
        key: &E::Key,
        now: Timestamp,
    ) -> Vec<SessionEvent<E::Key, P::Message>> {
        self.parsers.remove(key);
        let driver_events = self.driver.force_close(key, now);
        self.translate_events(&driver_events, None)
    }

    fn translate_events(
        &mut self,
        flow_events: &[FlowEvent<E::Key>],
        udp_payload: Option<&[u8]>,
    ) -> Vec<SessionEvent<E::Key, P::Message>> {
        let mut out = Vec::new();
        self.translate_events_into(flow_events, udp_payload, &mut out);
        out
    }

    fn translate_events_into(
        &mut self,
        flow_events: &[FlowEvent<E::Key>],
        udp_payload: Option<&[u8]>,
        out: &mut Vec<SessionEvent<E::Key, P::Message>>,
    ) {
        for ev in flow_events {
            match ev {
                FlowEvent::Started { key, ts, .. } => {
                    self.parsers
                        .entry(key.clone())
                        .or_insert_with_key(|k| (self.parser_factory)(k));
                    out.push(SessionEvent::Started {
                        key: key.clone(),
                        ts: *ts,
                    });
                }
                FlowEvent::Packet { key, side, ts, .. } => {
                    let Some(payload) = udp_payload else {
                        // Non-datagram packet or no payload — skip silently.
                        continue;
                    };
                    let Some(parser) = self.parsers.get_mut(key) else {
                        continue;
                    };
                    let kind = parser.parser_kind();
                    let mut messages: Vec<P::Message> = Vec::new();
                    parser.parse(payload, *side, *ts, &mut messages);
                    for m in messages.drain(..) {
                        crate::obs::trace_session_message(*side, &m);
                        out.push(SessionEvent::Application {
                            key: key.clone(),
                            side: *side,
                            message: m,
                            ts: *ts,
                            parser_kind: kind,
                        });
                    }
                    // Plan 55 — parser poison check.
                    // Plan 80 — parser-driven graceful close; poison wins.
                    if parser.is_poisoned() {
                        let reason = parser.poison_reason().map(truncate_reason);
                        self.synthesise_parser_poison(key, *side, reason, *ts, out);
                    } else if parser.is_done() {
                        self.synthesise_parser_done(key, *ts, out);
                    }
                }
                FlowEvent::Ended {
                    key,
                    reason,
                    stats,
                    l4,
                    ..
                } => {
                    self.parsers.remove(key);
                    out.push(SessionEvent::Closed {
                        key: key.clone(),
                        reason: *reason,
                        stats: stats.clone(),
                        l4: *l4,
                    });
                }
                FlowEvent::FlowAnomaly { key, kind, ts } => {
                    out.push(SessionEvent::FlowAnomaly {
                        key: key.clone(),
                        kind: kind.clone(),
                        ts: *ts,
                    });
                }
                FlowEvent::TrackerAnomaly { kind, ts } => {
                    out.push(SessionEvent::TrackerAnomaly {
                        kind: kind.clone(),
                        ts: *ts,
                    });
                }
                FlowEvent::Tick { key, stats, ts } => {
                    out.push(SessionEvent::Tick {
                        key: key.clone(),
                        stats: stats.clone(),
                        ts: *ts,
                    });
                }
                FlowEvent::Established { .. } | FlowEvent::StateChange { .. } => {
                    // TCP-only events; ignored.
                }
            }
        }
    }

    fn synthesise_parser_poison(
        &mut self,
        key: &E::Key,
        side: FlowSide,
        reason: Option<String>,
        ts: Timestamp,
        out: &mut Vec<SessionEvent<E::Key, P::Message>>,
    ) {
        if self.driver.emits_anomalies() {
            out.push(SessionEvent::FlowAnomaly {
                key: key.clone(),
                kind: AnomalyKind::SessionParseError { side, reason },
                ts,
            });
        }
        let stats = self
            .driver
            .tracker()
            .snapshot_stats(key)
            .unwrap_or_default();
        let l4 = self.driver.tracker().snapshot_l4(key);
        crate::obs::record_flow_ended(EndReason::ParseError, &stats);
        crate::obs::trace_flow_ended(EndReason::ParseError, &stats);
        out.push(SessionEvent::Closed {
            key: key.clone(),
            reason: EndReason::ParseError,
            stats,
            l4,
        });
        self.parsers.remove(key);
        self.driver.tracker_mut().forget(key);
    }

    /// Plan 80: parser-driven graceful close (mirror of
    /// [`Self::synthesise_parser_poison`] without the anomaly).
    fn synthesise_parser_done(
        &mut self,
        key: &E::Key,
        ts: Timestamp,
        out: &mut Vec<SessionEvent<E::Key, P::Message>>,
    ) {
        let _ = ts;
        let stats = self
            .driver
            .tracker()
            .snapshot_stats(key)
            .unwrap_or_default();
        let l4 = self.driver.tracker().snapshot_l4(key);
        crate::obs::record_flow_ended(EndReason::ParserDone, &stats);
        crate::obs::trace_flow_ended(EndReason::ParserDone, &stats);
        out.push(SessionEvent::Closed {
            key: key.clone(),
            reason: EndReason::ParserDone,
            stats,
            l4,
        });
        self.parsers.remove(key);
        self.driver.tracker_mut().forget(key);
    }
}

/// Extract the datagram payload a [`DatagramParser`] should see from an
/// Ethernet-framed view. Returns `None` if the frame can't be parsed or
/// isn't a datagram-shaped transport.
///
/// - **UDP** → the UDP payload (after the 8-byte UDP header).
/// - **ICMPv4 / ICMPv6** → the *full* ICMP message bytes (type + code +
///   rest), which is what [`crate::icmp::IcmpParser`] parses. Without
///   this, `datagram_broadcast(IcmpParser::new())` silently delivered
///   nothing (the ICMP datagram path was UDP-only before 0.14.1).
fn extract_udp_payload(view: PacketView<'_>) -> Option<&[u8]> {
    let sp = etherparse::SlicedPacket::from_ethernet(view.frame).ok()?;
    match sp.transport? {
        etherparse::TransportSlice::Udp(udp) => Some(udp.payload()),
        etherparse::TransportSlice::Icmpv4(icmp) => Some(icmp.slice()),
        etherparse::TransportSlice::Icmpv6(icmp) => Some(icmp.slice()),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::extract::{FiveTuple, parse::test_frames::*};

    fn view(frame: &[u8], sec: u32) -> PacketView<'_> {
        PacketView::new(frame, Timestamp::new(sec, 0))
    }

    /// Echo parser: emits one Vec<u8> per packet, side-tagged.
    #[derive(Default, Clone)]
    struct EchoUdp;
    impl DatagramParser for EchoUdp {
        type Message = (FlowSide, Vec<u8>);
        fn parse(
            &mut self,
            payload: &[u8],
            side: FlowSide,
            _ts: Timestamp,
            out: &mut Vec<Self::Message>,
        ) {
            out.push((side, payload.to_vec()));
        }
    }

    /// Plan 33: `finish()` closes every still-open UDP flow.
    #[test]
    fn finish_closes_open_flows() {
        let mut d = FlowDatagramDriver::new(FiveTuple::bidirectional(), EchoUdp);
        let f = ipv4_udp([10, 0, 0, 1], [10, 0, 0, 2], 1, 53, b"q");
        d.track(view(&f, 0));
        let closed = d
            .finish()
            .into_iter()
            .filter(|e| matches!(e, SessionEvent::Closed { .. }))
            .count();
        assert_eq!(closed, 1, "finish() must close the open flow");
        assert!(d.finish().is_empty(), "second finish() yields nothing");
    }

    /// Plan 36: `on_tick` fires on `finish` for a live UDP flow,
    /// before that flow's `Closed`.
    #[test]
    fn on_tick_fires_on_finish() {
        #[derive(Default, Clone)]
        struct TickParser;
        impl DatagramParser for TickParser {
            type Message = u8;
            fn parse(&mut self, _p: &[u8], _s: FlowSide, _ts: Timestamp, _out: &mut Vec<u8>) {}
            fn on_tick(&mut self, _now: Timestamp, out: &mut Vec<u8>) {
                out.push(7);
            }
        }
        let mut d = FlowDatagramDriver::new(FiveTuple::bidirectional(), TickParser);
        let f = ipv4_udp([10, 0, 0, 1], [10, 0, 0, 2], 1, 53, b"q");
        d.track(view(&f, 0));
        let count = |evs: Vec<SessionEvent<_, u8>>| {
            evs.iter()
                .filter(|e| matches!(e, SessionEvent::Application { message: 7, .. }))
                .count()
        };
        // finish() closes the UDP flow but drives a final on_tick.
        assert_eq!(count(d.finish()), 1);
        // No flows remain → no further ticks.
        assert_eq!(count(d.sweep(Timestamp::new(99, 0))), 0);
    }

    #[test]
    fn started_and_application_for_udp_packet() {
        let mut d = FlowDatagramDriver::new(FiveTuple::bidirectional(), EchoUdp);
        let f = ipv4_udp([10, 0, 0, 1], [10, 0, 0, 2], 1, 53, b"query");
        let events = d.track(view(&f, 0));
        assert!(
            events
                .iter()
                .any(|e| matches!(e, SessionEvent::Started { .. }))
        );
        let app = events.iter().find_map(|e| match e {
            SessionEvent::Application {
                message: (s, b), ..
            } => Some((*s, b.clone())),
            _ => None,
        });
        assert_eq!(app, Some((FlowSide::Initiator, b"query".to_vec())));
    }

    #[test]
    fn closed_event_on_idle_timeout() {
        let cfg = FlowTrackerConfig {
            idle_timeout_udp: std::time::Duration::from_secs(1),
            ..FlowTrackerConfig::default()
        };
        let mut d = FlowDatagramDriver::with_config(FiveTuple::bidirectional(), EchoUdp, cfg);
        let f = ipv4_udp([10, 0, 0, 1], [10, 0, 0, 2], 1, 53, b"q");
        d.track(view(&f, 0));
        let ended = d.sweep(Timestamp::new(10, 0));
        assert!(
            ended
                .iter()
                .any(|e| matches!(e, SessionEvent::Closed { .. }))
        );
    }

    #[test]
    fn tcp_packets_do_not_fire_application_events() {
        let mut d = FlowDatagramDriver::new(FiveTuple::bidirectional(), EchoUdp);
        let syn = ipv4_tcp(
            [0; 6],
            [0; 6],
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            0,
            0,
            0x02,
            b"",
        );
        let events = d.track(view(&syn, 0));
        assert!(
            events
                .iter()
                .any(|e| matches!(e, SessionEvent::Started { .. }))
        );
        assert!(
            !events
                .iter()
                .any(|e| matches!(e, SessionEvent::Application { .. })),
            "TCP packet produced an Application event in the UDP driver"
        );
    }

    /// Parser that poisons after >5 bytes on either side.
    #[derive(Default, Clone)]
    struct PoisonAfterBytes {
        seen: usize,
        poisoned: bool,
    }
    impl DatagramParser for PoisonAfterBytes {
        type Message = ();
        fn parse(&mut self, payload: &[u8], _side: FlowSide, _ts: Timestamp, _out: &mut Vec<()>) {
            self.seen += payload.len();
            if self.seen > 5 {
                self.poisoned = true;
            }
        }
        fn is_poisoned(&self) -> bool {
            self.poisoned
        }
        fn poison_reason(&self) -> Option<&str> {
            if self.poisoned {
                Some("too many bytes")
            } else {
                None
            }
        }
    }

    #[test]
    fn datagram_parser_poison_synthesises_parse_error_closed() {
        let mut d =
            FlowDatagramDriver::new(FiveTuple::bidirectional(), PoisonAfterBytes::default());
        let f = ipv4_udp([10, 0, 0, 1], [10, 0, 0, 2], 1, 53, b"0123456789");
        let events = d.track(view(&f, 0));
        let closed = events.iter().find_map(|e| match e {
            SessionEvent::Closed { reason, .. } => Some(*reason),
            _ => None,
        });
        assert_eq!(closed, Some(EndReason::ParseError));
    }
}