clawft_kernel/model_manifest/
adopt.rs

1//! Adopt-in-place scanning (ADR-101 section 3): hash model files where they
2//! lie and describe them as a [`ModelPackageBody`]. Nothing is copied or
3//! moved; the result records absolute locations only in the local registry,
4//! never in the signed body (which carries paths relative to the root).
5//!
6//! Recognised layouts: an HF cache snapshot or MLX quant directory (a tree
7//! of weight files plus tokenizer files, symlinks followed), a single GGUF
8//! file, and Ollama blobs addressed through an Ollama manifest.
9
10use std::fs::{self, File};
11use std::io::Read;
12use std::path::{Path, PathBuf};
13use std::time::UNIX_EPOCH;
14
15use serde::Deserialize;
16
17use super::body::{ModelError, ModelFormat, ModelPackageBody, ModelSource, MAX_SHARDS};
18use crate::workload_pkg::manifest::{valid_token, validate_rel_path};
19use crate::workload_pkg::FileRef;
20
21/// Most directory levels scanned below the root.
22const MAX_DEPTH: usize = 4;
23/// Most entries visited in one scan (guards a mis-pointed root).
24const MAX_VISITED: usize = 4096;
25/// Extensions treated as weight shards.
26const WEIGHT_EXTS: &[&str] = &["safetensors", "gguf", "bin", "pt", "pth", "npz", "ckpt"];
27/// Tokenizer files, first match wins.
28const TOKENIZER_FILES: &[&str] = &["tokenizer.json", "tokenizer.model"];
29/// Chat template files, or the config that carries the template.
30const TEMPLATE_FILES: &[&str] = &["chat_template.jinja", "chat_template.json", "tokenizer_config.json"];
31
32/// What a file is for.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, Deserialize)]
34#[serde(rename_all = "snake_case")]
35pub enum FileRole {
36    /// A weight shard.
37    Shard,
38    /// The tokenizer.
39    Tokenizer,
40    /// The chat template (or its carrier).
41    Template,
42}
43
44/// A file as found: relative path, size, mtime and content hash.
45#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, Deserialize)]
46pub struct ScannedFile {
47    /// Path relative to the root (forward slashes).
48    pub path: String,
49    /// Role.
50    pub role: FileRole,
51    /// Size in bytes.
52    pub size: u64,
53    /// Modification time, nanoseconds since the epoch (0 when unknown).
54    pub mtime_ns: u64,
55    /// BLAKE3 hex.
56    pub blake3: String,
57}
58
59/// Operator-supplied description of what is being adopted.
60#[derive(Debug, Clone)]
61pub struct AdoptInput {
62    /// Model name.
63    pub name: String,
64    /// Layout.
65    pub format: ModelFormat,
66    /// Provenance.
67    pub source: ModelSource,
68    /// The operator explicitly opts in to sharing the weights.
69    pub redistributable: bool,
70}
71
72/// A scanned model, ready to be attested and registered.
73#[derive(Debug, Clone)]
74pub struct ScannedModel {
75    /// Canonical model root.
76    pub root: PathBuf,
77    /// The describing body.
78    pub body: ModelPackageBody,
79    /// Every hashed file with its stat stamp.
80    pub files: Vec<ScannedFile>,
81}
82
83pub(crate) fn io_err(path: &Path, e: impl std::fmt::Display) -> ModelError {
84    ModelError::Io {
85        path: path.display().to_string(),
86        msg: e.to_string(),
87    }
88}
89
90/// Stat stamp of a file: `(size, mtime_ns)`. Follows symlinks.
91pub fn stamp(path: &Path) -> std::io::Result<(u64, u64)> {
92    let m = fs::metadata(path)?;
93    let mtime = m
94        .modified()
95        .ok()
96        .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
97        .map(|d| u64::try_from(d.as_nanos()).unwrap_or(u64::MAX))
98        .unwrap_or(0);
99    Ok((m.len(), mtime))
100}
101
102/// Stream-hash a file with BLAKE3 (1 MiB buffer; never read whole).
103pub fn hash_file(path: &Path) -> std::io::Result<String> {
104    let mut f = File::open(path)?;
105    let mut h = blake3::Hasher::new();
106    let mut buf = vec![0u8; 1024 * 1024];
107    loop {
108        let n = f.read(&mut buf)?;
109        if n == 0 {
110            break;
111        }
112        h.update(&buf[..n]);
113    }
114    Ok(h.finalize().to_hex().to_string())
115}
116
117/// Directories a model's files may resolve into: the canonical root, and for
118/// an HF cache snapshot (`.../snapshots/<rev>`) the sibling `blobs` directory
119/// of the same repo cache, which is where its files are links to.
120pub fn allowed_roots(root: &Path) -> Vec<PathBuf> {
121    let mut v = vec![root.to_path_buf()];
122    if let Some(snaps) = root.parent()
123        && snaps.file_name().is_some_and(|n| n == "snapshots")
124        && let Some(repo) = snaps.parent()
125        && let Ok(blobs) = repo.join("blobs").canonicalize()
126        && blobs.is_dir()
127    {
128        v.push(blobs);
129    }
130    v
131}
132
133/// Resolve `path` (following symlinks) and require the result to lie under
134/// one of `allowed`. Fails closed: a link that leaves the model root is an
135/// error naming the path, not a silently followed file.
136pub fn contained(path: &Path, allowed: &[PathBuf]) -> Result<PathBuf, ModelError> {
137    let real = path.canonicalize().map_err(|e| io_err(path, e))?;
138    if allowed.iter().any(|a| real.starts_with(a)) {
139        Ok(real)
140    } else {
141        Err(ModelError::Scan(format!(
142            "{} resolves outside the model root; refusing to read it",
143            path.display()
144        )))
145    }
146}
147
148fn scan_file_entry(
149    root: &Path,
150    allowed: &[PathBuf],
151    rel: &str,
152    role: FileRole,
153) -> Result<ScannedFile, ModelError> {
154    let abs = contained(&root.join(rel), allowed)?;
155    let (size, mtime_ns) = stamp(&abs).map_err(|e| io_err(&abs, e))?;
156    let blake3 = hash_file(&abs).map_err(|e| io_err(&abs, e))?;
157    Ok(ScannedFile {
158        path: rel.to_string(),
159        role,
160        size,
161        mtime_ns,
162        blake3,
163    })
164}
165
166fn is_weight(name: &str) -> bool {
167    name.rsplit_once('.')
168        .is_some_and(|(_, ext)| WEIGHT_EXTS.contains(&ext.to_ascii_lowercase().as_str()))
169}
170
171fn walk(
172    root: &Path,
173    allowed: &[PathBuf],
174    dir: &Path,
175    depth: usize,
176    found: &mut Vec<String>,
177    visited: &mut usize,
178) -> Result<(), ModelError> {
179    let Ok(rd) = fs::read_dir(dir) else { return Ok(()) };
180    for entry in rd.flatten() {
181        *visited += 1;
182        if *visited > MAX_VISITED {
183            return Ok(());
184        }
185        let name = entry.file_name().to_string_lossy().into_owned();
186        if name.starts_with('.') || !valid_token(&name, 128) {
187            continue;
188        }
189        let path = entry.path();
190        let Ok(link) = fs::symlink_metadata(&path) else { continue };
191        if link.file_type().is_symlink() {
192            // A symlinked directory is never followed. A symlinked file is
193            // followed only when it stays inside the allowed roots; a
194            // dangling or looping link is skipped.
195            match fs::metadata(&path) {
196                Ok(m) if m.is_dir() => {
197                    return Err(ModelError::Scan(format!(
198                        "{} is a symlinked directory; refusing to follow it",
199                        path.display()
200                    )));
201                }
202                Ok(_) => {
203                    contained(&path, allowed)?;
204                }
205                Err(_) => continue,
206            }
207        }
208        let Ok(meta) = fs::metadata(&path) else { continue };
209        if meta.is_dir() {
210            if depth < MAX_DEPTH {
211                walk(root, allowed, &path, depth + 1, found, visited)?;
212            }
213        } else if meta.is_file()
214            && let Ok(rel) = path.strip_prefix(root)
215        {
216            found.push(rel.to_string_lossy().replace('\\', "/"));
217        }
218    }
219    Ok(())
220}
221
222fn finish(
223    root: PathBuf,
224    input: AdoptInput,
225    shards: Vec<ScannedFile>,
226    tokenizer: Option<ScannedFile>,
227    template: Option<ScannedFile>,
228) -> Result<ScannedModel, ModelError> {
229    if shards.is_empty() {
230        return Err(ModelError::Scan(format!(
231            "no model weight files ({}) found under {}",
232            WEIGHT_EXTS.join(", "),
233            root.display()
234        )));
235    }
236    if shards.len() > MAX_SHARDS {
237        return Err(ModelError::Scan(format!(
238            "{} weight files exceeds the {MAX_SHARDS} shard limit",
239            shards.len()
240        )));
241    }
242    let mut body = ModelPackageBody {
243        name: input.name,
244        format: input.format,
245        shards: shards
246            .iter()
247            .map(|s| FileRef {
248                path: s.path.clone(),
249                size: s.size,
250                blake3: s.blake3.clone(),
251            })
252            .collect(),
253        tokenizer_blake3: tokenizer.as_ref().map(|t| t.blake3.clone()),
254        tokenizer_path: tokenizer.as_ref().map(|t| t.path.clone()),
255        template_blake3: template.as_ref().map(|t| t.blake3.clone()),
256        template_path: template.as_ref().map(|t| t.path.clone()),
257        source: input.source,
258        redistributable: input.redistributable,
259    };
260    body.shards.sort_by(|a, b| a.path.cmp(&b.path));
261    body.validate()?;
262    let mut files = shards;
263    files.extend(tokenizer);
264    files.extend(template);
265    Ok(ScannedModel { root, body, files })
266}
267
268/// Scan a model directory in place: weight files become shards, a tokenizer
269/// and a template file (when present) are hashed too.
270pub fn scan_dir(root: &Path, input: AdoptInput) -> Result<ScannedModel, ModelError> {
271    let root = root.canonicalize().map_err(|e| io_err(root, e))?;
272    if !root.is_dir() {
273        return Err(ModelError::Scan(format!("{} is not a directory", root.display())));
274    }
275    let allowed = allowed_roots(&root);
276    let mut found = Vec::new();
277    walk(&root, &allowed, &root, 0, &mut found, &mut 0)?;
278    found.sort();
279    let mut shards = Vec::new();
280    for rel in found.iter().filter(|r| is_weight(r.rsplit('/').next().unwrap_or(r))) {
281        shards.push(scan_file_entry(&root, &allowed, rel, FileRole::Shard)?);
282    }
283    let pick = |names: &[&str], role| -> Result<Option<ScannedFile>, ModelError> {
284        names
285            .iter()
286            .find(|n| found.iter().any(|f| f == *n))
287            .map(|n| scan_file_entry(&root, &allowed, n, role))
288            .transpose()
289    };
290    let tokenizer = pick(TOKENIZER_FILES, FileRole::Tokenizer)?;
291    let template = pick(TEMPLATE_FILES, FileRole::Template)?;
292    finish(root, input, shards, tokenizer, template)
293}
294
295/// Scan one weight file (a GGUF) in place. The root is its directory and
296/// the single shard is its file name.
297pub fn scan_file(path: &Path, input: AdoptInput) -> Result<ScannedModel, ModelError> {
298    let abs = path.canonicalize().map_err(|e| io_err(path, e))?;
299    let name = abs
300        .file_name()
301        .and_then(|n| n.to_str())
302        .ok_or_else(|| ModelError::Scan("file name is not valid UTF-8".into()))?
303        .to_string();
304    let root = abs
305        .parent()
306        .ok_or_else(|| ModelError::Scan("file has no parent directory".into()))?
307        .to_path_buf();
308    validate_rel_path(&name).map_err(|e| ModelError::Scan(e.to_string()))?;
309    let shard = scan_file_entry(&root, &allowed_roots(&root), &name, FileRole::Shard)?;
310    finish(root, input, vec![shard], None, None)
311}
312
313#[derive(Deserialize)]
314struct OllamaManifest {
315    layers: Vec<OllamaLayer>,
316}
317
318#[derive(Deserialize)]
319struct OllamaLayer {
320    #[serde(rename = "mediaType")]
321    media_type: String,
322    digest: String,
323    size: u64,
324}
325
326fn ollama_blob_rel(digest: &str) -> Result<String, ModelError> {
327    let hex = digest
328        .strip_prefix("sha256:")
329        .filter(|h| h.len() == 64 && h.bytes().all(|b| b.is_ascii_hexdigit()))
330        .ok_or_else(|| ModelError::Scan(format!("bad ollama digest {digest:?}")))?;
331    Ok(format!("blobs/sha256-{}", hex.to_ascii_lowercase()))
332}
333
334/// Scan an Ollama model in place through its manifest under `models_dir`
335/// (`manifests/registry.ollama.ai/library/<name>/<tag>`). The model layer
336/// becomes the shard; the template layer, when present, the template hash.
337/// `input.source.ollama_tag` is set to `name:tag`.
338pub fn scan_ollama(
339    models_dir: &Path,
340    name: &str,
341    tag: &str,
342    mut input: AdoptInput,
343) -> Result<ScannedModel, ModelError> {
344    for part in [name, tag] {
345        if !valid_token(part, 64) {
346            return Err(ModelError::Scan(format!("bad ollama name or tag {part:?}")));
347        }
348    }
349    let root = models_dir.canonicalize().map_err(|e| io_err(models_dir, e))?;
350    let allowed = allowed_roots(&root);
351    let mpath = contained(
352        &root
353            .join("manifests/registry.ollama.ai/library")
354            .join(name)
355            .join(tag),
356        &allowed,
357    )?;
358    let bytes = fs::read(&mpath).map_err(|e| io_err(&mpath, e))?;
359    if bytes.len() > 1024 * 1024 {
360        return Err(ModelError::Scan("ollama manifest too large".into()));
361    }
362    let m: OllamaManifest = serde_json::from_slice(&bytes)
363        .map_err(|e| ModelError::Scan(format!("ollama manifest: {e}")))?;
364    let mut shards = Vec::new();
365    let mut template = None;
366    for layer in &m.layers {
367        let (role, slot) = match layer.media_type.as_str() {
368            "application/vnd.ollama.image.model" => (FileRole::Shard, true),
369            "application/vnd.ollama.image.template" => (FileRole::Template, false),
370            _ => continue,
371        };
372        let rel = ollama_blob_rel(&layer.digest)?;
373        let f = scan_file_entry(&root, &allowed, &rel, role)?;
374        if f.size != layer.size {
375            return Err(ModelError::HashMismatch {
376                path: rel,
377                expected: format!("{} bytes", layer.size),
378                actual: format!("{} bytes", f.size),
379            });
380        }
381        if slot {
382            shards.push(f);
383        } else {
384            template = Some(f);
385        }
386    }
387    input.format = ModelFormat::Ollama;
388    input.source.ollama_tag = Some(format!("{name}:{tag}"));
389    finish(root, input, shards, None, template)
390}
clawft_kernel/model_manifest/adopt.rs

clawft_kernel/model_manifest/
adopt.rs
