// 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.
//
// Unit-test reproduction of a heap-use-after-free in Skia's Graphite Dawn
// backend (third_party/skia/src/gpu/graphite/dawn/DawnGraphicsPipeline.cpp).
//
// DawnGraphicsPipeline::Make() declares a default-constructed
// skia_private::TArray<wgpu::StaticSamplerBindingLayout> and, in its
// per-sampler loop, stores an interior pointer (`&emplace_back()`) into
// entries[i].nextInChain. When the TArray exceeds its initial capacity (8
// elements after the first emplace_back), it reallocates and frees its old
// buffer, leaving every previously-stored entries[*].nextInChain dangling.
// Those dangling pointers are then dereferenced inside Dawn's
// CreateBindGroupLayout (Unpack walks the nextInChain list).
//
// In production this path is reachable only on Android (Graphite-on-Dawn-on-
// Vulkan with AHardwareBuffer-backed YCbCr textures). This test forces the
// same code path on Linux by using Dawn's Null backend (which enables all
// features, including StaticSamplers and YCbCrVulkanSamplers) and composing
// >=9 YCbCr-backed image shaders into a single Graphite pipeline.

#include <vector>

#include "base/logging.h"
#include "gpu/command_buffer/service/graphite_shared_context.h"
#include "gpu/command_buffer/service/skia_utils.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/SkBlendMode.h"
#include "third_party/skia/include/core/SkCanvas.h"
#include "third_party/skia/include/core/SkColorSpace.h"
#include "third_party/skia/include/core/SkImage.h"
#include "third_party/skia/include/core/SkPaint.h"
#include "third_party/skia/include/core/SkShader.h"
#include "third_party/skia/include/core/SkTileMode.h"
#include "third_party/skia/include/gpu/graphite/BackendTexture.h"
#include "third_party/skia/include/gpu/graphite/Context.h"
#include "third_party/skia/include/gpu/graphite/Image.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"
#include "third_party/skia/include/gpu/graphite/dawn/DawnGraphiteTypes.h"

namespace gpu {
namespace {

// Number of immutable (YCbCr) samplers to bind in a single pipeline. The
// first emplace_back() on the TArray rounds capacity up to 8; the 9th forces
// a realloc. Use a generous margin in case allocator slack differs.
constexpr int kNumYcbcrImages = 12;

TEST(DawnGraphiteStaticSamplerUAFTest, ManyImmutableSamplersInOnePipeline) {
  // ---- Set up a Dawn Null device with the YCbCr / StaticSamplers features.
  dawnProcSetProcs(&dawn::native::GetProcs());

  // StaticSamplers / YCbCrVulkanSamplers are experimental Dawn features and
  // are only exposed when the allow_unsafe_apis toggle is set.
  static constexpr const char* kEnabledToggles[] = {"allow_unsafe_apis"};
  wgpu::DawnTogglesDescriptor toggles = {};
  toggles.enabledToggleCount = std::size(kEnabledToggles);
  toggles.enabledToggles = kEnabledToggles;

  wgpu::InstanceDescriptor instance_desc = {};
  static constexpr auto kTimedWaitAny =
      wgpu::InstanceFeatureName::TimedWaitAny;
  instance_desc.requiredFeatureCount = 1;
  instance_desc.requiredFeatures = &kTimedWaitAny;
  instance_desc.nextInChain = &toggles;
  dawn::native::Instance dawn_instance(&instance_desc);

  wgpu::RequestAdapterOptions options = {};
  options.backendType = wgpu::BackendType::Null;
  options.featureLevel = wgpu::FeatureLevel::Core;
  options.nextInChain = &toggles;
  std::vector<dawn::native::Adapter> adapters =
      dawn_instance.EnumerateAdapters(&options);
  ASSERT_FALSE(adapters.empty()) << "No Dawn Null adapter found";

  wgpu::Adapter adapter(adapters[0].Get());
  ASSERT_TRUE(adapter.HasFeature(wgpu::FeatureName::StaticSamplers));
  ASSERT_TRUE(adapter.HasFeature(wgpu::FeatureName::YCbCrVulkanSamplers));

  static constexpr wgpu::FeatureName kFeatures[] = {
      wgpu::FeatureName::StaticSamplers,
      wgpu::FeatureName::YCbCrVulkanSamplers,
  };
  wgpu::DeviceDescriptor device_desc = {};
  device_desc.requiredFeatureCount = std::size(kFeatures);
  device_desc.requiredFeatures = kFeatures;
  device_desc.nextInChain = &toggles;
  device_desc.SetUncapturedErrorCallback(
      [](const wgpu::Device&, wgpu::ErrorType, wgpu::StringView msg) {
        LOG(ERROR) << "Dawn device error: "
                   << std::string_view(msg.data, msg.length);
      });
  wgpu::Device device = adapter.CreateDevice(&device_desc);
  ASSERT_TRUE(device);

  // ---- Create the Graphite-on-Dawn context.
  skgpu::graphite::DawnBackendContext backend_context = {};
  backend_context.fInstance = wgpu::Instance(dawn_instance.Get());
  backend_context.fDevice = device;
  backend_context.fQueue = device.GetQueue();

  skgpu::graphite::ContextOptions context_options;
  std::unique_ptr<skgpu::graphite::Context> context =
      skgpu::graphite::ContextFactory::MakeDawn(backend_context,
                                                context_options);
  ASSERT_TRUE(context) << "Failed to create Graphite-on-Dawn(Null) context";

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

  // ---- Build a YCbCr Dawn texture-info that Graphite will treat as needing
  // an immutable sampler (DawnCaps::getImmutableSamplerInfo() returns
  // non-default whenever DawnDescriptorIsValid() is true).
  wgpu::YCbCrVkDescriptor ycbcr = {};
  // VK_FORMAT_G8_B8R8_2PLANE_420_UNORM (a "known" Vulkan YCbCr format so that
  // usesExternalFormat()==false and the texture is treated as a normal 2D
  // sampled texture in WGSL).
  ycbcr.vkFormat = 1000156003u;
  ycbcr.vkYCbCrModel = 1;  // YCBCR_IDENTITY
  ycbcr.vkYCbCrRange = 0;  // ITU_FULL
  ycbcr.vkXChromaOffset = 0;
  ycbcr.vkYChromaOffset = 0;
  ycbcr.vkChromaFilter = wgpu::FilterMode::Nearest;

  skgpu::graphite::DawnTextureInfo tex_info;
  tex_info.fSampleCount = skgpu::graphite::SampleCount::k1;
  tex_info.fMipmapped = skgpu::Mipmapped::kNo;
  tex_info.fFormat = wgpu::TextureFormat::OpaqueYCbCrAndroid;
  tex_info.fViewFormat = wgpu::TextureFormat::OpaqueYCbCrAndroid;
  tex_info.fUsage = wgpu::TextureUsage::TextureBinding;
  tex_info.fAspect = wgpu::TextureAspect::All;
  tex_info.fYcbcrVkDescriptor = ycbcr;

  // ---- Create N OpaqueYCbCrAndroid textures (and views) on the Null device,
  // wrap each one as a Graphite BackendTexture/SkImage, and turn it into an
  // image shader.
  const SkISize kSize = SkISize::Make(8, 8);
  wgpu::TextureDescriptor wgpu_tex_desc = {};
  wgpu_tex_desc.usage = wgpu::TextureUsage::TextureBinding;
  wgpu_tex_desc.size = {static_cast<uint32_t>(kSize.width()),
                        static_cast<uint32_t>(kSize.height()), 1};
  wgpu_tex_desc.format = wgpu::TextureFormat::OpaqueYCbCrAndroid;

  std::vector<wgpu::Texture> textures;
  std::vector<wgpu::TextureView> views;
  std::vector<sk_sp<SkImage>> images;
  sk_sp<SkShader> composed;

  for (int i = 0; i < kNumYcbcrImages; ++i) {
    wgpu::Texture tex = device.CreateTexture(&wgpu_tex_desc);
    ASSERT_TRUE(tex) << "CreateTexture(OpaqueYCbCrAndroid) failed (i=" << i
                     << ")";

    wgpu::TextureViewDescriptor view_desc = {};
    view_desc.dimension = wgpu::TextureViewDimension::e2D;
    view_desc.nextInChain = &ycbcr;
    wgpu::TextureView view = tex.CreateView(&view_desc);
    ASSERT_TRUE(view) << "CreateView(YCbCr) failed (i=" << i << ")";

    skgpu::graphite::BackendTexture be_tex =
        skgpu::graphite::BackendTextures::MakeDawn(kSize, tex_info, view.Get());
    ASSERT_TRUE(be_tex.isValid()) << "BackendTexture invalid (i=" << i << ")";

    sk_sp<SkImage> image = SkImages::WrapTexture(
        recorder.get(), be_tex, kPremul_SkAlphaType, SkColorSpace::MakeSRGB());
    ASSERT_TRUE(image) << "WrapTexture failed (i=" << i << ")";

    sk_sp<SkShader> shader = image->makeShader(
        SkTileMode::kClamp, SkTileMode::kClamp, SkSamplingOptions());
    ASSERT_TRUE(shader);

    composed = composed ? SkShaders::Blend(SkBlendMode::kPlus,
                                           std::move(composed), shader)
                        : shader;
    textures.push_back(std::move(tex));
    views.push_back(std::move(view));
    images.push_back(std::move(image));
  }

  // ---- Draw with all N image shaders so that they end up in a single
  // Graphite fragment pipeline.
  SkImageInfo target_ii =
      SkImageInfo::Make(32, 32, kRGBA_8888_SkColorType, kPremul_SkAlphaType);
  sk_sp<SkSurface> surface =
      SkSurfaces::RenderTarget(recorder.get(), target_ii);
  ASSERT_TRUE(surface);

  SkPaint paint;
  paint.setShader(std::move(composed));
  surface->getCanvas()->drawRect(SkRect::MakeWH(32, 32), paint);

  // Pipeline creation (DawnGraphicsPipeline::Make) is triggered here. With
  // >=9 immutable samplers, staticSamplerLayouts reallocates mid-loop,
  // leaving every earlier entries[*].nextInChain dangling. Dawn's
  // CreateBindGroupLayout then dereferences the freed allocation.
  std::unique_ptr<skgpu::graphite::Recording> recording = recorder->snap();
  ASSERT_TRUE(recording);

  skgpu::graphite::InsertRecordingInfo info = {};
  info.fRecording = recording.get();
  context->insertRecording(info);
  context->submit(skgpu::graphite::SyncToCpu::kYes);
}

}  // namespace
}  // namespace gpu