[AArch64][PAC][ELF] Lower constant and code references

Author: kovdan01

llvm/include/llvm/CodeGen/MachineModuleInfoImpls.h

@@ -83,10 +83,18 @@ class MachineModuleInfoMachO : public MachineModuleInfoImpl {
 /// MachineModuleInfoELF - This is a MachineModuleInfoImpl implementation
 /// for ELF targets.
 class MachineModuleInfoELF : public MachineModuleInfoImpl {
+public:
+  struct AuthStubInfo {
+    const MCExpr *AuthPtrRef;
+  };
+
+private:
   /// GVStubs - These stubs are used to materialize global addresses in PIC
   /// mode.
   DenseMap<MCSymbol *, StubValueTy> GVStubs;
 
+  DenseMap<MCSymbol *, AuthStubInfo> AuthPtrStubs;
+
   virtual void anchor(); // Out of line virtual method.
 
 public:
@@ -97,9 +105,19 @@ class MachineModuleInfoELF : public MachineModuleInfoImpl {
     return GVStubs[Sym];
   }
 
+  AuthStubInfo &getAuthPtrStubEntry(MCSymbol *Sym) {
+    assert(Sym && "Key cannot be null");
+    return AuthPtrStubs[Sym];
+  }
+
   /// Accessor methods to return the set of stubs in sorted order.
 
   SymbolListTy GetGVStubList() { return getSortedStubs(GVStubs); }
+
+  typedef std::pair<MCSymbol *, AuthStubInfo> AuthStubPairTy;
+  typedef std::vector<AuthStubPairTy> AuthStubListTy;
+
+  AuthStubListTy getAuthGVStubList();
 };
 
 /// MachineModuleInfoCOFF - This is a MachineModuleInfoImpl implementation

llvm/lib/CodeGen/MachineModuleInfoImpls.cpp

@@ -42,16 +42,29 @@ MachineModuleInfoImpl::SymbolListTy MachineModuleInfoImpl::getSortedStubs(
   return List;
 }
 
-MachineModuleInfoMachO::AuthStubListTy
-MachineModuleInfoMachO::getAuthGVStubList() {
-  AuthStubListTy List(AuthPtrStubs.begin(), AuthPtrStubs.end());
+template <typename MachineModuleInfoTarget>
+static typename MachineModuleInfoTarget::AuthStubListTy getAuthGVStubListHelper(
+    DenseMap<MCSymbol *, typename MachineModuleInfoTarget::AuthStubInfo>
+        &AuthPtrStubs) {
+  typename MachineModuleInfoTarget::AuthStubListTy List(AuthPtrStubs.begin(),
+                                                        AuthPtrStubs.end());
 
   if (!List.empty())
     std::sort(List.begin(), List.end(),
-              [](const AuthStubPairTy &LHS, const AuthStubPairTy &RHS) {
-              return LHS.first->getName() < RHS.first->getName();
+              [](const typename MachineModuleInfoTarget::AuthStubPairTy &LHS,
+                 const typename MachineModuleInfoTarget::AuthStubPairTy &RHS) {
+                return LHS.first->getName() < RHS.first->getName();
               });
 
   AuthPtrStubs.clear();
   return List;
 }
+
+MachineModuleInfoMachO::AuthStubListTy
+MachineModuleInfoMachO::getAuthGVStubList() {
+  return getAuthGVStubListHelper<MachineModuleInfoMachO>(AuthPtrStubs);
+}
+
+MachineModuleInfoELF::AuthStubListTy MachineModuleInfoELF::getAuthGVStubList() {
+  return getAuthGVStubListHelper<MachineModuleInfoELF>(AuthPtrStubs);
+}

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -860,9 +860,10 @@ void AArch64AsmPrinter::emitHwasanMemaccessSymbols(Module &M) {
   }
 }
 
-static void
-emitAuthenticatedPointer(MCStreamer &OutStreamer, MCSymbol *StubLabel,
-                         const MachineModuleInfoMachO::AuthStubInfo &StubInfo) {
+template <typename MachineModuleInfoTarget>
+static void emitAuthenticatedPointer(
+    MCStreamer &OutStreamer, MCSymbol *StubLabel,
+    const typename MachineModuleInfoTarget::AuthStubInfo &StubInfo) {
   // L_foo$addend$auth_ptr$ib$23:
   OutStreamer.emitLabel(StubLabel);
   OutStreamer.emitValue(StubInfo.AuthPtrRef, /*size=*/8);
@@ -888,7 +889,8 @@ void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
       emitAlignment(Align(8));
 
       for (auto &Stub : Stubs)
-        emitAuthenticatedPointer(*OutStreamer, Stub.first, Stub.second);
+        emitAuthenticatedPointer<MachineModuleInfoMachO>(
+            *OutStreamer, Stub.first, Stub.second);
 
       OutStreamer->addBlankLine();
     }
@@ -901,9 +903,27 @@ void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
     OutStreamer->emitAssemblerFlag(MCAF_SubsectionsViaSymbols);
   }
 
+  if (TT.isOSBinFormatELF()) {
+    // Output authenticated pointers as indirect symbols, if we have any.
+    MachineModuleInfoELF &MMIELF = MMI->getObjFileInfo<MachineModuleInfoELF>();
+
+    auto Stubs = MMIELF.getAuthGVStubList();
+
+    if (!Stubs.empty()) {
+      const TargetLoweringObjectFile &TLOF = getObjFileLowering();
+      OutStreamer->switchSection(TLOF.getDataSection());
+      emitAlignment(Align(8));
+
+      for (auto &Stub : Stubs)
+        emitAuthenticatedPointer<MachineModuleInfoELF>(*OutStreamer, Stub.first,
+                                                       Stub.second);
+
+      OutStreamer->addBlankLine();
+    }
+  }
+
   // Emit stack and fault map information.
   FM.serializeToFaultMapSection();
-
 }
 
 void AArch64AsmPrinter::emitLOHs() {
@@ -1877,8 +1897,8 @@ AArch64AsmPrinter::lowerPtrAuthGlobalConstant(const GlobalPtrAuthInfo &PAI) {
 
     std::string Buf;
     raw_string_ostream OS(Buf);
-    OS << "Couldn't resolve target base/addend of llvm.ptrauth global '"
-      << *BaseGVB << "'";
+    OS << "couldn't resolve target base/addend of llvm.ptrauth global '"
+       << *BaseGVB << "'";
     BaseGV->getContext().emitError(OS.str());
   }
 
@@ -1892,16 +1912,20 @@ AArch64AsmPrinter::lowerPtrAuthGlobalConstant(const GlobalPtrAuthInfo &PAI) {
         Sym, MCConstantExpr::create((-Offset).getSExtValue(), Ctx), Ctx);
 
   uint64_t KeyID = PAI.getKey()->getZExtValue();
-  if (!isUInt<2>(KeyID))
+  if (KeyID > AArch64PACKey::LAST)
     BaseGV->getContext().emitError(
-        "Invalid AArch64 PAC Key ID '" + utostr(KeyID) + "' in llvm.ptrauth global '" +
+        "AArch64 PAC key ID '" + Twine(KeyID) + "' out of range [0, " +
+        Twine(static_cast<int>(AArch64PACKey::LAST)) +
+        "] in llvm.ptrauth "
+        "global '" +
         BaseGV->getName() + "'");
 
   uint64_t Disc = PAI.getDiscriminator()->getZExtValue();
   if (!isUInt<16>(Disc))
-    BaseGV->getContext().emitError("Invalid AArch64 Discriminator '" +
-                                   utostr(Disc) + "' in llvm.ptrauth global '" +
-                                   BaseGV->getName() + "'");
+    BaseGV->getContext().emitError(
+        "AArch64 PAC discriminator '" + Twine(Disc) +
+        "' out of range [0, 0xFFFF] in llvm.ptrauth global '" +
+        BaseGV->getName() + "'");
 
   // Finally build the complete @AUTH expr.
   return AArch64AuthMCExpr::create(Sym, Disc, AArch64PACKey::ID(KeyID),

llvm/lib/Target/AArch64/AArch64ExpandHardenedPseudos.cpp

@@ -191,12 +191,22 @@ bool AArch64ExpandHardenedPseudos::expandPtrAuthPseudo(MachineInstr &MI) {
     // Where the $auth_ptr$ symbol is the stub slot containing the signed pointer
     // to _symbol.
     // We defined the stub ourselves, so we don't need a GOT access.
-    assert(TM.getTargetTriple().isOSBinFormatMachO() &&
-           "ptrauth chkstk_darwin only implemented on mach-o");
-    auto *TLOF =
-      static_cast<AArch64_MachoTargetObjectFile *>(TM.getObjFileLowering());
-    MCSymbol *AuthPtrStubSym =
-        TLOF->getAuthPtrSlotSymbol(TM, &MMI, GASym, Offset, Key, Disc);
+
+    MCSymbol *AuthPtrStubSym;
+    if (TM.getTargetTriple().isOSBinFormatMachO()) {
+      auto *TLOF =
+          static_cast<AArch64_MachoTargetObjectFile *>(TM.getObjFileLowering());
+      AuthPtrStubSym =
+          TLOF->getAuthPtrSlotSymbol(TM, &MMI, GASym, Offset, Key, Disc);
+    } else if (TM.getTargetTriple().isOSBinFormatELF()) {
+      auto *TLOF =
+          static_cast<AArch64_ELFTargetObjectFile *>(TM.getObjFileLowering());
+      AuthPtrStubSym =
+          TLOF->getAuthPtrSlotSymbol(TM, &MMI, GASym, Offset, Key, Disc);
+    } else {
+      llvm_unreachable(
+          "ptrauth static materialization only implemented on mach-o and elf");
+    }
 
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADRP), DstReg)
       .addSym(AuthPtrStubSym, AArch64II::MO_PAGE);
@@ -278,30 +288,33 @@ bool AArch64ExpandHardenedPseudos::expandPtrAuthPseudo(MachineInstr &MI) {
     }
   }
 
+  auto *SignMBB = &MBB;
+  auto SignMBBI = MBBI;
+
   unsigned DiscReg = AddrDisc;
   if (Disc) {
     assert(isUInt<16>(Disc) && "Constant discriminator is too wide");
 
     if (AddrDisc != AArch64::XZR) {
-      BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), AArch64::X17)
-        .addReg(AArch64::XZR)
-        .addReg(AddrDisc)
-        .addImm(0);
-      BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVKXi), AArch64::X17)
-        .addReg(AArch64::X17)
-        .addImm(Disc)
-        .addImm(/*shift=*/48);
+      BuildMI(*SignMBB, SignMBBI, DL, TII->get(AArch64::ORRXrs), AArch64::X17)
+          .addReg(AArch64::XZR)
+          .addReg(AddrDisc)
+          .addImm(0);
+      BuildMI(*SignMBB, SignMBBI, DL, TII->get(AArch64::MOVKXi), AArch64::X17)
+          .addReg(AArch64::X17)
+          .addImm(Disc)
+          .addImm(/*shift=*/48);
     } else {
-      BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVZXi), AArch64::X17)
-        .addImm(Disc)
-        .addImm(/*shift=*/0);
+      BuildMI(*SignMBB, SignMBBI, DL, TII->get(AArch64::MOVZXi), AArch64::X17)
+          .addImm(Disc)
+          .addImm(/*shift=*/0);
     }
     DiscReg = AArch64::X17;
   }
 
   unsigned PACOpc = getPACOpcodeForKey(Key, DiscReg == AArch64::XZR);
-  auto MIB = BuildMI(MBB, MBBI, DL, TII->get(PACOpc), AArch64::X16)
-      .addReg(AArch64::X16);
+  auto MIB = BuildMI(*SignMBB, SignMBBI, DL, TII->get(PACOpc), AArch64::X16)
+                 .addReg(AArch64::X16);
   if (DiscReg != AArch64::XZR)
     MIB.addReg(DiscReg);
 
@@ -448,8 +461,12 @@ bool AArch64ExpandHardenedPseudos::runOnMachineFunction(MachineFunction &MF) {
   bool Modified = false;
   for (auto &MBB : MF) {
     for (auto MBBI = MBB.instr_begin(), MBBE = MBB.instr_end(); MBBI != MBBE; ) {
-      auto &MI = *MBBI++;
-      if (expandMI(MI)) {
+      auto &MI = *MBBI;
+      bool IsExpanded = expandMI(MI);
+      // Previous expandMI invocation might split the current basic block, so we
+      // increment the instruction iterator after that.
+      ++MBBI;
+      if (IsExpanded) {
         MI.eraseFromBundle();
         Modified |= true;
       }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8601,6 +8601,8 @@ SDValue AArch64TargetLowering::LowerGlobalTLSAddress(SDValue Op,
 SDValue AArch64TargetLowering::LowerPtrAuthGlobalAddressViaGOT(
     SDValue TGA, SDLoc DL, EVT VT, AArch64PACKey::ID KeyC,
     SDValue Discriminator, SDValue AddrDiscriminator, SelectionDAG &DAG) const {
+  assert((Subtarget->isTargetMachO() || Subtarget->isTargetELF()) &&
+         "LOADauthptrgot is only supported for MachO and ELF");
 
   auto *TGN = cast<GlobalAddressSDNode>(TGA.getNode());
   const GlobalValue *GV = TGN->getGlobal();
@@ -8672,7 +8674,7 @@ AArch64TargetLowering::LowerPtrAuthGlobalAddress(SDValue Op,
   SDLoc DL(Op);
 
   // Both key allocation and the wrapper usage support are target-specific.
-  if (!Subtarget->isTargetMachO())
+  if (!Subtarget->isTargetMachO() && !Subtarget->isTargetELF())
     llvm_unreachable("Unimplemented ptrauth global lowering");
 
   uint64_t PtrOffsetC = 0;
@@ -8697,12 +8699,14 @@ AArch64TargetLowering::LowerPtrAuthGlobalAddress(SDValue Op,
   assert(PtrN->getTargetFlags() == 0 && "Unsupported tflags on ptrauth global");
 
   // None of our lowerings support an offset larger than 32-bit.
+  // TODO: rela elf supports 64-bit?
   if (!isUInt<32>(PtrOffsetC))
     report_fatal_error("Unsupported >32-bit-wide offset in ptrauth global");
 
   // The shared cache severely constrains the offset, to 5 bits, which prevents
-  // us from using LOADauthptrgot, our only option for extern_weak.
-  if (!isUInt<5>(PtrOffsetC) && PtrGV->hasExternalWeakLinkage())
+  // us from using LOADauthptrgot, our only option for extern_weak on MachO.
+  if (!isUInt<5>(PtrOffsetC) && PtrGV->hasExternalWeakLinkage() &&
+      (Subtarget->isTargetMachO() || Subtarget->isTargetELF()))
     report_fatal_error("Offset in weak ptrauth global reference is too large");
 
   // Blend only works if the integer discriminator is 16-bit wide.
@@ -8718,7 +8722,9 @@ AArch64TargetLowering::LowerPtrAuthGlobalAddress(SDValue Op,
 
   // No GOT load needed -> MOVaddrPAC
   if (!NeedsGOTLoad) {
-    assert(!PtrGV->hasExternalWeakLinkage() && "extern_weak should use GOT");
+    assert(!(PtrGV->hasExternalWeakLinkage() &&
+             (Subtarget->isTargetMachO() || Subtarget->isTargetELF())) &&
+           "extern_weak on MachO should use GOT via LOADauthptrgot");
     return SDValue(
         DAG.getMachineNode(AArch64::MOVaddrPAC, DL, MVT::i64,
                            {TPtr, Key, TAddrDiscriminator, Discriminator}),

llvm/lib/Target/AArch64/AArch64MCInstLower.cpp

@@ -39,13 +39,19 @@ AArch64MCInstLower::AArch64MCInstLower(MCContext &ctx, AsmPrinter &printer)
     : Ctx(ctx), Printer(printer) {}
 
 static MCSymbol *getAuthGVStub(const GlobalVariable *GVB, AsmPrinter &Printer) {
-  assert(Printer.TM.getTargetTriple().isOSBinFormatMachO() &&
-         "auth_ptr stubs only implemented on macho");
-  auto &TLOF = static_cast<const AArch64_MachoTargetObjectFile &>(
-      Printer.getObjFileLowering());
-
-  return TLOF.getAuthPtrSlotSymbol(Printer.TM, Printer.MMI,
-                                   *GlobalPtrAuthInfo::analyze(GVB));
+  if (Printer.TM.getTargetTriple().isOSBinFormatMachO()) {
+    auto &TLOF = static_cast<const AArch64_MachoTargetObjectFile &>(
+        Printer.getObjFileLowering());
+    return TLOF.getAuthPtrSlotSymbol(Printer.TM, Printer.MMI,
+                                     *GlobalPtrAuthInfo::analyze(GVB));
+  } else if (Printer.TM.getTargetTriple().isOSBinFormatELF()) {
+    auto &TLOF = static_cast<const AArch64_ELFTargetObjectFile &>(
+        Printer.getObjFileLowering());
+    return TLOF.getAuthPtrSlotSymbol(Printer.TM, Printer.MMI,
+                                     *GlobalPtrAuthInfo::analyze(GVB));
+  } else {
+    llvm_unreachable("auth_ptr stubs only implemented on macho and elf");
+  }
 }
 
 MCSymbol *