[msan] Wrap indirect functions.

Adds a flag to the MemorySanitizer pass that enables runtime rewriting of
indirect calls. This is part of MSanDR implementation and is needed to return
control to the DynamiRio-based helper tool on transition between instrumented
and non-instrumented modules. Disabled by default.


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@191006 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -157,6 +157,14 @@ static cl::opt<std::string>  ClBlacklistFile("msan-blacklist",
        cl::desc("File containing the list of functions where MemorySanitizer "
                 "should not report bugs"), cl::Hidden);
 
+// Experimental. Wraps all indirect calls in the instrumented code with
+// a call to the given function. This is needed to assist the dynamic
+// helper tool (MSanDR) to regain control on transition between instrumented and
+// non-instrumented code.
+static cl::opt<std::string> ClWrapIndirectCalls("msan-wrap-indirect-calls",
+       cl::desc("Wrap indirect calls with a given function"),
+       cl::Hidden);
+
 namespace {
 
 /// \brief An instrumentation pass implementing detection of uninitialized
@@ -172,8 +180,8 @@ class MemorySanitizer : public FunctionPass {
         TrackOrigins(TrackOrigins || ClTrackOrigins),
         TD(0),
         WarningFn(0),
-      BlacklistFile(BlacklistFile.empty() ? ClBlacklistFile
+        BlacklistFile(BlacklistFile.empty() ? ClBlacklistFile : BlacklistFile),
-                                          : BlacklistFile) { }
+        WrapIndirectCalls(!ClWrapIndirectCalls.empty()) {}
   const char *getPassName() const { return "MemorySanitizer"; }
   bool runOnFunction(Function &F);
   bool doInitialization(Module &M);
@@ -236,6 +244,12 @@ class MemorySanitizer : public FunctionPass {
   /// \brief An empty volatile inline asm that prevents callback merge.
   InlineAsm *EmptyAsm;
 
+  bool WrapIndirectCalls;
+  /// \brief Run-time wrapper for indirect calls.
+  Value *IndirectCallWrapperFn;
+  // Argument and return type of IndirectCallWrapperFn: void (*f)(void).
+  Type *AnyFunctionPtrTy;
+
   friend struct MemorySanitizerVisitor;
   friend struct VarArgAMD64Helper;
 };
@@ -329,6 +343,13 @@ void MemorySanitizer::initializeCallbacks(Module &M) {
   EmptyAsm = InlineAsm::get(FunctionType::get(IRB.getVoidTy(), false),
                             StringRef(""), StringRef(""),
                             /*hasSideEffects=*/true);
+
+  if (WrapIndirectCalls) {
+    AnyFunctionPtrTy =
+        PointerType::getUnqual(FunctionType::get(IRB.getVoidTy(), false));
+    IndirectCallWrapperFn = M.getOrInsertFunction(
+        ClWrapIndirectCalls, AnyFunctionPtrTy, AnyFunctionPtrTy, NULL);
+  }
 }
 
 /// \brief Module-level initialization.
@@ -1570,6 +1591,17 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     }
   }
 
+  // Replace call to (*Fn) with a call to (*IndirectCallWrapperFn(Fn)).
+  void wrapIndirectCall(IRBuilder<> &IRB, CallSite CS) {
+    Value *Fn = CS.getCalledValue();
+    Value *NewFn = IRB.CreateBitCast(
+        IRB.CreateCall(MS.IndirectCallWrapperFn,
+                       IRB.CreateBitCast(Fn, MS.AnyFunctionPtrTy)),
+        Fn->getType());
+    setShadow(NewFn, getShadow(Fn));
+    CS.setCalledFunction(NewFn);
+  }
+
   void visitCallSite(CallSite CS) {
     Instruction &I = *CS.getInstruction();
     assert((CS.isCall() || CS.isInvoke()) && "Unknown type of CallSite");
@@ -1608,6 +1640,10 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       }
     }
     IRBuilder<> IRB(&I);
+
+    if (MS.WrapIndirectCalls && !CS.getCalledFunction())
+      wrapIndirectCall(IRB, CS);
+
     unsigned ArgOffset = 0;
     DEBUG(dbgs() << "  CallSite: " << I << "\n");
     for (CallSite::arg_iterator ArgIt = CS.arg_begin(), End = CS.arg_end();
@@ -1651,7 +1687,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     DEBUG(dbgs() << "  done with call args\n");
 
     FunctionType *FT =
-      cast<FunctionType>(CS.getCalledValue()->getType()-> getContainedType(0));
+      cast<FunctionType>(CS.getCalledValue()->getType()->getContainedType(0));
     if (FT->isVarArg()) {
       VAHelper->visitCallSite(CS, IRB);
     }

test/Instrumentation/MemorySanitizer/wrap_indirect_calls.ll

@@ -0,0 +1,21 @@
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-wrap-indirect-calls=zzz -S | FileCheck %s
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Test for -msan-wrap-indirect-calls functionality.
+; Replaces indirect call to %f with a call to whatever is returned from the
+; wrapper function.
+
+; This does not depend on the sanitize_memory attribute.
+define i32 @func(i32 (i32, i32)* nocapture %f, i32 %x, i32 %y) {
+entry:
+  %call = tail call i32 %f(i32 %x, i32 %y)
+  ret i32 %call
+}
+
+; CHECK: @func
+; CHECK: bitcast i32 (i32, i32)* %f to void ()*
+; CHECK: call void ()* (void ()*)* @zzz(void ()*
+; CHECK: [[A:%[01-9a-z]+]] = bitcast void ()* {{.*}} to i32 (i32, i32)*
+; CHECK: call i32 {{.*}}[[A]](i32 {{.*}}, i32 {{.*}})
+; CHECK: ret i32