hwasan: Compatibility fixes for short granules.

We can't use short granules with stack instrumentation when targeting older
API levels because the rest of the system won't understand the short granule
tags stored in shadow memory.

Moreover, we need to be able to let old binaries (which won't understand
short granule tags) run on a new system that supports short granule
tags. Such binaries will call the __hwasan_tag_mismatch function when their
outlined checks fail. We can compensate for the binary's lack of support
for short granules by implementing the short granule part of the check in
the __hwasan_tag_mismatch function. Unfortunately we can't do anything about
inline checks, but I don't believe that we can generate these by default on
aarch64, nor did we do so when the ABI was fixed.

A new function, __hwasan_tag_mismatch_v2, is introduced that lets code
targeting the new runtime avoid redoing the short granule check. Because tag
mismatches are rare this isn't important from a performance perspective; the
main benefit is that it introduces a symbol dependency that prevents binaries
targeting the new runtime from running on older (i.e. incompatible) runtimes.

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

llvm-svn: 373035

3 files changed, 113 insertions, 80 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index c9ab3e37b57..1c40adc9362 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -99,7 +99,7 @@ public:
   void LowerPATCHABLE_FUNCTION_EXIT(const MachineInstr &MI);
   void LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI);
 
-  std::map<std::pair<unsigned, uint32_t>, MCSymbol *> HwasanMemaccessSymbols;
+  std::map<std::tuple<unsigned, bool, uint32_t>, MCSymbol *> HwasanMemaccessSymbols;
   void LowerHWASAN_CHECK_MEMACCESS(const MachineInstr &MI);
   void EmitHwasanMemaccessSymbols(Module &M);
 
@@ -237,8 +237,10 @@ void AArch64AsmPrinter::EmitSled(const MachineInstr &MI, SledKind Kind)
 
 void AArch64AsmPrinter::LowerHWASAN_CHECK_MEMACCESS(const MachineInstr &MI) {
   Register Reg = MI.getOperand(0).getReg();
+  bool IsShort =
+      MI.getOpcode() == AArch64::HWASAN_CHECK_MEMACCESS_SHORTGRANULES;
   uint32_t AccessInfo = MI.getOperand(1).getImm();
-  MCSymbol *&Sym = HwasanMemaccessSymbols[{Reg, AccessInfo}];
+  MCSymbol *&Sym = HwasanMemaccessSymbols[{Reg, IsShort, AccessInfo}];
   if (!Sym) {
     // FIXME: Make this work on non-ELF.
     if (!TM.getTargetTriple().isOSBinFormatELF())
@@ -246,6 +248,8 @@ void AArch64AsmPrinter::LowerHWASAN_CHECK_MEMACCESS(const MachineInstr &MI) {
 
     std::string SymName = "__hwasan_check_x" + utostr(Reg - AArch64::X0) + "_" +
                           utostr(AccessInfo);
+    if (IsShort)
+      SymName += "_short";
     Sym = OutContext.getOrCreateSymbol(SymName);
   }
 
@@ -263,15 +267,22 @@ void AArch64AsmPrinter::EmitHwasanMemaccessSymbols(Module &M) {
   std::unique_ptr<MCSubtargetInfo> STI(
       TM.getTarget().createMCSubtargetInfo(TT.str(), "", ""));
 
-  MCSymbol *HwasanTagMismatchSym =
+  MCSymbol *HwasanTagMismatchV1Sym =
       OutContext.getOrCreateSymbol("__hwasan_tag_mismatch");
+  MCSymbol *HwasanTagMismatchV2Sym =
+      OutContext.getOrCreateSymbol("__hwasan_tag_mismatch_v2");
 
-  const MCSymbolRefExpr *HwasanTagMismatchRef =
-      MCSymbolRefExpr::create(HwasanTagMismatchSym, OutContext);
+  const MCSymbolRefExpr *HwasanTagMismatchV1Ref =
+      MCSymbolRefExpr::create(HwasanTagMismatchV1Sym, OutContext);
+  const MCSymbolRefExpr *HwasanTagMismatchV2Ref =
+      MCSymbolRefExpr::create(HwasanTagMismatchV2Sym, OutContext);
 
   for (auto &P : HwasanMemaccessSymbols) {
-    unsigned Reg = P.first.first;
-    uint32_t AccessInfo = P.first.second;
+    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;
 
     OutStreamer->SwitchSection(OutContext.getELFSection(
@@ -304,82 +315,86 @@ void AArch64AsmPrinter::EmitHwasanMemaccessSymbols(Module &M) {
             .addReg(Reg)
             .addImm(AArch64_AM::getShifterImm(AArch64_AM::LSR, 56)),
         *STI);
-    MCSymbol *HandlePartialSym = OutContext.createTempSymbol();
+    MCSymbol *HandleMismatchOrPartialSym = OutContext.createTempSymbol();
     OutStreamer->EmitInstruction(
         MCInstBuilder(AArch64::Bcc)
             .addImm(AArch64CC::NE)
-            .addExpr(MCSymbolRefExpr::create(HandlePartialSym, OutContext)),
+            .addExpr(MCSymbolRefExpr::create(HandleMismatchOrPartialSym,
+                                             OutContext)),
         *STI);
     MCSymbol *ReturnSym = OutContext.createTempSymbol();
     OutStreamer->EmitLabel(ReturnSym);
     OutStreamer->EmitInstruction(
         MCInstBuilder(AArch64::RET).addReg(AArch64::LR), *STI);
+    OutStreamer->EmitLabel(HandleMismatchOrPartialSym);
 
-    OutStreamer->EmitLabel(HandlePartialSym);
-    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)
+    if (IsShort) {
+      OutStreamer->EmitInstruction(MCInstBuilder(AArch64::SUBSWri)
+                                       .addReg(AArch64::WZR)
+                                       .addReg(AArch64::W16)
+                                       .addImm(15)
                                        .addImm(0),
                                    *STI);
-    OutStreamer->EmitInstruction(MCInstBuilder(AArch64::SUBSWrs)
-                                     .addReg(AArch64::WZR)
-                                     .addReg(AArch64::W16)
-                                     .addReg(AArch64::W17)
-                                     .addImm(0),
-                                 *STI);
-    OutStreamer->EmitInstruction(
-        MCInstBuilder(AArch64::Bcc)
-            .addImm(AArch64CC::LS)
-            .addExpr(MCSymbolRefExpr::create(HandleMismatchSym, OutContext)),
-        *STI);
-
-    OutStreamer->EmitInstruction(
-        MCInstBuilder(AArch64::ORRXri)
-            .addReg(AArch64::X16)
-            .addReg(Reg)
-            .addImm(AArch64_AM::encodeLogicalImmediate(0xf, 64)),
-        *STI);
-    OutStreamer->EmitInstruction(MCInstBuilder(AArch64::LDRBBui)
-                                     .addReg(AArch64::W16)
-                                     .addReg(AArch64::X16)
-                                     .addImm(0),
-                                 *STI);
-    OutStreamer->EmitInstruction(
-        MCInstBuilder(AArch64::SUBSXrs)
-            .addReg(AArch64::XZR)
-            .addReg(AArch64::X16)
-            .addReg(Reg)
-            .addImm(AArch64_AM::getShifterImm(AArch64_AM::LSR, 56)),
-        *STI);
-    OutStreamer->EmitInstruction(
-        MCInstBuilder(AArch64::Bcc)
-            .addImm(AArch64CC::EQ)
-            .addExpr(MCSymbolRefExpr::create(ReturnSym, OutContext)),
-        *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)
+                                       .addReg(AArch64::WZR)
+                                       .addReg(AArch64::W16)
+                                       .addReg(AArch64::W17)
+                                       .addImm(0),
+                                   *STI);
+      OutStreamer->EmitInstruction(
+          MCInstBuilder(AArch64::Bcc)
+              .addImm(AArch64CC::LS)
+              .addExpr(MCSymbolRefExpr::create(HandleMismatchSym, OutContext)),
+          *STI);
+
+      OutStreamer->EmitInstruction(
+          MCInstBuilder(AArch64::ORRXri)
+              .addReg(AArch64::X16)
+              .addReg(Reg)
+              .addImm(AArch64_AM::encodeLogicalImmediate(0xf, 64)),
+          *STI);
+      OutStreamer->EmitInstruction(MCInstBuilder(AArch64::LDRBBui)
+                                       .addReg(AArch64::W16)
+                                       .addReg(AArch64::X16)
+                                       .addImm(0),
+                                   *STI);
+      OutStreamer->EmitInstruction(
+          MCInstBuilder(AArch64::SUBSXrs)
+              .addReg(AArch64::XZR)
+              .addReg(AArch64::X16)
+              .addReg(Reg)
+              .addImm(AArch64_AM::getShifterImm(AArch64_AM::LSR, 56)),
+          *STI);
+      OutStreamer->EmitInstruction(
+          MCInstBuilder(AArch64::Bcc)
+              .addImm(AArch64CC::EQ)
+              .addExpr(MCSymbolRefExpr::create(ReturnSym, OutContext)),
+          *STI);
+
+      OutStreamer->EmitLabel(HandleMismatchSym);
+    }
 
-    OutStreamer->EmitLabel(HandleMismatchSym);
     OutStreamer->EmitInstruction(MCInstBuilder(AArch64::STPXpre)
                                      .addReg(AArch64::SP)
                                      .addReg(AArch64::X0)
@@ -414,16 +429,16 @@ void AArch64AsmPrinter::EmitHwasanMemaccessSymbols(Module &M) {
         MCInstBuilder(AArch64::ADRP)
             .addReg(AArch64::X16)
             .addExpr(AArch64MCExpr::create(
-                HwasanTagMismatchRef,
-                AArch64MCExpr::VariantKind::VK_GOT_PAGE, OutContext)),
+                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)),
+                HwasanTagMismatchRef, AArch64MCExpr::VariantKind::VK_GOT_LO12,
+                OutContext)),
         *STI);
     OutStreamer->EmitInstruction(
         MCInstBuilder(AArch64::BR).addReg(AArch64::X16), *STI);
@@ -1096,6 +1111,7 @@ void AArch64AsmPrinter::EmitInstruction(const MachineInstr *MI) {
     return;
 
   case AArch64::HWASAN_CHECK_MEMACCESS:
+  case AArch64::HWASAN_CHECK_MEMACCESS_SHORTGRANULES:
     LowerHWASAN_CHECK_MEMACCESS(*MI);
     return;
 
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index ce6d032481c..661d089688b 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -800,6 +800,10 @@ def HWASAN_CHECK_MEMACCESS : Pseudo<
   (outs), (ins GPR64noip:$ptr, i32imm:$accessinfo),
   [(int_hwasan_check_memaccess X9, GPR64noip:$ptr, (i32 timm:$accessinfo))]>,
   Sched<[]>;
+def HWASAN_CHECK_MEMACCESS_SHORTGRANULES : Pseudo<
+  (outs), (ins GPR64noip:$ptr, i32imm:$accessinfo),
+  [(int_hwasan_check_memaccess_shortgranules X9, GPR64noip:$ptr, (i32 timm:$accessinfo))]>,
+  Sched<[]>;
 }
 
 // The cycle counter PMC register is PMCCNTR_EL0.
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index df7606dfa24..2bc127bc982 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -166,6 +166,11 @@ static cl::opt<bool>
                             cl::desc("instrument landing pads"), cl::Hidden,
                             cl::init(false), cl::ZeroOrMore);
 
+static cl::opt<bool> ClUseShortGranules(
+    "hwasan-use-short-granules",
+    cl::desc("use short granules in allocas and outlined checks"), cl::Hidden,
+    cl::init(false), cl::ZeroOrMore);
+
 static cl::opt<bool> ClInstrumentPersonalityFunctions(
     "hwasan-instrument-personality-functions",
     cl::desc("instrument personality functions"), cl::Hidden, cl::init(false),
@@ -269,6 +274,7 @@ private:
 
   bool CompileKernel;
   bool Recover;
+  bool UseShortGranules;
   bool InstrumentLandingPads;
 
   Function *HwasanCtorFunction;
@@ -372,11 +378,14 @@ void HWAddressSanitizer::initializeModule() {
   HwasanCtorFunction = nullptr;
 
   // Older versions of Android do not have the required runtime support for
-  // global or personality function instrumentation. On other platforms we
-  // currently require using the latest version of the runtime.
+  // 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;
+
   // If we don't have personality function support, fall back to landing pads.
   InstrumentLandingPads = ClInstrumentLandingPads.getNumOccurrences()
                               ? ClInstrumentLandingPads
@@ -608,9 +617,11 @@ void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,
       TargetTriple.isOSBinFormatELF() && !Recover) {
     Module *M = IRB.GetInsertBlock()->getParent()->getParent();
     Ptr = IRB.CreateBitCast(Ptr, Int8PtrTy);
-    IRB.CreateCall(
-        Intrinsic::getDeclaration(M, Intrinsic::hwasan_check_memaccess),
-        {shadowBase(), Ptr, ConstantInt::get(Int32Ty, AccessInfo)});
+    IRB.CreateCall(Intrinsic::getDeclaration(
+                       M, UseShortGranules
+                              ? Intrinsic::hwasan_check_memaccess_shortgranules
+                              : Intrinsic::hwasan_check_memaccess),
+                   {shadowBase(), Ptr, ConstantInt::get(Int32Ty, AccessInfo)});
     return;
   }
 
@@ -763,6 +774,8 @@ static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
 bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI,
                                    Value *Tag, size_t Size) {
   size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
+  if (!UseShortGranules)
+    Size = AlignedSize;
 
   Value *JustTag = IRB.CreateTrunc(Tag, IRB.getInt8Ty());
   if (ClInstrumentWithCalls) {