jax-serialize-executable-extension-confusion-poc

Benign JAX serialize_executable Extension-Confusion PoC

This repository contains a safe research PoC for a JAX serialized-executable loading gap.

Files

• artifact.jaxexec: benign PoC artifact using a JAX-looking extension
• verify_poc.py: reproduces runtime and scanner behavior
• results.json: captured verifier and scanner results
• versions.txt: package versions used during validation

What it demonstrates

JAX documents jax.experimental.serialize_executable as pickling support for precompiled binaries. The loader entrypoint deserialize_and_load() will unpickle attacker-controlled bytes.

The same payload bytes:

• are flagged by ModelScan when saved as artifact.pkl
• are skipped by ModelScan 0.8.8 when saved as artifact.jaxexec
• still trigger benign marker-file creation when passed to deserialize_and_load()

Reproduction

Run:

python .\verify_poc.py

Expected behavior:

• artifact.jaxexec creates marker.txt
• JAX then raises a benign TypeError because the payload is not a real serialized executable tuple
• ModelScan reports:
  • artifact.pkl: 1 critical pickle issue
  • artifact.jaxexec: 0 scanned files / 0 issues

Scanner summary

• ModelScan version: 0.8.8
• Extension routing is filename-based for supported formats
• No JAX serialized-executable format handler is present in the tested default settings

Security impact

• Impact class: loader-time code execution if a consumer treats a serialized JAX executable blob as trusted input
• PoC impact here is limited to creation of a local marker file

Limitations

• This is not a novel safe-format parser bug.
• The underlying format is still pickle-based, and JAX documentation explicitly says so.
• As a result, this is best treated as extension-confusion / scanner-coverage research, not a high-confidence bounty submission.

SHA256 / size

• artifact.jaxexec
  • SHA256: af07950d7052e718dc656231b8c12e4176d9ed1428245d3815c5183b72352701
  • Size: 201 bytes

Mitigations

• Never call deserialize_and_load() on untrusted bytes.
• Treat JAX serialized executables as unsafe pickle artifacts for scanning policy.
• Add explicit scanner support for JAX serialized-executable extensions or inspect magic/content rather than extension alone.