mscgo77/jax-serialize-executable-unpickler-rce-poc

JAX serialize_executable unrestricted pickle.Unpickler RCE

PoC for a huntr Model File Vulnerability (MFV) report against the JAX format (jax.experimental.serialize_executable).

The bug

jax/experimental/serialize_executable.py defines:

class _JaxPjrtUnpickler(pickle.Unpickler):
    def persistent_load(self, pid):
        ...  # only handles XLA-specific ('exec'/'device'/'client') tags

It never overrides find_class(). persistent_load only intercepts the special persistent_id tags this module's own _JaxPjrtPickler emits for XLA objects -- every other pickle opcode (GLOBAL/REDUCE, i.e. arbitrary class/function references) falls through to the stock pickle.Unpickler.find_class, which has no restriction at all. Calling the documented public function deserialize_and_load(serialized, in_tree, out_tree, backend=...) on untrusted serialized bytes runs arbitrary code the instant pickle.Unpickler.load() executes -- before a jax.stages.Compiled object is ever built, let alone run.

JAX's own docs do warn that the serialized bytes "must be trusted input," but the stated reason is narrower than the actual risk: the docs say running the deserialized executable afterward "may execute any custom call registered in the jaxlib." The PoC here proves code execution happens at .load() time itself, with no XLA custom call involved and no compiled function ever invoked.

Files

• malicious_serialized_executable.pkl -- the malicious payload. 98 bytes. Pickled as a (unloaded_executable, args_info_flat, no_kwargs) 3-tuple, matching exactly what .load() is supposed to return, so deserialize_and_load() runs to completion with zero exceptions. unloaded_executable's __reduce__ runs subprocess.Popen(["touch", MARKER]) as a side effect while reconstructing, then returns a stub object with a no-op .load().
• build_malicious_payload.py -- builds the payload above.
• reproduce_full_chain.py -- end-to-end repro: compiles a trivial real JAX function to get a legitimate (in_tree, out_tree) pair, then calls the public deserialize_and_load() with the malicious bytes swapped in. Returns a completely normal-looking jax.stages.Compiled object and creates /tmp/jax_pjrt_unpickler_pwned on load.

Reproduce

pip install jax
python3 reproduce_full_chain.py
ls -la /tmp/jax_pjrt_unpickler_pwned   # created by the RCE

Full report: huntr.com MFV submission (JAX format).