imsi

hygoni · Sep 5, 2024 · 104c791 · 104c791
1 parent 51586b2
commit 104c791
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/llvm/include/llvm/Transforms/Instrumentation/PreciseLeakSanitizer.h b/llvm/include/llvm/Transforms/Instrumentation/PreciseLeakSanitizer.h
@@ -85,6 +85,7 @@ class PreciseLeakSanVisitor : public InstVisitor<PreciseLeakSanVisitor> {
   void visitStoreInst(StoreInst &I);
   void visitReturnInst(ReturnInst &I);
   void visitCallInst(CallInst &I);
+  void visitInvokeInst(InvokeInst &I);
   void visitMemIntrinsics(MemIntrinsic &I);
   void pushNewLocalVarListStack();
   void pushNewBuiltinAllocaStack();

diff --git a/llvm/lib/Transforms/Instrumentation/PreciseLeakSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/PreciseLeakSanitizer.cpp
@@ -3,6 +3,7 @@
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
@@ -186,6 +187,31 @@ void PreciseLeakSanVisitor::visitCallInst(CallInst &I) {
     IntrinToInstrument.push_back(MI);
 }
 
+void PreciseLeakSanVisitor::visitInvokeInst(InvokeInst &I) {
+  if (I.getType()->isPointerTy()) {
+    Value *RetAddr = &I;
+    Instruction *LastInst = InstructionTraceTopDown(&I);
+    IRBuilder<> Builder(LastInst);
+    if (StoreInst *Inst = dyn_cast<StoreInst>(LastInst)) {
+      Value *CompareAddr = Inst->getValueOperand();
+      Plsan.CreateCallWithMetaData(Builder, Plsan.CheckReturnedOrStoredValueFn,
+                                   {RetAddr, CompareAddr});
+    } else if (ReturnInst *Inst = dyn_cast<ReturnInst>(LastInst)) {
+      Value *CompareAddr = Inst->getReturnValue();
+      // if return value is void
+      if (CompareAddr == NULL)
+        CompareAddr = ConstantPointerNull::get(Plsan.VoidPtrTy);
+      Plsan.CreateCallWithMetaData(Builder, Plsan.CheckReturnedOrStoredValueFn,
+                                   {RetAddr, CompareAddr});
+    } else {
+      Instruction *nextInst = I.getNextNode();
+      if (nextInst) 
+        Builder.SetInsertPoint(I.getNextNode());
+      Plsan.CreateCallWithMetaData(Builder, Plsan.CheckMemoryLeakFn, {RetAddr});
+    }
+  }
+}
+
 void PreciseLeakSanVisitor::visitMemIntrinsics(MemIntrinsic &I) {
   IRBuilder<> Builder(&I);
   if (MemSetInst *Inst = dyn_cast<MemSetInst>(&I)) {