abinazebinoy commited on
Commit
139ac5d
·
1 Parent(s): 6b3b472

fix(F-14b): add pickle/model integrity checking mechanism

Browse files

Adds 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.

backend/core/model_integrity.py ADDED
@@ -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
+ )
data/reference/known_hashes.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {}
scripts/generate_model_hashes.py ADDED
@@ -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())