sanmd: improve precision of UAR analysis

[dvyukov]

Summary

Only mark functions that have address-taken locals
as requiring UAR checking.

On a large internal app this reduces number of marked functions
from 78441 to 66618. Mostly small, trivial getter/setter-type
functions are unmarked, but also some amount of larger
number-crunching-type functions are unmarked as well.

[dvyukov]

llvm::isAllocaPromotable

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

274

I was looking at isSafeAndProfitableToSinkLoad:
https://github.com/llvm/llvm-project/blob/881076cde29699ef2a458c9d957ebcd49ded8893/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp#L621

and isAllocaPromotable:
https://github.com/llvm/llvm-project/blob/af4e856fa71c6b5086aeda79bfcd954e98cef591/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp#L63

[dvyukov]

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

274

Btw, isSafeAndProfitableToSinkLoad is broken. All allocas now sink to lifetime intrinsics. So it will conclude all allocas are address-taken. This immediately popped up in our end-to-end C++ tests, but wasn't detected by the InstCombine bitcode tests.

[melver]

compiler-rt/test/metadata/uar.cpp

94

Add test for explicit __builtin_alloca ?

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

274

All this is already in a big anonymous namespace, so static not needed.

274

Could file bug (https://github.com/llvm/llvm-project/issues) and point out that it needs the isLifetimeStartOrEnd() check, but we don't quite know if it's intended?

280–293

LoadInst, StoreInst, GetElementPtrInst and BitCastInst all are derived from Instruction, so these checks can be moved in the 'if (... dyn_cast<Instruction>' branch.

[dvyukov]

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

280–293

Can users be a non-instruction?
I assumed that I need a cast to Instruction just to get access to its methods.

[melver]

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

280–293

Yes, all I'm suggesting is a small optimization, because the additional casts to *Inst will fail if the cast to Instruction above failed.

[dvyukov]

addressed the comments

[dvyukov]

PTAL

compiler-rt/test/metadata/covered.cpp

 // RUN: %clangxx %s -o %t && %t | FileCheck %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=covered && %t | FileCheck -check-prefix=CHECK-C %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=atomics && %t | FileCheck -check-prefix=CHECK-A %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=uar && %t | FileCheck -check-prefix=CHECK-U %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=covered,atomics && %t | FileCheck -check-prefix=CHECK-CA %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=covered,uar && %t | FileCheck -check-prefix=CHECK-CU %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=atomics,uar && %t | FileCheck -check-prefix=CHECK-AU %s
 // RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=covered,atomics,uar && %t | FileCheck -check-prefix=CHECK-CAU %s
 
+__attribute__((noinline, not_tail_called)) void escape(const volatile void *p) {
+  static const volatile void *sink;
+  sink = p;
+}
+
 // CHECK-NOT: metadata add
 // CHECK: main
 // CHECK-NOT: metadata del
 
 // CHECK-C:      empty: features=0
 // CHECK-A-NOT:  empty:
 // CHECK-U-NOT:  empty:
 // CHECK-CA:     empty: features=1
 // CHECK-CU:     empty: features=0
 // CHECK-AU-NOT: empty:
 // CHECK-CAU:    empty: features=1
 void empty() {}
 
 // CHECK-C:  normal: features=0
 // CHECK-A:  normal: features=1
 // CHECK-U:  normal: features=2
 // CHECK-CA: normal: features=1
 // CHECK-CU: normal: features=2
 // CHECK-AU: normal: features=3
 // CHECK-CAU:normal: features=3
 void normal() {
-  volatile int x;
-  x = 0;
+  int x;
+  escape(&x);
 }
 
-// CHECK-C:   with_atomic: features=0
-// CHECK-A:   with_atomic: features=1
-// CHECK-U:   with_atomic: features=2
-// CHECK-CA:  with_atomic: features=1
-// CHECK-CU:  with_atomic: features=2
-// CHECK-AU:  with_atomic: features=3
-// CHECK-CAU: with_atomic: features=3
+// CHECK-C:     with_atomic: features=0
+// CHECK-A:     with_atomic: features=1
+// CHECK-U-NOT: with_atomic:
+// CHECK-CA:    with_atomic: features=1
+// CHECK-CU:    with_atomic: features=0
+// CHECK-AU:    with_atomic: features=1
+// CHECK-CAU:   with_atomic: features=1
 int with_atomic(int *p) { return __atomic_load_n(p, __ATOMIC_RELAXED); }
 
+// CHECK-C:   with_atomic_escape: features=0
+// CHECK-A:   with_atomic_escape: features=1
+// CHECK-U:   with_atomic_escape: features=2
+// CHECK-CA:  with_atomic_escape: features=1
+// CHECK-CU:  with_atomic_escape: features=2
+// CHECK-AU:  with_atomic_escape: features=3
+// CHECK-CAU: with_atomic_escape: features=3
+int with_atomic_escape(int *p) {
+  escape(&p);
+  return __atomic_load_n(p, __ATOMIC_RELAXED);
+}
+
 // CHECK-C:     ellipsis: features=0
 // CHECK-A:     ellipsis: features=1
 // CHECK-U-NOT: ellipsis:
 // CHECK-CA:    ellipsis: features=1
 // CHECK-CU:    ellipsis: features=0
 // CHECK-AU:    ellipsis: features=1
 // CHECK-CAU:   ellipsis: features=1
 void ellipsis(int *p, ...) {
+  escape(&p);
   volatile int x;
   x = 0;
 }
 
 // CHECK-C:     ellipsis_with_atomic: features=0
 // CHECK-A:     ellipsis_with_atomic: features=1
 // CHECK-U-NOT: ellipsis_with_atomic:
 // CHECK-CA:    ellipsis_with_atomic: features=1
 // CHECK-CU:    ellipsis_with_atomic: features=0
 // CHECK-AU:    ellipsis_with_atomic: features=1
 // CHECK-CAU:   ellipsis_with_atomic: features=1
 int ellipsis_with_atomic(int *p, ...) {
+  escape(&p);
   return __atomic_load_n(p, __ATOMIC_RELAXED);
 }
 
 #define FUNCTIONS                                                              \
   FN(empty);                                                                   \
   FN(normal);                                                                  \
   FN(with_atomic);                                                             \
+  FN(with_atomic_escape);                                                      \
   FN(ellipsis);                                                                \
   FN(ellipsis_with_atomic);                                                    \
   /**/
 
 #include "common.h"

compiler-rt/test/metadata/uar.cpp

-// RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=covered,uar && %t | FileCheck %s
+// RUN: %clangxx %s -O1 -o %t -fexperimental-sanitize-metadata=covered,uar && %t | FileCheck %s
 
 // CHECK: metadata add version 1
 
+__attribute__((noinline, not_tail_called)) void escape(const volatile void *p) {
+  static const volatile void *sink;
+  sink = p;
+}
+
+__attribute__((noinline, not_tail_called)) void use(int x) {
+  static volatile int sink;
+  sink += x;
+}
+
 // CHECK: empty: features=0 stack_args=0
 void empty() {}
 
 // CHECK: ellipsis: features=0 stack_args=0
 void ellipsis(const char *fmt, ...) {
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
 }
 
 // CHECK: non_empty_function: features=2 stack_args=0
 void non_empty_function() {
-  // Completely empty functions don't get uar metadata.
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
 }
 
 // CHECK: no_stack_args: features=2 stack_args=0
 void no_stack_args(long a0, long a1, long a2, long a3, long a4, long a5) {
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
 }
 
 // CHECK: stack_args: features=2 stack_args=16
 void stack_args(long a0, long a1, long a2, long a3, long a4, long a5, long a6) {
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
 }
 
 // CHECK: more_stack_args: features=2 stack_args=32
 void more_stack_args(long a0, long a1, long a2, long a3, long a4, long a5,
                      long a6, long a7, long a8) {
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
 }
 
 // CHECK: struct_stack_args: features=2 stack_args=144
 struct large {
   char x[131];
 };
 void struct_stack_args(large a) {
-  volatile int x;
-  x = 1;
+  int x;
+  escape(&x);
+}
+
+__attribute__((noinline)) int tail_called(int x) { return x; }
+
+// CHECK: with_tail_call: features=2
+int with_tail_call(int x) { [[clang::musttail]] return tail_called(x); }
+
+// CHECK: local_array: features=0
+void local_array(int x) {
+  int data[10];
+  use(data[x]);
+}
+
+// CHECK: local_alloca: features=0
+void local_alloca(int size, int i, int j) {
+  volatile int *p = static_cast<int *>(__builtin_alloca(size));
+  p[i] = 0;
+  use(p[j]);
+}
+
+// CHECK: escaping_alloca: features=2
+void escaping_alloca(int size, int i) {
+  volatile int *p = static_cast<int *>(__builtin_alloca(size));
+  escape(&p[i]);
 }
 
 #define FUNCTIONS                                                              \
   FN(empty);                                                                   \
   FN(ellipsis);                                                                \
   FN(non_empty_function);                                                      \
   FN(no_stack_args);                                                           \
   FN(stack_args);                                                              \
   FN(more_stack_args);                                                         \
   FN(struct_stack_args);                                                       \
+  FN(with_tail_call);                                                          \
+  FN(local_array);                                                             \
+  FN(local_alloca);                                                            \
+  FN(escaping_alloca);                                                         \
   /**/
 
 #include "common.h"

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

void SanitizerBinaryMetadata::runOn(Function &F, MetadataInfoSet &MIS) {
   if (FeatureMask || Options.UAR) {
     for (BasicBlock &BB : F)
       for (Instruction &I : BB)
         RequiresCovered |= runOn(I, MIS, MDB, FeatureMask);
   }
 
   if (F.isVarArg())
     FeatureMask &= ~kSanitizerBinaryMetadataUAR;
-  if (FeatureMask & kSanitizerBinaryMetadataUAR)
+  if (FeatureMask & kSanitizerBinaryMetadataUAR) {
+    RequiresCovered = true;
     NumMetadataUAR++;
+  }
 
   // Covered metadata is always emitted if explicitly requested, otherwise only
   // if some other metadata requires it to unambiguously interpret it for
   // modules compiled with SanitizerBinaryMetadata.
   if (Options.Covered || (FeatureMask && RequiresCovered)) {
     NumMetadataCovered++;
     const auto *MI = &MetadataInfo::Covered;
     MIS.insert(MI);
     const StringRef Section = getSectionName(MI->SectionSuffix);
     // The feature mask will be placed after the size (32 bit) of the function,
     // so in total one covered entry will use `sizeof(void*) + 4 + 4`.
     Constant *CFM = IRB.getInt32(FeatureMask);
     F.setMetadata(LLVMContext::MD_pcsections,
                   MDB.createPCSections({{Section, {CFM}}}));
   }
 }
 
+bool hasUseAfterReturnUnsafeUses(Value &V) {
+  for (User *U : V.users()) {
+    if (auto *I = dyn_cast<Instruction>(U)) {
+      if (I->isLifetimeStartOrEnd() || I->isDroppable())
+        continue;
+      if (isa<LoadInst>(U))
+        continue;
+      if (auto *SI = dyn_cast<StoreInst>(U)) {
+        // If storing TO the alloca, then the address isn't taken.
+        if (SI->getOperand(1) == &V)
+          continue;
+      }
+      if (auto *GEPI = dyn_cast<GetElementPtrInst>(U)) {
+        if (!hasUseAfterReturnUnsafeUses(*GEPI))
+          continue;
+      } else if (auto *BCI = dyn_cast<BitCastInst>(U)) {
+        if (!hasUseAfterReturnUnsafeUses(*BCI))
+          continue;
+      }
+    }
+    return true;
+  }
+  return false;
+}
+
+bool useAfterReturnUnsafe(Instruction &I) {
+  if (isa<AllocaInst>(I))
+    return hasUseAfterReturnUnsafeUses(I);
+  // Tail-called functions are not necessary intercepted
+  // at runtime because there is no call instruction.
+  // So conservatively mark the caller as requiring checking.
+  else if (auto *CI = dyn_cast<CallInst>(&I))
+    return CI->isTailCall();
+  return false;
+}
+
 bool SanitizerBinaryMetadata::runOn(Instruction &I, MetadataInfoSet &MIS,
                                     MDBuilder &MDB, uint32_t &FeatureMask) {
   SmallVector<const MetadataInfo *, 1> InstMetadata;
   bool RequiresCovered = false;
 
-  if (Options.UAR) {
-    for (unsigned i = 0; i < I.getNumOperands(); ++i) {
-      const Value *V = I.getOperand(i);
-      // TODO(dvyukov): check if V is an address of alloca/function arg.
-      // See isSafeAndProfitableToSinkLoad for addr-taken allocas
-      // and DeadArgumentEliminationPass::removeDeadStuffFromFunction
-      // for iteration over function args.
-      if (V) {
-        RequiresCovered = true;
-        FeatureMask |= kSanitizerBinaryMetadataUAR;
-      }
-    }
-  }
+  if (Options.UAR && !(FeatureMask & kSanitizerBinaryMetadataUAR)) {
+    if (useAfterReturnUnsafe(I))
 
   if (Options.Atomics && I.mayReadOrWriteMemory()) {
     auto SSID = getAtomicSyncScopeID(&I);
     if (SSID.has_value() && SSID.value() != SyncScope::SingleThread) {
       NumMetadataAtomics++;
       InstMetadata.push_back(&MetadataInfo::Atomics);
     }
     RequiresCovered = true;