diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index acf993c..a3530c7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -176,7 +176,7 @@ jobs: - name: Run the feature examples run: | - for s in closed_model_e2e multi_stage_byom transparency_log post_quantum channel_binding revocation_kill_switch quote_verification; do + for s in closed_model_e2e multi_stage_byom transparency_log post_quantum channel_binding revocation_kill_switch quote_verification load_guard; do echo "== $s ==" python "$s.py" done diff --git a/weight-custody-manifest/README.md b/weight-custody-manifest/README.md index 605b0db..f0c41f7 100644 --- a/weight-custody-manifest/README.md +++ b/weight-custody-manifest/README.md @@ -34,6 +34,17 @@ The demos below go deeper on the same machinery. --- +## Gate your model load (drop-in) + +`load_guard.py` is the smallest useful integration: verify the checkpoint before your loader ever touches it. It checks the manifest's joint signature and that the bytes on disk hash to exactly what the builder bound, then refuses a tampered or swapped fork **before** the load. Drop `guarded_load()` in front of your existing `safetensors` or `transformers` load. + +```bash +python load_guard.py # offline: certified loads, tampered refused +python load_guard.py --load --model model.safetensors # gate + actually load a real file (run-local) +``` + +--- + ## The demos Each script is one runnable feature with mock (software) attestation, so they run anywhere. `refuse_and_wipe.py` above is the 30-second version. @@ -78,7 +89,7 @@ python real_open_model.py ## What runs in CI -Every offline example runs in CI against the published PyPI package and must exit 0: `refuse_and_wipe`, the closed/open e2e pair, `multi_stage_byom`, `revocation_kill_switch`, `channel_binding`, `sovereign_self_custody`, `transparency_log`, `post_quantum`, `quote_verification`, `snp_replay`, and `provenance_model_signing` (with the `[model-signing]` extra). Only `real_open_model.py` is excluded (it downloads a model). +Every offline example runs in CI against the published PyPI package and must exit 0: `refuse_and_wipe`, the closed/open e2e pair, `multi_stage_byom`, `revocation_kill_switch`, `channel_binding`, `sovereign_self_custody`, `transparency_log`, `post_quantum`, `quote_verification`, `load_guard`, `snp_replay`, and `provenance_model_signing` (with the `[model-signing]` extra). Only `real_open_model.py` is excluded (it downloads a model). ## Reference diff --git a/weight-custody-manifest/load_guard.py b/weight-custody-manifest/load_guard.py new file mode 100644 index 0000000..4342827 --- /dev/null +++ b/weight-custody-manifest/load_guard.py @@ -0,0 +1,199 @@ +"""Gate your model load: verify the weights before you load them. + +A drop-in load guard. Before a checkpoint is loaded, prove two things: the +manifest that certifies it is jointly signed (builder + custodian), and the +bytes on disk hash to exactly what that manifest binds. A tampered or swapped +fork is refused BEFORE it ever reaches your loader. + +This is the integrity and provenance gate (Layer 1): "is this the checkpoint the +builder shipped." Honest scope: it catches tampering and swaps against software +and remote adversaries; it is accountability-grade, not silicon-proof, against an +operator who physically owns the hardware. + + pip install weight-custody-manifest + python load_guard.py # offline demo: certified loads, tampered refused + python load_guard.py --load --model model.safetensors # gate + actually load a real file +""" +from __future__ import annotations + +import argparse +import hashlib +import pathlib +import tempfile + +from wcm import ( + Ed25519Signer, + VerificationContext, + WeightCustodyManifest, + generate_ed25519, + verify_manifest, +) + + +def sha256_bytes(data: bytes) -> str: + return "sha256:" + hashlib.sha256(data).hexdigest() + + +def sha256_file(path: pathlib.Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return "sha256:" + h.hexdigest() + + +def banner(text: str) -> None: + print(f"\n{'=' * 64}\n{text}\n{'=' * 64}") + + +def build_manifest(weights_hash: str, org: str) -> dict: + serving = sha256_bytes(b"vllm + policy-bundle (the builder-signed serving stack)") + return { + "manifest_version": "0.1", + "weights_hash": weights_hash, + "builder": {"identity": org, "signing_key": "ed25519:demo"}, + "release_terms": { + "license": "Frontier-Model-License", + "permitted_derivatives": "fine-tune-only", + "derivatives": "fine-tune-only", + "permitted_environments": ["enterprise-governed-enclave"], + }, + "release_policy": { + "required_assurance_tier": "hardware-attested", + "trusted_time_source": "secure-tsc", + "required_hw_platform": ["amd-sev-snp", "nvidia-cc-gpu"], + "required_gpu_measurement": {"rim_pin": "nvidia-rim:golden"}, + "required_serving_image": { + "signer": "ed25519:demo", + "release_rule": "prefer-current", + "accepted_measurements": [{"measurement": serving, "status": "current"}], + }, + "attestation_revocation_check": "live-per-release, max-cache-age: short-window", + "revocation_authority": "builder-and-opaque-joint", + }, + "custody": { + "custodian": org, + "custodian_type": "customer-self-custody", + "kbs_image": {"measurement": sha256_bytes(b"reference-kbs-image"), "signer": "ed25519:demo"}, + "enclave_id": "did:example:enclave-01", + "attestation_cadence": "1h", + }, + "base_confidentiality": "gated-open", + "deployment_model": "builder-to-customer", + } + + +class RefusedToLoad(Exception): + """The checkpoint failed verification and must not be loaded.""" + + +def guarded_load(model_path, manifest, ctx, *, do_load: bool = False) -> str: + """Verify the manifest and the bytes on disk BEFORE loading; refuse on any mismatch. + + The whole point: both checks run before a single byte reaches your loader. + Returns the verified digest, or raises RefusedToLoad. + """ + model_path = pathlib.Path(model_path) + + # 1. The manifest itself must be jointly signed and valid. + if not verify_manifest(manifest, ctx).ok: + raise RefusedToLoad("manifest signature invalid (not the certified manifest)") + + # 2. The bytes on disk must hash to exactly what the manifest binds. + digest = sha256_file(model_path) + if digest != manifest.weights_hash: + raise RefusedToLoad( + "weights hash mismatch\n" + f" manifest binds : {manifest.weights_hash}\n" + f" file on disk : {digest}" + ) + + # Only now is it safe to load. + if do_load: + _real_load(model_path) + return digest + + +def _real_load(model_path: pathlib.Path) -> None: + """Actually load the verified file (run-local; needs the infer extras).""" + from safetensors import safe_open # lazy: pip install -r requirements-infer.txt + + with safe_open(str(model_path), framework="pt") as f: + keys = list(f.keys()) + print(f"loaded {len(keys)} tensors from {model_path.name} (verified first)") + + +def _signed_manifest(weights_hash: str, org: str = "frontier-labs"): + """Return (manifest, verification_context) for a jointly-signed manifest.""" + builder, custodian = generate_ed25519(), generate_ed25519() + manifest = WeightCustodyManifest.model_validate(build_manifest(weights_hash, org)) + manifest = manifest.with_signatures([ + Ed25519Signer(builder).sign(manifest.unsigned_dict(), role="builder", signer=org), + Ed25519Signer(custodian).sign(manifest.unsigned_dict(), role="custodian", signer=org), + ]) + ctx = VerificationContext() + ctx.add_key(builder.public_bytes) + ctx.add_key(custodian.public_bytes) + return manifest, ctx + + +def demo() -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp = pathlib.Path(tmp) + weights = b"" + certified = tmp / "model.certified.bin" + certified.write_bytes(weights) + + # The builder signs a manifest binding this exact checkpoint. + manifest, ctx = _signed_manifest(sha256_bytes(weights)) + + banner("1. The certified checkpoint. Verified, then loaded.") + digest = guarded_load(certified, manifest, ctx) + print("manifest signature :", "valid (builder + custodian)") + print("weights hash :", digest[:23], "... matches the manifest") + print("-> load proceeds") + + banner("2. A tampered fork. Refused before it loads.") + b = bytearray(weights) + b[0] ^= 0x01 # flip a single byte + tampered = tmp / "model.tampered.bin" + tampered.write_bytes(bytes(b)) + try: + guarded_load(tampered, manifest, ctx) + print("-> load proceeds") # not reached + except RefusedToLoad as exc: + print("REFUSED:", exc) + print("-> the loader never sees the bytes. One flipped byte is enough.") + + banner("Verify the weights before you load them.") + print("Integrity gate only (Layer 1). Honest scope: this catches tampering and swaps") + print("against software and remote adversaries. It is accountability-grade, not") + print("silicon-proof, against an operator who physically owns the hardware.") + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Gate a model load on a signed Weight Custody Manifest." + ) + ap.add_argument("--model", help="path to a real checkpoint to gate (run-local)") + ap.add_argument( + "--load", action="store_true", + help="actually load the file after verifying (needs --model and the infer extras)", + ) + args = ap.parse_args() + + if not args.model: + demo() + return + + # Real run: a builder signs a manifest over this file, then the guard runs + # before the load. Swap in a manifest you received from the builder instead. + path = pathlib.Path(args.model) + manifest, ctx = _signed_manifest(sha256_file(path)) + banner(f"Gating {path.name}") + guarded_load(path, manifest, ctx, do_load=args.load) + print("verified" + (" and loaded" if args.load else " (pass --load to load it)")) + + +if __name__ == "__main__": + main()