[llvm][hwasan] Decouple use of the TLS global for getting the shadow base and using the frame record feature

[leonardchan]

Summary

This allows for using the frame record feature (which uses __hwasan_tls) independently from however the user wants to access the shadow base, which prior was only usable if shadow wasn't accessed through the TLS variable or ifuncs.

Frame recording can be explicitly set according to ShadowMapping::WithFrameRecord in ShadowMapping::init. Currently, it is only enabled on Fuchsia and if TLS is used, so this should mimic the old behavior.

Added an extra case to prologue.ll that covers this new case.

[mcgrathr]

lgtm

[eugenis]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1054

I think this will break Linux kernel build.
They use -fsanitize=kernel-hwaddress with one of

-hwasan-instrument-with-calls=1

or

-hwasan-mapping-offset=

The assumption here is that InTls == 0 means there is no thread slot, and therefore no frame recording.

[leonardchan]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1054

Perhaps we could set the default value of -hwasan-record-stack-history to false so users who don't expect frame recording don't suddenly see TLS being used? It would mean people who do use frame recording won't see it anymore and will need to manually add -hwasan-record-stack-history=1, but maybe that won't break existing users.

[eugenis]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1054

These are internal compiler flags, they should not have users, ideally. The right thing should happen based on -fsanitize flag and the target triple, and the right thing in this case is to enable frame recording on supported platforms.

I'd rather put this check under InTls, or add something like ShadowMapping::WithFrameRecord and make the decision in ShadowMapping::init.

[leonardchan]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

1054

Updated to add ShadowMapping::WithFrameRecord.

[eugenis]

LGTM with a nit

llvm/test/Instrumentation/HWAddressSanitizer/prologue.ll

30

CHECK-ZERO-OFFSET is never used

[leonardchan]

llvm/test/Instrumentation/HWAddressSanitizer/prologue.ll

30

Woops. Meant to add a RUN case for testing a fuchsia target which should use this. Updated.

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

private:
   /// This struct defines the shadow mapping using the rule:
   ///   shadow = (mem >> Scale) + Offset.
   /// If InGlobal is true, then
   ///   extern char __hwasan_shadow[];
   ///   shadow = (mem >> Scale) + &__hwasan_shadow
   /// If InTls is true, then
   ///   extern char *__hwasan_tls;
   ///   shadow = (mem>>Scale) + align_up(__hwasan_shadow, kShadowBaseAlignment)
+  ///
+  /// If WithFrameRecord is true, then __hwasan_tls will be used to access the
+  /// ring buffer for storing stack allocations on targets that support it.
   struct ShadowMapping {
     int Scale;
     uint64_t Offset;
     bool InGlobal;
     bool InTls;
+    bool WithFrameRecord;
 
     void init(Triple &TargetTriple, bool InstrumentWithCalls);
     unsigned getObjectAlignment() const { return 1U << Scale; }
   };
   ShadowMapping Mapping;
 
   Type *VoidTy = Type::getVoidTy(M.getContext());
   Type *IntptrTy;
Value *HWAddressSanitizer::getHwasanThreadSlotPtr(IRBuilder<> &IRB, Type *Ty) {
   }
   if (ThreadPtrGlobal)
     return ThreadPtrGlobal;
 
   return nullptr;
 }
 
 void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
-  if (!Mapping.InTls) {
+  if (!Mapping.InTls)
     ShadowBase = getShadowNonTls(IRB);
-    return;
-  }
-
-  if (!WithFrameRecord && TargetTriple.isAndroid()) {
+  else if (!WithFrameRecord && TargetTriple.isAndroid())
     ShadowBase = getDynamicShadowIfunc(IRB);
+
+  if (!WithFrameRecord && ShadowBase)
     return;
-  }
 
   Value *SlotPtr = getHwasanThreadSlotPtr(IRB, IntptrTy);
   assert(SlotPtr);
 
   Value *ThreadLong = IRB.CreateLoad(IntptrTy, SlotPtr);
   // Extract the address field from ThreadLong. Unnecessary on AArch64 with TBI.
   Value *ThreadLongMaybeUntagged =
       TargetTriple.isAArch64() ? ThreadLong : untagPointer(IRB, ThreadLong);
if (WithFrameRecord) {
     Value *WrapMask = IRB.CreateXor(
         IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
         ConstantInt::get(IntptrTy, (uint64_t)-1));
     Value *ThreadLongNew = IRB.CreateAnd(
         IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 8)), WrapMask);
     IRB.CreateStore(ThreadLongNew, SlotPtr);
   }
 
+  if (!ShadowBase) {
     // Get shadow base address by aligning RecordPtr up.
     // Note: this is not correct if the pointer is already aligned.
     // Runtime library will make sure this never happens.
     ShadowBase = IRB.CreateAdd(
         IRB.CreateOr(
             ThreadLongMaybeUntagged,
             ConstantInt::get(IntptrTy, (1ULL << kShadowBaseAlignment) - 1)),
         ConstantInt::get(IntptrTy, 1), "hwasan.shadow");
     ShadowBase = IRB.CreateIntToPtr(ShadowBase, Int8PtrTy);
   }
+}
 
 Value *HWAddressSanitizer::readRegister(IRBuilder<> &IRB, StringRef Name) {
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
   Function *ReadRegister =
       Intrinsic::getDeclaration(M, Intrinsic::read_register, IntptrTy);
   MDNode *MD = MDNode::get(*C, {MDString::get(*C, Name)});
   Value *Args[] = {MetadataAsValue::get(*C, MD)};
   return IRB.CreateCall(ReadRegister, Args);
if (AllocasToInstrument.empty() && OperandsToInstrument.empty() &&
     return Changed;
 
   assert(!ShadowBase);
 
   Instruction *InsertPt = &*F.getEntryBlock().begin();
   IRBuilder<> EntryIRB(InsertPt);
   emitPrologue(EntryIRB,
                /*WithFrameRecord*/ ClRecordStackHistory &&
-                   !AllocasToInstrument.empty());
+                   Mapping.WithFrameRecord && !AllocasToInstrument.empty());
 
   if (!AllocasToInstrument.empty()) {
     Value *StackTag =
         ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB);
     instrumentStack(AllocasToInstrument, AllocaDbgMap, RetVec, StackTag);
   }
   // Pad and align each of the allocas that we instrumented to stop small
   // uninteresting allocas from hiding in instrumented alloca's padding and so
void HWAddressSanitizer::ShadowMapping::init(Triple &TargetTriple,
                                              bool InstrumentWithCalls) {
   Scale = kDefaultShadowScale;
   if (TargetTriple.isOSFuchsia()) {
     // Fuchsia is always PIE, which means that the beginning of the address
     // space is always available.
     InGlobal = false;
     InTls = false;
     Offset = 0;
+    WithFrameRecord = true;
   } else if (ClMappingOffset.getNumOccurrences() > 0) {
     InGlobal = false;
     InTls = false;
     Offset = ClMappingOffset;
+    WithFrameRecord = false;
   } else if (ClEnableKhwasan || InstrumentWithCalls) {
     InGlobal = false;
     InTls = false;
     Offset = 0;
+    WithFrameRecord = false;
   } else if (ClWithIfunc) {
     InGlobal = true;
     InTls = false;
     Offset = kDynamicShadowSentinel;
+    WithFrameRecord = false;
   } else if (ClWithTls) {
     InGlobal = false;
     InTls = true;
     Offset = kDynamicShadowSentinel;
+    WithFrameRecord = true;
   } else {
     InGlobal = false;
     InTls = false;
     Offset = kDynamicShadowSentinel;
+    WithFrameRecord = false;
   }
 }

llvm/test/Instrumentation/HWAddressSanitizer/prologue.ll

 ; Test -hwasan-with-ifunc flag.
 ;
 ; RUN: opt -hwasan -S < %s | \
-; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOGLOBAL,CHECK-TLS,CHECK-HISTORY
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOGLOBAL,CHECK-TLS-SLOT,CHECK-HISTORY,CHECK-HISTORY-TLS-SLOT
 ; RUN: opt -hwasan -S -hwasan-with-ifunc=0 -hwasan-with-tls=1 -hwasan-record-stack-history=1 < %s | \
-; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOGLOBAL,CHECK-TLS,CHECK-HISTORY
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOGLOBAL,CHECK-TLS-SLOT,CHECK-HISTORY,CHECK-HISTORY-TLS-SLOT
 ; RUN: opt -hwasan -S -hwasan-with-ifunc=0 -hwasan-with-tls=1 -hwasan-record-stack-history=0 < %s | \
 ; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOGLOBAL,CHECK-IFUNC,CHECK-NOHISTORY
 ; RUN: opt -hwasan -S -hwasan-with-ifunc=0 -hwasan-with-tls=0 < %s | \
 ; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-GLOBAL,CHECK-NOHISTORY
 ; RUN: opt -hwasan -S -hwasan-with-ifunc=1  -hwasan-with-tls=0 < %s | \
 ; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-IFUNC,CHECK-NOHISTORY
+; RUN: opt -hwasan -S -mtriple=aarch64-fuchsia < %s | \
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-ZERO-OFFSET,CHECK-SHORT-GRANULES,CHECK-HISTORY,CHECK-HWASAN-TLS,CHECK-HISTORY-HWASAN-TLS
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-android22"
 
 ; CHECK-IFUNC: @__hwasan_shadow = external global [0 x i8]
 ; CHECK-NOIFUNC: @__hwasan_shadow_memory_dynamic_address = external global i64
 
 define i32 @test_load(i32* %a) sanitize_hwaddress {
 ; First instrumentation in the function must be to load the dynamic shadow
 ; address into a local variable.
 ; CHECK-LABEL: @test_load
 ; CHECK: entry:
 
 ; CHECK-NOGLOBAL:   %[[A:[^ ]*]] = call i8* asm "", "=r,0"([0 x i8]* @__hwasan_shadow)
 ; CHECK-NOGLOBAL:   @llvm.hwasan.check.memaccess(i8* %[[A]]
+; CHECK-ZERO-OFFSET:  %[[A:[^ ]*]] = call i8* asm "", "=r,0"(i8* null)
+; CHECK-SHORT-GRANULES:  @llvm.hwasan.check.memaccess.shortgranules(i8* %[[A]]
 
 ; CHECK-GLOBAL: load i8*, i8** @__hwasan_shadow_memory_dynamic_address
 
 ; "store i64" is only used to update stack history (this input IR intentionally does not use any i64)
 ; W/o any allocas, the history is not updated, even if it is enabled explicitly with -hwasan-record-stack-history=1
 ; CHECK-NOT: store i64
 
 ; CHECK: ret i32

 ; CHECK-LABEL: @test_alloca
 ; CHECK: entry:
 
 ; CHECK-IFUNC:   %[[A:[^ ]*]] = call i8* asm "", "=r,0"([0 x i8]* @__hwasan_shadow)
 ; CHECK-IFUNC:   getelementptr i8, i8* %[[A]]
 
 ; CHECK-GLOBAL: load i8*, i8** @__hwasan_shadow_memory_dynamic_address
 
-; CHECK-TLS:   %[[A:[^ ]*]] = call i8* @llvm.thread.pointer()
-; CHECK-TLS:   %[[B:[^ ]*]] = getelementptr i8, i8* %[[A]], i32 48
-; CHECK-TLS:   %[[C:[^ ]*]] = bitcast i8* %[[B]] to i64*
-; CHECK-TLS:   %[[D:[^ ]*]] = load i64, i64* %[[C]]
-; CHECK-TLS:   %[[E:[^ ]*]] = ashr i64 %[[D]], 3
+; CHECK-TLS-SLOT:   %[[A:[^ ]*]] = call i8* @llvm.thread.pointer()
+; CHECK-TLS-SLOT:   %[[B:[^ ]*]] = getelementptr i8, i8* %[[A]], i32 48
+; CHECK-TLS-SLOT:   %[[C:[^ ]*]] = bitcast i8* %[[B]] to i64*
+; CHECK-TLS-SLOT:   %[[D:[^ ]*]] = load i64, i64* %[[C]]
+; CHECK-TLS-SLOT:   %[[E:[^ ]*]] = ashr i64 %[[D]], 3
+; CHECK-HWASAN-TLS: %[[D:[^ ]*]] = load i64, i64* @__hwasan_tls, align 8
+; CHECK-HWASAN-TLS: %[[E:[^ ]*]] = ashr i64 %[[D]], 3
 
 ; CHECK-NOHISTORY-NOT: store i64
 
 ; CHECK-HISTORY: call i64 @llvm.read_register.i64(metadata [[MD:![0-9]*]])
 ; CHECK-HISTORY: %[[PTR:[^ ]*]] = inttoptr i64 %[[D]] to i64*
 ; CHECK-HISTORY: store i64 %{{.*}}, i64* %[[PTR]]
 ; CHECK-HISTORY: %[[D1:[^ ]*]] = ashr i64 %[[D]], 56
 ; CHECK-HISTORY: %[[D2:[^ ]*]] = shl nuw nsw i64 %[[D1]], 12
 ; CHECK-HISTORY: %[[D3:[^ ]*]] = xor i64 %[[D2]], -1
 ; CHECK-HISTORY: %[[D4:[^ ]*]] = add i64 %[[D]], 8
 ; CHECK-HISTORY: %[[D5:[^ ]*]] = and i64 %[[D4]], %[[D3]]
-; CHECK-HISTORY: store i64 %[[D5]], i64* %[[C]]
+; CHECK-HISTORY-TLS-SLOT: store i64 %[[D5]], i64* %[[C]]
+; CHECK-HISTORY-HWASAN-TLS: store i64 %[[D5]], i64* @__hwasan_tls
 
 ; CHECK-TLS:   %[[F:[^ ]*]] = or i64 %[[D]], 4294967295
 ; CHECK-TLS:   = add i64 %[[F]], 1
 
 ; CHECK-HISTORY: = xor i64 %[[E]], 0
 
 ; CHECK-NOHISTORY-NOT: store i64
 
 
 entry:
   %x = alloca i32, align 4
   call void @use(i32* %x)
   ret void
 }
 
 ; CHECK-HISTORY: [[MD]] = !{!"pc"}