candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation
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"))
  );
}

/// hf-hub wrote `.no_exist/<X-Repo-Commit>/<file>` on a 404, truncating whatever the header led to.
#[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"]);
}