[BOLT] Gadget scanner: clarify MCPlusBuilder callbacks interface (#136147)

Clarify the semantics of `getAuthenticatedReg` and remove a redundant
`isAuthenticationOfReg` method, as combined auth+something instructions
(such as `retaa` on AArch64) should be handled carefully, especially
when searching for authentication oracles: usually, such instructions
cannot be authentication oracles and only some of them actually write an
authenticated pointer to a register (such as "ldra x0, [x1]!").

Use `std::optional<MCPhysReg>` returned type instead of plain MCPhysReg
and returning `getNoRegister()` as a "not applicable" indication.

Document a few existing methods, add information about preconditions.

Author: atrosinenko

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -562,35 +562,56 @@ class MCPlusBuilder {
     return {};
   }
 
-  virtual ErrorOr<MCPhysReg> getAuthenticatedReg(const MCInst &Inst) const {
-    llvm_unreachable("not implemented");
-    return getNoRegister();
-  }
-
-  virtual bool isAuthenticationOfReg(const MCInst &Inst,
-                                     MCPhysReg AuthenticatedReg) const {
+  /// Returns the register where an authenticated pointer is written to by Inst,
+  /// or std::nullopt if not authenticating any register.
+  ///
+  /// Sets IsChecked if the instruction always checks authenticated pointer,
+  /// i.e. it either writes a successfully authenticated pointer or terminates
+  /// the program abnormally (such as "ldra x0, [x1]!" on AArch64, which crashes
+  /// on authentication failure even if FEAT_FPAC is not implemented).
+  virtual std::optional<MCPhysReg>
+  getWrittenAuthenticatedReg(const MCInst &Inst, bool &IsChecked) const {
     llvm_unreachable("not implemented");
-    return false;
+    return std::nullopt;
   }
 
-  virtual MCPhysReg getSignedReg(const MCInst &Inst) const {
+  /// Returns the register signed by Inst, or std::nullopt if not signing any
+  /// register.
+  ///
+  /// The returned register is assumed to be both input and output operand,
+  /// as it is done on AArch64.
+  virtual std::optional<MCPhysReg> getSignedReg(const MCInst &Inst) const {
     llvm_unreachable("not implemented");
-    return getNoRegister();
+    return std::nullopt;
   }
 
-  virtual ErrorOr<MCPhysReg> getRegUsedAsRetDest(const MCInst &Inst) const {
+  /// 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.
+  ///
+  /// Sets IsAuthenticatedInternally if the instruction accepts a signed
+  /// pointer as its operand and authenticates it internally.
+  ///
+  /// Should only be called when isReturn(Inst) is true.
+  virtual std::optional<MCPhysReg>
+  getRegUsedAsRetDest(const MCInst &Inst,
+                      bool &IsAuthenticatedInternally) const {
     llvm_unreachable("not implemented");
-    return getNoRegister();
+    return std::nullopt;
   }
 
   /// Returns the register used as the destination of an indirect branch or call
   /// instruction. Sets IsAuthenticatedInternally if the instruction accepts
   /// a signed pointer as its operand and authenticates it internally.
+  ///
+  /// Should only be called if isIndirectCall(Inst) or isIndirectBranch(Inst)
+  /// returns true.
   virtual MCPhysReg
   getRegUsedAsIndirectBranchDest(const MCInst &Inst,
                                  bool &IsAuthenticatedInternally) const {
     llvm_unreachable("not implemented");
-    return getNoRegister();
+    return 0; // Unreachable. A valid register should be returned by the
+              // target implementation.
   }
 
   /// Returns the register containing an address safely materialized by `Inst`
@@ -602,14 +623,14 @@ class MCPlusBuilder {
   ///    controlled, under the Pointer Authentication threat model.
   ///
   /// If the instruction does not write to any register satisfying the above
-  /// two conditions, NoRegister is returned.
+  /// two conditions, std::nullopt is returned.
   ///
   /// The Pointer Authentication threat model assumes an attacker is able to
   /// modify any writable memory, but not executable code (due to W^X).
-  virtual MCPhysReg
+  virtual std::optional<MCPhysReg>
   getMaterializedAddressRegForPtrAuth(const MCInst &Inst) const {
     llvm_unreachable("not implemented");
-    return getNoRegister();
+    return std::nullopt;
   }
 
   /// Analyzes if this instruction can safely perform address arithmetics
@@ -622,10 +643,13 @@ class MCPlusBuilder {
   /// controlled, provided InReg and executable code are not. Please note that
   /// registers other than InReg as well as the contents of memory which is
   /// writable by the process should be considered attacker-controlled.
+  ///
+  /// The instruction should not write any values derived from InReg anywhere,
+  /// except for OutReg.
   virtual std::optional<std::pair<MCPhysReg, MCPhysReg>>
   analyzeAddressArithmeticsForPtrAuth(const MCInst &Inst) const {
     llvm_unreachable("not implemented");
-    return std::make_pair(getNoRegister(), getNoRegister());
+    return std::nullopt;
   }
 
   /// Analyzes if a pointer is checked to be authenticated successfully
@@ -670,10 +694,10 @@ class MCPlusBuilder {
   ///
   /// Use this function for simple, single-instruction patterns instead of
   /// its getAuthCheckedReg(BB) counterpart.
-  virtual MCPhysReg getAuthCheckedReg(const MCInst &Inst,
-                                      bool MayOverwrite) const {
+  virtual std::optional<MCPhysReg> getAuthCheckedReg(const MCInst &Inst,
+                                                     bool MayOverwrite) const {
     llvm_unreachable("not implemented");
-    return getNoRegister();
+    return std::nullopt;
   }
 
   virtual bool isTerminator(const MCInst &Inst) const;

bolt/lib/Passes/PAuthGadgetScanner.cpp

@@ -66,14 +66,12 @@ namespace PAuthGadgetScanner {
 }
 
 [[maybe_unused]] static void traceReg(const BinaryContext &BC, StringRef Label,
-                                      ErrorOr<MCPhysReg> Reg) {
+                                      MCPhysReg Reg) {
   dbgs() << "    " << Label << ": ";
-  if (Reg.getError())
-    dbgs() << "(error)";
-  else if (*Reg == BC.MIB->getNoRegister())
+  if (Reg == BC.MIB->getNoRegister())
     dbgs() << "(none)";
   else
-    dbgs() << BC.MRI->getName(*Reg);
+    dbgs() << BC.MRI->getName(Reg);
   dbgs() << "\n";
 }
 
@@ -365,17 +363,15 @@ class SrcSafetyAnalysis {
   SmallVector<MCPhysReg> getRegsMadeSafeToDeref(const MCInst &Point,
                                                 const SrcState &Cur) const {
     SmallVector<MCPhysReg> Regs;
-    const MCPhysReg NoReg = BC.MIB->getNoRegister();
 
     // A signed pointer can be authenticated, or
-    ErrorOr<MCPhysReg> AutReg = BC.MIB->getAuthenticatedReg(Point);
-    if (AutReg && *AutReg != NoReg)
+    bool Dummy = false;
+    if (auto AutReg = BC.MIB->getWrittenAuthenticatedReg(Point, Dummy))
       Regs.push_back(*AutReg);
 
     // ... a safe address can be materialized, or
-    MCPhysReg NewAddrReg = BC.MIB->getMaterializedAddressRegForPtrAuth(Point);
-    if (NewAddrReg != NoReg)
-      Regs.push_back(NewAddrReg);
+    if (auto NewAddrReg = BC.MIB->getMaterializedAddressRegForPtrAuth(Point))
+      Regs.push_back(*NewAddrReg);
 
     // ... an address can be updated in a safe manner, producing the result
     // which is as trusted as the input address.
@@ -391,13 +387,20 @@ class SrcSafetyAnalysis {
   SmallVector<MCPhysReg> getRegsMadeTrusted(const MCInst &Point,
                                             const SrcState &Cur) const {
     SmallVector<MCPhysReg> Regs;
-    const MCPhysReg NoReg = BC.MIB->getNoRegister();
 
     // An authenticated pointer can be checked, or
-    MCPhysReg CheckedReg =
+    std::optional<MCPhysReg> CheckedReg =
         BC.MIB->getAuthCheckedReg(Point, /*MayOverwrite=*/false);
-    if (CheckedReg != NoReg && Cur.SafeToDerefRegs[CheckedReg])
-      Regs.push_back(CheckedReg);
+    if (CheckedReg && Cur.SafeToDerefRegs[*CheckedReg])
+      Regs.push_back(*CheckedReg);
+
+    // ... a pointer can be authenticated by an instruction that always checks
+    // the pointer, or
+    bool IsChecked = false;
+    std::optional<MCPhysReg> AutReg =
+        BC.MIB->getWrittenAuthenticatedReg(Point, IsChecked);
+    if (AutReg && IsChecked)
+      Regs.push_back(*AutReg);
 
     if (CheckerSequenceInfo.contains(&Point)) {
       MCPhysReg CheckedReg;
@@ -413,9 +416,8 @@ class SrcSafetyAnalysis {
     }
 
     // ... a safe address can be materialized, or
-    MCPhysReg NewAddrReg = BC.MIB->getMaterializedAddressRegForPtrAuth(Point);
-    if (NewAddrReg != NoReg)
-      Regs.push_back(NewAddrReg);
+    if (auto NewAddrReg = BC.MIB->getMaterializedAddressRegForPtrAuth(Point))
+      Regs.push_back(*NewAddrReg);
 
     // ... an address can be updated in a safe manner, producing the result
     // which is as trusted as the input address.
@@ -738,25 +740,27 @@ shouldReportReturnGadget(const BinaryContext &BC, const MCInstReference &Inst,
   if (!BC.MIB->isReturn(Inst))
     return std::nullopt;
 
-  ErrorOr<MCPhysReg> MaybeRetReg = BC.MIB->getRegUsedAsRetDest(Inst);
-  if (MaybeRetReg.getError()) {
+  bool IsAuthenticated = false;
+  std::optional<MCPhysReg> RetReg =
+      BC.MIB->getRegUsedAsRetDest(Inst, IsAuthenticated);
+  if (!RetReg) {
     return make_generic_report(
         Inst, "Warning: pac-ret analysis could not analyze this return "
               "instruction");
   }
-  MCPhysReg RetReg = *MaybeRetReg;
+  if (IsAuthenticated)
+    return std::nullopt;
+
   LLVM_DEBUG({
     traceInst(BC, "Found RET inst", Inst);
-    traceReg(BC, "RetReg", RetReg);
-    traceReg(BC, "Authenticated reg", BC.MIB->getAuthenticatedReg(Inst));
+    traceReg(BC, "RetReg", *RetReg);
+    traceRegMask(BC, "SafeToDerefRegs", S.SafeToDerefRegs);
   });
-  if (BC.MIB->isAuthenticationOfReg(Inst, RetReg))
-    return std::nullopt;
-  LLVM_DEBUG({ traceRegMask(BC, "SafeToDerefRegs", S.SafeToDerefRegs); });
-  if (S.SafeToDerefRegs[RetReg])
+
+  if (S.SafeToDerefRegs[*RetReg])
     return std::nullopt;
 
-  return make_gadget_report(RetKind, Inst, RetReg);
+  return make_gadget_report(RetKind, Inst, *RetReg);
 }
 
 static std::optional<PartialReport<MCPhysReg>>
@@ -772,7 +776,7 @@ shouldReportCallGadget(const BinaryContext &BC, const MCInstReference &Inst,
   if (IsAuthenticated)
     return std::nullopt;
 
-  assert(DestReg != BC.MIB->getNoRegister());
+  assert(DestReg != BC.MIB->getNoRegister() && "Valid register expected");
   LLVM_DEBUG({
     traceInst(BC, "Found call inst", Inst);
     traceReg(BC, "Call destination reg", DestReg);
@@ -789,19 +793,19 @@ shouldReportSigningOracle(const BinaryContext &BC, const MCInstReference &Inst,
                           const SrcState &S) {
   static const GadgetKind SigningOracleKind("signing oracle found");
 
-  MCPhysReg SignedReg = BC.MIB->getSignedReg(Inst);
-  if (SignedReg == BC.MIB->getNoRegister())
+  std::optional<MCPhysReg> SignedReg = BC.MIB->getSignedReg(Inst);
+  if (!SignedReg)
     return std::nullopt;
 
   LLVM_DEBUG({
     traceInst(BC, "Found sign inst", Inst);
-    traceReg(BC, "Signed reg", SignedReg);
+    traceReg(BC, "Signed reg", *SignedReg);
     traceRegMask(BC, "TrustedRegs", S.TrustedRegs);
   });
-  if (S.TrustedRegs[SignedReg])
+  if (S.TrustedRegs[*SignedReg])
     return std::nullopt;
 
-  return make_gadget_report(SigningOracleKind, Inst, SignedReg);
+  return make_gadget_report(SigningOracleKind, Inst, *SignedReg);
 }
 
 template <typename T> static void iterateOverInstrs(BinaryFunction &BF, T Fn) {