259 lines
7.9 KiB
Rust
259 lines
7.9 KiB
Rust
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'");
|
|
}
|
|
}
|