Address review comments

Author: atrosinenko

bolt/include/bolt/Passes/PAuthGadgetScanner.h

@@ -199,18 +199,25 @@ 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 std::vector<MCInstReference> &Instrs) {}
+  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;
+  // The set of registers related to this gadget report (possibly empty).
   SmallVector<MCPhysReg> AffectedRegisters;
-  std::vector<MCInstReference> OverwritingInstrs;
+  // 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,
                const BitVector &AffectedRegisters)
@@ -223,9 +230,8 @@ struct GadgetReport : public Report {
     return AffectedRegisters;
   }
 
-  void
-  setOverwritingInstrs(const std::vector<MCInstReference> &Instrs) override {
-    OverwritingInstrs = Instrs;
+  void setOverwritingInstrs(const ArrayRef<MCInstReference> Instrs) override {
+    OverwritingInstrs.assign(Instrs.begin(), Instrs.end());
   }
 };
 
@@ -245,8 +251,12 @@ struct FunctionAnalysisResult {
 class Analysis : public BinaryFunctionPass {
   void runOnFunction(BinaryFunction &Function,
                      MCPlusBuilder::AllocatorIdTy AllocatorId);
-  FunctionAnalysisResult
-  computeDfState(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;

bolt/lib/Passes/PAuthGadgetScanner.cpp

@@ -376,9 +376,9 @@ class PacRetAnalysis
   }
 };
 
-static std::shared_ptr<Report> tryCheckReturn(const BinaryContext &BC,
-                                              const MCInstReference &Inst,
-                                              const State &S) {
+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;
@@ -408,8 +408,8 @@ static std::shared_ptr<Report> tryCheckReturn(const BinaryContext &BC,
 }
 
 FunctionAnalysisResult
-Analysis::computeDfState(BinaryFunction &BF,
-                         MCPlusBuilder::AllocatorIdTy AllocatorId) {
+Analysis::findGadgets(BinaryFunction &BF,
+                      MCPlusBuilder::AllocatorIdTy AllocatorId) {
   FunctionAnalysisResult Result;
 
   PacRetAnalysis PRA(BF, AllocatorId, {});
@@ -425,40 +425,39 @@ Analysis::computeDfState(BinaryFunction &BF,
       MCInstReference Inst(&BB, I);
       const State &S = *PRA.getStateAt(Inst);
 
-      if (auto Report = tryCheckReturn(BC, Inst, S))
+      if (auto Report = shouldReportReturnGadget(BC, Inst, S))
         Result.Diagnostics.push_back(Report);
     }
   }
+  return Result;
+}
 
-  if (Result.Diagnostics.empty())
-    return Result;
-
-  // 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.
+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)
-    for (MCPhysReg Reg : Report->getAffectedRegisters())
-      RegsToTrack.insert(Reg);
-
+    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()));
   }
-
-  return Result;
 }
 
 void Analysis::runOnFunction(BinaryFunction &BF,
@@ -472,10 +471,16 @@ void Analysis::runOnFunction(BinaryFunction &BF,
   if (!BF.hasCFG())
     return;
 
-  FunctionAnalysisResult FAR = computeDfState(BF, AllocatorId);
+  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.
@@ -540,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();
   });