Index: llvm/trunk/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp =================================================================== --- llvm/trunk/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp +++ llvm/trunk/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp @@ -211,6 +211,72 @@ } }; +/// 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, + std::vector 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. Argument positions are zero-indexed. If the transformation + /// from F to F' made the first argument of F into the third argument of F', + /// then ArgumentIndexMapping[0] will equal 2. + const std::vector 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& TransformedFunction, + LLVMContext& Ctx, AttributeList CallSiteAttrs) { + + // Construct a vector of AttributeSet for each function argument. + std::vector ArgumentAttributes( + TransformedFunction.TransformedType->getNumParams()); + + // Copy attributes from the parameter of the original function to the + // transformed version. 'ArgumentIndexMapping' holds the mapping from + // old argument position to new. + for (unsigned i=0, ie = TransformedFunction.ArgumentIndexMapping.size(); + i < ie; ++i) { + unsigned TransformedIndex = TransformedFunction.ArgumentIndexMapping[i]; + ArgumentAttributes[TransformedIndex] = CallSiteAttrs.getParamAttributes(i); + } + + // Copy annotations on varargs arguments. + for (unsigned i = TransformedFunction.OriginalType->getNumParams(), + ie = CallSiteAttrs.getNumAttrSets(); igetReturnType(), ArgTypes, false); } -FunctionType *DataFlowSanitizer::getCustomFunctionType(FunctionType *T) { +TransformedFunction DataFlowSanitizer::getCustomFunctionType(FunctionType *T) { SmallVector 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. + std::vector ArgumentIndexMapping; + for (unsigned i = 0, ie = T->getNumParams(); i != ie; ++i) { + Type* param_type = T->getParamType(i); FunctionType *FT; - if (isa(*i) && (FT = dyn_cast(cast( - *i)->getElementType()))) { + if (isa(param_type) && (FT = dyn_cast( + cast(param_type)->getElementType()))) { + ArgumentIndexMapping.push_back(ArgTypes.size()); ArgTypes.push_back(getTrampolineFunctionType(FT)->getPointerTo()); ArgTypes.push_back(Type::getInt8PtrTy(*Ctx)); } else { - ArgTypes.push_back(*i); + ArgumentIndexMapping.push_back(ArgTypes.size()); + ArgTypes.push_back(param_type); } } for (unsigned i = 0, e = T->getNumParams(); i != e; ++i) @@ -457,7 +531,9 @@ 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) { @@ -1459,11 +1535,11 @@ // wrapper. if (CallInst *CI = dyn_cast(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(CustomF)) { CustomFn->copyAttributesFrom(F); @@ -1531,7 +1607,8 @@ 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 Index: llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll =================================================================== --- llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll +++ llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll @@ -0,0 +1,37 @@ +; 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 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) + +; CHECK-LABEL: @"dfs$call_custom_funs_with_callbacks" +define void @call_custom_funs_with_callbacks(i8 (i32, double)* %callback_arg) { + ;; 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. Check their annotations. + ; 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 +} Index: llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_varargs_attributes.ll =================================================================== --- llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_varargs_attributes.ll +++ llvm/trunk/test/Instrumentation/DataFlowSanitizer/custom_fun_varargs_attributes.ll @@ -0,0 +1,27 @@ +; 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, ...) + +; CHECK-LABEL: @"dfs$call_custom_varargs" +define void @call_custom_varargs(i8* %buf) { + ;; 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 void +}