Hironabe333/mlflow-registry-artifact-integrity-gap-poc

MLflow Model Registry — Artifact Integrity Gap PoC

This repository contains a proof-of-concept demonstrating that MLflow's Model Registry records model_size_bytes in the MLmodel manifest at log time, but mlflow.pyfunc.load_model() does not verify this recorded size or any artifact content hash before loading a registered model artifact.

Finding Description

When a model is registered in the MLflow Model Registry:

• The MLmodel manifest records model_size_bytes (e.g., 1386)
• MlflowClient().get_model_version() returns provenance metadata (run_id, source, status, version, model_id) — no hash, checksum, or digest field
• If the registered artifact file (model.pkl) is replaced in the artifact store, the MLmodel manifest retains the original model_size_bytes (now stale), and the registry metadata remains unchanged
• mlflow.pyfunc.load_model() loads the replaced artifact without reading model_size_bytes, without computing a content hash, and without emitting any warning

Key Evidence

MODEL_SIZE_BYTES_RECORDED=1386       ← in MLmodel at log time
MODEL_PKL_SIZE_AFTER=47              ← actual file after replacement  
MODEL_SIZE_BYTES_STALE_AFTER_TAMPER=True
MODEL_SIZE_BYTES_VERIFIED_ON_LOAD=False
HASH_FIELDS_IN_MODEL_VERSION=NONE
REGISTRY_METADATA_UNCHANGED=True
BASELINE_PREDICTION=[0.0]
TAMPERED_PREDICTION=[1.0]
DECISION_FLIP_CONFIRMED=True
WARNING_EMITTED=False
INTEGRITY_GAP_COMPLETE=True

MLmodel_before.yaml and MLmodel_after.yaml have identical SHA256 — the manifest is not updated after artifact replacement, confirming the staleness of model_size_bytes.

Reproduction

pip install mlflow>=3.0.0 scikit-learn>=1.0.0 numpy>=1.20.0

python create_mlflow_registry_gap.py --outdir artifacts --workdir work
python inspect_mlflow_registry_metadata.py --artifacts artifacts --workdir work
python reproduce_mlflow_load_model.py --artifacts artifacts --workdir work

Files

File	Description
create_mlflow_registry_gap.py	Registers benign model, replaces artifact, confirms staleness
inspect_mlflow_registry_metadata.py	Confirms zero hash fields; stale model_size_bytes
reproduce_mlflow_load_model.py	Loads replaced artifact; confirms prediction flip, no warning
requirements.txt	Python dependencies
expected_output.txt	Expected key output values
SHA256SUMS_T1.txt	SHA256 hashes for all PoC files
artifacts/baseline_model.pkl	Original registered artifact (predicts 0.0)
artifacts/replacement_model.pkl	Replacement artifact (predicts 1.0)
artifacts/MLmodel_before.yaml	MLmodel manifest before replacement
artifacts/MLmodel_after.yaml	MLmodel manifest after replacement (unchanged, model_size_bytes stale)
artifacts/registry_metadata_before.json	Registry metadata before replacement
artifacts/registry_metadata_after.json	Registry metadata after replacement (unchanged)
artifacts/prediction_results.json	Prediction results JSON

Notes

• Artifact store write access is required to reproduce this finding.
• This is not a pickle deserialization vulnerability.
• This is not a scanner bypass finding.
• model_id (m-{hex}) is a random identifier, not a content hash.
• Tested with MLflow 3.11.1.