nu1lp7r/mfv-megengine-shape-oob-write

PoC: MegEngine .mge loader — tensor shape dim-count unchecked → stack OOB write with attacker-controlled values (CWE-787)

Target

MegEngine (Megvii's deep-learning framework, ~4.7k stars) — .mge is its native serialized model format (flatbuffers). Loading .mge files (model zoo / shared checkpoints / load_graph) is the standard workflow.

Vulnerability

The v2 model loader (src/serialization/impl/serializer_oss_v2.cpp:586-590, master verified 2026-07-18; also present in the v1 loader serializer_oss.cpp:615-619):

TensorLayout load_tensor_layout_without_format(const fbs::v2::Tensor* tensor) {
    TensorLayout layout;                       // TensorShape { size_t shape[7]; size_t ndim; ... }
    if (tensor->shape()) {
        layout.ndim = tensor->shape()->size(); // attacker-controlled, no <= MAX_NDIM check
        std::copy(tensor->shape()->begin(), tensor->shape()->end(), layout.shape);  // OOB write
    }

TensorShape::shape is a fixed size_t[MEGDNN_MAX_NDIM] (7) array inside a stack-allocated TensorLayout. A model file whose tensor declares more than 7 shape dims makes std::copy write ndim × 8 attacker-controlled bytes past the struct. The flatbuffers Verifier cannot catch this — a 1000-element shape vector is structurally valid fbs; the bug is verifier-past logic.

Proof (official MegEngine 1.13.0 wheel, pip install megengine, Python 3.10)

Two files in this repo differ only in the shape vector of a single ImmutableTensor:

• benign.mge — shape=[2,2] → load_graph loads cleanly
• malicious.mge — shape=[1 × 1000] → SIGSEGV in the official libmegengine_shared.so:

Thread 1 received signal SIGSEGV, Segmentation fault.
0x... in mgb::serialization::GraphLoaderOSSV2::OprLoadContextImpl::load_tensor_shared(bool)
    from libmegengine_shared.so
#1 0x0000000000000001 in ?? ()
#2 0x0000000000000001 in ?? ()
#3 0x0000000000000001 in ?? ()

Frames #1-#5 = 0x1 are the overwritten return addresses: the 1000 copied dims (all 1, from our file) smashed the call stack. The dim values are fully attacker-controlled 32-bit values (zero-extended to 64-bit size_t), i.e. this is a controlled stack OOB write primitive — not just a crash.

Reproduction

pip install megengine==1.13.0   # official build (manylinux wheel, py3.10)
python load_poc.py              # loads malicious.mge via megengine.core.tensor.megbrain_graph.load_graph
# → SIGSEGV (load_tensor_shared). Control: load benign.mge → succeeds.

craft_malicious.cpp regenerates both files using flatc-generated code from the in-repo schema (src/serialization/impl/schema_v2.fbs); real_model.mge is a genuine model dumped by mge.jit.trace.dump for reference.

Impact

Any application that loads untrusted .mge models (conversion/serving pipelines, CI model checks, megengine.load_graph on user uploads) is exposed to a controlled stack write at load time → denial of service and potential control-flow hijack. Same model-file trust class as CVE-2026-27940 (llama.cpp) / CVE-2026-7482 ("Bleeding Llama", CVSS 9.1), here in Megvii's framework, unpatched on master (both v1 and v2 loaders).

Suggested fix

In load_tensor_layout_without_format (and the v1 load_tensor_layout): before the copy, if (tensor->shape()->size() > MEGDNN_MAX_NDIM) mgb_throw(SerializationError, "tensor ndim too large");

Files

• malicious.mge / benign.mge — PoC pair (differ only in shape dim count)
• craft_malicious.cpp / inspect.cpp — generator + inspector (flatc-generated from in-repo schema)
• load_poc.py — loader driver
• segv_trace.txt — captured crash evidence (incl. control run)
• real_model.mge — genuine reference model dumped by megengine