[hwasan] Stack instrumentation.

[eugenis]

Summary

Very basic stack instrumentation using tagged pointers.
Tag for N'th alloca in a function is built as XOR of:

• base tag for the function, which is just some bits of SP (poor man's random)
• small constant which is a function of N.

Allocas are aligned to 16 bytes. On every ReturnInst allocas are
re-tagged to catch use-after-return.

This implementation has a bunch of issues that will be taken care of
later:

1. lifetime intrinsics referring to tagged pointers are not recognized in SDAG. This effectively disables stack coloring.
2. Generated code is quite inefficient. There is one extra instruction at each memory access that adds the base tag to the untagged alloca address. It would be better to keep tagged SP in a callee-saved register and address allocas as an offset of that XOR retag, but that needs better coordination between hwasan instrumentation pass and prologue/epilogue insertion.
3. Lifetime instrinsics are ignored and use-after-scope is not implemented. This would be harder to do than in ASan, because we need to use a differently tagged pointer depending on which lifetime.start / lifetime.end the current instruction is dominated / post-dominated.

[kcc]

LGTM with two nits, feel free to address them separately.

compiler-rt/test/hwasan/TestCases/stack-oob.cc

1 ↗

(On Diff #128245)

I'd add more sizes here.

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

355 ↗

(On Diff #128245)

Will this work if

• the size is large and memset doesn't get inlined
• hwasan has the msan interceptor

?

355 ↗

(On Diff #128245)

I mean,

• hwasan has the memset interceptor

[eugenis]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

355 ↗

(On Diff #128245)

That's why I'm changing the memset interceptor to skips checks if the address is in the shadow range.

[eugenis]

extended the test

[kcc]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

355 ↗

(On Diff #128245)

of course! I didn't connect the dots.
Then just leave a comment about this here, and remove the commented code above.

[eugenis]

added a comment

compiler-rt/trunk/lib/hwasan/hwasan.h

 // Reasonable values are 4 (for 1/16th shadow) and 6 (for 1/64th).
 const uptr kShadowScale = 4;
 const uptr kShadowAlignment = 1UL << kShadowScale;
 
 #define MEM_TO_SHADOW_OFFSET(mem) ((uptr)(mem) >> kShadowScale)
 #define MEM_TO_SHADOW(mem) ((uptr)(mem) >> kShadowScale)
 #define SHADOW_TO_MEM(shadow) ((uptr)(shadow) << kShadowScale)
 
-#define MEM_IS_APP(mem) true
+#define MEM_IS_APP(mem) MemIsApp((uptr)(mem))
 
 // TBI (Top Byte Ignore) feature of AArch64: bits [63:56] are ignored in address
 // translation and can be used to store a tag.
 const unsigned kAddressTagShift = 56;
 const uptr kAddressTagMask = 0xFFUL << kAddressTagShift;
 
 static inline tag_t GetTagFromPointer(uptr p) {
   return p >> kAddressTagShift;

 }
 
 namespace __hwasan {
 
 extern int hwasan_inited;
 extern bool hwasan_init_is_running;
 extern int hwasan_report_count;
 
+bool MemIsApp(uptr p);
+
 bool ProtectRange(uptr beg, uptr end);
 bool InitShadow();
 char *GetProcSelfMaps();
 void InitializeInterceptors();
 
 void HwasanAllocatorInit();
 void HwasanAllocatorThreadFinish();
 void HwasanDeallocate(StackTrace *stack, void *ptr);

compiler-rt/trunk/lib/hwasan/hwasan.cc

 }
 void __hwasan_store8_noabort(uptr p) {
   CheckAddress<ErrorAction::Recover, AccessType::Store, 3>(p);
 }
 void __hwasan_store16_noabort(uptr p) {
   CheckAddress<ErrorAction::Recover, AccessType::Store, 4>(p);
 }
 
+void __hwasan_tag_memory(uptr p, u8 tag, uptr sz) {
+  TagMemoryAligned(p, sz, tag);
+}
+
 #if !SANITIZER_SUPPORTS_WEAK_HOOKS
 extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
 const char* __hwasan_default_options() { return ""; }
 }  // extern "C"
 #endif
 
 extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_print_stack_trace() {
   GET_FATAL_STACK_TRACE_PC_BP(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME());
   stack.Print();
 }
 } // extern "C"

compiler-rt/trunk/lib/hwasan/hwasan_interceptors.cc

 #define COMMON_INTERCEPTOR_GET_TLS_RANGE(begin, end)                           \
   if (HwasanThread *t = GetCurrentThread()) {                                    \
     *begin = t->tls_begin();                                                   \
     *end = t->tls_end();                                                       \
   } else {                                                                     \
     *begin = *end = 0;                                                         \
   }
 
+#define COMMON_INTERCEPTOR_MEMSET_IMPL(ctx, dst, v, size) \
+  {                                                       \
+    COMMON_INTERCEPTOR_ENTER(ctx, memset, dst, v, size);  \
+    if (common_flags()->intercept_intrin &&               \
+        MEM_IS_APP(GetAddressFromPointer(dst)))           \
+      COMMON_INTERCEPTOR_WRITE_RANGE(ctx, dst, size);     \
+    return REAL(memset)(dst, v, size);                    \
+  }
+
 #include "sanitizer_common/sanitizer_platform_interceptors.h"
 #include "sanitizer_common/sanitizer_common_interceptors.inc"
 #include "sanitizer_common/sanitizer_signal_interceptors.inc"
 
 #define COMMON_SYSCALL_PRE_READ_RANGE(p, s) CHECK_UNPOISONED(p, s)
 #define COMMON_SYSCALL_PRE_WRITE_RANGE(p, s) \
   do {                                       \
     (void)(p);                               \

compiler-rt/trunk/lib/hwasan/hwasan_interface_internal.h

 void __hwasan_store2_noabort(uptr);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __hwasan_store4_noabort(uptr);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __hwasan_store8_noabort(uptr);
 SANITIZER_INTERFACE_ATTRIBUTE
 void __hwasan_store16_noabort(uptr);
 
+SANITIZER_INTERFACE_ATTRIBUTE
+void __hwasan_tag_memory(uptr p, u8 tag, uptr sz);
+
 // Returns the offset of the first tag mismatch or -1 if the whole range is
 // good.
 SANITIZER_INTERFACE_ATTRIBUTE
 sptr __hwasan_test_shadow(const void *x, uptr size);
 
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
 /* OPTIONAL */ const char* __hwasan_default_options();
 

compiler-rt/trunk/lib/hwasan/hwasan_linux.cc

static void ProtectGap(uptr addr, uptr size) {
 
   Report(
       "ERROR: Failed to protect the shadow gap. "
       "ASan cannot proceed correctly. ABORTING.\n");
   DumpProcessMap();
   Die();
 }
 
-bool InitShadow() {
-  const uptr maxVirtualAddress = GetMaxUserVirtualAddress();
-
-  // LowMem covers as much of the first 4GB as possible.
-  const uptr kLowMemEnd = 1UL<<32;
-  const uptr kLowShadowEnd = kLowMemEnd >> kShadowScale;
-  const uptr kLowShadowStart = kLowShadowEnd >> kShadowScale;
+static uptr kHighShadowStart;
+static uptr kHighShadowEnd;
+static uptr kHighMemStart;
+
+bool InitShadow() {
+  const uptr maxVirtualAddress = GetMaxUserVirtualAddress();
+
 
   // HighMem covers the upper part of the address space.
-  const uptr kHighShadowEnd = (maxVirtualAddress >> kShadowScale) + 1;
-  const uptr kHighShadowStart = Max(kLowMemEnd, kHighShadowEnd >> kShadowScale);
+  kHighShadowEnd = (maxVirtualAddress >> kShadowScale) + 1;
+  kHighShadowStart = Max(kLowMemEnd, kHighShadowEnd >> kShadowScale);
   CHECK(kHighShadowStart < kHighShadowEnd);
 
-  const uptr kHighMemStart = kHighShadowStart << kShadowScale;
+  kHighMemStart = kHighShadowStart << kShadowScale;
   CHECK(kHighShadowEnd <= kHighMemStart);
 
   if (Verbosity()) {
     Printf("|| `[%p, %p]` || HighMem    ||\n", (void *)kHighMemStart,
            (void *)maxVirtualAddress);
     if (kHighMemStart > kHighShadowEnd)
       Printf("|| `[%p, %p]` || ShadowGap2 ||\n", (void *)kHighShadowEnd,
              (void *)kHighMemStart);
bool InitShadow() {
   if (kHighShadowStart > kLowMemEnd)
     ProtectGap(kLowMemEnd, kHighShadowStart - kLowMemEnd);
   if (kHighMemStart > kHighShadowEnd)
     ProtectGap(kHighShadowEnd, kHighMemStart - kHighShadowEnd);
 
   return true;
 }
 
+bool MemIsApp(uptr p) {
+  CHECK(GetTagFromPointer(p) == 0);
+  return p >= kHighMemStart || (p >= kLowShadowEnd && p < kLowMemEnd);
+}
+
 static void HwasanAtExit(void) {
   if (flags()->print_stats && (flags()->atexit || hwasan_report_count > 0))
     ReportStats();
   if (hwasan_report_count > 0) {
     // ReportAtExitStatistics();
     if (common_flags()->exitcode)
       internal__exit(common_flags()->exitcode);
   }

compiler-rt/trunk/test/hwasan/TestCases/stack-oob.cc

+// RUN: %clangxx_hwasan -DSIZE=16 -O0 %s -o %t && not %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_hwasan -DSIZE=64 -O0 %s -o %t && not %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_hwasan -DSIZE=0x1000 -O0 %s -o %t && not %run %t 2>&1 | FileCheck %s
+
+// REQUIRES: stable-runtime
+
+#include <stdlib.h>
+#include <sanitizer/hwasan_interface.h>
+
+__attribute__((noinline))
+int f() {
+  char z[SIZE];
+  char *volatile p = z;
+  return p[SIZE];
+}
+
+int main() {
+  return f();
+  // CHECK: READ of size 1 at
+  // CHECK: #0 {{.*}} in f{{.*}}stack-oob.cc:14
+
+  // CHECK: HWAddressSanitizer can not describe address in more detail.
+
+  // CHECK: SUMMARY: HWAddressSanitizer: tag-mismatch {{.*}} in f
+}

compiler-rt/trunk/test/hwasan/TestCases/stack-uar.cc

+// RUN: %clangxx_hwasan -O0 %s -o %t && not %run %t 2>&1 | FileCheck %s
+
+// REQUIRES: stable-runtime
+
+#include <stdlib.h>
+#include <sanitizer/hwasan_interface.h>
+
+__attribute__((noinline))
+char *f() {
+  char z[0x1000];
+  char *volatile p = z;
+  return p;
+}
+
+int main() {
+  return *f();
+  // CHECK: READ of size 1 at
+  // CHECK: #0 {{.*}} in main{{.*}}stack-uar.cc:16
+
+  // CHECK: HWAddressSanitizer can not describe address in more detail.
+
+  // CHECK: SUMMARY: HWAddressSanitizer: tag-mismatch {{.*}} in main
+}

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue getTargetNode(BlockAddressSDNode *N, EVT Ty, SelectionDAG &DAG,
                         unsigned Flag) const;
   template <class NodeTy>
   SDValue getGOT(NodeTy *N, SelectionDAG &DAG, unsigned Flags = 0) const;
   template <class NodeTy>
   SDValue getAddrLarge(NodeTy *N, SelectionDAG &DAG, unsigned Flags = 0) const;
   template <class NodeTy>
   SDValue getAddr(NodeTy *N, SelectionDAG &DAG, unsigned Flags = 0) const;
+  SDValue LowerADDROFRETURNADDR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDarwinGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerELFGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerELFTLSDescCallSeq(SDValue SymAddr, const SDLoc &DL,
                                  SelectionDAG &DAG) const;
   SDValue LowerSETCC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBR_CC(SDValue Op, SelectionDAG &DAG) const;

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

 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
-#include "llvm/IR/MDBuilder.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/IR/Function.h"
-#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Instrumentation.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "hwasan"
 
 static const char *const kHwasanModuleCtorName = "hwasan.module_ctor";
 static const char *const kHwasanInitName = "__hwasan_init";
 
 // Accesses sizes are powers of two: 1, 2, 4, 8, 16.
 static const size_t kNumberOfAccessSizes = 5;
 
 static const size_t kShadowScale = 4;
+static const unsigned kAllocaAlignment = 1U << kShadowScale;
 static const unsigned kPointerTagShift = 56;
 
 static cl::opt<std::string> ClMemoryAccessCallbackPrefix(
     "hwasan-memory-access-callback-prefix",
     cl::desc("Prefix for memory access callbacks"), cl::Hidden,
     cl::init("__hwasan_"));
 
 static cl::opt<bool>
static cl::opt<bool> ClInstrumentAtomics(
     cl::desc("instrument atomic instructions (rmw, cmpxchg)"), cl::Hidden,
     cl::init(true));
 
 static cl::opt<bool> ClRecover(
     "hwasan-recover",
     cl::desc("Enable recovery mode (continue-after-error)."),
     cl::Hidden, cl::init(false));
 
+static cl::opt<bool> ClInstrumentStack("hwasan-instrument-stack",
+                                       cl::desc("instrument stack (allocas)"),
+                                       cl::Hidden, cl::init(true));
+
 namespace {
 
 /// \brief An instrumentation pass implementing detection of addressability bugs
 /// using tagged pointers.
 class HWAddressSanitizer : public FunctionPass {
 public:
   // Pass identification, replacement for typeid.
   static char ID;
public:
   void instrumentMemAccessInline(Value *PtrLong, bool IsWrite,
                                  unsigned AccessSizeIndex,
                                  Instruction *InsertBefore);
   bool instrumentMemAccess(Instruction *I);
   Value *isInterestingMemoryAccess(Instruction *I, bool *IsWrite,
                                    uint64_t *TypeSize, unsigned *Alignment,
                                    Value **MaybeMask);
 
+  bool isInterestingAlloca(const AllocaInst &AI);
+  bool tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag);
+  bool instrumentStack(SmallVectorImpl<AllocaInst *> &Allocas,
+                       SmallVectorImpl<Instruction *> &RetVec);
+
 private:
   LLVMContext *C;
   Type *IntptrTy;
+  Type *Int8Ty;
 
   bool Recover;
 
   Function *HwasanCtorFunction;
 
   Function *HwasanMemoryAccessCallback[2][kNumberOfAccessSizes];
   Function *HwasanMemoryAccessCallbackSized[2];
+
+  Function *HwasanTagMemoryFunc;
 };
 
 } // end anonymous namespace
 
 char HWAddressSanitizer::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     HWAddressSanitizer, "hwasan",
bool HWAddressSanitizer::doInitialization(Module &M) {
   DEBUG(dbgs() << "Init " << M.getName() << "\n");
   auto &DL = M.getDataLayout();
 
   Triple TargetTriple(M.getTargetTriple());
 
   C = &(M.getContext());
   IRBuilder<> IRB(*C);
   IntptrTy = IRB.getIntPtrTy(DL);
+  Int8Ty = IRB.getInt8Ty();
 
   std::tie(HwasanCtorFunction, std::ignore) =
       createSanitizerCtorAndInitFunctions(M, kHwasanModuleCtorName,
                                           kHwasanInitName,
                                           /*InitArgTypes=*/{},
                                           /*InitArgs=*/{});
   appendToGlobalCtors(M, HwasanCtorFunction, 0);
   return true;
for (size_t AccessSizeIndex = 0; AccessSizeIndex < kNumberOfAccessSizes;
          AccessSizeIndex++) {
       HwasanMemoryAccessCallback[AccessIsWrite][AccessSizeIndex] =
           checkSanitizerInterfaceFunction(M.getOrInsertFunction(
               ClMemoryAccessCallbackPrefix + TypeStr +
                   itostr(1ULL << AccessSizeIndex) + EndingStr,
               FunctionType::get(IRB.getVoidTy(), {IntptrTy}, false)));
     }
   }
+
+  HwasanTagMemoryFunc = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
+      "__hwasan_tag_memory", IRB.getVoidTy(), IntptrTy, Int8Ty, IntptrTy));
 }
 
 Value *HWAddressSanitizer::isInterestingMemoryAccess(Instruction *I,
                                                    bool *IsWrite,
                                                    uint64_t *TypeSize,
                                                    unsigned *Alignment,
                                                    Value **MaybeMask) {
   // Skip memory accesses inserted by another instrumentation.
bool HWAddressSanitizer::instrumentMemAccess(Instruction *I) {
   } else {
     IRB.CreateCall(HwasanMemoryAccessCallbackSized[IsWrite],
                    {AddrLong, ConstantInt::get(IntptrTy, TypeSize / 8)});
   }
 
   return true;
 }
 
+static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
+  uint64_t ArraySize = 1;
+  if (AI.isArrayAllocation()) {
+    const ConstantInt *CI = dyn_cast<ConstantInt>(AI.getArraySize());
+    assert(CI && "non-constant array size");
+    ArraySize = CI->getZExtValue();
+  }
+  Type *Ty = AI.getAllocatedType();
+  uint64_t SizeInBytes = AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
+  return SizeInBytes * ArraySize;
+}
+
+bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI,
+                                   Value *Tag) {
+  size_t Size = (getAllocaSizeInBytes(*AI) + kAllocaAlignment - 1) &
+                ~(kAllocaAlignment - 1);
+
+  Value *JustTag = IRB.CreateTrunc(Tag, IRB.getInt8Ty());
+  if (ClInstrumentWithCalls) {
+    IRB.CreateCall(HwasanTagMemoryFunc,
+                   {IRB.CreatePointerCast(AI, IntptrTy), JustTag,
+                    ConstantInt::get(IntptrTy, Size)});
+  } else {
+    size_t ShadowSize = Size >> kShadowScale;
+    Value *ShadowPtr = IRB.CreateIntToPtr(
+        IRB.CreateLShr(IRB.CreatePointerCast(AI, IntptrTy), kShadowScale),
+        IRB.getInt8PtrTy());
+    // If this memset is not inlined, it will be intercepted in the hwasan
+    // runtime library. That's OK, because the interceptor skips the checks if
+    // the address is in the shadow region.
+    // FIXME: the interceptor is not as fast as real memset. Consider lowering
+    // llvm.memset right here into either a sequence of stores, or a call to
+    // hwasan_tag_memory.
+    IRB.CreateMemSet(ShadowPtr, JustTag, ShadowSize, /*Align=*/1);
+  }
+  return true;
+}
+
+static unsigned RetagMask(unsigned AllocaNo) {
+  // A list of 8-bit numbers that have at most one run of non-zero bits.
+  // x = x ^ (mask << 56) can be encoded as a single armv8 instruction for these
+  // masks.
+  // The list does not include the value 255, which is used for UAR.
+  static unsigned FastMasks[] = {
+      0,   1,   2,   3,   4,   6,   7,   8,   12,  14,  15, 16,  24,
+      28,  30,  31,  32,  48,  56,  60,  62,  63,  64,  96, 112, 120,
+      124, 126, 127, 128, 192, 224, 240, 248, 252, 254};
+  return FastMasks[AllocaNo % (sizeof(FastMasks) / sizeof(FastMasks[0]))];
+}
+
+bool HWAddressSanitizer::instrumentStack(
+    SmallVectorImpl<AllocaInst *> &Allocas,
+    SmallVectorImpl<Instruction *> &RetVec) {
+  Function *F = Allocas[0]->getParent()->getParent();
+  Module *M = F->getParent();
+  Instruction *InsertPt = &*F->getEntryBlock().begin();
+  IRBuilder<> IRB(InsertPt);
+
+  // FIXME: use addressofreturnaddress (but implement it in aarch64 backend
+  // first).
+  auto GetStackPointerFn = Intrinsic::getDeclaration(M, Intrinsic::frameaddress);
+  Value *StackPointer = IRB.CreateCall(GetStackPointerFn, {Constant::getNullValue(IRB.getInt32Ty())});
+
+  // Extract some entropy from the stack pointer for the tags.
+  // Take bits 20..28 (ASLR entropy) and xor with bits 0..8 (these differ
+  // between functions).
+  Value *StackPointerLong = IRB.CreatePointerCast(StackPointer, IntptrTy);
+  Value *StackTag =
+      IRB.CreateXor(StackPointerLong, IRB.CreateLShr(StackPointerLong, 20),
+                    "hwasan.stack.base.tag");
+
+  // Ideally, we want to calculate tagged stack base pointer, and rewrite all
+  // alloca addresses using that. Unfortunately, offsets are not known yet
+  // (unless we use ASan-style mega-alloca). Instead we keep the base tag in a
+  // temp, shift-OR it into each alloca address and xor with the retag mask.
+  // This generates one extra instruction per alloca use.
+  for (unsigned N = 0; N < Allocas.size(); ++N) {
+    auto *AI = Allocas[N];
+    IRB.SetInsertPoint(AI->getNextNode());
+
+    // Replace uses of the alloca with tagged address.
+    std::string Name =
+        AI->hasName() ? AI->getName().str() : "alloca." + itostr(N);
+    Value *Tag =
+        IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, RetagMask(N)));
+    Value *AILong = IRB.CreatePointerCast(AI, IntptrTy);
+    Value *Replacement = IRB.CreateIntToPtr(
+        IRB.CreateOr(AILong, IRB.CreateShl(Tag, kPointerTagShift)),
+        AI->getType(), Name + ".hwasan");
+
+    for (auto UI = AI->use_begin(), UE = AI->use_end();
+         UI != UE;) {
+      Use &U = *UI++;
+      if (U.getUser() != AILong)
+        U.set(Replacement);
+    }
+
+    tagAlloca(IRB, AI, Tag);
+
+    for (auto RI : RetVec) {
+      IRB.SetInsertPoint(RI);
+
+      // Re-tag alloca memory with the special UAR tag.
+      Value *Tag = IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, 0xFFU));
+      tagAlloca(IRB, AI, Tag);
+    }
+  }
+
+  return true;
+}
+
+bool HWAddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
+  return (AI.getAllocatedType()->isSized() &&
+          // FIXME: instrument dynamic allocas, too
+          AI.isStaticAlloca() &&
+          // alloca() may be called with 0 size, ignore it.
+          getAllocaSizeInBytes(AI) > 0 &&
+          // We are only interested in allocas not promotable to registers.
+          // Promotable allocas are common under -O0.
+          !isAllocaPromotable(&AI) &&
+          // inalloca allocas are not treated as static, and we don't want
+          // dynamic alloca instrumentation for them as well.
+          !AI.isUsedWithInAlloca() &&
+          // swifterror allocas are register promoted by ISel
+          !AI.isSwiftError());
+}
+
 bool HWAddressSanitizer::runOnFunction(Function &F) {
   if (&F == HwasanCtorFunction)
     return false;
 
   if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
     return false;
 
   DEBUG(dbgs() << "Function: " << F.getName() << "\n");
 
   initializeCallbacks(*F.getParent());
 
   bool Changed = false;
   SmallVector<Instruction*, 16> ToInstrument;
+  SmallVector<AllocaInst*, 8> AllocasToInstrument;
+  SmallVector<Instruction*, 8> RetVec;
   for (auto &BB : F) {
     for (auto &Inst : BB) {
+      if (ClInstrumentStack)
+        if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
+          // Realign all allocas. We don't want small uninteresting allocas to
+          // hide in instrumented alloca's padding.
+          if (AI->getAlignment() < kAllocaAlignment)
+            AI->setAlignment(kAllocaAlignment);
+          // Instrument some of them.
+          if (isInterestingAlloca(*AI))
+            AllocasToInstrument.push_back(AI);
+          continue;
+        }
+
+      if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst) || isa<CleanupReturnInst>(Inst))
+        RetVec.push_back(&Inst);
+
       Value *MaybeMask = nullptr;
       bool IsWrite;
       unsigned Alignment;
       uint64_t TypeSize;
       Value *Addr = isInterestingMemoryAccess(&Inst, &IsWrite, &TypeSize,
                                               &Alignment, &MaybeMask);
       if (Addr || isa<MemIntrinsic>(Inst))
         ToInstrument.push_back(&Inst);
     }
   }
 
+  if (!AllocasToInstrument.empty())
+    Changed |= instrumentStack(AllocasToInstrument, RetVec);
+
   for (auto Inst : ToInstrument)
     Changed |= instrumentMemAccess(Inst);
 
   return Changed;
 }

llvm/trunk/test/Instrumentation/HWAddressSanitizer/alloca.ll

+; Test basic address sanitizer instrumentation.
+;
+; RUN: opt < %s -hwasan -S | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-android"
+
+declare void @use32(i32*)
+
+define void @test_alloca() sanitize_hwaddress {
+; CHECK-LABEL: @test_alloca(
+; CHECK: %[[FP:[^ ]*]] = call i8* @llvm.frameaddress(i32 0)
+; CHECK: %[[A:[^ ]*]] = ptrtoint i8* %[[FP]] to i64
+; CHECK: %[[B:[^ ]*]] = lshr i64 %[[A]], 20
+; CHECK: %[[BASE_TAG:[^ ]*]] = xor i64 %[[A]], %[[B]]
+
+; CHECK: %[[X:[^ ]*]] = alloca i32, align 16
+; CHECK: %[[X_TAG:[^ ]*]] = xor i64 %[[BASE_TAG]], 0
+; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X]] to i64
+; CHECK: %[[C:[^ ]*]] = shl i64 %[[X_TAG]], 56
+; CHECK: %[[D:[^ ]*]] = or i64 %[[X1]], %[[C]]
+; CHECK: %[[X_HWASAN:[^ ]*]] = inttoptr i64 %[[D]] to i32*
+
+; CHECK: %[[X_TAG2:[^ ]*]] = trunc i64 %[[X_TAG]] to i8
+; CHECK: %[[E:[^ ]*]] = ptrtoint i32* %[[X]] to i64
+; CHECK: %[[F:[^ ]*]] = lshr i64 %[[E]], 4
+; CHECK: %[[X_SHADOW:[^ ]*]] = inttoptr i64 %[[F]] to i8*
+; CHECK: call void @llvm.memset.p0i8.i64(i8* %[[X_SHADOW]], i8 %[[X_TAG2]], i64 1, i32 1, i1 false)
+; CHECK: call void @use32(i32* nonnull %[[X_HWASAN]])
+
+; CHECK: %[[X_TAG_UAR:[^ ]*]] = xor i64 %[[BASE_TAG]], 255
+; CHECK: %[[X_TAG_UAR2:[^ ]*]] = trunc i64 %[[X_TAG_UAR]] to i8
+; CHECK: %[[E2:[^ ]*]] = ptrtoint i32* %[[X]] to i64
+; CHECK: %[[F2:[^ ]*]] = lshr i64 %[[E2]], 4
+; CHECK: %[[X_SHADOW2:[^ ]*]] = inttoptr i64 %[[F2]] to i8*
+; CHECK: call void @llvm.memset.p0i8.i64(i8* %[[X_SHADOW2]], i8 %[[X_TAG_UAR2]], i64 1, i32 1, i1 false)
+; CHECK: ret void
+
+
+entry:
+  %x = alloca i32, align 4
+  call void @use32(i32* nonnull %x)
+  ret void
+}
+