[dfsan] Add explicit zero extensions for shadow parameters in function wrappers.

[sdardis]

Summary

In the case where dfsan provides a custom wrapper for a function,
shadow parameters are added for each parameter of the function.
These parameters are i16s. For targets which do not consider this
a legal type, the lack of sign extension information would cause
LLVM to generate anyexts around their usage with phi variables
and calling convention logic.

Address this by introducing zero exts for each shadow parameter.

[pcc]

Test case?

[sdardis]

Add test case, change logic for add zext attributes so that all shadow parameters are covered, not just the first.

[sdardis]

Ping.

[pcc]

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1529

I'm a little uncomfortable with this because it seems like we're effectively dropping attributes from the call instruction. If we don't expect to create a function definition the attributes on the declaration shouldn't matter, so maybe we should be adding the zeroext attributes just to the call rather than to the declaration.

[sdardis]

Address review comments.

[pcc]

LGTM

test/Instrumentation/DataFlowSanitizer/shadow-args-zext.ll

10

Prefer to move this and the two below inline with the original IR.

[sdardis]

Thanks for the review.

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

class DataFlowSanitizer : public ModulePass {
   AttrBuilder ReadOnlyNoneAttrs;
   bool DFSanRuntimeShadowMask;
 
   Value *getShadowAddress(Value *Addr, Instruction *Pos);
   bool isInstrumented(const Function *F);
   bool isInstrumented(const GlobalAlias *GA);
   FunctionType *getArgsFunctionType(FunctionType *T);
   FunctionType *getTrampolineFunctionType(FunctionType *T);
-  FunctionType *getCustomFunctionType(FunctionType *T);
+  Constant *getCustomFunction(Function *F);
   InstrumentedABI getInstrumentedABI();
   WrapperKind getWrapperKind(Function *F);
   void addGlobalNamePrefix(GlobalValue *GV);
   Function *buildWrapperFunction(Function *F, StringRef NewFName,
                                  GlobalValue::LinkageTypes NewFLink,
                                  FunctionType *NewFT);
   Constant *getOrBuildTrampolineFunction(FunctionType *FT, StringRef FName);
 
FunctionType *DataFlowSanitizer::getTrampolineFunctionType(FunctionType *T) {
   ArgTypes.append(T->param_begin(), T->param_end());
   ArgTypes.append(T->getNumParams(), ShadowTy);
   Type *RetType = T->getReturnType();
   if (!RetType->isVoidTy())
     ArgTypes.push_back(ShadowPtrTy);
   return FunctionType::get(T->getReturnType(), ArgTypes, false);
 }
 
-FunctionType *DataFlowSanitizer::getCustomFunctionType(FunctionType *T) {
+Constant * DataFlowSanitizer::getCustomFunction(Function *F) {
+  FunctionType * T = F->getFunctionType();
   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);
     }
   }
+
+  // ShadowTy may not be legal on all targets. Using it as a function
+  // parameter may lead to anyextends being generated which may cause
+  // unpredictable effects on some targets where i16 is not a legal
+  // type, e.g. MIPS.
+  const unsigned ShadowParamStart = ArgTypes.size() + 1;
   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, T->isVarArg());
+
+  FunctionType * CustomFT =
+      FunctionType::get(T->getReturnType(), ArgTypes, T->isVarArg());
+  std::string CustomFName = "__dfsw_";
+  CustomFName += F->getName();
+
+  Constant *CustomF = Mod->getOrInsertFunction(CustomFName, CustomFT);
+  if (Function *CustomFn = dyn_cast<Function>(CustomF)) {
+    CustomFn->copyAttributesFrom(F);
+
+    // Enforce zero extensions on shadow parameters.
+    for (unsigned i = 0, e = T->getNumParams(); i != e; ++i)
+      CustomFn->addAttribute(ShadowParamStart + i, Attribute::ZExt);
+
+    // Custom functions returning non-void will write to the return label.
+    if (!T->getReturnType()->isVoidTy()) {
+      CustomFn->removeAttributes(AttributeList::FunctionIndex,
+                                 ReadOnlyNoneAttrs);
+    }
+  }
+
+  return CustomF;
 }
 
 bool DataFlowSanitizer::doInitialization(Module &M) {
   llvm::Triple TargetTriple(M.getTargetTriple());
   bool IsX86_64 = TargetTriple.getArch() == llvm::Triple::x86_64;
   bool IsMIPS64 = TargetTriple.getArch() == llvm::Triple::mips64 ||
                   TargetTriple.getArch() == llvm::Triple::mips64el;
   bool IsAArch64 = TargetTriple.getArch() == llvm::Triple::aarch64 ||
case DataFlowSanitizer::WK_Functional: {
       visitOperandShadowInst(*CS.getInstruction());
       return;
     }
     case DataFlowSanitizer::WK_Custom: {
       // Don't try to handle invokes of custom functions, it's too complicated.
       // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
       // wrapper.
       if (CallInst *CI = dyn_cast<CallInst>(CS.getInstruction())) {
-        FunctionType *FT = F->getFunctionType();
-        FunctionType *CustomFT = DFSF.DFS.getCustomFunctionType(FT);
-        std::string CustomFName = "__dfsw_";
-        CustomFName += F->getName();
-        Constant *CustomF =
-            DFSF.DFS.Mod->getOrInsertFunction(CustomFName, CustomFT);
-        if (Function *CustomFn = dyn_cast<Function>(CustomF)) {
-          CustomFn->copyAttributesFrom(F);
-
-          // Custom functions returning non-void will write to the return label.
-          if (!FT->getReturnType()->isVoidTy()) {
-            CustomFn->removeAttributes(AttributeList::FunctionIndex,
-                                       DFSF.DFS.ReadOnlyNoneAttrs);
-          }
-        }
+        FunctionType * FT = F->getFunctionType();
+        Constant *CustomF = DFSF.DFS.getCustomFunction(F);
 
         std::vector<Value *> Args;
 
         CallSite::arg_iterator i = CS.arg_begin();
         for (unsigned n = FT->getNumParams(); n != 0; ++i, --n) {
           Type *T = (*i)->getType();
           FunctionType *ParamFT;
           if (isa<PointerType>(T) &&
case DataFlowSanitizer::WK_Custom: {
                              "labelreturn", &DFSF.F->getEntryBlock().front());
           }
           Args.push_back(DFSF.LabelReturnAlloca);
         }
 
         for (i = CS.arg_begin() + FT->getNumParams(); i != CS.arg_end(); ++i)
           Args.push_back(*i);
 
+        // Create a call to the custom function, matching the calling
+        // convention and call attributes. We use the attributes of the custom
+        // function rather the existing call instruction, as the addition of shadow
+        // parameters would not have the necessary ZExt attributes.
         CallInst *CustomCI = IRB.CreateCall(CustomF, Args);
+        Function *CustomFn = dyn_cast<Function>(CustomF);
         CustomCI->setCallingConv(CI->getCallingConv());
-        CustomCI->setAttributes(CI->getAttributes());
+        CustomCI->setAttributes(CustomFn ? CustomFn->getAttributes()
+                                         : CI->getAttributes());
 
         if (!FT->getReturnType()->isVoidTy()) {
           LoadInst *LabelLoad = IRB.CreateLoad(DFSF.LabelReturnAlloca);
           DFSF.setShadow(CustomCI, LabelLoad);
         }
 
         CI->replaceAllUsesWith(CustomCI);
         CI->eraseFromParent();

test/Instrumentation/DataFlowSanitizer/Inputs/shadow-args-abilist.txt

+fun:dfsan_get_label=uninstrumented
+fun:dfsan_get_label=custom
+
+fun:k2=uninstrumented
+fun:k2=custom
+
+fun:k4=uninstrumented
+fun:k4=custom

test/Instrumentation/DataFlowSanitizer/abilist.ll

 declare i32 @cb(i32)
 
 ; CHECK: @"dfs$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)
+  ; CHECK: call void @__dfsw_custom1(i32 1, i32 2, i16 zeroext 0, i16 zeroext 0)
   call void @custom1(i32 1, i32 2)
 
-  ; CHECK: call i32 @__dfsw_custom2(i32 1, i32 2, i16 0, i16 0, i16* %[[LABELRETURN]])
+  ; CHECK: call i32 @__dfsw_custom2(i32 1, i32 2, i16 zeroext 0, i16 zeroext 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)
+  ; CHECK: call void @__dfsw_customcb({{.*}} @"dfst0$customcb", i8* bitcast ({{.*}} @"dfs$cb" to i8*), i16 zeroext 0)
   call void @customcb(i32 (i32)* @cb)
 
   ; CHECK: %[[LABELVA1_0:.*]] = getelementptr inbounds [2 x i16], [2 x i16]* %[[LABELVA1]], i32 0, i32 0
   ; CHECK: store i16 0, i16* %[[LABELVA1_0]]
   ; CHECK: %[[LABELVA1_1:.*]] = getelementptr inbounds [2 x i16], [2 x i16]* %[[LABELVA1]], i32 0, i32 1
   ; CHECK: store i16 %{{.*}}, i16* %[[LABELVA1_1]]
   ; CHECK: %[[LABELVA1_0A:.*]] = getelementptr inbounds [2 x i16], [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 %{{.*}})
+  ; CHECK: call void (i32, i16, i16*, ...) @__dfsw_custom3(i32 1, i16 zeroext 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], [2 x i16]* %[[LABELVA2]], i32 0, i32 0
   ; CHECK: %[[LABELVA2_0A:.*]] = getelementptr inbounds [2 x i16], [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)
+  ; CHECK: call i32 (i32, i16, i16*, i16*, ...) @__dfsw_custom4(i32 1, i16 zeroext 0, i16* %[[LABELVA2_0A]], i16* %[[LABELRETURN]], i32 2, i32 3)
   call i32 (i32, ...) @custom4(i32 1, i32 2, i32 3)
 
   ret void
 }
 
 ; 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), 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_custom1(i32, i32, i16 zeroext, i16 zeroext)
+; CHECK: declare i32 @__dfsw_custom2(i32, i32, i16 zeroext, i16 zeroext, 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]]
 
-; CHECK: declare void @__dfsw_custom3(i32, i16, i16*, ...)
-; CHECK: declare i32 @__dfsw_custom4(i32, i16, i16*, i16*, ...)
+; CHECK: declare void @__dfsw_custom3(i32, i16 zeroext, i16*, ...)
+; CHECK: declare i32 @__dfsw_custom4(i32, i16 zeroext, i16*, i16*, ...)

test/Instrumentation/DataFlowSanitizer/shadow-args-zext.ll

+; RUN: opt < %s -dfsan -S --dfsan-abilist=%S/Inputs/shadow-args-abilist.txt  | 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 that the custom abi marks shadow parameters as zero extended.
+
+; CHECK-LABEL: @"dfs$m"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_dfsan_get_label(i64 signext 56, i16 zeroext 0, i16* %{{.*}})
+
+; CHECK-LABEL: @"dfs$k"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_k2(i64 signext 56, i64 signext 67, i16 zeroext {{.*}}, i16 zeroext {{.*}}, i16* %{{.*}})
+
+; CHECK-LABEL: @"dfs$k3"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_k4(i64 signext 56, i64 signext 67, i64 signext 78, i64 signext 89, i16 zeroext {{.*}}, i16 zeroext {{.*}}, i16 zeroext {{.*}}, i16 zeroext {{.*}}, i16* %{{.*}})
+
+
+; CHECK-LABEL: @"dfsw$dfsan_get_label"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_dfsan_get_label(i64 signext %0, i16 zeroext %1, i16* %{{.*}})
+
+; CHECK-LABEL: @"dfsw$k2"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_k2(i64 signext %{{.*}}, i64 signext %{{.*}}, i16 zeroext %{{.*}}, i16 zeroext %{{.*}}, i16* %{{.*}})
+
+; CHECK-LABEL: @"dfsw$k4"
+; CHECK: %{{.*}} = call zeroext i16 @__dfsw_k4(i64 signext %{{.*}}, i64 signext %{{.*}}, i64 signext %{{.*}}, i64 signext %{{.*}}, i16 zeroext %{{.*}}, i16 zeroext %{{.*}}, i16 zeroext %{{.*}}, i16 zeroext %{{.*}}, i16* %{{.*}})
+
+
+; CHECK: declare zeroext i16 @__dfsw_dfsan_get_label(i64 signext, i16 zeroext, i16*)
+; CHECK: declare zeroext i16 @__dfsw_k2(i64 signext, i64 signext, i16 zeroext, i16 zeroext, i16*)
+; CHECK: declare zeroext i16 @__dfsw_k4(i64 signext, i64 signext, i64 signext, i64 signext, i16 zeroext, i16 zeroext, i16 zeroext, i16 zeroext, i16*)
+
+define i32 @m() {
+entry:
+  %call = call zeroext i16 @dfsan_get_label(i64 signext 56)
+  %conv = zext i16 %call to i32
+  ret i32 %conv
+}
+
+define i32 @k() {
+entry:
+  %call = call zeroext i16 @k2(i64 signext 56, i64 signext 67)
+  %conv = zext i16 %call to i32
+  ret i32 %conv
+}
+
+define i32 @k3() {
+entry:
+  %call = call zeroext i16 @k4(i64 signext 56, i64 signext 67, i64 signext 78, i64 signext 89)
+  %conv = zext i16 %call to i32
+  ret i32 %conv
+}
+
+declare zeroext i16 @dfsan_get_label(i64 signext)
+declare zeroext i16 @k2(i64 signext, i64 signext)
+declare zeroext i16 @k4(i64 signext, i64 signext, i64 signext, i64 signext)