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