Fix DataFlowSanitizer instrumentation pass to take parameter position changes into account for custom functions.

[skerner]

Summary

When DataFlowSanitizer transforms a call to a custom function, the new call has extra parameters. The attributes on parameters must be updated to take the new position of each parameter into account.

[skerner]

Polish

[pcc]

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

243

I think this can just be a std::vector<unsigned>.

251

Please use llvm naming conventions (i.e. CallSiteAttrs instead of call_site_attrs). Also elsewhere.

[skerner]

Fix review comments.

[skerner]

Address review comments.

[skerner]

Peter, thanks for the review!

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

243

I chose a map so that arguments to the original function which are not included in the transformed function can be excluded. However, I now see that this can not happen. Changed.

[pcc]

LGTM

[skerner]

Peter, can you please commit this change?

lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

public:
   }
 
   /// Returns whether this module is listed in the given category.
   bool isIn(const Module &M, StringRef Category) const {
     return SCL->inSection("dataflow", "src", M.getModuleIdentifier(), Category);
   }
 };
 
+/// TransformedFunction is used to express the result of transforming one
+/// function type into another.  This struct is immutable.  It holds metadata
+/// useful for updating calls of the old function to the new type.
+struct TransformedFunction {
+  TransformedFunction(FunctionType* OriginalType,
+                      FunctionType* TransformedType,
+                      llvm::DenseMap<unsigned, unsigned> ArgumentIndexMapping)
+      : OriginalType(OriginalType),
+        TransformedType(TransformedType),
+        ArgumentIndexMapping(ArgumentIndexMapping) {}
+
+  // Disallow copies.
+  TransformedFunction(const TransformedFunction&) = delete;
+  TransformedFunction& operator=(const TransformedFunction&) = delete;
+
+  // Allow moves.
+  TransformedFunction(TransformedFunction&&) = default;
+  TransformedFunction& operator=(TransformedFunction&&) = default;
+
+  /// Type of the function before the transformation.
+  FunctionType * const OriginalType;
+
+  /// Type of the function after the transformation.
+  FunctionType * const TransformedType;
+
+  /// Transforming a function may change the position of arguments.  This
+  /// member records the mapping from each argument's old position to its new
+  /// position.  If the transformation from F to F' made the first argument
+  /// of F into the third argument of F', then this map will contain (0->2).
+  const llvm::DenseMap<unsigned, unsigned> ArgumentIndexMapping;
+};
+
+/// Given function attributes from a call site for the original function,
+/// return function attributes appropriate for a call to the transformed
+/// function.
+AttributeList TransformFunctionAttributes(
+    const TransformedFunction& transformed_function,
+    LLVMContext& Ctx, AttributeList call_site_attrs) {
+
+  // Construct a vector of AttributeSet for each function argument.
+  std::vector<llvm::AttributeSet> argument_attributes(
+      transformed_function.TransformedType->getNumParams());
+
+  // 'ArgumentIndexMapping' holds a map from the original argument
+  // index to the index of the same argument in the transformed function.
+  // Use that mapping to
+  for (const auto& old_to_new : transformed_function.ArgumentIndexMapping) {
+    const unsigned old_param_index = old_to_new.first;
+    const unsigned new_param_index = old_to_new.second;
+    argument_attributes[new_param_index] =
+        call_site_attrs.getParamAttributes(old_param_index);
+  }
+
+  // Copy annotations on varargs arguments.
+  for (unsigned i = transformed_function.OriginalType->getNumParams(),
+           ie = call_site_attrs.getNumAttrSets(); i<ie; ++i) {
+    argument_attributes.push_back(call_site_attrs.getParamAttributes(i));
+  }
+
+  return AttributeList::get(
+      Ctx,
+      call_site_attrs.getFnAttributes(),
+      call_site_attrs.getRetAttributes(),
+      llvm::makeArrayRef(argument_attributes));
+}
+
 class DataFlowSanitizer : public ModulePass {
   friend struct DFSanFunction;
   friend class DFSanVisitor;
 
   enum {
     ShadowWidth = 16
   };
 
class DataFlowSanitizer : public ModulePass {
   AttrBuilder ReadOnlyNoneAttrs;
   bool DFSanRuntimeShadowMask = false;
 
   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);
+  TransformedFunction getCustomFunctionType(FunctionType *T);
   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) {
+TransformedFunction DataFlowSanitizer::getCustomFunctionType(FunctionType *T) {
   SmallVector<Type *, 4> ArgTypes;
-  for (FunctionType::param_iterator i = T->param_begin(), e = T->param_end();
-       i != e; ++i) {
+
+  // Some parameters of the custom function being constructed are
+  // parameters of T.  Record the mapping from parameters of T to
+  // parameters of the custom function, so that parameter attributes
+  // at call sites can be updated.
+  DenseMap<unsigned, unsigned> ArgumentIndexMapping;
+
+  for (unsigned i = 0, e = T->getNumParams(); i != e; ++i) {
+    Type* param_type = T->getParamType(i);
     FunctionType *FT;
-    if (isa<PointerType>(*i) && (FT = dyn_cast<FunctionType>(cast<PointerType>(
-                                     *i)->getElementType()))) {
+    if (isa<PointerType>(param_type) && (FT = dyn_cast<FunctionType>(cast<PointerType>(
+        param_type)->getElementType()))) {
+      ArgumentIndexMapping[i] = ArgTypes.size();
       ArgTypes.push_back(getTrampolineFunctionType(FT)->getPointerTo());
       ArgTypes.push_back(Type::getInt8PtrTy(*Ctx));
     } else {
-      ArgTypes.push_back(*i);
+      ArgumentIndexMapping[i] = ArgTypes.size();
+      ArgTypes.push_back(param_type);
     }
   }
   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());
+  return TransformedFunction(
+      T, FunctionType::get(T->getReturnType(), ArgTypes, T->isVarArg()),
+      ArgumentIndexMapping);
 }
 
 bool DataFlowSanitizer::doInitialization(Module &M) {
   Triple TargetTriple(M.getTargetTriple());
   bool IsX86_64 = TargetTriple.getArch() == Triple::x86_64;
   bool IsMIPS64 = TargetTriple.getArch() == Triple::mips64 ||
                   TargetTriple.getArch() == Triple::mips64el;
   bool IsAArch64 = TargetTriple.getArch() == 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);
+        TransformedFunction CustomFn = DFSF.DFS.getCustomFunctionType(FT);
         std::string CustomFName = "__dfsw_";
         CustomFName += F->getName();
-        Constant *CustomF =
-            DFSF.DFS.Mod->getOrInsertFunction(CustomFName, CustomFT);
+        Constant *CustomF = DFSF.DFS.Mod->getOrInsertFunction(
+            CustomFName, CustomFn.TransformedType);
         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);
           }
case DataFlowSanitizer::WK_Custom:
           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());
+        CustomCI->setAttributes(TransformFunctionAttributes(CustomFn,
+            CI->getContext(), CI->getAttributes()));
 
         // Update the parameter attributes of the custom call instruction to
         // zero extend the shadow parameters. This is required for targets
         // which consider ShadowTy an illegal type.
         for (unsigned n = 0; n < FT->getNumParams(); n++) {
           const unsigned ArgNo = ShadowArgStart + n;
           if (CustomCI->getArgOperand(ArgNo)->getType() == DFSF.DFS.ShadowTy)
             CustomCI->addParamAttr(ArgNo, Attribute::ZExt);

test/Instrumentation/DataFlowSanitizer/custom_fn_varargs_attributes.ll

+; RUN: opt < %s -dfsan -dfsan-args-abi -dfsan-abilist=%S/Inputs/abilist.txt -S | FileCheck %s
+; RUN: opt < %s -dfsan                 -dfsan-abilist=%S/Inputs/abilist.txt -S | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+; Declare a custom varargs function.
+declare i16 @custom_varargs(i64, ...)
+
+define i16 @call_custom_varargs(i8* %buf) {
+  ; CHECK-LABEL: @"dfs$call_custom_varargs"
+
+  ;; All arguments have an annotation.  Check that the transformed function
+  ;; preserves each annotation.
+
+  ; CHECK: call zeroext i16 (i64, i16, i16*, i16*, ...)
+  ; CHECK: @__dfsw_custom_varargs
+  ; CHECK: i64 signext 200
+  ; CHECK: i8* nonnull
+  ; CHECK: i64 zeroext 20
+  ; CHECK: i32 signext 1
+  %call = call zeroext i16 (i64, ...) @custom_varargs(
+    i64 signext 200,
+    i8* nonnull %buf,
+    i64 zeroext 20,
+    i32 signext 1
+  )
+  ret i16 %call
+}

test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll

+; RUN: opt < %s -dfsan -dfsan-args-abi -dfsan-abilist=%S/Inputs/abilist.txt -S | FileCheck %s
+; RUN: opt < %s -dfsan                 -dfsan-abilist=%S/Inputs/abilist.txt -S | FileCheck %s
+
+;; TODO: name this custom_function_attributes.ll
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@a_string = internal constant [9 x i8] c"a string\00", align 1
+
+; Declare custom functions.  Inputs/abilist.txt causes any function with a
+; name matching /custom.*/ to be a custom function.
+declare i32 @custom_fun_one_callback(i8 (i32, double)* %callback_arg)
+declare i32 @custom_fun_two_callbacks(
+  i8 (i32, double)* %callback_arg1,
+  i64 %an_int,
+  i8 (i32, double)* %callback_arg2
+)
+
+declare i8 @a_callback_fun(i32, double)
+
+define void @call_custom_funs_with_callbacks(i8 (i32, double)* %callback_arg) {
+  ; CHECK-LABEL: @"dfs$call_custom_funs_with_callbacks"
+
+  ;; The callback should have attribute 'nonnull':
+  ; CHECK: call signext i32 @__dfsw_custom_fun_one_callback(
+  ; CHECK: nonnull @"dfst0$custom_fun_one_callback"
+  %call1 = call signext i32 @custom_fun_one_callback(i8 (i32, double)* nonnull @a_callback_fun)
+
+  ;; Call a custom function with two callbacks.  Test that they have
+  ;; annotations 'nonnull' and 'noalias':
+  ; CHECK: call i32 @__dfsw_custom_fun_two_callbacks(
+  ; CHECK: nonnull @"dfst0$custom_fun_two_callbacks"
+  ; CHECK: i64 12345
+  ; CHECK: noalias @"dfst2$custom_fun_two_callbacks"
+  %call2 = call i32 @custom_fun_two_callbacks(
+    i8 (i32, double)* nonnull @a_callback_fun,
+    i64 12345,
+    i8 (i32, double)* noalias @a_callback_fun
+  )
+
+  ret void
+}
 No newline at end of file