Support atomic write operations in stack safety

[fmayer]

Summary

This has two benefits:

• we can now mark allocas that are used in atomic operations as safe
• this fixes a bug that would incorrectly mark all atomic writes as safe in HWASan instrumentation. this is because stack safety keeps a list of all *unsafe* operations that are reachable from an alloca, but it did not analyze atomic writes, so it would always mark them as safe.

[fmayer]

tests

[fmayer]

tests

[vitalybuka]

llvm/lib/Analysis/StackSafetyAnalysis.cpp

412

maybe instead of introducing getNewValOperandIndex
just change RecordWrite(Value* operand) ?

[fmayer]

simplify

[fmayer]

upd

llvm/lib/Analysis/StackSafetyAnalysis.cpp

bool StackSafetyLocalAnalysis::isSafeAccess(const Use &U, AllocaInst *AI,
   const SCEV *SV = SE.getConstant(CalculationTy, TS.getFixedValue());
   return isSafeAccess(U, AI, SV);
 }
 
 bool StackSafetyLocalAnalysis::isSafeAccess(const Use &U, AllocaInst *AI,
                                             const SCEV *AccessSize) {
 
   if (!AI)
-    return true;
+    return true; // This only judges whether it is a safe *stack* access.
   if (isa<SCEVCouldNotCompute>(AccessSize))
     return false;
 
   const auto *I = cast<Instruction>(U.getUser());
 
   auto ToCharPtr = [&](const SCEV *V) {
     auto *PtrTy = IntegerType::getInt8PtrTy(SE.getContext());
     return SE.getTruncateOrZeroExtend(V, PtrTy);
while (!WorkList.empty()) {
     const Value *V = WorkList.pop_back_val();
     for (const Use &UI : V->uses()) {
       const auto *I = cast<Instruction>(UI.getUser());
       if (!SL.isReachable(I))
         continue;
 
       assert(V == UI.get());
 
+      auto RecordStore = [&](const Value* StoredVal) {
+        if (V == StoredVal) {
+          // Stored the pointer - conservatively assume it may be unsafe.
+          US.addRange(I, UnknownRange, /*IsSafe=*/false);
+          return;
+        }
+        if (AI && !SL.isAliveAfter(AI, I)) {
+          US.addRange(I, UnknownRange, /*IsSafe=*/false);
+          return;
+        }
+        auto TypeSize = DL.getTypeStoreSize(StoredVal->getType());
+        auto AccessRange = getAccessRange(UI, Ptr, TypeSize);
+        bool Safe = isSafeAccess(UI, AI, TypeSize);
+        US.addRange(I, AccessRange, Safe);
+        return;
+      };
+
       switch (I->getOpcode()) {
       case Instruction::Load: {
         if (AI && !SL.isAliveAfter(AI, I)) {
           US.addRange(I, UnknownRange, /*IsSafe=*/false);
           break;
         }
         auto TypeSize = DL.getTypeStoreSize(I->getType());
         auto AccessRange = getAccessRange(UI, Ptr, TypeSize);
         bool Safe = isSafeAccess(UI, AI, TypeSize);
         US.addRange(I, AccessRange, Safe);
         break;
       }
 
       case Instruction::VAArg:
         // "va-arg" from a pointer is safe.
         break;
-      case Instruction::Store: {
-        if (V == I->getOperand(0)) {
-          // Stored the pointer - conservatively assume it may be unsafe.
-          US.addRange(I, UnknownRange, /*IsSafe=*/false);
-          break;
-        }
-        if (AI && !SL.isAliveAfter(AI, I)) {
-          US.addRange(I, UnknownRange, /*IsSafe=*/false);
-          break;
-        }
-        auto TypeSize = DL.getTypeStoreSize(I->getOperand(0)->getType());
-        auto AccessRange = getAccessRange(UI, Ptr, TypeSize);
-        bool Safe = isSafeAccess(UI, AI, TypeSize);
-        US.addRange(I, AccessRange, Safe);
+      case Instruction::Store:
+        RecordStore(cast<StoreInst>(I)->getValueOperand());
+      case Instruction::AtomicCmpXchg:
+        RecordStore(cast<AtomicCmpXchgInst>(I)->getNewValOperand());
+      case Instruction::AtomicRMW:
+        RecordStore(cast<AtomicRMWInst>(I)->getValOperand());
         break;
-      }
 
       case Instruction::Ret:
         // Information leak.
         // FIXME: Process parameters correctly. This is a leak only if we return
         // alloca.
         US.addRange(I, UnknownRange, /*IsSafe=*/false);
         break;
 
void StackSafetyGlobalInfo::print(raw_ostream &O) const {
   for (const auto &F : M.functions()) {
     if (!F.isDeclaration()) {
       SSI.find(&F)->second.print(O, F.getName(), &F);
       O << "    safe accesses:"
         << "\n";
       for (const auto &I : instructions(F)) {
         const CallInst *Call = dyn_cast<CallInst>(&I);
         if ((isa<StoreInst>(I) || isa<LoadInst>(I) || isa<MemIntrinsic>(I) ||
+             isa<AtomicCmpXchgInst>(I) || isa<AtomicRMWInst>(I) ||
              (Call && Call->hasByValArgument())) &&
             stackAccessIsSafe(I)) {
           O << "     " << I << "\n";
         }
       }
       O << "\n";
     }
   }

llvm/test/Analysis/StackSafetyAnalysis/local.ll

entry:
   %a = alloca i32, align 4
   call void @llvm.lifetime.start.p0(i64 4, ptr %a)
   call void @llvm.memset.p0.i32(ptr %a, i8 1, i32 4, i1 false)
   call void @llvm.lifetime.end.p0(i64 4, ptr %a)
   call void @unknown_call(ptr %a)
   ret void
 }
 
+define void @Cmpxchg4Arg(ptr %p) {
+; CHECK-LABEL: @Cmpxchg4Arg
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: p[]: [0,4){{$}}
+; CHECK-NEXT: allocas uses:
+; GLOBAL-NEXT: safe accesses:
+; GLOBAL-NEXT: cmpxchg ptr %p, i32 0, i32 1 monotonic monotonic, align 1
+; CHECK-EMPTY:
+entry:
+  cmpxchg ptr %p, i32 0, i32 1 monotonic monotonic, align 1
+  ret void
+}
+
+define void @AtomicRMW4Arg(ptr %p) {
+; CHECK-LABEL: @AtomicRMW4Arg
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: p[]: [0,4){{$}}
+; CHECK-NEXT: allocas uses:
+; GLOBAL-NEXT: safe accesses:
+; GLOBAL-NEXT: atomicrmw add ptr %p, i32 1 monotonic, align 1
+; CHECK-EMPTY:
+entry:
+  atomicrmw add ptr %p, i32 1 monotonic, align 1
+  ret void
+}
+
+define void @Cmpxchg4Alloca() {
+; CHECK-LABEL: @Cmpxchg4Alloca
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: allocas uses:
+; CHECK-NEXT: x[4]: [0,4){{$}}
+; GLOBAL-NEXT: safe accesses:
+; GLOBAL-NEXT: cmpxchg ptr %x, i32 0, i32 1 monotonic monotonic, align 1
+; CHECK-EMPTY:
+entry:
+  %x = alloca i32, align 4
+  cmpxchg ptr %x, i32 0, i32 1 monotonic monotonic, align 1
+  ret void
+}
+
+define void @AtomicRMW4Alloca() {
+; CHECK-LABEL: @AtomicRMW4Alloca
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: allocas uses:
+; CHECK-NEXT: x[4]: [0,4){{$}}
+; GLOBAL-NEXT: safe accesses:
+; GLOBAL-NEXT: atomicrmw add ptr %x, i32 1 monotonic, align 1
+; CHECK-EMPTY:
+entry:
+  %x = alloca i32, align 4
+  atomicrmw add ptr %x, i32 1 monotonic, align 1
+  ret void
+}
+
+define void @StoreArg(ptr %p) {
+; CHECK-LABEL: @StoreArg
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: p[]: [0,4){{$}}
+; CHECK-NEXT: allocas uses:
+; GLOBAL-NEXT: safe accesses:
+; GLOBAL-NEXT: store i32 1, ptr %p
+; CHECK-EMPTY:
+entry:
+  store i32 1, ptr %p
+  ret void
+}
+
 declare void @llvm.lifetime.start.p0(i64, ptr nocapture)
 declare void @llvm.lifetime.end.p0(i64, ptr nocapture)

llvm/test/Instrumentation/HWAddressSanitizer/stack-safety-analysis.ll

entry:
   ; NOSTACK-NOT: call {{.*}}__hwasan_store
   %buf.sroa.0 = alloca i8, align 4
   call void @llvm.lifetime.start.p0(i64 1, ptr nonnull %buf.sroa.0)
   store volatile i8 0, ptr %buf.sroa.0, align 4, !tbaa !8
   call void @llvm.lifetime.end.p0(i64 1, ptr nonnull %buf.sroa.0)
   ret i32 0
 }
 
+; Check a safe alloca to ensure it does not get a tag.
+define i32 @test_cmpxchg(ptr %a) sanitize_hwaddress {
+entry:
+  ; CHECK-LABEL: @test_cmpxchg
+  ; NOSAFETY: call {{.*}}__hwasan_generate_tag
+  ; NOSAFETY: call {{.*}}__hwasan_store
+  ; SAFETY-NOT: call {{.*}}__hwasan_generate_tag
+  ; SAFETY-NOT: call {{.*}}__hwasan_store
+  ; NOSTACK-NOT: call {{.*}}__hwasan_generate_tag
+  ; NOSTACK-NOT: call {{.*}}__hwasan_store
+  %buf.sroa.0 = alloca i8, align 4
+  call void @llvm.lifetime.start.p0(i64 1, ptr nonnull %buf.sroa.0)
+  %0 = cmpxchg ptr %buf.sroa.0, i8 1, i8 2 monotonic monotonic, align 4
+  call void @llvm.lifetime.end.p0(i64 1, ptr nonnull %buf.sroa.0)
+  ret i32 0
+}
+
+; Check a safe alloca to ensure it does not get a tag.
+define i32 @test_atomicrwm(ptr %a) sanitize_hwaddress {
+entry:
+  ; CHECK-LABEL: @test_atomicrwm
+  ; NOSAFETY: call {{.*}}__hwasan_generate_tag
+  ; NOSAFETY: call {{.*}}__hwasan_store
+  ; SAFETY-NOT: call {{.*}}__hwasan_generate_tag
+  ; SAFETY-NOT: call {{.*}}__hwasan_store
+  ; NOSTACK-NOT: call {{.*}}__hwasan_generate_tag
+  ; NOSTACK-NOT: call {{.*}}__hwasan_store
+  %buf.sroa.0 = alloca i8, align 4
+  call void @llvm.lifetime.start.p0(i64 1, ptr nonnull %buf.sroa.0)
+  %0 = atomicrmw add ptr %buf.sroa.0, i8 1 monotonic, align 4
+  call void @llvm.lifetime.end.p0(i64 1, ptr nonnull %buf.sroa.0)
+  ret i32 0
+}
+
 ; Check a non-safe alloca to ensure it gets a tag.
 define i32 @test_use(ptr %a) sanitize_hwaddress {
 entry:
   ; CHECK-LABEL: @test_use
   ; NOSAFETY: call {{.*}}__hwasan_generate_tag
   ; NOSAFETY: call {{.*}}__hwasan_store
   ; SAFETY: call {{.*}}__hwasan_generate_tag
   ; SAFETY-NOT: call {{.*}}__hwasan_store
entry:
   %buf.sroa.0 = alloca [10 x i8], align 4
   %ptr = getelementptr [10 x i8], ptr %buf.sroa.0, i32 0, i32 10
   call void @llvm.lifetime.start.p0(i64 10, ptr nonnull %buf.sroa.0)
   store volatile i8 0, ptr %ptr, align 4, !tbaa !8
   call void @llvm.lifetime.end.p0(i64 10, ptr nonnull %buf.sroa.0)
   ret i32 0
 }
 
+define i32 @test_out_of_range2(ptr %a) sanitize_hwaddress {
+entry:
+  ; CHECK-LABEL: @test_out_of_range2
+  ; NOSAFETY: call {{.*}}__hwasan_generate_tag
+  ; NOSAFETY: call {{.*}}__hwasan_store
+  ; SAFETY: call {{.*}}__hwasan_generate_tag
+  ; SAFETY: call {{.*}}__hwasan_store
+  ; NOSTACK-NOT: call {{.*}}__hwasan_generate_tag
+  ; NOSTACK-NOT: call {{.*}}__hwasan_store
+  %buf.sroa.0 = alloca [10 x i8], align 4
+  %ptr = getelementptr [10 x i8], ptr %buf.sroa.0, i32 0, i32 10
+  call void @llvm.lifetime.start.p0(i64 10, ptr nonnull %buf.sroa.0)
+  %0 = cmpxchg ptr %ptr, i8 1, i8 2 monotonic monotonic, align 4
+  call void @llvm.lifetime.end.p0(i64 10, ptr nonnull %buf.sroa.0)
+  ret i32 0
+}
+
 define i32 @test_out_of_range3(ptr %a) sanitize_hwaddress {
 entry:
   ; CHECK-LABEL: @test_out_of_range3
   ; NOSAFETY: call {{.*}}__hwasan_generate_tag
   ; NOSAFETY: call {{.*}}__hwasan_memset
   ; SAFETY: call {{.*}}__hwasan_generate_tag
   ; SAFETY: call {{.*}}__hwasan_memset
   ; NOSTACK-NOT: call {{.*}}__hwasan_generate_tag
entry:
   call void @llvm.lifetime.start.p0(i64 10, ptr nonnull %buf.sroa.0)
   call void @llvm.lifetime.end.p0(i64 10, ptr nonnull %buf.sroa.0)
   call void @llvm.lifetime.start.p0(i64 10, ptr nonnull %buf.sroa.1)
   call void @llvm.memmove.p0.p0.i32(ptr %ptr, ptr %ptr1, i32 1, i1 true)
   call void @llvm.lifetime.end.p0(i64 10, ptr nonnull %buf.sroa.1)
   ret i32 0
 }
 
+define i32 @test_out_of_range6(ptr %a) sanitize_hwaddress {
+entry:
+  ; CHECK-LABEL: @test_out_of_range6
+  ; NOSAFETY: call {{.*}}__hwasan_generate_tag
+  ; NOSAFETY: call {{.*}}__hwasan_store
+  ; SAFETY: call {{.*}}__hwasan_generate_tag
+  ; SAFETY: call {{.*}}__hwasan_store
+  ; NOSTACK-NOT: call {{.*}}__hwasan_generate_tag
+  ; NOSTACK-NOT: call {{.*}}__hwasan_store
+  %buf.sroa.0 = alloca [10 x i8], align 4
+  %ptr = getelementptr [10 x i8], ptr %buf.sroa.0, i32 0, i32 10
+  call void @llvm.lifetime.start.p0(i64 10, ptr nonnull %buf.sroa.0)
+  %0 = atomicrmw add ptr %ptr, i32 1 monotonic, align 4
+  call void @llvm.lifetime.end.p0(i64 10, ptr nonnull %buf.sroa.0)
+  ret i32 0
+}
+
 ; Check an alloca with potentially out of range GEP to ensure it gets a tag and
 ; check.
 define i32 @test_potentially_out_of_range(ptr %a) sanitize_hwaddress {
 entry:
   ; CHECK-LABEL: @test_potentially_out_of_range
   ; NOSAFETY: call {{.*}}__hwasan_generate_tag
   ; NOSAFETY: call {{.*}}__hwasan_store
   ; SAFETY: call {{.*}}__hwasan_generate_tag