Kanisia/jax-orbax-metadata-absolute-path-read-poc

JAX / Orbax PyTree metadata absolute-path array read PoC

This CPU-only proof of concept demonstrates that an untrusted legacy Orbax PyTree checkpoint can redirect a default restore to a readable Zarr-v2 array directory outside the checkpoint root.

The malicious input is the generated malicious_checkpoint directory. Its _METADATA file retains the normal outer tree_metadata entry for safe; only key_metadata[0].key is replaced with a runtime-generated absolute path. The local malicious_checkpoint/safe directory is then removed. A default restore therefore returns the array from the external directory rather than data inside the checkpoint.

Preconditions

• CPU only; no GPU is required.
• Orbax restores the checkpoint without an explicit expected PyTree item.
• The target process can read the selected external Zarr / Orbax array directory.
• This PoC demonstrates external array-directory disclosure, not arbitrary byte/text-file disclosure, code execution, or a relative traversal.

Reproduction

Use Python 3.10+ in a clean virtual environment:

python -m venv .venv
.venv\Scripts\python.exe -m pip install -r requirements.txt
.venv\Scripts\python.exe build_and_verify.py

Expected result:

CONTROL_DEFAULT_RESTORE=[7]
PYTREE_DEFAULT_RESTORE=[31337]
STANDARD_DEFAULT_RESTORE=[31337]
LOCAL_SAFE_ARRAY_PRESENT=False
VERIFIED

The builder regenerates the checkpoint before every run. That is required because the working payload contains an absolute path, which must match the verifier's local fixture location. The generated verify_report.json records versions, CPU devices, and observed values without storing the local path.

Root cause

During default restore, Orbax reconstructs the PyTree from the checkpoint's _METADATA JSON. The file-controlled key becomes a parameter name. For a legacy checkpoint (no manifest.ocdbt), Orbax passes that name to os.path.join(checkpoint_directory, name) when building the TensorStore KvStore. Python discards the checkpoint directory when name is absolute, so TensorStore opens the external Zarr array.

An explicit expected target item rejects this altered metadata tree. The vulnerable load path is therefore the normal metadata-driven default restore:

import orbax.checkpoint as ocp
restored = ocp.PyTreeCheckpointer().restore("malicious_checkpoint")

Suggested fix

Reject absolute paths in checkpoint-derived parameter names before building a KvStore spec. For legacy checkpoints, resolve the candidate path and require it to remain contained within the trusted checkpoint directory. Apply the same validation to all metadata-derived TensorStore paths.