From 97ee8b5ffb61b2befeafb98930920de15b44800d Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko@accesssoftek.com>
Date: Thu, 26 Oct 2023 17:56:38 +0300
Subject: [PATCH] [AArch64][PAC] Refactor aarch64-ptrauth pass for better
 extensibility

Refactor Pointer Authentication pass in preparation for adding more
PAUTH_* pseudo instructions.

Fix handling of bundled TCRETURN* instructions (known to be generated
by KCFI). As other PAUTH_* instructions may need expansion even when
pac-ret is disabled, it is not generally possible to skip the whole
function easily. While this pass still does not support pac-ret being
enabled at the same time with KCFI, it should not crash if no checks are
actually emitted prior to TCRETURN instruction.
---
 .../lib/Target/AArch64/AArch64PointerAuth.cpp | 53 ++++++++++++-------
 1 file changed, 34 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 5d11f0d22574c..7576d2a899d1a 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -297,52 +297,67 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
 
 bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
   const auto *MFnI = MF.getInfo<AArch64FunctionInfo>();
-  if (!MFnI->shouldSignReturnAddress(true))
-    return false;
 
   Subtarget = &MF.getSubtarget<AArch64Subtarget>();
   TII = Subtarget->getInstrInfo();
   TRI = Subtarget->getRegisterInfo();
 
-  SmallVector<MachineBasicBlock::iterator> DeletedInstrs;
-  SmallVector<MachineBasicBlock::iterator> TailCallInstrs;
+  SmallVector<MachineBasicBlock::instr_iterator> PAuthPseudoInstrs;
+  SmallVector<MachineBasicBlock::instr_iterator> TailCallInstrs;
 
   bool Modified = false;
   bool HasAuthenticationInstrs = false;
 
   for (auto &MBB : MF) {
-    for (auto &MI : MBB) {
-      auto It = MI.getIterator();
+    // Using instr_iterator to catch unsupported bundled TCRETURN* instructions
+    // instead of just skipping them.
+    for (auto &MI : MBB.instrs()) {
       switch (MI.getOpcode()) {
       default:
+        // Bundled TCRETURN* instructions (such as created by KCFI)
+        // are not supported yet, but no support is required if no
+        // PAUTH_EPILOGUE instructions exist in the same function.
+        // Skip the BUNDLE instruction itself (actual bundled instructions
+        // follow it in the instruction list).
+        if (MI.isBundle())
+          continue;
         if (AArch64InstrInfo::isTailCallReturnInst(MI))
-          TailCallInstrs.push_back(It);
+          TailCallInstrs.push_back(MI.getIterator());
         break;
       case AArch64::PAUTH_PROLOGUE:
-        signLR(MF, It);
-        DeletedInstrs.push_back(It);
-        Modified = true;
-        break;
       case AArch64::PAUTH_EPILOGUE:
-        authenticateLR(MF, It);
-        DeletedInstrs.push_back(It);
-        Modified = true;
-        HasAuthenticationInstrs = true;
+        assert(!MI.isBundled());
+        PAuthPseudoInstrs.push_back(MI.getIterator());
         break;
       }
     }
   }
 
+  for (auto It : PAuthPseudoInstrs) {
+    switch (It->getOpcode()) {
+    case AArch64::PAUTH_PROLOGUE:
+      signLR(MF, It);
+      break;
+    case AArch64::PAUTH_EPILOGUE:
+      authenticateLR(MF, It);
+      HasAuthenticationInstrs = true;
+      break;
+    default:
+      llvm_unreachable("Unhandled opcode");
+    }
+    It->eraseFromParent();
+    Modified = true;
+  }
+
   // FIXME Do we need to emit any PAuth-related epilogue code at all
   //       when SCS is enabled?
   if (HasAuthenticationInstrs &&
       !MFnI->needsShadowCallStackPrologueEpilogue(MF)) {
-    for (auto TailCall : TailCallInstrs)
+    for (auto TailCall : TailCallInstrs) {
+      assert(!TailCall->isBundled() && "Not yet supported");
       Modified |= checkAuthenticatedLR(TailCall);
+    }
   }
 
-  for (auto MI : DeletedInstrs)
-    MI->eraseFromParent();
-
   return Modified;
 }