Spaces:
Running
fix(F-14b): add pickle/model integrity checking mechanism
Browse filesAdds backend/core/model_integrity.py: a SHA-256 checksum gate in front
of pickle.load() calls, checked against data/reference/known_hashes.json
-- pinned in this repo, not fetched from the same untrusted channel the
model files themselves come from (main.py's startup step downloads them
from a Hugging Face Space named by the attacker-influenceable SPACE_ID
env var, with no verification before unpickling).
Honest about current scope: this session's uploaded snapshot only has
Git-LFS pointer stubs for data/reference/*.pkl (132-byte text files),
so there are no real files to hash yet -- known_hashes.json starts
empty and verify_integrity() logs a loud once-per-file warning rather
than hard-failing (refusing to start with zero pinned hashes would
break every deployment today). scripts/generate_model_hashes.py lets
you populate real hashes once you have real model files; from that
point on, any mismatch fail-closes automatically with no further code
changes. This commit adds the mechanism only -- wiring it into the
actual pickle.load() call sites is the next commits.
|
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Model/reference-file integrity checking (F-14).
|
| 3 |
+
|
| 4 |
+
pickle.load() on a file whose provenance is "whatever a remote channel
|
| 5 |
+
currently contains" is a real code-execution surface if that channel is
|
| 6 |
+
ever compromised -- and main.py's startup step downloads
|
| 7 |
+
own_embedding_model.pt, clip_database.pkl, own_centroids.pkl, and
|
| 8 |
+
ensemble_xgb.pkl from a Hugging Face Space named by the SPACE_ID
|
| 9 |
+
environment variable, with no verification before those .pkl files are
|
| 10 |
+
unpickled. SPACE_ID is attacker-influenceable in a misconfigured
|
| 11 |
+
deployment.
|
| 12 |
+
|
| 13 |
+
This module provides a SHA-256 checksum gate in front of every
|
| 14 |
+
pickle.load() call site, checked against a value pinned in THIS repo
|
| 15 |
+
(KNOWN_HASHES below / data/reference/known_hashes.json) -- not fetched
|
| 16 |
+
from the same untrusted channel as the files themselves, which is the
|
| 17 |
+
whole point.
|
| 18 |
+
|
| 19 |
+
Current status (see REMAINING_FIXES.md): KNOWN_HASHES starts empty.
|
| 20 |
+
This session's uploaded snapshot only contains Git-LFS pointer stubs
|
| 21 |
+
for data/reference/*.pkl (132-byte text files, not real model data), so
|
| 22 |
+
there is nothing real to hash yet. Once real model files exist in your
|
| 23 |
+
environment, run:
|
| 24 |
+
|
| 25 |
+
python scripts/generate_model_hashes.py
|
| 26 |
+
|
| 27 |
+
...to populate data/reference/known_hashes.json with real SHA-256
|
| 28 |
+
values, then commit that file. Until it's populated, verify_integrity()
|
| 29 |
+
logs a loud warning (once per file) instead of hard-failing -- refusing
|
| 30 |
+
to start the app with no pinned hashes at all would break every
|
| 31 |
+
deployment today, which is worse than the gap it would close. Once
|
| 32 |
+
real hashes are committed, this becomes fail-closed automatically with
|
| 33 |
+
no further code changes: any mismatch raises ModelIntegrityError.
|
| 34 |
+
"""
|
| 35 |
+
import hashlib
|
| 36 |
+
import json
|
| 37 |
+
from pathlib import Path
|
| 38 |
+
from typing import Optional
|
| 39 |
+
|
| 40 |
+
from backend.core.logger import setup_logger
|
| 41 |
+
|
| 42 |
+
logger = setup_logger(__name__)
|
| 43 |
+
|
| 44 |
+
_HASHES_PATH = Path(__file__).parent.parent.parent / "data" / "reference" / "known_hashes.json"
|
| 45 |
+
_warned_files: set = set()
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
class ModelIntegrityError(Exception):
|
| 49 |
+
"""Raised when a reference/model file's SHA-256 does not match the
|
| 50 |
+
value pinned in data/reference/known_hashes.json."""
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _load_known_hashes() -> dict:
|
| 54 |
+
if not _HASHES_PATH.exists():
|
| 55 |
+
return {}
|
| 56 |
+
try:
|
| 57 |
+
with open(_HASHES_PATH, "r", encoding="utf-8") as f:
|
| 58 |
+
return json.load(f)
|
| 59 |
+
except Exception as e:
|
| 60 |
+
logger.warning(f"Could not parse {_HASHES_PATH}: {e}")
|
| 61 |
+
return {}
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def _sha256_of(path: Path) -> str:
|
| 65 |
+
h = hashlib.sha256()
|
| 66 |
+
with open(path, "rb") as f:
|
| 67 |
+
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
| 68 |
+
h.update(chunk)
|
| 69 |
+
return h.hexdigest()
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def verify_integrity(path: Path, known_hashes: Optional[dict] = None) -> None:
|
| 73 |
+
"""Verify path's SHA-256 against the pinned value for its filename.
|
| 74 |
+
|
| 75 |
+
- No pinned hash on file for this filename -> log a warning once,
|
| 76 |
+
then allow the load (see module docstring for why this isn't a
|
| 77 |
+
hard failure yet).
|
| 78 |
+
- Pinned hash present and it matches -> silent success.
|
| 79 |
+
- Pinned hash present and it does NOT match -> raises
|
| 80 |
+
ModelIntegrityError. Callers should let this propagate (do not
|
| 81 |
+
swallow it into the generic "file missing/corrupt, fall back to
|
| 82 |
+
neutral" except-branches these call sites already have) since a
|
| 83 |
+
hash mismatch is a materially different, higher-severity signal
|
| 84 |
+
than "file absent" or "file is an LFS pointer stub".
|
| 85 |
+
"""
|
| 86 |
+
known = known_hashes if known_hashes is not None else _load_known_hashes()
|
| 87 |
+
expected = known.get(path.name)
|
| 88 |
+
if not expected:
|
| 89 |
+
if path.name not in _warned_files:
|
| 90 |
+
_warned_files.add(path.name)
|
| 91 |
+
logger.warning(
|
| 92 |
+
"No pinned SHA-256 for %s in %s -- integrity NOT verified "
|
| 93 |
+
"before loading. Run scripts/generate_model_hashes.py once "
|
| 94 |
+
"you have real model files, then commit the result.",
|
| 95 |
+
path.name, _HASHES_PATH,
|
| 96 |
+
)
|
| 97 |
+
return
|
| 98 |
+
actual = _sha256_of(path)
|
| 99 |
+
if actual != expected:
|
| 100 |
+
raise ModelIntegrityError(
|
| 101 |
+
f"{path.name}: SHA-256 mismatch. Expected {expected}, got {actual}. "
|
| 102 |
+
f"Refusing to load -- this file's contents do not match what is "
|
| 103 |
+
f"pinned in {_HASHES_PATH}."
|
| 104 |
+
)
|
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{}
|
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Generate data/reference/known_hashes.json from the real reference/model
|
| 3 |
+
files currently on disk (F-14).
|
| 4 |
+
|
| 5 |
+
Run this once you have real (non-Git-LFS-stub) copies of:
|
| 6 |
+
data/reference/own_embedding_model.pt
|
| 7 |
+
data/reference/clip_database.pkl
|
| 8 |
+
data/reference/own_centroids.pkl
|
| 9 |
+
data/reference/ensemble_xgb.pkl
|
| 10 |
+
|
| 11 |
+
...then commit the resulting known_hashes.json. From that point on,
|
| 12 |
+
backend/core/model_integrity.py's verify_integrity() fail-closes on any
|
| 13 |
+
mismatch instead of just logging a warning -- no code changes needed,
|
| 14 |
+
it reads this file automatically.
|
| 15 |
+
|
| 16 |
+
Usage:
|
| 17 |
+
python scripts/generate_model_hashes.py
|
| 18 |
+
"""
|
| 19 |
+
import hashlib
|
| 20 |
+
import json
|
| 21 |
+
import sys
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
|
| 24 |
+
REPO_ROOT = Path(__file__).parent.parent
|
| 25 |
+
REF_DIR = REPO_ROOT / "data" / "reference"
|
| 26 |
+
HASHES_PATH = REF_DIR / "known_hashes.json"
|
| 27 |
+
|
| 28 |
+
# Keep in sync with backend/main.py's _model_files list.
|
| 29 |
+
TRACKED_FILES = [
|
| 30 |
+
"own_embedding_model.pt",
|
| 31 |
+
"clip_database.pkl",
|
| 32 |
+
"own_centroids.pkl",
|
| 33 |
+
"ensemble_xgb.pkl",
|
| 34 |
+
]
|
| 35 |
+
|
| 36 |
+
_LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec"
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def sha256_of(path: Path) -> str:
|
| 40 |
+
h = hashlib.sha256()
|
| 41 |
+
with open(path, "rb") as f:
|
| 42 |
+
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
| 43 |
+
h.update(chunk)
|
| 44 |
+
return h.hexdigest()
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def main() -> int:
|
| 48 |
+
hashes = {}
|
| 49 |
+
if HASHES_PATH.exists():
|
| 50 |
+
try:
|
| 51 |
+
hashes = json.loads(HASHES_PATH.read_text(encoding="utf-8"))
|
| 52 |
+
except Exception:
|
| 53 |
+
hashes = {}
|
| 54 |
+
|
| 55 |
+
skipped = []
|
| 56 |
+
updated = []
|
| 57 |
+
for fname in TRACKED_FILES:
|
| 58 |
+
path = REF_DIR / fname
|
| 59 |
+
if not path.exists():
|
| 60 |
+
skipped.append((fname, "not found"))
|
| 61 |
+
continue
|
| 62 |
+
with open(path, "rb") as f:
|
| 63 |
+
head = f.read(len(_LFS_POINTER_PREFIX))
|
| 64 |
+
if head == _LFS_POINTER_PREFIX:
|
| 65 |
+
skipped.append((fname, "still a Git-LFS pointer stub -- run `git lfs pull` first"))
|
| 66 |
+
continue
|
| 67 |
+
hashes[fname] = sha256_of(path)
|
| 68 |
+
updated.append(fname)
|
| 69 |
+
|
| 70 |
+
HASHES_PATH.parent.mkdir(parents=True, exist_ok=True)
|
| 71 |
+
HASHES_PATH.write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 72 |
+
|
| 73 |
+
for fname in updated:
|
| 74 |
+
print(f" hashed: {fname}")
|
| 75 |
+
for fname, reason in skipped:
|
| 76 |
+
print(f" skipped: {fname} ({reason})")
|
| 77 |
+
print(f"\nWrote {HASHES_PATH} with {len(hashes)} entries.")
|
| 78 |
+
if skipped:
|
| 79 |
+
print(
|
| 80 |
+
f"\n{len(skipped)} file(s) were skipped -- those files will still "
|
| 81 |
+
f"load with only a warning (no integrity check) until you re-run "
|
| 82 |
+
f"this script after obtaining the real files."
|
| 83 |
+
)
|
| 84 |
+
return 0
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
if __name__ == "__main__":
|
| 88 |
+
sys.exit(main())
|