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

[Andrew via Phabricator via llvm-commits]

browneee updated this revision to Diff 422925.
browneee added a comment.

Rebase


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D123701/new/

https://reviews.llvm.org/D123701

Files:
  llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
  llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
  llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll


Index: llvm/test/Instrumentation/DataFlowSanitizer/extern_weak.ll
===================================================================
--- /dev/null
+++ 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
+}
+
Index: llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
===================================================================
--- llvm/test/Instrumentation/DataFlowSanitizer/Inputs/abilist.txt
+++ 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
Index: llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
===================================================================
--- llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ 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<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;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D123701.422925.patch
Type: text/x-patch
Size: 3997 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220414/2291ebaa/attachment.bin>