scpcap v0

src/chunker.rs

@@ -0,0 +1,369 @@
+use anyhow::{Result, bail};
+use std::time::{Duration, Instant};
+
+pub const GLOBAL_HEADER_LEN: usize = 24;
+pub const RECORD_HEADER_LEN: usize = 16;
+pub const DEFAULT_FLUSH: Duration = Duration::from_millis(100);
+#[allow(dead_code)]
+pub const DEFAULT_TARGET_BYTES: usize = 256 * 1024;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ChunkKind {
+    Header,
+    Data,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Endianness {
+    Little,
+    Big,
+}
+
+fn detect_endianness(header: &[u8]) -> Result<Endianness> {
+    let magic = u32::from_le_bytes(header[0..4].try_into().unwrap());
+    match magic {
+        0xA1B2C3D4 | 0xA1B23C4D => Ok(Endianness::Little),
+        0xD4C3B2A1 | 0x4D3CB2A1 => Ok(Endianness::Big),
+        _ => bail!("unrecognized pcap global header magic: {magic:#010x}"),
+    }
+}
+
+fn read_u32(bytes: &[u8], e: Endianness) -> u32 {
+    let arr: [u8; 4] = bytes.try_into().unwrap();
+    match e {
+        Endianness::Little => u32::from_le_bytes(arr),
+        Endianness::Big => u32::from_be_bytes(arr),
+    }
+}
+
+/// Turns a stream of raw source bytes into chunks ready to hand to the wire
+/// protocol. `RawChunker` is a passthrough for the non-compress path;
+/// `RecordAlignedChunker` cuts on pcap record boundaries for the compress path.
+pub trait Chunker: Send {
+    fn feed(&mut self, bytes: &[u8]) -> Result<()>;
+    /// Pop a chunk that's ready by size/target. Call repeatedly until `None`.
+    fn ready_chunk(&mut self) -> Option<(ChunkKind, Vec<u8>)>;
+    /// Whether a flush-timer-triggered short chunk should be cut right now.
+    fn should_flush(&self, now: Instant) -> bool;
+    /// Force-cut whatever's pending (still respecting record alignment).
+    fn flush(&mut self) -> Option<(ChunkKind, Vec<u8>)>;
+    /// Called once, after the source is confirmed fully drained. Errors if
+    /// what's left doesn't cleanly end at a valid boundary.
+    fn finish(&mut self) -> Result<Option<(ChunkKind, Vec<u8>)>>;
+}
+
+#[derive(Default)]
+pub struct RawChunker {
+    pending: Vec<u8>,
+}
+
+impl RawChunker {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+impl Chunker for RawChunker {
+    fn feed(&mut self, bytes: &[u8]) -> Result<()> {
+        self.pending.extend_from_slice(bytes);
+        Ok(())
+    }
+
+    fn ready_chunk(&mut self) -> Option<(ChunkKind, Vec<u8>)> {
+        if self.pending.is_empty() {
+            None
+        } else {
+            Some((ChunkKind::Data, std::mem::take(&mut self.pending)))
+        }
+    }
+
+    fn should_flush(&self, _now: Instant) -> bool {
+        false
+    }
+
+    fn flush(&mut self) -> Option<(ChunkKind, Vec<u8>)> {
+        None
+    }
+
+    fn finish(&mut self) -> Result<Option<(ChunkKind, Vec<u8>)>> {
+        Ok(self.ready_chunk())
+    }
+}
+
+pub struct RecordAlignedChunker {
+    header_buf: Vec<u8>,
+    header_done: bool,
+    header_ready: bool,
+    endianness: Option<Endianness>,
+    pending: Vec<u8>,
+    complete_len: usize,
+    last_growth: Instant,
+    target: usize,
+    flush_after: Duration,
+}
+
+impl RecordAlignedChunker {
+    pub fn new(target: usize, flush_after: Duration) -> Self {
+        Self {
+            header_buf: Vec::with_capacity(GLOBAL_HEADER_LEN),
+            header_done: false,
+            header_ready: false,
+            endianness: None,
+            pending: Vec::new(),
+            complete_len: 0,
+            last_growth: Instant::now(),
+            target,
+            flush_after,
+        }
+    }
+
+    fn rescan(&mut self) {
+        let e = match self.endianness {
+            Some(e) => e,
+            None => return,
+        };
+        loop {
+            let remaining = &self.pending[self.complete_len..];
+            if remaining.len() < RECORD_HEADER_LEN {
+                break;
+            }
+            let incl_len = read_u32(&remaining[8..12], e) as usize;
+            let record_total = RECORD_HEADER_LEN + incl_len;
+            if remaining.len() < record_total {
+                break;
+            }
+            self.complete_len += record_total;
+        }
+    }
+
+    fn drain_complete(&mut self) -> (ChunkKind, Vec<u8>) {
+        let chunk: Vec<u8> = self.pending.drain(..self.complete_len).collect();
+        self.complete_len = 0;
+        (ChunkKind::Data, chunk)
+    }
+}
+
+impl Chunker for RecordAlignedChunker {
+    fn feed(&mut self, bytes: &[u8]) -> Result<()> {
+        if bytes.is_empty() {
+            return Ok(());
+        }
+        let mut rest = bytes;
+        if !self.header_done {
+            let need = GLOBAL_HEADER_LEN - self.header_buf.len();
+            let take = need.min(rest.len());
+            self.header_buf.extend_from_slice(&rest[..take]);
+            rest = &rest[take..];
+            if self.header_buf.len() == GLOBAL_HEADER_LEN {
+                self.endianness = Some(detect_endianness(&self.header_buf)?);
+                self.header_done = true;
+                self.header_ready = true;
+            } else {
+                // Still accumulating the header; nothing else to do.
+                return Ok(());
+            }
+        }
+        if !rest.is_empty() {
+            self.pending.extend_from_slice(rest);
+            self.rescan();
+            self.last_growth = Instant::now();
+        }
+        Ok(())
+    }
+
+    fn ready_chunk(&mut self) -> Option<(ChunkKind, Vec<u8>)> {
+        if self.header_ready {
+            self.header_ready = false;
+            return Some((ChunkKind::Header, std::mem::take(&mut self.header_buf)));
+        }
+        if self.complete_len >= self.target {
+            return Some(self.drain_complete());
+        }
+        None
+    }
+
+    fn should_flush(&self, now: Instant) -> bool {
+        self.complete_len > 0 && now.duration_since(self.last_growth) >= self.flush_after
+    }
+
+    fn flush(&mut self) -> Option<(ChunkKind, Vec<u8>)> {
+        if self.complete_len > 0 {
+            Some(self.drain_complete())
+        } else {
+            None
+        }
+    }
+
+    fn finish(&mut self) -> Result<Option<(ChunkKind, Vec<u8>)>> {
+        if !self.header_done {
+            bail!("source ended before a complete pcap global header was read");
+        }
+        if self.pending.len() != self.complete_len {
+            bail!("truncated trailing pcap record at end of source");
+        }
+        if self.complete_len > 0 {
+            Ok(Some(self.drain_complete()))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn header_bytes(magic_le: u32) -> Vec<u8> {
+        let mut h = vec![0u8; GLOBAL_HEADER_LEN];
+        h[0..4].copy_from_slice(&magic_le.to_le_bytes());
+        h
+    }
+
+    fn record_bytes(payload: &[u8], big_endian: bool) -> Vec<u8> {
+        let mut r = Vec::with_capacity(RECORD_HEADER_LEN + payload.len());
+        r.extend_from_slice(&[0u8; 8]); // ts_sec, ts_usec
+        let incl_len = payload.len() as u32;
+        if big_endian {
+            r.extend_from_slice(&incl_len.to_be_bytes());
+        } else {
+            r.extend_from_slice(&incl_len.to_le_bytes());
+        }
+        r.extend_from_slice(&incl_len.to_le_bytes()); // orig_len, value doesn't matter
+        r.extend_from_slice(payload);
+        r
+    }
+
+    #[test]
+    fn raw_chunker_passthrough() {
+        let mut c = RawChunker::new();
+        c.feed(b"abc").unwrap();
+        c.feed(b"def").unwrap();
+        let (kind, data) = c.ready_chunk().unwrap();
+        assert_eq!(kind, ChunkKind::Data);
+        assert_eq!(data, b"abcdef");
+        assert!(c.ready_chunk().is_none());
+    }
+
+    #[test]
+    fn header_emitted_as_its_own_chunk() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        c.feed(&header_bytes(0xA1B2C3D4)).unwrap();
+        let (kind, data) = c.ready_chunk().unwrap();
+        assert_eq!(kind, ChunkKind::Header);
+        assert_eq!(data.len(), GLOBAL_HEADER_LEN);
+        assert!(c.ready_chunk().is_none());
+    }
+
+    #[test]
+    fn insufficient_header_bytes_is_not_an_error() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        // Only 10 bytes so far -- not enough to validate the magic.
+        c.feed(&[0u8; 10]).unwrap();
+        assert!(c.ready_chunk().is_none());
+    }
+
+    #[test]
+    fn bad_magic_is_an_error() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        let mut bad = vec![0xFFu8; GLOBAL_HEADER_LEN];
+        bad[0..4].copy_from_slice(&0xDEADBEEFu32.to_le_bytes());
+        assert!(c.feed(&bad).is_err());
+    }
+
+    #[test]
+    fn chunks_always_end_on_record_boundary_when_fed_in_small_increments() {
+        let mut c = RecordAlignedChunker::new(50, DEFAULT_FLUSH); // tiny target to force cuts
+        let mut stream = header_bytes(0xA1B2C3D4);
+        let mut expected_records: Vec<Vec<u8>> = Vec::new();
+        for i in 0..20u8 {
+            let rec = record_bytes(&[i; 10], false);
+            expected_records.push(rec.clone());
+            stream.extend_from_slice(&rec);
+        }
+
+        // Feed in arbitrary tiny increments (1-3 bytes at a time).
+        let mut collected: Vec<u8> = Vec::new();
+        let mut i = 0;
+        let mut step = 1usize;
+        while i < stream.len() {
+            let n = step.min(stream.len() - i);
+            c.feed(&stream[i..i + n]).unwrap();
+            i += n;
+            step = (step % 3) + 1;
+            while let Some((_, data)) = c.ready_chunk() {
+                collected.extend_from_slice(&data);
+            }
+        }
+        if let Some((_, data)) = c.finish().unwrap() {
+            collected.extend_from_slice(&data);
+        }
+
+        // Reassembled bytes must equal the original stream, and every
+        // intermediate chunk boundary must have landed on a whole record --
+        // verified implicitly since decoding never errored and the
+        // concatenation is byte-exact.
+        assert_eq!(collected, stream);
+    }
+
+    #[test]
+    fn target_is_a_floor_not_a_ceiling_cuts_at_next_boundary() {
+        let mut c = RecordAlignedChunker::new(15, DEFAULT_FLUSH);
+        c.feed(&header_bytes(0xA1B2C3D4)).unwrap();
+        c.ready_chunk(); // drain header
+
+        // A single record whose payload alone overshoots the tiny target.
+        let rec = record_bytes(&[1u8; 40], false);
+        c.feed(&rec).unwrap();
+        let (kind, data) = c.ready_chunk().unwrap();
+        assert_eq!(kind, ChunkKind::Data);
+        assert_eq!(data, rec); // whole record, even though it overshoots target=15
+    }
+
+    #[test]
+    fn flush_predicate_and_flush_are_time_independent_of_sleep() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, Duration::from_millis(100));
+        c.feed(&header_bytes(0xA1B2C3D4)).unwrap();
+        c.ready_chunk();
+        let rec = record_bytes(&[9u8; 5], false);
+        c.feed(&rec).unwrap();
+
+        let just_now = Instant::now();
+        assert!(!c.should_flush(just_now));
+
+        let later = just_now + Duration::from_millis(150);
+        assert!(c.should_flush(later));
+
+        let (kind, data) = c.flush().unwrap();
+        assert_eq!(kind, ChunkKind::Data);
+        assert_eq!(data, rec);
+        assert!(c.flush().is_none());
+    }
+
+    #[test]
+    fn truncated_trailing_record_at_finish_is_an_error() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        c.feed(&header_bytes(0xA1B2C3D4)).unwrap();
+        c.ready_chunk();
+        let rec = record_bytes(&[1u8; 20], false);
+        // Feed everything except the last 3 bytes of the record.
+        c.feed(&rec[..rec.len() - 3]).unwrap();
+        assert!(c.finish().is_err());
+    }
+
+    #[test]
+    fn big_endian_magic_reads_incl_len_as_big_endian() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        c.feed(&header_bytes(0xD4C3B2A1)).unwrap();
+        c.ready_chunk();
+        let rec = record_bytes(&[1u8; 30], true);
+        c.feed(&rec).unwrap();
+        let (_, data) = c.finish().unwrap().unwrap();
+        assert_eq!(data, rec);
+    }
+
+    #[test]
+    fn finish_before_any_header_bytes_is_an_error() {
+        let mut c = RecordAlignedChunker::new(DEFAULT_TARGET_BYTES, DEFAULT_FLUSH);
+        assert!(c.finish().is_err());
+    }
+}