[DFSan] Avoid replacing uses of functions in comparisions.

[browneee]

Summary

This can cause crashes by accidentally optimizing out checks for
extern_weak_func != nullptr, when replaced with a known-not-null wrapper.

This solution isn't perfect (only avoids replacement on specific patterns)
but should address common cases.

Internal reference: b/185245029

[browneee]

Rebase

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

if (isInstrumented(&F)) {
           &F,
           (shouldTrackOrigins() ? std::string("dfso$") : std::string("dfsw$")) +
               std::string(F.getName()),
           WrapperLinkage, FT);
       NewF->removeFnAttrs(ReadOnlyNoneAttrs);
 
       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<ConstantExpr>(Usr)) {
+          // This is the most common case for icmp ne null
+          if (CE->getOpcode() == Instruction::ICmp) {
+            return false;
+          }
+        }
+        if (Instruction *I = dyn_cast<Instruction>(Usr)) {
+          if (I->getOpcode() == Instruction::ICmp) {
+            return false;
+          }
+        }
+        return true;
+      };
+      F.replaceUsesWithIf(WrappedFnCst, IsNotCmpUse);
 
       UnwrappedFnMap[WrappedFnCst] = &F;
       *FI = NewF;
 
       if (!F.isDeclaration()) {
         // This function is probably defining an interposition of an
         // uninstrumented function and hence needs to keep the original ABI.
         // But any functions it may call need to use the instrumented ABI, so

llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt

 fun:discard*=uninstrumented
 fun:discard*=discard
 
 fun:functional=uninstrumented
 fun:functional=functional
 
 fun:custom*=uninstrumented
 fun:custom*=custom
 
 fun:uninstrumented*=uninstrumented
 
 fun:function_to_force_zero=force_zero_labels
+
+fun:ExternWeak=uninstrumented

llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll

+; 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
+}
+