pcapsql_core/io/

mmap.rs

//! Memory-mapped packet source for efficient packet access.
//!
//! This module is only available when the `mmap` feature is enabled.
//!
//! Uses `memmap2` crate for platform-independent memory mapping.
//! Provides efficient access to packet data for large files.
//!
//! ## Supported Formats
//!
//! - Classic PCAP (little/big endian, micro/nanosecond timestamps)
//! - PCAPNG
//! - All above formats with compression (gzip, zstd, lz4, bzip2, xz)
//!
//! ## Design
//!
//! This module uses the unified layer stack:
//! ```text
//! MmapPacketReader
//!     └── GenericPcapReader<DecompressReader<Cursor<MmapSlice>>>
//!             └── DecompressReader handles compression
//!                     └── Cursor<MmapSlice> provides Read over mmap
//! ```
//!
//! Benchmarks show this approach (using pcap_parser) is 30% faster than
//! custom byte parsing while being more maintainable.

use std::fs::File;
use std::io::Cursor;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use memmap2::Mmap;

use crate::error::{Error, PcapError};

use super::decompress::{Compression, DecompressReader, MmapSlice};
use super::pcap_stream::{GenericPcapReader, PcapFormat};
use super::{
    PacketPosition, PacketRange, PacketReader, PacketRef, PacketSource, PacketSourceMetadata,
};

/// Memory-mapped packet source.
///
/// Maps the entire file into virtual memory, allowing the OS to handle
/// caching and paging. Supports both uncompressed and compressed files.
#[derive(Clone)]
pub struct MmapPacketSource {
    /// Path to the file (for error messages)
    path: PathBuf,
    /// Memory-mapped region (shared via Arc)
    mmap: Arc<Mmap>,
    /// Cached metadata
    metadata: PacketSourceMetadata,
    /// Detected compression format
    compression: Compression,
    /// PCAP format (after decompression)
    pcap_format: PcapFormat,
}

impl MmapPacketSource {
    /// Open a PCAP or PCAPNG file with memory mapping.
    ///
    /// Automatically detects and handles compression.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let path = path.as_ref().to_path_buf();
        let file = File::open(&path).map_err(Error::Io)?;

        // Create memory mapping
        let mmap = unsafe { Mmap::map(&file).map_err(Error::Io)? };

        // Detect compression
        let compression = Compression::detect(&mmap);

        // Detect PCAP format (may need to decompress first bytes)
        let (pcap_format, link_type) = if compression.is_compressed() {
            Self::detect_format_compressed(&mmap, compression)?
        } else {
            Self::detect_format_uncompressed(&mmap)?
        };

        let metadata = PacketSourceMetadata {
            link_type,
            snaplen: 65535,
            size_bytes: Some(mmap.len() as u64),
            packet_count: None,
            seekable: !compression.is_compressed(), // Can only seek in uncompressed
        };

        Ok(Self {
            path,
            mmap: Arc::new(mmap),
            metadata,
            compression,
            pcap_format,
        })
    }

    /// Detect format from uncompressed data.
    fn detect_format_uncompressed(data: &[u8]) -> Result<(PcapFormat, u32), Error> {
        if data.len() < 24 {
            return Err(Error::Pcap(PcapError::InvalidFormat {
                reason: "File too small for PCAP header".into(),
            }));
        }

        let format = PcapFormat::detect(data)?;
        let link_type = if format.is_pcapng() {
            1 // Will be updated from interface description block
        } else {
            Self::link_type_from_header(data, &format)
        };

        Ok((format, link_type))
    }

    /// Extract link type from legacy PCAP header bytes.
    fn link_type_from_header(data: &[u8], format: &PcapFormat) -> u32 {
        if data.len() < 24 {
            return 1; // Default to Ethernet
        }
        // Link type is at offset 20 in PCAP global header
        let byte_swap = matches!(format, PcapFormat::LegacyBeMicro | PcapFormat::LegacyBeNano);
        if byte_swap {
            u32::from_be_bytes([data[20], data[21], data[22], data[23]])
        } else {
            u32::from_le_bytes([data[20], data[21], data[22], data[23]])
        }
    }

    /// Detect format from compressed data (decompress first bytes).
    fn detect_format_compressed(
        data: &[u8],
        compression: Compression,
    ) -> Result<(PcapFormat, u32), Error> {
        use std::io::Read;

        // Decompress enough bytes to detect the format and link type
        let mut decoder: Box<dyn Read> = match compression {
            Compression::None => unreachable!(),
            Compression::Gzip => {
                let cursor = Cursor::new(data);
                let gz = flate2::read::GzDecoder::new(cursor);
                Box::new(gz) as Box<dyn Read>
            }
            #[cfg(feature = "compress-zstd")]
            Compression::Zstd => {
                let cursor = Cursor::new(data);
                let zstd = zstd::Decoder::new(cursor)?;
                Box::new(zstd) as Box<dyn Read>
            }
            #[cfg(feature = "compress-lz4")]
            Compression::Lz4 => {
                let cursor = Cursor::new(data);
                let lz4 = lz4_flex::frame::FrameDecoder::new(cursor);
                Box::new(lz4) as Box<dyn Read>
            }
            #[cfg(feature = "compress-bzip2")]
            Compression::Bzip2 => {
                let cursor = Cursor::new(data);
                let bz2 = bzip2::read::BzDecoder::new(cursor);
                Box::new(bz2) as Box<dyn Read>
            }
            #[cfg(feature = "compress-xz")]
            Compression::Xz => {
                let cursor = Cursor::new(data);
                let xz = xz2::read::XzDecoder::new(cursor);
                Box::new(xz) as Box<dyn Read>
            }
        };

        // Read enough bytes to detect format and link type
        let mut header = [0u8; 24];
        decoder.read_exact(&mut header).map_err(|e| {
            Error::Pcap(PcapError::InvalidFormat {
                reason: format!("Failed to read compressed header: {e}"),
            })
        })?;

        let format = PcapFormat::detect(&header)?;
        let link_type = if format.is_pcapng() {
            1 // Will be updated during reading
        } else {
            Self::link_type_from_header(&header, &format)
        };

        Ok((format, link_type))
    }

    /// Get the path to the file.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Get the detected compression format.
    pub fn compression(&self) -> Compression {
        self.compression
    }

    /// Get the PCAP format.
    pub fn pcap_format(&self) -> PcapFormat {
        self.pcap_format
    }

    /// Check if this is a PCAPNG file.
    pub fn is_pcapng(&self) -> bool {
        self.pcap_format.is_pcapng()
    }

    /// Check if this file is compressed.
    pub fn is_compressed(&self) -> bool {
        self.compression.is_compressed()
    }

    /// Get the link type.
    pub fn link_type(&self) -> u32 {
        self.metadata.link_type
    }
}

impl PacketSource for MmapPacketSource {
    type Reader = MmapPacketReader;

    fn metadata(&self) -> &PacketSourceMetadata {
        &self.metadata
    }

    fn reader(&self, range: Option<&PacketRange>) -> Result<Self::Reader, Error> {
        MmapPacketReader::new(
            self.mmap.clone(),
            self.compression,
            self.pcap_format,
            self.metadata.link_type,
            range.cloned(),
        )
    }

    fn partitions(&self, _max_partitions: usize) -> Result<Vec<PacketRange>, Error> {
        // For now, return single partition
        // Future optimization: scan for packet boundaries at byte offsets
        Ok(vec![PacketRange::whole()])
    }
}

impl std::fmt::Debug for MmapPacketSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MmapPacketSource")
            .field("path", &self.path)
            .field("size_bytes", &self.metadata.size_bytes)
            .field("link_type", &self.metadata.link_type)
            .field("compression", &self.compression)
            .field("pcap_format", &self.pcap_format)
            .finish()
    }
}

/// Memory-mapped packet reader.
///
/// Uses the unified layer stack:
/// - `GenericPcapReader` for PCAP parsing
/// - `DecompressReader` for transparent decompression
/// - `Cursor<MmapSlice>` for reading from mmap
pub struct MmapPacketReader {
    /// The unified PCAP reader
    inner: GenericPcapReader<DecompressReader<Cursor<MmapSlice>>>,
    /// Link type (may be updated from PCAPNG interface description)
    link_type: u32,
    /// Current byte offset (for position tracking)
    byte_offset: u64,
    /// Optional range restriction
    range: Option<PacketRange>,
}

impl MmapPacketReader {
    fn new(
        mmap: Arc<Mmap>,
        compression: Compression,
        pcap_format: PcapFormat,
        link_type: u32,
        range: Option<PacketRange>,
    ) -> Result<Self, Error> {
        // Create the layer stack: Cursor -> DecompressReader -> GenericPcapReader
        let slice = MmapSlice::new(mmap);
        let cursor = Cursor::new(slice);
        let decompress = DecompressReader::new(cursor, compression).map_err(|e| {
            Error::Pcap(PcapError::InvalidFormat {
                reason: format!("Failed to create decompressor: {e}"),
            })
        })?;
        let inner = GenericPcapReader::with_format(decompress, pcap_format)?;

        Ok(Self {
            inner,
            link_type,
            byte_offset: 0,
            range,
        })
    }

    /// Check if we've passed the end of our range.
    #[inline]
    fn past_range_end(&self) -> bool {
        if let Some(ref range) = self.range {
            if let Some(ref end) = range.end {
                return self.inner.frame_count() >= end.frame_number;
            }
        }
        false
    }
}

impl PacketReader for MmapPacketReader {
    fn process_packets<F>(&mut self, max: usize, mut f: F) -> Result<usize, Error>
    where
        F: FnMut(PacketRef<'_>) -> Result<(), Error>,
    {
        if self.past_range_end() {
            return Ok(0);
        }

        // Calculate how many packets we can process before hitting range end
        let effective_max = if let Some(ref range) = self.range {
            if let Some(ref end) = range.end {
                let remaining = end.frame_number.saturating_sub(self.inner.frame_count());
                max.min(remaining as usize)
            } else {
                max
            }
        } else {
            max
        };

        let count = self.inner.process_packets(effective_max, &mut f)?;

        // Update link type from reader (may have been updated from PCAPNG IDB)
        self.link_type = self.inner.link_type();

        Ok(count)
    }

    fn position(&self) -> PacketPosition {
        PacketPosition {
            byte_offset: self.byte_offset,
            frame_number: self.inner.frame_count(),
        }
    }

    fn link_type(&self) -> u32 {
        self.link_type
    }
}

// Implement Unpin for async compatibility
impl Unpin for MmapPacketReader {}

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

    fn test_pcap_path(name: &str) -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("testdata")
            .join("corpus")
            .join(name)
    }

    #[test]
    fn test_mmap_source_opens() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return; // Skip if test file doesn't exist
        }

        let source = MmapPacketSource::open(&path);
        assert!(source.is_ok(), "Failed to open: {:?}", source.err());
    }

    #[test]
    fn test_mmap_reader_reads_packets() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        let mut found_packet = false;
        reader
            .process_packets(1, |packet| {
                assert_eq!(packet.frame_number, 1);
                assert!(!packet.data.is_empty());
                found_packet = true;
                Ok(())
            })
            .unwrap();
        assert!(found_packet);
    }

    #[test]
    fn test_mmap_reader_counts_match_file() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        let mut count = 0;
        loop {
            let processed = reader
                .process_packets(100, |_| {
                    count += 1;
                    Ok(())
                })
                .unwrap();
            if processed == 0 {
                break;
            }
        }

        // dns.cap has 38 packets
        assert!(count > 0, "Should have read some packets");
    }

    #[test]
    fn test_mmap_position_tracking() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        let pos1 = reader.position();
        assert_eq!(pos1.frame_number, 0);

        reader.process_packets(1, |_| Ok(())).unwrap();
        let pos2 = reader.position();
        assert_eq!(pos2.frame_number, 1);
    }

    #[test]
    fn test_mmap_link_type() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }
        let source = MmapPacketSource::open(&path).unwrap();
        // dns.cap should have Ethernet link type (1)
        assert_eq!(source.metadata().link_type, 1);
    }

    #[test]
    fn test_mmap_netlink_link_type() {
        let path = test_pcap_path("nlmon-big.pcap");
        if !path.exists() {
            eprintln!("Skipping test - nlmon-big.pcap not found");
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        // nlmon-big.pcap should have NETLINK link type (253)
        assert_eq!(
            source.link_type(),
            253,
            "Expected LINKTYPE_NETLINK (253), got {}",
            source.link_type()
        );
    }

    #[test]
    fn test_mmap_debug_format() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let debug_str = format!("{:?}", source);
        assert!(debug_str.contains("MmapPacketSource"));
        assert!(debug_str.contains("dns.cap"));
    }

    #[test]
    fn test_mmap_packet_data_integrity() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        // Read first packet and validate Ethernet header structure
        reader
            .process_packets(1, |packet| {
                // Minimum Ethernet frame is 14 bytes (MAC dst + MAC src + EtherType)
                assert!(packet.data.len() >= 14, "Packet too small for Ethernet");

                // captured_len should match data length
                assert_eq!(packet.captured_len as usize, packet.data.len());

                // original_len should be >= captured_len
                assert!(packet.original_len >= packet.captured_len);

                // Frame number should start at 1
                assert_eq!(packet.frame_number, 1);

                // Timestamp should be positive (after Unix epoch)
                assert!(packet.timestamp_us > 0);
                Ok(())
            })
            .unwrap();
    }

    #[test]
    fn test_mmap_all_packets_valid() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        let mut prev_frame = 0u64;
        let mut count = 0;

        loop {
            let processed = reader
                .process_packets(100, |packet| {
                    // Frame numbers should be sequential
                    assert_eq!(packet.frame_number, prev_frame + 1);
                    prev_frame = packet.frame_number;

                    // Basic sanity checks
                    assert!(packet.data.len() <= 65535, "Packet exceeds max size");
                    assert_eq!(packet.captured_len as usize, packet.data.len());

                    count += 1;
                    Ok(())
                })
                .unwrap();
            if processed == 0 {
                break;
            }
        }

        assert!(count > 0, "Should have read packets");
    }

    #[test]
    fn test_mmap_timestamp_reasonableness() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let mut reader = source.reader(None).unwrap();

        let mut prev_timestamp = 0i64;

        loop {
            let processed = reader
                .process_packets(100, |packet| {
                    // Timestamps should be non-negative
                    assert!(packet.timestamp_us >= 0);

                    // Timestamps should generally not go backwards (within reason)
                    // Allow small backward jumps for clock drift
                    if prev_timestamp > 0 {
                        let diff = packet.timestamp_us - prev_timestamp;
                        assert!(
                            diff >= -1_000_000, // Allow up to 1 second backward
                            "Timestamp went backwards by {} us",
                            diff.abs()
                        );
                    }
                    prev_timestamp = packet.timestamp_us;
                    Ok(())
                })
                .unwrap();
            if processed == 0 {
                break;
            }
        }
    }

    #[test]
    fn test_mmap_range_reading() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();

        // Create a range ending at frame 10
        let range = PacketRange {
            start: PacketPosition {
                byte_offset: 0,
                frame_number: 0,
            },
            end: Some(PacketPosition {
                byte_offset: 0,
                frame_number: 10,
            }),
        };

        let mut reader = source.reader(Some(&range)).unwrap();

        let mut count = 0;
        loop {
            let processed = reader
                .process_packets(100, |_| {
                    count += 1;
                    Ok(())
                })
                .unwrap();
            if processed == 0 || count > 100 {
                break; // EOF or safety limit
            }
        }

        // Should have read frames 1-10 (10 frames, since range.end is exclusive)
        assert!(count <= 10, "Should respect range end limit, got {}", count);
    }

    #[test]
    fn test_mmap_metadata() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let meta = source.metadata();

        assert_eq!(meta.link_type, 1); // Ethernet
        assert!(meta.size_bytes.is_some());
        assert!(meta.size_bytes.unwrap() > 0);
    }

    #[test]
    fn test_mmap_clone() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source1 = MmapPacketSource::open(&path).unwrap();
        let source2 = source1.clone();

        // Both sources should work independently
        let mut reader1 = source1.reader(None).unwrap();
        let mut reader2 = source2.reader(None).unwrap();

        let mut frame1 = 0u64;
        let mut len1 = 0usize;
        let mut frame2 = 0u64;
        let mut len2 = 0usize;

        reader1
            .process_packets(1, |p| {
                frame1 = p.frame_number;
                len1 = p.data.len();
                Ok(())
            })
            .unwrap();
        reader2
            .process_packets(1, |p| {
                frame2 = p.frame_number;
                len2 = p.data.len();
                Ok(())
            })
            .unwrap();

        // Should read identical first packets
        assert_eq!(frame1, frame2);
        assert_eq!(len1, len2);
    }

    #[test]
    fn test_mmap_partitions() {
        let path = test_pcap_path("dns.cap");
        if !path.exists() {
            return;
        }

        let source = MmapPacketSource::open(&path).unwrap();
        let partitions = source.partitions(4).unwrap();

        // Currently returns single partition
        assert_eq!(partitions.len(), 1);
    }
}