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 @@ -1428,7 +1428,40 @@ Value *WrappedFnCst = ConstantExpr::getBitCast(NewF, PointerType::getUnqual(FT)); - F.replaceAllUsesWith(WrappedFnCst); + + // Extern weak functions can sometimes be null at execution time. + // Code will sometimes check if an extern weak function is null. + // This could look something like: + // declare extern_weak i8 @my_func(i8) + // br i1 icmp ne (i8 (i8)* @my_func, i8 (i8)* null), label %use_my_func, + // label %avoid_my_func + // The @"dfsw$my_func" wrapper is never null, so if we replace this use + // in the comparision, the icmp will simplify to false and we have + // accidentially optimized away a null check that is necessary. + // This can lead to a crash when the null extern_weak my_func is called. + // + // To prevent (the most common pattern of) this problem, + // do not replace uses in comparisons with the wrapper. + // We definitely want to replace uses in call instructions. + // Other uses (e.g. store the function address somewhere) might be + // called or compared or both - this case may not be handled correctly. + // We will default to replacing with wrapper in cases we are unsure. + auto IsNotCmpUse = [](Use &U) -> bool { + User *Usr = U.getUser(); + if (ConstantExpr *CE = dyn_cast(Usr)) { + // This is the most common case for icmp ne null + if (CE->getOpcode() == Instruction::ICmp) { + return false; + } + } + if (Instruction *I = dyn_cast(Usr)) { + if (I->getOpcode() == Instruction::ICmp) { + return false; + } + } + return true; + }; + F.replaceUsesWithIf(WrappedFnCst, IsNotCmpUse); UnwrappedFnMap[WrappedFnCst] = &F; *FI = NewF; diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt b/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt --- a/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt +++ b/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt @@ -10,3 +10,5 @@ fun:uninstrumented*=uninstrumented fun:function_to_force_zero=force_zero_labels + +fun:ExternWeak=uninstrumented diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll b/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll @@ -0,0 +1,32 @@ +; Check that we don't replace uses in cmp with wrapper (which would accidentally optimize out the cmp). +; RUN: opt < %s -dfsan -dfsan-abilist=%S/Inputs/abilist.txt -S | 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" + +; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]] +; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]] + +; CHECK: declare extern_weak i8 @ExternWeak(i8) +declare extern_weak i8 @ExternWeak(i8) + +define noundef i8 @call_if_exists() local_unnamed_addr { + ; CHECK-LABEL: @call_if_exists.dfsan + ; 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 {{.*}}) + %1 = call i8 @ExternWeak(i8 4) + br label %end + +avoid_func: + ; CHECK: avoid_func: + br label %end + +end: + ; CHECK: end: + %r = phi i8 [ %1, %use_func ], [ 0, %avoid_func ] + ret i8 %r +} +