[dfsan] New calling convention for custom functions with variadic arguments.

[pcc]

Summary

The previous calling convention prevented custom functions from being able
to access argument labels unless it knew how many variadic arguments there
were, and of which type. This restriction made it impossible to correctly
model functions in the printf family, as it is legal to pass more arguments
than required to those functions. We now pass arguments in the following order:

non-vararg arguments
labels for non-vararg arguments
[if vararg function, pointer to array of labels for vararg arguments]
[if non-void function, pointer to label for return value]
vararg arguments

[martignlo]

LGTM

[pcc]

Closed by commit rL220906 (authored by @pcc).

llvm/trunk/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) {
-  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));
     } else {
       ArgTypes.push_back(*i);
     }
   }
   for (unsigned i = 0, e = T->getNumParams(); i != e; ++i)
     ArgTypes.push_back(ShadowTy);
+  if (T->isVarArg())
+    ArgTypes.push_back(ShadowPtrTy);
   Type *RetType = T->getReturnType();
   if (!RetType->isVoidTy())
     ArgTypes.push_back(ShadowPtrTy);
-  return FunctionType::get(T->getReturnType(), ArgTypes, false);
+  return FunctionType::get(T->getReturnType(), ArgTypes, T->isVarArg());
 }
 
 bool DataFlowSanitizer::doInitialization(Module &M) {
   DataLayoutPass *DLP = getAnalysisIfAvailable<DataLayoutPass>();
   if (!DLP)
     report_fatal_error("data layout missing");
   DL = &DLP->getDataLayout();
 
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 = CS.arg_size(); n != 0; ++i, --n) {
+        for (unsigned n = FT->getNumParams(); 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 = CS.arg_size(); n != 0; ++i, --n)
+        for (unsigned n = FT->getNumParams(); n != 0; ++i, --n)
           Args.push_back(DFSF.getShadow(*i));
 
+        if (FT->isVarArg()) {
+          auto LabelVAAlloca =
+              new AllocaInst(ArrayType::get(DFSF.DFS.ShadowTy,
+                                            CS.arg_size() - FT->getNumParams()),
+                             "labelva", DFSF.F->getEntryBlock().begin());
+
+          for (unsigned n = 0; i != CS.arg_end(); ++i, ++n) {
+            auto LabelVAPtr = IRB.CreateStructGEP(LabelVAAlloca, n);
+            IRB.CreateStore(DFSF.getShadow(*i), LabelVAPtr);
+          }
+
+          Args.push_back(IRB.CreateStructGEP(LabelVAAlloca, 0));
+        }
+
         if (!FT->getReturnType()->isVoidTy()) {
           if (!DFSF.LabelReturnAlloca) {
             DFSF.LabelReturnAlloca =
                 new AllocaInst(DFSF.DFS.ShadowTy, "labelreturn",
                                DFSF.F->getEntryBlock().begin());
           }
           Args.push_back(DFSF.LabelReturnAlloca);
         }
 
+        for (i = CS.arg_begin() + FT->getNumParams(); i != CS.arg_end(); ++i)
+          Args.push_back(*i);
+
         CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
         CustomCI->setCallingConv(CI->getCallingConv());
         CustomCI->setAttributes(CI->getAttributes());
 
         if (!FT->getReturnType()->isVoidTy()) {
           LoadInst *LabelLoad = IRB.CreateLoad(DFSF.LabelReturnAlloca);
           DFSF.setShadow(CustomCI, LabelLoad);
         }

llvm/trunk/test/Instrumentation/DataFlowSanitizer/abilist.ll

 ; RUN: opt < %s -dfsan -dfsan-args-abi -dfsan-abilist=%S/Inputs/abilist.txt -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"
 
 ; CHECK: i32 @discard(i32 %a, i32 %b)
 define i32 @discard(i32 %a, i32 %b) {
   ret i32 0
 }
 
 ; CHECK: i32 @functional(i32 %a, i32 %b)
 define i32 @functional(i32 %a, i32 %b) {
   %c = add i32 %a, %b
   ret i32 %c
 }
 
+; 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]]
+@discardg = alias i32 (i32, i32)* (i32)* @g
+
 declare void @custom1(i32 %a, i32 %b)
 
+; CHECK: define linkonce_odr { i32, i16 } @"dfsw$custom2"(i32, i32, i16, i16)
+; CHECK: %[[LABELRETURN2:.*]] = alloca i16
+; CHECK: %[[RV:.*]] = call i32 @__dfsw_custom2
+; CHECK: %[[RVSHADOW:.*]] = load i16* %[[LABELRETURN2]]
+; CHECK: insertvalue {{.*}}[[RV]], 0
+; CHECK: insertvalue {{.*}}[[RVSHADOW]], 1
+; CHECK: ret { i32, i16 }
 declare i32 @custom2(i32 %a, i32 %b)
 
 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() {
+define void @f(i32 %x) {
+  ; CHECK: %[[LABELVA2:.*]] = alloca [2 x i16]
+  ; CHECK: %[[LABELVA1:.*]] = alloca [2 x i16]
   ; 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 (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]])
+  ; CHECK: %[[LABELVA1_0:.*]] = getelementptr inbounds [2 x i16]* %[[LABELVA1]], i32 0, i32 0
+  ; CHECK: store i16 0, i16* %[[LABELVA1_0]]
+  ; CHECK: %[[LABELVA1_1:.*]] = getelementptr inbounds [2 x i16]* %[[LABELVA1]], i32 0, i32 1
+  ; CHECK: store i16 %{{.*}}, i16* %[[LABELVA1_1]]
+  ; CHECK: %[[LABELVA1_0A:.*]] = getelementptr inbounds [2 x i16]* %[[LABELVA1]], i32 0, i32 0
+  ; CHECK: call void (i32, i16, i16*, ...)* @__dfsw_custom3(i32 1, i16 0, i16* %[[LABELVA1_0A]], i32 2, i32 %{{.*}})
+  call void (i32, ...)* @custom3(i32 1, i32 2, i32 %x)
+
+  ; CHECK: %[[LABELVA2_0:.*]] = getelementptr inbounds [2 x i16]* %[[LABELVA2]], i32 0, i32 0
+  ; CHECK: %[[LABELVA2_0A:.*]] = getelementptr inbounds [2 x i16]* %[[LABELVA2]], i32 0, i32 0
+  ; CHECK: call i32 (i32, i16, i16*, i16*, ...)* @__dfsw_custom4(i32 1, i16 0, i16* %[[LABELVA2_0A]], i16* %[[LABELRETURN]], i32 2, i32 3)
   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]]
-@discardg = alias i32 (i32, i32)* (i32)* @g
-
-; CHECK: define linkonce_odr { i32, i16 } @"dfsw$custom2"(i32, i32, i16, i16)
-; CHECK: %[[LABELRETURN2:.*]] = alloca i16
-; CHECK: %[[RV:.*]] = call i32 @__dfsw_custom2
-; CHECK: %[[RVSHADOW:.*]] = load i16* %[[LABELRETURN2]]
-; CHECK: insertvalue {{.*}}[[RV]], 0
-; CHECK: insertvalue {{.*}}[[RVSHADOW]], 1
-; CHECK: ret { i32, i16 }
-
 ; CHECK: @"dfs$g"
 define i32 (i32, i32)* @g(i32) {
   ; CHECK: ret {{.*}} @"dfsw$custom2"
   ret i32 (i32, i32)* @custom2
 }
 
 ; CHECK: define { i32, i16 } @"dfs$adiscard"(i32, i32, i16, i16)
 ; 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: declare void @__dfsw_custom3(i32, i16, i16*, ...)
+; CHECK: declare i32 @__dfsw_custom4(i32, i16, i16*, i16*, ...)
 
 ; 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]]
+