server/lib/src/scanner/model.rs
//! AI models as first-class artifacts: GGUF, Safetensors and ONNX are told by
//! their content, checked and described; the pickle family is blocked.
//!
//! A model file lives in a `raw` repository like any other blob; what makes it a
//! model is this scanner, which fills the [`ScannerKind::Model`] slot of the
//! scanner framework on the intake chain's SCANNED transition.
//!
//! ## What the content decides
//!
//! The file's name is never trusted for a pass. The bytes are sniffed in this order:
//!
//! * **GGUF**: the magic `GGUF` and version 2 or 3, then the key/value metadata
//! and the tensor table, every count and length checked against the bytes that
//! are left.
//! * **Safetensors**: the 8-byte little-endian header length, then a JSON header
//! that must fit inside the file and stay under [`MAX_SAFETENSORS_HEADER`].
//! * **PyTorch zip checkpoint**: an archive that holds `data.pkl`, found by
//! S3's archive walker (`holger-inspect`: zip, tar and gzip, nested, under its
//! bomb and traversal limits).
//! * **Pickle**: a stream of valid pickle opcodes that ends in `STOP`.
//! * **ONNX**: a protobuf `ModelProto` (field 1 `ir_version` first, a `graph`
//! field), parsed best-effort for its header only.
//!
//! ## What blocks
//!
//! * A pickle (by its opcodes), a PyTorch zip with `data.pkl`, or a file whose
//! name carries a pickle-family extension (`pkl`, `pickle`, `pt`, `pth`,
//! `joblib`, also under a compression suffix). Loading a pickle runs code.
//! `models.allow_pickle` turns this into a warning for the repositories it
//! names, and only for them.
//! * A GGUF or Safetensors header that is malformed, that breaks a limit, whose
//! tensors lie outside the file, or whose tensors overlap. These block in every
//! repository.
//! * A name that claims `.gguf`, `.safetensors` or `.onnx` over content that is
//! not that format.
//!
//! Every block cites its reason by name with a frozen `holger-model-<n>` code.
//!
//! ## What a model file gets
//!
//! A passing model file's evidence record carries a document with its format,
//! architecture, parameter count, tensor count, dtypes and quantization. The
//! intake chain also writes these as `model.*` properties, so the property
//! search axis finds them (`model.architecture=llama`).
use std::collections::{BTreeMap, BTreeSet};
use holger_errcode as ec;
use serde::{Deserialize, Serialize};
use super::{FindingSeverity, ScanFinding, ScanOutcome, ScanSubject, Scanner, ScannerKind};
/// The scanner id the evidence names.
pub const MODEL_SCANNER_ID: &str = "holger-model";
/// The largest Safetensors JSON header accepted, in bytes (the reference
/// implementation's limit).
pub const MAX_SAFETENSORS_HEADER: u64 = 100 * 1024 * 1024;
/// The most tensors one GGUF or Safetensors file may declare.
pub const MAX_TENSORS: u64 = 1 << 20;
/// The most metadata key/value pairs one GGUF file may declare.
pub const MAX_GGUF_KV: u64 = 1 << 20;
/// The longest GGUF string (a key or a string value), in bytes.
pub const MAX_GGUF_STRING: u64 = 64 * 1024 * 1024;
/// The deepest GGUF array nesting accepted.
pub const MAX_GGUF_ARRAY_DEPTH: u32 = 4;
/// GGML's own limit on tensor dimensions.
pub const MAX_TENSOR_DIMS: u32 = 4;
/// The most of a zip checkpoint's `data.pkl` that is read to name its imports.
pub const MAX_DATA_PKL_BYTES: usize = 64 * 1024 * 1024;
/// The most imports one pickle finding lists.
const MAX_GLOBALS_LISTED: usize = 16;
/// The model formats Holger treats as first-class.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ModelFormat {
Gguf,
Safetensors,
Onnx,
}
impl ModelFormat {
/// The stable label (`gguf`, `safetensors`, `onnx`).
pub fn as_str(self) -> &'static str {
match self {
ModelFormat::Gguf => "gguf",
ModelFormat::Safetensors => "safetensors",
ModelFormat::Onnx => "onnx",
}
}
/// The format a file name claims, by its last extension.
pub fn claimed_by(file_name: &str) -> Option<Self> {
match extension(file_name).as_deref() {
Some("gguf") => Some(ModelFormat::Gguf),
Some("safetensors") => Some(ModelFormat::Safetensors),
Some("onnx") => Some(ModelFormat::Onnx),
_ => None,
}
}
}
/// What a model file says about itself.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct ModelMetadata {
pub format: Option<ModelFormat>,
/// `general.architecture` (GGUF), `architecture` in the metadata (Safetensors,
/// ONNX), else the ONNX graph name.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub architecture: Option<String>,
/// The parameter count: the file's own `general.parameter_count` when it has
/// one, else the sum of the tensors' element counts.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parameters: Option<u64>,
/// How many tensors (ONNX: initializers) the file declares.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tensors: Option<u64>,
/// The tensor dtypes, sorted (`F16`, `Q4_K`, `FLOAT`).
pub dtypes: Vec<String>,
/// The quantization: GGUF's `general.file_type`, a Safetensors metadata
/// field or its FP8/FP4 dtypes, ONNX's QDQ/QLinear operators.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub quantization: Option<String>,
/// Format-specific facts: GGUF version, ONNX IR version, opset, producer.
pub extra: BTreeMap<String, String>,
}
impl ModelMetadata {
fn of(format: ModelFormat) -> Self {
Self {
format: Some(format),
..Self::default()
}
}
/// The `model.*` properties this metadata becomes. Keys are fixed; a field
/// the file does not have is left out.
pub fn properties(&self) -> Vec<(String, Vec<String>)> {
let mut out = Vec::new();
if let Some(f) = self.format {
out.push(("model.format".to_string(), vec![f.as_str().to_string()]));
}
if let Some(a) = &self.architecture {
out.push(("model.architecture".to_string(), vec![a.clone()]));
}
if let Some(p) = self.parameters {
out.push(("model.parameters".to_string(), vec![p.to_string()]));
}
if let Some(t) = self.tensors {
out.push(("model.tensors".to_string(), vec![t.to_string()]));
}
if !self.dtypes.is_empty() {
out.push(("model.dtype".to_string(), self.dtypes.clone()));
}
if let Some(q) = &self.quantization {
out.push(("model.quantization".to_string(), vec![q.clone()]));
}
out
}
}
/// What one pickle stream does on load: its protocol and the globals it imports.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct PickleInfo {
pub protocol: u8,
/// `module.name` for every GLOBAL, INST and STACK_GLOBAL, first seen first.
pub globals: Vec<String>,
}
/// Why a header was refused, with the code that names the class.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HeaderFault {
pub code: &'static ec::ErrCode,
pub why: String,
}
impl HeaderFault {
fn malformed(why: impl Into<String>) -> Self {
Self {
code: &ec::MODEL_HEADER_MALFORMED,
why: why.into(),
}
}
fn out_of_bounds(why: impl Into<String>) -> Self {
Self {
code: &ec::MODEL_TENSOR_OUT_OF_BOUNDS,
why: why.into(),
}
}
fn overlap(why: impl Into<String>) -> Self {
Self {
code: &ec::MODEL_TENSOR_OVERLAP,
why: why.into(),
}
}
}
/// What the bytes are, by their content alone.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Sniffed {
/// A GGUF, Safetensors or ONNX file whose header passed every check.
Model(ModelMetadata),
/// A GGUF or Safetensors file (by its magic or header shape) whose header failed a check.
Malformed {
format: ModelFormat,
fault: HeaderFault,
},
/// A pickle stream.
Pickle(PickleInfo),
/// An archive in the PyTorch checkpoint layout: it holds `data.pkl`, at any
/// nesting S3's walker opens. `imports` names what that pickle imports, when
/// it could be read.
TorchZip {
member: String,
imports: Option<String>,
},
/// None of the above.
NotAModel,
}
/// ★ Tell what the bytes are.
pub fn sniff(bytes: &[u8]) -> Sniffed {
if bytes.starts_with(b"GGUF") {
return match gguf::parse(bytes) {
Ok(m) => Sniffed::Model(m),
Err(fault) => Sniffed::Malformed {
format: ModelFormat::Gguf,
fault,
},
};
}
if safetensors::looks_like(bytes) {
return match safetensors::parse(bytes) {
Ok(m) => Sniffed::Model(m),
Err(fault) => Sniffed::Malformed {
format: ModelFormat::Safetensors,
fault,
},
};
}
if holger_inspect::magic::detect(bytes).is_unpacked() {
if let Some((member, imports)) = torch_zip(bytes) {
return Sniffed::TorchZip { member, imports };
}
return Sniffed::NotAModel;
}
if let Some(p) = pickle::walk(bytes) {
return Sniffed::Pickle(p);
}
if let Some(m) = onnx::parse(bytes) {
return Sniffed::Model(m);
}
Sniffed::NotAModel
}
/// The probe S3's archive walker runs on every member: a member named
/// `data.pkl` is reported with what its pickle imports.
struct DataPklProbe;
/// The probe's name in the walker's findings.
const DATA_PKL_PROBE: &str = "holger-model-data-pkl";
impl holger_inspect::malware::MalwareScanner for DataPklProbe {
fn name(&self) -> &str {
DATA_PKL_PROBE
}
fn required(&self) -> bool {
false
}
fn reads_members(&self) -> bool {
true
}
fn scan(&self, bytes: &[u8], name: &str) -> holger_inspect::malware::ScanOutcome {
use holger_inspect::malware::ScanOutcome;
if name != "data.pkl" && !name.ends_with("/data.pkl") {
return ScanOutcome::Clean;
}
ScanOutcome::Detected(vec![match pickle::walk(bytes) {
Some(p) => describe_globals(&p),
None => "its pickle could not be read".to_string(),
}])
}
}
/// The `data.pkl` member of a PyTorch checkpoint archive, found by S3's archive
/// walker (zip, tar and gzip, nested, under its bomb and traversal limits), and
/// what it imports.
fn torch_zip(bytes: &[u8]) -> Option<(String, Option<String>)> {
let mut limits = holger_inspect::Limits::default();
limits.max_scan_member_bytes = limits.max_scan_member_bytes.max(MAX_DATA_PKL_BYTES as u64);
let found = holger_inspect::Inspector::new(limits)
.with_scanner(std::sync::Arc::new(DataPklProbe))
.inspect(bytes, holger_inspect::Format::Generic, "checkpoint");
let hit = found.findings.iter().find(|f| {
f.kind == holger_inspect::FindingKind::Malware && f.detail.starts_with(DATA_PKL_PROBE)
})?;
let member = hit
.path
.split_once('!')
.map(|(_, m)| m)
.unwrap_or(&hit.path)
.to_string();
let imports = hit
.detail
.strip_prefix(DATA_PKL_PROBE)
.map(|d| d.trim_start_matches(": ").to_string());
Some((member, imports))
}
/// The lowercased last extension, skipping one compression suffix
/// (`model.pkl.gz` → `pkl`).
fn extension(file_name: &str) -> Option<String> {
let base = file_name
.rsplit(['/', '\\'])
.next()
.unwrap_or(file_name)
.to_ascii_lowercase();
let mut parts: Vec<&str> = base.split('.').skip(1).collect();
if parts.len() > 1
&& matches!(
*parts.last()?,
"gz" | "bz2" | "xz" | "lz4" | "zst" | "z" | "lzma"
)
{
parts.pop();
}
parts.last().map(|s| s.to_string())
}
/// The pickle-family extensions, blocked by name even when the content is not
/// recognized (a compressed joblib dump, a legacy tar checkpoint).
pub const PICKLE_EXTENSIONS: [&str; 5] = ["pkl", "pickle", "pt", "pth", "joblib"];
/// Does the name carry a pickle-family extension?
pub fn has_pickle_name(file_name: &str) -> bool {
extension(file_name).is_some_and(|e| PICKLE_EXTENSIONS.contains(&e.as_str()))
}
/// Refuse a `models.allow_pickle` list that names a repository the config does
/// not declare: an allowance for a name nobody serves is a typo that would
/// silently allow nothing, or allow a repository created later.
pub fn validate_allow_pickle<'a>(
allow: &[String],
declared: impl IntoIterator<Item = &'a str>,
) -> anyhow::Result<()> {
let declared: BTreeSet<&str> = declared.into_iter().collect();
for r in allow {
if !declared.contains(r.as_str()) {
anyhow::bail!(
"{}",
ec::MODEL_ALLOW_UNKNOWN_REPO.tag(format!(
"config `models.allow_pickle` names the repository '{r}', which the config does not declare"
))
);
}
}
Ok(())
}
/// The `models:` block of the config.
///
/// ```ron
/// models: (allow_pickle: ["ml-research-incoming"]),
/// ```
///
/// Absent ⇒ every pickle-family file is blocked in every repository.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelSettings {
/// The repositories where a pickle-family file is a warning instead of a
/// block. Each must be a repository the config declares.
#[serde(default)]
pub allow_pickle: Vec<String>,
}
/// ★ The model scanner. Fills the [`ScannerKind::Model`] slot.
#[derive(Debug, Clone, Default)]
pub struct ModelScanner {
allow_pickle: BTreeSet<String>,
}
impl ModelScanner {
/// The default: the pickle family is blocked everywhere.
pub fn new() -> Self {
Self::default()
}
/// Allow the pickle family, as a warning, in these repositories only.
pub fn allow_pickle_in<I, S>(mut self, repos: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.allow_pickle.extend(repos.into_iter().map(Into::into));
self
}
fn pickle_allowed(&self, subject: &ScanSubject<'_>) -> bool {
subject.repo.is_some_and(|r| self.allow_pickle.contains(r))
}
}
fn describe_globals(p: &PickleInfo) -> String {
if p.globals.is_empty() {
return "it imports nothing".to_string();
}
let mut shown: Vec<&str> = p
.globals
.iter()
.take(MAX_GLOBALS_LISTED)
.map(String::as_str)
.collect();
if p.globals.len() > MAX_GLOBALS_LISTED {
shown.push("…");
}
format!("it imports {}", shown.join(", "))
}
impl Scanner for ModelScanner {
fn id(&self) -> &str {
MODEL_SCANNER_ID
}
fn version(&self) -> String {
env!("CARGO_PKG_VERSION").to_string()
}
fn kind(&self) -> ScannerKind {
ScannerKind::Model
}
fn scan(&self, subject: &ScanSubject<'_>) -> anyhow::Result<ScanOutcome> {
let name = subject.file_name;
let sniffed = sniff(subject.bytes);
// The pickle-family reason, when there is one: (code, sentence).
let pickle_reason: Option<(&'static ec::ErrCode, String)> = match &sniffed {
Sniffed::Pickle(p) => Some((
&ec::MODEL_PICKLE_BLOCKED,
format!("'{name}' is a pickle (protocol {}) by its opcodes; loading it runs code, and {}", p.protocol, describe_globals(p)),
)),
Sniffed::TorchZip { member, imports } => Some((
&ec::MODEL_TORCH_ZIP_BLOCKED,
format!(
"'{name}' is a PyTorch zip checkpoint: it holds the pickle '{member}', and loading it runs code{}",
imports.as_ref().map(|i| format!("; {i}")).unwrap_or_default()
),
)),
_ if has_pickle_name(name) => Some((
&ec::MODEL_PICKLE_BLOCKED,
format!(
"'{name}' carries a pickle-family extension (.{}); such files run code when they are loaded",
extension(name).unwrap_or_default()
),
)),
_ => None,
};
let mut blocks: Vec<ScanFinding> = Vec::new();
let mut warns: Vec<ScanFinding> = Vec::new();
let mut sentences: Vec<String> = Vec::new();
if let Some((code, why)) = pickle_reason {
if self.pickle_allowed(subject) {
let s = ec::MODEL_PICKLE_ALLOWED.tag(format!(
"{why}; `models.allow_pickle` allows the pickle family in '{}', so this is a warning ({})",
subject.repo.unwrap_or_default(),
code.code
));
warns.push(ScanFinding::new(
ec::MODEL_PICKLE_ALLOWED.code,
FindingSeverity::High,
s.clone(),
));
sentences.push(s);
} else {
let s = code.tag(format!("{why}; the pickle family is blocked by default (`models.allow_pickle` names no such repository)"));
blocks.push(ScanFinding::new(
code.code,
FindingSeverity::Critical,
s.clone(),
));
sentences.push(s);
}
}
if let Sniffed::Malformed { format, fault } = &sniffed {
let s = fault.code.tag(format!(
"'{name}' is a {} file whose header fails a check: {}",
format.as_str(),
fault.why
));
blocks.push(ScanFinding::new(
fault.code.code,
FindingSeverity::High,
s.clone(),
));
sentences.push(s);
}
let actual = match &sniffed {
Sniffed::Model(m) => m.format,
Sniffed::Malformed { format, .. } => Some(*format),
_ => None,
};
if let Some(claimed) = ModelFormat::claimed_by(name) {
if actual != Some(claimed) {
let is = match &sniffed {
Sniffed::Pickle(_) => "a pickle".to_string(),
Sniffed::TorchZip { .. } => "a PyTorch zip checkpoint".to_string(),
_ => actual
.map(|f| format!("a {} file", f.as_str()))
.unwrap_or_else(|| "no model format Holger knows".to_string()),
};
let s = ec::MODEL_FORMAT_MISMATCH.tag(format!(
"'{name}' is named as {} but its content is {is}",
claimed.as_str()
));
blocks.push(ScanFinding::new(
ec::MODEL_FORMAT_MISMATCH.code,
FindingSeverity::High,
s.clone(),
));
sentences.push(s);
}
}
if !blocks.is_empty() {
blocks.extend(warns);
return Ok(ScanOutcome::block(blocks, sentences.join(" ")));
}
let document = match &sniffed {
Sniffed::Model(m) => Some(serde_json::to_value(m)?),
_ => None,
};
let outcome = if !warns.is_empty() {
ScanOutcome::warn(warns, sentences.join(" "))
} else if let Sniffed::Model(m) = &sniffed {
ScanOutcome::pass(format!(
"'{name}' is a well-formed {} model ({} tensors{})",
m.format.map(|f| f.as_str()).unwrap_or("?"),
m.tensors
.map(|t| t.to_string())
.unwrap_or_else(|| "?".into()),
m.architecture
.as_ref()
.map(|a| format!(", architecture {a}"))
.unwrap_or_default()
))
} else {
ScanOutcome::pass(format!("'{name}' is not a model file"))
};
Ok(match document {
Some(d) => outcome.with_document(d),
None => outcome,
})
}
}
/// The metadata of an artifact's scan evidence, when its model record carries one.
pub fn metadata_of(evidence: &super::Evidence) -> Option<ModelMetadata> {
evidence
.records_of(ScannerKind::Model)
.find_map(|r| r.document.as_ref())
.and_then(|d| serde_json::from_value(d.clone()).ok())
}
/// A bounds-checked little-endian reader over untrusted bytes.
struct Cursor<'a> {
b: &'a [u8],
pos: usize,
}
impl<'a> Cursor<'a> {
fn new(b: &'a [u8], pos: usize) -> Self {
Self { b, pos }
}
fn left(&self) -> usize {
self.b.len().saturating_sub(self.pos)
}
fn take(&mut self, n: usize) -> Result<&'a [u8], HeaderFault> {
if n > self.left() {
return Err(HeaderFault::malformed(format!(
"the header needs {n} more bytes at offset {} and the file has {}",
self.pos,
self.left()
)));
}
let s = &self.b[self.pos..self.pos + n];
self.pos += n;
Ok(s)
}
fn u32(&mut self) -> Result<u32, HeaderFault> {
Ok(u32::from_le_bytes(
self.take(4)?.try_into().expect("4 bytes"),
))
}
fn u64(&mut self) -> Result<u64, HeaderFault> {
Ok(u64::from_le_bytes(
self.take(8)?.try_into().expect("8 bytes"),
))
}
}
fn checked_product(dims: &[u64]) -> Option<u64> {
dims.iter().try_fold(1u64, |acc, d| acc.checked_mul(*d))
}
/// One tensor's byte range, for the bounds and overlap checks.
struct Range {
name: String,
start: u64,
end: u64,
}
/// Sort by start and refuse any two ranges that share a byte.
fn refuse_overlap(ranges: &mut [Range]) -> Result<(), HeaderFault> {
ranges.sort_by_key(|r| (r.start, r.end));
for w in ranges.windows(2) {
let (a, b) = (&w[0], &w[1]);
if b.start < a.end {
return Err(HeaderFault::overlap(format!(
"tensor '{}' [{}, {}) overlaps tensor '{}' [{}, {})",
b.name, b.start, b.end, a.name, a.start, a.end
)));
}
}
Ok(())
}
mod gguf {
use super::*;
/// Bytes per block and elements per block of each GGML tensor type.
fn ggml_type(t: u32) -> Option<(&'static str, u64, u64)> {
Some(match t {
0 => ("F32", 4, 1),
1 => ("F16", 2, 1),
2 => ("Q4_0", 18, 32),
3 => ("Q4_1", 20, 32),
6 => ("Q5_0", 22, 32),
7 => ("Q5_1", 24, 32),
8 => ("Q8_0", 34, 32),
9 => ("Q8_1", 36, 32),
10 => ("Q2_K", 84, 256),
11 => ("Q3_K", 110, 256),
12 => ("Q4_K", 144, 256),
13 => ("Q5_K", 176, 256),
14 => ("Q6_K", 210, 256),
15 => ("Q8_K", 292, 256),
16 => ("IQ2_XXS", 66, 256),
17 => ("IQ2_XS", 74, 256),
18 => ("IQ3_XXS", 98, 256),
19 => ("IQ1_S", 50, 256),
20 => ("IQ4_NL", 18, 32),
21 => ("IQ3_S", 110, 256),
22 => ("IQ2_S", 82, 256),
23 => ("IQ4_XS", 136, 256),
24 => ("I8", 1, 1),
25 => ("I16", 2, 1),
26 => ("I32", 4, 1),
27 => ("I64", 8, 1),
28 => ("F64", 8, 1),
29 => ("IQ1_M", 56, 256),
30 => ("BF16", 2, 1),
34 => ("TQ1_0", 54, 256),
35 => ("TQ2_0", 66, 256),
39 => ("MXFP4", 17, 32),
_ => return None,
})
}
/// llama.cpp's `general.file_type` names.
fn file_type(t: u64) -> String {
match t {
0 => "ALL_F32",
1 => "MOSTLY_F16",
2 => "MOSTLY_Q4_0",
3 => "MOSTLY_Q4_1",
7 => "MOSTLY_Q8_0",
8 => "MOSTLY_Q5_0",
9 => "MOSTLY_Q5_1",
10 => "MOSTLY_Q2_K",
11 => "MOSTLY_Q3_K_S",
12 => "MOSTLY_Q3_K_M",
13 => "MOSTLY_Q3_K_L",
14 => "MOSTLY_Q4_K_S",
15 => "MOSTLY_Q4_K_M",
16 => "MOSTLY_Q5_K_S",
17 => "MOSTLY_Q5_K_M",
18 => "MOSTLY_Q6_K",
19 => "MOSTLY_IQ2_XXS",
20 => "MOSTLY_IQ2_XS",
21 => "MOSTLY_Q2_K_S",
22 => "MOSTLY_IQ3_XS",
23 => "MOSTLY_IQ3_XXS",
24 => "MOSTLY_IQ1_S",
25 => "MOSTLY_IQ4_NL",
26 => "MOSTLY_IQ3_S",
27 => "MOSTLY_IQ3_M",
28 => "MOSTLY_IQ2_S",
29 => "MOSTLY_IQ2_M",
30 => "MOSTLY_IQ4_XS",
31 => "MOSTLY_IQ1_M",
32 => "MOSTLY_BF16",
36 => "MOSTLY_TQ1_0",
37 => "MOSTLY_TQ2_0",
38 => "MOSTLY_MXFP4_MOE",
other => return format!("FILE_TYPE_{other}"),
}
.to_string()
}
/// The scalar values the metadata extraction needs; arrays are walked and dropped.
enum Val {
Uint(u64),
Int(i64),
Str(String),
Other,
}
fn string(c: &mut Cursor<'_>, what: &str) -> Result<String, HeaderFault> {
let n = c.u64()?;
if n > MAX_GGUF_STRING {
return Err(HeaderFault::malformed(format!(
"{what} is {n} bytes long, over the {MAX_GGUF_STRING}-byte limit"
)));
}
let bytes = c.take(n as usize)?;
Ok(String::from_utf8_lossy(bytes).into_owned())
}
fn fixed_size(t: u32) -> Option<usize> {
Some(match t {
0 | 1 | 7 => 1,
2 | 3 => 2,
4..=6 => 4,
10..=12 => 8,
_ => return None,
})
}
fn take(c: &mut Cursor<'_>, n: usize) -> Result<u64, HeaderFault> {
let b = c.take(n)?;
let mut buf = [0u8; 8];
buf[..n].copy_from_slice(b);
Ok(u64::from_le_bytes(buf))
}
fn value(c: &mut Cursor<'_>, t: u32, depth: u32) -> Result<Val, HeaderFault> {
Ok(match t {
0 => Val::Uint(take(c, 1)?),
1 => Val::Int(take(c, 1)? as u8 as i8 as i64),
2 => Val::Uint(take(c, 2)?),
3 => Val::Int(take(c, 2)? as u16 as i16 as i64),
4 => Val::Uint(take(c, 4)?),
5 => Val::Int(take(c, 4)? as u32 as i32 as i64),
10 => Val::Uint(take(c, 8)?),
11 => Val::Int(take(c, 8)? as i64),
6 | 7 | 12 => {
c.take(fixed_size(t).expect("fixed"))?;
Val::Other
}
8 => Val::Str(string(c, "a string value")?),
9 => {
if depth >= MAX_GGUF_ARRAY_DEPTH {
return Err(HeaderFault::malformed(format!(
"arrays nest deeper than {MAX_GGUF_ARRAY_DEPTH}"
)));
}
let et = c.u32()?;
let n = c.u64()?;
match fixed_size(et) {
Some(sz) => {
let total = n
.checked_mul(sz as u64)
.filter(|t| *t <= c.left() as u64)
.ok_or_else(|| {
HeaderFault::malformed(format!(
"an array declares {n} elements, more than the file holds"
))
})?;
c.take(total as usize)?;
}
None => {
// Every element takes at least one byte, so a count past the
// bytes left is a lie told before any is read.
if n > c.left() as u64 {
return Err(HeaderFault::malformed(format!(
"an array declares {n} elements, more than the file holds"
)));
}
for _ in 0..n {
value(c, et, depth + 1)?;
}
}
}
Val::Other
}
other => {
return Err(HeaderFault::malformed(format!(
"unknown metadata value type {other}"
)))
}
})
}
pub(super) fn parse(b: &[u8]) -> Result<ModelMetadata, HeaderFault> {
let mut c = Cursor::new(b, 4);
let version = c.u32()?;
if !(2..=3).contains(&version) {
return Err(HeaderFault::malformed(format!(
"GGUF version {version} is not 2 or 3"
)));
}
let n_tensors = c.u64()?;
let n_kv = c.u64()?;
if n_tensors > MAX_TENSORS {
return Err(HeaderFault::malformed(format!(
"{n_tensors} tensors declared, over the limit of {MAX_TENSORS}"
)));
}
if n_kv > MAX_GGUF_KV {
return Err(HeaderFault::malformed(format!(
"{n_kv} metadata keys declared, over the limit of {MAX_GGUF_KV}"
)));
}
let mut m = ModelMetadata::of(ModelFormat::Gguf);
m.extra.insert("gguf_version".into(), version.to_string());
let mut alignment: u64 = 32;
let mut declared_params: Option<u64> = None;
let mut keys = BTreeSet::new();
for _ in 0..n_kv {
let key = string(&mut c, "a metadata key")?;
let t = c.u32()?;
let v = value(&mut c, t, 0)?;
if !keys.insert(key.clone()) {
return Err(HeaderFault::malformed(format!(
"the metadata key '{key}' appears twice"
)));
}
match (key.as_str(), v) {
("general.architecture", Val::Str(s)) => m.architecture = Some(s),
("general.name", Val::Str(s)) => {
m.extra.insert("name".into(), s);
}
("general.parameter_count", Val::Uint(n)) => declared_params = Some(n),
("general.parameter_count", Val::Int(n)) if n >= 0 => {
declared_params = Some(n as u64)
}
("general.file_type", Val::Uint(n)) => m.quantization = Some(file_type(n)),
("general.file_type", Val::Int(n)) if n >= 0 => {
m.quantization = Some(file_type(n as u64))
}
("general.alignment", Val::Uint(n)) => {
if n == 0 || !n.is_power_of_two() || n > 1 << 20 {
return Err(HeaderFault::malformed(format!(
"general.alignment {n} is not a power of two up to 1 MiB"
)));
}
alignment = n;
}
("general.alignment", _) => {
return Err(HeaderFault::malformed(
"general.alignment is not an unsigned integer",
));
}
_ => {}
}
}
struct Info {
name: String,
offset: u64,
size: u64,
}
let mut infos = Vec::new();
let mut names = BTreeSet::new();
let mut dtypes = BTreeSet::new();
let mut elements: u64 = 0;
let mut counts: BTreeMap<&'static str, u64> = BTreeMap::new();
for _ in 0..n_tensors {
let name = string(&mut c, "a tensor name")?;
let n_dims = c.u32()?;
if n_dims > MAX_TENSOR_DIMS {
return Err(HeaderFault::malformed(format!(
"tensor '{name}' has {n_dims} dimensions, over GGML's {MAX_TENSOR_DIMS}"
)));
}
let mut dims = Vec::with_capacity(n_dims as usize);
for _ in 0..n_dims {
dims.push(c.u64()?);
}
let t = c.u32()?;
let offset = c.u64()?;
let (tname, block_bytes, block_elems) = ggml_type(t).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}' has the unknown GGML type {t}"))
})?;
let n = checked_product(&dims).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}''s element count overflows"))
})?;
if dims.first().copied().unwrap_or(1) % block_elems != 0 {
return Err(HeaderFault::malformed(format!(
"tensor '{name}''s first dimension is not a multiple of the {tname} block size {block_elems}"
)));
}
let size = (n / block_elems).checked_mul(block_bytes).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}''s byte size overflows"))
})?;
if offset % alignment != 0 {
return Err(HeaderFault::malformed(format!(
"tensor '{name}''s offset {offset} is not aligned to {alignment}"
)));
}
if !names.insert(name.clone()) {
return Err(HeaderFault::malformed(format!(
"the tensor name '{name}' appears twice"
)));
}
dtypes.insert(tname);
*counts.entry(tname).or_default() += 1;
elements = elements.saturating_add(n);
infos.push(Info { name, offset, size });
}
let data_start = (c.pos as u64).div_ceil(alignment) * alignment;
let len = b.len() as u64;
if data_start > len && n_tensors > 0 {
return Err(HeaderFault::out_of_bounds(format!(
"the tensor data would start at {data_start}, past the end of the {len}-byte file"
)));
}
let mut ranges = Vec::with_capacity(infos.len());
for i in infos {
let start = data_start.checked_add(i.offset);
let end = start.and_then(|s| s.checked_add(i.size));
match (start, end) {
(Some(s), Some(e)) if e <= len => ranges.push(Range { name: i.name, start: s, end: e }),
_ => {
return Err(HeaderFault::out_of_bounds(format!(
"tensor '{}' at offset {} with {} bytes ends past the end of the {len}-byte file",
i.name, i.offset, i.size
)))
}
}
}
refuse_overlap(&mut ranges)?;
m.tensors = Some(n_tensors);
m.parameters = declared_params.or(Some(elements));
m.dtypes = dtypes.into_iter().map(str::to_string).collect();
if m.quantization.is_none() {
// The type most tensors use, F32 aside (norms stay F32 in every quant).
m.quantization = counts
.iter()
.filter(|(t, _)| **t != "F32")
.max_by_key(|(_, n)| **n)
.map(|(t, _)| t.to_string());
}
Ok(m)
}
}
mod safetensors {
use super::*;
/// Bits per element of each Safetensors dtype.
fn dtype_bits(d: &str) -> Option<u64> {
Some(match d {
"BOOL" | "U8" | "I8" | "F8_E5M2" | "F8_E4M3" | "F8_E8M0" => 8,
"I16" | "U16" | "F16" | "BF16" => 16,
"I32" | "U32" | "F32" => 32,
"I64" | "U64" | "F64" => 64,
"F4" => 4,
"F6_E2M3" | "F6_E3M2" => 6,
_ => return None,
})
}
/// The shape of a Safetensors file: a plausible header length followed by `{`.
pub(super) fn looks_like(b: &[u8]) -> bool {
b.len() >= 10
&& b[8] == b'{'
&& u64::from_le_bytes(b[..8].try_into().expect("8 bytes")) >= 2
}
pub(super) fn parse(b: &[u8]) -> Result<ModelMetadata, HeaderFault> {
let n = u64::from_le_bytes(b[..8].try_into().expect("8 bytes"));
if n > MAX_SAFETENSORS_HEADER {
return Err(HeaderFault::malformed(format!(
"the header length {n} is over the {MAX_SAFETENSORS_HEADER}-byte limit"
)));
}
let len = b.len() as u64;
if 8 + n > len {
return Err(HeaderFault::malformed(format!(
"the header length {n} runs past the end of the {len}-byte file"
)));
}
let header = &b[8..8 + n as usize];
let text = std::str::from_utf8(header)
.map_err(|_| HeaderFault::malformed("the header is not UTF-8"))?;
let json: serde_json::Value = serde_json::from_str(text)
.map_err(|e| HeaderFault::malformed(format!("the header is not JSON: {e}")))?;
let obj = json
.as_object()
.ok_or_else(|| HeaderFault::malformed("the header is not a JSON object"))?;
let data_len = len - 8 - n;
let mut m = ModelMetadata::of(ModelFormat::Safetensors);
let mut ranges = Vec::new();
let mut dtypes = BTreeSet::new();
let mut elements: u64 = 0;
if obj.len() as u64 > MAX_TENSORS + 1 {
return Err(HeaderFault::malformed(format!(
"{} tensors declared, over the limit of {MAX_TENSORS}",
obj.len()
)));
}
for (name, v) in obj {
if name == "__metadata__" {
let meta = v
.as_object()
.ok_or_else(|| HeaderFault::malformed("__metadata__ is not an object"))?;
for (k, v) in meta {
let s = v.as_str().ok_or_else(|| {
HeaderFault::malformed(format!("__metadata__ value '{k}' is not a string"))
})?;
match k.as_str() {
"architecture" | "general.architecture" | "model_type" => {
m.architecture = Some(s.to_string())
}
"quantization" | "quant_method" => m.quantization = Some(s.to_string()),
"format" => {
m.extra.insert("format".into(), s.to_string());
}
_ => {}
}
}
continue;
}
let t = v.as_object().ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}' is not an object"))
})?;
let dtype = t
.get("dtype")
.and_then(|d| d.as_str())
.ok_or_else(|| HeaderFault::malformed(format!("tensor '{name}' has no dtype")))?;
let bits = dtype_bits(dtype).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}' has the unknown dtype '{dtype}'"))
})?;
let shape: Vec<u64> = t
.get("shape")
.and_then(|s| s.as_array())
.and_then(|a| a.iter().map(|d| d.as_u64()).collect::<Option<Vec<_>>>())
.ok_or_else(|| {
HeaderFault::malformed(format!(
"tensor '{name}' has no shape of unsigned integers"
))
})?;
let offsets: Vec<u64> = t
.get("data_offsets")
.and_then(|s| s.as_array())
.and_then(|a| a.iter().map(|d| d.as_u64()).collect::<Option<Vec<_>>>())
.filter(|o| o.len() == 2)
.ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}' has no data_offsets pair"))
})?;
let (start, end) = (offsets[0], offsets[1]);
if start > end {
return Err(HeaderFault::malformed(format!(
"tensor '{name}' ends at {end}, before it starts at {start}"
)));
}
let n_el = checked_product(&shape).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}''s element count overflows"))
})?;
let nbits = n_el.checked_mul(bits).ok_or_else(|| {
HeaderFault::malformed(format!("tensor '{name}''s byte size overflows"))
})?;
if nbits % 8 != 0 || nbits / 8 != end - start {
return Err(HeaderFault::malformed(format!(
"tensor '{name}' ({dtype}, shape {shape:?}) needs {} bytes and its offsets give {}",
nbits.div_ceil(8),
end - start
)));
}
if end > data_len {
return Err(HeaderFault::out_of_bounds(format!(
"tensor '{name}' [{start}, {end}) ends past the {data_len} bytes of tensor data the file holds"
)));
}
dtypes.insert(dtype.to_string());
elements = elements.saturating_add(n_el);
ranges.push(Range {
name: name.clone(),
start,
end,
});
}
refuse_overlap(&mut ranges)?;
// The data must be exactly the tensors, back to back: a hole or a tail is
// bytes nothing describes, which is where a payload would hide.
let mut at = 0u64;
for r in &ranges {
if r.start != at {
return Err(HeaderFault::malformed(format!(
"{} undescribed bytes before tensor '{}' at {}",
r.start - at,
r.name,
r.start
)));
}
at = r.end;
}
if at != data_len {
return Err(HeaderFault::malformed(format!(
"{} undescribed bytes after the last tensor",
data_len - at
)));
}
m.tensors = Some(ranges.len() as u64);
m.parameters = Some(elements);
if m.quantization.is_none() {
if dtypes.iter().any(|d| d.starts_with("F8_")) {
m.quantization = Some("fp8".into());
} else if dtypes.iter().any(|d| d == "F4" || d.starts_with("F6_")) {
m.quantization = Some("fp4/fp6".into());
}
}
m.dtypes = dtypes.into_iter().collect();
Ok(m)
}
}
mod pickle {
use super::*;
/// How an opcode's argument is laid out.
enum Arg {
None,
Fixed(usize),
/// A length prefix of this many bytes, then that many bytes.
Len(usize),
/// A newline-terminated text line, checked by kind.
Line(LineKind),
/// Two newline-terminated lines: a module and a name.
TwoLines,
}
#[derive(Clone, Copy)]
enum LineKind {
Int,
Float,
Long,
Quoted,
Any,
Digits,
}
fn arg(op: u8) -> Option<Arg> {
use Arg::*;
use LineKind::*;
Some(match op {
b'(' | b'.' | b'0' | b'1' | b'2' | b'N' | b'Q' | b'R' | b'a' | b'b' | b'd' | b'}'
| b'e' | b'l' | b']' | b'o' | b's' | b't' | b')' | b'u' => None,
b'F' => Line(Float),
b'I' => Line(Int),
b'L' => Line(Long),
b'S' => Line(Quoted),
b'V' | b'P' => Line(Any),
b'g' | b'p' => Line(Digits),
b'c' | b'i' => TwoLines,
b'J' => Fixed(4),
b'K' | b'h' | b'q' => Fixed(1),
b'M' => Fixed(2),
b'j' | b'r' => Fixed(4),
b'G' => Fixed(8),
b'T' | b'X' | b'B' => Len(4),
b'U' | b'C' => Len(1),
// protocol 2
0x80 => Fixed(1),
0x81 | 0x85 | 0x86 | 0x87 | 0x88 | 0x89 => None,
0x82 => Fixed(1),
0x83 => Fixed(2),
0x84 => Fixed(4),
0x8a => Len(1),
0x8b => Len(4),
// protocol 4
0x8c => Len(1),
0x8d | 0x8e => Len(8),
0x8f..=0x94 => None,
0x95 => Fixed(8),
// protocol 5
0x96 => Len(8),
0x97 | 0x98 => None,
_ => return Option::None,
})
}
const MAX_LINE: usize = 64 * 1024;
fn line(b: &[u8], pos: usize) -> Option<(&[u8], usize)> {
let rest = b.get(pos..)?;
let n = rest.iter().take(MAX_LINE).position(|c| *c == b'\n')?;
Some((&rest[..n], pos + n + 1))
}
fn line_ok(l: &[u8], k: LineKind) -> bool {
let s = match std::str::from_utf8(l) {
Ok(s) => s,
Err(_) => return matches!(k, LineKind::Any),
};
match k {
LineKind::Int => {
s == "00"
|| s == "01"
|| s.trim_start_matches('-')
.bytes()
.all(|c| c.is_ascii_digit())
&& !s.trim_start_matches('-').is_empty()
}
LineKind::Long => {
let t = s.trim_end_matches('L').trim_start_matches('-');
!t.is_empty() && t.bytes().all(|c| c.is_ascii_digit())
}
LineKind::Float => s.parse::<f64>().is_ok(),
LineKind::Quoted => {
(s.starts_with('\'') && s.ends_with('\'') || s.starts_with('"') && s.ends_with('"'))
&& s.len() >= 2
}
LineKind::Digits => !s.is_empty() && s.bytes().all(|c| c.is_ascii_digit()),
LineKind::Any => true,
}
}
fn ident_ok(l: &[u8]) -> bool {
!l.is_empty() && l.iter().all(|c| c.is_ascii_graphic())
}
/// Walk the bytes as one pickle stream. `Some` when every opcode from the
/// first byte is valid and the stream reaches `STOP`; trailing bytes are
/// allowed (a legacy PyTorch file is several pickles then raw storage).
pub(super) fn walk(b: &[u8]) -> Option<PickleInfo> {
let mut pos = 0usize;
let mut ops = 0usize;
let mut info = PickleInfo::default();
// The last two strings pushed, for STACK_GLOBAL.
let mut recent: Vec<String> = Vec::new();
let push_global = |info: &mut PickleInfo, g: String| {
if !info.globals.contains(&g) {
info.globals.push(g);
}
};
loop {
let op = *b.get(pos)?;
pos += 1;
if ops == 0 && op == 0x80 {
let proto = *b.get(pos)?;
if !(2..=5).contains(&proto) {
return None;
}
info.protocol = proto;
} else if op == 0x80 {
// PROTO only opens a stream.
return None;
}
ops += 1;
if op == b'.' {
return (ops >= 2).then_some(info);
}
match arg(op)? {
Arg::None => {
if op == 0x93 {
if recent.len() >= 2 {
let name = recent.pop().expect("two");
let module = recent.pop().expect("two");
push_global(&mut info, format!("{module}.{name}"));
} else {
push_global(&mut info, "<computed>".to_string());
}
}
}
Arg::Fixed(n) => {
b.get(pos..pos + n)?;
pos += n;
}
Arg::Len(w) => {
let lb = b.get(pos..pos + w)?;
let mut buf = [0u8; 8];
buf[..w].copy_from_slice(lb);
let n = usize::try_from(u64::from_le_bytes(buf)).ok()?;
pos += w;
let end = pos.checked_add(n)?;
let body = b.get(pos..end)?;
if matches!(op, b'X' | 0x8c | 0x8d) {
recent.push(String::from_utf8_lossy(body).into_owned());
if recent.len() > 2 {
recent.remove(0);
}
}
pos = end;
}
Arg::Line(k) => {
let (l, next) = line(b, pos)?;
if !line_ok(l, k) {
return None;
}
pos = next;
}
Arg::TwoLines => {
let (module, next) = line(b, pos)?;
let (name, next) = line(b, next)?;
if !ident_ok(module) || !ident_ok(name) {
return None;
}
push_global(
&mut info,
format!(
"{}.{}",
String::from_utf8_lossy(module),
String::from_utf8_lossy(name)
),
);
pos = next;
}
}
}
}
}
mod onnx {
use super::*;
fn varint(b: &[u8], pos: &mut usize) -> Option<u64> {
let mut v: u64 = 0;
for i in 0..10 {
let byte = *b.get(*pos)?;
*pos += 1;
v |= u64::from(byte & 0x7f) << (7 * i);
if byte & 0x80 == 0 {
return Some(v);
}
}
None
}
/// A field's value: a varint, or the bytes of a length-delimited field.
enum F<'a> {
Varint(u64),
Bytes(&'a [u8]),
Fixed,
}
/// Every top-level field of one protobuf message, or `None` when the bytes are
/// not a well-formed message.
fn fields(b: &[u8]) -> Option<Vec<(u64, F<'_>)>> {
let mut pos = 0;
let mut out = Vec::new();
while pos < b.len() {
let key = varint(b, &mut pos)?;
let (field, wt) = (key >> 3, key & 7);
if field == 0 {
return None;
}
let v = match wt {
0 => F::Varint(varint(b, &mut pos)?),
1 => {
b.get(pos..pos + 8)?;
pos += 8;
F::Fixed
}
2 => {
let n = usize::try_from(varint(b, &mut pos)?).ok()?;
let end = pos.checked_add(n)?;
let s = b.get(pos..end)?;
pos = end;
F::Bytes(s)
}
5 => {
b.get(pos..pos + 4)?;
pos += 4;
F::Fixed
}
_ => return None,
};
out.push((field, v));
}
Some(out)
}
fn text(b: &[u8]) -> String {
String::from_utf8_lossy(b).into_owned()
}
fn data_type(t: u64) -> String {
match t {
1 => "FLOAT",
2 => "UINT8",
3 => "INT8",
4 => "UINT16",
5 => "INT16",
6 => "INT32",
7 => "INT64",
8 => "STRING",
9 => "BOOL",
10 => "FLOAT16",
11 => "DOUBLE",
12 => "UINT32",
13 => "UINT64",
14 => "COMPLEX64",
15 => "COMPLEX128",
16 => "BFLOAT16",
17 => "FLOAT8E4M3FN",
18 => "FLOAT8E4M3FNUZ",
19 => "FLOAT8E5M2",
20 => "FLOAT8E5M2FNUZ",
21 => "UINT4",
22 => "INT4",
23 => "FLOAT4E2M1",
other => return format!("TYPE_{other}"),
}
.to_string()
}
/// ★ Parse an ONNX `ModelProto` header, best-effort. `None` when the bytes are
/// not one: the first field must be `ir_version` (field 1, varint, 1..=100),
/// every top-level field must be well formed, and there must be a `graph`.
pub(super) fn parse(b: &[u8]) -> Option<ModelMetadata> {
if b.first() != Some(&0x08) {
return None;
}
let top = fields(b)?;
let ir = match top.first() {
Some((1, F::Varint(v))) if (1..=100).contains(v) => *v,
_ => return None,
};
let graph = top.iter().find_map(|(f, v)| match (f, v) {
(7, F::Bytes(g)) => Some(*g),
_ => None,
})?;
let mut m = ModelMetadata::of(ModelFormat::Onnx);
m.extra.insert("ir_version".into(), ir.to_string());
let mut meta_arch = None;
for (f, v) in &top {
match (f, v) {
(2, F::Bytes(s)) => {
m.extra.insert("producer".into(), text(s));
}
(3, F::Bytes(s)) => {
m.extra.insert("producer_version".into(), text(s));
}
(8, F::Bytes(s)) => {
// OperatorSetIdProto: 1 domain, 2 version.
if let Some(op) = fields(s) {
let domain = op.iter().find_map(|(f, v)| match (f, v) {
(1, F::Bytes(d)) => Some(text(d)),
_ => None,
});
let version = op.iter().find_map(|(f, v)| match (f, v) {
(2, F::Varint(n)) => Some(*n),
_ => None,
});
if let (None | Some(""), Some(n)) = (domain.as_deref(), version) {
m.extra.insert("opset".into(), n.to_string());
}
}
}
(14, F::Bytes(s)) => {
// StringStringEntryProto: 1 key, 2 value.
if let Some(kv) = fields(s) {
let get = |want: u64| {
kv.iter().find_map(|(f, v)| match (f, v) {
(f, F::Bytes(x)) if *f == want => Some(text(x)),
_ => None,
})
};
if let (Some(k), Some(v)) = (get(1), get(2)) {
if k == "architecture" {
meta_arch = Some(v);
}
}
}
}
_ => {}
}
}
// The graph, best-effort: a graph that does not parse leaves the counts unknown.
let mut graph_name = None;
if let Some(g) = fields(graph) {
let mut tensors = 0u64;
let mut elements = 0u64;
let mut dtypes = BTreeSet::new();
let mut quantized = false;
for (f, v) in &g {
match (f, v) {
(2, F::Bytes(s)) => graph_name = Some(text(s)),
(1, F::Bytes(node)) => {
if let Some(nf) = fields(node) {
for (f, v) in nf {
if let (4, F::Bytes(op)) = (f, v) {
let op = text(op);
if op.contains("QuantizeLinear")
|| op.starts_with("QLinear")
|| op.ends_with("Integer")
{
quantized = true;
}
}
}
}
}
(5, F::Bytes(t)) => {
tensors += 1;
if let Some(tf) = fields(t) {
let mut dims = Vec::new();
for (f, v) in tf {
match (f, v) {
(1, F::Varint(d)) => dims.push(d),
(1, F::Bytes(packed)) => {
let mut p = 0;
while p < packed.len() {
match varint(packed, &mut p) {
Some(d) => dims.push(d),
None => break,
}
}
}
(2, F::Varint(dt)) => {
dtypes.insert(data_type(dt));
}
_ => {}
}
}
elements = elements.saturating_add(checked_product(&dims).unwrap_or(0));
}
}
_ => {}
}
}
m.tensors = Some(tensors);
m.parameters = Some(elements);
m.dtypes = dtypes.into_iter().collect();
if quantized {
m.quantization = Some("qdq".into());
}
}
m.architecture = meta_arch.or(graph_name);
Some(m)
}
}
/// Small, generated model files for tests. Benign by construction: the pickles
/// build plain dicts and import only `collections.OrderedDict`.
pub mod fixtures {
/// A GGUF v3 file: `general.architecture = "llama"`, `general.file_type = 15`
/// (Q4_K_M), and two tensors (an F32 norm, an F16 matrix).
pub fn gguf() -> Vec<u8> {
fn s(out: &mut Vec<u8>, v: &str) {
out.extend_from_slice(&(v.len() as u64).to_le_bytes());
out.extend_from_slice(v.as_bytes());
}
let mut b = b"GGUF".to_vec();
b.extend_from_slice(&3u32.to_le_bytes());
b.extend_from_slice(&2u64.to_le_bytes()); // tensors
b.extend_from_slice(&3u64.to_le_bytes()); // kv
s(&mut b, "general.architecture");
b.extend_from_slice(&8u32.to_le_bytes());
s(&mut b, "llama");
s(&mut b, "general.file_type");
b.extend_from_slice(&4u32.to_le_bytes());
b.extend_from_slice(&15u32.to_le_bytes());
s(&mut b, "tokenizer.ggml.tokens");
b.extend_from_slice(&9u32.to_le_bytes());
b.extend_from_slice(&8u32.to_le_bytes());
b.extend_from_slice(&2u64.to_le_bytes());
s(&mut b, "<s>");
s(&mut b, "</s>");
// tensor 0: norm, 1 dim [4], F32, offset 0 (16 bytes)
s(&mut b, "blk.0.attn_norm.weight");
b.extend_from_slice(&1u32.to_le_bytes());
b.extend_from_slice(&4u64.to_le_bytes());
b.extend_from_slice(&0u32.to_le_bytes());
b.extend_from_slice(&0u64.to_le_bytes());
// tensor 1: matrix [4, 2], F16, offset 32 (16 bytes)
s(&mut b, "blk.0.attn_q.weight");
b.extend_from_slice(&2u32.to_le_bytes());
b.extend_from_slice(&4u64.to_le_bytes());
b.extend_from_slice(&2u64.to_le_bytes());
b.extend_from_slice(&1u32.to_le_bytes());
b.extend_from_slice(&32u64.to_le_bytes());
while !b.len().is_multiple_of(32) {
b.push(0);
}
b.extend_from_slice(&[0u8; 48]);
b
}
/// A Safetensors file with `header` as its JSON header and `data` after it.
pub fn safetensors_with(header: &str, data: &[u8]) -> Vec<u8> {
let mut b = (header.len() as u64).to_le_bytes().to_vec();
b.extend_from_slice(header.as_bytes());
b.extend_from_slice(data);
b
}
/// A well-formed Safetensors file: two F16/F32 tensors and a `__metadata__`.
pub fn safetensors() -> Vec<u8> {
safetensors_with(
r#"{"__metadata__":{"format":"pt","architecture":"bert"},"embed":{"dtype":"F16","shape":[2,3],"data_offsets":[0,12]},"bias":{"dtype":"F32","shape":[3],"data_offsets":[12,24]}}"#,
&[0u8; 24],
)
}
/// A Safetensors file whose header claims tensor data past the end of the file.
pub fn safetensors_lying_offsets() -> Vec<u8> {
safetensors_with(
r#"{"w":{"dtype":"F32","shape":[1024],"data_offsets":[0,4096]}}"#,
&[0u8; 16],
)
}
/// A Safetensors file whose two tensors claim the same bytes.
pub fn safetensors_overlapping() -> Vec<u8> {
safetensors_with(
r#"{"a":{"dtype":"F32","shape":[4],"data_offsets":[0,16]},"b":{"dtype":"F32","shape":[4],"data_offsets":[8,24]}}"#,
&[0u8; 24],
)
}
fn varint(out: &mut Vec<u8>, mut v: u64) {
loop {
let byte = (v & 0x7f) as u8;
v >>= 7;
if v == 0 {
out.push(byte);
return;
}
out.push(byte | 0x80);
}
}
fn field_bytes(out: &mut Vec<u8>, field: u64, body: &[u8]) {
varint(out, field << 3 | 2);
varint(out, body.len() as u64);
out.extend_from_slice(body);
}
fn field_varint(out: &mut Vec<u8>, field: u64, v: u64) {
varint(out, field << 3);
varint(out, v);
}
/// A small ONNX ModelProto: ir_version 9, opset 17, producer `holger-test`,
/// a graph `tiny-mlp` with one MatMul node and one FLOAT initializer [3, 4].
pub fn onnx() -> Vec<u8> {
let mut node = Vec::new();
field_bytes(&mut node, 1, b"x");
field_bytes(&mut node, 1, b"W");
field_bytes(&mut node, 2, b"y");
field_bytes(&mut node, 4, b"MatMul");
let mut init = Vec::new();
field_varint(&mut init, 1, 3);
field_varint(&mut init, 1, 4);
field_varint(&mut init, 2, 1);
field_bytes(&mut init, 8, b"W");
field_bytes(&mut init, 9, &[0u8; 48]);
let mut graph = Vec::new();
field_bytes(&mut graph, 1, &node);
field_bytes(&mut graph, 2, b"tiny-mlp");
field_bytes(&mut graph, 5, &init);
let mut opset = Vec::new();
field_bytes(&mut opset, 1, b"");
field_varint(&mut opset, 2, 17);
let mut m = Vec::new();
field_varint(&mut m, 1, 9);
field_bytes(&mut m, 2, b"holger-test");
field_bytes(&mut m, 3, b"1.0");
field_bytes(&mut m, 7, &graph);
field_bytes(&mut m, 8, &opset);
m
}
/// `pickle.dumps({"a": 1}, protocol=2)`: a benign dict.
pub fn pickle() -> Vec<u8> {
b"\x80\x02}q\x00X\x01\x00\x00\x00aq\x01K\x01s.".to_vec()
}
/// A protocol-4 pickle of an empty `collections.OrderedDict`: benign, but it
/// imports a global, which the finding names.
pub fn pickle_with_global() -> Vec<u8> {
let mut b = vec![0x80, 0x04];
b.push(0x8c);
b.push(11);
b.extend_from_slice(b"collections");
b.push(0x94);
b.push(0x8c);
b.push(11);
b.extend_from_slice(b"OrderedDict");
b.push(0x94);
b.push(0x93);
b.push(0x94);
b.push(b')');
b.push(b'R');
b.push(0x94);
b.push(b'.');
b
}
/// A protocol-0 pickle of `{"a": 1}`, as text opcodes.
pub fn pickle_protocol0() -> Vec<u8> {
b"(dp0\nVa\np1\nI1\ns.".to_vec()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::scanner::{ScanVerdict, ScannerSet};
use std::sync::Arc;
fn scan(name: &str, bytes: &[u8]) -> ScanOutcome {
ModelScanner::new()
.scan(&ScanSubject::new(name, "raw", bytes))
.unwrap()
}
#[test]
fn a_gguf_file_passes_and_shows_its_metadata() {
let Sniffed::Model(m) = sniff(&fixtures::gguf()) else {
panic!("{:?}", sniff(&fixtures::gguf()))
};
assert_eq!(m.format, Some(ModelFormat::Gguf));
assert_eq!(m.architecture.as_deref(), Some("llama"));
assert_eq!(m.tensors, Some(2));
assert_eq!(m.parameters, Some(12), "4 + 4×2 elements");
assert_eq!(m.dtypes, vec!["F16", "F32"]);
assert_eq!(m.quantization.as_deref(), Some("MOSTLY_Q4_K_M"));
assert_eq!(m.extra.get("gguf_version").map(String::as_str), Some("3"));
let out = scan("tiny.gguf", &fixtures::gguf());
assert_eq!(out.verdict, ScanVerdict::Pass, "{}", out.detail);
assert_eq!(out.document.as_ref().unwrap()["architecture"], "llama");
}
#[test]
fn a_safetensors_file_passes_and_shows_its_metadata() {
let out = scan("model.safetensors", &fixtures::safetensors());
assert_eq!(out.verdict, ScanVerdict::Pass, "{}", out.detail);
let m: ModelMetadata = serde_json::from_value(out.document.unwrap()).unwrap();
assert_eq!(m.format, Some(ModelFormat::Safetensors));
assert_eq!(m.architecture.as_deref(), Some("bert"));
assert_eq!(m.tensors, Some(2));
assert_eq!(m.parameters, Some(9));
assert_eq!(m.dtypes, vec!["F16", "F32"]);
}
#[test]
fn an_onnx_model_passes_and_shows_its_metadata() {
let out = scan("tiny.onnx", &fixtures::onnx());
assert_eq!(out.verdict, ScanVerdict::Pass, "{}", out.detail);
let m: ModelMetadata = serde_json::from_value(out.document.unwrap()).unwrap();
assert_eq!(m.format, Some(ModelFormat::Onnx));
assert_eq!(m.architecture.as_deref(), Some("tiny-mlp"));
assert_eq!(m.tensors, Some(1));
assert_eq!(m.parameters, Some(12));
assert_eq!(m.dtypes, vec!["FLOAT"]);
assert_eq!(m.extra.get("opset").map(String::as_str), Some("17"));
assert_eq!(
m.extra.get("producer").map(String::as_str),
Some("holger-test")
);
}
#[test]
fn a_model_is_told_by_its_content_whatever_its_name() {
let out = scan("weights.bin", &fixtures::gguf());
assert_eq!(out.verdict, ScanVerdict::Pass);
assert_eq!(out.document.unwrap()["format"], "gguf");
}
#[test]
fn a_tiny_pickle_is_blocked_by_default_with_its_code() {
for (name, bytes) in [
("model.pkl", fixtures::pickle()),
("data.bin", fixtures::pickle()),
("od.pickle", fixtures::pickle_with_global()),
("p0.joblib", fixtures::pickle_protocol0()),
] {
let out = scan(name, &bytes);
assert_eq!(out.verdict, ScanVerdict::Block, "{name}: {}", out.detail);
assert_eq!(out.findings[0].id, "holger-model-1", "{name}");
assert!(out.detail.starts_with("[holger-model-1]"), "{}", out.detail);
assert!(out.detail.contains("pickle"), "{}", out.detail);
}
let out = scan("od.pickle", &fixtures::pickle_with_global());
assert!(
out.detail.contains("collections.OrderedDict"),
"{}",
out.detail
);
}
#[test]
fn a_fake_pt_zip_holding_data_pkl_is_blocked_by_its_layout() {
let pt = zip_entry::fixture::zip(&[
("archive/data.pkl", &fixtures::pickle()[..]),
("archive/version", &b"3\n"[..]),
("archive/data/0", &[0u8; 16][..]),
]);
// Named .pt, or named anything: the layout decides.
for name in ["model.pt", "checkpoint.zip", "weights.bin"] {
let out = scan(name, &pt);
assert_eq!(out.verdict, ScanVerdict::Block, "{name}: {}", out.detail);
assert_eq!(out.findings[0].id, "holger-model-2", "{name}");
assert!(out.detail.contains("archive/data.pkl"), "{}", out.detail);
}
}
#[test]
fn a_pickle_family_name_blocks_even_when_the_content_is_not_recognized() {
let out = scan("dump.joblib.gz", b"\x1f\x8b\x08\x00 compressed");
assert_eq!(out.verdict, ScanVerdict::Block);
assert_eq!(out.findings[0].id, "holger-model-1");
assert!(out.detail.contains(".joblib"), "{}", out.detail);
}
#[test]
fn a_pickle_renamed_safetensors_is_blocked_by_its_content() {
let out = scan("model.safetensors", &fixtures::pickle());
assert_eq!(out.verdict, ScanVerdict::Block, "{}", out.detail);
let ids: Vec<&str> = out.findings.iter().map(|f| f.id.as_str()).collect();
assert!(
ids.contains(&"holger-model-1"),
"the pickle is named: {ids:?}"
);
assert!(
ids.contains(&"holger-model-6"),
"the false name is named: {ids:?}"
);
}
#[test]
fn a_safetensors_header_that_lies_about_its_offsets_is_blocked() {
let out = scan("model.safetensors", &fixtures::safetensors_lying_offsets());
assert_eq!(out.verdict, ScanVerdict::Block, "{}", out.detail);
assert_eq!(out.findings[0].id, "holger-model-4", "{}", out.detail);
let out = scan("model.safetensors", &fixtures::safetensors_overlapping());
assert_eq!(out.verdict, ScanVerdict::Block, "{}", out.detail);
assert_eq!(out.findings[0].id, "holger-model-5", "{}", out.detail);
}
#[test]
fn a_malformed_safetensors_header_blocks() {
let cases: Vec<(&str, Vec<u8>)> = vec![
("length past the file", {
let mut b = fixtures::safetensors();
b[..8].copy_from_slice(&(1u64 << 20).to_le_bytes());
b
}),
("length over the limit", {
let mut b = fixtures::safetensors();
b[..8].copy_from_slice(&(MAX_SAFETENSORS_HEADER + 1).to_le_bytes());
b
}),
("not json", fixtures::safetensors_with("{not json}", &[])),
(
"size disagrees",
fixtures::safetensors_with(
r#"{"w":{"dtype":"F32","shape":[2],"data_offsets":[0,4]}}"#,
&[0u8; 4],
),
),
(
"unknown dtype",
fixtures::safetensors_with(
r#"{"w":{"dtype":"Q9","shape":[1],"data_offsets":[0,1]}}"#,
&[0u8; 1],
),
),
(
"trailing bytes",
fixtures::safetensors_with(
r#"{"w":{"dtype":"U8","shape":[1],"data_offsets":[0,1]}}"#,
&[0u8; 9],
),
),
(
"a hole",
fixtures::safetensors_with(
r#"{"w":{"dtype":"U8","shape":[1],"data_offsets":[4,5]}}"#,
&[0u8; 5],
),
),
];
for (what, bytes) in cases {
let out = scan("m.safetensors", &bytes);
assert_eq!(out.verdict, ScanVerdict::Block, "{what}: {}", out.detail);
assert_eq!(
out.findings[0].id, "holger-model-3",
"{what}: {}",
out.detail
);
}
}
#[test]
fn a_malformed_gguf_header_blocks() {
let good = fixtures::gguf();
let mut bad_version = good.clone();
bad_version[4..8].copy_from_slice(&9u32.to_le_bytes());
let mut huge_count = good.clone();
huge_count[8..16].copy_from_slice(&u64::MAX.to_le_bytes());
let truncated = good[..60].to_vec();
for (what, b, code) in [
("version", bad_version, "holger-model-3"),
("tensor count", huge_count, "holger-model-3"),
("truncated", truncated, "holger-model-3"),
(
"data cut off",
good[..good.len() - 8].to_vec(),
"holger-model-4",
),
] {
let out = scan("m.gguf", &b);
assert_eq!(out.verdict, ScanVerdict::Block, "{what}: {}", out.detail);
assert_eq!(out.findings[0].id, code, "{what}: {}", out.detail);
}
}
#[test]
fn gguf_tensors_that_overlap_are_blocked() {
let mut b = fixtures::gguf();
// The second tensor's offset (32) is the last 8 bytes of its info; move it onto the first.
let at = b
.windows(19)
.position(|w| w == b"blk.0.attn_q.weight")
.unwrap()
+ 19
+ 4
+ 16
+ 4;
b[at..at + 8].copy_from_slice(&0u64.to_le_bytes());
let out = scan("m.gguf", &b);
assert_eq!(out.verdict, ScanVerdict::Block, "{}", out.detail);
assert_eq!(out.findings[0].id, "holger-model-5", "{}", out.detail);
}
#[test]
fn allow_pickle_warns_in_the_named_repository_only() {
let s = ModelScanner::new().allow_pickle_in(["ml-research"]);
let p = fixtures::pickle();
let there = s
.scan(&ScanSubject::new("m.pkl", "raw", &p).with_repo("ml-research"))
.unwrap();
assert_eq!(there.verdict, ScanVerdict::Warn, "{}", there.detail);
assert_eq!(there.findings[0].id, "holger-model-7");
assert!(
there.detail.contains("holger-model-1"),
"the reason is still named: {}",
there.detail
);
let elsewhere = s
.scan(&ScanSubject::new("m.pkl", "raw", &p).with_repo("incoming"))
.unwrap();
assert_eq!(elsewhere.verdict, ScanVerdict::Block);
let no_repo = s.scan(&ScanSubject::new("m.pkl", "raw", &p)).unwrap();
assert_eq!(
no_repo.verdict,
ScanVerdict::Block,
"no repository ⇒ the default"
);
// A malformed header blocks even where pickles are allowed.
let lying = fixtures::safetensors_lying_offsets();
let out = s
.scan(&ScanSubject::new("m.safetensors", "raw", &lying).with_repo("ml-research"))
.unwrap();
assert_eq!(out.verdict, ScanVerdict::Block);
}
#[test]
fn an_unknown_allow_pickle_repository_is_refused_by_name() {
let err =
validate_allow_pickle(&["typo".to_string()], ["incoming", "approved"]).unwrap_err();
assert!(err.to_string().starts_with("[holger-model-8]"), "{err}");
assert!(err.to_string().contains("'typo'"), "{err}");
validate_allow_pickle(&["incoming".to_string()], ["incoming"]).unwrap();
}
#[test]
fn ordinary_files_pass_and_carry_no_model_document() {
for (name, bytes) in [
("README.md", &b"# hello\n"[..]),
(
"demo-1.0-py3-none-any.whl",
&zip_entry::fixture::zip(&[("demo/__init__.py", &b""[..])])[..],
),
("empty", &b""[..]),
("noise", &b"\x08\x96\x01 not protobuf"[..]),
] {
let out = scan(name, bytes);
assert_eq!(out.verdict, ScanVerdict::Pass, "{name}: {}", out.detail);
assert!(out.document.is_none(), "{name}");
}
}
#[test]
fn the_model_slot_blocks_the_fold_and_its_metadata_is_read_back() {
let set = ScannerSet::new().with(Arc::new(ModelScanner::new()));
let ev = set.scan(&ScanSubject::new("m.pt", "raw", &fixtures::pickle()));
assert_eq!(ev.verdict, ScanVerdict::Block);
assert!(
ev.reasons
.iter()
.any(|r| r.contains("model (holger-model) blocks: [holger-model-1]")),
"{:?}",
ev.reasons
);
let ev = set.scan(&ScanSubject::new("m.gguf", "raw", &fixtures::gguf()));
assert_eq!(ev.verdict, ScanVerdict::Pass, "{:?}", ev.reasons);
let m = metadata_of(&ev).unwrap();
let props = m.properties();
assert!(
props.contains(&("model.architecture".into(), vec!["llama".into()])),
"{props:?}"
);
assert!(
props.contains(&("model.format".into(), vec!["gguf".into()])),
"{props:?}"
);
}
}