diff --git a/compiler-rt/lib/dfsan/dfsan.cpp b/compiler-rt/lib/dfsan/dfsan.cpp --- a/compiler-rt/lib/dfsan/dfsan.cpp +++ b/compiler-rt/lib/dfsan/dfsan.cpp @@ -128,6 +128,17 @@ fname); } +extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __dfsan_wrapper_extern_weak_null( + const void *addr, char *fname) { + if (!addr) + Report( + "ERROR: DataFlowSanitizer: dfsan generated wrapper calling null " + "extern_weak function %s\nIf this only happens with dfsan, the " + "dfsan instrumentation pass may be accidentally optimizing out a " + "null check\n", + fname); +} + // Use '-mllvm -dfsan-debug-nonzero-labels' and break on this function // to try to figure out where labels are being introduced in a nominally // label-free program. diff --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp @@ -439,6 +439,7 @@ FunctionType *DFSanUnionLoadFnTy; FunctionType *DFSanLoadLabelAndOriginFnTy; FunctionType *DFSanUnimplementedFnTy; + FunctionType *DFSanWrapperExternWeakNullFnTy; FunctionType *DFSanSetLabelFnTy; FunctionType *DFSanNonzeroLabelFnTy; FunctionType *DFSanVarargWrapperFnTy; @@ -454,6 +455,7 @@ FunctionCallee DFSanUnionLoadFn; FunctionCallee DFSanLoadLabelAndOriginFn; FunctionCallee DFSanUnimplementedFn; + FunctionCallee DFSanWrapperExternWeakNullFn; FunctionCallee DFSanSetLabelFn; FunctionCallee DFSanNonzeroLabelFn; FunctionCallee DFSanVarargWrapperFn; @@ -490,6 +492,7 @@ TransformedFunction getCustomFunctionType(FunctionType *T); WrapperKind getWrapperKind(Function *F); void addGlobalNameSuffix(GlobalValue *GV); + void buildExternWeakCheckIfNeeded(IRBuilder<> &IRB, Function *F); Function *buildWrapperFunction(Function *F, StringRef NewFName, GlobalValue::LinkageTypes NewFLink, FunctionType *NewFT); @@ -1041,6 +1044,10 @@ /*isVarArg=*/false); DFSanUnimplementedFnTy = FunctionType::get( Type::getVoidTy(*Ctx), Type::getInt8PtrTy(*Ctx), /*isVarArg=*/false); + Type *DFSanWrapperExternWeakNullArgs[2] = {Int8Ptr, Int8Ptr}; + DFSanWrapperExternWeakNullFnTy = + FunctionType::get(Type::getVoidTy(*Ctx), DFSanWrapperExternWeakNullArgs, + /*isVarArg=*/false); Type *DFSanSetLabelArgs[4] = {PrimitiveShadowTy, OriginTy, Type::getInt8PtrTy(*Ctx), IntptrTy}; DFSanSetLabelFnTy = FunctionType::get(Type::getVoidTy(*Ctx), @@ -1132,6 +1139,23 @@ } } +void DataFlowSanitizer::buildExternWeakCheckIfNeeded(IRBuilder<> &IRB, + Function *F) { + // If the function we are wrapping was ExternWeak, it may be null. + // The original code before calling this wrapper may have checked for null, + // but replacing with a known-to-not-be-null wrapper can break this check. + // When replacing uses of the extern weak function with the wrapper we try + // to avoid replacing uses in conditionals, but this is not perfect. + // In the case where we fail, and accidentially optimize out a null check + // for a extern weak function, add a check here to help identify the issue. + if (GlobalValue::isExternalWeakLinkage(F->getLinkage())) { + std::vector Args; + Args.push_back(IRB.CreatePointerCast(F, IRB.getInt8PtrTy())); + Args.push_back(IRB.CreateGlobalStringPtr(F->getName())); + IRB.CreateCall(DFSanWrapperExternWeakNullFn, Args); + } +} + Function * DataFlowSanitizer::buildWrapperFunction(Function *F, StringRef NewFName, GlobalValue::LinkageTypes NewFLink, @@ -1184,6 +1208,8 @@ } DFSanUnimplementedFn = Mod->getOrInsertFunction("__dfsan_unimplemented", DFSanUnimplementedFnTy); + DFSanWrapperExternWeakNullFn = Mod->getOrInsertFunction( + "__dfsan_wrapper_extern_weak_null", DFSanWrapperExternWeakNullFnTy); { AttributeList AL; AL = AL.addParamAttribute(M.getContext(), 0, Attribute::ZExt); @@ -1227,6 +1253,8 @@ DFSanLoadLabelAndOriginFn.getCallee()->stripPointerCasts()); DFSanRuntimeFunctions.insert( DFSanUnimplementedFn.getCallee()->stripPointerCasts()); + DFSanRuntimeFunctions.insert( + DFSanWrapperExternWeakNullFn.getCallee()->stripPointerCasts()); DFSanRuntimeFunctions.insert( DFSanSetLabelFn.getCallee()->stripPointerCasts()); DFSanRuntimeFunctions.insert( @@ -2860,16 +2888,19 @@ CB.setCalledFunction(&F); IRB.CreateCall(DFSF.DFS.DFSanUnimplementedFn, IRB.CreateGlobalStringPtr(F.getName())); + DFSF.DFS.buildExternWeakCheckIfNeeded(IRB, &F); DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB)); DFSF.setOrigin(&CB, DFSF.DFS.ZeroOrigin); return true; case DataFlowSanitizer::WK_Discard: CB.setCalledFunction(&F); + DFSF.DFS.buildExternWeakCheckIfNeeded(IRB, &F); DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB)); DFSF.setOrigin(&CB, DFSF.DFS.ZeroOrigin); return true; case DataFlowSanitizer::WK_Functional: CB.setCalledFunction(&F); + DFSF.DFS.buildExternWeakCheckIfNeeded(IRB, &F); visitInstOperands(CB); return true; case DataFlowSanitizer::WK_Custom: diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll b/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll --- a/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll +++ b/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll @@ -11,12 +11,15 @@ define noundef i8 @call_if_exists() local_unnamed_addr { ; CHECK-LABEL: @call_if_exists.dfsan + ; Ensure comparison is preserved ; CHECK: br i1 icmp ne ([[FUNCPTRTY:.*]] @ExternWeak, [[FUNCPTRTY]] null), label %use_func, label %avoid_func br i1 icmp ne (i8 (i8)* @ExternWeak, i8 (i8)* null), label %use_func, label %avoid_func use_func: ; CHECK: use_func: - ; CHECK: call i8 @ExternWeak(i8 {{.*}}) + ; Ensure extern weak function is validated before being called. + ; CHECK: call void @__dfsan_wrapper_extern_weak_null({{[^,]*}}@ExternWeak{{[^,]*}}, {{.*}}) + ; CHECK-NEXT: call i8 @ExternWeak(i8 {{.*}}) %1 = call i8 @ExternWeak(i8 4) br label %end