Kanisia/pytorch-mobile-flatbuffer-union-oob-dispatch-poc

PyTorch Mobile FlatBuffer unknown-union OOB dispatch PoC

Summary

PyTorch Mobile validates a .ptl FlatBuffer with VerifyModuleBuffer() before loading it. The generated VerifyIValueUnion() function accepts an unknown IValueUnion discriminator through its default: return true branch. The CPU FlatbufferLoader then uses the untrusted discriminator directly as an index into its fixed-size ivalue_parsers_ array and calls the resulting value as an IValueParser function pointer.

The parser table has entries only for union values 0 through 16. The PoC changes a valid Int IValue discriminator from 1 to 255, while preserving the rest of a valid TorchScript Mobile FlatBuffer. The malformed .ptl passes FlatBuffer verification, reaches the real mobile loader, and performs an out-of-bounds indirect dispatch before model loading completes.

Files

• control_flatbuffer.ptl — valid CPU mobile model; loading and forward pass return tensor([3.]).
• malicious_unknown_ivalue_union.ptl — the same model with one union-tag byte changed from 0x01 to 0xff.
• build_and_verify.py — recreates both files and runs each load in a separate process.
• gdb_trace_linux.txt — Linux CPU backtrace showing the invalid indirect call from FlatbufferLoader::parseAndPopulate.

Reproduction

Use a CPU-only PyTorch installation that exposes the mobile module loader. The recorded runs use PyTorch 2.13.0+cpu.

C:\Users\mean\Desktop\bbg\.venv\Scripts\python.exe .\build_and_verify.py

Expected control result:

control return code: 0
MODEL_LOADED
FORWARD_RESULT tensor([3.])

Expected malicious result on Windows is a non-zero access-violation exit code (0xC0000005, represented by this Python runtime as 3221225477) before MODEL_LOADED. On Linux CPU, the same file terminates with SIGSEGV (exit 139 when run from a shell). The accompanying GDB trace records the instruction pointer as 0x0000000000000001 and identifies FlatbufferLoader::parseAndPopulate as the first PyTorch frame.

Impact

An attacker who can supply a .ptl model to a PyTorch Mobile CPU consumer can trigger an out-of-bounds function-pointer dispatch during model loading. This is memory corruption in the model parser rather than a handled malformed-file error, and it happens before inference begins. The provided PoC demonstrates the invalid indirect call; it does not claim a complete standalone RCE chain.

Fix recommendation

Reject an IValueUnion discriminator greater than IValueUnion::MAX before indexing ivalue_parsers_, and make the generated union verifier reject unknown discriminator values for this non-forward-compatible model format. Add regression tests for unknown union tags such as 17 and 255.