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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

`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

---
Full diff: https://github.com/llvm/llvm-project/pull/229735.diff


4 Files Affected:

- (modified) mlir/lib/Bytecode/Reader/BytecodeReader.cpp (+8) 
- (added) mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-attr_group_count.mlirbc () 
- (added) mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-type_group_count.mlirbc () 
- (modified) mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 96ead28d0aa40..b1168eb5cd712 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 0000000000000..a9fe238deac68
Binary files /dev/null and b/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-attr_group_count.mlirbc differ
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 0000000000000..33f708b3486d9
Binary files /dev/null and b/mlir/test/Bytecode/invalid/invalid-attr_type_offset_section-type_group_count.mlirbc differ
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 9a8bc08e6c25b..31e5941cd1796 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)

``````````

</details>


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