interscript_rs/imf/
registry.rs

1//! models.yaml index resolution + cached downloads (SPEC.md §1/§2,
2//! `interscript_ml/registry.py` parity).
3//!
4//! Conformance points implemented here:
5//! * C1 — model ids resolve against a `models.yaml` index of
6//!   `version: 1`, honoring `SECRYST_INDEX`.
7//! * C2 — the whole-artifact sha256 is verified before installation
8//!   AND on every cache hit; a mismatch fails loudly.
9//! * §2 — resolution order: verified cache hit -> download (single URL
10//!   or parts in index order, each part verified as it lands) ->
11//!   whole-artifact verify -> atomic rename into place.
12//! * `renames:` — WO87 target-accurate rename aliases (old id -> new
13//!   id); old ids stay resolvable.
14
15use crate::error::{Error, Result};
16use crate::imf::loader::sha256_file;
17use std::collections::BTreeMap;
18use std::io::{Read, Write};
19use std::path::{Path, PathBuf};
20
21const DEFAULT_INDEX_URL: &str = "https://github.com/interscript/interscript-models/releases/download/index-v12/models-index.yaml";
22const ENV_INDEX: &str = "SECRYST_INDEX";
23const ENV_CACHE: &str = "SECRYST_CACHE";
24
25fn model_err(what: impl Into<String>) -> Error {
26    Error::Model(what.into())
27}
28
29#[derive(Debug, Clone)]
30pub struct Part {
31    pub url: String,
32    pub sha256: String,
33    pub size: u64,
34}
35
36#[derive(Debug, Clone)]
37pub struct IndexEntry {
38    pub id: String,
39    pub filename: String,
40    pub url: String,
41    pub sha256: String,
42    pub size: u64,
43    pub precision: String,
44    pub task: String,
45    pub parts: Vec<Part>,
46}
47
48#[derive(Debug, Clone, Default)]
49pub struct Index {
50    pub models: BTreeMap<String, IndexEntry>,
51    /// WO87 renames: old id -> current id (old ids stay resolvable).
52    pub renames: BTreeMap<String, String>,
53}
54
55impl Index {
56    /// Resolve a possibly-renamed id to its current entry id.
57    pub fn canonical_id(&self, id: &str) -> String {
58        let mut cur = id.to_string();
59        // Follow the rename chain (bounded to reject cycles).
60        for _ in 0..8 {
61            match self.renames.get(&cur) {
62                Some(next) => cur = next.clone(),
63                None => break,
64            }
65        }
66        cur
67    }
68}
69
70/// `SECRYST_CACHE` or `~/.cache/secryst` (the shared cache layout).
71pub fn cache_dir() -> PathBuf {
72    if let Some(dir) = std::env::var_os(ENV_CACHE) {
73        return PathBuf::from(dir);
74    }
75    let home = std::env::var_os("HOME")
76        .map(PathBuf::from)
77        .unwrap_or_else(|| PathBuf::from("."));
78    home.join(".cache").join("secryst")
79}
80
81/// The index source: `SECRYST_INDEX`, then the default index-v12 URL.
82fn index_source(explicit: Option<&str>) -> String {
83    explicit
84        .map(str::to_string)
85        .or_else(|| std::env::var(ENV_INDEX).ok())
86        .unwrap_or_else(|| DEFAULT_INDEX_URL.to_string())
87}
88
89#[derive(serde::Deserialize)]
90struct IndexYaml {
91    version: i32,
92    #[serde(default)]
93    models: BTreeMap<String, IndexEntryYaml>,
94    #[serde(default)]
95    renames: BTreeMap<String, String>,
96}
97
98#[derive(serde::Deserialize)]
99struct IndexEntryYaml {
100    filename: String,
101    #[serde(default)]
102    url: String,
103    sha256: String,
104    #[serde(default)]
105    size: u64,
106    #[serde(default)]
107    precision: String,
108    #[serde(default)]
109    task: String,
110    #[serde(default)]
111    parts: Vec<PartYaml>,
112}
113
114#[derive(serde::Deserialize)]
115struct PartYaml {
116    url: String,
117    sha256: String,
118    #[serde(default)]
119    size: u64,
120}
121
122/// Load a `models.yaml` index from an http(s) URL or a local path.
123pub fn load_index(explicit: Option<&str>) -> Result<Index> {
124    let source = index_source(explicit);
125    let text = if source.starts_with("http://") || source.starts_with("https://") {
126        let mut res = ureq::get(&source)
127            .call()
128            .map_err(|e| model_err(format!("index fetch failed ({source}): {e}")))?;
129        let buf = res
130            .body_mut()
131            .read_to_string()
132            .map_err(|e| model_err(format!("index read failed ({source}): {e}")))?;
133        buf
134    } else {
135        std::fs::read_to_string(&source)
136            .map_err(|e| model_err(format!("index read failed ({source}): {e}")))?
137    };
138
139    let raw: IndexYaml = serde_yaml::from_str(&text)
140        .map_err(|e| model_err(format!("index parse failed ({source}): {e}")))?;
141    if raw.version != 1 {
142        return Err(model_err(format!(
143            "index must be a mapping with version: 1, got {}",
144            raw.version
145        )));
146    }
147
148    let mut models = BTreeMap::new();
149    for (model_id, spec) in raw.models {
150        models.insert(
151            model_id.clone(),
152            IndexEntry {
153                id: model_id,
154                filename: spec.filename,
155                url: spec.url,
156                sha256: spec.sha256.to_lowercase(),
157                size: spec.size,
158                precision: if spec.precision.is_empty() {
159                    "fp32".into()
160                } else {
161                    spec.precision
162                },
163                task: spec.task,
164                parts: spec
165                    .parts
166                    .into_iter()
167                    .map(|p| Part {
168                        url: p.url,
169                        sha256: p.sha256.to_lowercase(),
170                        size: p.size,
171                    })
172                    .collect(),
173            },
174        );
175    }
176    Ok(Index {
177        models,
178        renames: raw.renames,
179    })
180}
181
182/// §2: return a verified local zip path for `model_id`, downloading
183/// into the cache when needed. Never returns an unverified file.
184pub fn resolve(model_id: &str, index: &Index) -> Result<PathBuf> {
185    let canonical = index.canonical_id(model_id);
186    let entry = index.models.get(&canonical).ok_or_else(|| {
187        model_err(format!(
188            "unknown model id {model_id:?} (known: {} ids; use --artifact-dir or SECRYST_CACHE)",
189            index.models.len()
190        ))
191    })?;
192
193    let target = cache_dir()
194        .join("models")
195        .join(&canonical)
196        .join(&entry.filename);
197    // Cache hits are re-verified, never trusted blindly (C2).
198    if target.is_file() && sha256_file(&target)? == entry.sha256 {
199        return Ok(target);
200    }
201
202    if let Some(parent) = target.parent() {
203        std::fs::create_dir_all(parent).map_err(|e| {
204            model_err(format!("cache dir {} create failed: {e}", parent.display()))
205        })?;
206    }
207    let tmp = target.with_extension(format!(
208        "part-{}",
209        std::process::id()
210    ));
211
212    let result = (|| -> Result<()> {
213        if !entry.parts.is_empty() {
214            download_parts(&entry.parts, &tmp)?;
215        } else if entry.url.is_empty() {
216            return Err(model_err(format!(
217                "index entry {canonical:?} has no url and no cached copy"
218            )));
219        } else if let Some(path) = entry.url.strip_prefix("file://") {
220            // file:// is a mirror, not a move.
221            std::fs::copy(path, &tmp)
222                .map_err(|e| model_err(format!("channel file copy failed ({path}): {e}")))?;
223        } else {
224            download_to(&entry.url, &tmp)?;
225        }
226        let actual = sha256_file(&tmp)?;
227        if actual != entry.sha256 {
228            return Err(model_err(format!(
229                "downloaded {} sha256 mismatch: got {actual}, index says {}",
230                entry.filename, entry.sha256
231            )));
232        }
233        // C6 (SHOULD): install atomically.
234        std::fs::rename(&tmp, &target)
235            .map_err(|e| model_err(format!("cache install failed: {e}")))?;
236        Ok(())
237    })();
238
239    if result.is_err() {
240        let _ = std::fs::remove_file(&tmp);
241    }
242    result?;
243    Ok(target)
244}
245
246fn open_channel(url: &str) -> Result<Box<dyn std::io::Read>> {
247    if let Some(path) = url.strip_prefix("file://") {
248        let f = std::fs::File::open(path)
249            .map_err(|e| model_err(format!("channel file missing ({path}): {e}")))?;
250        return Ok(Box::new(f));
251    }
252    let res = ureq::get(url)
253        .call()
254        .map_err(|e| model_err(format!("download failed ({url}): {e}")))?;
255    Ok(Box::new(
256        res.into_body()
257            .into_reader(),
258    ))
259}
260
261fn download_to(url: &str, dest: &Path) -> Result<()> {
262    let mut out = std::fs::File::create(dest)
263        .map_err(|e| model_err(format!("temp file create failed: {e}")))?;
264    let mut stream = open_channel(url)?;
265    std::io::copy(&mut stream, &mut out)
266        .map_err(|e| model_err(format!("download stream failed ({url}): {e}")))?;
267    Ok(())
268}
269
270/// §1: parts stream into one file in index order, each verified as it
271/// lands; the caller then checks the assembled file as usual.
272fn download_parts(parts: &[Part], dest: &Path) -> Result<()> {
273    use sha2::{Digest, Sha256};
274    let mut out = std::fs::OpenOptions::new()
275        .create(true)
276        .append(true)
277        .open(dest)
278        .map_err(|e| model_err(format!("temp file open failed: {e}")))?;
279    for (i, part) in parts.iter().enumerate() {
280        let mut hasher = Sha256::new();
281        let mut stream = open_channel(&part.url)?;
282        let mut buf = vec![0u8; 1024 * 1024];
283        loop {
284            let n = stream
285                .read(&mut buf)
286                .map_err(|e| model_err(format!("part stream failed: {e}")))?;
287            if n == 0 {
288                break;
289            }
290            out.write_all(&buf[..n])
291                .map_err(|e| model_err(format!("part write failed: {e}")))?;
292            hasher.update(&buf[..n]);
293        }
294        let actual = crate::imf::loader::to_hex(&hasher.finalize());
295        if actual != part.sha256 {
296            return Err(model_err(format!(
297                "part {i} sha256 mismatch: got {actual}, index says {}",
298                part.sha256
299            )));
300        }
301    }
302    Ok(())
303}
interscript_rs/imf/registry.rs

interscript_rs/imf/
registry.rs
