[BOLT][AArch64] Fix which PSign and PAuth variants are used (#120064)

- only the ones operating on LR should be marked
    with .cfi_cfi_negate_ra_state
- added support for fused PtrAuth and Ret instructions,
    e.g. RETAA.

Author: bgergely0

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -611,6 +611,21 @@ class MCPlusBuilder {
     return std::nullopt;
   }
 
+  virtual bool isPSignOnLR(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  virtual bool isPAuthOnLR(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  virtual bool isPAuthAndRet(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
   /// Returns the register used as a return address. Returns std::nullopt if
   /// not applicable, such as reading the return address from a system register
   /// or from the stack.
@@ -848,13 +863,6 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
     return false;
   }
-  virtual bool isPAuth(MCInst &Inst) const {
-    llvm_unreachable("not implemented");
-  }
-
-  virtual bool isPSign(MCInst &Inst) const {
-    llvm_unreachable("not implemented");
-  }
 
   virtual bool isCleanRegXOR(const MCInst &Inst) const {
     llvm_unreachable("not implemented");

bolt/include/bolt/Passes/InsertNegateRAStatePass.h

@@ -30,7 +30,9 @@ class InsertNegateRAState : public BinaryFunctionPass {
 
 private:
   /// Loops over all instructions and adds OpNegateRAState CFI
-  /// after any pointer signing or authenticating instructions.
+  /// after any pointer signing or authenticating instructions,
+  /// which operate on the LR, except fused ptrauth + ret instructions
+  /// (such as RETAA).
   /// Returns true, if any OpNegateRAState CFIs were added.
   bool addNegateRAStateAfterPacOrAuth(BinaryFunction &BF);
   /// Because states are tracked as MCAnnotations on individual instructions,

bolt/lib/Passes/InsertNegateRAStatePass.cpp

@@ -46,14 +46,16 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
   bool FirstIter = true;
   MCInst PrevInst;
   BinaryBasicBlock *PrevBB = nullptr;
+  // We need to iterate on BBs in the Layout order
+  // not in the order they are stored in the BF class.
   auto *Begin = BF.getLayout().block_begin();
   auto *End = BF.getLayout().block_end();
   for (auto *BB = Begin; BB != End; BB++) {
 
     // Support for function splitting:
-    // if two consecutive BBs are going to end up in different functions,
-    // we have to negate the RA State, so the new function starts with a Signed
-    // state.
+    // if two consecutive BBs with Signed state are going to end up in different
+    // functions, we have to add a OpNegateRAState to the beginning of the newly
+    // split function, so it starts with a Signed state.
     if (PrevBB != nullptr &&
         PrevBB->getFragmentNum() != (*BB)->getFragmentNum() &&
         BC.MIB->isRASigned(*((*BB)->begin()))) {
@@ -68,6 +70,8 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
         continue;
 
       if (!FirstIter) {
+        // Consecutive instructions with different RAState means we need to add
+        // a OpNegateRAState.
         if ((BC.MIB->isRASigned(PrevInst) && BC.MIB->isRAUnsigned(Inst)) ||
             (BC.MIB->isRAUnsigned(PrevInst) && BC.MIB->isRASigned(Inst))) {
 
@@ -90,7 +94,8 @@ bool InsertNegateRAState::addNegateRAStateAfterPacOrAuth(BinaryFunction &BF) {
   for (BinaryBasicBlock &BB : BF) {
     for (auto Iter = BB.begin(); Iter != BB.end(); ++Iter) {
       MCInst &Inst = *Iter;
-      if (BC.MIB->isPSign(Inst) || BC.MIB->isPAuth(Inst)) {
+      if (BC.MIB->isPSignOnLR(Inst) ||
+          (BC.MIB->isPAuthOnLR(Inst) && !BC.MIB->isPAuthAndRet(Inst))) {
         Iter = BF.addCFIInstruction(
             &BB, Iter + 1, MCCFIInstruction::createNegateRAState(nullptr));
         FoundAny = true;

bolt/lib/Passes/MarkRAStates.cpp

@@ -46,7 +46,8 @@ void MarkRAStates::runOnFunction(BinaryFunction &BF) {
   for (BinaryBasicBlock &BB : BF) {
     for (auto It = BB.begin(); It != BB.end(); ++It) {
       MCInst &Inst = *It;
-      if ((BC.MIB->isPSign(Inst) || BC.MIB->isPAuth(Inst)) &&
+      if ((BC.MIB->isPSignOnLR(Inst) ||
+           (BC.MIB->isPAuthOnLR(Inst) && !BC.MIB->isPAuthAndRet(Inst))) &&
           !BC.MIB->hasNegateRAState(Inst)) {
         // no .cfi_negate_ra_state attached to signing or authenticating instr
         // means, that this is a function with handwritten assembly, which might
@@ -71,7 +72,7 @@ void MarkRAStates::runOnFunction(BinaryFunction &BF) {
       if (BC.MIB->isCFI(Inst))
         continue;
 
-      if (BC.MIB->isPSign(Inst)) {
+      if (BC.MIB->isPSignOnLR(Inst)) {
         if (RAState) {
           // RA signing instructions should only follow unsigned RA state.
           BC.outs() << "BOLT-INFO: inconsistent RAStates in function "
@@ -80,7 +81,7 @@ void MarkRAStates::runOnFunction(BinaryFunction &BF) {
           return;
         }
         BC.MIB->setRASigning(Inst);
-      } else if (BC.MIB->isPAuth(Inst)) {
+      } else if (BC.MIB->isPAuthOnLR(Inst)) {
         if (!RAState) {
           // RA authenticating instructions should only follow signed RA state.
           BC.outs() << "BOLT-INFO: inconsistent RAStates in function "

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -231,6 +231,28 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
+  bool isPSignOnLR(const MCInst &Inst) const override {
+    std::optional<MCPhysReg> SignReg = getSignedReg(Inst);
+    return SignReg && *SignReg == AArch64::LR;
+  }
+
+  bool isPAuthOnLR(const MCInst &Inst) const override {
+    // LDR(A|B) should not be covered.
+    bool IsChecked;
+    std::optional<MCPhysReg> AuthReg =
+        getWrittenAuthenticatedReg(Inst, IsChecked);
+    return !IsChecked && AuthReg && *AuthReg == AArch64::LR;
+  }
+
+  bool isPAuthAndRet(const MCInst &Inst) const override {
+    return Inst.getOpcode() == AArch64::RETAA ||
+           Inst.getOpcode() == AArch64::RETAB ||
+           Inst.getOpcode() == AArch64::RETAASPPCi ||
+           Inst.getOpcode() == AArch64::RETABSPPCi ||
+           Inst.getOpcode() == AArch64::RETAASPPCr ||
+           Inst.getOpcode() == AArch64::RETABSPPCr;
+  }
+
   std::optional<MCPhysReg> getSignedReg(const MCInst &Inst) const override {
     switch (Inst.getOpcode()) {
     case AArch64::PACIA:
@@ -893,12 +915,6 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
     return false;
   }
-  bool isPAuth(MCInst &Inst) const override {
-    return Inst.getOpcode() == AArch64::AUTIASP;
-  }
-  bool isPSign(MCInst &Inst) const override {
-    return Inst.getOpcode() == AArch64::PACIASP;
-  }
 
   bool isRegToRegMove(const MCInst &Inst, MCPhysReg &From,
                       MCPhysReg &To) const override {