[BOLT] Gadget scanner: prevent false positives due to jump tables

As part of PAuth hardening, AArch64 LLVM backend can use a special
BR_JumpTable pseudo (enabled by -faarch64-jump-table-hardening
Clang option) which is expanded in the AsmPrinter into a contiguous
sequence without unsafe instructions in the middle.

This commit adds another target-specific callback to MCPlusBuilder
to make it possible to inhibit false positives for known-safe jump
table dispatch sequences. Without special handling, the branch
instruction is likely to be reported as a non-protected call (as its
destination is not produced by an auth instruction, PC-relative address
materialization, etc.) and possibly as a tail call being performed with
unsafe link register (as the detection whether the branch instruction
is a tail call is an heuristic).

For now, only the specific instruction sequence used by the AArch64
LLVM backend is matched.

Author: atrosinenko

bolt/include/bolt/Core/MCInstUtils.h

@@ -101,6 +101,15 @@ class MCInstReference {
   /// this function may be called from multithreaded code.
   uint64_t computeAddress(const MCCodeEmitter *Emitter = nullptr) const;
 
+  /// Returns the only preceding instruction, or std::nullopt if multiple or no
+  /// predecessors are possible.
+  ///
+  /// If CFG information is available, basic block boundary can be crossed,
+  /// provided there is exactly one predecessor. If CFG is not available, the
+  /// preceding instruction in the offset order is returned, unless this is the
+  /// first instruction of the function.
+  std::optional<MCInstReference> getSinglePredecessor();
+
   raw_ostream &print(raw_ostream &OS) const;
 
 private:

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -15,6 +15,7 @@
 #define BOLT_CORE_MCPLUSBUILDER_H
 
 #include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/MCInstUtils.h"
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -718,6 +719,19 @@ class MCPlusBuilder {
     return std::nullopt;
   }
 
+  /// Tests if BranchInst corresponds to an instruction sequence which is known
+  /// to be a safe dispatch via jump table.
+  ///
+  /// The target can decide which instruction sequences to consider "safe" from
+  /// the Pointer Authentication point of view, such as any jump table dispatch
+  /// sequence without function calls inside, any sequence which is contiguous,
+  /// or only some specific well-known sequences.
+  virtual bool
+  isSafeJumpTableBranchForPtrAuth(MCInstReference BranchInst) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
   virtual bool isTerminator(const MCInst &Inst) const;
 
   virtual bool isNoop(const MCInst &Inst) const {

bolt/lib/Core/MCInstUtils.cpp

@@ -84,3 +84,23 @@ raw_ostream &MCInstReference::print(raw_ostream &OS) const {
   OS << ">";
   return OS;
 }
+
+std::optional<MCInstReference> MCInstReference::getSinglePredecessor() {
+  if (const RefInBB *Ref = tryGetRefInBB()) {
+    if (Ref->Index != 0)
+      return MCInstReference(*Ref->BB, Ref->Index - 1);
+
+    if (Ref->BB->pred_size() != 1)
+      return std::nullopt;
+
+    BinaryBasicBlock &PredBB = **Ref->BB->pred_begin();
+    assert(!PredBB.empty() && "Empty basic blocks are not supported yet");
+    return MCInstReference(PredBB, *PredBB.rbegin());
+  }
+
+  const RefInBF &Ref = getRefInBF();
+  if (Ref.It == Ref.BF->instrs().begin())
+    return std::nullopt;
+
+  return MCInstReference(*Ref.BF, std::prev(Ref.It));
+}

bolt/lib/Passes/PAuthGadgetScanner.cpp

@@ -1364,6 +1364,11 @@ shouldReportUnsafeTailCall(const BinaryContext &BC, const BinaryFunction &BF,
     return std::nullopt;
   }
 
+  if (BC.MIB->isSafeJumpTableBranchForPtrAuth(Inst)) {
+    LLVM_DEBUG({ dbgs() << "  Safe jump table detected, skipping.\n"; });
+    return std::nullopt;
+  }
+
   // Returns at most one report per instruction - this is probably OK...
   for (auto Reg : RegsToCheck)
     if (!S.TrustedRegs[Reg])
@@ -1394,6 +1399,11 @@ shouldReportCallGadget(const BinaryContext &BC, const MCInstReference &Inst,
   if (S.SafeToDerefRegs[DestReg])
     return std::nullopt;
 
+  if (BC.MIB->isSafeJumpTableBranchForPtrAuth(Inst)) {
+    LLVM_DEBUG({ dbgs() << "  Safe jump table detected, skipping.\n"; });
+    return std::nullopt;
+  }
+
   return make_gadget_report(CallKind, Inst, DestReg);
 }
 

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -545,6 +545,79 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return std::nullopt;
   }
 
+  bool
+  isSafeJumpTableBranchForPtrAuth(MCInstReference BranchInst) const override {
+    MCInstReference CurRef = BranchInst;
+    auto StepBack = [&]() {
+      do {
+        auto PredInst = CurRef.getSinglePredecessor();
+        if (!PredInst)
+          return false;
+        CurRef = *PredInst;
+      } while (isCFI(CurRef));
+
+      return true;
+    };
+
+    // Match this contiguous sequence:
+    //    cmp   Xm, #count
+    //    csel  Xm, Xm, xzr, ls
+    //    adrp  Xn, .LJTIxyz
+    //    add   Xn, Xn, :lo12:.LJTIxyz
+    //    ldrsw Xm, [Xn, Xm, lsl #2]
+    //  .Ltmp:
+    //    adr   Xn, .Ltmp
+    //    add   Xm, Xn, Xm
+    //    br    Xm
+
+    // FIXME: Check label operands of ADR/ADRP+ADD and #count operand of CMP.
+
+    using namespace LowLevelInstMatcherDSL;
+    Reg Xm, Xn;
+
+    if (!matchInst(CurRef, AArch64::BR, Xm) || !StepBack())
+      return false;
+
+    if (!matchInst(CurRef, AArch64::ADDXrs, Xm, Xn, Xm, Imm(0)) || !StepBack())
+      return false;
+
+    if (!matchInst(CurRef, AArch64::ADR, Xn /*, .Ltmp*/) || !StepBack())
+      return false;
+
+    if (!matchInst(CurRef, AArch64::LDRSWroX, Xm, Xn, Xm, Imm(0), Imm(1)) ||
+        !StepBack())
+      return false;
+
+    if (matchInst(CurRef, AArch64::ADR, Xn /*, .LJTIxyz*/)) {
+      if (!StepBack())
+        return false;
+      if (!matchInst(CurRef, AArch64::HINT, Imm(0)) || !StepBack())
+        return false;
+    } else if (matchInst(CurRef, AArch64::ADDXri, Xn,
+                         Xn /*, :lo12:.LJTIxyz*/)) {
+      if (!StepBack())
+        return false;
+      if (!matchInst(CurRef, AArch64::ADRP, Xn /*, .LJTIxyz*/) || !StepBack())
+        return false;
+    } else {
+      return false;
+    }
+
+    if (!matchInst(CurRef, AArch64::CSELXr, Xm, Xm, Reg(AArch64::XZR),
+                   Imm(AArch64CC::LS)) ||
+        !StepBack())
+      return false;
+
+    if (!matchInst(CurRef, AArch64::SUBSXri, Reg(AArch64::XZR),
+                   Xm /*, #count*/))
+      return false;
+
+    // Some platforms treat X16 and X17 as more protected registers, others
+    // do not make such distinction. So far, accept any registers as Xm and Xn.
+
+    return true;
+  }
+
   bool isADRP(const MCInst &Inst) const override {
     return Inst.getOpcode() == AArch64::ADRP;
   }