hwasan: Support for outlined checks in the Linux kernel.

[pcc]

Summary

Add support for match-all tags and GOT-free runtime calls, which
are both required for the kernel to be able to support outlined
checks. This requires extending the access info to let the backend
know when to enable these features. To make the code easier to maintain
introduce an enum with the bit field positions for the access info.

Allow outlined checks to be enabled with -mllvm
-hwasan-inline-all-checks=0. Kernels that contain runtime support for
outlined checks may pass this flag. Kernels lacking runtime support
will continue to link because they do not pass the flag. Old versions
of LLVM will ignore the flag and continue to use inline checks.

With a separate kernel patch [1] I measured the code size of defconfig
+ tag-based KASAN, as well as boot time (i.e. time to init launch)
on a DragonBoard 845c with an Android arm64 GKI kernel. The results
are below:

code size    boot time

before 92824064 6.18s
after 38822400 6.65s

[1] https://linux-review.googlesource.com/id/I1a30036c70ab3c3ee78d75ed9b87ef7cdc3fdb76

Depends on D90425

[pcc]

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

361

This was meant to stay where it was (I had originally moved it here because I had the match-all check first, but benchmarking revealed that moving the match-all check after the shadow check was faster). I'll move it back in the next upload.

[pcc]

Revert unnecessary change, fix tests

[eugenis]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

505

I don't understand the impact of this diff.
Now, Recover==true no longer overrides an explicit -hwasan-inline-all-checks=1.
Why?
Is that what prevented outlining in the kernel before?

510

ClMatchTag & 0xF
either here or when constructing AccessInfo

[hctim]

llvm/include/llvm/Transforms/Instrumentation/HWAddressSanitizer.h

40

why not enum HWASanAccessInfo { ...?

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

396

IIUC, this is being optimised to

lsr x16, x1, #56
cmp x16, MatchAllTag
b.eq .ret

and this seems clearer to me. Is it possible to update the codegen here to be the same?

513

B instead of BR implies that the kernel will never return from __hwasan_tag_mismatch (which is still reachable as it looks like it's still possible for the kernel to compile without short granules). Does this assumption always hold?

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

796–798

can we declare this as a const kAccessInfoRuntimeMask (and elsewhere)?

[pcc]

llvm/include/llvm/Transforms/Instrumentation/HWAddressSanitizer.h

40

Because we don't want these to go straight into the llvm namespace. enum class avoids this but then the enumerators would be more awkward to work with in this context (e.g. need to be static_casted into integers).

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

396

There is no AArch64::LSR or AArch64::CMP. When the assembler sees one of these it rewrites the instruction into the more general form (i.e. UBFM, SUBS). So we can't make it the same.

513

B instead of BR implies that the kernel will never return from __hwasan_tag_mismatch

No it doesn't. With either B or BR x30 is not overwritten so a RET in __hwasan_tag_mismatch would go straight back to the caller of __hwasan_check_*. This is how tail calls work for example.

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

505

Yes, outlining was being prevented in the kernel because it always had Recover enabled (see addKernelHWAddressSanitizerPasses in clang). What this diff lets you do is say -hwasan-inline-all-checks=0 and as a result you are promising that your runtime supports recovery from outlined checks.

510

You mean ClMatchAllTag & 0xFF, right? This isn't MTE :)

It is unnecessary because the field is a uint8_t but I suppose we could avoid an implicit truncation that way.

796–798

Okay, I'll add it to HWAddressSanitizer.h.

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

510

Right, I missed that. I was concerned about overwriting other fields of AccessInfo if the tag value is too large.

llvm/include/llvm/Transforms/Instrumentation/HWAddressSanitizer.h

 private:
   bool CompileKernel;
   bool Recover;
 };
 
 FunctionPass *createHWAddressSanitizerLegacyPassPass(bool CompileKernel = false,
                                                      bool Recover = false);
 
+namespace HWASanAccessInfo {
+
+// Bit field positions for the accessinfo parameter to
+// llvm.hwasan.check.memaccess. Shared between the pass and the backend. Bits
+// 0-15 are also used by the runtime.
+enum {
+  AccessSizeShift = 0, // 4 bits
+  IsWriteShift = 4,
+  RecoverShift = 5,
+  MatchAllShift = 16, // 8 bits
+  HasMatchAllShift = 24,
+  CompileKernelShift = 25,
+};
+
+enum { RuntimeMask = 0xffff };
+
+} // namespace HWASanAccessInfo
+
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

 #include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <map>
 #include <memory>
 
 using namespace llvm;
 
void AArch64AsmPrinter::EmitHwasanMemaccessSymbols(Module &M) {
   for (auto &P : HwasanMemaccessSymbols) {
     unsigned Reg = std::get<0>(P.first);
     bool IsShort = std::get<1>(P.first);
     uint32_t AccessInfo = std::get<2>(P.first);
     const MCSymbolRefExpr *HwasanTagMismatchRef =
         IsShort ? HwasanTagMismatchV2Ref : HwasanTagMismatchV1Ref;
     MCSymbol *Sym = P.second;
 
+    bool HasMatchAllTag =
+        (AccessInfo >> HWASanAccessInfo::HasMatchAllShift) & 1;
+    uint8_t MatchAllTag =
+        (AccessInfo >> HWASanAccessInfo::MatchAllShift) & 0xff;
+    unsigned Size =
+        1 << ((AccessInfo >> HWASanAccessInfo::AccessSizeShift) & 0xf);
+    bool CompileKernel =
+        (AccessInfo >> HWASanAccessInfo::CompileKernelShift) & 1;
+
     OutStreamer->SwitchSection(OutContext.getELFSection(
         ".text.hot", ELF::SHT_PROGBITS,
         ELF::SHF_EXECINSTR | ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
         Sym->getName()));
 
     OutStreamer->emitSymbolAttribute(Sym, MCSA_ELF_TypeFunction);
     OutStreamer->emitSymbolAttribute(Sym, MCSA_Weak);
     OutStreamer->emitSymbolAttribute(Sym, MCSA_Hidden);
     OutStreamer->emitLabel(Sym);
 
     OutStreamer->emitInstruction(MCInstBuilder(AArch64::SBFMXri)
                                      .addReg(AArch64::X16)
                                      .addReg(Reg)
                                      .addImm(4)
                                      .addImm(55),
                                  *STI);
     OutStreamer->emitInstruction(
         MCInstBuilder(AArch64::LDRBBroX)
             .addReg(AArch64::W16)
OutStreamer->emitInstruction(
                                              OutContext)),
         *STI);
     MCSymbol *ReturnSym = OutContext.createTempSymbol();
     OutStreamer->emitLabel(ReturnSym);
     OutStreamer->emitInstruction(
         MCInstBuilder(AArch64::RET).addReg(AArch64::LR), *STI);
     OutStreamer->emitLabel(HandleMismatchOrPartialSym);
 
+    if (HasMatchAllTag) {
+      OutStreamer->emitInstruction(MCInstBuilder(AArch64::UBFMXri)
+                                       .addReg(AArch64::X16)
+                                       .addReg(Reg)
+                                       .addImm(56)
+                                       .addImm(63),
+                                   *STI);
+      OutStreamer->emitInstruction(MCInstBuilder(AArch64::SUBSXri)
+                                       .addReg(AArch64::XZR)
+                                       .addReg(AArch64::X16)
+                                       .addImm(MatchAllTag)
+                                       .addImm(0),
+                                   *STI);
+      OutStreamer->emitInstruction(
+          MCInstBuilder(AArch64::Bcc)
+              .addImm(AArch64CC::EQ)
+              .addExpr(MCSymbolRefExpr::create(ReturnSym, OutContext)),
+          *STI);
+    }
+
     if (IsShort) {
       OutStreamer->emitInstruction(MCInstBuilder(AArch64::SUBSWri)
                                        .addReg(AArch64::WZR)
                                        .addReg(AArch64::W16)
                                        .addImm(15)
                                        .addImm(0),
                                    *STI);
       MCSymbol *HandleMismatchSym = OutContext.createTempSymbol();
       OutStreamer->emitInstruction(
           MCInstBuilder(AArch64::Bcc)
               .addImm(AArch64CC::HI)
               .addExpr(MCSymbolRefExpr::create(HandleMismatchSym, OutContext)),
           *STI);
 
       OutStreamer->emitInstruction(
           MCInstBuilder(AArch64::ANDXri)
               .addReg(AArch64::X17)
               .addReg(Reg)
               .addImm(AArch64_AM::encodeLogicalImmediate(0xf, 64)),
           *STI);
-      unsigned Size = 1 << (AccessInfo & 0xf);
       if (Size != 1)
         OutStreamer->emitInstruction(MCInstBuilder(AArch64::ADDXri)
                                          .addReg(AArch64::X17)
                                          .addReg(AArch64::X17)
                                          .addImm(Size - 1)
                                          .addImm(0),
                                      *STI);
       OutStreamer->emitInstruction(MCInstBuilder(AArch64::SUBSWrs)
for (auto &P : HwasanMemaccessSymbols) {
 
     if (Reg != AArch64::X0)
       OutStreamer->emitInstruction(MCInstBuilder(AArch64::ORRXrs)
                                        .addReg(AArch64::X0)
                                        .addReg(AArch64::XZR)
                                        .addReg(Reg)
                                        .addImm(0),
                                    *STI);
-    OutStreamer->emitInstruction(MCInstBuilder(AArch64::MOVZXi)
-                                     .addReg(AArch64::X1)
-                                     .addImm(AccessInfo)
-                                     .addImm(0),
-                                 *STI);
+    OutStreamer->emitInstruction(
+        MCInstBuilder(AArch64::MOVZXi)
+            .addImm(AccessInfo & HWASanAccessInfo::RuntimeMask)
 
+    if (CompileKernel) {
+      // The Linux kernel's dynamic loader doesn't support GOT relative
+      // relocations, but it doesn't support late binding either, so just call
+      // the function directly.
+      OutStreamer->emitInstruction(
+          MCInstBuilder(AArch64::B).addExpr(HwasanTagMismatchRef), *STI);
+    } else {
       // Intentionally load the GOT entry and branch to it, rather than possibly
-    // late binding the function, which may clobber the registers before we have
-    // a chance to save them.
+      // late binding the function, which may clobber the registers before we
+      // have a chance to save them.
       OutStreamer->emitInstruction(
           MCInstBuilder(AArch64::ADRP)
               .addReg(AArch64::X16)
               .addExpr(AArch64MCExpr::create(
                   HwasanTagMismatchRef, AArch64MCExpr::VariantKind::VK_GOT_PAGE,
                   OutContext)),
           *STI);
       OutStreamer->emitInstruction(
           MCInstBuilder(AArch64::LDRXui)
               .addReg(AArch64::X16)
               .addReg(AArch64::X16)
               .addExpr(AArch64MCExpr::create(
                   HwasanTagMismatchRef, AArch64MCExpr::VariantKind::VK_GOT_LO12,
                   OutContext)),
           *STI);
       OutStreamer->emitInstruction(
           MCInstBuilder(AArch64::BR).addReg(AArch64::X16), *STI);
     }
   }
+}
 
 void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
   EmitHwasanMemaccessSymbols(M);
 
   const Triple &TT = TM.getTargetTriple();
   if (TT.isOSBinFormatMachO()) {
     // Funny Darwin hack: This flag tells the linker that no global symbols
     // contain code that falls through to other global symbols (e.g. the obvious

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

private:
   Type *IntptrTy;
   Type *Int8PtrTy;
   Type *Int8Ty;
   Type *Int32Ty;
   Type *Int64Ty = Type::getInt64Ty(M.getContext());
 
   bool CompileKernel;
   bool Recover;
+  bool OutlinedChecks;
   bool UseShortGranules;
   bool InstrumentLandingPads;
 
+  bool HasMatchAllTag = false;
+  uint8_t MatchAllTag = 0;
+
   Function *HwasanCtorFunction;
 
   FunctionCallee HwasanMemoryAccessCallback[2][kNumberOfAccessSizes];
   FunctionCallee HwasanMemoryAccessCallbackSized[2];
 
   FunctionCallee HwasanTagMemoryFunc;
   FunctionCallee HwasanGenerateTagFunc;
 
void HWAddressSanitizer::initializeModule() {
   // Older versions of Android do not have the required runtime support for
   // short granules, global or personality function instrumentation. On other
   // platforms we currently require using the latest version of the runtime.
   bool NewRuntime =
       !TargetTriple.isAndroid() || !TargetTriple.isAndroidVersionLT(30);
 
   UseShortGranules =
       ClUseShortGranules.getNumOccurrences() ? ClUseShortGranules : NewRuntime;
+  OutlinedChecks =
+      TargetTriple.isAArch64() && TargetTriple.isOSBinFormatELF() &&
+      (ClInlineAllChecks.getNumOccurrences() ? !ClInlineAllChecks : !Recover);
+
+  if (ClMatchAllTag.getNumOccurrences()) {
+    if (ClMatchAllTag != -1) {
+      HasMatchAllTag = true;
+      MatchAllTag = ClMatchAllTag & 0xFF;
+    }
+  } else if (CompileKernel) {
+    HasMatchAllTag = true;
+    MatchAllTag = 0xFF;
+  }
 
   // If we don't have personality function support, fall back to landing pads.
   InstrumentLandingPads = ClInstrumentLandingPads.getNumOccurrences()
                               ? ClInstrumentLandingPads
                               : !NewRuntime;
 
   if (!CompileKernel) {
     createHwasanCtorComdat();
if (Mapping.Offset == 0)
     return IRB.CreateIntToPtr(Shadow, Int8PtrTy);
   // (Mem >> Scale) + Offset
   return IRB.CreateGEP(Int8Ty, ShadowBase, Shadow);
 }
 
 void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,
                                                    unsigned AccessSizeIndex,
                                                    Instruction *InsertBefore) {
-  const int64_t AccessInfo = Recover * 0x20 + IsWrite * 0x10 + AccessSizeIndex;
+  const int64_t AccessInfo =
+      (CompileKernel << HWASanAccessInfo::CompileKernelShift) +
+      (HasMatchAllTag << HWASanAccessInfo::HasMatchAllShift) +
+      (MatchAllTag << HWASanAccessInfo::MatchAllShift) +
+      (Recover << HWASanAccessInfo::RecoverShift) +
+      (IsWrite << HWASanAccessInfo::IsWriteShift) +
+      (AccessSizeIndex << HWASanAccessInfo::AccessSizeShift);
   IRBuilder<> IRB(InsertBefore);
 
-  if (!ClInlineAllChecks && TargetTriple.isAArch64() &&
-      TargetTriple.isOSBinFormatELF() && !Recover) {
+  if (OutlinedChecks) {
     Module *M = IRB.GetInsertBlock()->getParent()->getParent();
     Ptr = IRB.CreateBitCast(Ptr, Int8PtrTy);
     IRB.CreateCall(Intrinsic::getDeclaration(
                        M, UseShortGranules
                               ? Intrinsic::hwasan_check_memaccess_shortgranules
                               : Intrinsic::hwasan_check_memaccess),
                    {ShadowBase, Ptr, ConstantInt::get(Int32Ty, AccessInfo)});
     return;
   }
 
   Value *PtrLong = IRB.CreatePointerCast(Ptr, IntptrTy);
   Value *PtrTag = IRB.CreateTrunc(IRB.CreateLShr(PtrLong, kPointerTagShift),
                                   IRB.getInt8Ty());
   Value *AddrLong = untagPointer(IRB, PtrLong);
   Value *Shadow = memToShadow(AddrLong, IRB);
   Value *MemTag = IRB.CreateLoad(Int8Ty, Shadow);
   Value *TagMismatch = IRB.CreateICmpNE(PtrTag, MemTag);
 
-  int matchAllTag = ClMatchAllTag.getNumOccurrences() > 0 ?
-      ClMatchAllTag : (CompileKernel ? 0xFF : -1);
-  if (matchAllTag != -1) {
-    Value *TagNotIgnored = IRB.CreateICmpNE(PtrTag,
-        ConstantInt::get(PtrTag->getType(), matchAllTag));
+  if (HasMatchAllTag) {
+    Value *TagNotIgnored = IRB.CreateICmpNE(
+        PtrTag, ConstantInt::get(PtrTag->getType(), MatchAllTag));
     TagMismatch = IRB.CreateAnd(TagMismatch, TagNotIgnored);
   }
 
   Instruction *CheckTerm =
       SplitBlockAndInsertIfThen(TagMismatch, InsertBefore, false,
                                 MDBuilder(*C).createBranchWeights(1, 100000));
 
   IRB.SetInsertPoint(CheckTerm);
void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,
 
   IRB.SetInsertPoint(CheckFailTerm);
   InlineAsm *Asm;
   switch (TargetTriple.getArch()) {
     case Triple::x86_64:
       // The signal handler will find the data address in rdi.
       Asm = InlineAsm::get(
           FunctionType::get(IRB.getVoidTy(), {PtrLong->getType()}, false),
-          "int3\nnopl " + itostr(0x40 + AccessInfo) + "(%rax)",
+          "int3\nnopl " +
+              itostr(0x40 + (AccessInfo & HWASanAccessInfo::RuntimeMask)) +
+              "(%rax)",
           "{rdi}",
           /*hasSideEffects=*/true);
       break;
     case Triple::aarch64:
     case Triple::aarch64_be:
       // The signal handler will find the data address in x0.
       Asm = InlineAsm::get(
           FunctionType::get(IRB.getVoidTy(), {PtrLong->getType()}, false),
-          "brk #" + itostr(0x900 + AccessInfo),
+          "brk #" +
+              itostr(0x900 + (AccessInfo & HWASanAccessInfo::RuntimeMask)),
           "{x0}",
           /*hasSideEffects=*/true);
       break;
     default:
       report_fatal_error("unsupported architecture");
   }
   IRB.CreateCall(Asm, PtrLong);
   if (Recover)

llvm/test/CodeGen/AArch64/hwasan-check-memaccess.ll

define i8* @f2(i8* %x0, i8* %x1) {
   ; CHECK-NEXT: mov x20, x1
   ; CHECK-NEXT: bl __hwasan_check_x0_2_short_v2
   ; CHECK-NEXT: ldp x30, x20, [sp], #16
   ; CHECK-NEXT: ret
   call void @llvm.hwasan.check.memaccess.shortgranules(i8* %x1, i8* %x0, i32 2)
   ret i8* %x0
 }
 
+define void @f3(i8* %x0, i8* %x1) {
+  ; 0x3ff0000 (kernel, match-all = 0xff)
+  call void @llvm.hwasan.check.memaccess(i8* %x0, i8* %x1, i32 67043328)
+  ret void
+}
+
 declare void @llvm.hwasan.check.memaccess(i8*, i8*, i32)
 declare void @llvm.hwasan.check.memaccess.shortgranules(i8*, i8*, i32)
 
 ; CHECK:      .section .text.hot,"axG",@progbits,__hwasan_check_x0_2_short_v2,comdat
 ; CHECK-NEXT: .type __hwasan_check_x0_2_short_v2,@function
 ; CHECK-NEXT: .weak __hwasan_check_x0_2_short_v2
 ; CHECK-NEXT: .hidden __hwasan_check_x0_2_short_v2
 ; CHECK-NEXT: __hwasan_check_x0_2_short_v2:

 ; CHECK-NEXT: .Ltmp3:
 ; CHECK-NEXT: stp x0, x1, [sp, #-256]!
 ; CHECK-NEXT: stp x29, x30, [sp, #232]
 ; CHECK-NEXT: mov x0, x1
 ; CHECK-NEXT: mov x1, #1
 ; CHECK-NEXT: adrp  x16, :got:__hwasan_tag_mismatch
 ; CHECK-NEXT: ldr x16, [x16, :got_lo12:__hwasan_tag_mismatch]
 ; CHECK-NEXT: br  x16
+
+; CHECK:      __hwasan_check_x1_67043328:
+; CHECK-NEXT: sbfx x16, x1, #4, #52
+; CHECK-NEXT: ldrb w16, [x9, x16]
+; CHECK-NEXT: cmp x16, x1, lsr #56
+; CHECK-NEXT: b.ne .Ltmp5
+; CHECK-NEXT: .Ltmp6:
+; CHECK-NEXT: ret
+; CHECK-NEXT: .Ltmp5:
+; CHECK-NEXT: lsr x16, x1, #56
+; CHECK-NEXT: cmp x16, #255
+; CHECK-NEXT: b.eq .Ltmp6
+; CHECK-NEXT: stp x0, x1, [sp, #-256]!
+; CHECK-NEXT: stp x29, x30, [sp, #232]
+; CHECK-NEXT: mov x0, x1
+; CHECK-NEXT: mov x1, #0
+; CHECK-NEXT: b __hwasan_tag_mismatch

llvm/test/Instrumentation/HWAddressSanitizer/kernel.ll

 ; Test KHWASan instrumentation.
 ;
 ; RUN: opt < %s -hwasan -hwasan-kernel=1 -hwasan-recover=1 -S | FileCheck %s --allow-empty --check-prefixes=INIT
 ; RUN: opt < %s -hwasan -hwasan-kernel=1 -hwasan-recover=1 -S | FileCheck %s  --check-prefixes=CHECK,NOOFFSET,MATCH-ALL
 ; RUN: opt < %s -hwasan -hwasan-kernel=1 -hwasan-recover=1 -hwasan-mapping-offset=12345678 -S | FileCheck %s  --check-prefixes=CHECK,OFFSET,MATCH-ALL
 ; RUN: opt < %s -hwasan -hwasan-kernel=1 -hwasan-recover=1 -hwasan-match-all-tag=-1 -S | FileCheck %s  --check-prefixes=CHECK,NOOFFSET,NO-MATCH-ALL
+; RUN: opt < %s -hwasan -hwasan-kernel=1 -hwasan-recover=1 -hwasan-inline-all-checks=0 -hwasan-mapping-offset=12345678 -S | FileCheck %s  --check-prefixes=OUTLINE
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-android"
 
 define i8 @test_load(i8* %a) sanitize_hwaddress {
 ; CHECK-LABEL: @test_load(
 ; OFFSET: %[[SHADOW:[^ ]*]] = call i8* asm "", "=r,0"(i8* inttoptr (i64 12345678 to i8*))
 ; CHECK: %[[A:[^ ]*]] = ptrtoint i8* %a to i64

 ; NO-MATCH-ALL: br i1 %[[F]], label {{.*}}, label {{.*}}, !prof {{.*}}
 
 ; CHECK: call void asm sideeffect "brk #2336", "{x0}"(i64 %[[A]])
 ; CHECK: br label
 
 ; CHECK: %[[G:[^ ]*]] = load i8, i8* %a, align 4
 ; CHECK: ret i8 %[[G]]
 
+; OUTLINE: %[[SHADOW:[^ ]*]] = call i8* asm "", "=r,0"(i8* inttoptr (i64 12345678 to i8*))
+; OUTLINE: call void @llvm.hwasan.check.memaccess(i8* %[[SHADOW]], i8* %a, i32 67043360)
 entry:
   %b = load i8, i8* %a, align 4
   ret i8 %b
 }
 
 ; INIT-NOT: call void @__hwasan_init