use crate::{
event::{AnomalyKind, EndReason, FlowSide, FlowStats},
extractor::{L4Proto, Orientation},
parser_kind::ParserKind,
timestamp::Timestamp,
};
#[cfg(all(feature = "extractors", feature = "reassembler"))]
pub(crate) mod core;
#[cfg(all(feature = "extractors", feature = "reassembler"))]
mod driver;
#[cfg(all(feature = "extractors", feature = "reassembler"))]
pub(crate) mod engine;
#[cfg(all(feature = "extractors", feature = "reassembler"))]
pub use driver::{DatagramDriver, SessionDriver};
pub const DEFAULT_FRAME_DRAIN_MAX_BUFFER: usize = 64 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum FrameDrainError {
BufferFull,
ZeroByteAdvance,
}
impl std::fmt::Display for FrameDrainError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FrameDrainError::BufferFull => f.write_str("parser buffer exceeded max_buffer cap"),
FrameDrainError::ZeroByteAdvance => {
f.write_str("parse_one returned Some((_, 0)) — zero-byte advance")
}
}
}
}
impl std::error::Error for FrameDrainError {}
#[derive(Debug)]
pub struct BufferedFrameDrain<M> {
buf: Vec<u8>,
out: Vec<M>,
max_buffer: usize,
poisoned: Option<FrameDrainError>,
}
impl<M> Default for BufferedFrameDrain<M> {
fn default() -> Self {
Self::new()
}
}
impl<M> Clone for BufferedFrameDrain<M> {
fn clone(&self) -> Self {
Self {
buf: Vec::new(),
out: Vec::new(),
max_buffer: self.max_buffer,
poisoned: None,
}
}
}
impl<M> BufferedFrameDrain<M> {
pub fn new() -> Self {
Self::with_max_buffer(DEFAULT_FRAME_DRAIN_MAX_BUFFER)
}
pub fn with_max_buffer(max_buffer: usize) -> Self {
Self {
buf: Vec::new(),
out: Vec::new(),
max_buffer,
poisoned: None,
}
}
pub fn extend(&mut self, bytes: &[u8]) -> Result<(), FrameDrainError> {
if self.poisoned.is_some() {
return Err(self.poisoned.clone().unwrap());
}
let new_len = self.buf.len().saturating_add(bytes.len());
if new_len > self.max_buffer {
let room = self.max_buffer.saturating_sub(self.buf.len());
self.buf.extend_from_slice(&bytes[..room.min(bytes.len())]);
self.poisoned = Some(FrameDrainError::BufferFull);
return Err(FrameDrainError::BufferFull);
}
self.buf.extend_from_slice(bytes);
Ok(())
}
pub fn drain_with<F>(&mut self, mut parse_one: F)
where
F: FnMut(&[u8]) -> Option<(M, usize)>,
{
while self.poisoned.is_none() {
match parse_one(&self.buf) {
Some((msg, 0)) => {
self.out.push(msg);
self.poisoned = Some(FrameDrainError::ZeroByteAdvance);
return;
}
Some((msg, n)) => {
self.out.push(msg);
if n >= self.buf.len() {
self.buf.clear();
} else {
self.buf.drain(..n);
}
}
None => return,
}
}
}
pub fn take_messages(&mut self) -> Vec<M> {
std::mem::take(&mut self.out)
}
pub fn buffered_len(&self) -> usize {
self.buf.len()
}
pub fn is_poisoned(&self) -> bool {
self.poisoned.is_some()
}
pub fn poison_reason(&self) -> Option<&FrameDrainError> {
self.poisoned.as_ref()
}
}
pub struct AccumulatingSessionParser<F, M>
where
F: Fn(&[u8]) -> Option<(M, usize)> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
parser_kind: ParserKind,
parse_one: F,
max_buffer: usize,
init: BufferedFrameDrain<M>,
resp: BufferedFrameDrain<M>,
}
impl<F, M> std::fmt::Debug for AccumulatingSessionParser<F, M>
where
F: Fn(&[u8]) -> Option<(M, usize)> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AccumulatingSessionParser")
.field("parser_kind", &self.parser_kind)
.field("init_buffered", &self.init.buffered_len())
.field("resp_buffered", &self.resp.buffered_len())
.field(
"poisoned",
&(self.init.is_poisoned() || self.resp.is_poisoned()),
)
.finish()
}
}
impl<F, M> Clone for AccumulatingSessionParser<F, M>
where
F: Fn(&[u8]) -> Option<(M, usize)> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
fn clone(&self) -> Self {
Self {
parser_kind: self.parser_kind,
parse_one: self.parse_one.clone(),
max_buffer: self.max_buffer,
init: BufferedFrameDrain::with_max_buffer(self.max_buffer),
resp: BufferedFrameDrain::with_max_buffer(self.max_buffer),
}
}
}
impl<F, M> AccumulatingSessionParser<F, M>
where
F: Fn(&[u8]) -> Option<(M, usize)> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
pub fn new(parser_kind: ParserKind, parse_one: F) -> Self {
Self::with_max_buffer(parser_kind, parse_one, DEFAULT_FRAME_DRAIN_MAX_BUFFER)
}
pub fn with_max_buffer(parser_kind: ParserKind, parse_one: F, max_buffer: usize) -> Self {
Self {
parser_kind,
parse_one,
max_buffer,
init: BufferedFrameDrain::with_max_buffer(max_buffer),
resp: BufferedFrameDrain::with_max_buffer(max_buffer),
}
}
}
impl<F, M> SessionParser for AccumulatingSessionParser<F, M>
where
F: Fn(&[u8]) -> Option<(M, usize)> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
type Message = M;
fn feed_initiator(&mut self, bytes: &[u8], _ts: Timestamp, out: &mut Vec<M>) {
if self.init.extend(bytes).is_err() {
out.append(&mut self.init.take_messages());
return;
}
let parse_one = self.parse_one.clone();
self.init.drain_with(parse_one);
out.append(&mut self.init.take_messages());
}
fn feed_responder(&mut self, bytes: &[u8], _ts: Timestamp, out: &mut Vec<M>) {
if self.resp.extend(bytes).is_err() {
out.append(&mut self.resp.take_messages());
return;
}
let parse_one = self.parse_one.clone();
self.resp.drain_with(parse_one);
out.append(&mut self.resp.take_messages());
}
fn parser_kind(&self) -> ParserKind {
self.parser_kind
}
fn is_poisoned(&self) -> bool {
self.init.is_poisoned() || self.resp.is_poisoned()
}
fn poison_reason(&self) -> Option<&str> {
let err = self.init.poison_reason().or(self.resp.poison_reason())?;
Some(match err {
FrameDrainError::BufferFull => "buffer cap exceeded",
FrameDrainError::ZeroByteAdvance => "parse_one returned zero-byte advance",
})
}
}
pub struct PerDatagramParser<F, M>
where
F: Fn(&[u8]) -> Option<M> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
parser_kind: ParserKind,
parse_one: F,
}
impl<F, M> std::fmt::Debug for PerDatagramParser<F, M>
where
F: Fn(&[u8]) -> Option<M> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PerDatagramParser")
.field("parser_kind", &self.parser_kind)
.finish()
}
}
impl<F, M> Clone for PerDatagramParser<F, M>
where
F: Fn(&[u8]) -> Option<M> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
fn clone(&self) -> Self {
Self {
parser_kind: self.parser_kind,
parse_one: self.parse_one.clone(),
}
}
}
impl<F, M> PerDatagramParser<F, M>
where
F: Fn(&[u8]) -> Option<M> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
pub fn new(parser_kind: ParserKind, parse_one: F) -> Self {
Self {
parser_kind,
parse_one,
}
}
}
impl<F, M> DatagramParser for PerDatagramParser<F, M>
where
F: Fn(&[u8]) -> Option<M> + Clone + Send + 'static,
M: Send + std::fmt::Debug + 'static,
{
type Message = M;
fn parse(&mut self, payload: &[u8], _side: FlowSide, _ts: Timestamp, out: &mut Vec<M>) {
if let Some(msg) = (self.parse_one)(payload) {
out.push(msg);
}
}
fn parser_kind(&self) -> ParserKind {
self.parser_kind
}
}
pub trait SessionParser: Send + 'static {
type Message: Send + std::fmt::Debug + 'static;
fn feed_initiator(&mut self, bytes: &[u8], ts: Timestamp, out: &mut Vec<Self::Message>);
fn feed_responder(&mut self, bytes: &[u8], ts: Timestamp, out: &mut Vec<Self::Message>);
fn fin_initiator(&mut self, _out: &mut Vec<Self::Message>) {}
fn fin_responder(&mut self, _out: &mut Vec<Self::Message>) {}
fn rst_initiator(&mut self) {}
fn rst_responder(&mut self) {}
fn on_tick(&mut self, _now: Timestamp, _out: &mut Vec<Self::Message>) {}
fn on_gap(
&mut self,
_side: FlowSide,
_missing: u64,
_ts: Timestamp,
_out: &mut Vec<Self::Message>,
) -> GapResponse {
GapResponse::StopSide
}
fn is_poisoned(&self) -> bool {
false
}
fn poison_reason(&self) -> Option<&str> {
None
}
fn is_done(&self) -> bool {
false
}
fn parser_kind(&self) -> ParserKind {
ParserKind::Unspecified
}
}
pub trait SessionParserFactory<K>: Send + 'static {
type Parser: SessionParser;
fn new_parser(&mut self, key: &K) -> Self::Parser;
}
impl<K, P> SessionParserFactory<K> for P
where
P: SessionParser + Default + Clone,
{
type Parser = P;
fn new_parser(&mut self, _key: &K) -> P {
self.clone()
}
}
pub trait DatagramParser: Send + 'static {
type Message: Send + std::fmt::Debug + 'static;
fn parse(
&mut self,
payload: &[u8],
side: FlowSide,
ts: Timestamp,
out: &mut Vec<Self::Message>,
);
fn on_tick(&mut self, _now: Timestamp, _out: &mut Vec<Self::Message>) {}
fn is_poisoned(&self) -> bool {
false
}
fn poison_reason(&self) -> Option<&str> {
None
}
fn is_done(&self) -> bool {
false
}
fn parser_kind(&self) -> ParserKind {
ParserKind::Unspecified
}
fn transports(&self) -> Transports {
Transports::UDP
}
}
bitflags::bitflags! {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Transports: u8 {
const UDP = 1;
const ICMP = 1 << 1;
const ICMPV6 = 1 << 2;
const SCTP = 1 << 3;
const OTHER = 1 << 4;
}
}
impl Transports {
pub const ICMP_ANY: Transports = Transports::ICMP.union(Transports::ICMPV6);
pub fn admits(self, l4: Option<crate::L4Proto>) -> bool {
use crate::L4Proto;
let bit = match l4 {
Some(L4Proto::Udp) => Transports::UDP,
Some(L4Proto::Icmp) => Transports::ICMP,
Some(L4Proto::IcmpV6) => Transports::ICMPV6,
Some(L4Proto::Sctp) => Transports::SCTP,
Some(L4Proto::Tcp) => return false,
_ => Transports::OTHER,
};
self.contains(bit)
}
}
pub trait DatagramParserFactory<K>: Send + 'static {
type Parser: DatagramParser;
fn new_parser(&mut self, key: &K) -> Self::Parser;
fn transports(&self) -> Transports {
Transports::UDP
}
}
impl<K, P> DatagramParserFactory<K> for P
where
P: DatagramParser + Default + Clone,
{
type Parser = P;
fn new_parser(&mut self, _key: &K) -> P {
self.clone()
}
fn transports(&self) -> Transports {
DatagramParser::transports(self)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum GapResponse {
Continue,
#[default]
StopSide,
Stop,
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(tag = "type", rename_all = "snake_case"))]
#[cfg_attr(
feature = "serde",
serde(bound(
serialize = "K: serde::Serialize, M: serde::Serialize",
deserialize = "K: serde::de::DeserializeOwned, M: serde::de::DeserializeOwned"
))
)]
#[non_exhaustive]
pub enum SessionEvent<K, M> {
Started {
key: K,
side: FlowSide,
orientation: Orientation,
ts: Timestamp,
l4: Option<L4Proto>,
},
Application {
key: K,
side: FlowSide,
orientation: Orientation,
message: M,
ts: Timestamp,
parser_kind: ParserKind,
},
ParserClosed {
key: K,
parser_kind: ParserKind,
reason: EndReason,
detail: Option<String>,
ts: Timestamp,
},
ParserSideStopped {
key: K,
parser_kind: ParserKind,
side: FlowSide,
reason: EndReason,
detail: Option<String>,
ts: Timestamp,
},
Closed {
key: K,
reason: EndReason,
stats: FlowStats,
l4: Option<L4Proto>,
ts: Timestamp,
},
FlowAnomaly {
key: K,
kind: AnomalyKind,
ts: Timestamp,
},
TrackerAnomaly { kind: AnomalyKind, ts: Timestamp },
Tick {
key: K,
stats: FlowStats,
ts: Timestamp,
},
}
impl<K, M> SessionEvent<K, M> {
pub fn from_flow_event(ev: crate::event::FlowEvent<K>) -> Option<Self> {
use crate::event::FlowEvent;
Some(match ev {
FlowEvent::Started {
key,
side,
orientation,
ts,
l4,
} => SessionEvent::Started {
key,
side,
orientation,
ts,
l4,
},
FlowEvent::Ended {
key,
reason,
stats,
l4,
..
} => {
let ts = stats.last_seen;
SessionEvent::Closed {
key,
reason,
stats,
l4,
ts,
}
}
FlowEvent::FlowAnomaly { key, kind, ts } => SessionEvent::FlowAnomaly { key, kind, ts },
FlowEvent::TrackerAnomaly { kind, ts } => SessionEvent::TrackerAnomaly { kind, ts },
FlowEvent::Tick { key, stats, ts } => SessionEvent::Tick { key, stats, ts },
FlowEvent::Packet { .. }
| FlowEvent::Established { .. }
| FlowEvent::StateChange { .. } => return None,
})
}
pub fn key(&self) -> Option<&K> {
match self {
SessionEvent::Started { key, .. }
| SessionEvent::Application { key, .. }
| SessionEvent::ParserClosed { key, .. }
| SessionEvent::ParserSideStopped { key, .. }
| SessionEvent::Closed { key, .. }
| SessionEvent::FlowAnomaly { key, .. }
| SessionEvent::Tick { key, .. } => Some(key),
SessionEvent::TrackerAnomaly { .. } => None,
}
}
pub fn timestamp(&self) -> Timestamp {
match self {
SessionEvent::Started { ts, .. }
| SessionEvent::Application { ts, .. }
| SessionEvent::ParserClosed { ts, .. }
| SessionEvent::ParserSideStopped { ts, .. }
| SessionEvent::Closed { ts, .. }
| SessionEvent::FlowAnomaly { ts, .. }
| SessionEvent::TrackerAnomaly { ts, .. }
| SessionEvent::Tick { ts, .. } => *ts,
}
}
pub fn anomaly_kind(&self) -> Option<&AnomalyKind> {
match self {
SessionEvent::FlowAnomaly { kind, .. } | SessionEvent::TrackerAnomaly { kind, .. } => {
Some(kind)
}
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub struct TemplateFactory<P>(pub P);
impl<K, P> SessionParserFactory<K> for TemplateFactory<P>
where
P: SessionParser + Clone,
{
type Parser = P;
fn new_parser(&mut self, _key: &K) -> P {
self.0.clone()
}
}
impl<K, P> DatagramParserFactory<K> for TemplateFactory<P>
where
P: DatagramParser + Clone,
{
type Parser = P;
fn new_parser(&mut self, _key: &K) -> P {
self.0.clone()
}
fn transports(&self) -> Transports {
self.0.transports()
}
}
#[cfg(any(feature = "ftp", feature = "smtp"))]
pub(crate) fn skip_partial_line<'a>(pending: &mut bool, bytes: &'a [u8]) -> &'a [u8] {
if !*pending {
return bytes;
}
match bytes.iter().position(|&b| b == b'\n') {
Some(p) => {
*pending = false;
&bytes[p + 1..]
}
None => &[],
}
}
#[cfg(any(feature = "smb", feature = "modbus"))]
pub(crate) fn resync_frames(
pending: &mut bool,
buf: &mut bytes::BytesMut,
keep: usize,
find: impl Fn(&[u8]) -> Option<usize>,
) {
if !*pending {
return;
}
match find(buf) {
Some(at) => {
let _ = buf.split_to(at);
*pending = false;
}
None => {
let drop = buf.len().saturating_sub(keep);
let _ = buf.split_to(drop);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn flow_events_map_to_their_session_form() {
use crate::event::FlowEvent;
let stats = FlowStats {
last_seen: Timestamp::new(7, 0),
..Default::default()
};
let closed: Option<SessionEvent<u8, ()>> =
SessionEvent::from_flow_event(FlowEvent::Ended {
key: 1u8,
reason: EndReason::Fin,
stats,
history: Default::default(),
l4: None,
});
assert!(matches!(
closed,
Some(SessionEvent::Closed { key: 1, reason: EndReason::Fin, ts, .. }) if ts == Timestamp::new(7, 0)
));
let tick: Option<SessionEvent<u8, ()>> = SessionEvent::from_flow_event(FlowEvent::Tick {
key: 2u8,
stats: FlowStats::default(),
ts: Timestamp::new(8, 0),
});
assert!(matches!(tick, Some(SessionEvent::Tick { key: 2, .. })));
}
#[derive(Default, Clone)]
struct CountParser {
init_bytes: usize,
resp_bytes: usize,
}
impl SessionParser for CountParser {
type Message = (FlowSide, usize);
fn feed_initiator(&mut self, b: &[u8], _ts: Timestamp, out: &mut Vec<Self::Message>) {
self.init_bytes += b.len();
out.push((FlowSide::Initiator, self.init_bytes));
}
fn feed_responder(&mut self, b: &[u8], _ts: Timestamp, out: &mut Vec<Self::Message>) {
self.resp_bytes += b.len();
out.push((FlowSide::Responder, self.resp_bytes));
}
}
#[test]
fn auto_impl_session_parser_factory() {
let mut f: CountParser = CountParser::default();
let mut p: CountParser = SessionParserFactory::<u32>::new_parser(&mut f, &7);
let mut m = Vec::new();
p.feed_initiator(b"abc", Timestamp::default(), &mut m);
assert_eq!(m, vec![(FlowSide::Initiator, 3)]);
}
#[derive(Default, Clone)]
struct EchoDgram;
impl DatagramParser for EchoDgram {
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()));
}
}
#[test]
fn auto_impl_datagram_parser_factory() {
let mut f = EchoDgram;
let mut p: EchoDgram = DatagramParserFactory::<()>::new_parser(&mut f, &());
let mut m = Vec::new();
p.parse(b"hello", FlowSide::Responder, Timestamp::default(), &mut m);
assert_eq!(m, vec![(FlowSide::Responder, b"hello".to_vec())]);
}
}