[Security] Arbitrary host-file disclosure when loading an untrusted checkpoint: tensorstore file kvstore follows symlinked chunk files with no containment

[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

A Flax/Orbax checkpoint is a directory the publisher fully controls. Array leaves are stored as tensorstore zarr3 chunk files and read back through the tensorstore file kvstore, which follows filesystem symlinks with no containment to the checkpoint directory. If a chunk file inside a published checkpoint is a symlink to an arbitrary host path, calling restore_checkpoint reads that file's bytes directly into the restored array.

Loading a malicious checkpoint therefore silently discloses arbitrary host files (SSH private keys, cloud-credential files, /etc/passwd, tokens) into the model — no code execution required, no trust_remote_code-style opt-in required. This is the checkpoint-directory analog of the tar/zip symlink-traversal class (cf. CVE-2007-4559, Python tarfile filter='data').

Details

flax.training.checkpoints.restore_checkpoint dispatches modern checkpoints to Orbax:

restored = orbax_checkpointer.restore(ckpt_path, item=target, **restore_kwargs)

[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.

[vfdev-5]

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 (GHSA-gw9q-c7gh-j9vm family).

I'm still trying to understand the attack angle. If the victim reads a poisoned checkpoint, the model finally contains some data from the host machine. So, the model contains garbage values and does not do what it was intended for. tar/zip symlink traversal CVE tells that:

Directory traversal vulnerability in the (1) extract and (2) extractall functions in the tarfile module in Python allows user-assisted remote attackers to overwrite arbitrary files via a .. (dot dot) sequence in filenames in a TAR archive, a related issue to GHSA-p23r-j983-4557.

Can we also overwrite arbitrary files with the current flax code?

[geo-chen]

no, we can't. The vulnerability is read-only: an attacker-controlled checkpoint can exfiltrate arbitrary file contents into restored arrays, but cannot overwrite files on the host.

wrt garbage values, the attacker would symlink only one small array in a checkpoint that is otherwise a legitimate, working model. Exfiltration is the security impact here.

[vfdev-5]

Can we read other user's data or something that current non-root user can not read by default?

[geo-chen]

no, it's not privilege escalation, but rather, trust-boundary violation.

[vfdev-5]

os.symlink("/etc/passwd", chunk)

This is the important part of the poisoned checkpoint and if symlink does not persist after the distribution, the attack is not working, IMO.

On the other hand, I still do not understand how _restore_mpas is related to the issue. I run the debugger and see that

    restored = orbax_checkpointer.restore(
      ckpt_path, item=target, **restore_kwargs
    )

# {'w': array([116, 104, 105, 115,  45, 105, 115,  45, 115, 101,  99, 114, 101, 116,  10], dtype=uint8)}

is responsible to the checkpoint loading. If checkpoint is made with symlinks etc, flax can't prevent from loading them.

[geo-chen]

thanks for confirming the read - the array decodes to this-is-secret.

On _restore_mpas: you're right that it's not relevant here. That path was from an earlier observation about the //GDAPlaceholder: multiprocess-array format.
I've updated the issue to focus on symlinked chunk files being read directly by the tensorstore file kvstore, independently of _restore_mpas.

On symlink persistence: the attack target is users who download and load checkpoints from model hubs or shared storage. Git preserves symlinks - HuggingFace Hub repos are git repositories, so a symlink committed to the repo survives as a symlink when cloned or downloaded by the victim. The symlink target (/etc/passwd, ~/.ssh/id_rsa, etc.) resolves on the victim's machine at restore time, not the attacker's. The attacker only needs to publish the symlink; they don't need it to point anywhere meaningful on their own host.

Shared filesystems (NFS, FUSE mounts, cloud storage FUSE) and direct checkpoint directory transfers also preserve symlinks. The scenario where distribution strips symlinks (e.g., zip without -y) is the narrower case.

On responsibility: a containment check doesn't have to live in tensorstore. Flax's restore path (and orbax's) opens the checkpoint directory - at that layer, before materialising any array, it's straightforward to verify that no chunk file is a symlink or resolves outside the checkpoint root (os.path.realpath + os.path.commonpath). Archive extractors apply exactly this check at the extraction layer, not inside the compression library. The same pattern applies here.

[vfdev-5]

Flax's restore path (and orbax's) opens the checkpoint directory - at that layer, before materialising any array, it's straightforward to verify that no chunk file is a symlink or resolves outside the checkpoint root (os.path.realpath + os.path.commonpath).

yes, I agree, we can do that in Flax if Orbax wont do that.