Index: lib/Transforms/Instrumentation/DataFlowSanitizer.cpp =================================================================== --- lib/Transforms/Instrumentation/DataFlowSanitizer.cpp +++ 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, + llvm::DenseMap 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 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 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(); 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. + DenseMap ArgumentIndexMapping; + + for (unsigned i = 0, e = T->getNumParams(); i != e; ++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[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) @@ -457,7 +532,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 +1536,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 +1608,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: test/Instrumentation/DataFlowSanitizer/custom_fn_varargs_attributes.ll =================================================================== --- /dev/null +++ test/Instrumentation/DataFlowSanitizer/custom_fn_varargs_attributes.ll @@ -0,0 +1,28 @@ +; 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 +} Index: test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll =================================================================== --- /dev/null +++ test/Instrumentation/DataFlowSanitizer/custom_fun_callback_attributes.ll @@ -0,0 +1,42 @@ +; 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