[msan] Don't modify CFG iterating it

[vitalybuka]

In rare cases SplitBlockAndInsertSimpleForLoop in paintOrigin crashes outsize iterators.

Somehow existing SplitBlockAndInsertIfThen do not invalidate iterators.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Vitaly Buka (vitalybuka)

Changes

In rare cases SplitBlockAndInsertSimpleForLoop in paintOrigin crashes outsize iterators.

Somehow existing SplitBlockAndInsertIfThen do not invalidate iterators.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp (+19-8)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index cc2295c44023c4..4a5f4a40726574 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -1135,6 +1135,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   std::unique_ptr<VarArgHelper> VAHelper;
   const TargetLibraryInfo *TLI;
   Instruction *FnPrologueEnd;
+  SmallVector<Instruction *, 128> Instructions;
 
   // The following flags disable parts of MSan instrumentation based on
   // exclusion list contents and command-line options.
@@ -1520,6 +1521,11 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     for (BasicBlock *BB : depth_first(FnPrologueEnd->getParent()))
       visit(*BB);
 
+    // `visit` above only collects instructions. Process them after iterating
+    // CFG to avoid requirement on CFG transformations.
+    for (Instruction *I : Instructions)
+      instrument(*I);
+
     // Finalize PHI nodes.
     for (PHINode *PN : ShadowPHINodes) {
       PHINode *PNS = cast<PHINode>(getShadow(PN));
@@ -2181,14 +2187,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     return ConstantDataVector::get(IRB.getContext(), OrderingTable);
   }
 
-  // ------------------- Visitors.
-  using InstVisitor<MemorySanitizerVisitor>::visit;
-  void visit(Instruction &I) {
-    if (I.getMetadata(LLVMContext::MD_nosanitize))
-      return;
-    // Don't want to visit if we're in the prologue
-    if (isInPrologue(I))
-      return;
+  void instrument(Instruction &I) {
     if (!DebugCounter::shouldExecute(DebugInstrumentInstruction)) {
       LLVM_DEBUG(dbgs() << "Skipping instruction: " << I << "\n");
       // We still need to set the shadow and origin to clean values.
@@ -2199,6 +2198,18 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     InstVisitor<MemorySanitizerVisitor>::visit(I);
   }
 
+  // ------------------- Visitors.
+  using InstVisitor<MemorySanitizerVisitor>::visit;
+  void visit(Instruction &I) {
+    if (I.getMetadata(LLVMContext::MD_nosanitize))
+      return;
+    // Don't want to visit if we're in the prologue
+    if (isInPrologue(I))
+      return;
+
+    Instructions.push_back(&I);
+  }
+
   /// Instrument LoadInst
   ///
   /// Loads the corresponding shadow and (optionally) origin.