[PAC][CodeGen][ELF][AArch64] Support signed GOT (llvm#113811)

This re-applies llvm#96164 after revert in llvm#102434.

Support the following relocations and assembly operators:

- `R_AARCH64_AUTH_ADR_GOT_PAGE` (`:got_auth:` for `adrp`)
- `R_AARCH64_AUTH_LD64_GOT_LO12_NC` (`:got_auth_lo12:` for `ldr`)
- `R_AARCH64_AUTH_GOT_ADD_LO12_NC` (`:got_auth_lo12:` for `add`)

`LOADgotAUTH` pseudo-instruction is introduced which is later expanded to
actual instruction sequence like the following.

```
adrp x16, :got_auth:sym
add x16, x16, :got_auth_lo12:sym
ldr x0, [x16]
autia x0, x16
```

If a resign is requested, like below, `LOADgotPAC` pseudo is used, and GOT
load is lowered similarly to `LOADgotAUTH`.

```
@var = global i32 0
define ptr @resign_globalvar() {
  ret ptr ptrauth (ptr @var, i32 3, i64 43)
}
```

If FPAC bit is not set and auth instruction is emitted, a check+trap sequence
similar to one used for `AUT` pseudo is emitted to ensure auth success.

Both SelectionDAG and GlobalISel are suppported.
For FastISel, we fall back to SelectionDAG.

Tests starting with 'ptrauth-' have corresponding variants w/o this prefix.

See also specification
https://github.com/ARM-software/abi-aa/blob/main/pauthabielf64/pauthabielf64.rst#appendix-signed-got

Author: kovdan01

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -169,6 +169,11 @@ class AArch64AsmPrinter : public AsmPrinter {
   // adrp-add followed by PAC sign)
   void LowerMOVaddrPAC(const MachineInstr &MI);
 
+  // Emit the sequence for LOADgotAUTH (load signed pointer from signed ELF GOT
+  // and authenticate it with, if FPAC bit is not set, check+trap sequence after
+  // authenticating)
+  void LowerLOADgotAUTH(const MachineInstr &MI);
+
   /// tblgen'erated driver function for lowering simple MI->MC
   /// pseudo instructions.
   bool lowerPseudoInstExpansion(const MachineInstr *MI, MCInst &Inst);
@@ -873,6 +878,22 @@ void AArch64AsmPrinter::emitEndOfAsmFile(Module &M) {
 
       OutStreamer->addBlankLine();
     }
+
+    // With signed ELF GOT enabled, the linker looks at the symbol type to
+    // choose between keys IA (for STT_FUNC) and DA (for other types). Symbols
+    // for functions not defined in the module have STT_NOTYPE type by default.
+    // This makes linker to emit signing schema with DA key (instead of IA) for
+    // corresponding R_AARCH64_AUTH_GLOB_DAT dynamic reloc. To avoid that, force
+    // all function symbols used in the module to have STT_FUNC type. See
+    // https://github.com/ARM-software/abi-aa/blob/main/pauthabielf64/pauthabielf64.rst#default-signing-schema
+    const auto *PtrAuthELFGOTFlag = mdconst::extract_or_null<ConstantInt>(
+        M.getModuleFlag("ptrauth-elf-got"));
+    if (PtrAuthELFGOTFlag && PtrAuthELFGOTFlag->getZExtValue() == 1)
+      for (const GlobalValue &GV : M.global_values())
+        if (!GV.use_empty() && isa<Function>(GV) &&
+            !GV.getName().starts_with("llvm."))
+          OutStreamer->emitSymbolAttribute(getSymbol(&GV),
+                                           MCSA_ELF_TypeFunction);
   }
 
   // Emit stack and fault map information.
@@ -2068,6 +2089,10 @@ void AArch64AsmPrinter::LowerLOADauthptrstatic(const MachineInstr &MI) {
 
 void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   const bool IsGOTLoad = MI.getOpcode() == AArch64::LOADgotPAC;
+  const bool IsELFSignedGOT = MI.getParent()
+                                  ->getParent()
+                                  ->getInfo<AArch64FunctionInfo>()
+                                  ->hasELFSignedGOT();
   MachineOperand GAOp = MI.getOperand(0);
   const uint64_t KeyC = MI.getOperand(1).getImm();
   assert(KeyC <= AArch64PACKey::LAST &&
@@ -2084,9 +2109,17 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   // Emit:
   // target materialization:
   // - via GOT:
-  //     adrp x16, :got:target
-  //     ldr x16, [x16, :got_lo12:target]
-  //     add offset to x16 if offset != 0
+  //   - unsigned GOT:
+  //       adrp x16, :got:target
+  //       ldr x16, [x16, :got_lo12:target]
+  //       add offset to x16 if offset != 0
+  //   - ELF signed GOT:
+  //       adrp x17, :got:target
+  //       add x17, x17, :got_auth_lo12:target
+  //       ldr x16, [x17]
+  //       aut{i|d}a x16, x17
+  //       check+trap sequence (if no FPAC)
+  //       add offset to x16 if offset != 0
   //
   // - direct:
   //     adrp x16, target
@@ -2129,13 +2162,48 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   MCInstLowering.lowerOperand(GAMOLo, GAMCLo);
 
   EmitToStreamer(
-      MCInstBuilder(AArch64::ADRP).addReg(AArch64::X16).addOperand(GAMCHi));
+      MCInstBuilder(AArch64::ADRP)
+          .addReg(IsGOTLoad && IsELFSignedGOT ? AArch64::X17 : AArch64::X16)
+          .addOperand(GAMCHi));
 
   if (IsGOTLoad) {
-    EmitToStreamer(MCInstBuilder(AArch64::LDRXui)
-                       .addReg(AArch64::X16)
-                       .addReg(AArch64::X16)
-                       .addOperand(GAMCLo));
+    if (IsELFSignedGOT) {
+      EmitToStreamer(MCInstBuilder(AArch64::ADDXri)
+                         .addReg(AArch64::X17)
+                         .addReg(AArch64::X17)
+                         .addOperand(GAMCLo)
+                         .addImm(0));
+
+      EmitToStreamer(MCInstBuilder(AArch64::LDRXui)
+                         .addReg(AArch64::X16)
+                         .addReg(AArch64::X17)
+                         .addImm(0));
+
+      assert(GAOp.isGlobal());
+      assert(GAOp.getGlobal()->getValueType() != nullptr);
+      unsigned AuthOpcode = GAOp.getGlobal()->getValueType()->isFunctionTy()
+                                ? AArch64::AUTIA
+                                : AArch64::AUTDA;
+
+      EmitToStreamer(MCInstBuilder(AuthOpcode)
+                         .addReg(AArch64::X16)
+                         .addReg(AArch64::X16)
+                         .addReg(AArch64::X17));
+
+      if (!STI->hasFPAC()) {
+        auto AuthKey = (AuthOpcode == AArch64::AUTIA ? AArch64PACKey::IA
+                                                     : AArch64PACKey::DA);
+
+        emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AuthKey,
+                                           /*ShouldTrap=*/true,
+                                           /*OnFailure=*/nullptr);
+      }
+    } else {
+      EmitToStreamer(MCInstBuilder(AArch64::LDRXui)
+                         .addReg(AArch64::X16)
+                         .addReg(AArch64::X16)
+                         .addOperand(GAMCLo));
+    }
   } else {
     EmitToStreamer(MCInstBuilder(AArch64::ADDXri)
                        .addReg(AArch64::X16)
@@ -2203,6 +2271,69 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   EmitToStreamer(MIB);
 }
 
+void AArch64AsmPrinter::LowerLOADgotAUTH(const MachineInstr &MI) {
+  Register DstReg = MI.getOperand(0).getReg();
+  Register AuthResultReg = STI->hasFPAC() ? DstReg : AArch64::X16;
+  const MachineOperand &GAMO = MI.getOperand(1);
+  assert(GAMO.getOffset() == 0);
+
+  MachineOperand GAHiOp(GAMO);
+  MachineOperand GALoOp(GAMO);
+  GAHiOp.addTargetFlag(AArch64II::MO_PAGE);
+  GALoOp.addTargetFlag(AArch64II::MO_PAGEOFF | AArch64II::MO_NC);
+
+  MCOperand GAMCHi, GAMCLo;
+  MCInstLowering.lowerOperand(GAHiOp, GAMCHi);
+  MCInstLowering.lowerOperand(GALoOp, GAMCLo);
+
+  EmitToStreamer(
+      MCInstBuilder(AArch64::ADRP).addReg(AArch64::X17).addOperand(GAMCHi));
+
+  EmitToStreamer(MCInstBuilder(AArch64::ADDXri)
+                     .addReg(AArch64::X17)
+                     .addReg(AArch64::X17)
+                     .addOperand(GAMCLo)
+                     .addImm(0));
+
+  EmitToStreamer(MCInstBuilder(AArch64::LDRXui)
+                     .addReg(AuthResultReg)
+                     .addReg(AArch64::X17)
+                     .addImm(0));
+
+  assert(GAMO.isGlobal());
+  MCSymbol *UndefWeakSym;
+  if (GAMO.getGlobal()->hasExternalWeakLinkage()) {
+    UndefWeakSym = createTempSymbol("undef_weak");
+    EmitToStreamer(
+        MCInstBuilder(AArch64::CBZX)
+            .addReg(AuthResultReg)
+            .addExpr(MCSymbolRefExpr::create(UndefWeakSym, OutContext)));
+  }
+
+  assert(GAMO.getGlobal()->getValueType() != nullptr);
+  unsigned AuthOpcode = GAMO.getGlobal()->getValueType()->isFunctionTy()
+                            ? AArch64::AUTIA
+                            : AArch64::AUTDA;
+  EmitToStreamer(MCInstBuilder(AuthOpcode)
+                     .addReg(AuthResultReg)
+                     .addReg(AuthResultReg)
+                     .addReg(AArch64::X17));
+
+  if (GAMO.getGlobal()->hasExternalWeakLinkage())
+    OutStreamer->emitLabel(UndefWeakSym);
+
+  if (!STI->hasFPAC()) {
+    auto AuthKey =
+        (AuthOpcode == AArch64::AUTIA ? AArch64PACKey::IA : AArch64PACKey::DA);
+
+    emitPtrauthCheckAuthenticatedValue(AuthResultReg, AArch64::X17, AuthKey,
+                                       /*ShouldTrap=*/true,
+                                       /*OnFailure=*/nullptr);
+
+    emitMovXReg(DstReg, AuthResultReg);
+  }
+}
+
 const MCExpr *
 AArch64AsmPrinter::lowerBlockAddressConstant(const BlockAddress &BA) {
   const MCExpr *BAE = AsmPrinter::lowerBlockAddressConstant(BA);
@@ -2381,6 +2512,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     LowerMOVaddrPAC(*MI);
     return;
 
+  case AArch64::LOADgotAUTH:
+    LowerLOADgotAUTH(*MI);
+    return;
+
   case AArch64::BRA:
   case AArch64::BLRA:
     emitPtrauthBranch(MI);

llvm/lib/Target/AArch64/AArch64FastISel.cpp

@@ -454,6 +454,9 @@ unsigned AArch64FastISel::materializeGV(const GlobalValue *GV) {
   if (!Subtarget->useSmallAddressing() && !Subtarget->isTargetMachO())
     return 0;
 
+  if (FuncInfo.MF->getInfo<AArch64FunctionInfo>()->hasELFSignedGOT())
+    return 0;
+
   unsigned OpFlags = Subtarget->ClassifyGlobalReference(GV, TM);
 
   EVT DestEVT = TLI.getValueType(DL, GV->getType(), true);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -9599,6 +9599,11 @@ SDValue AArch64TargetLowering::getGOT(NodeTy *N, SelectionDAG &DAG,
   SDValue GotAddr = getTargetNode(N, Ty, DAG, AArch64II::MO_GOT | Flags);
   // FIXME: Once remat is capable of dealing with instructions with register
   // operands, expand this into two nodes instead of using a wrapper node.
+  if (DAG.getMachineFunction()
+          .getInfo<AArch64FunctionInfo>()
+          ->hasELFSignedGOT())
+    return SDValue(DAG.getMachineNode(AArch64::LOADgotAUTH, DL, Ty, GotAddr),
+                   0);
   return DAG.getNode(AArch64ISD::LOADgot, DL, Ty, GotAddr);
 }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -1942,8 +1942,15 @@ let Predicates = [HasPAuth] in {
                Sched<[WriteI, ReadI]> {
     let isReMaterializable = 1;
     let isCodeGenOnly = 1;
-    let Size = 40; // 12 fixed + 28 variable, for pointer offset, and discriminator
-    let Defs = [X16,X17];
+    let Size = 68; // 12 fixed + 56 variable, for pointer offset, discriminator and
+                   // ELF signed GOT signed pointer authentication (if no FPAC)
+    let Defs = [X16,X17,NZCV];
+  }
+
+  def LOADgotAUTH : Pseudo<(outs GPR64common:$dst), (ins i64imm:$addr), []>,
+               Sched<[WriteI, ReadI]> {
+    let Defs = [X16,X17,NZCV];
+    let Size = 44;
   }
 
   // Load a signed global address from a special $auth_ptr$ stub slot.

llvm/lib/Target/AArch64/AArch64MCInstLower.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "AArch64MCInstLower.h"
+#include "AArch64MachineFunctionInfo.h"
 #include "MCTargetDesc/AArch64MCExpr.h"
 #include "Utils/AArch64BaseInfo.h"
 #include "llvm/CodeGen/AsmPrinter.h"
@@ -185,9 +186,12 @@ MCOperand AArch64MCInstLower::lowerSymbolOperandELF(const MachineOperand &MO,
                                                     MCSymbol *Sym) const {
   uint32_t RefFlags = 0;
 
-  if (MO.getTargetFlags() & AArch64II::MO_GOT)
-    RefFlags |= AArch64MCExpr::VK_GOT;
-  else if (MO.getTargetFlags() & AArch64II::MO_TLS) {
+  if (MO.getTargetFlags() & AArch64II::MO_GOT) {
+    const MachineFunction *MF = MO.getParent()->getParent()->getParent();
+    RefFlags |= (MF->getInfo<AArch64FunctionInfo>()->hasELFSignedGOT()
+                     ? AArch64MCExpr::VK_GOT_AUTH
+                     : AArch64MCExpr::VK_GOT);
+  } else if (MO.getTargetFlags() & AArch64II::MO_TLS) {
     TLSModel::Model Model;
     if (MO.isGlobal()) {
       const GlobalValue *GV = MO.getGlobal();

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

@@ -72,6 +72,18 @@ static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) {
   return Key == "b_key";
 }
 
+static bool hasELFSignedGOTHelper(const Function &F,
+                                  const AArch64Subtarget *STI) {
+  if (!Triple(STI->getTargetTriple()).isOSBinFormatELF())
+    return false;
+  const Module *M = F.getParent();
+  const auto *Flag = mdconst::extract_or_null<ConstantInt>(
+      M->getModuleFlag("ptrauth-elf-got"));
+  if (Flag && Flag->getZExtValue() == 1)
+    return true;
+  return false;
+}
+
 AArch64FunctionInfo::AArch64FunctionInfo(const Function &F,
                                          const AArch64Subtarget *STI) {
   // If we already know that the function doesn't have a redzone, set
@@ -80,6 +92,7 @@ AArch64FunctionInfo::AArch64FunctionInfo(const Function &F,
     HasRedZone = false;
   std::tie(SignReturnAddress, SignReturnAddressAll) = GetSignReturnAddress(F);
   SignWithBKey = ShouldSignWithBKey(F, *STI);
+  HasELFSignedGOT = hasELFSignedGOTHelper(F, STI);
   // TODO: skip functions that have no instrumented allocas for optimization
   IsMTETagged = F.hasFnAttribute(Attribute::SanitizeMemTag);
 

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -177,6 +177,11 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// SignWithBKey modifies the default PAC-RET mode to signing with the B key.
   bool SignWithBKey = false;
 
+  /// HasELFSignedGOT is true if the target binary format is ELF and the IR
+  /// module containing the corresponding function has "ptrauth-elf-got" flag
+  /// set to 1.
+  bool HasELFSignedGOT = false;
+
   /// SigningInstrOffset captures the offset of the PAC-RET signing instruction
   /// within the prologue, so it can be re-used for authentication in the
   /// epilogue when using PC as a second salt (FEAT_PAuth_LR)
@@ -509,6 +514,8 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
 
   bool shouldSignWithBKey() const { return SignWithBKey; }
 
+  bool hasELFSignedGOT() const { return HasELFSignedGOT; }
+
   MCSymbol *getSigningInstrLabel() const { return SignInstrLabel; }
   void setSigningInstrLabel(MCSymbol *Label) { SignInstrLabel = Label; }
 

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

@@ -897,6 +897,7 @@ class AArch64Operand : public MCParsedAsmOperand {
     if (DarwinRefKind == MCSymbolRefExpr::VK_PAGEOFF ||
         ELFRefKind == AArch64MCExpr::VK_LO12 ||
         ELFRefKind == AArch64MCExpr::VK_GOT_LO12 ||
+        ELFRefKind == AArch64MCExpr::VK_GOT_AUTH_LO12 ||
         ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12 ||
         ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12_NC ||
         ELFRefKind == AArch64MCExpr::VK_TPREL_LO12 ||
@@ -1008,19 +1009,20 @@ class AArch64Operand : public MCParsedAsmOperand {
     int64_t Addend;
     if (AArch64AsmParser::classifySymbolRef(Expr, ELFRefKind,
                                           DarwinRefKind, Addend)) {
-      return DarwinRefKind == MCSymbolRefExpr::VK_PAGEOFF
-          || DarwinRefKind == MCSymbolRefExpr::VK_TLVPPAGEOFF
-          || (DarwinRefKind == MCSymbolRefExpr::VK_GOTPAGEOFF && Addend == 0)
-          || ELFRefKind == AArch64MCExpr::VK_LO12
-          || ELFRefKind == AArch64MCExpr::VK_DTPREL_HI12
-          || ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12
-          || ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12_NC
-          || ELFRefKind == AArch64MCExpr::VK_TPREL_HI12
-          || ELFRefKind == AArch64MCExpr::VK_TPREL_LO12
-          || ELFRefKind == AArch64MCExpr::VK_TPREL_LO12_NC
-          || ELFRefKind == AArch64MCExpr::VK_TLSDESC_LO12
-          || ELFRefKind == AArch64MCExpr::VK_SECREL_HI12
-          || ELFRefKind == AArch64MCExpr::VK_SECREL_LO12;
+      return DarwinRefKind == MCSymbolRefExpr::VK_PAGEOFF ||
+             DarwinRefKind == MCSymbolRefExpr::VK_TLVPPAGEOFF ||
+             (DarwinRefKind == MCSymbolRefExpr::VK_GOTPAGEOFF && Addend == 0) ||
+             ELFRefKind == AArch64MCExpr::VK_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_GOT_AUTH_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_DTPREL_HI12 ||
+             ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12_NC ||
+             ELFRefKind == AArch64MCExpr::VK_TPREL_HI12 ||
+             ELFRefKind == AArch64MCExpr::VK_TPREL_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_TPREL_LO12_NC ||
+             ELFRefKind == AArch64MCExpr::VK_TLSDESC_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_SECREL_HI12 ||
+             ELFRefKind == AArch64MCExpr::VK_SECREL_LO12;
     }
 
     // If it's a constant, it should be a real immediate in range.
@@ -3309,6 +3311,7 @@ ParseStatus AArch64AsmParser::tryParseAdrpLabel(OperandVector &Operands) {
                DarwinRefKind != MCSymbolRefExpr::VK_TLVPPAGE &&
                ELFRefKind != AArch64MCExpr::VK_ABS_PAGE_NC &&
                ELFRefKind != AArch64MCExpr::VK_GOT_PAGE &&
+               ELFRefKind != AArch64MCExpr::VK_GOT_AUTH_PAGE &&
                ELFRefKind != AArch64MCExpr::VK_GOT_PAGE_LO15 &&
                ELFRefKind != AArch64MCExpr::VK_GOTTPREL_PAGE &&
                ELFRefKind != AArch64MCExpr::VK_TLSDESC_PAGE) {
@@ -4428,6 +4431,8 @@ bool AArch64AsmParser::parseSymbolicImmVal(const MCExpr *&ImmVal) {
                   .Case("got", AArch64MCExpr::VK_GOT_PAGE)
                   .Case("gotpage_lo15", AArch64MCExpr::VK_GOT_PAGE_LO15)
                   .Case("got_lo12", AArch64MCExpr::VK_GOT_LO12)
+                  .Case("got_auth", AArch64MCExpr::VK_GOT_AUTH_PAGE)
+                  .Case("got_auth_lo12", AArch64MCExpr::VK_GOT_AUTH_LO12)
                   .Case("gottprel", AArch64MCExpr::VK_GOTTPREL_PAGE)
                   .Case("gottprel_lo12", AArch64MCExpr::VK_GOTTPREL_LO12_NC)
                   .Case("gottprel_g1", AArch64MCExpr::VK_GOTTPREL_G1)
@@ -5801,6 +5806,7 @@ bool AArch64AsmParser::validateInstruction(MCInst &Inst, SMLoc &IDLoc,
 
         // Only allow these with ADDXri/ADDWri
         if ((ELFRefKind == AArch64MCExpr::VK_LO12 ||
+             ELFRefKind == AArch64MCExpr::VK_GOT_AUTH_LO12 ||
              ELFRefKind == AArch64MCExpr::VK_DTPREL_HI12 ||
              ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12 ||
              ELFRefKind == AArch64MCExpr::VK_DTPREL_LO12_NC ||

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -2967,7 +2967,9 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     }
 
     if (OpFlags & AArch64II::MO_GOT) {
-      I.setDesc(TII.get(AArch64::LOADgot));
+      I.setDesc(TII.get(MF.getInfo<AArch64FunctionInfo>()->hasELFSignedGOT()
+                            ? AArch64::LOADgotAUTH
+                            : AArch64::LOADgot));
       I.getOperand(1).setTargetFlags(OpFlags);
     } else if (TM.getCodeModel() == CodeModel::Large &&
                !TM.isPositionIndependent()) {