EnigmaConsultant/huntr-poc-orbax-msgpack-tuple-recursion-segfault

orbax-checkpoint 0.12.1: uncontrolled recursion in msgpack TUPLE ext-type handler causes SIGSEGV (native stack overflow) when loading a <1KB malicious checkpoint file

• Target: orbax-checkpoint (PyPI), version 0.12.1 — confirmed current latest release at time of testing (pip index versions orbax-checkpoint → 0.12.1, matches installed version).
• Tested with: flax 0.12.7 + orbax-checkpoint 0.12.1, pip-installed in an isolated venv, Python 3.13.
• Category: Model File Format — native memory-safety crash (uncontrolled recursion → stack overflow → SIGSEGV) via a malicious checkpoint file.
• Related huntr program target: "Flax-second-bug".

Root cause

File: orbax/checkpoint/msgpack_utils.py

class _MsgpackExtType(enum.IntEnum):
    NDARRAY = 1
    NATIVE_COMPLEX = 2
    NPSCALAR = 3
    TUPLE = 4

def _msgpack_ext_unpack(code, data):
    ...
    elif code == _MsgpackExtType.TUPLE:
        return tuple(msgpack.unpackb(data, raw=False, ext_hook=_msgpack_ext_unpack))  # line 115

Each nested tuple, encoded as ExtType(4, ...), forces a brand-new, separate msgpack.unpackb() call from inside the ext_hook callback of the outer call. This chain of nested Python/C call-stack frames is not subject to msgpack's own internal container-nesting depth guard — that guard only bounds recursion within a single unpackb call over native array/map/str containers (confirmed separately: a raw 5000-level nested map triggers msgpack's own clean, catchable StackError, not a crash).

Because the TUPLE-ext recursion instead re-enters the C extension stack frame-by-frame across independent calls, it blows the real OS thread/C stack directly, causing a segmentation fault — well before Python's own sys.getrecursionlimit() (default 1000) would raise a catchable RecursionError.

Reachability — real public API

This is reachable via the actual public, documented checkpoint-loading API:

orbax.checkpoint.aggregate_handlers.MsgpackHandler.deserialize(path)

which reads the raw bytes of a legacy "aggregate" checkpoint file straight from disk and passes them unmodified to msgpack_utils.msgpack_restore() → msgpack.unpackb(data, ext_hook=_msgpack_ext_unpack, raw=False). Simply loading an attacker-supplied checkpoint file crashes the victim process — no special API misuse required.

Dedup note

This is a genuinely separate vulnerability from the already-filed flax.serialization msgpack_restore() / _unchunk() unbounded-chunk-allocation DoS:

	This finding	Prior flax finding
Package	orbax-checkpoint	flax (serialization module)
Vuln class	Native stack-overflow segfault via uncontrolled recursion	Unbounded heap allocation via np.concatenate
Function	_msgpack_ext_unpack's TUPLE-ext handling	_unchunk
Ext type	Only exists in orbax's derived/extended copy of the msgpack module — flax.serialization does not define a TUPLE ext type at all (only ndarray/native_complex/npscalar)	N/A

PoC

A small generator (orbax_gen_payload.py) constructs N levels of nested msgpack ExtType(code=4 / TUPLE) wrappers by hand — pure msgpack.ExtType/packb calls, no orbax code needed to author the malicious file, only to load it. Each level's ext payload is msgpack.packb([inner_ext_or_value]).

python3 orbax_gen_payload.py <N> <outfile.msgpack>

Victim scripts included:

• orbax_restore_victim.py <path> — calls the low-level orbax.checkpoint.msgpack_utils.msgpack_restore() directly.
• orbax_full_path_victim.py <path> — calls the real, end-to-end public API orbax.checkpoint.aggregate_handlers.MsgpackHandler().deserialize(path) (via etils.epath.Path).

Negative control (N=190, <1KB payload)

$ python3 orbax_restore_victim.py nested_190_negative_control.msgpack; echo $?
RESTORED OK, nesting depth measured: 190
0

Both msgpack_restore() and the full public MsgpackHandler.deserialize() API restore successfully, exit code 0.

Positive / attack (N=200, 937-byte file)

$ python3 -X faulthandler orbax_full_path_victim.py nested_200_positive.msgpack
Fatal Python error: Segmentation fault

Current thread 0x00007fb92b7cd200 (most recent call first):
  File ".../venv/lib/python3.13/site-packages/orbax/checkpoint/msgpack_utils.py", line 115 in _msgpack_ext_unpack
  File ".../venv/lib/python3.13/site-packages/orbax/checkpoint/msgpack_utils.py", line 115 in _msgpack_ext_unpack
  File ".../venv/lib/python3.13/site-packages/orbax/checkpoint/msgpack_utils.py", line 115 in _msgpack_ext_unpack
  [... repeated identical frames ...]

Extension modules: jaxlib.cpu_feature_guard, numpy._core._multiarray_umath, ..., msgpack._cmsgpack (total: 16)
$ echo $?
139

Both the low-level msgpack_utils.msgpack_restore(data) call and the real end-to-end orbax.checkpoint.aggregate_handlers.MsgpackHandler().deserialize(path) checkpoint-loading call crash the interpreter with SIGSEGV (process exit code 139), reproduced deterministically across repeated runs (2/2 — see evidence_run1_low.log, evidence_run2_low.log, and evidence_run_full_faulthandler.log for the full captured faulthandler output, including repeated stack frames pinned at msgpack_utils.py:115 in _msgpack_ext_unpack).

Threshold sweep

N (nesting levels)	file size	exit code
160	—	0
170	—	0
180	—	0
190	887 B	0
200	937 B	139 (SIGSEGV)

Environment confirmation

$ pip show orbax-checkpoint | head -2
Name: orbax-checkpoint
Version: 0.12.1
$ pip index versions orbax-checkpoint | head -1
orbax-checkpoint (0.12.1)     # matches installed version = current PyPI release
$ ulimit -s
8192
$ python3 -c "import sys; print(sys.getrecursionlimit())"
1000

Verified this crash occurs below Python's own default recursion limit (sys.getrecursionlimit() == 1000) and with a standard/default OS stack ulimit (ulimit -s == 8192 KB, not artificially lowered) — ruling out a simple catchable RecursionError and confirming a genuine native memory-safety crash.

Files in this repo

• orbax_gen_payload.py — malicious-checkpoint generator (no orbax dependency).
• orbax_restore_victim.py — low-level victim script (msgpack_utils.msgpack_restore).
• orbax_full_path_victim.py / victim_full.py, victim_low.py — real public-API victim scripts.
• nested_190_negative_control.msgpack — 887-byte benign file, restores successfully (negative control).
• nested_200_positive.msgpack — 937-byte malicious file, triggers SIGSEGV (positive/attack case).
• evidence_run_full_faulthandler.log, evidence_run1_low.log, evidence_run2_low.log — captured verbatim faulthandler crash output from repeated runs.

Impact

Any consumer of orbax-checkpoint's legacy aggregate-checkpoint / msgpack loading path (directly, or transitively via flax checkpoint restore helpers that delegate to it) can be crashed by an attacker-supplied checkpoint file under 1KB in size — a straightforward denial-of-service against any service or pipeline that loads checkpoints from an untrusted or shared source (e.g. a model hub, shared storage, or a checkpoint accepted from a third party).