[BOLT] Gadget scanner: refactor analysis of RET instructions (#131897)

In preparation for implementing detection of more gadget kinds,
refactor checking for non-protected return instructions.

3 files changed, 124 insertions, 77 deletions

diff --git a/bolt/include/bolt/Passes/PAuthGadgetScanner.h b/bolt/include/bolt/Passes/PAuthGadgetScanner.h
index 2d8109f..1a8abff 100644
--- a/bolt/include/bolt/Passes/PAuthGadgetScanner.h
+++ b/bolt/include/bolt/Passes/PAuthGadgetScanner.h
@@ -199,19 +199,40 @@ struct Report {
   virtual void generateReport(raw_ostream &OS,
                               const BinaryContext &BC) const = 0;
 
+  // The two methods below are called by Analysis::computeDetailedInfo when
+  // iterating over the reports.
+  virtual const ArrayRef<MCPhysReg> getAffectedRegisters() const { return {}; }
+  virtual void setOverwritingInstrs(const ArrayRef<MCInstReference> Instrs) {}
+
   void printBasicInfo(raw_ostream &OS, const BinaryContext &BC,
                       StringRef IssueKind) const;
 };
 
 struct GadgetReport : public Report {
+  // The particular kind of gadget that is detected.
   const GadgetKind &Kind;
-  std::vector<MCInstReference> OverwritingInstrs;
+  // The set of registers related to this gadget report (possibly empty).
+  SmallVector<MCPhysReg> AffectedRegisters;
+  // The instructions that clobber the affected registers.
+  // There is no one-to-one correspondence with AffectedRegisters: for example,
+  // the same register can be overwritten by different instructions in different
+  // preceding basic blocks.
+  SmallVector<MCInstReference> OverwritingInstrs;
 
   GadgetReport(const GadgetKind &Kind, MCInstReference Location,
-               std::vector<MCInstReference> OverwritingInstrs)
-      : Report(Location), Kind(Kind), OverwritingInstrs(OverwritingInstrs) {}
+               const BitVector &AffectedRegisters)
+      : Report(Location), Kind(Kind),
+        AffectedRegisters(AffectedRegisters.set_bits()) {}
 
   void generateReport(raw_ostream &OS, const BinaryContext &BC) const override;
+
+  const ArrayRef<MCPhysReg> getAffectedRegisters() const override {
+    return AffectedRegisters;
+  }
+
+  void setOverwritingInstrs(const ArrayRef<MCInstReference> Instrs) override {
+    OverwritingInstrs.assign(Instrs.begin(), Instrs.end());
+  }
 };
 
 /// Report with a free-form message attached.
@@ -224,16 +245,18 @@ struct GenericReport : public Report {
 };
 
 struct FunctionAnalysisResult {
-  SmallSet<MCPhysReg, 1> RegistersAffected;
   std::vector<std::shared_ptr<Report>> Diagnostics;
 };
 
 class Analysis : public BinaryFunctionPass {
   void runOnFunction(BinaryFunction &Function,
                      MCPlusBuilder::AllocatorIdTy AllocatorId);
-  FunctionAnalysisResult
-  computeDfState(PacRetAnalysis &PRA, BinaryFunction &BF,
-                 MCPlusBuilder::AllocatorIdTy AllocatorId);
+  FunctionAnalysisResult findGadgets(BinaryFunction &BF,
+                                     MCPlusBuilder::AllocatorIdTy AllocatorId);
+
+  void computeDetailedInfo(BinaryFunction &BF,
+                           MCPlusBuilder::AllocatorIdTy AllocatorId,
+                           FunctionAnalysisResult &Result);
 
   std::map<const BinaryFunction *, FunctionAnalysisResult> AnalysisResults;
   std::mutex AnalysisResultsMutex;
diff --git a/bolt/lib/Passes/PAuthGadgetScanner.cpp b/bolt/lib/Passes/PAuthGadgetScanner.cpp
index 6f29399..84c209a 100644
--- a/bolt/lib/Passes/PAuthGadgetScanner.cpp
+++ b/bolt/lib/Passes/PAuthGadgetScanner.cpp
@@ -353,7 +353,7 @@ protected:
 public:
   std::vector<MCInstReference>
   getLastClobberingInsts(const MCInst Ret, BinaryFunction &BF,
-                         const BitVector &UsedDirtyRegs) const {
+                         const ArrayRef<MCPhysReg> UsedDirtyRegs) const {
     if (RegsToTrackInstsFor.empty())
       return {};
     auto MaybeState = getStateAt(Ret);
@@ -362,7 +362,7 @@ public:
     const State &S = *MaybeState;
     // Due to aliasing registers, multiple registers may have been tracked.
     std::set<const MCInst *> LastWritingInsts;
-    for (MCPhysReg TrackedReg : UsedDirtyRegs.set_bits()) {
+    for (MCPhysReg TrackedReg : UsedDirtyRegs) {
       for (const MCInst *Inst : lastWritingInsts(S, TrackedReg))
         LastWritingInsts.insert(Inst);
     }
@@ -376,60 +376,90 @@ public:
   }
 };
 
+static std::shared_ptr<Report>
+shouldReportReturnGadget(const BinaryContext &BC, const MCInstReference &Inst,
+                         const State &S) {
+  static const GadgetKind RetKind("non-protected ret found");
+  if (!BC.MIB->isReturn(Inst))
+    return nullptr;
+
+  ErrorOr<MCPhysReg> MaybeRetReg = BC.MIB->getRegUsedAsRetDest(Inst);
+  if (MaybeRetReg.getError()) {
+    return std::make_shared<GenericReport>(
+        Inst, "Warning: pac-ret analysis could not analyze this return "
+              "instruction");
+  }
+  MCPhysReg RetReg = *MaybeRetReg;
+  LLVM_DEBUG({
+    traceInst(BC, "Found RET inst", Inst);
+    traceReg(BC, "RetReg", RetReg);
+    traceReg(BC, "Authenticated reg", BC.MIB->getAuthenticatedReg(Inst));
+  });
+  if (BC.MIB->isAuthenticationOfReg(Inst, RetReg))
+    return nullptr;
+  BitVector UsedDirtyRegs = S.NonAutClobRegs;
+  LLVM_DEBUG({ traceRegMask(BC, "NonAutClobRegs at Ret", UsedDirtyRegs); });
+  UsedDirtyRegs &= BC.MIB->getAliases(RetReg, /*OnlySmaller=*/true);
+  LLVM_DEBUG({ traceRegMask(BC, "Intersection with RetReg", UsedDirtyRegs); });
+  if (!UsedDirtyRegs.any())
+    return nullptr;
+
+  return std::make_shared<GadgetReport>(RetKind, Inst, UsedDirtyRegs);
+}
+
 FunctionAnalysisResult
-Analysis::computeDfState(PacRetAnalysis &PRA, BinaryFunction &BF,
-                         MCPlusBuilder::AllocatorIdTy AllocatorId) {
+Analysis::findGadgets(BinaryFunction &BF,
+                      MCPlusBuilder::AllocatorIdTy AllocatorId) {
+  FunctionAnalysisResult Result;
+
+  PacRetAnalysis PRA(BF, AllocatorId, {});
   PRA.run();
   LLVM_DEBUG({
     dbgs() << " After PacRetAnalysis:\n";
     BF.dump();
   });
 
-  FunctionAnalysisResult Result;
-  // Now scan the CFG for non-authenticating return instructions that use an
-  // overwritten, non-authenticated register as return address.
   BinaryContext &BC = BF.getBinaryContext();
   for (BinaryBasicBlock &BB : BF) {
-    for (int64_t I = BB.size() - 1; I >= 0; --I) {
-      MCInst &Inst = BB.getInstructionAtIndex(I);
-      if (BC.MIB->isReturn(Inst)) {
-        ErrorOr<MCPhysReg> MaybeRetReg = BC.MIB->getRegUsedAsRetDest(Inst);
-        if (MaybeRetReg.getError()) {
-          Result.Diagnostics.push_back(std::make_shared<GenericReport>(
-              MCInstInBBReference(&BB, I),
-              "Warning: pac-ret analysis could not analyze this return "
-              "instruction"));
-          continue;
-        }
-        MCPhysReg RetReg = *MaybeRetReg;
-        LLVM_DEBUG({
-          traceInst(BC, "Found RET inst", Inst);
-          traceReg(BC, "RetReg", RetReg);
-          traceReg(BC, "Authenticated reg", BC.MIB->getAuthenticatedReg(Inst));
-        });
-        if (BC.MIB->isAuthenticationOfReg(Inst, RetReg))
-          break;
-        BitVector UsedDirtyRegs = PRA.getStateAt(Inst)->NonAutClobRegs;
-        LLVM_DEBUG(
-            { traceRegMask(BC, "NonAutClobRegs at Ret", UsedDirtyRegs); });
-        UsedDirtyRegs &= BC.MIB->getAliases(RetReg, /*OnlySmaller=*/true);
-        LLVM_DEBUG(
-            { traceRegMask(BC, "Intersection with RetReg", UsedDirtyRegs); });
-        if (UsedDirtyRegs.any()) {
-          static const GadgetKind RetKind("non-protected ret found");
-          // This return instruction needs to be reported
-          Result.Diagnostics.push_back(std::make_shared<GadgetReport>(
-              RetKind, MCInstInBBReference(&BB, I),
-              PRA.getLastClobberingInsts(Inst, BF, UsedDirtyRegs)));
-          for (MCPhysReg RetRegWithGadget : UsedDirtyRegs.set_bits())
-            Result.RegistersAffected.insert(RetRegWithGadget);
-        }
-      }
+    for (int64_t I = 0, E = BB.size(); I < E; ++I) {
+      MCInstReference Inst(&BB, I);
+      const State &S = *PRA.getStateAt(Inst);
+
+      if (auto Report = shouldReportReturnGadget(BC, Inst, S))
+        Result.Diagnostics.push_back(Report);
     }
   }
   return Result;
 }
 
+void Analysis::computeDetailedInfo(BinaryFunction &BF,
+                                   MCPlusBuilder::AllocatorIdTy AllocatorId,
+                                   FunctionAnalysisResult &Result) {
+  BinaryContext &BC = BF.getBinaryContext();
+
+  // Collect the affected registers across all gadgets found in this function.
+  SmallSet<MCPhysReg, 4> RegsToTrack;
+  for (auto Report : Result.Diagnostics)
+    RegsToTrack.insert_range(Report->getAffectedRegisters());
+  std::vector<MCPhysReg> RegsToTrackVec(RegsToTrack.begin(), RegsToTrack.end());
+
+  // Re-compute the analysis with register tracking.
+  PacRetAnalysis PRWIA(BF, AllocatorId, RegsToTrackVec);
+  PRWIA.run();
+  LLVM_DEBUG({
+    dbgs() << " After detailed PacRetAnalysis:\n";
+    BF.dump();
+  });
+
+  // Augment gadget reports.
+  for (auto Report : Result.Diagnostics) {
+    LLVM_DEBUG(
+        { traceInst(BC, "Attaching clobbering info to", Report->Location); });
+    Report->setOverwritingInstrs(PRWIA.getLastClobberingInsts(
+        Report->Location, BF, Report->getAffectedRegisters()));
+  }
+}
+
 void Analysis::runOnFunction(BinaryFunction &BF,
                              MCPlusBuilder::AllocatorIdTy AllocatorId) {
   LLVM_DEBUG({
@@ -438,27 +468,25 @@ void Analysis::runOnFunction(BinaryFunction &BF,
     BF.dump();
   });
 
-  if (BF.hasCFG()) {
-    PacRetAnalysis PRA(BF, AllocatorId, {});
-    FunctionAnalysisResult FAR = computeDfState(PRA, BF, AllocatorId);
-    if (!FAR.RegistersAffected.empty()) {
-      // Redo the analysis, but now also track which instructions last wrote
-      // to any of the registers in RetRegsWithGadgets, so that better
-      // diagnostics can be produced.
-      std::vector<MCPhysReg> RegsToTrack;
-      for (MCPhysReg R : FAR.RegistersAffected)
-        RegsToTrack.push_back(R);
-      PacRetAnalysis PRWIA(BF, AllocatorId, RegsToTrack);
-      FAR = computeDfState(PRWIA, BF, AllocatorId);
-    }
+  if (!BF.hasCFG())
+    return;
 
-    // `runOnFunction` is typically getting called from multiple threads in
-    // parallel. Therefore, use a lock to avoid data races when storing the
-    // result of the analysis in the `AnalysisResults` map.
-    {
-      std::lock_guard<std::mutex> Lock(AnalysisResultsMutex);
-      AnalysisResults[&BF] = FAR;
-    }
+  FunctionAnalysisResult FAR = findGadgets(BF, AllocatorId);
+  if (FAR.Diagnostics.empty())
+    return;
+
+  // Redo the analysis, but now also track which instructions last wrote
+  // to any of the registers in RetRegsWithGadgets, so that better
+  // diagnostics can be produced.
+
+  computeDetailedInfo(BF, AllocatorId, FAR);
+
+  // `runOnFunction` is typically getting called from multiple threads in
+  // parallel. Therefore, use a lock to avoid data races when storing the
+  // result of the analysis in the `AnalysisResults` map.
+  {
+    std::lock_guard<std::mutex> Lock(AnalysisResultsMutex);
+    AnalysisResults[&BF] = FAR;
   }
 }
 
@@ -517,7 +545,7 @@ void GadgetReport::generateReport(raw_ostream &OS,
      << " instructions that write to the affected registers after any "
         "authentication are:\n";
   // Sort by address to ensure output is deterministic.
-  std::vector<MCInstReference> OI = OverwritingInstrs;
+  SmallVector<MCInstReference> OI = OverwritingInstrs;
   llvm::sort(OI, [](const MCInstReference &A, const MCInstReference &B) {
     return A.getAddress() < B.getAddress();
   });
diff --git a/bolt/test/binary-analysis/AArch64/gs-pauth-debug-output.s b/bolt/test/binary-analysis/AArch64/gs-pauth-debug-output.s
index a1546f6..18e0fbf 100644
--- a/bolt/test/binary-analysis/AArch64/gs-pauth-debug-output.s
+++ b/bolt/test/binary-analysis/AArch64/gs-pauth-debug-output.s
@@ -118,20 +118,16 @@ clobber:
 // CHECK-NEXT:   .. result: (pacret-state<NonAutClobRegs: W30 , Insts: [0](0x{{[0-9a-f]+}} )>)
 // CHECK-NEXT:  PacRetAnalysis::ComputeNext(   ret     x30, pacret-state<NonAutClobRegs: W30 , Insts: [0](0x{{[0-9a-f]+}} )>)
 // CHECK-NEXT:   .. result: (pacret-state<NonAutClobRegs: W30 , Insts: [0](0x{{[0-9a-f]+}} )>)
-// CHECK-NEXT:  After PacRetAnalysis:
+// CHECK-NEXT:  After detailed PacRetAnalysis:
 // CHECK-NEXT: Binary Function "clobber"  {
 // ...
 // CHECK:      End of Function "clobber"
 
 // The analysis was re-computed with register tracking, as an issue was found in this function.
-// Re-checking the instructions:
+// Iterating over the reports and attaching clobbering info:
 
 // CHECK-EMPTY:
-// CHECK-NEXT:   Found RET inst:     00000000:         ret # PacRetAnalysis: pacret-state<NonAutClobRegs: BitVector, Insts: [0](0x{{[0-9a-f]+}} )>
-// CHECK-NEXT:     RetReg: LR
-// CHECK-NEXT:     Authenticated reg: (none)
-// CHECK-NEXT:     NonAutClobRegs at Ret: W30
-// CHECK-NEXT:     Intersection with RetReg: W30
+// CHECK-NEXT:   Attaching clobbering info to:     00000000:         ret # PacRetAnalysis: pacret-state<NonAutClobRegs: BitVector, Insts: [0](0x{{[0-9a-f]+}} )>
 
 
 // CHECK-LABEL:Analyzing in function main, AllocatorId 1