[hwasan] Fix kernel instrumentation of stack.

[eugenis]

Summary

Kernel addresses have 0xFF in the most significant byte.
A tag can not be pushed there with OR (tag << 56);
use AND ((tag << 56) | 0x00FF..FF) instead.

[kcc]

LGTM with a nit

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

446

please add a comment, similar to the commit message.

[andreyknvl]

Thanks!

[andreyknvl]

I'll send my patch with more khwasan related changes, once this is committed.

[eugenis]

add a comment

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

void instrumentMemAccessInline(Value *PtrLong, bool IsWrite,
                                  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);
+  Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   bool instrumentStack(SmallVectorImpl<AllocaInst *> &Allocas,
                        SmallVectorImpl<Instruction *> &RetVec);
   Value *getNextTagWithCall(IRBuilder<> &IRB);
   Value *getStackBaseTag(IRBuilder<> &IRB);
   Value *getAllocaTag(IRBuilder<> &IRB, Value *StackTag, AllocaInst *AI,
                      unsigned AllocaNo);
   Value *getUARTag(IRBuilder<> &IRB, Value *StackTag);
 

 }
 
 Value *HWAddressSanitizer::getUARTag(IRBuilder<> &IRB, Value *StackTag) {
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
   return IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, 0xFFU));
 }
 
+Value *HWAddressSanitizer::tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong,
+                                      Value *Tag) {
+  Value *TaggedPtrLong;
+  if (ClEnableKhwasan) {
+    Value *ShiftedTag = IRB.CreateOr(
+        IRB.CreateShl(Tag, kPointerTagShift),
+        ConstantInt::get(IntptrTy, (1ULL << kPointerTagShift) - 1));
+    TaggedPtrLong = IRB.CreateAnd(PtrLong, ShiftedTag);
+  } else {
+    Value *ShiftedTag = IRB.CreateShl(Tag, kPointerTagShift);
+    TaggedPtrLong = IRB.CreateOr(PtrLong, ShiftedTag);
+  }
+  return IRB.CreateIntToPtr(TaggedPtrLong, Ty);
+}
+
 bool HWAddressSanitizer::instrumentStack(
     SmallVectorImpl<AllocaInst *> &Allocas,
     SmallVectorImpl<Instruction *> &RetVec) {
   Function *F = Allocas[0]->getParent()->getParent();
   Instruction *InsertPt = &*F->getEntryBlock().begin();
   IRBuilder<> IRB(InsertPt);
 
   Value *StackTag = getStackBaseTag(IRB);
 
   // 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.
     Value *Tag = getAllocaTag(IRB, StackTag, AI, N);
     Value *AILong = IRB.CreatePointerCast(AI, IntptrTy);
+    Value *Replacement = tagPointer(IRB, AI->getType(), AILong, Tag);
     std::string Name =
         AI->hasName() ? AI->getName().str() : "alloca." + itostr(N);
-    Value *Replacement = IRB.CreateIntToPtr(
-        IRB.CreateOr(AILong, IRB.CreateShl(Tag, kPointerTagShift)),
-        AI->getType(), Name + ".hwasan");
+    Replacement->setName(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);

llvm/test/Instrumentation/HWAddressSanitizer/kernel-alloca.ll

+; Test basic address sanitizer instrumentation.
+;
+; RUN: opt < %s -hwasan -hwasan-kernel=1 -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 %[[C]], 72057594037927935
+; CHECK: %[[E:[^ ]*]] = and i64 %[[X1]], %[[D]]
+; CHECK: %[[X_HWASAN:[^ ]*]] = inttoptr i64 %[[E]] to i32*
+
+entry:
+  %x = alloca i32, align 4
+  call void @use32(i32* nonnull %x)
+  ret void
+}