pcapsql_core/io/

pcap_stream.rs

//! Generic PCAP/PCAPNG reader over any Read source.
//!
//! This module provides a unified PCAP parser that works with any `R: Read` source,
//! using the battle-tested `pcap_parser` crate. Benchmarks show this approach is
//! actually 30% faster than custom byte parsing while being more maintainable.
//!
//! ## Usage
//!
//! ```ignore
//! use std::fs::File;
//! use std::io::Cursor;
//!
//! // From a file with known format
//! let file = File::open("capture.pcap")?;
//! let reader = GenericPcapReader::with_format(file, PcapFormat::LegacyLeMicro)?;
//!
//! // From memory-mapped data
//! let mmap = Arc::new(unsafe { Mmap::map(&file)? });
//! let cursor = Cursor::new(MmapSlice::new(mmap));
//! let reader = GenericPcapReader::with_format(cursor, PcapFormat::LegacyLeMicro)?;
//! ```

use std::io::{BufReader, Read};

use bytes::Bytes;
use pcap_parser::traits::PcapReaderIterator;
use pcap_parser::{LegacyPcapReader, PcapBlockOwned, PcapNGReader};

use crate::error::{Error, PcapError};
use crate::io::{PacketRef, RawPacket};

/// Buffer size for pcap_parser readers (64KB).
const BUFFER_SIZE: usize = 262144;

/// Format of the PCAP file.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PcapFormat {
    /// Classic PCAP (little-endian, microseconds)
    LegacyLeMicro,
    /// Classic PCAP (big-endian, microseconds)
    LegacyBeMicro,
    /// Classic PCAP (little-endian, nanoseconds)
    LegacyLeNano,
    /// Classic PCAP (big-endian, nanoseconds)
    LegacyBeNano,
    /// PCAPNG format
    PcapNg,
}

impl PcapFormat {
    /// Detect PCAP format from magic bytes.
    pub fn detect(data: &[u8]) -> Result<Self, Error> {
        if data.len() < 4 {
            return Err(Error::Pcap(PcapError::InvalidFormat {
                reason: "Data too small for PCAP magic".into(),
            }));
        }

        let magic = u32::from_ne_bytes([data[0], data[1], data[2], data[3]]);

        match magic {
            0xa1b2c3d4 => Ok(PcapFormat::LegacyLeMicro),
            0xd4c3b2a1 => Ok(PcapFormat::LegacyBeMicro),
            0xa1b23c4d => Ok(PcapFormat::LegacyLeNano),
            0x4d3cb2a1 => Ok(PcapFormat::LegacyBeNano),
            0x0a0d0d0a => Ok(PcapFormat::PcapNg),
            _ => Err(Error::Pcap(PcapError::InvalidFormat {
                reason: format!("Unknown PCAP magic: 0x{magic:08x}"),
            })),
        }
    }

    /// Whether this is a PCAPNG format.
    pub fn is_pcapng(&self) -> bool {
        matches!(self, PcapFormat::PcapNg)
    }

    /// Whether this is a legacy PCAP format.
    pub fn is_legacy(&self) -> bool {
        !self.is_pcapng()
    }

    /// Whether this format uses little-endian byte order.
    ///
    /// This is relevant for parsing header fields like link_type.
    /// For PCAPNG, the section header defines endianness, but we assume
    /// little-endian as it's the most common case.
    pub fn is_little_endian(&self) -> bool {
        match self {
            PcapFormat::LegacyLeMicro | PcapFormat::LegacyLeNano => true,
            PcapFormat::LegacyBeMicro | PcapFormat::LegacyBeNano => false,
            PcapFormat::PcapNg => true, // Most PCAPNG files are little-endian
        }
    }
}

/// Generic PCAP/PCAPNG reader over any Read source.
///
/// This is the unified PCAP parser that uses `pcap_parser` for all formats.
/// It provides a simple `next_packet()` interface that works with any byte source.
pub struct GenericPcapReader<R: Read> {
    inner: ReaderInner<R>,
    frame_number: u64,
    link_type: u32,
}

/// Inner reader using enum dispatch for format-specific handling.
enum ReaderInner<R: Read> {
    Legacy(LegacyPcapReader<BufReader<R>>),
    Ng(PcapNGReader<BufReader<R>>),
}

impl<R: Read> GenericPcapReader<R> {
    /// Create a reader with known format.
    ///
    /// This is the primary constructor. Use `PcapFormat::detect()` to determine
    /// the format from magic bytes before calling this.
    ///
    /// Note: For legacy PCAP, the header is read during construction to extract
    /// the link_type. For PCAPNG, link_type is initially 1 and gets updated when
    /// the Interface Description Block is read during packet processing.
    pub fn with_format(source: R, format: PcapFormat) -> Result<Self, Error> {
        let buf_reader = BufReader::with_capacity(BUFFER_SIZE, source);

        let (inner, link_type) = if format.is_pcapng() {
            let reader = PcapNGReader::new(BUFFER_SIZE, buf_reader).map_err(|e| {
                Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("Failed to parse PCAPNG: {e}"),
                })
            })?;
            // PCAPNG link_type comes from IDB block, read during packet processing
            (ReaderInner::Ng(reader), 1u32)
        } else {
            let mut reader = LegacyPcapReader::new(BUFFER_SIZE, buf_reader).map_err(|e| {
                Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("Failed to parse legacy PCAP: {e}"),
                })
            })?;
            // Read ahead to get link_type from the legacy PCAP header
            let link_type = Self::read_legacy_header_link_type(&mut reader)?;
            (ReaderInner::Legacy(reader), link_type)
        };

        Ok(GenericPcapReader {
            inner,
            frame_number: 0,
            link_type,
        })
    }

    /// Read the legacy PCAP header to extract link_type.
    ///
    /// This consumes the header block but leaves the reader positioned
    /// at the first packet.
    fn read_legacy_header_link_type<S: Read>(
        reader: &mut LegacyPcapReader<S>,
    ) -> Result<u32, Error> {
        use pcap_parser::PcapError as PcapParserError;

        loop {
            match reader.next() {
                Ok((offset, block)) => match block {
                    PcapBlockOwned::LegacyHeader(header) => {
                        let link_type = header.network.0 as u32;
                        reader.consume(offset);
                        return Ok(link_type);
                    }
                    PcapBlockOwned::Legacy(_) => {
                        // Shouldn't happen - header should come first
                        // Don't consume, let normal reading handle it
                        return Ok(1); // Default to Ethernet
                    }
                    _ => {
                        reader.consume(offset);
                        continue;
                    }
                },
                Err(PcapParserError::Eof) => {
                    return Ok(1); // Empty file, default to Ethernet
                }
                Err(PcapParserError::Incomplete(_)) => {
                    reader.refill().map_err(|e| {
                        Error::Pcap(PcapError::InvalidFormat {
                            reason: format!("Failed to read PCAP header: {e}"),
                        })
                    })?;
                    continue;
                }
                Err(e) => {
                    return Err(Error::Pcap(PcapError::InvalidFormat {
                        reason: format!("Failed to parse PCAP header: {e}"),
                    }));
                }
            }
        }
    }

    /// Read the next packet.
    ///
    /// Returns `Ok(None)` at end of file.
    pub fn next_packet(&mut self) -> Result<Option<RawPacket>, Error> {
        match &mut self.inner {
            ReaderInner::Legacy(reader) => {
                read_legacy_packet(reader, &mut self.frame_number, &mut self.link_type)
            }
            ReaderInner::Ng(reader) => {
                read_pcapng_packet(reader, &mut self.frame_number, &mut self.link_type)
            }
        }
    }

    /// Get the link type (e.g., 1 = Ethernet).
    pub fn link_type(&self) -> u32 {
        self.link_type
    }

    /// Get the current frame count.
    pub fn frame_count(&self) -> u64 {
        self.frame_number
    }

    /// Process packets with zero-copy borrowed data.
    ///
    /// The callback receives borrowed packet data. The borrow is valid
    /// only during the callback - data must be processed before returning.
    /// This eliminates the `Bytes::copy_from_slice()` overhead of `next_packet()`.
    ///
    /// Returns the number of packets processed.
    #[inline]
    pub fn process_packets<F>(&mut self, max: usize, f: F) -> Result<usize, Error>
    where
        F: FnMut(PacketRef<'_>) -> Result<(), Error>,
    {
        match &mut self.inner {
            ReaderInner::Legacy(reader) => {
                process_legacy_packets(reader, max, &mut self.frame_number, &mut self.link_type, f)
            }
            ReaderInner::Ng(reader) => {
                process_pcapng_packets(reader, max, &mut self.frame_number, &mut self.link_type, f)
            }
        }
    }
}

/// Read next packet from a legacy PCAP reader.
fn read_legacy_packet<S: Read>(
    reader: &mut LegacyPcapReader<S>,
    frame_number: &mut u64,
    link_type: &mut u32,
) -> Result<Option<RawPacket>, Error> {
    use pcap_parser::PcapError as PcapParserError;

    loop {
        match reader.next() {
            Ok((offset, block)) => match block {
                PcapBlockOwned::Legacy(packet) => {
                    *frame_number += 1;

                    let timestamp_us = (packet.ts_sec as i64) * 1_000_000 + (packet.ts_usec as i64);

                    let raw = RawPacket {
                        frame_number: *frame_number,
                        timestamp_us,
                        captured_length: packet.caplen,
                        original_length: packet.origlen,
                        link_type: *link_type as u16,
                        data: Bytes::copy_from_slice(packet.data),
                    };

                    reader.consume(offset);
                    return Ok(Some(raw));
                }
                PcapBlockOwned::LegacyHeader(header) => {
                    *link_type = header.network.0 as u32;
                    reader.consume(offset);
                    continue;
                }
                _ => {
                    reader.consume(offset);
                    continue;
                }
            },
            Err(PcapParserError::Eof) => return Ok(None),
            Err(PcapParserError::Incomplete(_)) => {
                reader.refill().map_err(|e| {
                    Error::Pcap(PcapError::InvalidFormat {
                        reason: format!("Legacy PCAP refill error: {e}"),
                    })
                })?;
                continue;
            }
            Err(e) => {
                return Err(Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("Legacy PCAP parse error: {e}"),
                }));
            }
        }
    }
}

/// Read next packet from a PCAPNG reader.
fn read_pcapng_packet<S: Read>(
    reader: &mut PcapNGReader<S>,
    frame_number: &mut u64,
    link_type: &mut u32,
) -> Result<Option<RawPacket>, Error> {
    use pcap_parser::PcapError as PcapParserError;

    loop {
        match reader.next() {
            Ok((offset, block)) => match block {
                PcapBlockOwned::NG(ng_block) => {
                    use pcap_parser::pcapng::*;

                    match ng_block {
                        Block::InterfaceDescription(idb) => {
                            *link_type = idb.linktype.0 as u32;
                            reader.consume(offset);
                            continue;
                        }
                        Block::EnhancedPacket(epb) => {
                            *frame_number += 1;

                            let timestamp_us = ((epb.ts_high as i64) << 32) | (epb.ts_low as i64);

                            let packet = RawPacket {
                                frame_number: *frame_number,
                                timestamp_us,
                                captured_length: epb.caplen,
                                original_length: epb.origlen,
                                link_type: *link_type as u16,
                                data: Bytes::copy_from_slice(epb.data),
                            };

                            reader.consume(offset);
                            return Ok(Some(packet));
                        }
                        Block::SimplePacket(spb) => {
                            *frame_number += 1;

                            let packet = RawPacket {
                                frame_number: *frame_number,
                                timestamp_us: 0,
                                captured_length: spb.data.len() as u32,
                                original_length: spb.origlen,
                                link_type: *link_type as u16,
                                data: Bytes::copy_from_slice(spb.data),
                            };

                            reader.consume(offset);
                            return Ok(Some(packet));
                        }
                        _ => {
                            reader.consume(offset);
                            continue;
                        }
                    }
                }
                _ => {
                    reader.consume(offset);
                    continue;
                }
            },
            Err(PcapParserError::Eof) => return Ok(None),
            Err(PcapParserError::Incomplete(_)) => {
                reader.refill().map_err(|e| {
                    Error::Pcap(PcapError::InvalidFormat {
                        reason: format!("PCAPNG refill error: {e}"),
                    })
                })?;
                continue;
            }
            Err(e) => {
                return Err(Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("PCAPNG parse error: {e}"),
                }));
            }
        }
    }
}

/// Process packets from a legacy PCAP reader with zero-copy.
///
/// This is the zero-copy version of `read_legacy_packet`. The callback receives
/// borrowed packet data that must be processed before returning.
fn process_legacy_packets<S: Read, F>(
    reader: &mut LegacyPcapReader<S>,
    max: usize,
    frame_number: &mut u64,
    link_type: &mut u32,
    mut f: F,
) -> Result<usize, Error>
where
    F: FnMut(PacketRef<'_>) -> Result<(), Error>,
{
    use pcap_parser::PcapError as PcapParserError;

    let mut count = 0;
    while count < max {
        match reader.next() {
            Ok((offset, block)) => {
                match block {
                    PcapBlockOwned::Legacy(packet) => {
                        *frame_number += 1;

                        let timestamp_us =
                            (packet.ts_sec as i64) * 1_000_000 + (packet.ts_usec as i64);

                        // Create borrowed packet reference - no copy!
                        let packet_ref = PacketRef {
                            frame_number: *frame_number,
                            timestamp_us,
                            captured_len: packet.caplen,
                            original_len: packet.origlen,
                            link_type: *link_type as u16,
                            data: packet.data, // Borrowed from pcap_parser buffer
                        };

                        // Call the callback with borrowed data
                        f(packet_ref)?;

                        // Only consume after callback completes
                        reader.consume(offset);
                        count += 1;
                    }
                    PcapBlockOwned::LegacyHeader(header) => {
                        *link_type = header.network.0 as u32;
                        reader.consume(offset);
                        continue;
                    }
                    _ => {
                        reader.consume(offset);
                        continue;
                    }
                }
            }
            Err(PcapParserError::Eof) => break,
            Err(PcapParserError::Incomplete(_)) => {
                reader.refill().map_err(|e| {
                    Error::Pcap(PcapError::InvalidFormat {
                        reason: format!("Legacy PCAP refill error: {e}"),
                    })
                })?;
                continue;
            }
            Err(e) => {
                return Err(Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("Legacy PCAP parse error: {e}"),
                }));
            }
        }
    }
    Ok(count)
}

/// Process packets from a PCAPNG reader with zero-copy.
///
/// This is the zero-copy version of `read_pcapng_packet`. The callback receives
/// borrowed packet data that must be processed before returning.
fn process_pcapng_packets<S: Read, F>(
    reader: &mut PcapNGReader<S>,
    max: usize,
    frame_number: &mut u64,
    link_type: &mut u32,
    mut f: F,
) -> Result<usize, Error>
where
    F: FnMut(PacketRef<'_>) -> Result<(), Error>,
{
    use pcap_parser::PcapError as PcapParserError;

    let mut count = 0;
    while count < max {
        match reader.next() {
            Ok((offset, block)) => {
                match block {
                    PcapBlockOwned::NG(ng_block) => {
                        use pcap_parser::pcapng::*;

                        match ng_block {
                            Block::InterfaceDescription(idb) => {
                                *link_type = idb.linktype.0 as u32;
                                reader.consume(offset);
                                continue;
                            }
                            Block::EnhancedPacket(epb) => {
                                *frame_number += 1;

                                let timestamp_us =
                                    ((epb.ts_high as i64) << 32) | (epb.ts_low as i64);

                                // Create borrowed packet reference - no copy!
                                let packet_ref = PacketRef {
                                    frame_number: *frame_number,
                                    timestamp_us,
                                    captured_len: epb.caplen,
                                    original_len: epb.origlen,
                                    link_type: *link_type as u16,
                                    data: epb.data, // Borrowed from pcap_parser buffer
                                };

                                // Call the callback with borrowed data
                                f(packet_ref)?;

                                // Only consume after callback completes
                                reader.consume(offset);
                                count += 1;
                            }
                            Block::SimplePacket(spb) => {
                                *frame_number += 1;

                                // Create borrowed packet reference - no copy!
                                let packet_ref = PacketRef {
                                    frame_number: *frame_number,
                                    timestamp_us: 0,
                                    captured_len: spb.data.len() as u32,
                                    original_len: spb.origlen,
                                    link_type: *link_type as u16,
                                    data: spb.data, // Borrowed from pcap_parser buffer
                                };

                                // Call the callback with borrowed data
                                f(packet_ref)?;

                                // Only consume after callback completes
                                reader.consume(offset);
                                count += 1;
                            }
                            _ => {
                                reader.consume(offset);
                                continue;
                            }
                        }
                    }
                    _ => {
                        reader.consume(offset);
                        continue;
                    }
                }
            }
            Err(PcapParserError::Eof) => break,
            Err(PcapParserError::Incomplete(_)) => {
                reader.refill().map_err(|e| {
                    Error::Pcap(PcapError::InvalidFormat {
                        reason: format!("PCAPNG refill error: {e}"),
                    })
                })?;
                continue;
            }
            Err(e) => {
                return Err(Error::Pcap(PcapError::InvalidFormat {
                    reason: format!("PCAPNG parse error: {e}"),
                }));
            }
        }
    }
    Ok(count)
}

// GenericPcapReader is Send when R is Send
unsafe impl<R: Read + Send> Send for GenericPcapReader<R> {}

// Required for async compatibility
impl<R: Read> Unpin for GenericPcapReader<R> {}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn test_pcap_format_detect() {
        // Magic bytes are stored as-written by the capturing system.
        // On a little-endian machine, 0xa1b2c3d4 is stored as [0xd4, 0xc3, 0xb2, 0xa1].
        // When read back on a little-endian machine with from_ne_bytes(), we get 0xa1b2c3d4.

        // Little-endian microseconds (stored as [0xd4, 0xc3, 0xb2, 0xa1])
        let le_micro = [0xd4, 0xc3, 0xb2, 0xa1];
        assert_eq!(
            PcapFormat::detect(&le_micro).unwrap(),
            PcapFormat::LegacyLeMicro
        );

        // Big-endian microseconds (stored as [0xa1, 0xb2, 0xc3, 0xd4])
        let be_micro = [0xa1, 0xb2, 0xc3, 0xd4];
        assert_eq!(
            PcapFormat::detect(&be_micro).unwrap(),
            PcapFormat::LegacyBeMicro
        );

        // PCAPNG
        let pcapng = [0x0a, 0x0d, 0x0d, 0x0a];
        assert_eq!(PcapFormat::detect(&pcapng).unwrap(), PcapFormat::PcapNg);

        // Unknown magic
        let unknown = [0xDE, 0xAD, 0xBE, 0xEF];
        assert!(PcapFormat::detect(&unknown).is_err());
    }

    #[test]
    fn test_pcap_format_properties() {
        assert!(PcapFormat::LegacyLeMicro.is_legacy());
        assert!(!PcapFormat::LegacyLeMicro.is_pcapng());

        assert!(PcapFormat::PcapNg.is_pcapng());
        assert!(!PcapFormat::PcapNg.is_legacy());
    }

    #[test]
    fn test_pcap_format_endianness() {
        // Little-endian formats
        assert!(PcapFormat::LegacyLeMicro.is_little_endian());
        assert!(PcapFormat::LegacyLeNano.is_little_endian());
        assert!(PcapFormat::PcapNg.is_little_endian()); // Most PCAPNG files are LE

        // Big-endian formats
        assert!(!PcapFormat::LegacyBeMicro.is_little_endian());
        assert!(!PcapFormat::LegacyBeNano.is_little_endian());
    }

    /// Create a minimal valid PCAP file for testing.
    fn create_minimal_pcap() -> Vec<u8> {
        let mut data = Vec::new();

        // PCAP global header
        data.extend_from_slice(&[0xd4, 0xc3, 0xb2, 0xa1]); // Magic (little endian)
        data.extend_from_slice(&[0x02, 0x00]); // Version major (2)
        data.extend_from_slice(&[0x04, 0x00]); // Version minor (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); // Thiszone
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); // Sigfigs
        data.extend_from_slice(&[0xff, 0xff, 0x00, 0x00]); // Snaplen (65535)
        data.extend_from_slice(&[0x01, 0x00, 0x00, 0x00]); // Network (Ethernet)

        // One packet header + minimal Ethernet frame
        let packet_data = [
            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // Dst MAC
            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, // Src MAC
            0x08, 0x00, // EtherType (IPv4)
        ];

        let ts_sec: u32 = 1000000000;
        let ts_usec: u32 = 500000;
        let caplen: u32 = packet_data.len() as u32;
        let origlen: u32 = packet_data.len() as u32;

        data.extend_from_slice(&ts_sec.to_le_bytes());
        data.extend_from_slice(&ts_usec.to_le_bytes());
        data.extend_from_slice(&caplen.to_le_bytes());
        data.extend_from_slice(&origlen.to_le_bytes());
        data.extend_from_slice(&packet_data);

        data
    }

    #[test]
    fn test_generic_reader_from_memory() {
        let pcap_data = create_minimal_pcap();

        // Detect format from magic bytes
        let format = PcapFormat::detect(&pcap_data).expect("Failed to detect format");

        let cursor = Cursor::new(pcap_data);
        let mut reader =
            GenericPcapReader::with_format(cursor, format).expect("Failed to create reader");

        // Read first packet
        let packet = reader.next_packet().expect("Read error");
        assert!(packet.is_some());

        let pkt = packet.unwrap();
        assert_eq!(pkt.frame_number, 1);
        assert_eq!(pkt.captured_length, 14);
        assert_eq!(pkt.original_length, 14);
        assert_eq!(pkt.link_type, 1); // Ethernet
        assert_eq!(pkt.timestamp_us, 1000000000_500000i64);
        assert_eq!(pkt.data.len(), 14);

        // No more packets
        let packet2 = reader.next_packet().expect("Read error");
        assert!(packet2.is_none());
    }

    #[test]
    fn test_generic_reader_link_type() {
        let pcap_data = create_minimal_pcap();
        let format = PcapFormat::detect(&pcap_data).expect("Failed to detect format");
        let cursor = Cursor::new(pcap_data);

        let mut reader =
            GenericPcapReader::with_format(cursor, format).expect("Failed to create reader");

        // Link type is set after reading header block
        reader.next_packet().ok();
        assert_eq!(reader.link_type(), 1); // Ethernet
    }

    #[test]
    fn test_generic_reader_frame_count() {
        let pcap_data = create_minimal_pcap();
        let format = PcapFormat::detect(&pcap_data).expect("Failed to detect format");
        let cursor = Cursor::new(pcap_data);

        let mut reader =
            GenericPcapReader::with_format(cursor, format).expect("Failed to create reader");
        assert_eq!(reader.frame_count(), 0);

        reader.next_packet().ok();
        assert_eq!(reader.frame_count(), 1);
    }
}