From aa654292bd9ca5a68e8b5b6d93ccd3fbee31ab1d Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Thu, 18 Aug 2016 23:51:15 +0000 Subject: [PATCH] [asan] Optimize store size in FunctionStackPoisoner::poisonRedZones Summary: Reduce store size to avoid leading and trailing zeros. Reviewers: kcc, eugenis Subscribers: llvm-commits Differential Revision: https://reviews.llvm.org/D23648 llvm-svn: 279178 --- .../llvm/Transforms/Utils/ASanStackFrameLayout.h | 1 + .../Instrumentation/AddressSanitizer.cpp | 95 +++++++++++++--------- .../Instrumentation/AddressSanitizer/lifetime.ll | 3 +- .../AddressSanitizer/stack-poisoning.ll | 11 +-- 4 files changed, 63 insertions(+), 47 deletions(-) diff --git a/llvm/include/llvm/Transforms/Utils/ASanStackFrameLayout.h b/llvm/include/llvm/Transforms/Utils/ASanStackFrameLayout.h index 0c9b72f..ef36b82 100644 --- a/llvm/include/llvm/Transforms/Utils/ASanStackFrameLayout.h +++ b/llvm/include/llvm/Transforms/Utils/ASanStackFrameLayout.h @@ -24,6 +24,7 @@ class AllocaInst; 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. diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp index 3ad8e14..67aef34 100644 --- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp @@ -795,8 +795,6 @@ struct FunctionStackPoisoner : public InstVisitor { 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, @@ -1926,30 +1924,62 @@ void FunctionStackPoisoner::initializeCallbacks(Module &M) { 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 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(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; + } + + uint64_t Val = 0; + 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; } } @@ -1963,26 +1993,6 @@ static int StackMallocSizeClass(uint64_t LocalStackSize) { 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, @@ -2199,6 +2209,8 @@ void FunctionStackPoisoner::poisonStack() { } }; + SmallVector ShadowBytesAfterReturn; + // (Un)poison the stack before all ret instructions. for (auto Ret : RetVec) { IRBuilder<> IRBRet(Ret); @@ -2225,8 +2237,11 @@ void FunctionStackPoisoner::poisonStack() { 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)); diff --git a/llvm/test/Instrumentation/AddressSanitizer/lifetime.ll b/llvm/test/Instrumentation/AddressSanitizer/lifetime.ll index cf75197..0eaceb8 100644 --- a/llvm/test/Instrumentation/AddressSanitizer/lifetime.ll +++ b/llvm/test/Instrumentation/AddressSanitizer/lifetime.ll @@ -108,10 +108,9 @@ entry: ; 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 } diff --git a/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning.ll b/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning.ll index 4823cdf..a3c852a 100644 --- a/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning.ll +++ b/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning.ll @@ -16,11 +16,12 @@ entry: ; 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 @@ -49,9 +50,9 @@ entry: ; 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. -- 2.7.4