From bfbd2ccf3ee14e66f162b333e9f0aa7cc6cb32c0 Mon Sep 17 00:00:00 2001
From: Le Hoang Quyen <[email protected]>
Date: Wed, 22 Jul 2026 14:56:29 +0800
Subject: [PATCH] Metal: Fix index buffer rewriting OOB write and rendering bug

When rewriting index buffers for provoking vertex conventions (e.g.,
flat shading with primitive restart enabled), the output buffer was
previously sized using primCountForIndexCount(count).

When primitive restart markers appeared at the beginning of the draw
window or between primitive ranges, clippedRange.begin shifted the
starting offset forward, causing the compute shader to write past the
end of the pre-allocated buffer and corrupt trailing primitive indices.

This patch extracts the dispatch ranges in a single pass into an
angle::FastVector and computes the exact required buffer capacity using
CheckedNumeric::Max, ensuring full buffer bounds coverage and correct
rendering.

Bug: chromium:536626343
Change-Id: Iecd8a3b9a280cb4b330a6d9373996839b2299250
Reviewed-on: https://chromium-review.googlesource.com/c/angle/angle/+/8129619
Reviewed-by: Geoff Lang <[email protected]>
Commit-Queue: Quyen Le <[email protected]>
---

diff --git a/src/libANGLE/renderer/metal/ProvokingVertexHelper.mm b/src/libANGLE/renderer/metal/ProvokingVertexHelper.mm
index 26d0f7e..7577677 100644
--- a/src/libANGLE/renderer/metal/ProvokingVertexHelper.mm
+++ b/src/libANGLE/renderer/metal/ProvokingVertexHelper.mm
@@ -9,6 +9,7 @@
 
 #include "libANGLE/renderer/metal/ProvokingVertexHelper.h"
 #import <Foundation/Foundation.h>
+#include "common/FastVector.h"
 #include "common/base/anglebase/numerics/checked_math.h"
 #include "libANGLE/Display.h"
 #include "libANGLE/renderer/metal/ContextMtl.h"
@@ -21,6 +22,13 @@
 
 namespace
 {
+struct IndexRewriteRange
+{
+    uint32_t indexCount;
+    uint32_t primitiveCount;
+    size_t srcOffset;
+    size_t dstOffset;
+};
 constexpr size_t kInitialIndexBufferSize = 0xFFFF;  // Initial 64k pool.
 }
 static inline uint32_t primCountForIndexCount(const uint fixIndexBufferKey,
@@ -223,31 +231,9 @@
             return angle::Result::Stop;
     }
 
-    // Maximum primitive/index count needed for buffer allocation, based on the full draw window.
-    uint32_t totalPrimCount     = primCountForIndexCount(indexBufferKey, count);
-    uint32_t totalNewIndexCount = 0;
-    ANGLE_CHECK_GL_MATH(
-        context, indexCountForPrimCount(indexBufferKey, totalPrimCount, &totalNewIndexCount));
-
     const size_t indexTypeShift = gl::GetDrawElementsTypeShift(indexBufferType);
-    size_t firstIndexOffset     = firstIndex << indexTypeShift;
-    size_t newFirstIndexOffset  = firstIndexOffset;
-    if (mode != newMode)
-    {
-        newFirstIndexOffset *= perPrimitiveIndexCount;
-    }
-    angle::CheckedNumeric<size_t> checkedBufferSize(totalNewIndexCount);
-    checkedBufferSize <<= indexTypeShift;
-    checkedBufferSize += newFirstIndexOffset;
-    ANGLE_CHECK_GL_MATH(context, checkedBufferSize.IsValid());
-    mtl::BufferSlice newBuffer;
-    ANGLE_TRY(mIndexBuffers.allocate(context, checkedBufferSize.ValueOrDie(), &newBuffer));
-
-    mtl::ComputeCommandEncoder *encoder =
-        context->getComputeCommandEncoderWithoutEndingRenderEncoder();
-    const bool isForGenerateIndices = false;
-    ANGLE_TRY(
-        prepareCommandEncoderForFunction(context, encoder, indexBufferKey, isForGenerateIndices));
+    angle::FastVector<IndexRewriteRange, 4> rewriteRanges;
+    angle::CheckedNumeric<size_t> checkedBufferSize = 0;
 
     const size_t lastIndex = firstIndex + count - 1;
     for (const auto &range : drawIndexRanges)
@@ -267,28 +253,52 @@
         {
             continue;
         }
-        size_t beginOffset = clippedRange.begin << indexTypeShift;
+        angle::CheckedNumeric<size_t> srcOffset = clippedRange.begin;
+        srcOffset <<= indexTypeShift;
+
         uint32_t primitiveCount;
-        size_t newBeginOffset;
+        angle::CheckedNumeric<size_t> dstOffset = srcOffset;
         if (mode == newMode)
         {
             primitiveCount = indexCount / perPrimitiveIndexCount;
-            newBeginOffset = clippedRange.begin << indexTypeShift;
         }
         else
         {
             // Expanded modes: `N` source indices produce `(N - perPrimitiveIndexCount + 1)`
             // primitives.
             primitiveCount = indexCount - perPrimitiveIndexCount + 1;
-            newBeginOffset = (clippedRange.begin << indexTypeShift) * perPrimitiveIndexCount;
+            dstOffset *= perPrimitiveIndexCount;
         }
 
-        auto threadsPerThreadgroup = MTLSizeMake(MIN(primitiveCount, 64u), 1, 1);
-        encoder->setBuffer(indexBuffer.buffer(), indexBuffer.offset() + beginOffset, 0);
-        encoder->setBufferForWrite(newBuffer.buffer(), newBuffer.offset() + newBeginOffset, 1);
-        encoder->setData(indexCount, 2);
-        encoder->setData(primitiveCount, 3);
-        encoder->dispatch(MTLSizeMake((static_cast<NSUInteger>(primitiveCount) +
+        ANGLE_CHECK_GL_MATH(context, srcOffset.IsValid() && dstOffset.IsValid());
+        rewriteRanges.push_back(
+            {indexCount, primitiveCount, srcOffset.ValueOrDie(), dstOffset.ValueOrDie()});
+
+        angle::CheckedNumeric<size_t> rangeEndOffset(primitiveCount);
+        rangeEndOffset *= perPrimitiveIndexCount;
+        rangeEndOffset <<= indexTypeShift;
+        rangeEndOffset += dstOffset;
+        checkedBufferSize = checkedBufferSize.Max(rangeEndOffset);
+    }
+
+    ANGLE_CHECK_GL_MATH(context, checkedBufferSize.IsValid());
+    mtl::BufferSlice newBuffer;
+    ANGLE_TRY(mIndexBuffers.allocate(context, checkedBufferSize.ValueOrDie(), &newBuffer));
+
+    mtl::ComputeCommandEncoder *encoder =
+        context->getComputeCommandEncoderWithoutEndingRenderEncoder();
+    const bool isForGenerateIndices = false;
+    ANGLE_TRY(
+        prepareCommandEncoderForFunction(context, encoder, indexBufferKey, isForGenerateIndices));
+
+    for (const IndexRewriteRange &rangeInfo : rewriteRanges)
+    {
+        auto threadsPerThreadgroup = MTLSizeMake(MIN(rangeInfo.primitiveCount, 64u), 1, 1);
+        encoder->setBuffer(indexBuffer.buffer(), indexBuffer.offset() + rangeInfo.srcOffset, 0);
+        encoder->setBufferForWrite(newBuffer.buffer(), newBuffer.offset() + rangeInfo.dstOffset, 1);
+        encoder->setData(rangeInfo.indexCount, 2);
+        encoder->setData(rangeInfo.primitiveCount, 3);
+        encoder->dispatch(MTLSizeMake((static_cast<NSUInteger>(rangeInfo.primitiveCount) +
                                        threadsPerThreadgroup.width - 1) /
                                           threadsPerThreadgroup.width,
                                       1, 1),
diff --git a/src/tests/gl_tests/DrawElementsTest.cpp b/src/tests/gl_tests/DrawElementsTest.cpp
index 9c9ef46..ace2d81 100644
--- a/src/tests/gl_tests/DrawElementsTest.cpp
+++ b/src/tests/gl_tests/DrawElementsTest.cpp
@@ -967,6 +967,103 @@
     ASSERT_GL_NO_ERROR();
 }
 
+// Test a large flat-shaded GL_TRIANGLES draw with primitive restart markers at the
+// beginning of the index buffer. Previously, the Metal backend would size the
+// rewritten index buffer for the primitive-aligned count while writing at absolute
+// per-range offsets, so the last write for the shifted range landed past the allocation.
+// Uses distinct vertices for every quad so the final quad specifically validates that the
+// rewritten provoking vertex indices at the end of the buffer are rendered correctly.
+TEST_P(DrawElementsTest, FlatTrianglesLargePrimitiveRestartAtBegin)
+{
+    constexpr char kFlatVS[] = R"(#version 300 es
+in vec4 a_position;
+in float a_mark;
+flat out float v_mark;
+void main()
+{
+    v_mark = a_mark;
+    gl_Position = a_position;
+})";
+    constexpr char kFlatFS[] = R"(#version 300 es
+precision highp float;
+flat in float v_mark;
+out vec4 fragColor;
+void main()
+{
+    fragColor = v_mark >= 0.5 ? vec4(0.0, 1.0, 0.0, 1.0) : vec4(1.0, 0.0, 0.0, 1.0);
+})";
+    ANGLE_GL_PROGRAM(program, kFlatVS, kFlatFS);
+    glUseProgram(program);
+
+    constexpr GLsizei kQuads = 5000;
+    std::vector<Vector3> vertices(kQuads * 4, Vector3(0.0f, 0.0f, 0.0f));
+    std::vector<GLfloat> marks(kQuads * 4, 0.0f);
+
+    // Quad kQuads - 1 covers the full screen and has provoking marks set to 1.0 (green).
+    size_t lastQuadBase        = (kQuads - 1) * 4;
+    vertices[lastQuadBase + 0] = Vector3(-1.0f, -1.0f, 0.0f);
+    vertices[lastQuadBase + 1] = Vector3(1.0f, -1.0f, 0.0f);
+    vertices[lastQuadBase + 2] = Vector3(-1.0f, 1.0f, 0.0f);
+    vertices[lastQuadBase + 3] = Vector3(1.0f, 1.0f, 0.0f);
+
+    // Last vertex convention, triangles {0,1,2, 2,1,3}: 2 and 3 are provoking.
+    marks[lastQuadBase + 2] = 1.0f;
+    marks[lastQuadBase + 3] = 1.0f;
+

diff --git a/src/tests/gl_tests/DrawElementsTest.cpp b/src/tests/gl_tests/DrawElementsTest.cpp
index 9c9ef46..ace2d81 100644
--- a/src/tests/gl_tests/DrawElementsTest.cpp
+++ b/src/tests/gl_tests/DrawElementsTest.cpp
@@ -967,6 +967,103 @@
     ASSERT_GL_NO_ERROR();
 }
 
+// Test a large flat-shaded GL_TRIANGLES draw with primitive restart markers at the
+// beginning of the index buffer. Previously, the Metal backend would size the
+// rewritten index buffer for the primitive-aligned count while writing at absolute
+// per-range offsets, so the last write for the shifted range landed past the allocation.
+// Uses distinct vertices for every quad so the final quad specifically validates that the
+// rewritten provoking vertex indices at the end of the buffer are rendered correctly.
+TEST_P(DrawElementsTest, FlatTrianglesLargePrimitiveRestartAtBegin)
+{
+    constexpr char kFlatVS[] = R"(#version 300 es
+in vec4 a_position;
+in float a_mark;
+flat out float v_mark;
+void main()
+{
+    v_mark = a_mark;
+    gl_Position = a_position;
+})";
+    constexpr char kFlatFS[] = R"(#version 300 es
+precision highp float;
+flat in float v_mark;
+out vec4 fragColor;
+void main()
+{
+    fragColor = v_mark >= 0.5 ? vec4(0.0, 1.0, 0.0, 1.0) : vec4(1.0, 0.0, 0.0, 1.0);
+})";
+    ANGLE_GL_PROGRAM(program, kFlatVS, kFlatFS);
+    glUseProgram(program);
+
+    constexpr GLsizei kQuads = 5000;
+    std::vector<Vector3> vertices(kQuads * 4, Vector3(0.0f, 0.0f, 0.0f));
+    std::vector<GLfloat> marks(kQuads * 4, 0.0f);
+
+    // Quad kQuads - 1 covers the full screen and has provoking marks set to 1.0 (green).
+    size_t lastQuadBase        = (kQuads - 1) * 4;
+    vertices[lastQuadBase + 0] = Vector3(-1.0f, -1.0f, 0.0f);
+    vertices[lastQuadBase + 1] = Vector3(1.0f, -1.0f, 0.0f);
+    vertices[lastQuadBase + 2] = Vector3(-1.0f, 1.0f, 0.0f);
+    vertices[lastQuadBase + 3] = Vector3(1.0f, 1.0f, 0.0f);
+
+    // Last vertex convention, triangles {0,1,2, 2,1,3}: 2 and 3 are provoking.
+    marks[lastQuadBase + 2] = 1.0f;
+    marks[lastQuadBase + 3] = 1.0f;
+
+    GLBuffer vertexBuffer;
+    glBindBuffer(GL_ARRAY_BUFFER, vertexBuffer);
+    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices[0]) * vertices.size(), vertices.data(),
+                 GL_STATIC_DRAW);
+    GLint posLocation = glGetAttribLocation(program, "a_position");
+    ASSERT_NE(-1, posLocation);
+    glEnableVertexAttribArray(posLocation);
+    glVertexAttribPointer(posLocation, 3, GL_FLOAT, GL_FALSE, 0, 0);
+
+    GLBuffer markBuffer;
+    glBindBuffer(GL_ARRAY_BUFFER, markBuffer);
+    glBufferData(GL_ARRAY_BUFFER, sizeof(marks[0]) * marks.size(), marks.data(), GL_STATIC_DRAW);
+    GLint markLocation = glGetAttribLocation(program, "a_mark");
+    ASSERT_NE(-1, markLocation);
+    glEnableVertexAttribArray(markLocation);
+    glVertexAttribPointer(markLocation, 1, GL_FLOAT, GL_FALSE, 0, 0);
+
+    // Two restart markers followed by quads with distinct vertex indices. The index count is
+    // large enough to exceed backend staging buffer limits and count % 3 == 2 so primitive
+    // alignment differs from the raw count.
+    constexpr GLuint kRestart   = 0xFFFFFFFFu;
+    std::vector<GLuint> indices = {kRestart, kRestart};
+    indices.reserve(2 + static_cast<size_t>(kQuads) * 6);
+    for (GLuint q = 0; q < static_cast<GLuint>(kQuads); ++q)
+    {
+        GLuint base = q * 4;
+        indices.push_back(base + 0);
+        indices.push_back(base + 1);
+        indices.push_back(base + 2);
+        indices.push_back(base + 2);
+        indices.push_back(base + 1);
+        indices.push_back(base + 3);
+    }
+    ASSERT_EQ(2u, indices.size() % 3);
+
+    GLBuffer elementBuffer;
+    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, elementBuffer);
+    glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices[0]) * indices.size(), indices.data(),
+                 GL_STATIC_DRAW);
+    ASSERT_GL_NO_ERROR();
+
+    glEnable(GL_PRIMITIVE_RESTART_FIXED_INDEX);
+    glClearColor(1.f, 0.f, 0.f, 1.f);
+    glClear(GL_COLOR_BUFFER_BIT);
+    glDrawElements(GL_TRIANGLES, static_cast<GLsizei>(indices.size()), GL_UNSIGNED_INT, nullptr);
+    ASSERT_GL_NO_ERROR();
+
+    EXPECT_PIXEL_COLOR_EQ(0, 0, GLColor::green);
+    EXPECT_PIXEL_COLOR_EQ(getWindowWidth() - 1, 0, GLColor::green);
+    EXPECT_PIXEL_COLOR_EQ(0, getWindowHeight() - 1, GLColor::green);
+    EXPECT_PIXEL_COLOR_EQ(getWindowWidth() - 1, getWindowHeight() - 1, GLColor::green);
+    ASSERT_GL_NO_ERROR();
+}
+
 // Tests various draw element parameter, vertex buffer contents variants.
 // Does not yet test using GL_BYTE 0xFF, GL_SHORT 0xFFFF with primitive restart off.
 TEST_P(DrawElementsVariantsTest, Draw)