diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index 3cf14712..75e3d11b 100644 --- a/crates/engine/src/segment/mod.rs +++ b/crates/engine/src/segment/mod.rs @@ -262,7 +262,7 @@ impl Segmenter { if mirrors.is_empty() { return Err(lightcraft_segment::fetch::no_mirrors_message()); } - self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::Options::default()) + self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::options()) } #[cfg(not(feature = "sam"))] { diff --git a/crates/fetch/README.md b/crates/fetch/README.md new file mode 100644 index 00000000..60ac1d32 --- /dev/null +++ b/crates/fetch/README.md @@ -0,0 +1,7 @@ +# lightcraft-fetch + +Native model downloads in pure Rust (L0, Apache-2.0): HTTP/1.1 over std sockets and rustls with the RustCrypto provider, without ring, aws-lc, C or assembly build dependencies. The wasm32 crate is empty. + +Call `download` with bounded file specifications and ordered mirrors, or `download_url` for an exact catalog URL (including its query). Each file has a local filename independent of the URL, a size cap and optional pinned size/SHA-256. Transfers resume pinned `.part` files, verify before renaming, report progress and check cancellation during reads. Run on a worker thread. No proxy support; DNS resolution uses the platform resolver. Tokens are off by default and never forwarded to a different redirect host. + +Tests use local synthetic HTTP servers for redirects, resume, truncation, cancellation, stalls, hashes, size limits and hostile specifications; they make no internet requests. diff --git a/crates/fetch/src/http.rs b/crates/fetch/src/http.rs index 332f0147..e900850a 100644 --- a/crates/fetch/src/http.rs +++ b/crates/fetch/src/http.rs @@ -60,7 +60,10 @@ pub struct Url { impl Url { pub fn parse(s: &str) -> Result { - let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}")); + let bad = |why: &str| HttpError::BadUrl(why.to_string()); + if s.len() > 16 * 1024 { + return Err(bad("URL exceeds 16 KiB")); + } let s = s.trim(); if s.chars().any(|c| c.is_control() || c == ' ') { return Err(bad("contains spaces or control characters")); @@ -509,6 +512,12 @@ pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result { + pub name: &'a str, /// Exact size in bytes, when pinned (then the download resumes and is checked against it). pub size: Option, /// SHA-256 (lower-case hex), when pinned. - pub sha256: Option<&'static str>, + pub sha256: Option<&'a str>, /// Largest size accepted when `size` is not pinned. pub max: u64, } @@ -79,11 +79,14 @@ pub struct Options { pub stall_timeout: Duration, /// Tries per mirror and file (each resumes where the last one stopped). pub attempts: u32, + /// The environment variable holding a bearer token for the first mirror's host (sent only over https, never to a + /// host a redirect leads to). `None`: nothing is ever sent, so one model's token cannot leak to another's server. + pub token_env: Option<&'static str>, } impl Default for Options { fn default() -> Self { - Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 } + Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3, token_env: None } } } @@ -103,6 +106,8 @@ pub struct Progress { pub enum DownloadError { /// No mirror is configured. NoMirrors, + /// Invalid filename, size, hash or aggregate bound supplied by the caller. + Invalid(String), Cancelled, /// A local file problem (folder not writable, disk full…). Disk(String), @@ -116,6 +121,7 @@ pub enum DownloadError { impl std::fmt::Display for DownloadError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + DownloadError::Invalid(e) => write!(f, "invalid download: {e}"), DownloadError::NoMirrors => write!(f, "no download location is configured for this model in this build"), DownloadError::Cancelled => write!(f, "the download was cancelled"), DownloadError::Disk(e) => write!(f, "{e}"), @@ -132,16 +138,45 @@ enum Fail { Disk(String), } -/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there -/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it). +/// Fetch a single file from its exact URL, preserving any query string. The local filename +/// remains `file.name`; it never comes from redirects or response headers. +pub fn download_url( + file: &FileSpec<'_>, + url: &str, + dir: &Path, + opts: &Options, + cancel: &AtomicBool, + progress: &mut dyn FnMut(&Progress), +) -> Result<(), DownloadError> { + download_source(std::slice::from_ref(file), &[url.to_string()], dir, opts, cancel, progress, true) +} + +/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Verified files +/// are kept. `progress` is called as bytes arrive (often: throttle in it). pub fn download( - files: &[FileSpec], + files: &[FileSpec<'_>], mirrors: &[String], dir: &Path, opts: &Options, cancel: &AtomicBool, progress: &mut dyn FnMut(&Progress), ) -> Result<(), DownloadError> { + download_source(files, mirrors, dir, opts, cancel, progress, false) +} + +fn download_source( + files: &[FileSpec<'_>], + mirrors: &[String], + dir: &Path, + opts: &Options, + cancel: &AtomicBool, + progress: &mut dyn FnMut(&Progress), + exact_urls: bool, +) -> Result<(), DownloadError> { + let total = validate_files(files)?; + if mirrors.len() > MAX_MIRRORS { + return Err(DownloadError::Invalid("too many mirrors".into())); + } std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?; // what's left to fetch let mut todo = Vec::new(); @@ -156,14 +191,13 @@ pub fn download( if mirrors.is_empty() { return Err(DownloadError::NoMirrors); } - let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum(); let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum(); for f in todo { let mut errors = Vec::new(); let mut ok = false; 'mirrors: for m in mirrors { - let url = format!("{m}/{}", f.name); - for _ in 0..opts.attempts.max(1) { + let url = if exact_urls { m.clone() } else { format!("{m}/{}", f.name) }; + for _ in 0..opts.attempts.clamp(1, 16) { if cancel.load(Ordering::Relaxed) { return Err(DownloadError::Cancelled); } @@ -184,11 +218,11 @@ pub fn download( Err(Fail::Cancelled) => return Err(DownloadError::Cancelled), Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)), Err(Fail::Retry(e)) => { - log::warn!("download of {url}: {e} (retrying)"); + log::warn!("download of {}: {e} (retrying)", short(&url)); errors.push(format!("{}: {e}", short(m))); } Err(Fail::NextMirror(e)) => { - log::warn!("download of {url}: {e}"); + log::warn!("download of {}: {e}", short(&url)); errors.push(format!("{}: {e}", short(m))); continue 'mirrors; } @@ -203,13 +237,56 @@ pub fn download( Ok(()) } +/// Validate before touching disk, including names supplied by a user model catalog. +fn validate_files(files: &[FileSpec<'_>]) -> Result { + if files.len() > 256 { + return Err(DownloadError::Invalid("too many files".into())); + } + let mut seen = std::collections::HashSet::new(); + let mut total = 0u64; + for f in files { + let stem = f.name.split('.').next().unwrap_or_default().to_ascii_uppercase(); + let reserved = matches!(stem.as_str(), "CON" | "PRN" | "AUX" | "NUL") + || stem + .strip_prefix("COM") + .or_else(|| stem.strip_prefix("LPT")) + .is_some_and(|n| matches!(n, "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9")); + if reserved + || f.name.ends_with(['.', ' ']) + || f.name.ends_with(".part") + || f.name.is_empty() + || f.name.len() > 255 + || f.name.starts_with('.') + || f.name.contains(['/', '\\', ':']) + || f.name.chars().any(char::is_control) + || !seen.insert(f.name) + { + return Err(DownloadError::Invalid("a filename must be a unique, ordinary filename".into())); + } + if f.max == 0 + || f.size.is_some_and(|n| n == 0 || n > f.max) + || f.sha256.is_some_and(|h| h.len() != 64 || !h.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))) + { + return Err(DownloadError::Invalid("invalid file size or SHA-256".into())); + } + total = total.checked_add(f.size.unwrap_or(0)).ok_or_else(|| DownloadError::Invalid("total size overflow".into()))?; + } + Ok(total) +} + +/// The bearer token to send to `url`: only a non-empty one, only over https, and only to the host the download began +/// at (`origin`), never to the host a redirect leads to. +fn bearer<'a>(token: Option<&'a str>, url: &Url, origin: &str) -> Option<&'a str> { + token.map(str::trim).filter(|t| !t.is_empty() && url.tls && url.host == origin) +} + /// A mirror for messages (no query string: it may hold a token). fn short(m: &str) -> String { Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into()) } /// Whether `path` is already the right file (exact size and hash when pinned). -fn verified(f: &FileSpec, path: &Path) -> Result { +fn verified(f: &FileSpec<'_>, path: &Path) -> Result { let Ok(meta) = std::fs::metadata(path) else { return Ok(false) }; if !meta.is_file() || meta.len() == 0 { return Ok(false); @@ -243,7 +320,7 @@ fn part_path(dir: &Path, name: &str) -> PathBuf { /// Fetch one file from one mirror into `.part` (resuming it when its size is pinned), /// check it and move it into place. -fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> { +fn fetch_file(f: &FileSpec<'_>, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> { let part = part_path(dir, f.name); let limit = f.size.unwrap_or(f.max); let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display())); @@ -268,12 +345,8 @@ fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &Atom headers.push(("Range", format!("bytes={have}-"))); } // a token for the configured host only (never sent on to a redirect's host) - if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN") - && !t.trim().is_empty() - && url.tls - && url.host == origin - { - headers.push(("Authorization", format!("Bearer {}", t.trim()))); + if let Some(t) = bearer(opts.token_env.and_then(|name| std::env::var(name).ok()).as_deref(), &url, &origin) { + headers.push(("Authorization", format!("Bearer {t}"))); } let r = http::get(&url, &headers, &limits).map_err(http_fail)?; match r.status { diff --git a/crates/fetch/src/tests.rs b/crates/fetch/src/tests.rs index 57c168a2..4a828c6d 100644 --- a/crates/fetch/src/tests.rs +++ b/crates/fetch/src/tests.rs @@ -22,7 +22,7 @@ fn leak(s: String) -> &'static str { Box::leak(s.into_boxed_str()) } -fn pinned() -> FileSpec { +fn pinned() -> FileSpec<'static> { let p = payload(); FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 } } @@ -137,10 +137,10 @@ fn tmp(name: &str) -> PathBuf { } fn quick() -> Options { - Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 } + Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2, token_env: None } } -fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec) { +fn run(files: &[FileSpec<'_>], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec) { let mut seen = Vec::new(); let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone())); (r, seen) @@ -314,3 +314,75 @@ fn content_ranges() { assert_eq!(parse_content_range("bytes 9-1/300"), None); assert_eq!(parse_content_range("items 1-2/3"), None); } + +/// A bearer token goes only to https, only to the host the download began at, and only when there is one. +#[test] +fn a_token_is_sent_only_over_https_to_the_first_host() { + let u = |s: &str| Url::parse(s).unwrap(); + let first = u("https://huggingface.co/x/model.bin"); + assert_eq!(bearer(Some(" hf_abc \n"), &first, "huggingface.co"), Some("hf_abc")); + // the storage host a redirect leads to, plain http, no token, an empty one + assert_eq!(bearer(Some("hf_abc"), &u("https://cdn-lfs.example/x"), "huggingface.co"), None); + assert_eq!(bearer(Some("hf_abc"), &u("http://huggingface.co/x"), "huggingface.co"), None); + assert_eq!(bearer(None, &first, "huggingface.co"), None); + assert_eq!(bearer(Some(" "), &first, "huggingface.co"), None); + // a caller that names no variable sends nothing, whatever the environment holds + assert!(Options::default().token_env.is_none()); +} + +#[test] +fn exact_url_keeps_query_and_owned_catalog_values() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/download?file=weights&sig=secret", listener.local_addr().unwrap()); + let worker = std::thread::spawn(move || { + let (mut socket, _) = listener.accept().unwrap(); + let mut reader = BufReader::new(socket.try_clone().unwrap()); + let mut request = String::new(); + reader.read_line(&mut request).unwrap(); + assert_eq!(request.trim(), "GET /download?file=weights&sig=secret HTTP/1.1"); + loop { + let mut line = String::new(); + if reader.read_line(&mut line).unwrap() == 0 || line.trim().is_empty() { + break; + } + } + socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc").unwrap(); + }); + let name = String::from("catalog.onnx"); + let hash = sha_hex(b"abc"); + let spec = FileSpec { name: &name, size: Some(3), sha256: Some(&hash), max: 3 }; + let dir = tmp("exact"); + let mut progress = Vec::new(); + download_url(&spec, &url, &dir, &quick(), &AtomicBool::new(false), &mut |p| progress.push(p.clone())).unwrap(); + worker.join().unwrap(); + assert_eq!(std::fs::read(dir.join(name)).unwrap(), b"abc"); + assert!(progress.iter().all(|p| !p.mirror.contains("secret"))); + let _ = std::fs::remove_dir_all(dir); +} + +#[test] +fn hostile_file_specs_are_rejected_before_touching_disk() { + let dir = tmp("invalid"); + let base = FileSpec { name: "valid.onnx", size: Some(1), sha256: None, max: 1 }; + let cancel = AtomicBool::new(false); + for name in ["", ".hidden", "../escape", "a/b", "a\\b", "C:escape", "bad\nname", "trailing.", "CON.onnx", "a.part"] { + let spec = FileSpec { name, ..base.clone() }; + assert!(matches!(run(&[spec], &[], &dir, &cancel).0, Err(DownloadError::Invalid(_))), "{name}"); + assert!(!dir.exists()); + } + for spec in [ + FileSpec { max: 0, ..base.clone() }, + FileSpec { size: Some(0), ..base.clone() }, + FileSpec { size: Some(2), ..base.clone() }, + FileSpec { sha256: Some("not a SHA"), ..base.clone() }, + ] { + assert!(matches!(run(&[spec], &[], &dir, &cancel).0, Err(DownloadError::Invalid(_)))); + assert!(!dir.exists()); + } + let huge = FileSpec { size: Some(u64::MAX), max: u64::MAX, ..base.clone() }; + let other = FileSpec { name: "other.onnx", ..huge.clone() }; + for specs in [vec![huge, other], vec![base.clone(); 257], vec![base.clone(), base]] { + assert!(matches!(run(&specs, &[], &dir, &cancel).0, Err(DownloadError::Invalid(_)))); + assert!(!dir.exists()); + } +} diff --git a/crates/segment/src/fetch/mod.rs b/crates/segment/src/fetch/mod.rs index 55c083cc..35f2d2eb 100644 --- a/crates/segment/src/fetch/mod.rs +++ b/crates/segment/src/fetch/mod.rs @@ -51,6 +51,14 @@ pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec { lightcraft_fetch::mirrors(env, file, DEFAULT_MIRRORS) } +/// Download settings for SAM 3: the defaults, plus [`TOKEN_ENV`] for a gated mirror. +pub fn options() -> Options { + Options { token_env: Some(TOKEN_ENV), ..Options::default() } +} + +/// Environment variable with a bearer token for the SAM 3 mirror's own host (a gated Hugging Face repository). +pub const TOKEN_ENV: &str = "LIGHTCRAFT_SAM3_TOKEN"; + /// What to tell the user when there is no mirror for the SAM 3 model. pub fn no_mirrors_message() -> String { format!(