diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index 3cf14712e..75e3d11b7 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/src/lib.rs b/crates/fetch/src/lib.rs index 499a6eda3..12ad639c7 100644 --- a/crates/fetch/src/lib.rs +++ b/crates/fetch/src/lib.rs @@ -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 } } } @@ -203,6 +206,12 @@ pub fn download( Ok(()) } +/// 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()) @@ -268,12 +277,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 57c168a2d..ad3f1063d 100644 --- a/crates/fetch/src/tests.rs +++ b/crates/fetch/src/tests.rs @@ -137,7 +137,7 @@ 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) { @@ -314,3 +314,18 @@ 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()); +} diff --git a/crates/segment/src/fetch/mod.rs b/crates/segment/src/fetch/mod.rs index 55c083ccf..35f2d2eb5 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!(