[asan] Optimize store size in FunctionStackPoisoner::poisonRedZones

[vitalybuka]

Summary

Reduce store size to avoid leading and trailing zeros.

[vitalybuka]

Reused to poison stack-after-return

[eugenis]

LGTM

lib/Transforms/Instrumentation/AddressSanitizer.cpp

1959 ↗

(On Diff #68491)

don't do this if !DoPoison

llvm/trunk/include/llvm/Transforms/Utils/ASanStackFrameLayout.h

 
 class AllocaInst;
 
 // These magic constants should be the same as in
 // in asan_internal.h from ASan runtime in compiler-rt.
 static const int kAsanStackLeftRedzoneMagic = 0xf1;
 static const int kAsanStackMidRedzoneMagic = 0xf2;
 static const int kAsanStackRightRedzoneMagic = 0xf3;
+static const int kAsanStackUseAfterReturnMagic = 0xf5;
 static const int kAsanStackUseAfterScopeMagic = 0xf8;
 
 // Input/output data struct for ComputeASanStackFrameLayout.
 struct ASanStackVariableDescription {
   const char *Name;    // Name of the variable that will be displayed by asan
                        // if a stack-related bug is reported.
   uint64_t Size;       // Size of the variable in bytes.
   size_t LifetimeSize; // Size in bytes to use for lifetime analysis check.

llvm/trunk/lib/Transforms/Instrumentation/AddressSanitizer.cpp

struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   }
 
   /// Finds alloca where the value comes from.
   AllocaInst *findAllocaForValue(Value *V);
   void poisonRedZones(ArrayRef<uint8_t> ShadowBytes, IRBuilder<> &IRB,
                       Value *ShadowBase, bool DoPoison);
   void poisonAlloca(Value *V, uint64_t Size, IRBuilder<> &IRB, bool DoPoison);
 
-  void SetShadowToStackAfterReturnInlined(IRBuilder<> &IRB, Value *ShadowBase,
-                                          int Size);
   Value *createAllocaForLayout(IRBuilder<> &IRB, const ASanStackFrameLayout &L,
                                bool Dynamic);
   PHINode *createPHI(IRBuilder<> &IRB, Value *Cond, Value *ValueIfTrue,
                      Instruction *ThenTerm, Value *ValueIfFalse);
 };
 
 } // anonymous namespace
 
void FunctionStackPoisoner::initializeCallbacks(Module &M) {
 
   AsanAllocaPoisonFunc = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
       kAsanAllocaPoison, IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
   AsanAllocasUnpoisonFunc =
       checkSanitizerInterfaceFunction(M.getOrInsertFunction(
           kAsanAllocasUnpoison, IRB.getVoidTy(), IntptrTy, IntptrTy, nullptr));
 }
 
+// If DoPoison is true function copies ShadowBytes into shadow memory.
+// If DoPoison is false function sets 0s into shadow memory.
+// Function assumes that if ShadowBytes[i] is 0, then corresponding shadow
+// memory is constant for duration of the function and it contains 0s. So we
+// will try to minimize writes into corresponding addresses of the real shadow
+// memory.
 void FunctionStackPoisoner::poisonRedZones(ArrayRef<uint8_t> ShadowBytes,
                                            IRBuilder<> &IRB, Value *ShadowBase,
                                            bool DoPoison) {
-  size_t n = ShadowBytes.size();
-  size_t i = 0;
-  // We need to (un)poison n bytes of stack shadow. Poison as many as we can
-  // using 64-bit stores (if we are on 64-bit arch), then poison the rest
-  // with 32-bit stores, then with 16-byte stores, then with 8-byte stores.
-  for (size_t LargeStoreSizeInBytes = ASan.LongSize / 8;
-       LargeStoreSizeInBytes != 0; LargeStoreSizeInBytes /= 2) {
-    for (; i + LargeStoreSizeInBytes - 1 < n; i += LargeStoreSizeInBytes) {
-      uint64_t Val = 0;
-      for (size_t j = 0; j < LargeStoreSizeInBytes; j++) {
-        if (F.getParent()->getDataLayout().isLittleEndian())
+  const size_t End = ShadowBytes.size();
+
+  const size_t LargestStoreSizeInBytes =
+      std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);
+
+  const bool IsLittleEndian = F.getParent()->getDataLayout().isLittleEndian();
+
+  // Poison given range in shadow using larges store size with out leading and
+  // trailing zeros. Zeros never change, so they need neither poisoning nor
+  // up-poisoning, but we don't mind if some of them get into a middle of a
+  // store.
+  for (size_t i = 0; i < End;) {
+    if (!ShadowBytes[i]) {
+      ++i;
+      continue;
+    }
+
+    size_t StoreSizeInBytes = LargestStoreSizeInBytes;
+    // Fit store size into the range.
+    while (StoreSizeInBytes > End - i)
+      StoreSizeInBytes /= 2;
+
+    // Minimize store size by trimming trailing zeros.
+    for (size_t j = StoreSizeInBytes - 1; j && !ShadowBytes[i + j]; --j) {
+      while (j <= StoreSizeInBytes / 2)
+        StoreSizeInBytes /= 2;
+    }
+
+    if (DoPoison) {
+      for (size_t j = 0; j < StoreSizeInBytes; j++) {
+        if (IsLittleEndian)
           Val |= (uint64_t)ShadowBytes[i + j] << (8 * j);
         else
           Val = (Val << 8) | ShadowBytes[i + j];
       }
-      if (!Val) continue;
-      Value *Ptr = IRB.CreateAdd(ShadowBase, ConstantInt::get(IntptrTy, i));
-      Type *StoreTy = Type::getIntNTy(*C, LargeStoreSizeInBytes * 8);
-      Value *Poison = ConstantInt::get(StoreTy, DoPoison ? Val : 0);
-      IRB.CreateStore(Poison, IRB.CreateIntToPtr(Ptr, StoreTy->getPointerTo()));
+      assert(Val); // Impossible because ShadowBytes[i] != 0
+    } else {
+      Val = 0;
     }
+
+    Value *Ptr = IRB.CreateAdd(ShadowBase, ConstantInt::get(IntptrTy, i));
+    Value *Poison = IRB.getIntN(StoreSizeInBytes * 8, Val);
+    IRB.CreateStore(Poison,
+                    IRB.CreateIntToPtr(Ptr, Poison->getType()->getPointerTo()));
+
+    i += StoreSizeInBytes;
   }
 }
 
 // Fake stack allocator (asan_fake_stack.h) has 11 size classes
 // for every power of 2 from kMinStackMallocSize to kMaxAsanStackMallocSizeClass
 static int StackMallocSizeClass(uint64_t LocalStackSize) {
   assert(LocalStackSize <= kMaxStackMallocSize);
   uint64_t MaxSize = kMinStackMallocSize;
   for (int i = 0;; i++, MaxSize *= 2)
     if (LocalStackSize <= MaxSize) return i;
   llvm_unreachable("impossible LocalStackSize");
 }
 
-// Set Size bytes starting from ShadowBase to kAsanStackAfterReturnMagic.
-// We can not use MemSet intrinsic because it may end up calling the actual
-// memset. Size is a multiple of 8.
-// Currently this generates 8-byte stores on x86_64; it may be better to
-// generate wider stores.
-void FunctionStackPoisoner::SetShadowToStackAfterReturnInlined(
-    IRBuilder<> &IRB, Value *ShadowBase, int Size) {
-  assert(!(Size % 8));
-
-  // kAsanStackAfterReturnMagic is 0xf5.
-  const uint64_t kAsanStackAfterReturnMagic64 = 0xf5f5f5f5f5f5f5f5ULL;
-
-  for (int i = 0; i < Size; i += 8) {
-    Value *p = IRB.CreateAdd(ShadowBase, ConstantInt::get(IntptrTy, i));
-    IRB.CreateStore(
-        ConstantInt::get(IRB.getInt64Ty(), kAsanStackAfterReturnMagic64),
-        IRB.CreateIntToPtr(p, IRB.getInt64Ty()->getPointerTo()));
-  }
-}
-
 PHINode *FunctionStackPoisoner::createPHI(IRBuilder<> &IRB, Value *Cond,
                                           Value *ValueIfTrue,
                                           Instruction *ThenTerm,
                                           Value *ValueIfFalse) {
   PHINode *PHI = IRB.CreatePHI(IntptrTy, 2);
   BasicBlock *CondBlock = cast<Instruction>(Cond)->getParent();
   PHI->addIncoming(ValueIfFalse, CondBlock);
   BasicBlock *ThenBlock = ThenTerm->getParent();
auto UnpoisonStack = [&](IRBuilder<> &IRB) {
     poisonRedZones(L.ShadowBytes, IRB, ShadowBase, false);
     if (HavePoisonedStaticAllocas) {
       // If we poisoned some allocas in llvm.lifetime analysis,
       // unpoison whole stack frame now.
       poisonAlloca(LocalStackBase, LocalStackSize, IRB, false);
     }
   };
 
+  SmallVector<uint8_t, 64> ShadowBytesAfterReturn;
+
   // (Un)poison the stack before all ret instructions.
   for (auto Ret : RetVec) {
     IRBuilder<> IRBRet(Ret);
     // Mark the current frame as retired.
     IRBRet.CreateStore(ConstantInt::get(IntptrTy, kRetiredStackFrameMagic),
                        BasePlus0);
     if (DoStackMalloc) {
       assert(StackMallocIdx >= 0);
if (DoStackMalloc) {
       Value *Cmp =
           IRBRet.CreateICmpNE(FakeStack, Constant::getNullValue(IntptrTy));
       TerminatorInst *ThenTerm, *ElseTerm;
       SplitBlockAndInsertIfThenElse(Cmp, Ret, &ThenTerm, &ElseTerm);
 
       IRBuilder<> IRBPoison(ThenTerm);
       if (StackMallocIdx <= 4) {
         int ClassSize = kMinStackMallocSize << StackMallocIdx;
-        SetShadowToStackAfterReturnInlined(IRBPoison, ShadowBase,
-                                           ClassSize >> Mapping.Scale);
+        if ((int)ShadowBytesAfterReturn.size() != ClassSize) {
+          ShadowBytesAfterReturn.resize(ClassSize,
+                                        kAsanStackUseAfterReturnMagic);
+        }
+        poisonRedZones(ShadowBytesAfterReturn, IRBPoison, ShadowBase, true);
         Value *SavedFlagPtrPtr = IRBPoison.CreateAdd(
             FakeStack,
             ConstantInt::get(IntptrTy, ClassSize - ASan.LongSize / 8));
         Value *SavedFlagPtr = IRBPoison.CreateLoad(
             IRBPoison.CreateIntToPtr(SavedFlagPtrPtr, IntptrPtrTy));
         IRBPoison.CreateStore(
             Constant::getNullValue(IRBPoison.getInt8Ty()),
             IRBPoison.CreateIntToPtr(SavedFlagPtr, IRBPoison.getInt8PtrTy()));

llvm/trunk/test/Instrumentation/AddressSanitizer/lifetime.ll

entry:
   ; CHECK: __asan_unpoison_stack_memory
 
   store i8* %0, i8** %d, align 8
 
   call void @llvm.lifetime.end(i64 1024, i8* %0)
   ; CHECK: __asan_poison_stack_memory
 
   ret void
+  ; CHECK: store i32 0
   ; CHECK: store i64 0
   ; CHECK: store i64 0
-  ; CHECK: store i64 0
-  ; CHECK: store i32 0
   ; CHECK-NEXT: __asan_unpoison_stack_memory
 }
 
 define void @zero_sized(i64 %a) #0 {
 ; CHECK-LABEL: define void @zero_sized(i64 %a)
 
 entry:
   %a.addr = alloca i64, align 8

llvm/trunk/test/Instrumentation/AddressSanitizer/stack-poisoning.ll

 ; CHECK-PLAIN-NOT: label
 ; CHECK-PLAIN: ret void
 
 ; CHECK-UAR-LABEL: Bar
 ; CHECK-UAR: load i32, i32* @__asan_option_detect_stack_use_after_return
 ; CHECK-UAR: label
 ; CHECK-UAR: call i64 @__asan_stack_malloc_4
 ; CHECK-UAR: label
+; Poison red zones.
 ; CHECK-UAR: store i64 -1007680412564983311
 ; CHECK-UAR: store i64 72057598113936114
-; CHECK-UAR: store i64 4076008178
-; CHECK-UAR: store i64 -868082074072645632
-; CHECK-UAR: store i32 -202116109
+; CHECK-UAR: store i32 -218959118
+; CHECK-UAR: store i64 -868082074056920316
+; CHECK-UAR: store i16 -3085
 ; CHECK-UAR: call void @Foo
 ; CHECK-UAR: call void @Foo
 ; CHECK-UAR: call void @Foo
 ; If LocalStackBase != OrigStackBase
 ; CHECK-UAR: label
 ; Then Block: poison the entire frame.
   ; CHECK-UAR: store i64 -723401728380766731
   ; CHECK-UAR: store i64 -723401728380766731
; Then Block: poison the entire frame.
   ; CHECK-UAR: store i64 -723401728380766731
   ; CHECK-UAR: store i64 -723401728380766731
   ; CHECK-UAR: store i8 0
   ; CHECK-UAR-NOT: store
   ; CHECK-UAR: label
 ; Else Block: no UAR frame. Only unpoison the redzones.
   ; CHECK-UAR: store i64 0
   ; CHECK-UAR: store i64 0
-  ; CHECK-UAR: store i64 0
-  ; CHECK-UAR: store i64 0
   ; CHECK-UAR: store i32 0
+  ; CHECK-UAR: store i64 0
+  ; CHECK-UAR: store i16 0
   ; CHECK-UAR-NOT: store
   ; CHECK-UAR: label
 ; Done, no more stores.
 ; CHECK-UAR-NOT: store
 ; CHECK-UAR: ret void
 
   %x = alloca [20 x i8], align 16
   %y = alloca [25 x i8], align 1