[HWASan] Add basic stack tagging support for LAM.

Adds the basic instrumentation needed for stack tagging.

Currently does not support stack short granules or TLS stack histories,
since a different code path is followed for the callback instrumentation
we use.

We may simply wait to support these two features until we switch to
a custom calling convention.

Patch By: xiangzhangllvm, morehouse

Reviewed By: vitalybuka

Differential Revision: https://reviews.llvm.org/D102901

lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -69,7 +69,6 @@ static const size_t kNumberOfAccessSizes = 5;
 static const size_t kDefaultShadowScale = 4;
 static const uint64_t kDynamicShadowSentinel =
     std::numeric_limits<uint64_t>::max();
-static const unsigned kPointerTagShift = 56;
 
 static const unsigned kShadowBaseAlignment = 32;
 
@@ -186,6 +185,11 @@ static cl::opt<bool> ClInlineAllChecks("hwasan-inline-all-checks",
                                        cl::desc("inline all checks"),
                                        cl::Hidden, cl::init(false));
 
+// Enabled from clang by "-fsanitize-hwaddress-experimental-aliasing".
+static cl::opt<bool> ClUsePageAliases("hwasan-experimental-use-page-aliases",
+                                      cl::desc("Use page aliasing in HWASan"),
+                                      cl::Hidden, cl::init(false));
+
 namespace {
 
 /// An instrumentation pass implementing detection of addressability bugs
@@ -242,6 +246,9 @@ public:
   Value *getUARTag(IRBuilder<> &IRB, Value *StackTag);
 
   Value *getHwasanThreadSlotPtr(IRBuilder<> &IRB, Type *Ty);
+  Value *applyTagMask(IRBuilder<> &IRB, Value *OldTag);
+  unsigned retagMask(unsigned AllocaNo);
+
   void emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord);
 
   void instrumentGlobal(GlobalVariable *GV, uint8_t Tag);
@@ -298,6 +305,9 @@ private:
   bool HasMatchAllTag = false;
   uint8_t MatchAllTag = 0;
 
+  unsigned PointerTagShift;
+  uint64_t TagMaskByte;
+
   Function *HwasanCtorFunction;
 
   FunctionCallee HwasanMemoryAccessCallback[2][kNumberOfAccessSizes];
@@ -489,11 +499,17 @@ void HWAddressSanitizer::initializeModule() {
 
   TargetTriple = Triple(M.getTargetTriple());
 
-  // x86_64 uses userspace pointer aliases, currently heap-only with callback
-  // instrumentation only.
-  UsePageAliases = TargetTriple.getArch() == Triple::x86_64;
-  InstrumentWithCalls = UsePageAliases ? true : ClInstrumentWithCalls;
+  // x86_64 currently has two modes:
+  // - Intel LAM (default)
+  // - pointer aliasing
+  // Pointer aliasing mode is heap only.  LAM mode is heap+stack, with support
+  // planned for globals as well.
+  bool IsX86_64 = TargetTriple.getArch() == Triple::x86_64;
+  UsePageAliases = ClUsePageAliases && IsX86_64;
+  InstrumentWithCalls = IsX86_64 ? true : ClInstrumentWithCalls;
   InstrumentStack = UsePageAliases ? false : ClInstrumentStack;
+  PointerTagShift = IsX86_64 ? 57 : 56;
+  TagMaskByte = IsX86_64 ? 0x3F : 0xFF;
 
   Mapping.init(TargetTriple, InstrumentWithCalls);
 
@@ -537,7 +553,9 @@ void HWAddressSanitizer::initializeModule() {
     createHwasanCtorComdat();
     bool InstrumentGlobals =
         ClGlobals.getNumOccurrences() ? ClGlobals : NewRuntime;
-    if (InstrumentGlobals && !UsePageAliases)
+
+    // TODO: Support globals for x86_64 in non-aliasing mode.
+    if (InstrumentGlobals && !IsX86_64)
       instrumentGlobals();
 
     bool InstrumentPersonalityFunctions =
@@ -759,7 +777,7 @@ void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,
   }
 
   Value *PtrLong = IRB.CreatePointerCast(Ptr, IntptrTy);
-  Value *PtrTag = IRB.CreateTrunc(IRB.CreateLShr(PtrLong, kPointerTagShift),
+  Value *PtrTag = IRB.CreateTrunc(IRB.CreateLShr(PtrLong, PointerTagShift),
                                   IRB.getInt8Ty());
   Value *AddrLong = untagPointer(IRB, PtrLong);
   Value *Shadow = memToShadow(AddrLong, IRB);
@@ -927,7 +945,10 @@ bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag,
   return true;
 }
 
-static unsigned RetagMask(unsigned AllocaNo) {
+unsigned HWAddressSanitizer::retagMask(unsigned AllocaNo) {
+  if (TargetTriple.getArch() == Triple::x86_64)
+    return AllocaNo & TagMaskByte;
+
   // 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.
@@ -945,6 +966,16 @@ static unsigned RetagMask(unsigned AllocaNo) {
   return FastMasks[AllocaNo % (sizeof(FastMasks) / sizeof(FastMasks[0]))];
 }
 
+Value *HWAddressSanitizer::applyTagMask(IRBuilder<> &IRB, Value *OldTag) {
+  if (TargetTriple.getArch() == Triple::x86_64) {
+    Constant *TagMask = ConstantInt::get(IntptrTy, TagMaskByte);
+    Value *NewTag = IRB.CreateAnd(OldTag, TagMask);
+    return NewTag;
+  }
+  // aarch64 uses 8-bit tags, so no mask is needed.
+  return OldTag;
+}
+
 Value *HWAddressSanitizer::getNextTagWithCall(IRBuilder<> &IRB) {
   return IRB.CreateZExt(IRB.CreateCall(HwasanGenerateTagFunc), IntptrTy);
 }
@@ -968,8 +999,9 @@ Value *HWAddressSanitizer::getStackBaseTag(IRBuilder<> &IRB) {
   // between functions).
   Value *StackPointerLong = IRB.CreatePointerCast(StackPointer, IntptrTy);
   Value *StackTag =
-      IRB.CreateXor(StackPointerLong, IRB.CreateLShr(StackPointerLong, 20),
-                    "hwasan.stack.base.tag");
+      applyTagMask(IRB, IRB.CreateXor(StackPointerLong,
+                                      IRB.CreateLShr(StackPointerLong, 20)));
+  StackTag->setName("hwasan.stack.base.tag");
   return StackTag;
 }
 
@@ -978,7 +1010,7 @@ Value *HWAddressSanitizer::getAllocaTag(IRBuilder<> &IRB, Value *StackTag,
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
   return IRB.CreateXor(StackTag,
-                       ConstantInt::get(IntptrTy, RetagMask(AllocaNo)));
+                       ConstantInt::get(IntptrTy, retagMask(AllocaNo)));
 }
 
 Value *HWAddressSanitizer::getUARTag(IRBuilder<> &IRB, Value *StackTag) {
@@ -986,7 +1018,7 @@ Value *HWAddressSanitizer::getUARTag(IRBuilder<> &IRB, Value *StackTag) {
     return ConstantInt::get(IntptrTy, 0);
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
-  return IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, 0xFFU));
+  return IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, TagMaskByte));
 }
 
 // Add a tag to an address.
@@ -996,13 +1028,13 @@ Value *HWAddressSanitizer::tagPointer(IRBuilder<> &IRB, Type *Ty,
   Value *TaggedPtrLong;
   if (CompileKernel) {
     // Kernel addresses have 0xFF in the most significant byte.
-    Value *ShiftedTag = IRB.CreateOr(
-        IRB.CreateShl(Tag, kPointerTagShift),
-        ConstantInt::get(IntptrTy, (1ULL << kPointerTagShift) - 1));
+    Value *ShiftedTag =
+        IRB.CreateOr(IRB.CreateShl(Tag, PointerTagShift),
+                     ConstantInt::get(IntptrTy, (1ULL << PointerTagShift) - 1));
     TaggedPtrLong = IRB.CreateAnd(PtrLong, ShiftedTag);
   } else {
-    // Userspace can simply do OR (tag << 56);
-    Value *ShiftedTag = IRB.CreateShl(Tag, kPointerTagShift);
+    // Userspace can simply do OR (tag << PointerTagShift);
+    Value *ShiftedTag = IRB.CreateShl(Tag, PointerTagShift);
     TaggedPtrLong = IRB.CreateOr(PtrLong, ShiftedTag);
   }
   return IRB.CreateIntToPtr(TaggedPtrLong, Ty);
@@ -1016,12 +1048,12 @@ Value *HWAddressSanitizer::untagPointer(IRBuilder<> &IRB, Value *PtrLong) {
     // Kernel addresses have 0xFF in the most significant byte.
     UntaggedPtrLong =
         IRB.CreateOr(PtrLong, ConstantInt::get(PtrLong->getType(),
-                                               0xFFULL << kPointerTagShift));
+                                               0xFFULL << PointerTagShift));
   } else {
     // Userspace addresses have 0x00.
-    UntaggedPtrLong = IRB.CreateAnd(
-        PtrLong,
-        ConstantInt::get(PtrLong->getType(), ~(0xFFULL << kPointerTagShift)));
+    UntaggedPtrLong =
+        IRB.CreateAnd(PtrLong, ConstantInt::get(PtrLong->getType(),
+                                                ~(0xFFULL << PointerTagShift)));
   }
   return UntaggedPtrLong;
 }
@@ -1171,7 +1203,7 @@ bool HWAddressSanitizer::instrumentStack(
       // Tag offset logically applies to the alloca pointer, and it makes sense
       // to put it at the beginning of the expression.
       SmallVector<uint64_t, 8> NewOps = {dwarf::DW_OP_LLVM_tag_offset,
-                                         RetagMask(N)};
+                                         retagMask(N)};
       auto Locations = DDI->location_ops();
       unsigned LocNo = std::distance(Locations.begin(), find(Locations, AI));
       DDI->setExpression(
@@ -1428,7 +1460,7 @@ void HWAddressSanitizer::instrumentGlobal(GlobalVariable *GV, uint8_t Tag) {
   Constant *Aliasee = ConstantExpr::getIntToPtr(
       ConstantExpr::getAdd(
           ConstantExpr::getPtrToInt(NewGV, Int64Ty),
-          ConstantInt::get(Int64Ty, uint64_t(Tag) << kPointerTagShift)),
+          ConstantInt::get(Int64Ty, uint64_t(Tag) << PointerTagShift)),
       GV->getType());
   auto *Alias = GlobalAlias::create(GV->getValueType(), GV->getAddressSpace(),
                                     GV->getLinkage(), "", Aliasee, &M);

test/Instrumentation/HWAddressSanitizer/X86/alloca-array.ll

@@ -0,0 +1,15 @@
+; 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 = "x86_64-unknown-linux-gnu"
+
+declare void @use(i8*, i8*)
+
+define void @test_alloca() sanitize_hwaddress {
+  ; CHECK: alloca { [4 x i8], [12 x i8] }, align 16
+  %x = alloca i8, i64 4
+  ; CHECK: alloca i8, i64 16, align 16
+  %y = alloca i8, i64 16
+  call void @use(i8* %x, i8* %y)
+  ret void
+}

test/Instrumentation/HWAddressSanitizer/X86/alloca-with-calls.ll

@@ -0,0 +1,23 @@
+; Test alloca instrumentation when tags are generated by HWASan function.
+;
+; RUN: opt < %s -hwasan -hwasan-generate-tags-with-calls -S | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @use32(i32*)
+
+define void @test_alloca() sanitize_hwaddress {
+; CHECK-LABEL: @test_alloca(
+; CHECK: %[[BC:[^ ]*]] = bitcast { i32, [12 x i8] }* %x to i32*
+; CHECK: %[[T1:[^ ]*]] = call i8 @__hwasan_generate_tag()
+; CHECK: %[[A:[^ ]*]] = zext i8 %[[T1]] to i64
+; CHECK: %[[B:[^ ]*]] = ptrtoint i32* %[[BC]] to i64
+; CHECK: %[[C:[^ ]*]] = shl i64 %[[A]], 57
+; CHECK: or i64 %[[B]], %[[C]]
+
+entry:
+  %x = alloca i32, align 4
+  call void @use32(i32* nonnull %x)
+  ret void
+}

test/Instrumentation/HWAddressSanitizer/X86/alloca.ll

@@ -0,0 +1,45 @@
+; Test alloca instrumentation.
+;
+; RUN: opt < %s -hwasan -S | FileCheck %s --check-prefixes=CHECK,NO-UAR-TAGS
+; RUN: opt < %s -hwasan -hwasan-uar-retag-to-zero=0 -S | FileCheck %s --check-prefixes=CHECK,UAR-TAGS
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @use32(i32*)
+
+define void @test_alloca() sanitize_hwaddress {
+; CHECK-LABEL: @test_alloca(
+; CHECK: %[[FP:[^ ]*]] = call i8* @llvm.frameaddress.p0i8(i32 0)
+; CHECK: %[[A:[^ ]*]] = ptrtoint i8* %[[FP]] to i64
+; CHECK: %[[B:[^ ]*]] = lshr i64 %[[A]], 20
+; CHECK: %[[A_XOR_B:[^ ]*]] = xor i64 %[[A]], %[[B]]
+; CHECK: %[[BASE_TAG:[^ ]*]] = and i64 %[[A_XOR_B]], 63
+
+; CHECK: %[[X:[^ ]*]] = alloca { i32, [12 x i8] }, align 16
+; CHECK: %[[X_BC:[^ ]*]] = bitcast { i32, [12 x i8] }* %[[X]] to i32*
+; CHECK: %[[X_TAG:[^ ]*]] = xor i64 %[[BASE_TAG]], 0
+; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X_BC]] to i64
+; CHECK: %[[C:[^ ]*]] = shl i64 %[[X_TAG]], 57
+; CHECK: %[[D:[^ ]*]] = or i64 %[[X1]], %[[C]]
+; CHECK: %[[X_HWASAN:[^ ]*]] = inttoptr i64 %[[D]] to i32*
+
+; CHECK: %[[X_TAG2:[^ ]*]] = trunc i64 %[[X_TAG]] to i8
+; CHECK: %[[X_I8:[^ ]*]] = bitcast i32* %[[X_BC]] to i8*
+; CHECK: call void @__hwasan_tag_memory(i8* %[[X_I8]], i8 %[[X_TAG2]], i64 16)
+
+; CHECK: call void @use32(i32* nonnull %[[X_HWASAN]])
+
+; UAR-TAGS: %[[BASE_TAG_COMPL:[^ ]*]] = xor i64 %[[BASE_TAG]], 63
+; UAR-TAGS: %[[X_TAG_UAR:[^ ]*]] = trunc i64 %[[BASE_TAG_COMPL]] to i8
+; CHECK: %[[X_I8_2:[^ ]*]] = bitcast i32* %[[X_BC]] to i8*
+; NO-UAR-TAGS: call void @__hwasan_tag_memory(i8* %[[X_I8_2]], i8 0, i64 16)
+; UAR-TAGS: call void @__hwasan_tag_memory(i8* %[[X_I8_2]], i8 %[[X_TAG_UAR]], i64 16)
+; CHECK: ret void
+
+
+entry:
+  %x = alloca i32, align 4
+  call void @use32(i32* nonnull %x)
+  ret void
+}