server/lib/tests/ai_models_on_the_intake_chain.rs
//! **AI models are first-class artifacts on the intake chain, and the pickle
//! family is blocked by default.**
//!
//! A real config: four `raw` repositories, the promotion ladder over them, the
//! `intake:` block with `scan:` on, and the `models:` block. The chain is the one
//! [`Holger::intake_chain`] builds from that config, so the model scanner and the
//! property store are the server's own.
//!
//! * Small generated GGUF, Safetensors and ONNX files pass SCANNED, and their
//! metadata (format, architecture, parameters, tensors, dtype, quantization)
//! lands in the server's property store, where the property search axis finds it.
//! * A tiny benign pickle and a fake `.pt` zip holding `data.pkl` are refused by
//! default, each with its frozen code.
//! * A pickle renamed `.safetensors` is refused by its content.
//! * A Safetensors header that lies about its offsets is refused.
//! * `models.allow_pickle` naming a repository the config does not declare
//! refuses the boot (`holger-model-8`).
use std::io::Write;
use std::sync::Arc;
use holger_server_lib::audit::{AuditLog, MemoryAuditLog};
use holger_server_lib::promote::intake::{
ArtifactState, IntakeChain, IntakeError, TransitionRequest,
};
use holger_server_lib::scanner::model::fixtures;
use holger_server_lib::search::{search_repos_with_properties, SearchQuery};
use holger_server_lib::{read_ron_config, wire_holger, ArtifactId, Holger};
struct Server {
_data: tempfile::TempDir,
_cfg: tempfile::NamedTempFile,
holger: Holger,
}
fn config(models: &str) -> Server {
let data = tempfile::tempdir().unwrap();
let d = data.path().display().to_string();
let repo = |name: &str| {
format!(
r#"(ron_name: "{name}", ron_repo_type: "raw", ron_upstreams: [], ron_in: None,
ron_out: Some((ron_storage_endpoint: "storage", ron_exposed_endpoint: "door")),
ron_data_dir: Some("{d}/{name}")),"#
)
};
let ron = format!(
r#"(
repositories: [ {} {} {} {} ],
exposed_endpoints: [ (ron_name: "door", ron_url: "127.0.0.1:1", ron_http_url: Some("127.0.0.1:2")) ],
storage_endpoints: [ (ron_name: "storage", ron_storage_type: "znippy", ron_path: "{d}") ],
audit: (dir: "{d}/audit", enabled: true),
promotion: Some((
ladder_name: "intake",
stages: [
(name: "incoming", repo: "incoming"),
(name: "quarantine", repo: "quarantine"),
(name: "approved", repo: "approved"),
(name: "production", repo: "production"),
],
intake: Some((rules: (sandbox: off),
incoming: "incoming", quarantine: "quarantine",
approved: "approved", production: "production",
scan: Some((required: [model])),
)),
)),
{models}
)"#,
repo("incoming"),
repo("quarantine"),
repo("approved"),
repo("production"),
);
let mut cfg = tempfile::NamedTempFile::new().unwrap();
cfg.write_all(ron.as_bytes()).unwrap();
let mut holger: Holger = read_ron_config(cfg.path()).expect("config parses");
holger.instantiate_backends().expect("backends");
wire_holger(&mut holger).expect("wire");
Server {
_data: data,
_cfg: cfg,
holger,
}
}
fn chain(s: &Server) -> Arc<IntakeChain> {
let audit: Arc<dyn AuditLog> = Arc::new(MemoryAuditLog::new());
s.holger
.intake_chain(audit)
.expect("the chain builds")
.expect("the config declares one")
}
fn id(name: &str) -> ArtifactId {
ArtifactId {
namespace: None,
name: name.into(),
version: String::new(),
}
}
fn ingest_and_scan(
s: &Server,
chain: &IntakeChain,
name: &str,
bytes: &[u8],
) -> Result<(), IntakeError> {
let incoming = s.holger.get_repo("incoming").unwrap().clone();
let put = id(name);
chain.ingest(
"alice",
"127.0.0.1:9",
"incoming",
&id(name),
bytes,
move || incoming.put(&put, bytes),
)?;
chain
.advance(TransitionRequest::new(
id(name),
ArtifactState::Scanned,
"sec-officer",
))
.map(|_| ())
}
#[test]
fn gguf_safetensors_and_onnx_pass_and_their_metadata_is_searchable() {
let s = config("");
let chain = chain(&s);
for (name, bytes) in [
("llm/tiny.gguf", fixtures::gguf()),
("bert/model.safetensors", fixtures::safetensors()),
("mlp/tiny.onnx", fixtures::onnx()),
] {
ingest_and_scan(&s, &chain, name, &bytes)
.unwrap_or_else(|e| panic!("{name}: {}", e.tagged()));
assert_eq!(chain.state(&id(name)), ArtifactState::Scanned, "{name}");
}
let props = s.holger.property_store();
let gguf = props.get("production", &id("llm/tiny.gguf"));
assert_eq!(
gguf.get("model.architecture"),
Some(&vec!["llama".to_string()]),
"{gguf:?}"
);
assert_eq!(gguf.get("model.parameters"), Some(&vec!["12".to_string()]));
assert_eq!(
gguf.get("model.dtype"),
Some(&vec!["F16".to_string(), "F32".to_string()])
);
let st = props.get("incoming", &id("bert/model.safetensors"));
assert_eq!(
st.get("model.format"),
Some(&vec!["safetensors".to_string()]),
"{st:?}"
);
assert_eq!(st.get("model.tensors"), Some(&vec!["2".to_string()]));
let onnx = props.get("incoming", &id("mlp/tiny.onnx"));
assert_eq!(
onnx.get("model.architecture"),
Some(&vec!["tiny-mlp".to_string()]),
"{onnx:?}"
);
let repos = vec![(
"incoming".to_string(),
s.holger.get_repo("incoming").unwrap().clone(),
)];
for (key, value, want) in [
("model.format", "onnx", "mlp/tiny.onnx"),
("model.architecture", "llama", "llm/tiny.gguf"),
("model.architecture", "bert", "bert/model.safetensors"),
] {
let q = SearchQuery {
properties: vec![(key.into(), value.into())],
..Default::default()
};
let hits = search_repos_with_properties(&repos, &q, Some(props.as_ref()));
let names: Vec<&str> = hits.artifacts.iter().map(|h| h.id.name.as_str()).collect();
assert_eq!(names, vec![want], "{key}={value}");
}
}
#[test]
fn the_pickle_family_is_blocked_by_default_with_a_frozen_code() {
let s = config("");
let chain = chain(&s);
let pt = zip_entry::fixture::zip(&[
("archive/data.pkl", &fixtures::pickle()[..]),
("archive/version", &b"3\n"[..]),
]);
for (name, bytes, code) in [
("tiny.pkl", fixtures::pickle(), "holger-model-1"),
("fake.pt", pt, "holger-model-2"),
("renamed.safetensors", fixtures::pickle(), "holger-model-1"),
(
"lying.safetensors",
fixtures::safetensors_lying_offsets(),
"holger-model-4",
),
] {
let err = ingest_and_scan(&s, &chain, name, &bytes).expect_err(name);
assert!(
matches!(err, IntakeError::PolicyRefused { .. }),
"{name}: {err}"
);
assert!(
err.tagged().contains(code),
"{name} names {code}: {}",
err.tagged()
);
assert_eq!(
chain.state(&id(name)),
ArtifactState::Ingested,
"{name} did not move"
);
}
}
#[test]
fn allow_pickle_must_name_a_declared_repository() {
let s = config(r#"models: (allow_pickle: ["no-such-repo"]),"#);
let audit: Arc<dyn AuditLog> = Arc::new(MemoryAuditLog::new());
let err = s.holger.intake_chain(audit).expect_err("refused");
assert!(format!("{err:#}").contains("[holger-model-8]"), "{err:#}");
let s = config(r#"models: (allow_pickle: ["incoming"]),"#);
let chain = chain(&s);
ingest_and_scan(&s, &chain, "approved-research.pkl", &fixtures::pickle())
.expect("a pickle in an allowed repository is a warning, not a block");
}