From 501bdb2fed3132e53fa836e9f4b7da3acffb149f Mon Sep 17 00:00:00 2001
From: Swung0x48 <swung0x48@outlook.com>
Date: Wed, 23 Sep 2026 08:23:16 -0400
Subject: [PATCH] [MG_State, MG_Test] (P7): bound program archive vector counts

---
 .../ProgramState/ProgramArtifactsCodec.cpp    | 18 +++++++------
 .../Program/ProgramArtifactsCodecTest.cpp     | 25 +++++++++++++++++++
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/MobileGL/MG_State/GLState/ProgramState/ProgramArtifactsCodec.cpp b/MobileGL/MG_State/GLState/ProgramState/ProgramArtifactsCodec.cpp
index 5ad158f9..f1c16911 100644
--- a/MobileGL/MG_State/GLState/ProgramState/ProgramArtifactsCodec.cpp
+++ b/MobileGL/MG_State/GLState/ProgramState/ProgramArtifactsCodec.cpp
@@ -247,15 +247,16 @@ namespace MobileGL::MG_State::GLState {
             return true;
         }
 
-        // A COUNT IS CHECKED AGAINST THE BYTES THAT REMAIN BEFORE ANYTHING IS ALLOCATED. Every
-        // element this codec writes costs at least one byte, so a count larger than the
-        // remaining bytes cannot describe this stream - and refusing it here is what stops a
-        // corrupt or truncated archive from turning into a multi-gigabyte resize before the
-        // element loop notices it has run out.
-        Bool TakeCount(ReadCursor& in, SizeT& count) {
+        // A COUNT IS CHECKED AGAINST THE BYTES THAT REMAIN BEFORE ANYTHING IS ALLOCATED.
+        // For resizable vectors, also charge sizeof(value_type) per element: even when a
+        // serialized record is compact, its immediate vector reservation must not exceed the
+        // remaining untrusted archive bytes. The division form avoids multiplying attacker
+        // input and is safe on both 32-bit and 64-bit SizeT.
+        Bool TakeCount(ReadCursor& in, SizeT& count, SizeT allocationBytesPerElement = 1) {
             Uint64 raw = 0;
             if (!TakeRaw(in, raw)) return false;
-            if (raw > static_cast<Uint64>(in.Remaining())) {
+            if (allocationBytesPerElement == 0 ||
+                raw > static_cast<Uint64>(in.Remaining() / allocationBytesPerElement)) {
                 in.Ok = false;
                 return false;
             }
@@ -304,7 +305,8 @@ namespace MobileGL::MG_State::GLState {
                 }
             } else if constexpr (ArchiveVector<T>) {
                 SizeT count = 0;
-                if (!TakeCount(in, count)) return;
+                using Element = typename T::value_type;
+                if (!TakeCount(in, count, sizeof(Element))) return;
                 value.clear();
                 value.resize(count);
                 for (auto& element : value) {
diff --git a/MobileGL/MG_Test/Program/ProgramArtifactsCodecTest.cpp b/MobileGL/MG_Test/Program/ProgramArtifactsCodecTest.cpp
index 3290e206..e223e6f8 100644
--- a/MobileGL/MG_Test/Program/ProgramArtifactsCodecTest.cpp
+++ b/MobileGL/MG_Test/Program/ProgramArtifactsCodecTest.cpp
@@ -326,6 +326,31 @@ TEST(ProgramArtifactsCodec, ATruncatedStreamIsRefusedNotGuessed) {
 #endif
 }
 
+// A count is not permission to reserve a container larger than the bytes left in the
+// untrusted archive. This mutates the first LinkArtifacts vector count to the boundary value
+// that the former raw <= Remaining check accepted but sizeof(value_type) <= Remaining rejects.
+TEST(ProgramArtifactsCodec, AVectorCountCannotReservePastTheRemainingArchiveBytes) {
+#if MOBILEGL_PIPE_PUSH
+    Vector<Uint8> bytes;
+    EncodeProgramArtifacts(LinkArtifacts{}, SpirvArtifacts{}, bytes);
+    constexpr SizeT countOffset = sizeof(Uint32) + sizeof(Uint64);
+    ASSERT_GT(bytes.size(), countOffset + sizeof(Uint64));
+    const SizeT available = bytes.size() - countOffset - sizeof(Uint64);
+    const Uint64 forgedCount = static_cast<Uint64>(available / sizeof(UniformReflection)) + 1;
+    ASSERT_LT(forgedCount, static_cast<Uint64>(available))
+        << "fixture must pass the old count <= remaining check";
+    std::memcpy(bytes.data() + countOffset, &forgedCount, sizeof(forgedCount));
+
+    LinkArtifacts link;
+    SpirvArtifacts spirv;
+    EXPECT_FALSE(DecodeProgramArtifacts(bytes.data(), bytes.size(), link, spirv));
+    EXPECT_TRUE(link.uniformReflection.empty());
+    EXPECT_TRUE(spirv.generatedSpirv.empty());
+#else
+    GTEST_SKIP() << "MOBILEGL_PIPE_PUSH is off: the archive codec is push-only";
+#endif
+}
+
 // Negative control 2: both the version and schema word must refuse mismatches.
 // Monolith v1 retains the old native-size echo in that second word.
 TEST(ProgramArtifactsCodec, AVersionMismatchIsRefused) {