use crate::{
event::{AnomalyKind, EndReason, FlowSide, FlowStats},
parser_kind::ParserKind,
timestamp::Timestamp,
};
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 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
}
}
pub trait DatagramParserFactory<K>: Send + 'static {
type Parser: DatagramParser;
fn new_parser(&mut self, key: &K) -> Self::Parser;
}
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()
}
}
#[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]
#[allow(dead_code)]
pub(crate) enum SessionEvent<K, M> {
Started { key: K, ts: Timestamp },
Application {
key: K,
side: FlowSide,
message: M,
ts: Timestamp,
parser_kind: ParserKind,
},
Closed {
key: K,
reason: EndReason,
stats: FlowStats,
l4: Option<crate::extractor::L4Proto>,
},
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> {
#[cfg(test)]
pub(crate) fn anomaly_kind(&self) -> Option<&AnomalyKind> {
match self {
SessionEvent::FlowAnomaly { kind, .. } | SessionEvent::TrackerAnomaly { kind, .. } => {
Some(kind)
}
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[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())]);
}
}