|
| 1 | +"""Gate your model load: verify the weights before you load them. |
| 2 | +
|
| 3 | +A drop-in load guard. Before a checkpoint is loaded, prove two things: the |
| 4 | +manifest that certifies it is jointly signed (builder + custodian), and the |
| 5 | +bytes on disk hash to exactly what that manifest binds. A tampered or swapped |
| 6 | +fork is refused BEFORE it ever reaches your loader. |
| 7 | +
|
| 8 | +This is the integrity and provenance gate (Layer 1): "is this the checkpoint the |
| 9 | +builder shipped." Honest scope: it catches tampering and swaps against software |
| 10 | +and remote adversaries; it is accountability-grade, not silicon-proof, against an |
| 11 | +operator who physically owns the hardware. |
| 12 | +
|
| 13 | + pip install weight-custody-manifest |
| 14 | + python load_guard.py # offline demo: certified loads, tampered refused |
| 15 | + python load_guard.py --load --model model.safetensors # gate + actually load a real file |
| 16 | +""" |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import argparse |
| 20 | +import hashlib |
| 21 | +import pathlib |
| 22 | +import tempfile |
| 23 | + |
| 24 | +from wcm import ( |
| 25 | + Ed25519Signer, |
| 26 | + VerificationContext, |
| 27 | + WeightCustodyManifest, |
| 28 | + generate_ed25519, |
| 29 | + verify_manifest, |
| 30 | +) |
| 31 | + |
| 32 | + |
| 33 | +def sha256_bytes(data: bytes) -> str: |
| 34 | + return "sha256:" + hashlib.sha256(data).hexdigest() |
| 35 | + |
| 36 | + |
| 37 | +def sha256_file(path: pathlib.Path) -> str: |
| 38 | + h = hashlib.sha256() |
| 39 | + with open(path, "rb") as f: |
| 40 | + for chunk in iter(lambda: f.read(1 << 20), b""): |
| 41 | + h.update(chunk) |
| 42 | + return "sha256:" + h.hexdigest() |
| 43 | + |
| 44 | + |
| 45 | +def banner(text: str) -> None: |
| 46 | + print(f"\n{'=' * 64}\n{text}\n{'=' * 64}") |
| 47 | + |
| 48 | + |
| 49 | +def build_manifest(weights_hash: str, org: str) -> dict: |
| 50 | + serving = sha256_bytes(b"vllm + policy-bundle (the builder-signed serving stack)") |
| 51 | + return { |
| 52 | + "manifest_version": "0.1", |
| 53 | + "weights_hash": weights_hash, |
| 54 | + "builder": {"identity": org, "signing_key": "ed25519:demo"}, |
| 55 | + "release_terms": { |
| 56 | + "license": "Frontier-Model-License", |
| 57 | + "permitted_derivatives": "fine-tune-only", |
| 58 | + "derivatives": "fine-tune-only", |
| 59 | + "permitted_environments": ["enterprise-governed-enclave"], |
| 60 | + }, |
| 61 | + "release_policy": { |
| 62 | + "required_assurance_tier": "hardware-attested", |
| 63 | + "trusted_time_source": "secure-tsc", |
| 64 | + "required_hw_platform": ["amd-sev-snp", "nvidia-cc-gpu"], |
| 65 | + "required_gpu_measurement": {"rim_pin": "nvidia-rim:golden"}, |
| 66 | + "required_serving_image": { |
| 67 | + "signer": "ed25519:demo", |
| 68 | + "release_rule": "prefer-current", |
| 69 | + "accepted_measurements": [{"measurement": serving, "status": "current"}], |
| 70 | + }, |
| 71 | + "attestation_revocation_check": "live-per-release, max-cache-age: short-window", |
| 72 | + "revocation_authority": "builder-and-opaque-joint", |
| 73 | + }, |
| 74 | + "custody": { |
| 75 | + "custodian": org, |
| 76 | + "custodian_type": "customer-self-custody", |
| 77 | + "kbs_image": {"measurement": sha256_bytes(b"reference-kbs-image"), "signer": "ed25519:demo"}, |
| 78 | + "enclave_id": "did:example:enclave-01", |
| 79 | + "attestation_cadence": "1h", |
| 80 | + }, |
| 81 | + "base_confidentiality": "gated-open", |
| 82 | + "deployment_model": "builder-to-customer", |
| 83 | + } |
| 84 | + |
| 85 | + |
| 86 | +class RefusedToLoad(Exception): |
| 87 | + """The checkpoint failed verification and must not be loaded.""" |
| 88 | + |
| 89 | + |
| 90 | +def guarded_load(model_path, manifest, ctx, *, do_load: bool = False) -> str: |
| 91 | + """Verify the manifest and the bytes on disk BEFORE loading; refuse on any mismatch. |
| 92 | +
|
| 93 | + The whole point: both checks run before a single byte reaches your loader. |
| 94 | + Returns the verified digest, or raises RefusedToLoad. |
| 95 | + """ |
| 96 | + model_path = pathlib.Path(model_path) |
| 97 | + |
| 98 | + # 1. The manifest itself must be jointly signed and valid. |
| 99 | + if not verify_manifest(manifest, ctx).ok: |
| 100 | + raise RefusedToLoad("manifest signature invalid (not the certified manifest)") |
| 101 | + |
| 102 | + # 2. The bytes on disk must hash to exactly what the manifest binds. |
| 103 | + digest = sha256_file(model_path) |
| 104 | + if digest != manifest.weights_hash: |
| 105 | + raise RefusedToLoad( |
| 106 | + "weights hash mismatch\n" |
| 107 | + f" manifest binds : {manifest.weights_hash}\n" |
| 108 | + f" file on disk : {digest}" |
| 109 | + ) |
| 110 | + |
| 111 | + # Only now is it safe to load. |
| 112 | + if do_load: |
| 113 | + _real_load(model_path) |
| 114 | + return digest |
| 115 | + |
| 116 | + |
| 117 | +def _real_load(model_path: pathlib.Path) -> None: |
| 118 | + """Actually load the verified file (run-local; needs the infer extras).""" |
| 119 | + from safetensors import safe_open # lazy: pip install -r requirements-infer.txt |
| 120 | + |
| 121 | + with safe_open(str(model_path), framework="pt") as f: |
| 122 | + keys = list(f.keys()) |
| 123 | + print(f"loaded {len(keys)} tensors from {model_path.name} (verified first)") |
| 124 | + |
| 125 | + |
| 126 | +def _signed_manifest(weights_hash: str, org: str = "frontier-labs"): |
| 127 | + """Return (manifest, verification_context) for a jointly-signed manifest.""" |
| 128 | + builder, custodian = generate_ed25519(), generate_ed25519() |
| 129 | + manifest = WeightCustodyManifest.model_validate(build_manifest(weights_hash, org)) |
| 130 | + manifest = manifest.with_signatures([ |
| 131 | + Ed25519Signer(builder).sign(manifest.unsigned_dict(), role="builder", signer=org), |
| 132 | + Ed25519Signer(custodian).sign(manifest.unsigned_dict(), role="custodian", signer=org), |
| 133 | + ]) |
| 134 | + ctx = VerificationContext() |
| 135 | + ctx.add_key(builder.public_bytes) |
| 136 | + ctx.add_key(custodian.public_bytes) |
| 137 | + return manifest, ctx |
| 138 | + |
| 139 | + |
| 140 | +def demo() -> None: |
| 141 | + with tempfile.TemporaryDirectory() as tmp: |
| 142 | + tmp = pathlib.Path(tmp) |
| 143 | + weights = b"<the certified model weights the builder shipped>" |
| 144 | + certified = tmp / "model.certified.bin" |
| 145 | + certified.write_bytes(weights) |
| 146 | + |
| 147 | + # The builder signs a manifest binding this exact checkpoint. |
| 148 | + manifest, ctx = _signed_manifest(sha256_bytes(weights)) |
| 149 | + |
| 150 | + banner("1. The certified checkpoint. Verified, then loaded.") |
| 151 | + digest = guarded_load(certified, manifest, ctx) |
| 152 | + print("manifest signature :", "valid (builder + custodian)") |
| 153 | + print("weights hash :", digest[:23], "... matches the manifest") |
| 154 | + print("-> load proceeds") |
| 155 | + |
| 156 | + banner("2. A tampered fork. Refused before it loads.") |
| 157 | + b = bytearray(weights) |
| 158 | + b[0] ^= 0x01 # flip a single byte |
| 159 | + tampered = tmp / "model.tampered.bin" |
| 160 | + tampered.write_bytes(bytes(b)) |
| 161 | + try: |
| 162 | + guarded_load(tampered, manifest, ctx) |
| 163 | + print("-> load proceeds") # not reached |
| 164 | + except RefusedToLoad as exc: |
| 165 | + print("REFUSED:", exc) |
| 166 | + print("-> the loader never sees the bytes. One flipped byte is enough.") |
| 167 | + |
| 168 | + banner("Verify the weights before you load them.") |
| 169 | + print("Integrity gate only (Layer 1). Honest scope: this catches tampering and swaps") |
| 170 | + print("against software and remote adversaries. It is accountability-grade, not") |
| 171 | + print("silicon-proof, against an operator who physically owns the hardware.") |
| 172 | + |
| 173 | + |
| 174 | +def main() -> None: |
| 175 | + ap = argparse.ArgumentParser( |
| 176 | + description="Gate a model load on a signed Weight Custody Manifest." |
| 177 | + ) |
| 178 | + ap.add_argument("--model", help="path to a real checkpoint to gate (run-local)") |
| 179 | + ap.add_argument( |
| 180 | + "--load", action="store_true", |
| 181 | + help="actually load the file after verifying (needs --model and the infer extras)", |
| 182 | + ) |
| 183 | + args = ap.parse_args() |
| 184 | + |
| 185 | + if not args.model: |
| 186 | + demo() |
| 187 | + return |
| 188 | + |
| 189 | + # Real run: a builder signs a manifest over this file, then the guard runs |
| 190 | + # before the load. Swap in a manifest you received from the builder instead. |
| 191 | + path = pathlib.Path(args.model) |
| 192 | + manifest, ctx = _signed_manifest(sha256_file(path)) |
| 193 | + banner(f"Gating {path.name}") |
| 194 | + guarded_load(path, manifest, ctx, do_load=args.load) |
| 195 | + print("verified" + (" and loaded" if args.load else " (pass --load to load it)")) |
| 196 | + |
| 197 | + |
| 198 | +if __name__ == "__main__": |
| 199 | + main() |
0 commit comments