Add integrity verification for ML model artifacts before joblib.load #66

New issue

Closed

opened 2026-07-04 01:30:03 +00:00 by coding-agent-marvin8 · 1 comment

coding-agent-marvin8 commented

2026-07-04 01:30:03 +00:00

(Migrated from codeberg.org)

Problem

load_artifacts in services/ml_inference.py calls joblib.load on files from the models directory without any integrity check. joblib uses pickle internally, so a malicious .pkl file in the models directory would execute arbitrary code when the service starts or after nightly retraining.

Flagged by WuMing in PR #65 (inline comment on ml_inference.py).

Threat model

An attacker needs write access to the models volume to exploit this. In a typical self-hosted deployment that level of access already implies container exec or host FS access, so the blast radius is bounded — but the risk is still worth mitigating explicitly.

Proposed fix

Option A — SHA-256 sidecar (preferred): after save_artifacts writes model.pkl and pipeline.pkl, write a manifest.sha256 file containing their hex digests. Before load_artifacts deserialises anything, recompute and compare. Detects any tampering between write and load.

Option B — File permissions: document and enforce that the models directory is owned by the service user with chmod 700. Mount the volume read-only at container start; remount rw only during the nightly retraining window.

A combination of both is ideal: permissions reduce the attack surface; the sidecar provides a last-line-of-defence integrity check at deserialisation time.

Out of scope

SafeTensors: not applicable (serialises tensor buffers, not sklearn Pipeline objects)
ONNX conversion: technically possible via skl2onnx but adds significant dependency weight and TF-IDF vocabulary round-trip complexity — disproportionate for this project

## Problem `load_artifacts` in `services/ml_inference.py` calls `joblib.load` on files from the models directory without any integrity check. joblib uses pickle internally, so a malicious `.pkl` file in the models directory would execute arbitrary code when the service starts or after nightly retraining. Flagged by WuMing in PR #65 (inline comment on `ml_inference.py`). ## Threat model An attacker needs write access to the models volume to exploit this. In a typical self-hosted deployment that level of access already implies container exec or host FS access, so the blast radius is bounded — but the risk is still worth mitigating explicitly. ## Proposed fix **Option A — SHA-256 sidecar (preferred):** after `save_artifacts` writes `model.pkl` and `pipeline.pkl`, write a `manifest.sha256` file containing their hex digests. Before `load_artifacts` deserialises anything, recompute and compare. Detects any tampering between write and load. **Option B — File permissions:** document and enforce that the models directory is owned by the service user with `chmod 700`. Mount the volume read-only at container start; remount rw only during the nightly retraining window. A combination of both is ideal: permissions reduce the attack surface; the sidecar provides a last-line-of-defence integrity check at deserialisation time. ## Out of scope - SafeTensors: not applicable (serialises tensor buffers, not sklearn Pipeline objects) - ONNX conversion: technically possible via `skl2onnx` but adds significant dependency weight and TF-IDF vocabulary round-trip complexity — disproportionate for this project

coding-agent-marvin8 commented

2026-07-04 04:37:52 +00:00

(Migrated from codeberg.org)

Exploration findings

The vulnerability is in fenliu/training/model.py. save_artifacts (line 58–59) writes xgboost_model.pkl and feature_pipeline.pkl via joblib.dump with no integrity record. load_artifacts (lines 79–80) deserialises them with joblib.load directly — no check that the files are what was written.

Plan

Option A — SHA-256 sidecar (implemented):

Add a private _write_manifest(output_dir) helper that iterates over the two pkl files, computes SHA-256 of each, and writes manifest.sha256 in the format <hex> <filename> (one line per file).
Call it at the end of save_artifacts.
Add a private _verify_manifest(output_dir) helper that reads manifest.sha256, recomputes digests of both pkl files, and raises ValueError on mismatch or FileNotFoundError if the manifest is absent.
Call it at the start of load_artifacts, before either joblib.load call.

Option B — volume permissions (documented):

Add a security note to docs/getting-started/container-deployment.md recommending chmod 700 on the models directory (owned by the fenliu service user) and, where the orchestrator supports it, mounting the volume read-only after training and switching to read-write only for the nightly retraining window.

TDD order: write the three failing tests first (manifest absent, file tampered, manifest exists after save), then implement both helpers.

Key files: packages/fenliu/src/fenliu/training/model.py, packages/fenliu/tests/test_training.py, packages/fenliu/docs/getting-started/container-deployment.md

## Exploration findings The vulnerability is in `fenliu/training/model.py`. `save_artifacts` (line 58–59) writes `xgboost_model.pkl` and `feature_pipeline.pkl` via `joblib.dump` with no integrity record. `load_artifacts` (lines 79–80) deserialises them with `joblib.load` directly — no check that the files are what was written. ## Plan **Option A — SHA-256 sidecar (implemented):** - Add a private `_write_manifest(output_dir)` helper that iterates over the two pkl files, computes SHA-256 of each, and writes `manifest.sha256` in the format `<hex> <filename>` (one line per file). - Call it at the end of `save_artifacts`. - Add a private `_verify_manifest(output_dir)` helper that reads `manifest.sha256`, recomputes digests of both pkl files, and raises `ValueError` on mismatch or `FileNotFoundError` if the manifest is absent. - Call it at the start of `load_artifacts`, before either `joblib.load` call. **Option B — volume permissions (documented):** Add a security note to `docs/getting-started/container-deployment.md` recommending `chmod 700` on the models directory (owned by the `fenliu` service user) and, where the orchestrator supports it, mounting the volume read-only after training and switching to read-write only for the nightly retraining window. **TDD order:** write the three failing tests first (manifest absent, file tampered, manifest exists after save), then implement both helpers. **Key files:** `packages/fenliu/src/fenliu/training/model.py`, `packages/fenliu/tests/test_training.py`, `packages/fenliu/docs/getting-started/container-deployment.md`

No labels

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

marvin8/dujiangyan#66

No description provided.

Rows
Columns