[Mlir-commits] [mlir] [mlir][bytecode] Fix SourceMgr lifetime extension in BytecodeReader (PR #185321)

[Jacques Pienaar]

================
@@ -231,6 +233,38 @@ TEST(Bytecode, OpWithoutProperties) {
               OperationEquivalence::computeHash(roundtripped));
 }
 
+TEST(Bytecode, SourceMgrLifetimeExtendedByReader) {
+  MLIRContext context;
+
+  // Create a trivial module and serialize it to bytecode.
+  OwningOpRef<ModuleOp> module =
+      ModuleOp::create(UnknownLoc::get(&context));
+  std::string bytecode;
+  llvm::raw_string_ostream os(bytecode);
+  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), os)));
+
+  // Build a BytecodeReader whose SourceMgr goes out of scope immediately after
+  // construction. The reader must keep it alive via shared ownership, otherwise
+  // readTopLevel will access freed memory (UAF).
+  ParserConfig config(&context);
+  std::shared_ptr<BytecodeReader> reader;
+  {
+    auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+    auto buffer = llvm::MemoryBuffer::getMemBufferCopy(bytecode, "model");
+    sourceMgr->AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());
+    llvm::MemoryBufferRef bufRef =
+        *sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID());
+
+    reader = std::make_shared<BytecodeReader>(
+        bufRef, config, /*lazyLoad=*/true, sourceMgr);
+    // sourceMgr destroyed here — reader must have extended its lifetime.
+  }
+
+  Block block;
+  EXPECT_TRUE(succeeded(
+      reader->readTopLevel(&block, [](Operation *) { return true; })));
----------------
jpienaar wrote:

Would this reliably fail if this were wrong? (I'm assuming ASAN clearly shows it, but not sure regular fastbuild).

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