CodeGen: Optionally emit PAuth relocations as IRELATIVE relocations.

This supports the following use cases:
- ConstantPtrAuth expressions that are unrepresentable using standard PAuth
  relocations such as expressions involving an integer operand or
  deactivation symbols.
- libc implementations that do not support PAuth relocations.

For more information see the RFC:
https://discourse.llvm.org/t/rfc-structure-protection-a-family-of-uaf-mitigation-techniques/85555

TODO:
- Add tests.

Pull Request: llvm#133533

Author: pcc

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -49,12 +49,14 @@
 #include "llvm/IR/Module.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstBuilder.h"
 #include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCSectionMachO.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCValue.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
@@ -95,6 +97,7 @@ class AArch64AsmPrinter : public AsmPrinter {
   bool EnableImportCallOptimization = false;
   DenseMap<MCSection *, std::vector<std::pair<MCSymbol *, MCSymbol *>>>
       SectionToImportedFunctionCalls;
+  unsigned PAuthIFuncNextUniqueID = 1;
 
 public:
   static char ID;
@@ -209,6 +212,10 @@ class AArch64AsmPrinter : public AsmPrinter {
   // authenticating)
   void LowerLOADgotAUTH(const MachineInstr &MI);
 
+  const MCExpr *emitPAuthRelocationAsIRelative(
+      const MCExpr *Target, uint16_t Disc, AArch64PACKey::ID KeyID,
+      bool HasAddressDiversity, bool IsDSOLocal);
+
   /// tblgen'erated driver function for lowering simple MI->MC
   /// pseudo instructions.
   bool lowerPseudoInstExpansion(const MachineInstr *MI, MCInst &Inst);
@@ -2225,6 +2232,146 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   EmitToStreamer(*OutStreamer, BRInst);
 }
 
+static void emitAddress(MCStreamer &Streamer, MCRegister Reg,
+                        const MCExpr *Expr, bool DSOLocal,
+                        const MCSubtargetInfo &STI) {
+  MCValue Val;
+  if (!Expr->evaluateAsRelocatable(Val, nullptr))
+    report_fatal_error("emitAddress could not evaluate");
+  if (DSOLocal) {
+    Streamer.emitInstruction(
+        MCInstBuilder(AArch64::ADRP)
+            .addReg(Reg)
+            .addExpr(MCSpecifierExpr::create(Expr, AArch64::S_ABS_PAGE,
+                                             Streamer.getContext())),
+        STI);
+    Streamer.emitInstruction(
+        MCInstBuilder(AArch64::ADDXri)
+            .addReg(Reg)
+            .addReg(Reg)
+            .addExpr(MCSpecifierExpr::create(Expr, AArch64::S_LO12,
+                                             Streamer.getContext()))
+            .addImm(0),
+        STI);
+  } else {
+    auto *SymRef = MCSymbolRefExpr::create(Val.getAddSym(), Streamer.getContext());
+    Streamer.emitInstruction(
+        MCInstBuilder(AArch64::ADRP)
+            .addReg(Reg)
+            .addExpr(MCSpecifierExpr::create(SymRef, AArch64::S_GOT_PAGE,
+                                             Streamer.getContext())),
+        STI);
+    Streamer.emitInstruction(
+        MCInstBuilder(AArch64::LDRXui)
+            .addReg(Reg)
+            .addReg(Reg)
+            .addExpr(MCSpecifierExpr::create(SymRef, AArch64::S_GOT_LO12,
+                                             Streamer.getContext())),
+        STI);
+    if (Val.getConstant())
+      Streamer.emitInstruction(MCInstBuilder(AArch64::ADDXri)
+                                   .addReg(Reg)
+                                   .addReg(Reg)
+                                   .addImm(Val.getConstant())
+                                   .addImm(0),
+                               STI);
+  }
+}
+
+static bool targetSupportsPAuthRelocation(const Triple &TT,
+                                          const MCExpr *Target) {
+  // No released version of glibc supports PAuth relocations.
+  if (TT.isOSGlibc())
+    return false;
+
+  // We emit PAuth constants as IRELATIVE relocations in cases where the
+  // constant cannot be represented as a PAuth relocation:
+  // 1) The signed value is not a symbol.
+  return !isa<MCConstantExpr>(Target);
+}
+
+static bool targetSupportsIRelativeRelocation(const Triple &TT) {
+  // IFUNCs are ELF-only.
+  if (!TT.isOSBinFormatELF())
+    return false;
+
+  // musl doesn't support IFUNCs.
+  if (TT.isMusl())
+    return false;
+
+  return true;
+}
+
+const MCExpr *AArch64AsmPrinter::emitPAuthRelocationAsIRelative(
+    const MCExpr *Target, uint16_t Disc, AArch64PACKey::ID KeyID,
+    bool HasAddressDiversity, bool IsDSOLocal) {
+  const Triple &TT = TM.getTargetTriple();
+
+  // We only emit an IRELATIVE relocation if the target supports IRELATIVE and
+  // does not support the kind of PAuth relocation that we are trying to emit.
+  if (targetSupportsPAuthRelocation(TT, Target) ||
+      !targetSupportsIRelativeRelocation(TT))
+    return nullptr;
+
+  // For now, only the DA key is supported.
+  if (KeyID != AArch64PACKey::DA)
+    return nullptr;
+
+  std::unique_ptr<MCSubtargetInfo> STI(
+      TM.getTarget().createMCSubtargetInfo(TT.str(), "", ""));
+  assert(STI && "Unable to create subtarget info");
+
+  MCSymbol *Place = OutStreamer->getContext().createTempSymbol();
+  OutStreamer->emitLabel(Place);
+  OutStreamer->pushSection();
+
+  OutStreamer->switchSection(OutStreamer->getContext().getELFSection(
+      ".text.startup", ELF::SHT_PROGBITS, ELF::SHF_ALLOC | ELF::SHF_EXECINSTR,
+      0, "", true, PAuthIFuncNextUniqueID++, nullptr));
+
+  MCSymbol *IRelativeSym =
+      OutStreamer->getContext().createLinkerPrivateSymbol("pauth_ifunc");
+  OutStreamer->emitLabel(IRelativeSym);
+  if (isa<MCConstantExpr>(Target)) {
+    OutStreamer->emitInstruction(MCInstBuilder(AArch64::MOVZXi)
+                                     .addReg(AArch64::X0)
+                                     .addExpr(Target)
+                                     .addImm(0),
+                                 *STI);
+  } else {
+    emitAddress(*OutStreamer, AArch64::X0, Target, IsDSOLocal, *STI);
+  }
+  if (HasAddressDiversity) {
+    auto *PlacePlusDisc = MCBinaryExpr::createAdd(
+        MCSymbolRefExpr::create(Place, OutStreamer->getContext()),
+        MCConstantExpr::create(static_cast<int16_t>(Disc),
+                               OutStreamer->getContext()),
+        OutStreamer->getContext());
+    emitAddress(*OutStreamer, AArch64::X1, PlacePlusDisc, /*IsDSOLocal=*/true,
+                *STI);
+  } else {
+    emitMOVZ(AArch64::X1, Disc, 0);
+  }
+
+  MCSymbol *PrePACInst = OutStreamer->getContext().createTempSymbol();
+  OutStreamer->emitLabel(PrePACInst);
+
+  // We don't know the subtarget because this is being emitted for a global
+  // initializer. Because the performance of IFUNC resolvers is unimportant, we
+  // always call the EmuPAC runtime, which will end up using the PAC instruction
+  // if the target supports PAC.
+  MCSymbol *EmuPAC =
+      OutStreamer->getContext().getOrCreateSymbol("__emupac_pacda");
+  const MCSymbolRefExpr *EmuPACRef =
+      MCSymbolRefExpr::create(EmuPAC, OutStreamer->getContext());
+  OutStreamer->emitInstruction(MCInstBuilder(AArch64::B).addExpr(EmuPACRef),
+                               *STI);
+  OutStreamer->popSection();
+
+  return MCSpecifierExpr::create(IRelativeSym, AArch64::S_FUNCINIT,
+                                 OutStreamer->getContext());
+}
+
 const MCExpr *
 AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
   MCContext &Ctx = OutContext;
@@ -2236,23 +2383,20 @@ AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
 
   auto *BaseGVB = dyn_cast<GlobalValue>(BaseGV);
 
-  // If we can't understand the referenced ConstantExpr, there's nothing
-  // else we can do: emit an error.
-  if (!BaseGVB) {
-    BaseGV->getContext().emitError(
-        "cannot resolve target base/addend of ptrauth constant");
-    return nullptr;
+  const MCExpr *Sym;
+  if (BaseGVB) {
+    // If there is an addend, turn that into the appropriate MCExpr.
+    Sym = MCSymbolRefExpr::create(getSymbol(BaseGVB), Ctx);
+    if (Offset.sgt(0))
+      Sym = MCBinaryExpr::createAdd(
+          Sym, MCConstantExpr::create(Offset.getSExtValue(), Ctx), Ctx);
+    else if (Offset.slt(0))
+      Sym = MCBinaryExpr::createSub(
+          Sym, MCConstantExpr::create((-Offset).getSExtValue(), Ctx), Ctx);
+  } else {
+    Sym = MCConstantExpr::create(Offset.getSExtValue(), Ctx);
   }
 
-  // If there is an addend, turn that into the appropriate MCExpr.
-  const MCExpr *Sym = MCSymbolRefExpr::create(getSymbol(BaseGVB), Ctx);
-  if (Offset.sgt(0))
-    Sym = MCBinaryExpr::createAdd(
-        Sym, MCConstantExpr::create(Offset.getSExtValue(), Ctx), Ctx);
-  else if (Offset.slt(0))
-    Sym = MCBinaryExpr::createSub(
-        Sym, MCConstantExpr::create((-Offset).getSExtValue(), Ctx), Ctx);
-
   uint64_t KeyID = CPA.getKey()->getZExtValue();
   // We later rely on valid KeyID value in AArch64PACKeyIDToString call from
   // AArch64AuthMCExpr::printImpl, so fail fast.
@@ -2270,6 +2414,12 @@ AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
     Disc = 0;
   }
 
+  // Check if we need to represent this with an IRELATIVE and emit it if so.
+  if (auto *IFuncSym = emitPAuthRelocationAsIRelative(
+          Sym, Disc, AArch64PACKey::ID(KeyID), CPA.hasAddressDiscriminator(),
+          BaseGVB && BaseGVB->isDSOLocal()))
+    return IFuncSym;
+
   // Finally build the complete @AUTH expr.
   return AArch64AuthMCExpr::create(Sym, Disc, AArch64PACKey::ID(KeyID),
                                    CPA.hasAddressDiscriminator(), Ctx);