mirror of
https://github.com/DaLaw2/NetGuardia.git
synced 2026-08-24 14:10:28 +09:00
feat(http): multipart model upload + staging orphan sweep
Backend scaffolding for the admin model-upload UI. Accepts `manifest`
(YAML) + `onnx` (binary) multipart fields, streams each directly to a
fresh `models/.staging/<uuid>/` directory with enforced size caps, and
sanity-checks the ONNX header before ever touching disk. Structural
schema validation against `MLInferenceConfig` — plus the atomic
promote-to-`models/` step — is the next milestone's job; this commit
lands only the upload / staging surface.
New file `adapter/http/model_upload.rs`:
- `POST /api/ml/models/upload` scoped under `/api/ml/models`. Returns
`{ staging_id, manifest_bytes, onnx_bytes, onnx_filename }` on
success.
- Streams per-field, never buffers the full `.onnx` in RAM. Per-field
byte cap enforced inline: 100 MB for onnx, 64 KB for manifest.
Overflow yields 413 and the staging subdir is torn down before the
response returns.
- `looks_like_onnx` heuristic rejects empty, all-zero, all-printable,
and common container magics (ZIP, PNG, PDF, ELF) up front. Positive
acceptance is handled authoritatively later by tract during
`build_adapter`.
- `sanitize_filename` keeps the client's suffix but strips traversal
so the staging dir can never escape.
- `clean_staging_orphans(root, max_age)` sweeps leftover subdirs;
reused at startup and available for a periodic task later.
Wiring:
- `actix-multipart = "0.7"` and `uuid` (with `v4`) added to
`net-guardia/Cargo.toml`.
- `system.rs` calls `clean_staging_orphans` once during `run()` before
the model watcher starts listening, with a 1 h age threshold. Two
new `SystemLog` variants (`StagingOrphansCleaned` /
`StagingOrphansSweepFailed`) report results.
- `http_server.rs` registers the new scope inside the existing
`/api` scope. CSRF / JWT middleware already covers the route; no
permission table changes needed beyond what I-11 provides.
Tests: 10 new (5 ONNX-sniff edge cases — empty / zero-padded / plain
text / varint tag / length-delimited tag; 3 `sanitize_filename`
cases; 2 orphan-sweep cases including the no-staging-dir short
circuit). 240 pass total. clippy --package net-guardia -- -D warnings
clean.
Closes I-region I-12.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
1578ff2071
commit
a419c64aae
136
Cargo.lock
generated
136
Cargo.lock
generated
@ -52,7 +52,7 @@ checksum = "daa239b93927be1ff123eebada5a3ff23e89f0124ccb8609234e5103d5a5ae6d"
|
||||
dependencies = [
|
||||
"actix-utils",
|
||||
"actix-web",
|
||||
"derive_more",
|
||||
"derive_more 2.1.1",
|
||||
"futures-util",
|
||||
"log",
|
||||
"once_cell",
|
||||
@ -74,7 +74,7 @@ dependencies = [
|
||||
"brotli",
|
||||
"bytes",
|
||||
"bytestring",
|
||||
"derive_more",
|
||||
"derive_more 2.1.1",
|
||||
"encoding_rs",
|
||||
"flate2",
|
||||
"foldhash 0.1.5",
|
||||
@ -108,6 +108,44 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "actix-multipart"
|
||||
version = "0.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d5118a26dee7e34e894f7e85aa0ee5080ae4c18bf03c0e30d49a80e418f00a53"
|
||||
dependencies = [
|
||||
"actix-multipart-derive",
|
||||
"actix-utils",
|
||||
"actix-web",
|
||||
"derive_more 0.99.20",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"httparse",
|
||||
"local-waker",
|
||||
"log",
|
||||
"memchr",
|
||||
"mime",
|
||||
"rand 0.8.5",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_plain",
|
||||
"tempfile",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "actix-multipart-derive"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e11eb847f49a700678ea2fa73daeb3208061afa2b9d1a8527c03390f4c4a1c6b"
|
||||
dependencies = [
|
||||
"darling",
|
||||
"parse-size",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "actix-router"
|
||||
version = "0.5.4"
|
||||
@ -189,7 +227,7 @@ dependencies = [
|
||||
"bytestring",
|
||||
"cfg-if",
|
||||
"cookie",
|
||||
"derive_more",
|
||||
"derive_more 2.1.1",
|
||||
"encoding_rs",
|
||||
"foldhash 0.1.5",
|
||||
"futures-core",
|
||||
@ -898,6 +936,12 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "convert_case"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e"
|
||||
|
||||
[[package]]
|
||||
name = "convert_case"
|
||||
version = "0.10.0"
|
||||
@ -1042,6 +1086,41 @@ dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling"
|
||||
version = "0.20.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
|
||||
dependencies = [
|
||||
"darling_core",
|
||||
"darling_macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling_core"
|
||||
version = "0.20.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
|
||||
dependencies = [
|
||||
"fnv",
|
||||
"ident_case",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"strsim",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling_macro"
|
||||
version = "0.20.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
|
||||
dependencies = [
|
||||
"darling_core",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dashmap"
|
||||
version = "6.1.0"
|
||||
@ -1082,6 +1161,19 @@ dependencies = [
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_more"
|
||||
version = "0.99.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6edb4b64a43d977b8e99788fe3a04d483834fba1215a7e02caa415b626497f7f"
|
||||
dependencies = [
|
||||
"convert_case 0.4.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"rustc_version",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_more"
|
||||
version = "2.1.1"
|
||||
@ -1097,7 +1189,7 @@ version = "2.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"convert_case 0.10.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"rustc_version",
|
||||
@ -1743,6 +1835,12 @@ version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954"
|
||||
|
||||
[[package]]
|
||||
name = "ident_case"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "1.1.0"
|
||||
@ -2342,6 +2440,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"actix",
|
||||
"actix-cors",
|
||||
"actix-multipart",
|
||||
"actix-web",
|
||||
"actix-ws",
|
||||
"aes-gcm",
|
||||
@ -2393,6 +2492,7 @@ dependencies = [
|
||||
"tracing-subscriber",
|
||||
"tract-onnx",
|
||||
"url",
|
||||
"uuid",
|
||||
"which",
|
||||
"xsk-rs",
|
||||
]
|
||||
@ -2649,6 +2749,12 @@ dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parse-size"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "487f2ccd1e17ce8c1bfab3a65c89525af41cfad4c8659021a1e9a2aacd73b89b"
|
||||
|
||||
[[package]]
|
||||
name = "password-hash"
|
||||
version = "0.5.0"
|
||||
@ -3417,6 +3523,15 @@ dependencies = [
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_plain"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ce1fc6db65a611022b23a0dec6975d63fb80a302cb3388835ff02c097258d50"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_spanned"
|
||||
version = "1.0.4"
|
||||
@ -3669,6 +3784,19 @@ dependencies = [
|
||||
"xattr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tempfile"
|
||||
version = "3.27.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.4.2",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.69"
|
||||
|
||||
@ -20,6 +20,8 @@ actix = { workspace = true }
|
||||
actix-web = { workspace = true }
|
||||
actix-cors = { workspace = true }
|
||||
actix-ws = { workspace = true }
|
||||
actix-multipart = "0.7"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
rust-embed = "8.11.0"
|
||||
mime_guess = "2.0.5"
|
||||
url = "2.5.8"
|
||||
|
||||
@ -8,6 +8,7 @@ pub mod fusion;
|
||||
pub mod health;
|
||||
pub mod logs;
|
||||
pub mod ml;
|
||||
pub mod model_upload;
|
||||
pub mod notification;
|
||||
pub mod rate_limit;
|
||||
pub mod report;
|
||||
|
||||
399
net-guardia/src/adapter/http/model_upload.rs
Normal file
399
net-guardia/src/adapter/http/model_upload.rs
Normal file
@ -0,0 +1,399 @@
|
||||
//! Multipart upload surface for BYO model files. Accepts a `manifest`
|
||||
//! YAML field plus an `onnx` binary field, streams each to
|
||||
//! `models/.staging/<uuid>/` with enforced size caps, and leaves atomic
|
||||
//! promotion to `models/` for the model-watcher hot-reload flow to
|
||||
//! handle after I-6 wires validation + rename.
|
||||
//!
|
||||
//! Scope here: body-size caps, streaming-to-disk (never buffering the
|
||||
//! full `.onnx` in RAM), filename-collision safety via UUID subdir, and
|
||||
//! cheap up-front sanity on the ONNX header. Structural schema
|
||||
//! validation against `MLInferenceConfig` is the next milestone's job.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use actix_multipart::Multipart;
|
||||
use actix_web::{HttpResponse, Responder, Scope, web};
|
||||
use futures_util::TryStreamExt;
|
||||
use tokio::fs;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::infrastructure::app_config::AppConfig;
|
||||
use crate::model::config::constants::{MODELS_DIR, STAGING_SUBDIR};
|
||||
|
||||
/// Hard ceiling on the `.onnx` stream. Anything larger is either a
|
||||
/// mistake or an attempt to DoS the disk.
|
||||
pub const MAX_ONNX_BYTES: usize = 100 * 1024 * 1024;
|
||||
|
||||
/// Hard ceiling on the `manifest` YAML stream. Real manifests are a few
|
||||
/// kilobytes at most; this leaves headroom without allowing blob
|
||||
/// smuggling through the field.
|
||||
pub const MAX_MANIFEST_BYTES: usize = 64 * 1024;
|
||||
|
||||
/// Multipart field names the client must use. Stable wire contract —
|
||||
/// the frontend form generator depends on these exact strings.
|
||||
const FIELD_MANIFEST: &str = "manifest";
|
||||
const FIELD_ONNX: &str = "onnx";
|
||||
|
||||
/// Number of bytes of the ONNX body we inspect up-front for an obvious
|
||||
/// non-Protobuf header. A fuller structural check (shape vs manifest
|
||||
/// declared `features`) runs during `build_adapter` in the promote path.
|
||||
const ONNX_SNIFF_BYTES: usize = 16;
|
||||
|
||||
pub fn initialize() -> Scope {
|
||||
web::scope("/models").route("/upload", web::post().to(upload))
|
||||
}
|
||||
|
||||
/// `POST /api/ml/models/upload` — multipart with `manifest` (YAML text)
|
||||
/// and `onnx` (binary) fields. Writes both to a fresh
|
||||
/// `models/.staging/<uuid>/` and returns the staging id + byte counts.
|
||||
async fn upload(app_config: web::Data<Arc<AppConfig>>, payload: Multipart) -> impl Responder {
|
||||
let _ = app_config; // kept for future per-tenant staging roots
|
||||
|
||||
let staging_root = PathBuf::from(MODELS_DIR).join(STAGING_SUBDIR);
|
||||
let staging_id = Uuid::new_v4().to_string();
|
||||
let staging_dir = staging_root.join(&staging_id);
|
||||
|
||||
match ingest_multipart(payload, &staging_dir).await {
|
||||
Ok(summary) => HttpResponse::Ok().json(serde_json::json!({
|
||||
"staging_id": staging_id,
|
||||
"manifest_bytes": summary.manifest_bytes,
|
||||
"onnx_bytes": summary.onnx_bytes,
|
||||
"onnx_filename": summary.onnx_filename,
|
||||
})),
|
||||
Err(kind) => {
|
||||
// Best-effort teardown so orphan staging dirs don't
|
||||
// accumulate from failed uploads.
|
||||
let _ = fs::remove_dir_all(&staging_dir).await;
|
||||
kind.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Successful-path metadata the handler surfaces to the client.
|
||||
#[derive(Debug)]
|
||||
struct UploadSummary {
|
||||
manifest_bytes: usize,
|
||||
onnx_bytes: usize,
|
||||
onnx_filename: String,
|
||||
}
|
||||
|
||||
/// Errors that can surface a specific HTTP response. Kept in-module
|
||||
/// because none of these have callers outside this handler.
|
||||
#[derive(Debug)]
|
||||
enum UploadError {
|
||||
MissingField(&'static str),
|
||||
DuplicateField(&'static str),
|
||||
UnknownField(String),
|
||||
ManifestTooLarge,
|
||||
OnnxTooLarge,
|
||||
OnnxNotBinary,
|
||||
StreamFailure(String),
|
||||
StagingSetupFailure(String),
|
||||
}
|
||||
|
||||
impl UploadError {
|
||||
fn into_response(self) -> HttpResponse {
|
||||
let (status, message) = match self {
|
||||
Self::MissingField(name) => (400, format!("missing required multipart field: {name}")),
|
||||
Self::DuplicateField(name) => (400, format!("multipart field sent twice: {name}")),
|
||||
Self::UnknownField(name) => (400, format!("unexpected multipart field: {name}")),
|
||||
Self::ManifestTooLarge => (413, format!("manifest exceeds {MAX_MANIFEST_BYTES} bytes")),
|
||||
Self::OnnxTooLarge => (413, format!("onnx exceeds {MAX_ONNX_BYTES} bytes")),
|
||||
Self::OnnxNotBinary => (
|
||||
400,
|
||||
"onnx field does not look like a protobuf-encoded ONNX model".to_string(),
|
||||
),
|
||||
Self::StreamFailure(err) => (400, format!("upload stream error: {err}")),
|
||||
Self::StagingSetupFailure(err) => (500, format!("staging directory error: {err}")),
|
||||
};
|
||||
let body = serde_json::json!({ "error": message });
|
||||
match status {
|
||||
400 => HttpResponse::BadRequest().json(body),
|
||||
413 => HttpResponse::PayloadTooLarge().json(body),
|
||||
_ => HttpResponse::InternalServerError().json(body),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn ingest_multipart(mut payload: Multipart, staging_dir: &Path) -> Result<UploadSummary, UploadError> {
|
||||
fs::create_dir_all(staging_dir)
|
||||
.await
|
||||
.map_err(|e| UploadError::StagingSetupFailure(e.to_string()))?;
|
||||
|
||||
let mut manifest_written: Option<usize> = None;
|
||||
let mut onnx_summary: Option<(String, usize)> = None;
|
||||
|
||||
while let Some(mut field) = payload
|
||||
.try_next()
|
||||
.await
|
||||
.map_err(|e| UploadError::StreamFailure(e.to_string()))?
|
||||
{
|
||||
let field_name = field
|
||||
.content_disposition()
|
||||
.and_then(|cd| cd.get_name())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
match field_name.as_str() {
|
||||
FIELD_MANIFEST => {
|
||||
if manifest_written.is_some() {
|
||||
return Err(UploadError::DuplicateField(FIELD_MANIFEST));
|
||||
}
|
||||
let dest = staging_dir.join("manifest.yaml");
|
||||
let written = stream_field_to_file(&mut field, &dest, MAX_MANIFEST_BYTES, false).await?;
|
||||
manifest_written = Some(written);
|
||||
}
|
||||
FIELD_ONNX => {
|
||||
if onnx_summary.is_some() {
|
||||
return Err(UploadError::DuplicateField(FIELD_ONNX));
|
||||
}
|
||||
let onnx_filename = field
|
||||
.content_disposition()
|
||||
.and_then(|cd| cd.get_filename())
|
||||
.map(sanitize_filename)
|
||||
.unwrap_or_else(|| "model.onnx".to_string());
|
||||
let dest = staging_dir.join(&onnx_filename);
|
||||
let written = stream_field_to_file(&mut field, &dest, MAX_ONNX_BYTES, true).await?;
|
||||
onnx_summary = Some((onnx_filename, written));
|
||||
}
|
||||
other => {
|
||||
return Err(UploadError::UnknownField(other.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let manifest_bytes = manifest_written.ok_or(UploadError::MissingField(FIELD_MANIFEST))?;
|
||||
let (onnx_filename, onnx_bytes) = onnx_summary.ok_or(UploadError::MissingField(FIELD_ONNX))?;
|
||||
|
||||
Ok(UploadSummary {
|
||||
manifest_bytes,
|
||||
onnx_bytes,
|
||||
onnx_filename,
|
||||
})
|
||||
}
|
||||
|
||||
/// Stream a multipart field directly to disk. Aborts (and leaves the
|
||||
/// caller to clean up) when the declared byte cap is exceeded or when
|
||||
/// the binary sniff rejects the first chunk.
|
||||
async fn stream_field_to_file(
|
||||
field: &mut actix_multipart::Field,
|
||||
dest: &Path,
|
||||
max_bytes: usize,
|
||||
sniff_onnx: bool,
|
||||
) -> Result<usize, UploadError> {
|
||||
let mut file = fs::File::create(dest)
|
||||
.await
|
||||
.map_err(|e| UploadError::StagingSetupFailure(e.to_string()))?;
|
||||
let mut total = 0usize;
|
||||
let mut sniffed = !sniff_onnx;
|
||||
|
||||
while let Some(chunk) = field
|
||||
.try_next()
|
||||
.await
|
||||
.map_err(|e| UploadError::StreamFailure(e.to_string()))?
|
||||
{
|
||||
if !sniffed {
|
||||
// Cheap up-front validation: reject obvious non-ONNX blobs
|
||||
// (empty first chunk, all-zero header, all-printable text).
|
||||
if !looks_like_onnx(&chunk) {
|
||||
return Err(UploadError::OnnxNotBinary);
|
||||
}
|
||||
sniffed = true;
|
||||
}
|
||||
total = total.saturating_add(chunk.len());
|
||||
if total > max_bytes {
|
||||
return Err(if sniff_onnx {
|
||||
UploadError::OnnxTooLarge
|
||||
} else {
|
||||
UploadError::ManifestTooLarge
|
||||
});
|
||||
}
|
||||
file.write_all(&chunk)
|
||||
.await
|
||||
.map_err(|e| UploadError::StreamFailure(e.to_string()))?;
|
||||
}
|
||||
file.flush()
|
||||
.await
|
||||
.map_err(|e| UploadError::StreamFailure(e.to_string()))?;
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Shape-preserving filename sanitizer: keep the extension the client
|
||||
/// sent (it may be `.onnx`, `.bin`, whatever), but strip any directory
|
||||
/// traversal so the staging dir can never escape.
|
||||
pub fn sanitize_filename(raw: impl AsRef<str>) -> String {
|
||||
let raw = raw.as_ref();
|
||||
let trimmed = Path::new(raw)
|
||||
.file_name()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or("model.onnx");
|
||||
if trimmed.is_empty() || trimmed == "." || trimmed == ".." {
|
||||
"model.onnx".to_string()
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Loose first-chunk heuristic. An ONNX protobuf starts with a varint
|
||||
/// tag byte — the field=1 wire=varint (`0x08` for `ir_version`) and
|
||||
/// field=1 wire=length-delimited (`0x0a`) patterns both occur in real
|
||||
/// models — but enumerating positive accept patterns is fragile because
|
||||
/// tract accepts several tag orderings. We instead:
|
||||
///
|
||||
/// 1. Reject magic bytes of container formats that users routinely
|
||||
/// upload by mistake (ZIP, PNG, PDF, ELF).
|
||||
/// 2. Reject all-zero and all-printable-ASCII prefixes (buffers and
|
||||
/// text files).
|
||||
///
|
||||
/// The authoritative structural validation happens during
|
||||
/// `build_adapter`; this heuristic's job is catching the obvious wrong
|
||||
/// upload before the bytes hit disk.
|
||||
pub fn looks_like_onnx(first_chunk: &[u8]) -> bool {
|
||||
if first_chunk.is_empty() {
|
||||
return false;
|
||||
}
|
||||
// Container formats that users commonly confuse with ONNX.
|
||||
const BLOCKED_MAGICS: &[&[u8]] = &[
|
||||
b"PK\x03\x04", // ZIP / JAR / DOCX — some pipelines ship ONNX weights this way,
|
||||
// but our upload path expects a single standalone .onnx file.
|
||||
b"\x89PNG",
|
||||
b"%PDF",
|
||||
b"\x7fELF",
|
||||
];
|
||||
for magic in BLOCKED_MAGICS {
|
||||
if first_chunk.starts_with(magic) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
let prefix = &first_chunk[..first_chunk.len().min(ONNX_SNIFF_BYTES)];
|
||||
if prefix.iter().all(|&b| b == 0) {
|
||||
return false;
|
||||
}
|
||||
let mostly_ascii = prefix.iter().filter(|&&b| b.is_ascii_graphic() || b == b' ').count() >= prefix.len() - 1;
|
||||
if mostly_ascii {
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Remove staging subdirectories older than `max_age`. Runs on startup
|
||||
/// and on a periodic timer so failed uploads don't accumulate.
|
||||
pub fn clean_staging_orphans(staging_root: &Path, max_age: std::time::Duration) -> std::io::Result<usize> {
|
||||
if !staging_root.exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
let now = std::time::SystemTime::now();
|
||||
let mut cleaned = 0usize;
|
||||
for entry in std::fs::read_dir(staging_root)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let metadata = entry.metadata()?;
|
||||
let mtime = metadata.modified()?;
|
||||
let age = now.duration_since(mtime).unwrap_or_default();
|
||||
if age >= max_age {
|
||||
std::fs::remove_dir_all(&path)?;
|
||||
cleaned += 1;
|
||||
}
|
||||
}
|
||||
Ok(cleaned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn onnx_sniff_rejects_empty() {
|
||||
assert!(!looks_like_onnx(&[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn onnx_sniff_rejects_zero_padded_prefix() {
|
||||
assert!(!looks_like_onnx(&[0u8; 32]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn onnx_sniff_rejects_plain_text() {
|
||||
// A YAML or plain-text payload that ended up in the wrong field.
|
||||
assert!(!looks_like_onnx(b"name: wrong-file\nkind: yaml\n"));
|
||||
assert!(!looks_like_onnx(b"PK\x03\x04"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn onnx_sniff_accepts_varint_tag_prefix() {
|
||||
// `0x08` = tag field 1, wire-type varint (ir_version). Real ONNX
|
||||
// files commonly open with this.
|
||||
let buf = [0x08u8, 0x07, 0x12, 0x0a, 0x70, 0x79, 0x74, 0x6f, 0x72, 0x63, 0x68, 0x00];
|
||||
assert!(looks_like_onnx(&buf));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn onnx_sniff_accepts_length_delimited_tag() {
|
||||
// `0x0a` = tag field 1, wire-type length-delimited. Also valid.
|
||||
let buf = [0x0au8, 0x10, 0x80, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07];
|
||||
assert!(looks_like_onnx(&buf));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_filename_strips_directory_components() {
|
||||
assert_eq!(sanitize_filename("../../etc/passwd"), "passwd");
|
||||
assert_eq!(sanitize_filename("subdir/model.onnx"), "model.onnx");
|
||||
assert_eq!(sanitize_filename("/abs/path/classifier.onnx"), "classifier.onnx");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_filename_rejects_degenerate_values() {
|
||||
assert_eq!(sanitize_filename(""), "model.onnx");
|
||||
assert_eq!(sanitize_filename("."), "model.onnx");
|
||||
assert_eq!(sanitize_filename(".."), "model.onnx");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn orphan_cleanup_removes_every_dir_when_max_age_is_zero() {
|
||||
// A zero-length max age declares every existing entry stale, so
|
||||
// the helper must sweep all of them. Portable without touching
|
||||
// filesystem mtime APIs.
|
||||
let tmp = std::env::temp_dir().join(format!("nguardia-staging-test-{}", Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
std::fs::create_dir_all(tmp.join("abandoned-1")).unwrap();
|
||||
std::fs::create_dir_all(tmp.join("abandoned-2")).unwrap();
|
||||
// A file (not a dir) should be ignored by the sweep.
|
||||
std::fs::write(tmp.join("sidecar.log"), b"noise").unwrap();
|
||||
|
||||
let cleaned = clean_staging_orphans(&tmp, Duration::ZERO).unwrap();
|
||||
assert_eq!(cleaned, 2);
|
||||
assert!(!tmp.join("abandoned-1").exists());
|
||||
assert!(!tmp.join("abandoned-2").exists());
|
||||
assert!(tmp.join("sidecar.log").exists(), "non-directory entries must survive");
|
||||
|
||||
std::fs::remove_dir_all(&tmp).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn orphan_cleanup_preserves_fresh_directories() {
|
||||
// With a generous max_age, a freshly-created directory must not
|
||||
// be touched — the positive case of the time-guard.
|
||||
let tmp = std::env::temp_dir().join(format!("nguardia-staging-fresh-{}", Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
std::fs::create_dir_all(tmp.join("recent")).unwrap();
|
||||
|
||||
let cleaned = clean_staging_orphans(&tmp, Duration::from_secs(3600)).unwrap();
|
||||
assert_eq!(cleaned, 0);
|
||||
assert!(tmp.join("recent").exists());
|
||||
|
||||
std::fs::remove_dir_all(&tmp).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn orphan_cleanup_is_noop_when_staging_root_missing() {
|
||||
let missing = PathBuf::from("/nonexistent/staging/path/for/test");
|
||||
let cleaned = clean_staging_orphans(&missing, Duration::from_secs(60)).unwrap();
|
||||
assert_eq!(cleaned, 0);
|
||||
}
|
||||
}
|
||||
@ -11,8 +11,8 @@ use macros::log;
|
||||
use crate::adapter::ebpf::EbpfServices;
|
||||
use crate::adapter::http::{
|
||||
acl, api_keys, audit as audit_api, auth, default, filter, fusion, health as health_api, logs as logs_api, ml,
|
||||
notification as notification_api, rate_limit as rate_limit_api, report as report_api, setup as setup_api, soar,
|
||||
stats, system as system_api,
|
||||
model_upload, notification as notification_api, rate_limit as rate_limit_api, report as report_api,
|
||||
setup as setup_api, soar, stats, system as system_api,
|
||||
};
|
||||
use crate::adapter::persistence::Database;
|
||||
use crate::adapter::websocket::routes as ws;
|
||||
@ -281,6 +281,7 @@ pub async fn run(params: HttpServerParams) -> Result<(), Error> {
|
||||
.service(stats::initialize())
|
||||
.service(health_api::initialize())
|
||||
.service(ml::initialize())
|
||||
.service(model_upload::initialize())
|
||||
.service(fusion::initialize())
|
||||
.service(system_api::initialize())
|
||||
.service(soar::initialize())
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
@ -14,6 +15,7 @@ use tokio::sync::oneshot;
|
||||
use tokio::time::{interval, sleep};
|
||||
|
||||
use crate::adapter::ebpf::EbpfServices;
|
||||
use crate::adapter::http::model_upload;
|
||||
use crate::adapter::persistence::Database;
|
||||
use crate::core::acl_service::AclService;
|
||||
use crate::core::auth::jwt::JwtService;
|
||||
@ -44,6 +46,7 @@ use crate::infrastructure::suricata_monitor::SuricataMonitor;
|
||||
use crate::interface::port::audit::AuditRepo;
|
||||
use crate::interface::port::setting::SettingRepo;
|
||||
use crate::interface::port::stats::StatsRepo;
|
||||
use crate::model::config::constants::{MODELS_DIR, STAGING_SUBDIR};
|
||||
use crate::model::detection::ml_detection::AlertMessage;
|
||||
use crate::model::error::Error;
|
||||
use crate::model::error::system::SystemError;
|
||||
@ -156,6 +159,19 @@ impl System {
|
||||
pub async fn run(&mut self) -> Result<ShutdownMode, Error> {
|
||||
log!(SystemLog::Initializing);
|
||||
|
||||
// Sweep model-upload staging directories left over from failed
|
||||
// uploads before the model watcher starts listening. A stale
|
||||
// `.staging/<uuid>/` would otherwise outlive restarts and eat
|
||||
// disk if the admin repeatedly aborted uploads mid-stream.
|
||||
{
|
||||
let staging_root = PathBuf::from(MODELS_DIR).join(STAGING_SUBDIR);
|
||||
match model_upload::clean_staging_orphans(&staging_root, Duration::from_secs(3600)) {
|
||||
Ok(0) => {}
|
||||
Ok(n) => log!(SystemLog::StagingOrphansCleaned(n as u64)),
|
||||
Err(e) => log!(SystemLog::StagingOrphansSweepFailed(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
// ML source state snapshot. Day 1 with no manifest renders as Dormant
|
||||
// — the rest of the stack still runs (3-source fusion).
|
||||
{
|
||||
|
||||
@ -155,5 +155,11 @@ loggable! {
|
||||
|
||||
#[error("HTML report generated at {path}")]
|
||||
HtmlReportGenerated { path: String } => tracing::Level::INFO,
|
||||
|
||||
#[error("Cleaned {count} stale model-upload staging directories")]
|
||||
StagingOrphansCleaned { count: u64 } => tracing::Level::INFO,
|
||||
|
||||
#[error("Staging-orphan sweep failed: {error}")]
|
||||
StagingOrphansSweepFailed { error: String } => tracing::Level::WARN,
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user