Introduce support for custom wrappers for vararg functions

[pcc]

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1360

Can you give this a different name to avoid shadowing by the other FT on line 1394?

1412

Can't you just use CS.arg_size() here?

1432

Likewise.

[martignlo]

Done.

[martignlo]

Ping.

[pcc]

LGTM

[martignlo]

Closed by commit rL218671 (authored by @martignlo).

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

for (unsigned i = 0, e = T->getNumParams(); i != e; ++i)
     ArgTypes.push_back(ShadowTy);
   Type *RetType = T->getReturnType();
   if (!RetType->isVoidTy())
     ArgTypes.push_back(ShadowPtrTy);
   return FunctionType::get(T->getReturnType(), ArgTypes, false);
 }
 
 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) {
     FunctionType *FT;
     if (isa<PointerType>(*i) && (FT = dyn_cast<FunctionType>(cast<PointerType>(
                                      *i)->getElementType()))) {
       ArgTypes.push_back(getTrampolineFunctionType(FT)->getPointerTo());
       ArgTypes.push_back(Type::getInt8PtrTy(*Ctx));
DataFlowSanitizer::InstrumentedABI DataFlowSanitizer::getInstrumentedABI() {
   return ClArgsABI ? IA_Args : IA_TLS;
 }
 
 DataFlowSanitizer::WrapperKind DataFlowSanitizer::getWrapperKind(Function *F) {
   if (ABIList.isIn(*F, "functional"))
     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;
 }
 
 void DataFlowSanitizer::addGlobalNamePrefix(GlobalValue *GV) {
   std::string GVName = GV->getName(), Prefix = "dfs$";
   GV->setName(Prefix + GVName);
if (isInstrumented(&F)) {
             ConstantExpr::getBitCast(NewF, PointerType::getUnqual(FT)));
         NewF->takeName(&F);
         F.eraseFromParent();
         *i = NewF;
         addGlobalNamePrefix(NewF);
       } 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
       // WrapperKind is done in the second pass below.
       FunctionType *NewFT = getInstrumentedABI() == IA_Args
                                 ? getArgsFunctionType(FT)
                                 : FT;
       Function *NewF = buildWrapperFunction(
if (isInstrumented(&F)) {
         // by the push_back, taking care that the new range does not include
         // any functions added by this code.
         size_t N = i - FnsToInstrument.begin(),
                Count = e - FnsToInstrument.begin();
         FnsToInstrument.push_back(&F);
         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;
     }
   }
 
   for (std::vector<Function *>::iterator i = FnsToInstrument.begin(),
                                          e = FnsToInstrument.end();
        i != e; ++i) {
     if (!*i || (*i)->isDeclaration())
       continue;

 
 void DFSanVisitor::visitCallSite(CallSite CS) {
   Function *F = CS.getCalledFunction();
   if ((F && F->isIntrinsic()) || isa<InlineAsm>(CS.getCalledValue())) {
     visitOperandShadowInst(*CS.getInstruction());
     return;
   }
 
+  assert(!(cast<FunctionType>(
+      CS.getCalledValue()->getType()->getPointerElementType())->isVarArg() &&
+           dyn_cast<InvokeInst>(CS.getInstruction())));
+
   IRBuilder<> IRB(CS.getInstruction());
 
   DenseMap<Value *, Function *>::iterator i =
       DFSF.DFS.UnwrappedFnMap.find(CS.getCalledValue());
   if (i != DFSF.DFS.UnwrappedFnMap.end()) {
     Function *F = i->second;
     switch (DFSF.DFS.getWrapperKind(F)) {
     case DataFlowSanitizer::WK_Warning: {
case DataFlowSanitizer::WK_Custom: {
             CustomFn->removeAttributes(AttributeSet::FunctionIndex,
                                        DFSF.DFS.ReadOnlyNoneAttrs);
           }
         }
 
         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) &&
               (ParamFT = dyn_cast<FunctionType>(
                    cast<PointerType>(T)->getElementType()))) {
             std::string TName = "dfst";
             TName += utostr(FT->getNumParams() - n);
             TName += "$";
             TName += F->getName();
             Constant *T = DFSF.DFS.getOrBuildTrampolineFunction(ParamFT, TName);
             Args.push_back(T);
             Args.push_back(
                 IRB.CreateBitCast(*i, Type::getInt8PtrTy(*DFSF.DFS.Ctx)));
           } else {
             Args.push_back(*i);
           }
         }
 
         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()) {
           if (!DFSF.LabelReturnAlloca) {
             DFSF.LabelReturnAlloca =
                 new AllocaInst(DFSF.DFS.ShadowTy, "labelreturn",
                                DFSF.F->getEntryBlock().begin());
           }

test/Instrumentation/DataFlowSanitizer/abilist.ll

define i32 @functional(i32 %a, i32 %b) {
   %c = add i32 %a, %b
   ret i32 %c
 }
 
 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)
 
 declare i32 @cb(i32)
 
 ; CHECK: @"dfs$f"
 define void @f() {
   ; CHECK: %[[LABELRETURN:.*]] = alloca i16
 
   ; CHECK: call void @__dfsw_custom1(i32 1, i32 2, i16 0, i16 0)
   call void @custom1(i32 1, i32 2)
 
   ; CHECK: call i32 @__dfsw_custom2(i32 1, i32 2, i16 0, i16 0, i16* %[[LABELRETURN]])
   call i32 @custom2(i32 1, i32 2)
 
   ; 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
 }
 
 ; CHECK: define i32 (i32, i32)* @discardg(i32)
 ; CHECK: %[[CALL:.*]] = call { i32 (i32, i32)*, i16 } @"dfs$g"(i32 %0, i16 0)
 ; CHECK: %[[XVAL:.*]] = extractvalue { i32 (i32, i32)*, i16 } %[[CALL]], 0
 ; CHECK: ret {{.*}} %[[XVAL]]

 ; CHECK: %[[CALL:.*]] = call i32 @discard(i32 %0, i32 %1)
 ; CHECK: %[[IVAL0:.*]] = insertvalue { i32, i16 } undef, i32 %[[CALL]], 0
 ; CHECK: %[[IVAL1:.*]] = insertvalue { i32, i16 } %[[IVAL0]], i16 0, 1
 ; CHECK: ret { i32, i16 } %[[IVAL1]]
 @adiscard = alias i32 (i32, i32)* @discard
 
 ; 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)*
 ; CHECK: %[[CALL:.*]] = call { i32, i16 } %[[BC]](i32 %1, i16 %2)
 ; CHECK: %[[XVAL0:.*]] = extractvalue { i32, i16 } %[[CALL]], 0
 ; CHECK: %[[XVAL1:.*]] = extractvalue { i32, i16 } %[[CALL]], 1
 ; CHECK: store i16 %[[XVAL1]], i16* %3
 ; CHECK: ret i32 %[[XVAL0]]