150 lines
5.2 KiB
Rust
150 lines
5.2 KiB
Rust
//! Verify every isolated runtime code artifact before starting a worker.
|
|||
|
|
use super::{file_hash, Device, UPSTREAM};
|
||
|
|
use anyhow::{ensure, Context, Result};
|
||
|
|
use serde::Deserialize;
|
||
|
|
use sha2::{Digest, Sha256};
|
||
|
|
use std::{
|
||
|
|
collections::{BTreeMap, BTreeSet},
|
||
|
|
fs,
|
||
|
|
path::Path,
|
||
|
|
};
|
||
|
|
#[derive(Deserialize)]
|
||
|
|
struct Manifest {
|
||
|
|
schema: u32,
|
||
|
|
upstream_commit: String,
|
||
|
|
device: String,
|
||
|
|
distribution: bool,
|
||
|
|
source_hashes: BTreeMap<String, String>,
|
||
|
|
files: Vec<Entry>,
|
||
|
|
}
|
||
|
|
#[derive(Deserialize)]
|
||
|
|
struct Entry {
|
||
|
|
path: String,
|
||
|
|
bytes: u64,
|
||
|
|
sha256: String,
|
||
|
|
}
|
||
|
|
pub(crate) fn verify(worker: &Path, device: Device) -> Result<String> {
|
||
|
|
let directory = worker
|
||
|
|
.parent()
|
||
|
|
.context("Missing speech runtime directory")?;
|
||
|
|
let path = directory.join("runtime-manifest.json");
|
||
|
|
ensure!(
|
||
|
|
fs::symlink_metadata(&path)
|
||
|
|
.context("Speech runtime inventory missing")?
|
||
|
|
.file_type()
|
||
|
|
.is_file(),
|
||
|
|
"Runtime inventory must be a regular file"
|
||
|
|
);
|
||
|
|
ensure!(
|
||
|
|
fs::metadata(&path)?.len() <= 512 * 1024,
|
||
|
|
"Runtime inventory exceeds bound"
|
||
|
|
);
|
||
|
|
let bytes = fs::read(path)?;
|
||
|
|
let manifest: Manifest =
|
||
|
|
serde_json::from_slice(&bytes).context("Invalid speech runtime inventory")?;
|
||
|
|
ensure!(
|
||
|
|
manifest.schema == 1 && manifest.upstream_commit == UPSTREAM,
|
||
|
|
"Stale speech runtime inventory"
|
||
|
|
);
|
||
|
|
ensure!(
|
||
|
|
matches!(manifest.device.as_str(), "auto" | "cpu")
|
||
|
|
&& manifest.distribution == (manifest.device == "auto"),
|
||
|
|
"Inconsistent speech distribution inventory"
|
||
|
|
);
|
||
|
|
ensure!(
|
||
|
|
device == Device::Cpu || manifest.distribution,
|
||
|
|
"CPU-only speech runtime is an internal tool"
|
||
|
|
);
|
||
|
|
for name in ["CMakeLists.txt", "worker.cpp", "build.py"] {
|
||
|
|
ensure!(
|
||
|
|
manifest.source_hashes.get(name).is_some_and(|s| hash(s)),
|
||
|
|
"Missing speech build-input provenance"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
ensure!(
|
||
|
|
!manifest.files.is_empty() && manifest.files.len() <= 128,
|
||
|
|
"Unbounded speech runtime inventory"
|
||
|
|
);
|
||
|
|
let mut seen = BTreeSet::new();
|
||
|
|
let mut total = 0u64;
|
||
|
|
for entry in manifest.files {
|
||
|
|
ensure!(
|
||
|
|
!entry.path.is_empty()
|
||
|
|
&& entry.path != "."
|
||
|
|
&& entry.path != ".."
|
||
|
|
&& !entry.path.contains(['/', '\\', ':'])
|
||
|
|
&& seen.insert(entry.path.clone()),
|
||
|
|
"Invalid speech inventory filename"
|
||
|
|
);
|
||
|
|
ensure!(
|
||
|
|
entry.bytes > 0 && entry.bytes <= 256 * 1024 * 1024 && hash(&entry.sha256),
|
||
|
|
"Invalid speech inventory artifact"
|
||
|
|
);
|
||
|
|
total = total
|
||
|
|
.checked_add(entry.bytes)
|
||
|
|
.context("Speech runtime size overflow")?;
|
||
|
|
ensure!(
|
||
|
|
total <= 1024 * 1024 * 1024,
|
||
|
|
"Speech runtime exceeds size budget"
|
||
|
|
);
|
||
|
|
let path = directory.join(&entry.path);
|
||
|
|
ensure!(
|
||
|
|
fs::symlink_metadata(&path)?.file_type().is_file()
|
||
|
|
&& fs::metadata(&path)?.len() == entry.bytes
|
||
|
|
&& file_hash(&path)? == entry.sha256,
|
||
|
|
"Speech runtime artifact size/SHA changed"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
ensure!(
|
||
|
|
worker
|
||
|
|
.file_name()
|
||
|
|
.and_then(|s| s.to_str())
|
||
|
|
.is_some_and(|name| seen.contains(name)),
|
||
|
|
"Worker omitted from speech runtime inventory"
|
||
|
|
);
|
||
|
|
Ok(format!("{:x}", Sha256::digest(bytes)))
|
||
|
|
}
|
||
|
|
fn hash(value: &str) -> bool {
|
||
|
|
value.len() == 64 && value.bytes().all(|b| b.is_ascii_hexdigit())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use serde_json::json;
|
||
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||
|
|
#[test]
|
||
|
|
fn altered_dependency_or_missing_inventory_fails_before_process_launch() {
|
||
|
|
let unique = SystemTime::now()
|
||
|
|
.duration_since(UNIX_EPOCH)
|
||
|
|
.unwrap()
|
||
|
|
.as_nanos();
|
||
|
|
let root = std::env::temp_dir().join(format!(
|
||
|
|
"searchi-speech-inventory-{}-{unique}",
|
||
|
|
std::process::id()
|
||
|
|
));
|
||
|
|
fs::create_dir(&root).unwrap();
|
||
|
|
let worker = root.join("worker");
|
||
|
|
let library = root.join("ggml-cpu.dll");
|
||
|
|
fs::write(&worker, b"structural worker fixture").unwrap();
|
||
|
|
fs::write(&library, b"structural library fixture").unwrap();
|
||
|
|
assert!(verify(&worker, Device::Cpu).is_err());
|
||
|
|
let artifact = |path: &Path| json!({"path":path.file_name().unwrap().to_str().unwrap(),"bytes":fs::metadata(path).unwrap().len(),"sha256":file_hash(path).unwrap()});
|
||
|
|
let sources: BTreeMap<_, _> = ["CMakeLists.txt", "worker.cpp", "build.py"]
|
||
|
|
.into_iter()
|
||
|
|
.map(|s| (s, "a".repeat(64)))
|
||
|
|
.collect();
|
||
|
|
let manifest = json!({"schema":1,"upstream_commit":UPSTREAM,"device":"cpu","distribution":false,"source_hashes":sources,"files":[artifact(&worker),artifact(&library)]});
|
||
|
|
fs::write(
|
||
|
|
root.join("runtime-manifest.json"),
|
||
|
|
serde_json::to_vec(&manifest).unwrap(),
|
||
|
|
)
|
||
|
|
.unwrap();
|
||
|
|
let identity = verify(&worker, Device::Cpu).unwrap();
|
||
|
|
assert_eq!(identity.len(), 64);
|
||
|
|
assert!(verify(&worker, Device::Auto).is_err());
|
||
|
|
fs::write(library, b"same length changed bits").unwrap();
|
||
|
|
assert!(verify(&worker, Device::Cpu).is_err());
|
||
|
|
fs::remove_dir_all(root).unwrap();
|
||
|
|
}
|
||
|
|
}
|