[AddressSanitizer] Remove memory effects from functions

If left as-is, subsequent optimizations might utilize the possible
memory effects and optimize-out the instrumentation. Think of the
following case:
```
  store i8 4, ptr %shadow
  call void @llvm.lifetime.start.p0(i64 4, ptr %local)
  %28 = call void @foo(ptr %local)
  store i8 -8, ptr %shadow
  call void @llvm.lifetime.end.p0(i64 4, ptr %local)
```

where `foo` is an external function with `memory(argmem: write)`. A pass
such as DeadStoreElimination is allowed to remove the initial store,
which might fail sanitizer checks within `foo`.

Author: guy-david

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

@@ -58,6 +58,13 @@ void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
                                bool IsKasan, uint64_t *ShadowBase,
                                int *MappingScale, bool *OrShadowOffset);
 
+/// Remove memory attributes that are incompatible with the instrumentation
+/// added by AddressSanitizer and HWSanitizer. See more details in the
+/// implementation.
+/// \p ReadsArgMem - whether to remove memory attributes from function
+/// arguments as well.
+void removeASanIncompatibleFnAttributes(Function &F, bool ReadsArgMem);
+
 } // namespace llvm
 
 #endif

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -612,6 +612,48 @@ void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
   *OrShadowOffset = Mapping.OrShadowOffset;
 }
 
+void removeASanIncompatibleFnAttributes(Function &F, bool ReadsArgMem) {
+  // Remove memory attributes that are invalid with ASan and HWSan.
+  // ASan checks read from shadow, which invalidates memory(argmem: *)
+  // Short granule checks on function arguments read from the argument memory
+  // (last byte of the granule), which invalidates writeonly.
+  //
+  // This is not only true for sanitized functions, because AttrInfer can
+  // infer those attributes on libc functions, which is not true if those
+  // are instrumented (Android) or intercepted.
+  //
+  // We might want to model ASan shadow memory more opaquely to get rid of
+  // this problem altogether, by hiding the shadow memory write in an
+  // intrinsic, essentially like in the AArch64StackTagging pass. But that's
+  // for another day.
+
+  // The API is weird. `onlyReadsMemory` actually means "does not write", and
+  // `onlyWritesMemory` actually means "does not read". So we reconstruct
+  // "accesses memory" && "does not read" <=> "writes".
+  bool Changed = false;
+  if (!F.doesNotAccessMemory()) {
+    bool WritesMemory = !F.onlyReadsMemory();
+    bool ReadsMemory = !F.onlyWritesMemory();
+    if ((WritesMemory && !ReadsMemory) || F.onlyAccessesArgMemory()) {
+      F.removeFnAttr(Attribute::Memory);
+      Changed = true;
+    }
+  }
+  if (ReadsArgMem) {
+    for (Argument &A : F.args()) {
+      if (A.hasAttribute(Attribute::WriteOnly)) {
+        A.removeAttr(Attribute::WriteOnly);
+        Changed = true;
+      }
+    }
+  }
+  if (Changed) {
+    // nobuiltin makes sure later passes don't restore assumptions about
+    // the function.
+    F.addFnAttr(Attribute::NoBuiltin);
+  }
+}
+
 ASanAccessInfo::ASanAccessInfo(int32_t Packed)
     : Packed(Packed),
       AccessSizeIndex((Packed >> kAccessSizeIndexShift) & kAccessSizeIndexMask),
@@ -2733,6 +2775,9 @@ GlobalVariable *ModuleAddressSanitizer::getOrCreateModuleName() {
 bool ModuleAddressSanitizer::instrumentModule() {
   initializeCallbacks();
 
+  for (Function &F : M)
+    removeASanIncompatibleFnAttributes(F, /*ReadsArgMem=*/false);
+
   // Create a module constructor. A destructor is created lazily because not all
   // platforms, and not all modules need it.
   if (ConstructorKind == AsanCtorKind::Global) {

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -325,7 +325,6 @@ class HWAddressSanitizer {
                                           FunctionAnalysisManager &FAM) const;
   void initializeModule();
   void createHwasanCtorComdat();
-  void removeFnAttributes(Function *F);
 
   void initializeCallbacks(Module &M);
 
@@ -620,56 +619,13 @@ void HWAddressSanitizer::createHwasanCtorComdat() {
   appendToCompilerUsed(M, Dummy);
 }
 
-void HWAddressSanitizer::removeFnAttributes(Function *F) {
-  // Remove memory attributes that are invalid with HWASan.
-  // HWASan checks read from shadow, which invalidates memory(argmem: *)
-  // Short granule checks on function arguments read from the argument memory
-  // (last byte of the granule), which invalidates writeonly.
-  //
-  // This is not only true for sanitized functions, because AttrInfer can
-  // infer those attributes on libc functions, which is not true if those
-  // are instrumented (Android) or intercepted.
-  //
-  // We might want to model HWASan shadow memory more opaquely to get rid of
-  // this problem altogether, by hiding the shadow memory write in an
-  // intrinsic, essentially like in the AArch64StackTagging pass. But that's
-  // for another day.
-
-  // The API is weird. `onlyReadsMemory` actually means "does not write", and
-  // `onlyWritesMemory` actually means "does not read". So we reconstruct
-  // "accesses memory" && "does not read" <=> "writes".
-  bool Changed = false;
-  if (!F->doesNotAccessMemory()) {
-    bool WritesMemory = !F->onlyReadsMemory();
-    bool ReadsMemory = !F->onlyWritesMemory();
-    if ((WritesMemory && !ReadsMemory) || F->onlyAccessesArgMemory()) {
-      F->removeFnAttr(Attribute::Memory);
-      Changed = true;
-    }
-  }
-  for (Argument &A : F->args()) {
-    if (A.hasAttribute(Attribute::WriteOnly)) {
-      A.removeAttr(Attribute::WriteOnly);
-      Changed = true;
-    }
-  }
-  if (Changed) {
-    // nobuiltin makes sure later passes don't restore assumptions about
-    // the function.
-    F->addFnAttr(Attribute::NoBuiltin);
-  }
-}
-
 /// Module-level initialization.
 ///
 /// inserts a call to __hwasan_init to the module's constructor list.
 void HWAddressSanitizer::initializeModule() {
   LLVM_DEBUG(dbgs() << "Init " << M.getName() << "\n");
   TargetTriple = M.getTargetTriple();
 
-  for (Function &F : M.functions())
-    removeFnAttributes(&F);
-
   // x86_64 currently has two modes:
   // - Intel LAM (default)
   // - pointer aliasing (heap only)
@@ -718,6 +674,9 @@ void HWAddressSanitizer::initializeModule() {
   InstrumentGlobals =
       !CompileKernel && !UsePageAliases && optOr(ClGlobals, NewRuntime);
 
+  for (Function &F : M.functions())
+    removeASanIncompatibleFnAttributes(F, /*ReadsArgMem=*/UseShortGranules);
+
   if (!CompileKernel) {
     createHwasanCtorComdat();
 

llvm/test/Instrumentation/AddressSanitizer/AMDGPU/asan_instrument_mem_intrinsics.ll

@@ -186,7 +186,7 @@ declare void @llvm.memset.p5.i32(ptr addrspace(5) nocapture writeonly, i8, i32,
 
 define void @test_mem_intrinsic_memcpy(ptr %dest0,ptr %src0,ptr addrspace(1) %dest1,ptr addrspace(1) %src1,ptr addrspace(2) %dest2,ptr addrspace(2) %src2,ptr addrspace(3) %dest3,ptr addrspace(3) %src3,ptr addrspace(4) %dest4,ptr addrspace(4) %src4,ptr addrspace(5) %dest5,ptr addrspace(5) %src5) #0 {
 entry:
-  ;CHECK: define void @test_mem_intrinsic_memcpy(ptr [[DEST0:%.*]], ptr [[SRC0:%.*]], ptr addrspace(1) [[DEST1:%.*]], ptr addrspace(1) [[SRC1:%.*]], ptr addrspace(2) [[DEST2:%.*]], ptr addrspace(2) [[SRC2:%.*]], ptr addrspace(3) [[DEST3:%.*]], ptr addrspace(3) [[SRC3:%.*]], ptr addrspace(4) [[DEST4:%.*]], ptr addrspace(4) [[SRC4:%.*]], ptr addrspace(5) [[DEST5:%.*]], ptr addrspace(5) [[SRC5:%.*]]) #2 {
+  ;CHECK: define void @test_mem_intrinsic_memcpy(ptr [[DEST0:%.*]], ptr [[SRC0:%.*]], ptr addrspace(1) [[DEST1:%.*]], ptr addrspace(1) [[SRC1:%.*]], ptr addrspace(2) [[DEST2:%.*]], ptr addrspace(2) [[SRC2:%.*]], ptr addrspace(3) [[DEST3:%.*]], ptr addrspace(3) [[SRC3:%.*]], ptr addrspace(4) [[DEST4:%.*]], ptr addrspace(4) [[SRC4:%.*]], ptr addrspace(5) [[DEST5:%.*]], ptr addrspace(5) [[SRC5:%.*]]){{.*}} {
   ;CHECK-NEXT: entry:
   ;CHECK-NEXT: [[VR0:%.*]] = call ptr @__asan_memcpy(ptr [[DEST0]], ptr [[SRC0]], i64 64)
   ;CHECK-NEXT: [[VR1:%.*]] = addrspacecast ptr addrspace(1) [[SRC1]] to ptr
@@ -469,7 +469,7 @@ entry:
 
 define void @test_mem_intrinsic_memmove(ptr %dest0,ptr %src0,ptr addrspace(1) %dest1,ptr addrspace(1) %src1,ptr addrspace(2) %dest2,ptr addrspace(2) %src2,ptr addrspace(3) %dest3,ptr addrspace(3) %src3,ptr addrspace(4) %dest4,ptr addrspace(4) %src4,ptr addrspace(5) %dest5,ptr addrspace(5) %src5) #0 {
 entry:
-  ;CHECK: define void @test_mem_intrinsic_memmove(ptr [[DEST0:%.*]], ptr [[SRC0:%.*]], ptr addrspace(1) [[DEST1:%.*]], ptr addrspace(1) [[SRC1:%.*]], ptr addrspace(2) [[DEST2:%.*]], ptr addrspace(2) [[SRC2:%.*]], ptr addrspace(3) [[DEST3:%.*]], ptr addrspace(3) [[SRC3:%.*]], ptr addrspace(4) [[DEST4:%.*]], ptr addrspace(4) [[SRC4:%.*]], ptr addrspace(5) [[DEST5:%.*]], ptr addrspace(5) [[SRC5:%.*]]) #2 {
+  ;CHECK: define void @test_mem_intrinsic_memmove(ptr [[DEST0:%.*]], ptr [[SRC0:%.*]], ptr addrspace(1) [[DEST1:%.*]], ptr addrspace(1) [[SRC1:%.*]], ptr addrspace(2) [[DEST2:%.*]], ptr addrspace(2) [[SRC2:%.*]], ptr addrspace(3) [[DEST3:%.*]], ptr addrspace(3) [[SRC3:%.*]], ptr addrspace(4) [[DEST4:%.*]], ptr addrspace(4) [[SRC4:%.*]], ptr addrspace(5) [[DEST5:%.*]], ptr addrspace(5) [[SRC5:%.*]]){{.*}} {
   ;CHECK-NEXT: entry:
   ;CHECK-NEXT: [[VR0:%.*]] = call ptr @__asan_memmove(ptr [[DEST0]], ptr [[SRC0]], i64 64)
   ;CHECK-NEXT: [[VR1:%.*]] = addrspacecast ptr addrspace(1) [[SRC1]] to ptr
@@ -753,7 +753,7 @@ entry:
 
 define void @test_mem_intrinsic_memset(ptr %ptr0,ptr addrspace(1) %ptr1,ptr addrspace(2) %ptr2,ptr addrspace(3) %ptr3,ptr addrspace(4) %ptr4,ptr addrspace(5) %ptr5) #0{
 entry:
-  ;CHECK: define void @test_mem_intrinsic_memset(ptr [[PTR0:%.*]], ptr addrspace(1) [[PTR1:%.*]], ptr addrspace(2) [[PTR2:%.*]], ptr addrspace(3) [[PTR3:%.*]], ptr addrspace(4) [[PTR4:%.*]], ptr addrspace(5) [[PTR5:%.*]]) #2 {
+  ;CHECK: define void @test_mem_intrinsic_memset(ptr [[PTR0:%.*]], ptr addrspace(1) [[PTR1:%.*]], ptr addrspace(2) [[PTR2:%.*]], ptr addrspace(3) [[PTR3:%.*]], ptr addrspace(4) [[PTR4:%.*]], ptr addrspace(5) [[PTR5:%.*]]){{.*}} {
   ;CHECK-NEXT: entry:
   ;CHECK-NEXT: [[VR0:%.*]] = call ptr @__asan_memset(ptr [[PTR0]], i32 1, i64 128)
   ;CHECK-NEXT: [[VR1:%.*]] = addrspacecast ptr addrspace(1) [[PTR1]] to ptr

llvm/test/Instrumentation/AddressSanitizer/remove-memory-effects.ll

@@ -0,0 +1,23 @@
+; Remove possible memory effects from functions that are invalidated by
+; AddressSanitizer instrumentation.
+
+; RUN: opt -passes='asan<use-after-scope>' -S %s | FileCheck %s
+
+target datalayout = "e-i64:64-f80:128-s:64-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: @foo(ptr writeonly) #[[ATTRS_FOO:[0-9]+]]
+declare void @foo(ptr writeonly) memory(argmem: write)
+
+; CHECK: @bar() #[[ATTRS_BAR:[0-9]+]]
+define void @bar() sanitize_address {
+entry:
+  %x = alloca i32, align 4
+  call void @llvm.lifetime.start.p0(i64 4, ptr %x)
+  call void @foo(ptr %x)
+  call void @llvm.lifetime.end.p0(i64 4, ptr %x)
+  ret void
+}
+
+; CHECK: attributes #[[ATTRS_FOO]] = { nobuiltin }
+; CHECK: attributes #[[ATTRS_BAR]] = { sanitize_address }