[Mlir-commits] [mlir] [mlir][bytecode] Reject dialect groups with more entries than declared (PR #229735)

[Dhairyashil R G]

https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/229735

`AttrTypeReader::initialize` sizes the attribute and type tables from the counts at the start of the offset section, then fills them through `parseDialectGrouping`, where each group has its own entry count. Nothing checked that the group counts stay within the declared total, so a malformed file could write past the end of a table. With assertions enabled this fails the `SmallVector::operator[]` assertion; without them it corrupts memory.

This patch emits an error when a group has more entries than the table has room for. The writer always emits groups that sum to the declared count, so valid bytecode is not affected.

Found by fuzzing `mlir-opt` with bytecode inputs.

Assisted-by: Claude

>From d3fb416568b454aff5575d7d9517a1c573aea80f Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <18738735+dhairyashilRG at users.noreply.github.com>
Date: Tue, 29 Sep 2026 21:05:11 +0530
Subject: [PATCH] [mlir][bytecode] Reject dialect groups with more entries than
 declared

AttrTypeReader::initialize sizes the attribute and type tables from the
counts at the start of the offset section. It then fills them through
parseDialectGrouping, where each group has its own entry count. Nothing
checked that the group counts stay within the declared total, so a
malformed file could write past the end of a table. With assertions
enabled this fails the SmallVector::operator[] assertion; without them it
corrupts memory.

Emit an error when a group has more entries than the table has room for.
The writer always emits groups that sum to the declared count, so valid
bytecode is not affected.

Assisted-by: Claude
---
 mlir/lib/Bytecode/Reader/BytecodeReader.cpp       |   8 ++++++++
 ...tr_type_offset_section-attr_group_count.mlirbc | Bin 0 -> 87 bytes
 ...tr_type_offset_section-type_group_count.mlirbc | Bin 0 -> 131 bytes
 .../invalid/invalid_attr_type_offset_section.mlir |  13 +++++++++++++
 4 files changed, 21 insertions(+)
 create mode 100644 mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-attr_group_count.mlirbc
 create mode 100644 mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-type_group_count.mlirbc

diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 96ead28d0aa4065..b1168eb5cd712d5 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -1331,6 +1331,14 @@ LogicalResult AttrTypeReader::initialize(
 
     // Parse an individual entry.
     auto parseEntryFn = [&](BytecodeDialect *dialect) -> LogicalResult {
+      // Verify that the dialect groups don't contain more entries than were
+      // declared for this range.
+      if (currentIndex == endIndex) {
+        return offsetReader.emitError(
+            "Attribute or Type dialect group entries exceed the declared "
+            "number of entries (",
+            endIndex, ")");
+      }
       auto &entry = range[currentIndex++];
 
       uint64_t entrySize;
diff --git a/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-attr_group_count.mlirbc b/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-attr_group_count.mlirbc
new file mode 100644
index 0000000000000000000000000000000000000000..a9fe238deac680c9fb644f01deb36cf0632d8af2
GIT binary patch
literal 87
zcmebEc^|}El3HBCz{tzY2!!m+K$?+Nn45`LoSB)Km6b)Dl`(*kk)4GdD9*;r$jZoI
m&MwF=m{gjXQ<9m-kei=Unv=?4Q(TggnP<nq!O6_Z$Or%zvJG|s

literal 0
HcmV?d00001

diff --git a/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-type_group_count.mlirbc b/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-type_group_count.mlirbc
new file mode 100644
index 0000000000000000000000000000000000000000..33f708b3486d977c8d1274ca764c579113ebcb27
GIT binary patch
literal 131
zcmXAiF$%&!5Jl(jWE_`e7=lF*3vVFUizm=R8+%ELg&2$Mp4%%J at x9`e?|S*%tqgM^
z7-La+?+o9ig-wFEh9e7$MiGOkac9gD5khLVf?EQ#Q3h5F0X<Ud?(MVYvE+$f at 4k)$
T{SJrnPo>jb`aGR!9a(@c5TO&U

literal 0
HcmV?d00001

diff --git a/mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir b/mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir
index 9a8bc08e6c25b87..31e5941cd179615 100644
--- a/mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir
+++ b/mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir
@@ -14,3 +14,16 @@
 
 // RUN: not mlir-opt %S/invalid-attr_type_offset_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
 // TRAILING_DATA: unexpected trailing data in the Attribute/Type offset section
+
+//===--------------------------------------------------------------------===//
+// Group Count
+//===--------------------------------------------------------------------===//
+
+// The dialect groups contain more entries than the declared number of
+// attributes.
+// RUN: not mlir-opt %S/invalid-attr_type_offset_section-attr_group_count.mlirbc 2>&1 | FileCheck %s --check-prefix=ATTR_GROUP_COUNT
+// ATTR_GROUP_COUNT: Attribute or Type dialect group entries exceed the declared number of entries (1)
+
+// The dialect groups contain more entries than the declared number of types.
+// RUN: not mlir-opt %S/invalid-attr_type_offset_section-type_group_count.mlirbc 2>&1 | FileCheck %s --check-prefix=TYPE_GROUP_COUNT
+// TYPE_GROUP_COUNT: Attribute or Type dialect group entries exceed the declared number of entries (1)