[llvm] cddcf21 - [DFSan] Avoid replacing uses of functions in comparisions.

[Andrew Browne via llvm-commits]

Author: Andrew Browne
Date: 2022-04-14T14:14:52-07:00
New Revision: cddcf2170ae8d5a199bb99aac2fd27f520696efe

URL: https://github.com/llvm/llvm-project/commit/cddcf2170ae8d5a199bb99aac2fd27f520696efe
DIFF: https://github.com/llvm/llvm-project/commit/cddcf2170ae8d5a199bb99aac2fd27f520696efe.diff

LOG: [DFSan] Avoid replacing uses of functions in comparisions.

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

Reviewed By: vitalybuka

Differential Revision: https://reviews.llvm.org/D123701

Added: 
    llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll

Modified: 
    llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
    llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index 5b9e778363870..fe81ce0fc4efa 100644
--- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -1428,7 +1428,40 @@ bool DataFlowSanitizer::runImpl(Module &M) {
 
       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;

diff  --git a/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt b/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
index 712c9dbceea7e..ebd64b5f78c74 100644
--- a/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
+++ b/llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
@@ -10,3 +10,5 @@ fun:custom*=custom
 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
index 0000000000000..00f33687b745d
--- /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
+}
+