[AArch64][ELF] Support R_AARCH64_AUTH_ABS64 static relocation

The patch adds parser, MCExpr, and emitter support for the authenticated
pointer auth relocation.

In assembly, this is expressed using:
  .quad <symbol>@AUTH(<key>, <discriminator> [, addr])
For example:
  .quad _g3@AUTH(ib, 1234, addr)

The optional 'addr' specifier represents whether the generated pointer
authentication code will also include address diversity (by blending the
address of the storage location of the relocated pointer with the
user-specified constant discriminator).

The @AUTH expression lowers to R_AARCH64_AUTH_ABS64 ELF relocation.

The signing schema is encoded in the place of relocation to be applied
as follows:

```
| 63                | 62 | 61:60 | 59:48 |  47:32        | 31:0   |
| ----------------- | -- | ----- | ----- | ------------- | ------ |
| address diversity | 0  | key   | 0     | discriminator | addend |
```

See the following for details:
https://github.com/ARM-software/abi-aa/blob/main/pauthabielf64/pauthabielf64.rst#static-relocations

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

Co-authored-by: Ahmed Bougacha <ahmed@bougacha.org>
Co-authored-by: Peter Collingbourne <peter@pcc.me.uk>

llvm/docs/PointerAuth.md

@@ -300,3 +300,43 @@ instructions as such:
 * [``llvm.ptrauth.resign``](#llvm-ptrauth-resign): ``AUT*+PAC*``.  These are
   represented as a single pseudo-instruction in the backend to guarantee that
   the intermediate raw pointer value is not spilled and attackable.
+
+#### Assembly Representation
+
+At the assembly level,
+[Authenticated Relocations](#authenticated-global-relocation) are represented
+using the `@AUTH` modifier:
+
+```asm
+    .quad _target@AUTH(<key>,<discriminator>[,addr])
+```
+
+where:
+* `key` is the Armv8.3-A key identifier (`ia`, `ib`, `da`, `db`)
+* `discriminator` is the 16-bit unsigned discriminator value
+* `addr` signifies that the authenticated pointer is address-discriminated
+  (that is, that the relocation's target address is to be blended into the
+  `discriminator` before it is used in the sign operation.
+
+For example:
+```asm
+  _authenticated_reference_to_sym:
+    .quad _sym@AUTH(db,0)
+  _authenticated_reference_to_sym_addr_disc:
+    .quad _sym@AUTH(ia,12,addr)
+```
+
+#### ELF Object File Representation
+
+At the object file level,
+[Authenticated Relocations](#authenticated-global-relocation) are represented
+using the `R_AARCH64_AUTH_ABS64` relocation kind (with value `0xE100`).
+
+The signing schema is encoded in the place of relocation to be applied
+as follows:
+
+```
+| 63                | 62       | 61:60    | 59:48    |  47:32        | 31:0                |
+| ----------------- | -------- | -------- | -------- | ------------- | ------------------- |
+| address diversity | reserved | key      | reserved | discriminator | reserved for addend |
+```

llvm/include/llvm/BinaryFormat/ELFRelocs/AArch64.def

@@ -134,6 +134,7 @@ ELF_RELOC(R_AARCH64_TLS_DTPREL64,                    0x405)
 ELF_RELOC(R_AARCH64_TLS_TPREL64,                     0x406)
 ELF_RELOC(R_AARCH64_TLSDESC,                         0x407)
 ELF_RELOC(R_AARCH64_IRELATIVE,                       0x408)
+ELF_RELOC(R_AARCH64_AUTH_ABS64,                      0xe100)
 
 // ELF_RELOC(R_AARCH64_P32_NONE,                         0)
 ELF_RELOC(R_AARCH64_P32_ABS32,                       0x001)

llvm/lib/MC/MCExpr.cpp

@@ -738,17 +738,24 @@ static void AttemptToFoldSymbolOffsetDifference(
 /// They might look redundant, but this function can be used before layout
 /// is done (see the object streamer for example) and having the Asm argument
 /// lets us avoid relaxations early.
-static bool
-EvaluateSymbolicAdd(const MCAssembler *Asm, const MCAsmLayout *Layout,
-                    const SectionAddrMap *Addrs, bool InSet, const MCValue &LHS,
-                    const MCSymbolRefExpr *RHS_A, const MCSymbolRefExpr *RHS_B,
-                    int64_t RHS_Cst, MCValue &Res) {
+static bool EvaluateSymbolicAdd(const MCAssembler *Asm,
+                                const MCAsmLayout *Layout,
+                                const SectionAddrMap *Addrs, bool InSet,
+                                const MCValue &LHS, const MCValue &RHS,
+                                MCValue &Res) {
   // FIXME: This routine (and other evaluation parts) are *incredibly* sloppy
   // about dealing with modifiers. This will ultimately bite us, one day.
   const MCSymbolRefExpr *LHS_A = LHS.getSymA();
   const MCSymbolRefExpr *LHS_B = LHS.getSymB();
   int64_t LHS_Cst = LHS.getConstant();
 
+  const MCSymbolRefExpr *RHS_A = RHS.getSymA();
+  const MCSymbolRefExpr *RHS_B = RHS.getSymB();
+  int64_t RHS_Cst = RHS.getConstant();
+
+  if (LHS.getRefKind() != RHS.getRefKind())
+    return false;
+
   // Fold the result constant immediately.
   int64_t Result_Cst = LHS_Cst + RHS_Cst;
 
@@ -957,14 +964,19 @@ bool MCExpr::evaluateAsRelocatableImpl(MCValue &Res, const MCAssembler *Asm,
       case MCBinaryExpr::Sub:
         // Negate RHS and add.
         // The cast avoids undefined behavior if the constant is INT64_MIN.
-        return EvaluateSymbolicAdd(Asm, Layout, Addrs, InSet, LHSValue,
-                                   RHSValue.getSymB(), RHSValue.getSymA(),
-                                   -(uint64_t)RHSValue.getConstant(), Res);
+        return EvaluateSymbolicAdd(
+            Asm, Layout, Addrs, InSet, LHSValue,
+            MCValue::get(RHSValue.getSymB(), RHSValue.getSymA(),
+                         -(uint64_t)RHSValue.getConstant(),
+                         RHSValue.getRefKind()),
+            Res);
 
       case MCBinaryExpr::Add:
-        return EvaluateSymbolicAdd(Asm, Layout, Addrs, InSet, LHSValue,
-                                   RHSValue.getSymA(), RHSValue.getSymB(),
-                                   RHSValue.getConstant(), Res);
+        return EvaluateSymbolicAdd(
+            Asm, Layout, Addrs, InSet, LHSValue,
+            MCValue::get(RHSValue.getSymA(), RHSValue.getSymB(),
+                         RHSValue.getConstant(), RHSValue.getRefKind()),
+            Res);
       }
     }
 

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

@@ -180,6 +180,8 @@ class AArch64AsmParser : public MCTargetAsmParser {
   bool showMatchError(SMLoc Loc, unsigned ErrCode, uint64_t ErrorInfo,
                       OperandVector &Operands);
 
+  bool parseAuthExpr(const MCExpr *&Res, SMLoc &EndLoc);
+
   bool parseDirectiveArch(SMLoc L);
   bool parseDirectiveArchExtension(SMLoc L);
   bool parseDirectiveCPU(SMLoc L);
@@ -324,6 +326,8 @@ class AArch64AsmParser : public MCTargetAsmParser {
   unsigned validateTargetOperandClass(MCParsedAsmOperand &Op,
                                       unsigned Kind) override;
 
+  bool parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) override;
+
   static bool classifySymbolRef(const MCExpr *Expr,
                                 AArch64MCExpr::VariantKind &ELFRefKind,
                                 MCSymbolRefExpr::VariantKind &DarwinRefKind,
@@ -7451,6 +7455,112 @@ bool AArch64AsmParser::parseDirectiveSEHSaveAnyReg(SMLoc L, bool Paired,
   return false;
 }
 
+bool AArch64AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) {
+  // Try @AUTH expressions: they're more complex than the usual symbol variants.
+  if (!parseAuthExpr(Res, EndLoc))
+    return false;
+  return getParser().parsePrimaryExpr(Res, EndLoc, nullptr);
+}
+
+///  parseAuthExpr
+///  ::= _sym@AUTH(ib,123[,addr])
+///  ::= (_sym + 5)@AUTH(ib,123[,addr])
+///  ::= (_sym - 5)@AUTH(ib,123[,addr])
+bool AArch64AsmParser::parseAuthExpr(const MCExpr *&Res, SMLoc &EndLoc) {
+  MCAsmParser &Parser = getParser();
+  MCContext &Ctx = getContext();
+
+  AsmToken Tok = Parser.getTok();
+
+  // Look for '_sym@AUTH' ...
+  if (Tok.is(AsmToken::Identifier) && Tok.getIdentifier().endswith("@AUTH")) {
+    StringRef SymName = Tok.getIdentifier().drop_back(strlen("@AUTH"));
+    if (SymName.find('@') != StringRef::npos)
+      return TokError(
+          "combination of @AUTH with other modifiers not supported");
+    Res = MCSymbolRefExpr::create(Ctx.getOrCreateSymbol(SymName), Ctx);
+
+    Parser.Lex(); // Eat the identifier.
+  } else {
+    // ... or look for a more complex symbol reference, such as ...
+    SmallVector<AsmToken, 6> Tokens;
+
+    // ... '"_long sym"@AUTH' ...
+    if (Tok.is(AsmToken::String))
+      Tokens.resize(2);
+    // ... or '(_sym + 5)@AUTH'.
+    else if (Tok.is(AsmToken::LParen))
+      Tokens.resize(6);
+    else
+      return true;
+
+    if (Parser.getLexer().peekTokens(Tokens) != Tokens.size())
+      return true;
+
+    // In either case, the expression ends with '@' 'AUTH'.
+    if (Tokens[Tokens.size() - 2].isNot(AsmToken::At) ||
+        Tokens[Tokens.size() - 1].isNot(AsmToken::Identifier) ||
+        Tokens[Tokens.size() - 1].getIdentifier() != "AUTH")
+      return true;
+
+    if (Tok.is(AsmToken::String)) {
+      StringRef SymName;
+      if (Parser.parseIdentifier(SymName))
+        return true;
+      Res = MCSymbolRefExpr::create(Ctx.getOrCreateSymbol(SymName), Ctx);
+    } else {
+      if (Parser.parsePrimaryExpr(Res, EndLoc, nullptr))
+        return true;
+    }
+
+    Parser.Lex(); // '@'
+    Parser.Lex(); // 'AUTH'
+  }
+
+  // At this point, we encountered "<id>@AUTH". There is no fallback anymore.
+  if (parseToken(AsmToken::LParen, "expected '('"))
+    return true;
+
+  if (Parser.getTok().isNot(AsmToken::Identifier))
+    return TokError("expected key name");
+
+  StringRef KeyStr = Parser.getTok().getIdentifier();
+  auto KeyIDOrNone = AArch64StringToPACKeyID(KeyStr);
+  if (!KeyIDOrNone)
+    return TokError("invalid key '" + KeyStr + "'");
+  Parser.Lex();
+
+  if (parseToken(AsmToken::Comma, "expected ','"))
+    return true;
+
+  if (Parser.getTok().isNot(AsmToken::Integer))
+    return TokError("expected integer discriminator");
+  int64_t Discriminator = Parser.getTok().getIntVal();
+
+  if (!isUInt<16>(Discriminator))
+    return TokError("integer discriminator " + Twine(Discriminator) +
+                    " out of range [0, 0xFFFF]");
+  Parser.Lex();
+
+  bool UseAddressDiversity = false;
+  if (Parser.getTok().is(AsmToken::Comma)) {
+    Parser.Lex();
+    if (Parser.getTok().isNot(AsmToken::Identifier) ||
+        Parser.getTok().getIdentifier() != "addr")
+      return TokError("expected 'addr'");
+    UseAddressDiversity = true;
+    Parser.Lex();
+  }
+
+  EndLoc = Parser.getTok().getEndLoc();
+  if (parseToken(AsmToken::RParen, "expected ')'"))
+    return true;
+
+  Res = AArch64AuthMCExpr::create(Res, Discriminator, *KeyIDOrNone,
+                                  UseAddressDiversity, Ctx);
+  return false;
+}
+
 bool
 AArch64AsmParser::classifySymbolRef(const MCExpr *Expr,
                                     AArch64MCExpr::VariantKind &ELFRefKind,

llvm/lib/Target/AArch64/MCTargetDesc/AArch64AsmBackend.cpp

@@ -400,6 +400,19 @@ void AArch64AsmBackend::applyFixup(const MCAssembler &Asm, const MCFixup &Fixup,
                                    MutableArrayRef<char> Data, uint64_t Value,
                                    bool IsResolved,
                                    const MCSubtargetInfo *STI) const {
+  if (Fixup.getTargetKind() == FK_Data_8 && TheTriple.isOSBinFormatELF()) {
+    auto RefKind = static_cast<AArch64MCExpr::VariantKind>(Target.getRefKind());
+    AArch64MCExpr::VariantKind SymLoc = AArch64MCExpr::getSymbolLoc(RefKind);
+    if (SymLoc == AArch64AuthMCExpr::VK_AUTH ||
+        SymLoc == AArch64AuthMCExpr::VK_AUTHADDR) {
+      assert(Value == 0);
+      const auto *Expr = cast<AArch64AuthMCExpr>(Fixup.getValue());
+      Value = (uint64_t(Expr->getDiscriminator()) << 32) |
+              (uint64_t(Expr->getKey()) << 60) |
+              (uint64_t(Expr->hasAddressDiversity()) << 63);
+    }
+  }
+
   if (!Value)
     return; // Doesn't change encoding.
   unsigned Kind = Fixup.getKind();

llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFObjectWriter.cpp

@@ -207,8 +207,12 @@ unsigned AArch64ELFObjectWriter::getRelocType(MCContext &Ctx,
                         "ILP32 8 byte absolute data "
                         "relocation not supported (LP64 eqv: ABS64)");
         return ELF::R_AARCH64_NONE;
-      } else
+      } else {
+        if (RefKind == AArch64MCExpr::VK_AUTH ||
+            RefKind == AArch64MCExpr::VK_AUTHADDR)
+          return ELF::R_AARCH64_AUTH_ABS64;
         return ELF::R_AARCH64_ABS64;
+      }
     case AArch64::fixup_aarch64_add_imm12:
       if (RefKind == AArch64MCExpr::VK_DTPREL_HI12)
         return R_CLS(TLSLD_ADD_DTPREL_HI12);

llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCExpr.cpp

@@ -154,3 +154,47 @@ void AArch64MCExpr::fixELFSymbolsInTLSFixups(MCAssembler &Asm) const {
 
   fixELFSymbolsInTLSFixupsImpl(getSubExpr(), Asm);
 }
+
+const AArch64AuthMCExpr *AArch64AuthMCExpr::create(const MCExpr *Expr,
+                                                   uint16_t Discriminator,
+                                                   AArch64PACKey::ID Key,
+                                                   bool HasAddressDiversity,
+                                                   MCContext &Ctx) {
+  return new (Ctx)
+      AArch64AuthMCExpr(Expr, Discriminator, Key, HasAddressDiversity);
+}
+
+void AArch64AuthMCExpr::printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const {
+  bool WrapSubExprInParens = !isa<MCSymbolRefExpr>(getSubExpr());
+  if (WrapSubExprInParens)
+    OS << '(';
+  getSubExpr()->print(OS, MAI);
+  if (WrapSubExprInParens)
+    OS << ')';
+
+  OS << "@AUTH(" << AArch64PACKeyIDToString(Key) << ',' << Discriminator;
+  if (hasAddressDiversity())
+    OS << ",addr";
+  OS << ')';
+}
+
+void AArch64AuthMCExpr::visitUsedExpr(MCStreamer &Streamer) const {
+  Streamer.visitUsedExpr(*getSubExpr());
+}
+
+MCFragment *AArch64AuthMCExpr::findAssociatedFragment() const {
+  llvm_unreachable("FIXME: what goes here?");
+}
+
+bool AArch64AuthMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
+                                                  const MCAsmLayout *Layout,
+                                                  const MCFixup *Fixup) const {
+  if (!getSubExpr()->evaluateAsRelocatable(Res, Layout, Fixup))
+    return false;
+
+  if (Res.getSymB())
+    report_fatal_error("Auth relocation can't reference two symbols");
+
+  Res = MCValue::get(Res.getSymA(), nullptr, Res.getConstant(), getKind());
+  return true;
+}

llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCExpr.h

@@ -14,7 +14,9 @@
 #ifndef LLVM_LIB_TARGET_AARCH64_MCTARGETDESC_AARCH64MCEXPR_H
 #define LLVM_LIB_TARGET_AARCH64_MCTARGETDESC_AARCH64MCEXPR_H
 
+#include "Utils/AArch64BaseInfo.h"
 #include "llvm/MC/MCExpr.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 
 namespace llvm {
@@ -34,6 +36,8 @@ class AArch64MCExpr : public MCTargetExpr {
     VK_TPREL    = 0x007,
     VK_TLSDESC  = 0x008,
     VK_SECREL   = 0x009,
+    VK_AUTH     = 0x00a,
+    VK_AUTHADDR = 0x00b,
     VK_SymLocBits = 0x00f,
 
     // Variants specifying which part of the final address calculation is
@@ -116,6 +120,7 @@ class AArch64MCExpr : public MCTargetExpr {
   const MCExpr *Expr;
   const VariantKind Kind;
 
+protected:
   explicit AArch64MCExpr(const MCExpr *Expr, VariantKind Kind)
     : Expr(Expr), Kind(Kind) {}
 
@@ -171,6 +176,42 @@ class AArch64MCExpr : public MCTargetExpr {
     return E->getKind() == MCExpr::Target;
   }
 };
+
+class AArch64AuthMCExpr final : public AArch64MCExpr {
+  uint16_t Discriminator;
+  AArch64PACKey::ID Key;
+
+  explicit AArch64AuthMCExpr(const MCExpr *Expr, uint16_t Discriminator,
+                             AArch64PACKey::ID Key, bool HasAddressDiversity)
+      : AArch64MCExpr(Expr, HasAddressDiversity ? VK_AUTHADDR : VK_AUTH),
+        Discriminator(Discriminator), Key(Key) {}
+
+public:
+  static const AArch64AuthMCExpr *
+  create(const MCExpr *Expr, uint16_t Discriminator, AArch64PACKey::ID Key,
+         bool HasAddressDiversity, MCContext &Ctx);
+
+  AArch64PACKey::ID getKey() const { return Key; }
+  uint16_t getDiscriminator() const { return Discriminator; }
+  bool hasAddressDiversity() const { return getKind() == VK_AUTHADDR; }
+
+  void printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const override;
+
+  void visitUsedExpr(MCStreamer &Streamer) const override;
+
+  MCFragment *findAssociatedFragment() const override;
+
+  bool evaluateAsRelocatableImpl(MCValue &Res, const MCAsmLayout *Layout,
+                                 const MCFixup *Fixup) const override;
+
+  static bool classof(const MCExpr *E) {
+    return isa<AArch64MCExpr>(E) && classof(cast<AArch64MCExpr>(E));
+  }
+
+  static bool classof(const AArch64MCExpr *E) {
+    return E->getKind() == VK_AUTH || E->getKind() == VK_AUTHADDR;
+  }
+};
 } // end namespace llvm
 
 #endif