[AArch64][PAC] Factor out the emission of pointer check sequence (NFC)

[atrosinenko]

When pointer is authenticated or resigned, it may be required to
explicitly check the authenticated value to prevent introducing signing
or authentication oracles. While the check sequence is expensive in
general, a more efficient sequence can be emitted under specific
assumptions.

This commit factors out the emission of the code sequence to check the
authenticated pointer value in preparation for adding other variants
of checking code, as it is currently done when emitting tail calls.

[atrosinenko]

• [AArch64][PAC] Eliminate excessive MOVs when computing blend #115185
• [AArch64][PAC] Move emission of LR checks in tail calls to AsmPrinter #110705
• [AArch64][PAC] Factor out the emission of pointer check sequence (NFC) #110702 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

When pointer is authenticated or resigned, it may be required to
explicitly check the authenticated value to prevent introducing signing
or authentication oracles. While the check sequence is expensive in
general, a more efficient sequence can be emitted under specific
assumptions.

This commit factors out the emission of the code sequence to check the
authenticated pointer value in preparation for adding other variants
of checking code, as it is currently done when emitting tail calls.

---

Full diff: https://github.com/llvm/llvm-project/pull/110702.diff

1 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp (+91-77)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index a982ea67a0f279..6d2dd0ecbccf31 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -150,6 +150,12 @@ class AArch64AsmPrinter : public AsmPrinter {
   // Emit the sequence for BRA/BLRA (authenticate + branch/call).
   void emitPtrauthBranch(const MachineInstr *MI);
 
+  void emitPtrauthCheckAuthenticatedValue(Register TestedReg,
+                                          Register ScratchReg,
+                                          AArch64PACKey::ID Key,
+                                          bool ShouldTrap,
+                                          const MCSymbol *OnFailure);
+
   // Emit the sequence for AUT or AUTPAC.
   void emitPtrauthAuthResign(const MachineInstr *MI);
 
@@ -1719,45 +1725,37 @@ unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
   return AArch64::X17;
 }
 
-void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
-  const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
-
-  // We can expand AUT/AUTPAC into 3 possible sequences:
-  // - unchecked:
-  //      autia x16, x0
-  //      pacib x16, x1 ; if AUTPAC
+/// Emits a code sequence to check an authenticated pointer value.
+///
+/// If OnFailure argument is passed, jump there on check failure instead
+/// of proceeding to the next instruction (only if ShouldTrap is false).
+void AArch64AsmPrinter::emitPtrauthCheckAuthenticatedValue(
+    Register TestedReg, Register ScratchReg, AArch64PACKey::ID Key,
+    bool ShouldTrap, const MCSymbol *OnFailure) {
+  // Insert a sequence to check if authentication of TestedReg succeeded,
+  // such as:
   //
   // - checked and clearing:
-  //      mov x17, x0
-  //      movk x17, #disc, lsl #48
-  //      autia x16, x17
+  //      ; x16 is TestedReg, x17 is ScratchReg
   //      mov x17, x16
   //      xpaci x17
   //      cmp x16, x17
   //      b.eq Lsuccess
   //      mov x16, x17
   //      b Lend
-  //     Lsuccess:
-  //      mov x17, x1
-  //      movk x17, #disc, lsl #48
-  //      pacib x16, x17
-  //     Lend:
-  //   Where we only emit the AUT if we started with an AUT.
+  //    Lsuccess:
+  //      ; skipped if authentication failed
+  //    Lend:
+  //      ...
   //
   // - checked and trapping:
-  //      mov x17, x0
-  //      movk x17, #disc, lsl #48
-  //      autia x16, x0
   //      mov x17, x16
   //      xpaci x17
   //      cmp x16, x17
   //      b.eq Lsuccess
   //      brk #<0xc470 + aut key>
-  //     Lsuccess:
-  //      mov x17, x1
-  //      movk x17, #disc, lsl #48
-  //      pacib x16, x17 ; if AUTPAC
-  //   Where the b.eq skips over the trap if the PAC is valid.
+  //    Lsuccess:
+  //      ...
   //
   // This sequence is expensive, but we need more information to be able to
   // do better.
@@ -1770,6 +1768,71 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
   // Either way, we also don't always know whether TBI is enabled or not for
   // the specific target environment.
 
+  unsigned XPACOpc = getXPACOpcodeForKey(Key);
+
+  MCSymbol *SuccessSym = createTempSymbol("auth_success_");
+
+  //  mov Xscratch, Xtested
+  emitMovXReg(ScratchReg, TestedReg);
+
+  //  xpac(i|d) Xscratch
+  EmitToStreamer(MCInstBuilder(XPACOpc).addReg(ScratchReg).addReg(ScratchReg));
+
+  //  cmp Xtested, Xscratch
+  EmitToStreamer(MCInstBuilder(AArch64::SUBSXrs)
+                     .addReg(AArch64::XZR)
+                     .addReg(TestedReg)
+                     .addReg(ScratchReg)
+                     .addImm(0));
+
+  //  b.eq Lsuccess
+  EmitToStreamer(MCInstBuilder(AArch64::Bcc)
+                     .addImm(AArch64CC::EQ)
+                     .addExpr(MCSymbolRefExpr::create(SuccessSym, OutContext)));
+
+  if (ShouldTrap) {
+    assert(!OnFailure && "Cannot specify OnFailure with ShouldTrap");
+    // Trapping sequences do a 'brk'.
+    //  brk #<0xc470 + aut key>
+    EmitToStreamer(MCInstBuilder(AArch64::BRK).addImm(0xc470 | Key));
+  } else {
+    // Non-trapping checked sequences return the stripped result in TestedReg,
+    // skipping over success-only code (such as re-signing the pointer) if
+    // there is one.
+    // Note that this can introduce an authentication oracle (such as based on
+    // the high bits of the re-signed value).
+
+    // FIXME: Can we simply return the AUT result, already in TestedReg?
+    //  mov Xtested, Xscratch
+    emitMovXReg(TestedReg, ScratchReg);
+
+    if (OnFailure) {
+      //  b Lend
+      EmitToStreamer(
+          MCInstBuilder(AArch64::B)
+              .addExpr(MCSymbolRefExpr::create(OnFailure, OutContext)));
+    }
+  }
+
+  // If the auth check succeeds, we can continue.
+  // Lsuccess:
+  OutStreamer->emitLabel(SuccessSym);
+}
+
+void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
+  const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
+
+  // We expand AUT/AUTPAC into a sequence of the form
+  //
+  //      ; authenticate x16
+  //      ; check pointer in x16
+  //    Lsuccess:
+  //      ; sign x16 (if AUTPAC)
+  //    Lend:   ; if not trapping on failure
+  //
+  // with the checking sequence chosen depending on whether we should check
+  // the pointer and whether we should trap on failure.
+
   // By default, auth/resign sequences check for auth failures.
   bool ShouldCheck = true;
   // In the checked sequence, we only trap if explicitly requested.
@@ -1800,8 +1863,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
   uint64_t AUTDisc = MI->getOperand(1).getImm();
   unsigned AUTAddrDisc = MI->getOperand(2).getReg();
 
-  unsigned XPACOpc = getXPACOpcodeForKey(AUTKey);
-
   // Compute aut discriminator into x17
   assert(isUInt<16>(AUTDisc));
   unsigned AUTDiscReg = emitPtrauthDiscriminator(AUTDisc, AUTAddrDisc);
@@ -1824,59 +1885,12 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
 
   MCSymbol *EndSym = nullptr;
 
-  // Checked sequences do an additional strip-and-compare.
   if (ShouldCheck) {
-    MCSymbol *SuccessSym = createTempSymbol("auth_success_");
-
-    // XPAC has tied src/dst: use x17 as a temporary copy.
-    //  mov x17, x16
-    emitMovXReg(AArch64::X17, AArch64::X16);
-
-    //  xpaci x17
-    EmitToStreamer(
-        *OutStreamer,
-        MCInstBuilder(XPACOpc).addReg(AArch64::X17).addReg(AArch64::X17));
-
-    //  cmp x16, x17
-    EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::SUBSXrs)
-                                     .addReg(AArch64::XZR)
-                                     .addReg(AArch64::X16)
-                                     .addReg(AArch64::X17)
-                                     .addImm(0));
-
-    //  b.eq Lsuccess
-    EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::Bcc)
-                                     .addImm(AArch64CC::EQ)
-                                     .addExpr(MCSymbolRefExpr::create(
-                                         SuccessSym, OutContext)));
-
-    if (ShouldTrap) {
-      // Trapping sequences do a 'brk'.
-      //  brk #<0xc470 + aut key>
-      EmitToStreamer(*OutStreamer,
-                     MCInstBuilder(AArch64::BRK).addImm(0xc470 | AUTKey));
-    } else {
-      // Non-trapping checked sequences return the stripped result in x16,
-      // skipping over the PAC if there is one.
-
-      // FIXME: can we simply return the AUT result, already in x16? without..
-      //        ..traps this is usable as an oracle anyway, based on high bits
-      //  mov x17, x16
-      emitMovXReg(AArch64::X16, AArch64::X17);
-
-      if (IsAUTPAC) {
-        EndSym = createTempSymbol("resign_end_");
-
-        //  b Lend
-        EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::B)
-                                         .addExpr(MCSymbolRefExpr::create(
-                                             EndSym, OutContext)));
-      }
-    }
+    if (IsAUTPAC && !ShouldTrap)
+      EndSym = createTempSymbol("resign_end_");
 
-    // If the auth check succeeds, we can continue.
-    // Lsuccess:
-    OutStreamer->emitLabel(SuccessSym);
+    emitPtrauthCheckAuthenticatedValue(AArch64::X16, AArch64::X17, AUTKey,
+                                       ShouldTrap, EndSym);
   }
 
   // We already emitted unchecked and checked-but-non-trapping AUTs.

[asl]

@kbeyls @ahmedbougacha Will you please a look?

[kovdan01]

LGTM