[llvm] [BOLT][AArch64] Support for pointer authentication (PR #117578)

Gergely Bálint via llvm-commits llvm-commits at lists.llvm.org
Thu Nov 28 00:54:36 PST 2024


https://github.com/bgergely0 updated https://github.com/llvm/llvm-project/pull/117578

>From 61efea94c53c467567c1afbc5bb07bf381b7d48f Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Tue, 19 Nov 2024 09:43:25 +0100
Subject: [PATCH 1/7] [BOLT] Recognize paciasp and autiasp instructions

---
 bolt/include/bolt/Core/MCPlusBuilder.h           | 7 +++++++
 bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 3634fed9757ceb..b720a56806bdb7 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -648,6 +648,13 @@ 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");
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 7e08e5c81d26ff..c07a93c0a1460b 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -316,6 +316,12 @@ 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 {

>From 0353abbe4284960231b9ff6b689e0e5c7befbf18 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Mon, 25 Nov 2024 11:48:15 +0100
Subject: [PATCH 2/7] [BOLT] Support for OpNegateRAState - first half

    - when reading binary, drop OpNegateRAState CFIs
    - change CFI State calculation to ignore OpNegateRAState CFIs
---
 bolt/lib/Core/BinaryBasicBlock.cpp | 6 +++++-
 bolt/lib/Core/BinaryFunction.cpp   | 1 +
 bolt/lib/Core/Exceptions.cpp       | 6 ++++--
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Core/BinaryBasicBlock.cpp b/bolt/lib/Core/BinaryBasicBlock.cpp
index 2a2192b79bb4bf..bc1d3f112f2ed2 100644
--- a/bolt/lib/Core/BinaryBasicBlock.cpp
+++ b/bolt/lib/Core/BinaryBasicBlock.cpp
@@ -201,7 +201,11 @@ int32_t BinaryBasicBlock::getCFIStateAtInstr(const MCInst *Instr) const {
       InstrSeen = (&Inst == Instr);
       continue;
     }
-    if (Function->getBinaryContext().MIB->isCFI(Inst)) {
+    // Fix: ignoring OpNegateRAState CFIs here, as they dont have a "State"
+    // number associated with them.
+    if (Function->getBinaryContext().MIB->isCFI(Inst) &&
+        (Function->getCFIFor(Inst)->getOperation() !=
+         MCCFIInstruction::OpNegateRAState)) {
       LastCFI = &Inst;
       break;
     }
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 5da777411ba7a1..6eea398ba95ad3 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -2596,6 +2596,7 @@ struct CFISnapshot {
   void advanceTo(int32_t State) {
     for (int32_t I = CurState, E = State; I != E; ++I) {
       const MCCFIInstruction &Instr = FDE[I];
+      assert(Instr.getOperation() != MCCFIInstruction::OpNegateRAState);
       if (Instr.getOperation() != MCCFIInstruction::OpRestoreState) {
         update(Instr, I);
         continue;
diff --git a/bolt/lib/Core/Exceptions.cpp b/bolt/lib/Core/Exceptions.cpp
index 0b2e63b8ca6a79..f528db8449dbd2 100644
--- a/bolt/lib/Core/Exceptions.cpp
+++ b/bolt/lib/Core/Exceptions.cpp
@@ -632,8 +632,10 @@ bool CFIReaderWriter::fillCFIInfoFor(BinaryFunction &Function) const {
       // DW_CFA_GNU_window_save and DW_CFA_GNU_NegateRAState just use the same
       // id but mean different things. The latter is used in AArch64.
       if (Function.getBinaryContext().isAArch64()) {
-        Function.addCFIInstruction(
-            Offset, MCCFIInstruction::createNegateRAState(nullptr));
+        // Fix: not adding OpNegateRAState since the location they are needed
+        // depends on the order of BasicBlocks, which changes during
+        // optimizations. They are generated in InsertNegateRAStatePass after
+        // optimizations instead.
         break;
       }
       if (opts::Verbosity >= 1)

>From 6217b6664dfe12ce24bed465d30f4607b9b5040f Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Mon, 25 Nov 2024 11:57:15 +0100
Subject: [PATCH 3/7] [BOLT] Support for OpNegateRAState - second half

    - create InsertNegateRAStatePass
    - this pass explores the CFG, and finds if a BB has
        signed or unsigned return address state
    - after exploration, it inserts OpNegateRAState at the end of BBs,
        where the following BB has a different RA State (or where RA
        State boundaries are)
---
 bolt/include/bolt/Core/BinaryBasicBlock.h     |  16 ++
 .../bolt/Passes/InsertNegateRAStatePass.h     |  33 +++
 bolt/lib/Passes/CMakeLists.txt                |   1 +
 bolt/lib/Passes/InsertNegateRAStatePass.cpp   | 247 ++++++++++++++++++
 bolt/lib/Rewrite/BinaryPassManager.cpp        |   3 +
 5 files changed, 300 insertions(+)
 create mode 100644 bolt/include/bolt/Passes/InsertNegateRAStatePass.h
 create mode 100644 bolt/lib/Passes/InsertNegateRAStatePass.cpp

diff --git a/bolt/include/bolt/Core/BinaryBasicBlock.h b/bolt/include/bolt/Core/BinaryBasicBlock.h
index 25cccc4edecf68..fe48d3b4f740d6 100644
--- a/bolt/include/bolt/Core/BinaryBasicBlock.h
+++ b/bolt/include/bolt/Core/BinaryBasicBlock.h
@@ -37,6 +37,11 @@ class JumpTable;
 
 class BinaryBasicBlock {
 public:
+  enum class RAStateEnum : char {
+    Unknown, /// Not discovered yet
+    Signed,
+    Unsigned,
+  };
   /// Profile execution information for a given edge in CFG.
   ///
   /// If MispredictedCount equals COUNT_INFERRED, then we have a profile
@@ -350,6 +355,17 @@ class BinaryBasicBlock {
                                                       BranchInfo.end());
   }
 
+  RAStateEnum RAState{RAStateEnum::Unknown};
+  void setRASigned() { RAState = RAStateEnum::Signed; }
+  bool isRAStateUnknown() { return RAState == RAStateEnum::Unknown; }
+  bool isRAStateSigned() { return RAState == RAStateEnum::Signed; }
+  /// Unsigned should only overwrite Unknown state, and not Signed
+  void setRAUnsigned() {
+    if (RAState == RAStateEnum::Unknown) {
+      RAState = RAStateEnum::Unsigned;
+    }
+  }
+
   /// Get instruction at given index.
   MCInst &getInstructionAtIndex(unsigned Index) { return Instructions[Index]; }
 
diff --git a/bolt/include/bolt/Passes/InsertNegateRAStatePass.h b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h
new file mode 100644
index 00000000000000..ac1a2f4cf7327c
--- /dev/null
+++ b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h
@@ -0,0 +1,33 @@
+#ifndef BOLT_PASSES_INSERT_NEGATE_RA_STATE_PASS
+#define BOLT_PASSES_INSERT_NEGATE_RA_STATE_PASS
+
+#include "bolt/Passes/BinaryPasses.h"
+#include <stack>
+
+namespace llvm {
+namespace bolt {
+
+class InsertNegateRAState : public BinaryFunctionPass {
+public:
+  explicit InsertNegateRAState() : BinaryFunctionPass(false) {}
+
+  const char *getName() const override { return "insert-negate-ra-state-pass"; }
+
+  /// Pass entry point
+  Error runOnFunctions(BinaryContext &BC) override;
+  void runOnFunction(BinaryFunction &BF);
+  bool addNegateRAStateAfterPacOrAuth(BinaryFunction &BF);
+  bool BBhasAUTH(BinaryContext &BC, BinaryBasicBlock *BB);
+  bool BBhasSIGN(BinaryContext &BC, BinaryBasicBlock *BB);
+  void explore_call_graph(BinaryContext &BC, BinaryBasicBlock *BB);
+  void process_signed_BB(BinaryContext &BC, BinaryBasicBlock *BB,
+                         std::stack<BinaryBasicBlock *> *SignedStack,
+                         std::stack<BinaryBasicBlock *> *UnsignedStack);
+  void process_unsigned_BB(BinaryContext &BC, BinaryBasicBlock *BB,
+                           std::stack<BinaryBasicBlock *> *SignedStack,
+                           std::stack<BinaryBasicBlock *> *UnsignedStack);
+};
+
+} // namespace bolt
+} // namespace llvm
+#endif
diff --git a/bolt/lib/Passes/CMakeLists.txt b/bolt/lib/Passes/CMakeLists.txt
index 1c1273b3d2420d..d7864e30305116 100644
--- a/bolt/lib/Passes/CMakeLists.txt
+++ b/bolt/lib/Passes/CMakeLists.txt
@@ -17,6 +17,7 @@ add_llvm_library(LLVMBOLTPasses
   IdenticalCodeFolding.cpp
   IndirectCallPromotion.cpp
   Inliner.cpp
+  InsertNegateRAStatePass.cpp
   Instrumentation.cpp
   JTFootprintReduction.cpp
   LongJmp.cpp
diff --git a/bolt/lib/Passes/InsertNegateRAStatePass.cpp b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
new file mode 100644
index 00000000000000..c5db6df3a2b606
--- /dev/null
+++ b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
@@ -0,0 +1,247 @@
+#include "bolt/Passes/InsertNegateRAStatePass.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Utils/CommandLineOpts.h"
+#include <cstdlib>
+#include <fstream>
+#include <iterator>
+
+using namespace llvm;
+
+namespace llvm {
+namespace bolt {
+
+void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+
+  if (BF.getState() == BinaryFunction::State::Empty) {
+    return;
+  }
+
+  if (BF.getState() != BinaryFunction::State::CFG &&
+      BF.getState() != BinaryFunction::State::CFG_Finalized) {
+    BC.errs() << "BOLT-WARNING: No CFG for " << BF.getPrintName()
+              << " in InsertNegateRAStatePass\n";
+    return;
+  }
+
+  if (BF.getState() == BinaryFunction::State::CFG_Finalized) {
+    BC.errs() << "BOLT-WARNING: CFG finalized for " << BF.getPrintName()
+              << " in InsertNegateRAStatePass\n";
+    return;
+  }
+
+  if (BF.isIgnored())
+    return;
+
+  if (!addNegateRAStateAfterPacOrAuth(BF)) {
+    // none inserted, function doesn't need more work
+    return;
+  }
+
+  auto FirstBB = BF.begin();
+  explore_call_graph(BC, &(*FirstBB));
+
+  // We have to do the walk again, starting from any undiscovered autiasp
+  // instructions, because some autiasp might not be reachable because of
+  // indirect branches but we know that autiasp block should have a Signed
+  // state, so we can work out other Unkown states starting from these nodes.
+  for (BinaryBasicBlock &BB : BF) {
+    if (BBhasAUTH(BC, &BB) && BB.isRAStateUnknown()) {
+      BB.setRASigned();
+      explore_call_graph(BC, &BB);
+    }
+  }
+
+  // insert negateRAState-s where there is a State boundary:
+  // that is: two consecutive BBs have different RA State
+  BinaryFunction::iterator PrevBB;
+  bool FirstIter = true;
+  for (auto BB = BF.begin(); BB != BF.end(); ++BB) {
+    if (!FirstIter) {
+      if ((PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           (*BB).RAState == BinaryBasicBlock::RAStateEnum::Unsigned &&
+           !BBhasAUTH(BC, &(*PrevBB))) ||
+          (PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           (*BB).RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           BBhasAUTH(BC, &(*PrevBB)))) {
+        auto InstRevIter = PrevBB->getLastNonPseudo();
+        MCInst LastNonPseudo = *InstRevIter;
+        auto InstIter = InstRevIter.base();
+        BF.addCFIInstruction(&(*PrevBB), InstIter,
+                             MCCFIInstruction::createNegateRAState(nullptr));
+      }
+    } else {
+      FirstIter = false;
+    }
+    PrevBB = BB;
+  }
+}
+
+void InsertNegateRAState::explore_call_graph(BinaryContext &BC,
+                                             BinaryBasicBlock *BB) {
+  std::stack<BinaryBasicBlock *> SignedStack;
+  std::stack<BinaryBasicBlock *> UnsignedStack;
+
+  // start according to the first BB
+  if (BBhasSIGN(BC, BB)) {
+    SignedStack.push(BB);
+    process_signed_BB(BC, BB, &SignedStack, &UnsignedStack);
+  } else {
+    UnsignedStack.push(BB);
+    process_unsigned_BB(BC, BB, &SignedStack, &UnsignedStack);
+  }
+
+  while (!(SignedStack.empty() && UnsignedStack.empty())) {
+    if (!SignedStack.empty()) {
+      BB = SignedStack.top();
+      SignedStack.pop();
+      process_signed_BB(BC, BB, &SignedStack, &UnsignedStack);
+    } else if (!UnsignedStack.empty()) {
+      BB = UnsignedStack.top();
+      UnsignedStack.pop();
+      process_unsigned_BB(BC, BB, &SignedStack, &UnsignedStack);
+    }
+  }
+}
+void InsertNegateRAState::process_signed_BB(
+    BinaryContext &BC, BinaryBasicBlock *BB,
+    std::stack<BinaryBasicBlock *> *SignedStack,
+    std::stack<BinaryBasicBlock *> *UnsignedStack) {
+
+  BB->setRASigned();
+
+  if (BBhasAUTH(BC, BB)) {
+    // successors of block with autiasp are stored in the Unsigned Stack
+    for (BinaryBasicBlock *Succ : BB->successors()) {
+      if (Succ->getFunction() == BB->getFunction() &&
+          Succ->isRAStateUnknown()) {
+        UnsignedStack->push(Succ);
+      }
+    }
+  } else {
+    for (BinaryBasicBlock *Succ : BB->successors()) {
+      if (Succ->getFunction() == BB->getFunction() &&
+          !Succ->isRAStateSigned()) {
+        SignedStack->push(Succ);
+      }
+    }
+  }
+  // process predecessors
+  if (BBhasSIGN(BC, BB)) {
+    for (BinaryBasicBlock *Pred : BB->predecessors()) {
+      if (Pred->getFunction() == BB->getFunction() &&
+          Pred->isRAStateUnknown()) {
+        UnsignedStack->push(Pred);
+      }
+    }
+  } else {
+    for (BinaryBasicBlock *Pred : BB->predecessors()) {
+      if (Pred->getFunction() == BB->getFunction() &&
+          !Pred->isRAStateSigned()) {
+        SignedStack->push(Pred);
+      }
+    }
+  }
+}
+
+void InsertNegateRAState::process_unsigned_BB(
+    BinaryContext &BC, BinaryBasicBlock *BB,
+    std::stack<BinaryBasicBlock *> *SignedStack,
+    std::stack<BinaryBasicBlock *> *UnsignedStack) {
+
+  BB->setRAUnsigned();
+
+  if (BBhasSIGN(BC, BB)) {
+    BB->setRASigned();
+    // successors of block with paciasp are stored in the Signed Stack
+    for (BinaryBasicBlock *Succ : BB->successors()) {
+      if (Succ->getFunction() == BB->getFunction() &&
+          !Succ->isRAStateSigned()) {
+        SignedStack->push(Succ);
+      }
+    }
+  } else {
+    for (BinaryBasicBlock *Succ : BB->successors()) {
+      if (Succ->getFunction() == BB->getFunction() &&
+          Succ->isRAStateUnknown()) {
+        UnsignedStack->push(Succ);
+      }
+    }
+  }
+
+  // process predecessors
+  if (BBhasAUTH(BC, BB)) {
+    BB->setRASigned();
+    for (BinaryBasicBlock *Pred : BB->predecessors()) {
+      if (Pred->getFunction() == BB->getFunction() &&
+          !Pred->isRAStateSigned()) {
+        SignedStack->push(Pred);
+      }
+    }
+  } else {
+    for (BinaryBasicBlock *Pred : BB->predecessors()) {
+      if (Pred->getFunction() == BB->getFunction() &&
+          Pred->isRAStateUnknown()) {
+        UnsignedStack->push(Pred);
+      }
+    }
+  }
+}
+
+bool InsertNegateRAState::BBhasAUTH(BinaryContext &BC, BinaryBasicBlock *BB) {
+  for (auto Iter = BB->begin(); Iter != BB->end(); ++Iter) {
+    MCInst Inst = *Iter;
+    if (BC.MIB->isPAuth(Inst)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool InsertNegateRAState::BBhasSIGN(BinaryContext &BC, BinaryBasicBlock *BB) {
+  for (auto Iter = BB->begin(); Iter != BB->end(); ++Iter) {
+    MCInst Inst = *Iter;
+    if (BC.MIB->isPSign(Inst)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool InsertNegateRAState::addNegateRAStateAfterPacOrAuth(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+  bool FoundAny = false;
+  for (BinaryBasicBlock &BB : BF) {
+    for (auto Iter = BB.begin(); Iter != BB.end(); ++Iter) {
+      MCInst Inst = *Iter;
+      if (BC.MIB->isPSign(Inst)) {
+        Iter = BF.addCFIInstruction(
+            &BB, Iter + 1, MCCFIInstruction::createNegateRAState(nullptr));
+        FoundAny = true;
+      }
+
+      if (BC.MIB->isPAuth(Inst)) {
+        Iter = BF.addCFIInstruction(
+            &BB, Iter + 1, MCCFIInstruction::createNegateRAState(nullptr));
+        FoundAny = true;
+      }
+    }
+  }
+  return FoundAny;
+}
+
+Error InsertNegateRAState::runOnFunctions(BinaryContext &BC) {
+  ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+    runOnFunction(BF);
+  };
+
+  ParallelUtilities::runOnEachFunction(
+      BC, ParallelUtilities::SchedulingPolicy::SP_TRIVIAL, WorkFun, nullptr,
+      "InsertNegateRAStatePass");
+
+  return Error::success();
+}
+
+} // end namespace bolt
+} // end namespace llvm
diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp
index b0906041833484..d11321d8ef93ae 100644
--- a/bolt/lib/Rewrite/BinaryPassManager.cpp
+++ b/bolt/lib/Rewrite/BinaryPassManager.cpp
@@ -20,6 +20,7 @@
 #include "bolt/Passes/IdenticalCodeFolding.h"
 #include "bolt/Passes/IndirectCallPromotion.h"
 #include "bolt/Passes/Inliner.h"
+#include "bolt/Passes/InsertNegateRAStatePass.h"
 #include "bolt/Passes/Instrumentation.h"
 #include "bolt/Passes/JTFootprintReduction.h"
 #include "bolt/Passes/LongJmp.h"
@@ -499,6 +500,8 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
     // targets. No extra instructions after this pass, otherwise we may have
     // relocations out of range and crash during linking.
     Manager.registerPass(std::make_unique<LongJmpPass>(PrintLongJmp));
+
+    Manager.registerPass(std::make_unique<InsertNegateRAState>());
   }
 
   // This pass should always run last.*

>From 9f3d9bdf8238c90d426d2a93944dc9a09a39f9a6 Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at arm.com>
Date: Mon, 25 Sep 2023 20:54:36 +0200
Subject: [PATCH 4/7] [GadgetScanner/pacret] Handle noreturn functions.

---
 bolt/include/bolt/Core/MCPlusBuilder.h        | 23 ++++++++++++++
 bolt/include/bolt/Utils/CommandLineOpts.h     |  7 +++++
 bolt/lib/Core/BinaryFunction.cpp              | 10 ++++--
 bolt/lib/Passes/BinaryPasses.cpp              | 26 +++++++---------
 bolt/lib/Rewrite/RewriteInstance.cpp          | 21 ++++++-------
 bolt/lib/Utils/CommandLineOpts.cpp            | 24 +++++++++++++-
 bolt/test/gadget-scanner/gs-pacret-noreturn.s | 31 +++++++++++++++++++
 7 files changed, 112 insertions(+), 30 deletions(-)
 create mode 100644 bolt/test/gadget-scanner/gs-pacret-noreturn.s

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index b720a56806bdb7..937c5c668c9b27 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -16,6 +16,7 @@
 
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Core/Relocation.h"
+#include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/StringMap.h"
@@ -27,6 +28,7 @@
 #include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -546,6 +548,27 @@ class MCPlusBuilder {
     return Analysis->isCall(Inst) || isTailCall(Inst);
   }
 
+  virtual std::optional<StringRef> getCalleeName(const MCInst &Inst) const {
+    assert(isCall(Inst));
+    if (MCPlus::getNumPrimeOperands(Inst) != 1 || !Inst.getOperand(0).isExpr())
+      return {};
+
+    const MCSymbol *CalleeSymbol = getTargetSymbol(Inst);
+    assert(CalleeSymbol != nullptr);
+    return CalleeSymbol->getName();
+  }
+
+  virtual bool isNoReturnCall(const MCInst &Inst) const {
+    if (!isCall(Inst))
+      return false;
+    auto calleeName = getCalleeName(Inst);
+    if (calleeName)
+      for (std::string &Name : opts::AssumeNoReturnFunctions)
+        if (calleeName->equals(Name))
+          return true;
+    return false;
+  }
+
   virtual bool isReturn(const MCInst &Inst) const {
     return Analysis->isReturn(Inst);
   }
diff --git a/bolt/include/bolt/Utils/CommandLineOpts.h b/bolt/include/bolt/Utils/CommandLineOpts.h
index 04bf7db5de9527..dcdcb0ed84cfc9 100644
--- a/bolt/include/bolt/Utils/CommandLineOpts.h
+++ b/bolt/include/bolt/Utils/CommandLineOpts.h
@@ -62,6 +62,13 @@ extern llvm::cl::opt<bool> TimeOpts;
 extern llvm::cl::opt<bool> UseOldText;
 extern llvm::cl::opt<bool> UpdateDebugSections;
 
+extern llvm::cl::list<std::string> AssumeNoReturnFunctions;
+extern llvm::cl::opt<std::string> AssumeNoReturnFunctionsFile;
+
+/// Reads names from FunctionNamesFile and adds them to FunctionNames.
+void populateFunctionNames(const llvm::cl::opt<std::string> &FunctionNamesFile,
+                           llvm::cl::list<std::string> &FunctionNames);
+
 // The default verbosity level (0) is pretty terse, level 1 is fairly
 // verbose and usually prints some informational message for every
 // function processed.  Level 2 is for the noisiest of messages and
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 6eea398ba95ad3..34dd5136f9918e 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -2155,7 +2155,8 @@ Error BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
       addCFIPlaceholders(Offset, InsertBB);
     }
 
-    const bool IsBlockEnd = MIB->isTerminator(Instr);
+    const bool IsBlockEnd =
+        MIB->isTerminator(Instr) || MIB->isNoReturnCall(Instr);
     IsLastInstrNop = MIB->isNoop(Instr);
     if (!IsLastInstrNop)
       LastInstrOffset = Offset;
@@ -2242,8 +2243,11 @@ Error BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
       //
       // Conditional tail call is a special case since we don't add a taken
       // branch successor for it.
-      IsPrevFT = !MIB->isTerminator(*LastInstr) ||
-                 MIB->getConditionalTailCall(*LastInstr);
+      if (MIB->isNoReturnCall(*LastInstr))
+        IsPrevFT = false;
+      else
+        IsPrevFT = !MIB->isTerminator(*LastInstr) ||
+                   MIB->getConditionalTailCall(*LastInstr);
     } else if (BB->succ_size() == 1) {
       IsPrevFT = MIB->isConditionalBranch(*LastInstr);
     } else {
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 03d3dd75a03368..0cce08f650ae74 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -1852,17 +1852,16 @@ Error InlineMemcpy::runOnFunctions(BinaryContext &BC) {
       for (auto II = BB.begin(); II != BB.end(); ++II) {
         MCInst &Inst = *II;
 
-        if (!BC.MIB->isCall(Inst) || MCPlus::getNumPrimeOperands(Inst) != 1 ||
-            !Inst.getOperand(0).isExpr())
+        if (!BC.MIB->isCall(Inst))
           continue;
-
-        const MCSymbol *CalleeSymbol = BC.MIB->getTargetSymbol(Inst);
-        if (CalleeSymbol->getName() != "memcpy" &&
-            CalleeSymbol->getName() != "memcpy at PLT" &&
-            CalleeSymbol->getName() != "_memcpy8")
+        std::optional<StringRef> CalleeName = BC.MIB->getCalleeName(Inst);
+        if (!CalleeName)
+          continue;
+        if (*CalleeName != "memcpy" && *CalleeName != "memcpy at PLT" &&
+            *CalleeName != "_memcpy8")
           continue;
 
-        const bool IsMemcpy8 = (CalleeSymbol->getName() == "_memcpy8");
+        const bool IsMemcpy8 = (*CalleeName == "_memcpy8");
         const bool IsTailCall = BC.MIB->isTailCall(Inst);
 
         const InstructionListType NewCode =
@@ -1951,13 +1950,12 @@ Error SpecializeMemcpy1::runOnFunctions(BinaryContext &BC) {
       for (auto II = CurBB->begin(); II != CurBB->end(); ++II) {
         MCInst &Inst = *II;
 
-        if (!BC.MIB->isCall(Inst) || MCPlus::getNumPrimeOperands(Inst) != 1 ||
-            !Inst.getOperand(0).isExpr())
+        if (!BC.MIB->isCall(Inst))
           continue;
-
-        const MCSymbol *CalleeSymbol = BC.MIB->getTargetSymbol(Inst);
-        if (CalleeSymbol->getName() != "memcpy" &&
-            CalleeSymbol->getName() != "memcpy at PLT")
+        std::optional<StringRef> CalleeName = BC.MIB->getCalleeName(Inst);
+        if (!CalleeName)
+          continue;
+        if (*CalleeName != "memcpy" && *CalleeName != "memcpy at PLT")
           continue;
 
         if (BC.MIB->isTailCall(Inst))
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 7059a3dd231099..5352a6ec8a8ebc 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -269,6 +269,9 @@ MCPlusBuilder *createMCPlusBuilder(const Triple::ArchType Arch,
                                    const MCInstrInfo *Info,
                                    const MCRegisterInfo *RegInfo,
                                    const MCSubtargetInfo *STI) {
+  opts::populateFunctionNames(opts::AssumeNoReturnFunctionsFile,
+                              opts::AssumeNoReturnFunctions);
+
 #ifdef X86_AVAILABLE
   if (Arch == Triple::x86_64)
     return createX86MCPlusBuilder(Analysis, Info, RegInfo, STI);
@@ -2929,18 +2932,12 @@ void RewriteInstance::handleRelocation(const SectionRef &RelocatedSection,
 
 void RewriteInstance::selectFunctionsToProcess() {
   // Extend the list of functions to process or skip from a file.
-  auto populateFunctionNames = [](cl::opt<std::string> &FunctionNamesFile,
-                                  cl::list<std::string> &FunctionNames) {
-    if (FunctionNamesFile.empty())
-      return;
-    std::ifstream FuncsFile(FunctionNamesFile, std::ios::in);
-    std::string FuncName;
-    while (std::getline(FuncsFile, FuncName))
-      FunctionNames.push_back(FuncName);
-  };
-  populateFunctionNames(opts::FunctionNamesFile, opts::ForceFunctionNames);
-  populateFunctionNames(opts::SkipFunctionNamesFile, opts::SkipFunctionNames);
-  populateFunctionNames(opts::FunctionNamesFileNR, opts::ForceFunctionNamesNR);
+  opts::populateFunctionNames(opts::FunctionNamesFile,
+                              opts::ForceFunctionNames);
+  opts::populateFunctionNames(opts::SkipFunctionNamesFile,
+                              opts::SkipFunctionNames);
+  opts::populateFunctionNames(opts::FunctionNamesFileNR,
+                              opts::ForceFunctionNamesNR);
 
   // Make a set of functions to process to speed up lookups.
   std::unordered_set<std::string> ForceFunctionsNR(
diff --git a/bolt/lib/Utils/CommandLineOpts.cpp b/bolt/lib/Utils/CommandLineOpts.cpp
index de82420a167131..96df4b290b55ed 100644
--- a/bolt/lib/Utils/CommandLineOpts.cpp
+++ b/bolt/lib/Utils/CommandLineOpts.cpp
@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Utils/CommandLineOpts.h"
-#include "VCSVersion.inc"
+//#include "VCSVersion.inc"
+#include "llvm/Support/VCSRevision.h"
+#include <fstream>
 
 using namespace llvm;
 
@@ -206,11 +208,31 @@ cl::opt<bool> UpdateDebugSections(
     cl::desc("update DWARF debug sections of the executable"),
     cl::cat(BoltCategory));
 
+cl::list<std::string> AssumeNoReturnFunctions(
+    "noreturnfuncs", cl::CommaSeparated,
+    cl::desc("List which function names to assume are no-return"),
+    cl::value_desc("func1,func2,func3,..."), cl::Hidden, cl::cat(BoltCategory));
+
+cl::opt<std::string> AssumeNoReturnFunctionsFile(
+    "noreturnfuncs-file",
+    cl::desc("file with list of functions to assume are no-return"), cl::Hidden,
+    cl::cat(BoltCategory));
+
 cl::opt<unsigned>
     Verbosity("v", cl::desc("set verbosity level for diagnostic output"),
               cl::init(0), cl::ZeroOrMore, cl::cat(BoltCategory),
               cl::sub(cl::SubCommand::getAll()));
 
+void populateFunctionNames(const cl::opt<std::string> &FunctionNamesFile,
+                           cl::list<std::string> &FunctionNames) {
+  if (FunctionNamesFile.empty())
+    return;
+  std::ifstream FuncsFile(FunctionNamesFile, std::ios::in);
+  std::string FuncName;
+  while (std::getline(FuncsFile, FuncName))
+    FunctionNames.push_back(FuncName);
+}
+
 bool processAllFunctions() {
   if (opts::AggregateOnly)
     return false;
diff --git a/bolt/test/gadget-scanner/gs-pacret-noreturn.s b/bolt/test/gadget-scanner/gs-pacret-noreturn.s
new file mode 100644
index 00000000000000..d3e1c2272541da
--- /dev/null
+++ b/bolt/test/gadget-scanner/gs-pacret-noreturn.s
@@ -0,0 +1,31 @@
+// Check that there are no false positives related to no-return functions.
+
+// RUN: %clang %cflags -march=armv8.3-a -mbranch-protection=pac-ret %s %p/../Inputs/asm_main.c -o %t.exe
+// RUN: llvm-bolt-gadget-scanner %t.exe --noreturnfuncs="doesnotreturn/1" 2>&1 | FileCheck -check-prefix=CHECK --allow-empty %s
+
+
+// Verify that we can also detect gadgets across basic blocks
+
+        .globl f_call_returning
+        .type   f_call_returning, at function
+f_call_returning:
+        bl      call_returning
+        ret
+        .size f_call_returning, .-f_call_returning
+// CHECK-LABEL:     GS-PACRET: non-protected ret found in function f_call_returning, basic block .L{{[^,]+}}, at address
+// CHECK-NEXT:  The return instruction is     {{[0-9a-f]+}}:       ret
+// CHECK-NEXT:  The 1 instructions that write to the return register after any authentication are:
+// CHECK-NEXT:  1.     {{[0-9a-f]+}}:      bl call_returning
+
+        .type doesnotreturn, at function
+doesnotreturn:
+        brk 1
+        .size doesnotreturn, .-doesnotreturn
+
+        .globl f_call_noreturn
+        .type   f_call_noreturn, at function
+f_call_noreturn:
+        bl      doesnotreturn
+        ret
+        .size f_call_noreturn, .-f_call_noreturn
+// CHECK-NOT: function f_call_noreturn

>From 1ddcc0439c7186da3de670f5ea6f434fa3c23fe6 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Thu, 7 Nov 2024 14:30:42 +0100
Subject: [PATCH 5/7] Fix to build with "handle noret"

---
 llvm/include/llvm/ADT/StringRef.h | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/llvm/include/llvm/ADT/StringRef.h b/llvm/include/llvm/ADT/StringRef.h
index 5b525c8e56ecc9..381a254ce86348 100644
--- a/llvm/include/llvm/ADT/StringRef.h
+++ b/llvm/include/llvm/ADT/StringRef.h
@@ -177,6 +177,11 @@ namespace llvm {
       return size() == RHS.size() && compare_insensitive(RHS) == 0;
     }
 
+    /// Check for string equality, ignoring case.
+    [[nodiscard]] bool equals(StringRef RHS) const {
+      return Length == RHS.Length && compare(RHS) == 0;
+    }
+
     /// compare - Compare two strings; the result is negative, zero, or positive
     /// if this string is lexicographically less than, equal to, or greater than
     /// the \p RHS.

>From 8e9562b1a68e899f990dffbcd304c5cbb6d2e508 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Tue, 26 Nov 2024 08:36:02 +0100
Subject: [PATCH 6/7] Remove bolt-gadget-scanner specific test

---
 bolt/test/gadget-scanner/gs-pacret-noreturn.s | 31 -------------------
 1 file changed, 31 deletions(-)
 delete mode 100644 bolt/test/gadget-scanner/gs-pacret-noreturn.s

diff --git a/bolt/test/gadget-scanner/gs-pacret-noreturn.s b/bolt/test/gadget-scanner/gs-pacret-noreturn.s
deleted file mode 100644
index d3e1c2272541da..00000000000000
--- a/bolt/test/gadget-scanner/gs-pacret-noreturn.s
+++ /dev/null
@@ -1,31 +0,0 @@
-// Check that there are no false positives related to no-return functions.
-
-// RUN: %clang %cflags -march=armv8.3-a -mbranch-protection=pac-ret %s %p/../Inputs/asm_main.c -o %t.exe
-// RUN: llvm-bolt-gadget-scanner %t.exe --noreturnfuncs="doesnotreturn/1" 2>&1 | FileCheck -check-prefix=CHECK --allow-empty %s
-
-
-// Verify that we can also detect gadgets across basic blocks
-
-        .globl f_call_returning
-        .type   f_call_returning, at function
-f_call_returning:
-        bl      call_returning
-        ret
-        .size f_call_returning, .-f_call_returning
-// CHECK-LABEL:     GS-PACRET: non-protected ret found in function f_call_returning, basic block .L{{[^,]+}}, at address
-// CHECK-NEXT:  The return instruction is     {{[0-9a-f]+}}:       ret
-// CHECK-NEXT:  The 1 instructions that write to the return register after any authentication are:
-// CHECK-NEXT:  1.     {{[0-9a-f]+}}:      bl call_returning
-
-        .type doesnotreturn, at function
-doesnotreturn:
-        brk 1
-        .size doesnotreturn, .-doesnotreturn
-
-        .globl f_call_noreturn
-        .type   f_call_noreturn, at function
-f_call_noreturn:
-        bl      doesnotreturn
-        ret
-        .size f_call_noreturn, .-f_call_noreturn
-// CHECK-NOT: function f_call_noreturn

>From 93ab27d50158baca43deb2dddd6397428b268002 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Tue, 26 Nov 2024 08:39:56 +0100
Subject: [PATCH 7/7] [BOLT] InsertNegateRAStatePass: change CFG error
 conditions

    - allow running on finalized CFG
    - don't throw error if one functions doesn't have CFG, only skip that
      one
---
 bolt/lib/Passes/InsertNegateRAStatePass.cpp | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/bolt/lib/Passes/InsertNegateRAStatePass.cpp b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
index c5db6df3a2b606..e7c3bdb0508134 100644
--- a/bolt/lib/Passes/InsertNegateRAStatePass.cpp
+++ b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
@@ -20,13 +20,7 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
 
   if (BF.getState() != BinaryFunction::State::CFG &&
       BF.getState() != BinaryFunction::State::CFG_Finalized) {
-    BC.errs() << "BOLT-WARNING: No CFG for " << BF.getPrintName()
-              << " in InsertNegateRAStatePass\n";
-    return;
-  }
-
-  if (BF.getState() == BinaryFunction::State::CFG_Finalized) {
-    BC.errs() << "BOLT-WARNING: CFG finalized for " << BF.getPrintName()
+    BC.outs() << "BOLT-INFO: No CFG for " << BF.getPrintName()
               << " in InsertNegateRAStatePass\n";
     return;
   }



More information about the llvm-commits mailing list