use std::fs;
use serde_json::json;
use tempfile::tempdir;
use super::ID;
use super::REPO;
use super::SHA;
use super::body;
use super::config;
use super::layout::tree;
use super::oid;
use super::server::Reply;
use super::server::Server;
use crate::error::CanddingError;
use crate::hub::download::fetch;
use crate::registry::WeightsFormat;
fn listing(
file: &'static str,
reply: impl Fn() -> Reply + Send + Sync + 'static,
) -> impl Fn(&str, &str) -> Reply + Send + Sync + 'static {
move |_, path| {
let path = path.split('?').next().unwrap_or(path);
if path == format!("/api/models/{ID}/revision/main") {
return Reply::new(200, json!({ "id": ID, "sha": SHA }).to_string());
}
if path.starts_with(&format!("/api/models/{ID}/tree/")) {
let entry =
json!({ "type": "file", "oid": oid(file), "size": body(file).len(), "path": file });
return Reply::new(200, json!([entry]).to_string());
}
reply()
}
}
#[test]
fn hostile_commit_and_etag_headers_on_a_file_stay_out_of_every_path() {
let root = tempdir().unwrap();
let cache = root.path().join("a/b/cache");
fs::create_dir_all(&cache).unwrap();
let server = Server::start(listing("config.json", || {
Reply::new(200, body("config.json"))
.header("ETag", "\"../../../../../escaped-etag\"")
.header("X-Repo-Commit", "../../../../../escaped-commit")
}));
let config = config(&cache, &server.url);
let snapshot = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
assert_eq!(snapshot, cache.join(REPO).join("snapshots").join(SHA));
let outside: Vec<String> = tree(root.path(), "")
.into_iter()
.filter(|entry| !entry.starts_with("a/b/cache/"))
.collect();
assert_eq!(outside, Vec::<String>::new());
let pointer = fs::read_link(snapshot.join("config.json")).unwrap();
assert_eq!(
pointer.to_string_lossy(),
format!("../../blobs/{}", oid("config.json"))
);
}
#[test]
fn a_404_with_a_hostile_commit_header_truncates_nothing() {
let root = tempdir().unwrap();
let cache = root.path().join("a/b/cache");
fs::create_dir_all(&cache).unwrap();
let victim = root.path().join("victim/settings.json");
fs::create_dir_all(victim.parent().unwrap()).unwrap();
fs::write(&victim, "precious user settings\n").unwrap();
let server = Server::start(listing("settings.json", || {
Reply::new(404, "")
.header("X-Error-Code", "EntryNotFound")
.header("X-Repo-Commit", "../../../../../victim")
}));
let config = config(&cache, &server.url);
match fetch(&config, ID, "main", WeightsFormat::Safetensors) {
Err(CanddingError::DownloadFailed { id, .. }) => assert_eq!(id, ID),
other => panic!("{other:?}"),
}
assert_eq!(
fs::read_to_string(&victim).unwrap(),
"precious user settings\n"
);
let outside: Vec<String> = tree(root.path(), "")
.into_iter()
.filter(|entry| !entry.starts_with("a/b/cache/"))
.collect();
assert_eq!(outside, ["victim/settings.json"]);
}