flax-subfolder-shard-traversal-poc

Transformers Flax Sharded Checkpoint Index Escapes Requested Subfolder

Severity

Medium, 6.7/10.

Rationale: this is artifact-carried model weight and inference-output manipulation through a normal FlaxBertModel.from_pretrained(..., subfolder="model") load path. It is not arbitrary code execution, and it is scoped to the deprecated-but-still-used Transformers 4.x Flax/JAX loader line, but the uploaded model artifact is self-contained and the loader reads a shard outside the requested model subfolder.

Summary

This repository contains a benign proof of concept for a Transformers 4.x Flax sharded checkpoint path traversal.

The intended load target is the model/ subfolder:

FlaxBertModel.from_pretrained(source, subfolder="model")

That subfolder contains config.json, flax_model.msgpack.index.json, and a decoy shard:

model/config.json
model/flax_model.msgpack.index.json
model/decoy_inside_shard.msgpack

However, the index maps every weight to ../outside_shard.msgpack:

"embeddings/LayerNorm/bias": "../outside_shard.msgpack"

The runtime follows that relative shard path and loads the root-level shard:

outside_shard.msgpack

The decoy shard inside the requested subfolder contains marker value -31337.0. The escaped root-level shard contains marker value 4242.0. Runtime loads 4242.0, proving the loader read outside the requested subfolder.

No code execution payload is included. The PoC only demonstrates deterministic model weight and output manipulation.

Impact

• Artifact-carried subfolder boundary bypass through the Flax sharded checkpoint index.
• Reproduces through normal FlaxBertModel.from_pretrained(source, subfolder="model").
• The runtime-loaded parameter comes from a shard outside the requested subfolder.
• Inference changes: the verifier records a first-eight-output max absolute delta of about 3.46409 between the in-subfolder decoy parameters and the runtime-loaded escaped-shard parameters.
• ModelScan 0.8.8 and PickleScan 1.0.4 report no dangerous findings in this environment.

Affected Versions Tested

• Python 3.12.3
• transformers==4.57.6
• huggingface_hub==0.36.2
• flax==0.12.7
• jax==0.10.0
• jaxlib==0.10.0
• modelscan==0.8.8
• picklescan==1.0.4

Transformers v5 has removed the JAX/Flax backend. This report is therefore scoped to the latest tested 4.x Flax-capable loader, not the v5 API line.

Files

outside_shard.msgpack
  Root-level Flax msgpack shard with marker value 4242.0.

model/config.json
  Tiny Flax BERT config loaded from the requested subfolder.

model/flax_model.msgpack.index.json
  Sharded Flax checkpoint index inside the requested subfolder. It maps all
  weights to ../outside_shard.msgpack.

model/decoy_inside_shard.msgpack
  In-subfolder decoy shard with marker value -31337.0.

verify_flax_subfolder_shard_traversal_poc.py
  End-to-end verifier. Supports both local paths and direct Hugging Face repo
  IDs as --source.

scripts/build_flax_shard_traversal.py
  Original artifact builder.

scripts/original_inference_effect.py
  Original lab script showing inference output change.

evidence/fresh_verify.json
  Fresh local verification output from this staged repository.

evidence/hub_direct_verify.json
  Verification output after upload using the public Hugging Face repo ID as the
  verifier source.

evidence/lab_marker_verify.json
  Original lab marker-value verification.

evidence/lab_inference_effect.json
  Original lab inference-output verification.

evidence/sha256.txt
  Hashes for uploaded files.

Reproduction

Install matching dependencies in an isolated environment:

pip install "transformers==4.57.6" "huggingface_hub==0.36.2" "flax==0.12.7" "jax==0.10.0" "jaxlib==0.10.0" "modelscan==0.8.8" "picklescan==1.0.4"

Run against a local checkout:

python verify_flax_subfolder_shard_traversal_poc.py

Or run directly against the uploaded Hugging Face repo:

python verify_flax_subfolder_shard_traversal_poc.py --source pragnyanramtha/flax-subfolder-shard-traversal-poc

Expected output highlights:

{
  "subfolder": "model",
  "target_key": "embeddings/LayerNorm/bias",
  "weight_map_unique_values": ["../outside_shard.msgpack"],
  "decoy_marker_value": -31337.0,
  "outside_marker_value": 4242.0,
  "runtime_loaded_marker_value": 4242.0,
  "first_8_max_abs_delta_decoy_vs_runtime": 3.464094400405884,
  "verified": true
}

The critical observation is that runtime_loaded_marker_value comes from outside_shard.msgpack, even though the caller requested subfolder="model".

Scanner Behavior

ModelScan 0.8.8 reports zero issues for the requested model/ subfolder:

{
  "total_issues": 0,
  "scanned": {"total_scanned": 0},
  "skipped": {
    "total_skipped": 3,
    "skipped_files": [
      {"source": "decoy_inside_shard.msgpack"},
      {"source": "flax_model.msgpack.index.json"},
      {"source": "config.json"}
    ]
  }
}

ModelScan also reports zero issues for the repository root in this environment, and PickleScan reports zero infected files and zero dangerous globals.

This scanner output is supporting context. The core issue is the runtime loader following model-controlled shard paths outside the requested subfolder and using those tensors for inference.

Root Cause

The Flax sharded checkpoint index is read from the requested subfolder, but the weight_map shard filenames are joined with the requested model path without enforcing that the normalized result stays inside the requested subfolder.

In this PoC:

requested model directory:
  model/

index location:
  model/flax_model.msgpack.index.json

index shard filename:
  ../outside_shard.msgpack

runtime-used shard:
  outside_shard.msgpack

This lets a model artifact make the loader consume parameters outside the subfolder the caller intended to load.

Safety Notes

• No arbitrary code execution is used.
• No network callbacks, credential access, persistence, or destructive payloads are present.
• Marker values are numeric tensor values chosen only to make the loader mismatch obvious.

Limitations

• This is output/weight manipulation, not ACE.
• The validated affected loader is Transformers 4.x Flax/JAX support. The Transformers v5 API line removed the Flax backend.
• The traversal is bounded by the model repository or snapshot layout in the self-contained Hugging Face artifact.

Duplicate Checks

Local research did not find an obvious published advisory for this specific Flax sharded checkpoint subfolder traversal. This is distinct from generic unsafe pickle loading because the payload is a Flax msgpack shard and the behavior is triggered by flax_model.msgpack.index.json.