From f49592dddca90754b17e121f0166092ba3659f63 Mon Sep 17 00:00:00 2001
From: Lorenzo Martignoni <martignlo@google.com>
Date: Tue, 30 Sep 2014 12:33:16 +0000
Subject: [PATCH] Introduce support for custom wrappers for vararg functions.

Differential Revision: http://reviews.llvm.org/D5412



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@218671 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../Instrumentation/DataFlowSanitizer.cpp     | 27 ++++++++++++-------
 .../DataFlowSanitizer/abilist.ll              | 14 +++++++---
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index 771dcd7d9fe..446bcf79b55 100644
--- a/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -387,7 +387,12 @@ FunctionType *DataFlowSanitizer::getTrampolineFunctionType(FunctionType *T) {
 }
 
 FunctionType *DataFlowSanitizer::getCustomFunctionType(FunctionType *T) {
-  assert(!T->isVarArg());
+  if (T->isVarArg()) {
+    // The labels are passed after all the arguments so there is no need to
+    // adjust the function type.
+    return T;
+  }
+
   llvm::SmallVector<Type *, 4> ArgTypes;
   for (FunctionType::param_iterator i = T->param_begin(), e = T->param_end();
        i != e; ++i) {
@@ -476,7 +481,7 @@ DataFlowSanitizer::WrapperKind DataFlowSanitizer::getWrapperKind(Function *F) {
     return WK_Functional;
   if (ABIList.isIn(*F, "discard"))
     return WK_Discard;
-  if (ABIList.isIn(*F, "custom") && !F->isVarArg())
+  if (ABIList.isIn(*F, "custom"))
     return WK_Custom;
 
   return WK_Warning;
@@ -704,11 +709,6 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
       } else {
         addGlobalNamePrefix(&F);
       }
-               // Hopefully, nobody will try to indirectly call a vararg
-               // function... yet.
-    } else if (FT->isVarArg()) {
-      UnwrappedFnMap[&F] = &F;
-      *i = nullptr;
     } else if (!IsZeroArgsVoidRet || getWrapperKind(&F) == WK_Custom) {
       // Build a wrapper function for F.  The wrapper simply calls F, and is
       // added to FnsToInstrument so that any instrumentation according to its
@@ -744,6 +744,11 @@ bool DataFlowSanitizer::runOnModule(Module &M) {
         i = FnsToInstrument.begin() + N;
         e = FnsToInstrument.begin() + Count;
       }
+               // Hopefully, nobody will try to indirectly call a vararg
+               // function... yet.
+    } else if (FT->isVarArg()) {
+      UnwrappedFnMap[&F] = &F;
+      *i = nullptr;
     }
   }
 
@@ -1352,6 +1357,10 @@ void DFSanVisitor::visitCallSite(CallSite CS) {
     return;
   }
 
+  assert(!(cast<FunctionType>(
+      CS.getCalledValue()->getType()->getPointerElementType())->isVarArg() &&
+           dyn_cast<InvokeInst>(CS.getInstruction())));
+
   IRBuilder<> IRB(CS.getInstruction());
 
   DenseMap<Value *, Function *>::iterator i =
@@ -1400,7 +1409,7 @@ void DFSanVisitor::visitCallSite(CallSite CS) {
         std::vector<Value *> Args;
 
         CallSite::arg_iterator i = CS.arg_begin();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n) {
+        for (unsigned n = CS.arg_size(); n != 0; ++i, --n) {
           Type *T = (*i)->getType();
           FunctionType *ParamFT;
           if (isa<PointerType>(T) &&
@@ -1420,7 +1429,7 @@ void DFSanVisitor::visitCallSite(CallSite CS) {
         }
 
         i = CS.arg_begin();
-        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
+        for (unsigned n = CS.arg_size(); n != 0; ++i, --n)
           Args.push_back(DFSF.getShadow(*i));
 
         if (!FT->getReturnType()->isVoidTy()) {
diff --git a/test/Instrumentation/DataFlowSanitizer/abilist.ll b/test/Instrumentation/DataFlowSanitizer/abilist.ll
index 53b7d86dbda..ada5390057a 100644
--- a/test/Instrumentation/DataFlowSanitizer/abilist.ll
+++ b/test/Instrumentation/DataFlowSanitizer/abilist.ll
@@ -16,7 +16,9 @@ declare void @custom1(i32 %a, i32 %b)
 
 declare i32 @custom2(i32 %a, i32 %b)
 
-declare void @custom3(...)
+declare void @custom3(i32 %a, ...)
+
+declare i32 @custom4(i32 %a, ...)
 
 declare void @customcb(i32 (i32)* %cb)
 
@@ -35,9 +37,11 @@ define void @f() {
   ; CHECK: call void @__dfsw_customcb({{.*}} @"dfst0$customcb", i8* bitcast ({{.*}} @"dfs$cb" to i8*), i16 0)
   call void @customcb(i32 (i32)* @cb)
 
-  ; CHECK: call void @__dfsan_unimplemented
-  ; CHECK: call void (...)* @custom3()
-  call void (...)* @custom3()
+  ; CHECK: call void (i32, ...)* @__dfsw_custom3(i32 1, i32 2, i32 3, i16 0, i16 0, i16 0)
+  call void (i32, ...)* @custom3(i32 1, i32 2, i32 3)
+
+  ; CHECK: call i32 (i32, ...)* @__dfsw_custom4(i32 1, i32 2, i32 3, i16 0, i16 0, i16 0, i16* %[[LABELRETURN]])
+  call i32 (i32, ...)* @custom4(i32 1, i32 2, i32 3)
 
   ret void
 }
@@ -71,6 +75,8 @@ define i32 (i32, i32)* @g(i32) {
 
 ; CHECK: declare void @__dfsw_custom1(i32, i32, i16, i16)
 ; CHECK: declare i32 @__dfsw_custom2(i32, i32, i16, i16, i16*)
+; CHECK: declare void @__dfsw_custom3(i32, ...)
+; CHECK: declare i32 @__dfsw_custom4(i32, ...)
 
 ; CHECK-LABEL: define linkonce_odr i32 @"dfst0$customcb"(i32 (i32)*, i32, i16, i16*)
 ; CHECK: %[[BC:.*]] = bitcast i32 (i32)* %0 to { i32, i16 } (i32, i16)*