diff --git a/Cargo.lock b/Cargo.lock index 6c6bb36..dd40f54 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,351 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "cc" +version = "1.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "find-msvc-tools" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + [[package]] name = "scpcap" version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "tempfile", + "zstd", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/Cargo.toml b/Cargo.toml index f7cd8f5..5014212 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,3 +4,9 @@ version = "0.1.0" edition = "2024" [dependencies] +anyhow = "1.0.104" +clap = { version = "4.6.6", features = ["derive"] } +zstd = "0.13.3" + +[dev-dependencies] +tempfile = "3.27.0" diff --git a/src/channel.rs b/src/channel.rs new file mode 100644 index 0000000..f1ddbba --- /dev/null +++ b/src/channel.rs @@ -0,0 +1,52 @@ +use crate::protocol::{Message, MessageSink, MessageSource}; +use anyhow::Result; +use std::sync::mpsc::{Receiver, Sender}; + +pub struct ChannelSink(pub Sender); + +impl MessageSink for ChannelSink { + fn send_data(&mut self, payload: &[u8]) -> Result<()> { + self.0.send(Message::Data(payload.to_vec()))?; + Ok(()) + } + + fn send_finalize(&mut self) -> Result<()> { + self.0.send(Message::Finalize)?; + Ok(()) + } + + fn send_error(&mut self, msg: &str) -> Result<()> { + // No dedicated error variant needed in-process: the sender thread's + // Result propagates directly to the caller of run_send. + anyhow::bail!("{msg}"); + } +} + +pub struct ChannelSource(pub Receiver); + +impl MessageSource for ChannelSource { + fn recv(&mut self) -> Result> { + match self.0.recv() { + Ok(msg) => Ok(Some(msg)), + Err(_) => Ok(None), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn channel_round_trip() { + let (tx, rx) = std::sync::mpsc::channel(); + let mut sink = ChannelSink(tx); + let mut source = ChannelSource(rx); + sink.send_data(b"abc").unwrap(); + sink.send_finalize().unwrap(); + assert_eq!(source.recv().unwrap(), Some(Message::Data(b"abc".to_vec()))); + assert_eq!(source.recv().unwrap(), Some(Message::Finalize)); + drop(sink); + assert_eq!(source.recv().unwrap(), None); + } +} diff --git a/src/chunker.rs b/src/chunker.rs new file mode 100644 index 0000000..676c208 --- /dev/null +++ b/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 { + 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)>; + /// 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)>; + /// 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)>>; +} + +#[derive(Default)] +pub struct RawChunker { + pending: Vec, +} + +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)> { + 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)> { + None + } + + fn finish(&mut self) -> Result)>> { + Ok(self.ready_chunk()) + } +} + +pub struct RecordAlignedChunker { + header_buf: Vec, + header_done: bool, + header_ready: bool, + endianness: Option, + pending: Vec, + 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) { + let chunk: Vec = 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)> { + 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)> { + if self.complete_len > 0 { + Some(self.drain_complete()) + } else { + None + } + } + + fn finish(&mut self) -> Result)>> { + 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 { + 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 { + 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::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 = 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()); + } +} diff --git a/src/cli.rs b/src/cli.rs new file mode 100644 index 0000000..afeabf0 --- /dev/null +++ b/src/cli.rs @@ -0,0 +1,108 @@ +use clap::Parser; + +/// scp-like tool for copying (possibly still-growing) pcap capture files. +#[derive(Parser, Debug)] +#[command(name = "scpcap")] +pub struct ClientCli { + /// Source: local path or [user@]host:path. If the name ends in the + /// partial extension, it's treated as still growing. + pub source: String, + + /// Destination: local path or [user@]host:path. + pub dest: String, + + /// Compress in flight; destination is written as a streaming .pcap.zst. + /// Only "zstd" is supported. + #[arg(long, value_name = "CODEC", num_args = 0..=1, default_missing_value = "zstd", require_equals = true)] + pub compress: Option, + + /// Target chunk size for compression, e.g. "256k", "1m". + #[arg(long, default_value = "256k")] + pub chunk: String, + + /// zstd compression level (lower = faster; default tuned for throughput + /// to stay competitive with `scp -C`). + #[arg(long, default_value_t = crate::zstd_frame::DEFAULT_LEVEL)] + pub level: i32, + + /// Suffix identifying a still-growing file. + #[arg(long, default_value = "partial")] + pub extension: String, +} + +/// Remote helper mode, ssh-spawned by the client (like `rsync --server`). +/// Not meant to be invoked directly by a user. +#[derive(Parser, Debug)] +#[command(name = "scpcap --server")] +pub struct ServerCli { + #[command(flatten)] + pub role: ServerRoleFlags, + + /// Source path to read (--send mode). + #[arg(long)] + pub source: Option, + + /// Source's final (non-partial) path, for rename detection (--send mode). + #[arg(long)] + pub final_source: Option, + + /// Whether the source is a still-growing (partial) file (--send mode). + #[arg(long)] + pub partial: bool, + + /// Compress in flight (--send mode). No adjacent positional args in + /// --server mode, so (unlike ClientCli's --compress) there's no + /// ambiguity requiring `--compress=zstd` -- `--compress zstd` is fine. + #[arg(long, value_name = "CODEC", num_args = 0..=1, default_missing_value = "zstd")] + pub compress: Option, + + /// Target chunk size for compression (--send mode). + #[arg(long, default_value = "256k")] + pub chunk: String, + + /// zstd compression level (--send mode). + #[arg(long, default_value_t = crate::zstd_frame::DEFAULT_LEVEL)] + pub level: i32, + + /// Destination final path (--recv mode). + #[arg(long)] + pub dest_final: Option, + + /// Destination temp path, written while in flight (--recv mode). + #[arg(long)] + pub dest_temp: Option, +} + +/// The clap-facing shape of the role choice: two mutually exclusive, +/// individually-long-named flags, matching rsync's `--server --send`/ +/// `--server --recv` convention. Clap's derive has no direct way to map two +/// separate boolean flags onto one enum field, so this stays a flag pair at +/// the parsing boundary and converts to `ServerRole` via `role()` below. +#[derive(clap::Args, Debug)] +#[group(required = true, multiple = false)] +pub struct ServerRoleFlags { + #[arg(long)] + pub send: bool, + #[arg(long)] + pub recv: bool, +} + +/// The role a `--server` process plays, for the rest of the code to match on +/// instead of checking two booleans. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ServerRole { + Send, + Recv, +} + +impl ServerCli { + /// The `--group(required, exclusive)` on `ServerRoleFlags` guarantees + /// exactly one of `send`/`recv` is set by the time parsing succeeds. + pub fn role(&self) -> ServerRole { + if self.role.send { + ServerRole::Send + } else { + ServerRole::Recv + } + } +} diff --git a/src/client.rs b/src/client.rs new file mode 100644 index 0000000..8227cae --- /dev/null +++ b/src/client.rs @@ -0,0 +1,127 @@ +use crate::channel::{ChannelSink, ChannelSource}; +use crate::cli::ClientCli; +use crate::naming::{self, resolve_dest_paths, strip_name, Target}; +use crate::pipeline::{run_recv, run_send, RecvOptions, SendOptions}; +use crate::protocol::{MessageSink, MessageSource, WireSink, WireSource}; +use crate::ssh::{build_server_command, spawn_remote_server}; +use anyhow::{bail, Result}; +use std::path::PathBuf; +use std::sync::mpsc; +use std::thread; + +pub fn run(cli: ClientCli) -> Result<()> { + if let Some(codec) = &cli.compress + && codec != "zstd" + { + bail!("unsupported codec {codec:?}: only \"zstd\" is supported"); + } + + let source_target = Target::parse(&cli.source); + let dest_target = Target::parse(&cli.dest); + + // strip_name works purely on suffixes, so it's fine to run on the full + // path rather than just the basename; resolve_dest_paths below extracts + // the basename itself when DEST names a directory. + let source_stripped = strip_name(source_target.path(), &cli.extension); + // final_source_path only strips the partial suffix (not .zst -- an + // already-compressed source renames to .zst, not ). + let final_source_path = naming::final_source_path(source_target.path(), &cli.extension); + let is_partial = source_stripped.was_partial; + let effective_compress = cli.compress.is_some() && !source_stripped.was_zst; + let chunk_target = naming::parse_size(&cli.chunk)?; + + let dest = resolve_dest_paths(dest_target.path(), &source_stripped, effective_compress, &cli.extension); + + match (&source_target, &dest_target) { + (Target::Local(src), Target::Local(_)) => { + let send_opts = SendOptions { + source_path: PathBuf::from(src), + is_partial, + final_source_path: PathBuf::from(&final_source_path), + compress: effective_compress, + chunk_target, + level: cli.level, + }; + let recv_opts = RecvOptions { + dest_final: PathBuf::from(&dest.final_path), + dest_temp: PathBuf::from(&dest.temp_path), + }; + let (tx, rx) = mpsc::channel(); + let recv_handle = + thread::spawn(move || run_recv(Box::new(ChannelSource(rx)), recv_opts)); + run_send(send_opts, Box::new(ChannelSink(tx)))?; + recv_handle.join().expect("recv thread panicked")?; + Ok(()) + } + (Target::Local(src), Target::Remote { host, .. }) => { + let server_args = vec![ + "--recv".to_string(), + "--dest-final".to_string(), + dest.final_path.clone(), + "--dest-temp".to_string(), + dest.temp_path.clone(), + ]; + let remote_cmd = build_server_command(&server_args); + let mut child = spawn_remote_server(host, &remote_cmd)?; + let stdin = child.stdin.take().expect("piped stdin"); + + let send_opts = SendOptions { + source_path: PathBuf::from(src), + is_partial, + final_source_path: PathBuf::from(&final_source_path), + compress: effective_compress, + chunk_target, + level: cli.level, + }; + let sink: Box = Box::new(WireSink(stdin)); + let send_result = run_send(send_opts, sink); + let status = child.wait()?; + send_result?; + if !status.success() { + bail!("remote scpcap --server --recv exited with {status}"); + } + Ok(()) + } + (Target::Remote { host, .. }, Target::Local(_)) => { + let mut server_args = vec![ + "--send".to_string(), + "--source".to_string(), + source_target.path().to_string(), + "--final-source".to_string(), + final_source_path.clone(), + "--extension".to_string(), + cli.extension.clone(), + "--chunk".to_string(), + cli.chunk.clone(), + "--level".to_string(), + cli.level.to_string(), + ]; + if is_partial { + server_args.push("--partial".to_string()); + } + if effective_compress { + server_args.push("--compress".to_string()); + server_args.push("zstd".to_string()); + } + let remote_cmd = build_server_command(&server_args); + let mut child = spawn_remote_server(host, &remote_cmd)?; + let stdout = child.stdout.take().expect("piped stdout"); + + let recv_opts = RecvOptions { + dest_final: PathBuf::from(&dest.final_path), + dest_temp: PathBuf::from(&dest.temp_path), + }; + let source: Box = Box::new(WireSource(stdout)); + let recv_result = run_recv(source, recv_opts); + let status = child.wait()?; + recv_result?; + if !status.success() { + bail!("remote scpcap --server --send exited with {status}"); + } + Ok(()) + } + (Target::Remote { .. }, Target::Remote { .. }) => { + bail!("remote-to-remote copies are not supported") + } + } +} diff --git a/src/main.rs b/src/main.rs index e7a11a9..df7e669 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,25 @@ -fn main() { - println!("Hello, world!"); +mod channel; +mod chunker; +mod cli; +mod client; +mod naming; +mod pipeline; +mod protocol; +mod server; +mod ssh; +mod tail; +mod zstd_frame; + +use clap::Parser; + +fn main() -> anyhow::Result<()> { + let mut args: Vec = std::env::args().collect(); + if args.get(1).map(|s| s.as_str()) == Some("--server") { + args.remove(1); + let server_cli = cli::ServerCli::parse_from(args); + server::run(server_cli) + } else { + let client_cli = cli::ClientCli::parse_from(args); + client::run(client_cli) + } } diff --git a/src/naming.rs b/src/naming.rs new file mode 100644 index 0000000..089f286 --- /dev/null +++ b/src/naming.rs @@ -0,0 +1,258 @@ +use anyhow::{bail, Result}; +use std::path::Path; + +/// Either a local filesystem path or a `[user@]host:path` remote target. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Target { + Local(String), + Remote { host: String, path: String }, +} + +impl Target { + pub fn parse(arg: &str) -> Target { + // scp-style: split on the first ':' that comes before the first '/'. + // A bare Windows-style drive letter isn't a concern (Linux-only). + if let Some(colon) = arg.find(':') { + let before_colon = &arg[..colon]; + if !before_colon.is_empty() && !before_colon.contains('/') { + return Target::Remote { + host: before_colon.to_string(), + path: arg[colon + 1..].to_string(), + }; + } + } + Target::Local(arg.to_string()) + } + + pub fn path(&self) -> &str { + match self { + Target::Local(p) => p, + Target::Remote { path, .. } => path, + } + } +} + +/// Strip the partial-suffix (e.g. ".partial") first, then a trailing ".zst" if +/// present. Order matters: "cap.pcap.zst.partial" must be recognized as a +/// growing, zst-compressed source, not misclassified. +pub struct StrippedName { + pub logical: String, + pub was_partial: bool, + pub was_zst: bool, +} + +pub fn strip_name(name: &str, partial_ext: &str) -> StrippedName { + let suffix = format!(".{partial_ext}"); + let (after_partial, was_partial) = match name.strip_suffix(&suffix) { + Some(stripped) => (stripped, true), + None => (name, false), + }; + let (logical, was_zst) = match after_partial.strip_suffix(".zst") { + Some(stripped) => (stripped, true), + None => (after_partial, false), + }; + StrippedName { + logical: logical.to_string(), + was_partial, + was_zst, + } +} + +/// Final (non-partial) path for a source, given its literal argument path. +pub fn final_source_path(source_path: &str, partial_ext: &str) -> String { + let suffix = format!(".{partial_ext}"); + match source_path.strip_suffix(&suffix) { + Some(stripped) => stripped.to_string(), + None => source_path.to_string(), + } +} + +pub struct DestPaths { + pub final_path: String, + pub temp_path: String, +} + +/// Resolve the destination's final and temp (in-flight) paths from the literal +/// DEST argument path, the source's stripped logical name, and whether +/// compression is actually being applied. +pub fn resolve_dest_paths( + dest_path: &str, + source_stripped: &StrippedName, + effective_compress: bool, + partial_ext: &str, +) -> DestPaths { + // DEST may name a directory (trailing '/') -- in that case reuse the + // source's logical basename underneath it. Otherwise DEST names the file + // directly. + let dest_logical = if dest_path.ends_with('/') { + let base = Path::new(&source_stripped.logical) + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_else(|| source_stripped.logical.clone()); + format!("{dest_path}{base}") + } else { + dest_path.to_string() + }; + + let final_path = if effective_compress && !dest_logical.ends_with(".zst") { + format!("{dest_logical}.zst") + } else { + dest_logical + }; + + let temp_path = format!("{final_path}.{partial_ext}"); + + DestPaths { + final_path, + temp_path, + } +} + +/// Parse a human chunk-size string like "256k", "1m", or a bare byte count. +pub fn parse_size(s: &str) -> Result { + let s = s.trim(); + if s.is_empty() { + bail!("empty size"); + } + let (digits, mult) = match s.chars().last().unwrap() { + 'k' | 'K' => (&s[..s.len() - 1], 1024), + 'm' | 'M' => (&s[..s.len() - 1], 1024 * 1024), + 'g' | 'G' => (&s[..s.len() - 1], 1024 * 1024 * 1024), + _ => (s, 1), + }; + let n: usize = digits + .trim() + .parse() + .map_err(|_| anyhow::anyhow!("invalid size: {s:?}"))?; + Ok(n * mult) +} + +/// Single-quote a string for safe inclusion as one argument in a remote shell +/// command line, e.g. for `ssh host sh -c '...'`-style invocations. +pub fn shell_quote(s: &str) -> String { + let mut out = String::with_capacity(s.len() + 2); + out.push('\''); + for c in s.chars() { + if c == '\'' { + out.push_str("'\\''"); + } else { + out.push(c); + } + } + out.push('\''); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn target_parse_local() { + assert_eq!(Target::parse("/a/b/c"), Target::Local("/a/b/c".into())); + assert_eq!(Target::parse("relative/path"), Target::Local("relative/path".into())); + } + + #[test] + fn target_parse_remote() { + assert_eq!( + Target::parse("host:/a/b"), + Target::Remote { host: "host".into(), path: "/a/b".into() } + ); + assert_eq!( + Target::parse("user@host:/a/b"), + Target::Remote { host: "user@host".into(), path: "/a/b".into() } + ); + } + + #[test] + fn target_parse_local_with_colon_after_slash() { + // A colon appearing after a '/' isn't a host separator. + assert_eq!(Target::parse("/a/b:c"), Target::Local("/a/b:c".into())); + } + + #[test] + fn strip_order_partial_then_zst() { + let s = strip_name("cap.pcap.zst.partial", "partial"); + assert_eq!(s.logical, "cap.pcap"); + assert!(s.was_partial); + assert!(s.was_zst); + } + + #[test] + fn strip_plain_partial() { + let s = strip_name("cap.pcap.partial", "partial"); + assert_eq!(s.logical, "cap.pcap"); + assert!(s.was_partial); + assert!(!s.was_zst); + } + + #[test] + fn strip_already_final() { + let s = strip_name("cap.pcap", "partial"); + assert_eq!(s.logical, "cap.pcap"); + assert!(!s.was_partial); + assert!(!s.was_zst); + } + + #[test] + fn strip_custom_extension() { + let s = strip_name("cap.pcap.zst.inprogress", "inprogress"); + assert_eq!(s.logical, "cap.pcap"); + assert!(s.was_partial); + assert!(s.was_zst); + } + + #[test] + fn final_source_path_strips_suffix() { + assert_eq!(final_source_path("/x/cap.pcap.partial", "partial"), "/x/cap.pcap"); + assert_eq!(final_source_path("/x/cap.pcap", "partial"), "/x/cap.pcap"); + } + + #[test] + fn dest_paths_append_zst_when_compressing() { + let src = strip_name("cap.pcap.partial", "partial"); + let d = resolve_dest_paths("/out/cap.pcap", &src, true, "partial"); + assert_eq!(d.final_path, "/out/cap.pcap.zst"); + assert_eq!(d.temp_path, "/out/cap.pcap.zst.partial"); + } + + #[test] + fn dest_paths_no_double_zst() { + let src = strip_name("cap.pcap.zst.partial", "partial"); + let d = resolve_dest_paths("/out/cap.pcap.zst", &src, false, "partial"); + assert_eq!(d.final_path, "/out/cap.pcap.zst"); + assert_eq!(d.temp_path, "/out/cap.pcap.zst.partial"); + } + + #[test] + fn dest_paths_no_compress() { + let src = strip_name("cap.pcap.partial", "partial"); + let d = resolve_dest_paths("/out/cap.pcap", &src, false, "partial"); + assert_eq!(d.final_path, "/out/cap.pcap"); + assert_eq!(d.temp_path, "/out/cap.pcap.partial"); + } + + #[test] + fn dest_paths_directory_target() { + let src = strip_name("cap.pcap.partial", "partial"); + let d = resolve_dest_paths("/out/", &src, false, "partial"); + assert_eq!(d.final_path, "/out/cap.pcap"); + } + + #[test] + fn size_parsing() { + assert_eq!(parse_size("256k").unwrap(), 256 * 1024); + assert_eq!(parse_size("1m").unwrap(), 1024 * 1024); + assert_eq!(parse_size("2M").unwrap(), 2 * 1024 * 1024); + assert_eq!(parse_size("512").unwrap(), 512); + assert!(parse_size("").is_err()); + assert!(parse_size("abc").is_err()); + } + + #[test] + fn quoting() { + assert_eq!(shell_quote("/plain/path"), "'/plain/path'"); + assert_eq!(shell_quote("it's"), "'it'\\''s'"); + } +} diff --git a/src/pipeline.rs b/src/pipeline.rs new file mode 100644 index 0000000..9557dbb --- /dev/null +++ b/src/pipeline.rs @@ -0,0 +1,317 @@ +use crate::chunker::{Chunker, RawChunker, RecordAlignedChunker, DEFAULT_FLUSH}; +use crate::protocol::{Message, MessageSink, MessageSource}; +use crate::tail::{run_tail, SourceEvent}; +use crate::zstd_frame::{compress_frame, MARKER_FRAME}; +use anyhow::{bail, Result}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::sync::mpsc; +use std::thread; + +/// Bound on the reader->compressor and compressor->sender channels: enough to +/// let stages overlap (disk read, CPU compress, network send happening on +/// different chunks at once) without unbounded memory growth if one stage is +/// slower than another. +const PIPELINE_DEPTH: usize = 8; + +pub struct SendOptions { + pub source_path: PathBuf, + pub is_partial: bool, + pub final_source_path: PathBuf, + pub compress: bool, + pub chunk_target: usize, + pub level: i32, +} + +enum CompressedItem { + Frame(Vec), + Eof, + Error(String), +} + +/// Reads (and, in compress mode, pcap-record-aligned-chunks and zstd-compresses) +/// the source, sending the result through `sink`. Compress mode runs a 3-stage +/// pipeline (tail thread -> compress thread -> this thread sends) so that +/// disk I/O, CPU compression, and network writes overlap instead of +/// serializing -- see the plan's Performance section. Non-compress mode +/// collapses to 2 stages since there's no compression work to overlap. +pub fn run_send(opts: SendOptions, mut sink: Box) -> Result<()> { + let chunker: Box = if opts.compress { + Box::new(RecordAlignedChunker::new(opts.chunk_target, DEFAULT_FLUSH)) + } else { + Box::new(RawChunker::new()) + }; + + let (tail_tx, tail_rx) = mpsc::sync_channel::(PIPELINE_DEPTH); + let reader = thread::spawn({ + let source_path = opts.source_path.clone(); + let is_partial = opts.is_partial; + let final_source_path = opts.final_source_path.clone(); + move || run_tail(source_path, is_partial, final_source_path, chunker, tail_tx) + }); + + if opts.compress { + let (comp_tx, comp_rx) = mpsc::sync_channel::(PIPELINE_DEPTH); + let level = opts.level; + let compressor = thread::spawn(move || -> Result<()> { + for event in tail_rx { + match event { + SourceEvent::Chunk(_kind, data) => { + let frame = compress_frame(&data, level)?; + if comp_tx.send(CompressedItem::Frame(frame)).is_err() { + return Ok(()); // receiver gone, sender loop already errored + } + } + SourceEvent::Eof => { + let _ = comp_tx.send(CompressedItem::Eof); + return Ok(()); + } + SourceEvent::Error(e) => { + let _ = comp_tx.send(CompressedItem::Error(e.clone())); + bail!(e); + } + } + } + Ok(()) + }); + + sink.send_data(&MARKER_FRAME)?; + let mut send_err: Option = None; + for item in comp_rx { + match item { + CompressedItem::Frame(f) => { + if let Err(e) = sink.send_data(&f) { + send_err = Some(e); + break; + } + } + CompressedItem::Eof => { + if let Err(e) = sink.send_finalize() { + send_err = Some(e); + } + break; + } + CompressedItem::Error(e) => { + // Best-effort: tell the peer why, so it gets a real + // diagnostic instead of just "stream ended without + // FINALIZE". If this send also fails (peer already gone), + // that's fine -- we still report the original error below. + let _ = sink.send_error(&e); + send_err = Some(anyhow::anyhow!(e)); + break; + } + } + } + let reader_result = reader.join().unwrap(); + let compressor_result = compressor.join().unwrap(); + if let Some(e) = send_err { + return Err(e); + } + reader_result?; + compressor_result?; + } else { + let mut send_err: Option = None; + for event in tail_rx { + match event { + SourceEvent::Chunk(_kind, data) => { + if let Err(e) = sink.send_data(&data) { + send_err = Some(e); + break; + } + } + SourceEvent::Eof => { + if let Err(e) = sink.send_finalize() { + send_err = Some(e); + } + break; + } + SourceEvent::Error(e) => { + let _ = sink.send_error(&e); + send_err = Some(anyhow::anyhow!(e)); + break; + } + } + } + let reader_result = reader.join().unwrap(); + if let Some(e) = send_err { + return Err(e); + } + reader_result?; + } + Ok(()) +} + +pub struct RecvOptions { + pub dest_final: PathBuf, + pub dest_temp: PathBuf, +} + +/// Writes whatever bytes arrive verbatim to a temp path, then atomically +/// renames to the final path on `Finalize`. A stream that ends without +/// `Finalize` is an abort: the temp file is left in place and this returns an +/// error rather than renaming. +pub fn run_recv(mut source: Box, opts: RecvOptions) -> Result<()> { + let mut f = File::create(&opts.dest_temp)?; + loop { + match source.recv()? { + Some(Message::Data(payload)) => f.write_all(&payload)?, + Some(Message::Finalize) => { + f.sync_all()?; + std::fs::rename(&opts.dest_temp, &opts.dest_final)?; + return Ok(()); + } + None => bail!("transfer aborted: stream ended without FINALIZE"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::channel::{ChannelSink, ChannelSource}; + use tempfile::tempdir; + + fn header_bytes(magic_le: u32) -> Vec { + let mut h = vec![0u8; crate::chunker::GLOBAL_HEADER_LEN]; + h[0..4].copy_from_slice(&magic_le.to_le_bytes()); + h + } + + fn record_bytes(payload: &[u8]) -> Vec { + let mut r = Vec::new(); + r.extend_from_slice(&[0u8; 8]); + let incl_len = payload.len() as u32; + r.extend_from_slice(&incl_len.to_le_bytes()); + r.extend_from_slice(&incl_len.to_le_bytes()); + r.extend_from_slice(payload); + r + } + + fn run_pipe(opts: SendOptions, dest_final: PathBuf, dest_temp: PathBuf) -> Result<()> { + let (tx, rx) = mpsc::channel::(); + let recv_opts = RecvOptions { dest_final, dest_temp }; + let recv_handle = thread::spawn(move || run_recv(Box::new(ChannelSource(rx)), recv_opts)); + run_send(opts, Box::new(ChannelSink(tx)))?; + recv_handle.join().unwrap() + } + + #[test] + fn non_compress_already_final_source() { + let dir = tempdir().unwrap(); + let src = dir.path().join("cap.pcap"); + std::fs::write(&src, b"just some bytes, not necessarily pcap").unwrap(); + let dest_final = dir.path().join("out.pcap"); + let dest_temp = dir.path().join("out.pcap.partial"); + + let opts = SendOptions { + source_path: src.clone(), + is_partial: false, + final_source_path: src.clone(), + compress: false, + chunk_target: crate::chunker::DEFAULT_TARGET_BYTES, + level: crate::zstd_frame::DEFAULT_LEVEL, + }; + run_pipe(opts, dest_final.clone(), dest_temp.clone()).unwrap(); + + assert!(!dest_temp.exists()); + assert_eq!( + std::fs::read(&dest_final).unwrap(), + b"just some bytes, not necessarily pcap" + ); + } + + #[test] + fn non_compress_partial_source_waits_for_rename() { + let dir = tempdir().unwrap(); + let partial = dir.path().join("cap.pcap.partial"); + let final_src = dir.path().join("cap.pcap"); + std::fs::write(&partial, b"growing-data").unwrap(); + let dest_final = dir.path().join("out.pcap"); + let dest_temp = dir.path().join("out.pcap.partial"); + + let opts = SendOptions { + source_path: partial.clone(), + is_partial: true, + final_source_path: final_src.clone(), + compress: false, + chunk_target: crate::chunker::DEFAULT_TARGET_BYTES, + level: crate::zstd_frame::DEFAULT_LEVEL, + }; + + let (tx, rx) = mpsc::channel::(); + let recv_handle = thread::spawn({ + let dest_final = dest_final.clone(); + let dest_temp = dest_temp.clone(); + move || run_recv(Box::new(ChannelSource(rx)), RecvOptions { dest_final, dest_temp }) + }); + let send_handle = thread::spawn(move || run_send(opts, Box::new(ChannelSink(tx)))); + + thread::sleep(std::time::Duration::from_millis(60)); + std::fs::rename(&partial, &final_src).unwrap(); + + send_handle.join().unwrap().unwrap(); + recv_handle.join().unwrap().unwrap(); + + assert_eq!(std::fs::read(&dest_final).unwrap(), b"growing-data"); + } + + #[test] + fn compress_mode_produces_conformant_stream_and_round_trips() { + let dir = tempdir().unwrap(); + let src = dir.path().join("cap.pcap"); + let mut stream = header_bytes(0xA1B2C3D4); + for i in 0..50u8 { + stream.extend_from_slice(&record_bytes(&[i; 100])); + } + std::fs::write(&src, &stream).unwrap(); + + let dest_final = dir.path().join("out.pcap.zst"); + let dest_temp = dir.path().join("out.pcap.zst.partial"); + let opts = SendOptions { + source_path: src.clone(), + is_partial: false, + final_source_path: src.clone(), + compress: true, + chunk_target: 512, // small target to force multiple frames + level: crate::zstd_frame::DEFAULT_LEVEL, + }; + run_pipe(opts, dest_final.clone(), dest_temp.clone()).unwrap(); + + let out_bytes = std::fs::read(&dest_final).unwrap(); + + // Marker frame present at offset 0. + assert_eq!(&out_bytes[0..14], &MARKER_FRAME[..]); + + // Whole-file zstd -d output equals the original source, byte for byte. + let decoded = zstd::stream::decode_all(out_bytes.as_slice()).unwrap(); + assert_eq!(decoded, stream); + + // Header frame (first frame after the marker) declares exactly 24 bytes. + let after_marker = &out_bytes[14..]; + let header_content_size = zstd::zstd_safe::get_frame_content_size(after_marker) + .unwrap() + .unwrap(); + assert_eq!(header_content_size, 24); + } + + #[test] + fn compress_mode_bad_magic_errors_and_leaves_no_final_file() { + let dir = tempdir().unwrap(); + let src = dir.path().join("cap.pcap"); + std::fs::write(&src, b"not a pcap file at all, 24+ bytes long").unwrap(); + let dest_final = dir.path().join("out.pcap.zst"); + let dest_temp = dir.path().join("out.pcap.zst.partial"); + + let opts = SendOptions { + source_path: src.clone(), + is_partial: false, + final_source_path: src.clone(), + compress: true, + chunk_target: crate::chunker::DEFAULT_TARGET_BYTES, + level: crate::zstd_frame::DEFAULT_LEVEL, + }; + assert!(run_pipe(opts, dest_final.clone(), dest_temp).is_err()); + assert!(!dest_final.exists()); + } +} diff --git a/src/protocol.rs b/src/protocol.rs new file mode 100644 index 0000000..b7ea0a0 --- /dev/null +++ b/src/protocol.rs @@ -0,0 +1,153 @@ +use anyhow::Result; +use std::io::{Read, Write}; + +const MSG_DATA: u8 = 0x01; +const MSG_FINALIZE: u8 = 0x02; +const MSG_ERROR: u8 = 0x03; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Message { + Data(Vec), + Finalize, +} + +/// Sends framed messages to a peer -- either a serializing wire sink (ssh +/// child stdin) or an in-process channel sink. +pub trait MessageSink: Send { + fn send_data(&mut self, payload: &[u8]) -> Result<()>; + fn send_finalize(&mut self) -> Result<()>; + fn send_error(&mut self, msg: &str) -> Result<()>; +} + +/// Receives framed messages from a peer. `Ok(None)` means the stream ended +/// without a `Finalize` -- an abort, not a clean completion. +pub trait MessageSource: Send { + fn recv(&mut self) -> Result>; +} + +pub struct WireSink(pub W); + +impl MessageSink for WireSink { + fn send_data(&mut self, payload: &[u8]) -> Result<()> { + let len = payload.len() as u64; + self.0.write_all(&[MSG_DATA])?; + self.0.write_all(&len.to_le_bytes())?; + self.0.write_all(payload)?; + self.0.flush()?; + Ok(()) + } + + fn send_finalize(&mut self) -> Result<()> { + self.0.write_all(&[MSG_FINALIZE])?; + self.0.flush()?; + Ok(()) + } + + fn send_error(&mut self, msg: &str) -> Result<()> { + let bytes = msg.as_bytes(); + let len = bytes.len() as u32; + self.0.write_all(&[MSG_ERROR])?; + self.0.write_all(&len.to_le_bytes())?; + self.0.write_all(bytes)?; + self.0.flush()?; + Ok(()) + } +} + +pub struct WireSource(pub R); + +impl MessageSource for WireSource { + fn recv(&mut self) -> Result> { + let mut tag = [0u8; 1]; + match self.0.read_exact(&mut tag) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e.into()), + } + match tag[0] { + MSG_DATA => { + let mut len_buf = [0u8; 8]; + self.0.read_exact(&mut len_buf)?; + let len = u64::from_le_bytes(len_buf) as usize; + let mut payload = vec![0u8; len]; + self.0.read_exact(&mut payload)?; + Ok(Some(Message::Data(payload))) + } + MSG_FINALIZE => Ok(Some(Message::Finalize)), + MSG_ERROR => { + let mut len_buf = [0u8; 4]; + self.0.read_exact(&mut len_buf)?; + let len = u32::from_le_bytes(len_buf) as usize; + let mut msg_buf = vec![0u8; len]; + self.0.read_exact(&mut msg_buf)?; + let msg = String::from_utf8_lossy(&msg_buf).to_string(); + anyhow::bail!("remote error: {msg}"); + } + other => anyhow::bail!("unknown message tag: {other:#x}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn round_trip_data_and_finalize() { + let mut buf: Vec = Vec::new(); + { + let mut sink = WireSink(&mut buf); + sink.send_data(b"hello").unwrap(); + sink.send_data(b"").unwrap(); + sink.send_finalize().unwrap(); + } + let mut source = WireSource(buf.as_slice()); + assert_eq!(source.recv().unwrap(), Some(Message::Data(b"hello".to_vec()))); + assert_eq!(source.recv().unwrap(), Some(Message::Data(Vec::new()))); + assert_eq!(source.recv().unwrap(), Some(Message::Finalize)); + assert_eq!(source.recv().unwrap(), None); + } + + #[test] + fn error_message_surfaces_as_err() { + let mut buf: Vec = Vec::new(); + WireSink(&mut buf).send_error("boom").unwrap(); + let mut source = WireSource(buf.as_slice()); + let err = source.recv().unwrap_err(); + assert!(err.to_string().contains("boom")); + } + + /// A `Read` impl that yields bytes a handful at a time, to exercise + /// `read_exact`'s handling of short reads (ssh pipes short-read routinely). + struct ChunkyReader { + data: Vec, + pos: usize, + step: usize, + } + + impl Read for ChunkyReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let remaining = self.data.len() - self.pos; + let n = remaining.min(self.step).min(buf.len()); + buf[..n].copy_from_slice(&self.data[self.pos..self.pos + n]); + self.pos += n; + Ok(n) + } + } + + #[test] + fn survives_short_reads() { + let mut buf: Vec = Vec::new(); + { + let mut sink = WireSink(&mut buf); + sink.send_data(&vec![7u8; 5000]).unwrap(); + sink.send_finalize().unwrap(); + } + let mut source = WireSource(ChunkyReader { data: buf, pos: 0, step: 3 }); + match source.recv().unwrap() { + Some(Message::Data(d)) => assert_eq!(d, vec![7u8; 5000]), + other => panic!("unexpected: {other:?}"), + } + assert_eq!(source.recv().unwrap(), Some(Message::Finalize)); + } +} diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..1015a8b --- /dev/null +++ b/src/server.rs @@ -0,0 +1,57 @@ +use crate::cli::{ServerCli, ServerRole}; +use crate::naming; +use crate::pipeline::{run_recv, run_send, RecvOptions, SendOptions}; +use crate::protocol::{MessageSink, MessageSource, WireSink, WireSource}; +use anyhow::{bail, Result}; +use std::path::PathBuf; + +/// Entry point for `scpcap --server ...`, spawned remotely over ssh by a +/// client. Talks the wire protocol over its own stdin/stdout. +pub fn run(cli: ServerCli) -> Result<()> { + match cli.role() { + ServerRole::Send => run_send_server(cli), + ServerRole::Recv => run_recv_server(cli), + } +} + +fn run_send_server(cli: ServerCli) -> Result<()> { + let source = cli.source.ok_or_else(|| anyhow::anyhow!("--send requires --source"))?; + let final_source = cli + .final_source + .ok_or_else(|| anyhow::anyhow!("--send requires --final-source"))?; + let compress = match &cli.compress { + Some(codec) if codec != "zstd" => bail!("unsupported codec {codec:?}"), + Some(_) => true, + None => false, + }; + let chunk_target = naming::parse_size(&cli.chunk)?; + + let opts = SendOptions { + source_path: PathBuf::from(source), + is_partial: cli.partial, + final_source_path: PathBuf::from(final_source), + compress, + chunk_target, + level: cli.level, + }; + // `Stdout`/`Stdin` (not the `.lock()` guards) are used here: the guards + // aren't `Send`, but `MessageSink`/`MessageSource` require it. + let sink: Box = Box::new(WireSink(std::io::stdout())); + run_send(opts, sink) +} + +fn run_recv_server(cli: ServerCli) -> Result<()> { + let dest_final = cli + .dest_final + .ok_or_else(|| anyhow::anyhow!("--recv requires --dest-final"))?; + let dest_temp = cli + .dest_temp + .ok_or_else(|| anyhow::anyhow!("--recv requires --dest-temp"))?; + + let opts = RecvOptions { + dest_final: PathBuf::from(dest_final), + dest_temp: PathBuf::from(dest_temp), + }; + let source: Box = Box::new(WireSource(std::io::stdin())); + run_recv(source, opts) +} diff --git a/src/ssh.rs b/src/ssh.rs new file mode 100644 index 0000000..c1a679b --- /dev/null +++ b/src/ssh.rs @@ -0,0 +1,44 @@ +use crate::naming::shell_quote; +use anyhow::{Context, Result}; +use std::process::{Child, Command, Stdio}; + +/// Builds the argv for the remote `scpcap --server` invocation, run as a +/// single ssh remote-command argument (so it goes through the remote shell, +/// same as rsync's `ssh host rsync --server ...`). +pub fn build_server_command(args: &[String]) -> String { + let mut parts = vec!["scpcap".to_string(), "--server".to_string()]; + parts.extend(args.iter().map(|a| shell_quote(a))); + parts.join(" ") +} + +/// Spawns `ssh ` with piped stdin/stdout, inheriting +/// stderr so remote diagnostics are visible to the user directly. +pub fn spawn_remote_server(host: &str, remote_command: &str) -> Result { + Command::new("ssh") + .arg("--") + .arg(host) + .arg(remote_command) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .with_context(|| format!("failed to spawn ssh to {host}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builds_quoted_command() { + let cmd = build_server_command(&[ + "--recv".to_string(), + "--dest-final".to_string(), + "/a b/c.pcap".to_string(), + ]); + assert_eq!( + cmd, + "scpcap --server '--recv' '--dest-final' '/a b/c.pcap'" + ); + } +} diff --git a/src/tail.rs b/src/tail.rs new file mode 100644 index 0000000..8a349ac --- /dev/null +++ b/src/tail.rs @@ -0,0 +1,288 @@ +use crate::chunker::{Chunker, ChunkKind}; +use anyhow::{bail, Result}; +use std::fs::File; +use std::os::unix::fs::FileExt; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +/// Poll interval for growth/rename checks, per STREAMING_ZSTD_FORMAT.md section 9. +pub const POLL_INTERVAL: Duration = Duration::from_millis(20); +const READ_WINDOW: usize = 1024 * 1024; + +pub enum SourceEvent { + Chunk(ChunkKind, Vec), + Eof, + Error(String), +} + +/// Abstracts over `mpsc::Sender` and `mpsc::SyncSender` so `run_tail` can feed +/// either an unbounded or a bounded (backpressured) event channel. +pub trait EventSink: Send { + fn send(&self, event: SourceEvent) -> Result<(), SourceEvent>; +} + +impl EventSink for std::sync::mpsc::Sender { + fn send(&self, event: SourceEvent) -> Result<(), SourceEvent> { + std::sync::mpsc::Sender::send(self, event).map_err(|e| e.0) + } +} + +impl EventSink for std::sync::mpsc::SyncSender { + fn send(&self, event: SourceEvent) -> Result<(), SourceEvent> { + // A full bounded channel blocks here, which is the intended + // backpressure: the tailer won't outrun a slower downstream stage. + std::sync::mpsc::SyncSender::send(self, event).map_err(|e| e.0) + } +} + +/// Tails `source_path` (opened exactly once, never reopened by path -- even +/// after the rename below is observed) and feeds bytes to `chunker`, sending +/// ready chunks to `events`. Implements STREAMING_ZSTD_FORMAT.md section 7: +/// finalizes only once the rename to `final_path` has been observed AND every +/// byte up to that point has been drained. If `source_path` was not itself a +/// partial-suffixed name, it's treated as "born complete" (section 7.6). +/// +/// Per the confirmed design: if `source_path` doesn't exist at all when this +/// is called, that's a hard error (no polling for creation). +pub fn run_tail( + source_path: PathBuf, + is_partial: bool, + final_path: PathBuf, + mut chunker: Box, + events: impl EventSink, +) -> Result<()> { + if !source_path.exists() { + bail!("source not found: {}", source_path.display()); + } + let result = drain(&source_path, is_partial, &final_path, chunker.as_mut(), &events); + match &result { + Ok(()) => { + let _ = events.send(SourceEvent::Eof); + } + Err(e) => { + let _ = events.send(SourceEvent::Error(e.to_string())); + } + } + result +} + +fn drain( + source_path: &Path, + is_partial: bool, + final_path: &Path, + chunker: &mut dyn Chunker, + events: &impl EventSink, +) -> Result<()> { + let file = File::open(source_path)?; + let mut pos: u64 = 0; + let mut complete_signal = !is_partial; + let mut buf = vec![0u8; READ_WINDOW]; + + loop { + let file_len = file.metadata()?.len(); + // Whether this iteration did real work -- if so, loop again + // immediately rather than sleeping, so a large already-available + // file drains at disk/network speed instead of being capped at + // READ_WINDOW / POLL_INTERVAL. Only an iteration that found nothing + // to do (genuinely waiting for growth or the rename) sleeps. + let mut made_progress = false; + if file_len > pos { + let want = ((file_len - pos) as usize).min(buf.len()); + let n = file.read_at(&mut buf[..want], pos)?; + if n > 0 { + chunker.feed(&buf[..n])?; + pos += n as u64; + emit_ready(chunker, events)?; + made_progress = true; + } + } else if chunker.should_flush(Instant::now()) + && let Some((kind, data)) = chunker.flush() + { + send_chunk(events, kind, data)?; + } + // If file_len < pos, the writer truncated/restarted mid-file + // (STREAMING_ZSTD_FORMAT.md section 7.7): do nothing, never rewind + // `pos`, just wait for it to regrow past `pos` on a later poll. + + if !complete_signal && final_path.try_exists()? { + // Rename observed -- signal now, well before finalization, per + // section 7.4's repoint-then-drain-then-finalize ordering. + complete_signal = true; + } + + let current_len = file.metadata()?.len(); + if complete_signal && pos == current_len { + if let Some((kind, data)) = chunker.finish()? { + send_chunk(events, kind, data)?; + } + return Ok(()); + } + + if !made_progress { + std::thread::sleep(POLL_INTERVAL); + } + } +} + +fn emit_ready(chunker: &mut dyn Chunker, events: &impl EventSink) -> Result<()> { + while let Some((kind, data)) = chunker.ready_chunk() { + send_chunk(events, kind, data)?; + } + Ok(()) +} + +fn send_chunk(events: &impl EventSink, kind: ChunkKind, data: Vec) -> Result<()> { + events + .send(SourceEvent::Chunk(kind, data)) + .map_err(|_| anyhow::anyhow!("event channel closed")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chunker::RawChunker; + use std::io::Write; + use std::sync::mpsc; + use tempfile::tempdir; + + fn collect(rx: mpsc::Receiver) -> (Vec, bool) { + let mut data = Vec::new(); + let mut got_eof = false; + for ev in rx { + match ev { + SourceEvent::Chunk(_, d) => data.extend_from_slice(&d), + SourceEvent::Eof => got_eof = true, + SourceEvent::Error(e) => panic!("unexpected error event: {e}"), + } + } + (data, got_eof) + } + + #[test] + fn already_final_source_is_born_complete() { + let dir = tempdir().unwrap(); + let path = dir.path().join("cap.pcap"); + std::fs::write(&path, b"hello world").unwrap(); + + let (tx, rx) = mpsc::channel(); + run_tail(path.clone(), false, path.clone(), Box::new(RawChunker::new()), tx).unwrap(); + let (data, eof) = collect(rx); + assert_eq!(data, b"hello world"); + assert!(eof); + } + + #[test] + fn missing_source_is_a_hard_error() { + let dir = tempdir().unwrap(); + let path = dir.path().join("nope.pcap.partial"); + let final_path = dir.path().join("nope.pcap"); + let (tx, _rx) = mpsc::channel(); + let err = run_tail(path, true, final_path, Box::new(RawChunker::new()), tx).unwrap_err(); + assert!(err.to_string().contains("not found")); + } + + #[test] + fn partial_source_waits_for_rename_and_drain() { + let dir = tempdir().unwrap(); + let partial_path = dir.path().join("cap.pcap.partial"); + let final_path = dir.path().join("cap.pcap"); + let mut f = File::create(&partial_path).unwrap(); + f.write_all(b"chunk1").unwrap(); + f.sync_all().unwrap(); + + let (tx, rx) = mpsc::channel(); + let handle = std::thread::spawn({ + let partial_path = partial_path.clone(); + let final_path = final_path.clone(); + move || run_tail(partial_path, true, final_path, Box::new(RawChunker::new()), tx) + }); + + std::thread::sleep(Duration::from_millis(60)); + f.write_all(b"chunk2").unwrap(); + f.sync_all().unwrap(); + std::thread::sleep(Duration::from_millis(60)); + drop(f); + std::fs::rename(&partial_path, &final_path).unwrap(); + + handle.join().unwrap().unwrap(); + let (data, eof) = collect(rx); + assert_eq!(data, b"chunk1chunk2"); + assert!(eof); + } + + #[test] + fn held_fd_survives_rename_and_even_unlink() { + let dir = tempdir().unwrap(); + let partial_path = dir.path().join("cap.pcap.partial"); + let final_path = dir.path().join("cap.pcap"); + let mut f = File::create(&partial_path).unwrap(); + f.write_all(b"abc").unwrap(); + f.sync_all().unwrap(); + + let (tx, rx) = mpsc::channel(); + let handle = std::thread::spawn({ + let partial_path = partial_path.clone(); + let final_path = final_path.clone(); + move || run_tail(partial_path, true, final_path, Box::new(RawChunker::new()), tx) + }); + + std::thread::sleep(Duration::from_millis(60)); + std::fs::rename(&partial_path, &final_path).unwrap(); + // Continue writing through the original handle after the rename -- + // the reader must keep draining via its already-open fd. + f.write_all(b"def").unwrap(); + f.sync_all().unwrap(); + drop(f); + + handle.join().unwrap().unwrap(); + let (data, eof) = collect(rx); + assert_eq!(data, b"abcdef"); + assert!(eof); + } + + #[test] + fn truncation_then_regrowth_does_not_error_or_rewind() { + // Simulates a writer restart (section 7.7): the file shrinks below + // what we've already read, then regrows past it with the corrected + // continuation. The reader must not error and must not rewind `pos`. + let dir = tempdir().unwrap(); + let partial_path = dir.path().join("cap.pcap.partial"); + let final_path = dir.path().join("cap.pcap"); + std::fs::write(&partial_path, b"0123456789").unwrap(); + + let (tx, rx) = mpsc::channel(); + let handle = std::thread::spawn({ + let partial_path = partial_path.clone(); + let final_path = final_path.clone(); + move || run_tail(partial_path, true, final_path, Box::new(RawChunker::new()), tx) + }); + + // Let the reader fully drain the initial 10 bytes. + std::thread::sleep(Duration::from_millis(80)); + // Writer restarts: truncate below what's already been read, then + // regrow past it -- open with truncate + write fresh content that's + // shorter than the 10 bytes already consumed, then extend further. + { + let mut f = std::fs::OpenOptions::new().write(true).truncate(true).open(&partial_path).unwrap(); + f.write_all(b"abc").unwrap(); + f.sync_all().unwrap(); + } + std::thread::sleep(Duration::from_millis(80)); + { + let mut f = std::fs::OpenOptions::new().append(true).open(&partial_path).unwrap(); + f.write_all(b"defghijklmnop").unwrap(); + f.sync_all().unwrap(); + } + std::thread::sleep(Duration::from_millis(80)); + std::fs::rename(&partial_path, &final_path).unwrap(); + + handle.join().unwrap().unwrap(); + let (data, eof) = collect(rx); + // The reader never rewinds: it was already past byte 10 ("0123456789") + // when the truncate+regrow happened, so it waits at pos=10 until the + // file regrows past that point, then reads the correct continuation + // ("abcdefghijklmnop"[10..] == "klmnop"). + assert_eq!(data, b"0123456789klmnop"); + assert!(eof); + } +} diff --git a/src/zstd_frame.rs b/src/zstd_frame.rs new file mode 100644 index 0000000..7c0fc29 --- /dev/null +++ b/src/zstd_frame.rs @@ -0,0 +1,65 @@ +use anyhow::Result; +use std::io::Write; + +/// zstd skippable-frame marker identifying this as a pcap-stream .pcap.zst +/// container. See STREAMING_ZSTD_FORMAT.md section 2 for the byte layout. +pub const MARKER_FRAME: [u8; 14] = [ + 0x50, 0x2A, 0x4D, 0x18, // skippable magic, LE + 0x06, 0x00, 0x00, 0x00, // payload size = 6 + 0x50, 0x53, 0x5A, 0x31, // "PSZ1" + 0x01, // version + 0x00, // flags +]; + +/// Default zstd compression level: chosen for throughput, not ratio, since +/// each chunk is an independent frame that already sacrifices some ratio for +/// streamability -- lean into speed to stay competitive with `scp -C`. +pub const DEFAULT_LEVEL: i32 = 1; + +/// Compress `data` as a single independent zstd frame with a declared content +/// size and a content checksum, per STREAMING_ZSTD_FORMAT.md invariants 4/5. +/// The one-shot `zstd::bulk` helpers don't guarantee either, so this goes +/// through the streaming `Encoder` API explicitly. +pub fn compress_frame(data: &[u8], level: i32) -> Result> { + let mut encoder = zstd::stream::Encoder::new(Vec::new(), level)?; + encoder.set_pledged_src_size(Some(data.len() as u64))?; + encoder.include_checksum(true)?; + encoder.write_all(data)?; + Ok(encoder.finish()?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn round_trips_and_declares_size_and_checksum() { + let data = b"the quick brown fox jumps over the lazy dog".repeat(100); + let frame = compress_frame(&data, DEFAULT_LEVEL).unwrap(); + + let content_size = zstd::zstd_safe::get_frame_content_size(&frame) + .unwrap() + .expect("content size must be declared"); + assert_eq!(content_size, data.len() as u64); + + let decoded = zstd::stream::decode_all(frame.as_slice()).unwrap(); + assert_eq!(decoded, data); + } + + #[test] + fn corrupted_frame_fails_checksum() { + let data = b"some pcap-shaped bytes".repeat(50); + let mut frame = compress_frame(&data, DEFAULT_LEVEL).unwrap(); + let last = frame.len() - 1; + frame[last] ^= 0xFF; + assert!(zstd::stream::decode_all(frame.as_slice()).is_err()); + } + + #[test] + fn marker_frame_bytes() { + assert_eq!( + MARKER_FRAME, + [0x50, 0x2A, 0x4D, 0x18, 0x06, 0x00, 0x00, 0x00, 0x50, 0x53, 0x5A, 0x31, 0x01, 0x00] + ); + } +}