bolt/lib/Passes/NonPacProtectedRetAnalysis.cpp - llvm-project - Git at Google

 //===- bolt/Passes/NonPacProtectedRetAnalysis.cpp -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass that looks for any AArch64 return instructions
 // that may not be protected by PAuth authentication instructions when needed.
 //
 //===----------------------------------------------------------------------===//

 #include "bolt/Passes/NonPacProtectedRetAnalysis.h"
 #include "bolt/Core/ParallelUtilities.h"
 #include "bolt/Passes/DataflowAnalysis.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/Support/Format.h"
 #include <memory>

 #define DEBUG_TYPE "bolt-nonpacprotectedret"

 namespace llvm {
 namespace bolt {

 raw_ostream &operator<<(raw_ostream &OS, const MCInstInBBReference &Ref) {
   OS << "MCInstBBRef<";
   if (Ref.BB == nullptr)
     OS << "BB:(null)";
   else
     OS << "BB:" << Ref.BB->getName() << ":" << Ref.BBIndex;
   OS << ">";
   return OS;
 }

 raw_ostream &operator<<(raw_ostream &OS, const MCInstInBFReference &Ref) {
   OS << "MCInstBFRef<";
   if (Ref.BF == nullptr)
     OS << "BF:(null)";
   else
     OS << "BF:" << Ref.BF->getPrintName() << ":" << Ref.getOffset();
   OS << ">";
   return OS;
 }

 raw_ostream &operator<<(raw_ostream &OS, const MCInstReference &Ref) {
   switch (Ref.ParentKind) {
   case MCInstReference::BasicBlockParent:
     OS << Ref.U.BBRef;
     return OS;
   case MCInstReference::FunctionParent:
     OS << Ref.U.BFRef;
     return OS;
   }
   llvm_unreachable("");
 }

 namespace NonPacProtectedRetAnalysis {

 // The security property that is checked is:
 // When a register is used as the address to jump to in a return instruction,
 // that register must either:
 // (a) never be changed within this function, i.e. have the same value as when
 //     the function started, or
 // (b) the last write to the register must be by an authentication instruction.

 // This property is checked by using dataflow analysis to keep track of which
 // registers have been written (def-ed), since last authenticated. Those are
 // exactly the registers containing values that should not be trusted (as they
 // could have changed since the last time they were authenticated). For pac-ret,
 // any return instruction using such a register is a gadget to be reported. For
 // PAuthABI, probably at least any indirect control flow using such a register
 // should be reported.

 // Furthermore, when producing a diagnostic for a found non-pac-ret protected
 // return, the analysis also lists the last instructions that wrote to the
 // register used in the return instruction.
 // The total set of registers used in return instructions in a given function is
 // small. It almost always is just `X30`.
 // In order to reduce the memory consumption of storing this additional state
 // during the dataflow analysis, this is computed by running the dataflow
 // analysis twice:
 // 1. In the first run, the dataflow analysis only keeps track of the security
 //    property: i.e. which registers have been overwritten since the last
 //    time they've been authenticated.
 // 2. If the first run finds any return instructions using a register last
 //    written by a non-authenticating instruction, the dataflow analysis will
 //    be run a second time. The first run will return which registers are used
 //    in the gadgets to be reported. This information is used in the second run
 //    to also track which instructions last wrote to those registers.

 struct State {
   /// A BitVector containing the registers that have been clobbered, and
   /// not authenticated.
   BitVector NonAutClobRegs;
   /// A vector of sets, only used in the second data flow run.
   /// Each element in the vector represents one of the registers for which we
   /// track the set of last instructions that wrote to this register. For
   /// pac-ret analysis, the expectation is that almost all return instructions
   /// only use register `X30`, and therefore, this vector will probably have
   /// length 1 in the second run.
   std::vector<SmallPtrSet<const MCInst *, 4>> LastInstWritingReg;
   State() {}
   State(unsigned NumRegs, unsigned NumRegsToTrack)
       : NonAutClobRegs(NumRegs), LastInstWritingReg(NumRegsToTrack) {}
   State &operator|=(const State &StateIn) {
     NonAutClobRegs |= StateIn.NonAutClobRegs;
     for (unsigned I = 0; I < LastInstWritingReg.size(); ++I)
       for (const MCInst *J : StateIn.LastInstWritingReg[I])
         LastInstWritingReg[I].insert(J);
     return *this;
   }
   bool operator==(const State &RHS) const {
     return NonAutClobRegs == RHS.NonAutClobRegs &&
            LastInstWritingReg == RHS.LastInstWritingReg;
   }
   bool operator!=(const State &RHS) const { return !((*this) == RHS); }
 };

 static void printLastInsts(
     raw_ostream &OS,
     const std::vector<SmallPtrSet<const MCInst *, 4>> &LastInstWritingReg) {
   OS << "Insts: ";
   for (unsigned I = 0; I < LastInstWritingReg.size(); ++I) {
     auto &Set = LastInstWritingReg[I];
     OS << "[" << I << "](";
     for (const MCInst *MCInstP : Set)
       OS << MCInstP << " ";
     OS << ")";
   }
 }

 raw_ostream &operator<<(raw_ostream &OS, const State &S) {
   OS << "pacret-state<";
   OS << "NonAutClobRegs: " << S.NonAutClobRegs << ", ";
   printLastInsts(OS, S.LastInstWritingReg);
   OS << ">";
   return OS;
 }

 class PacStatePrinter {
 public:
   void print(raw_ostream &OS, const State &State) const;
   explicit PacStatePrinter(const BinaryContext &BC) : BC(BC) {}

 private:
   const BinaryContext &BC;
 };

 void PacStatePrinter::print(raw_ostream &OS, const State &S) const {
   RegStatePrinter RegStatePrinter(BC);
   OS << "pacret-state<";
   OS << "NonAutClobRegs: ";
   RegStatePrinter.print(OS, S.NonAutClobRegs);
   OS << ", ";
   printLastInsts(OS, S.LastInstWritingReg);
   OS << ">";
 }

 class PacRetAnalysis
     : public DataflowAnalysis<PacRetAnalysis, State, /*Backward=*/false,
                               PacStatePrinter> {
   using Parent =
       DataflowAnalysis<PacRetAnalysis, State, false, PacStatePrinter>;
   friend Parent;

 public:
   PacRetAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId,
                  const std::vector<MCPhysReg> &RegsToTrackInstsFor)
       : Parent(BF, AllocId), NumRegs(BF.getBinaryContext().MRI->getNumRegs()),
         RegsToTrackInstsFor(RegsToTrackInstsFor),
         TrackingLastInsts(!RegsToTrackInstsFor.empty()),
         Reg2StateIdx(RegsToTrackInstsFor.empty()
                          ? 0
                          : *llvm::max_element(RegsToTrackInstsFor) + 1,
                      -1) {
     for (unsigned I = 0; I < RegsToTrackInstsFor.size(); ++I)
       Reg2StateIdx[RegsToTrackInstsFor[I]] = I;
   }
   virtual ~PacRetAnalysis() {}

 protected:
   const unsigned NumRegs;
   /// RegToTrackInstsFor is the set of registers for which the dataflow analysis
   /// must compute which the last set of instructions writing to it are.
   const std::vector<MCPhysReg> RegsToTrackInstsFor;
   const bool TrackingLastInsts;
   /// Reg2StateIdx maps Register to the index in the vector used in State to
   /// track which instructions last wrote to this register.
   std::vector<uint16_t> Reg2StateIdx;

   SmallPtrSet<const MCInst *, 4> &lastWritingInsts(State &S,
                                                    MCPhysReg Reg) const {
     assert(Reg < Reg2StateIdx.size());
     assert(isTrackingReg(Reg));
     return S.LastInstWritingReg[Reg2StateIdx[Reg]];
   }
   const SmallPtrSet<const MCInst *, 4> &lastWritingInsts(const State &S,
                                                          MCPhysReg Reg) const {
     assert(Reg < Reg2StateIdx.size());
     assert(isTrackingReg(Reg));
     return S.LastInstWritingReg[Reg2StateIdx[Reg]];
   }

   bool isTrackingReg(MCPhysReg Reg) const {
     return llvm::is_contained(RegsToTrackInstsFor, Reg);
   }

   void preflight() {}

   State getStartingStateAtBB(const BinaryBasicBlock &BB) {
     return State(NumRegs, RegsToTrackInstsFor.size());
   }

   State getStartingStateAtPoint(const MCInst &Point) {
     return State(NumRegs, RegsToTrackInstsFor.size());
   }

   void doConfluence(State &StateOut, const State &StateIn) {
     PacStatePrinter P(BC);
     LLVM_DEBUG({
       dbgs() << " PacRetAnalysis::Confluence(\n";
       dbgs() << "   State 1: ";
       P.print(dbgs(), StateOut);
       dbgs() << "\n";
       dbgs() << "   State 2: ";
       P.print(dbgs(), StateIn);
       dbgs() << ")\n";
     });

     StateOut |= StateIn;

     LLVM_DEBUG({
       dbgs() << "   merged state: ";
       P.print(dbgs(), StateOut);
       dbgs() << "\n";
     });
   }

   State computeNext(const MCInst &Point, const State &Cur) {
     PacStatePrinter P(BC);
     LLVM_DEBUG({
       dbgs() << " PacRetAnalysis::ComputeNext(";
       BC.InstPrinter->printInst(&const_cast<MCInst &>(Point), 0, "", *BC.STI,
                                 dbgs());
       dbgs() << ", ";
       P.print(dbgs(), Cur);
       dbgs() << ")\n";
     });

     State Next = Cur;
     BitVector Written = BitVector(NumRegs, false);
     // Assume a call can clobber all registers, including callee-saved
     // registers. There's a good chance that callee-saved registers will be
     // saved on the stack at some point during execution of the callee.
     // Therefore they should also be considered as potentially modified by an
     // attacker/written to.
     // Also, not all functions may respect the AAPCS ABI rules about
     // caller/callee-saved registers.
     if (BC.MIB->isCall(Point))
       Written.set();
     else
       // FIXME: `getWrittenRegs` only sets the register directly written in the
       // instruction, and the smaller aliasing registers. It does not set the
       // larger aliasing registers. To also set the larger aliasing registers,
       // we'd have to call `getClobberedRegs`.
       // It is unclear if there is any test case which shows a different
       // behaviour between using `getWrittenRegs` vs `getClobberedRegs`. We'd
       // first would like to see such a test case before making a decision
       // on whether using `getClobberedRegs` below would be better.
       // Also see the discussion on this at
       // https://github.com/llvm/llvm-project/pull/122304#discussion_r1939511909
       BC.MIB->getWrittenRegs(Point, Written);
     Next.NonAutClobRegs |= Written;
     // Keep track of this instruction if it writes to any of the registers we
     // need to track that for:
     for (MCPhysReg Reg : RegsToTrackInstsFor)
       if (Written[Reg])
         lastWritingInsts(Next, Reg) = {&Point};

     ErrorOr<MCPhysReg> AutReg = BC.MIB->getAuthenticatedReg(Point);
     if (AutReg && *AutReg != BC.MIB->getNoRegister()) {
       // FIXME: should we use `OnlySmaller=false` below? See similar
       // FIXME about `getWrittenRegs` above and further discussion about this
       // at
       // https://github.com/llvm/llvm-project/pull/122304#discussion_r1939515516
       Next.NonAutClobRegs.reset(
           BC.MIB->getAliases(*AutReg, /*OnlySmaller=*/true));
       if (TrackingLastInsts && isTrackingReg(*AutReg))
         lastWritingInsts(Next, *AutReg).clear();
     }

     LLVM_DEBUG({
       dbgs() << "  .. result: (";
       P.print(dbgs(), Next);
       dbgs() << ")\n";
     });

     return Next;
   }

   StringRef getAnnotationName() const { return StringRef("PacRetAnalysis"); }

 public:
   std::vector<MCInstReference>
   getLastClobberingInsts(const MCInst Ret, BinaryFunction &BF,
                          const BitVector &UsedDirtyRegs) const {
     if (!TrackingLastInsts)
       return {};
     auto MaybeState = getStateAt(Ret);
     if (!MaybeState)
       llvm_unreachable("Expected State to be present");
     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 (const MCInst *Inst : lastWritingInsts(S, TrackedReg))
         LastWritingInsts.insert(Inst);
     }
     std::vector<MCInstReference> Result;
     for (const MCInst *Inst : LastWritingInsts) {
       MCInstInBBReference Ref = MCInstInBBReference::get(Inst, BF);
       assert(Ref.BB != nullptr && "Expected Inst to be found");
       Result.push_back(MCInstReference(Ref));
     }
     return Result;
   }
 };

 FunctionAnalysisResult
 Analysis::computeDfState(PacRetAnalysis &PRA, BinaryFunction &BF,
                          MCPlusBuilder::AllocatorIdTy 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<GenDiag>(
               MCInstInBBReference(&BB, I),
               "Warning: pac-ret analysis could not analyze this return "
               "instruction"));
           continue;
         }
         MCPhysReg RetReg = *MaybeRetReg;
         LLVM_DEBUG({
           dbgs() << "  Found RET inst: ";
           BC.printInstruction(dbgs(), Inst);
           dbgs() << "    RetReg: " << BC.MRI->getName(RetReg)
                  << "; authenticatesReg: "
                  << BC.MIB->isAuthenticationOfReg(Inst, RetReg) << "\n";
         });
         if (BC.MIB->isAuthenticationOfReg(Inst, RetReg))
           break;
         BitVector UsedDirtyRegs = PRA.getStateAt(Inst)->NonAutClobRegs;
         LLVM_DEBUG({
           dbgs() << "  NonAutClobRegs at Ret: ";
           RegStatePrinter RSP(BC);
           RSP.print(dbgs(), UsedDirtyRegs);
           dbgs() << "\n";
         });
         UsedDirtyRegs &= BC.MIB->getAliases(RetReg, /*OnlySmaller=*/true);
         LLVM_DEBUG({
           dbgs() << "  Intersection with RetReg: ";
           RegStatePrinter RSP(BC);
           RSP.print(dbgs(), UsedDirtyRegs);
           dbgs() << "\n";
         });
         if (UsedDirtyRegs.any()) {
           // This return instruction needs to be reported
           Result.Diagnostics.push_back(std::make_shared<Gadget>(
               MCInstInBBReference(&BB, I),
               PRA.getLastClobberingInsts(Inst, BF, UsedDirtyRegs)));
           for (MCPhysReg RetRegWithGadget : UsedDirtyRegs.set_bits())
             Result.RegistersAffected.insert(RetRegWithGadget);
         }
       }
     }
   }
   return Result;
 }

 void Analysis::runOnFunction(BinaryFunction &BF,
                              MCPlusBuilder::AllocatorIdTy AllocatorId) {
   LLVM_DEBUG({
     dbgs() << "Analyzing in function " << BF.getPrintName() << ", AllocatorId "
            << AllocatorId << "\n";
     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);
     }

     // `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;
     }
   }
 }

 static void printBB(const BinaryContext &BC, const BinaryBasicBlock *BB,
                     size_t StartIndex = 0, size_t EndIndex = -1) {
   if (EndIndex == (size_t)-1)
     EndIndex = BB->size() - 1;
   const BinaryFunction *BF = BB->getFunction();
   for (unsigned I = StartIndex; I <= EndIndex; ++I) {
     // FIXME: this assumes all instructions are 4 bytes in size. This is true
     // for AArch64, but it might be good to extract this function so it can be
     // used elsewhere and for other targets too.
     uint64_t Address = BB->getOffset() + BF->getAddress() + 4 * I;
     const MCInst &Inst = BB->getInstructionAtIndex(I);
     if (BC.MIB->isCFI(Inst))
       continue;
     BC.printInstruction(outs(), Inst, Address, BF);
   }
 }

 static void reportFoundGadgetInSingleBBSingleOverwInst(
     raw_ostream &OS, const BinaryContext &BC, const MCInstReference OverwInst,
     const MCInstReference RetInst) {
   BinaryBasicBlock *BB = RetInst.getBasicBlock();
   assert(OverwInst.ParentKind == MCInstReference::BasicBlockParent);
   assert(RetInst.ParentKind == MCInstReference::BasicBlockParent);
   MCInstInBBReference OverwInstBB = OverwInst.U.BBRef;
   if (BB == OverwInstBB.BB) {
     // overwriting inst and ret instruction are in the same basic block.
     assert(OverwInstBB.BBIndex < RetInst.U.BBRef.BBIndex);
     OS << "  This happens in the following basic block:\n";
     printBB(BC, BB);
   }
 }

 void Gadget::generateReport(raw_ostream &OS, const BinaryContext &BC) const {
   GenDiag(RetInst, "non-protected ret found").generateReport(OS, BC);

   BinaryFunction *BF = RetInst.getFunction();
   OS << "  The " << OverwritingRetRegInst.size()
      << " instructions that write to the return register after any "
         "authentication are:\n";
   // Sort by address to ensure output is deterministic.
   std::vector<MCInstReference> ORRI = OverwritingRetRegInst;
   llvm::sort(ORRI, [](const MCInstReference &A, const MCInstReference &B) {
     return A.getAddress() < B.getAddress();
   });
   for (unsigned I = 0; I < ORRI.size(); ++I) {
     MCInstReference InstRef = ORRI[I];
     OS << "  " << (I + 1) << ". ";
     BC.printInstruction(OS, InstRef, InstRef.getAddress(), BF);
   };
   LLVM_DEBUG({
     dbgs() << "  .. OverWritingRetRegInst:\n";
     for (MCInstReference Ref : OverwritingRetRegInst) {
       dbgs() << "    " << Ref << "\n";
     }
   });
   if (OverwritingRetRegInst.size() == 1) {
     const MCInstReference OverwInst = OverwritingRetRegInst[0];
     assert(OverwInst.ParentKind == MCInstReference::BasicBlockParent);
     reportFoundGadgetInSingleBBSingleOverwInst(OS, BC, OverwInst, RetInst);
   }
 }

 void GenDiag::generateReport(raw_ostream &OS, const BinaryContext &BC) const {
   BinaryFunction *BF = RetInst.getFunction();
   BinaryBasicBlock *BB = RetInst.getBasicBlock();

   OS << "\nGS-PACRET: " << Diag.Text;
   OS << " in function " << BF->getPrintName();
   if (BB)
     OS << ", basic block " << BB->getName();
   OS << ", at address " << llvm::format("%x", RetInst.getAddress()) << "\n";
   OS << "  The return instruction is ";
   BC.printInstruction(OS, RetInst, RetInst.getAddress(), BF);
 }

 Error Analysis::runOnFunctions(BinaryContext &BC) {
   ParallelUtilities::WorkFuncWithAllocTy WorkFun =
       [&](BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocatorId) {
         runOnFunction(BF, AllocatorId);
       };

   ParallelUtilities::PredicateTy SkipFunc = [&](const BinaryFunction &BF) {
     return false;
   };

   ParallelUtilities::runOnEachFunctionWithUniqueAllocId(
       BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
       SkipFunc, "NonPacProtectedRetAnalysis");

   for (BinaryFunction *BF : BC.getAllBinaryFunctions())
     if (AnalysisResults.count(BF) > 0) {
       for (const std::shared_ptr<Annotation> &A :
            AnalysisResults[BF].Diagnostics)
         A->generateReport(outs(), BC);
     }
   return Error::success();
 }

 } // namespace NonPacProtectedRetAnalysis
 } // namespace bolt
 } // namespace llvm
	//===- bolt/Passes/NonPacProtectedRetAnalysis.cpp -------------------------===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	//
	// This file implements a pass that looks for any AArch64 return instructions
	// that may not be protected by PAuth authentication instructions when needed.
	//
	//===----------------------------------------------------------------------===//

	#include "bolt/Passes/NonPacProtectedRetAnalysis.h"
	#include "bolt/Core/ParallelUtilities.h"
	#include "bolt/Passes/DataflowAnalysis.h"
	#include "llvm/ADT/SmallSet.h"
	#include "llvm/MC/MCInst.h"
	#include "llvm/Support/Format.h"
	#include <memory>

	#define DEBUG_TYPE "bolt-nonpacprotectedret"

	namespace llvm {
	namespace bolt {

	raw_ostream &operator<<(raw_ostream &OS, const MCInstInBBReference &Ref) {
	OS << "MCInstBBRef<";
	if (Ref.BB == nullptr)
	OS << "BB:(null)";
	else
	OS << "BB:" << Ref.BB->getName() << ":" << Ref.BBIndex;
	OS << ">";
	return OS;
	}

	raw_ostream &operator<<(raw_ostream &OS, const MCInstInBFReference &Ref) {
	OS << "MCInstBFRef<";
	if (Ref.BF == nullptr)
	OS << "BF:(null)";
	else
	OS << "BF:" << Ref.BF->getPrintName() << ":" << Ref.getOffset();
	OS << ">";
	return OS;
	}

	raw_ostream &operator<<(raw_ostream &OS, const MCInstReference &Ref) {
	switch (Ref.ParentKind) {
	case MCInstReference::BasicBlockParent:
	OS << Ref.U.BBRef;
	return OS;
	case MCInstReference::FunctionParent:
	OS << Ref.U.BFRef;
	return OS;
	}
	llvm_unreachable("");
	}

	namespace NonPacProtectedRetAnalysis {

	// The security property that is checked is:
	// When a register is used as the address to jump to in a return instruction,
	// that register must either:
	// (a) never be changed within this function, i.e. have the same value as when
	// the function started, or
	// (b) the last write to the register must be by an authentication instruction.

	// This property is checked by using dataflow analysis to keep track of which
	// registers have been written (def-ed), since last authenticated. Those are
	// exactly the registers containing values that should not be trusted (as they
	// could have changed since the last time they were authenticated). For pac-ret,
	// any return instruction using such a register is a gadget to be reported. For
	// PAuthABI, probably at least any indirect control flow using such a register
	// should be reported.

	// Furthermore, when producing a diagnostic for a found non-pac-ret protected
	// return, the analysis also lists the last instructions that wrote to the
	// register used in the return instruction.
	// The total set of registers used in return instructions in a given function is
	// small. It almost always is just `X30`.
	// In order to reduce the memory consumption of storing this additional state
	// during the dataflow analysis, this is computed by running the dataflow
	// analysis twice:
	// 1. In the first run, the dataflow analysis only keeps track of the security
	// property: i.e. which registers have been overwritten since the last
	// time they've been authenticated.
	// 2. If the first run finds any return instructions using a register last
	// written by a non-authenticating instruction, the dataflow analysis will
	// be run a second time. The first run will return which registers are used
	// in the gadgets to be reported. This information is used in the second run
	// to also track which instructions last wrote to those registers.

	struct State {
	/// A BitVector containing the registers that have been clobbered, and
	/// not authenticated.
	BitVector NonAutClobRegs;
	/// A vector of sets, only used in the second data flow run.
	/// Each element in the vector represents one of the registers for which we
	/// track the set of last instructions that wrote to this register. For
	/// pac-ret analysis, the expectation is that almost all return instructions
	/// only use register `X30`, and therefore, this vector will probably have
	/// length 1 in the second run.
	std::vector<SmallPtrSet<const MCInst *, 4>> LastInstWritingReg;
	State() {}
	State(unsigned NumRegs, unsigned NumRegsToTrack)
	: NonAutClobRegs(NumRegs), LastInstWritingReg(NumRegsToTrack) {}
	State &operator\|=(const State &StateIn) {
	NonAutClobRegs \|= StateIn.NonAutClobRegs;
	for (unsigned I = 0; I < LastInstWritingReg.size(); ++I)
	for (const MCInst *J : StateIn.LastInstWritingReg[I])
	LastInstWritingReg[I].insert(J);
	return *this;
	}
	bool operator==(const State &RHS) const {
	return NonAutClobRegs == RHS.NonAutClobRegs &&
	LastInstWritingReg == RHS.LastInstWritingReg;
	}
	bool operator!=(const State &RHS) const { return !((*this) == RHS); }
	};

	static void printLastInsts(
	raw_ostream &OS,
	const std::vector<SmallPtrSet<const MCInst *, 4>> &LastInstWritingReg) {
	OS << "Insts: ";
	for (unsigned I = 0; I < LastInstWritingReg.size(); ++I) {
	auto &Set = LastInstWritingReg[I];
	OS << "[" << I << "](";
	for (const MCInst *MCInstP : Set)
	OS << MCInstP << " ";
	OS << ")";
	}
	}

	raw_ostream &operator<<(raw_ostream &OS, const State &S) {
	OS << "pacret-state<";
	OS << "NonAutClobRegs: " << S.NonAutClobRegs << ", ";
	printLastInsts(OS, S.LastInstWritingReg);
	OS << ">";
	return OS;
	}

	class PacStatePrinter {
	public:
	void print(raw_ostream &OS, const State &State) const;
	explicit PacStatePrinter(const BinaryContext &BC) : BC(BC) {}

	private:
	const BinaryContext &BC;
	};

	void PacStatePrinter::print(raw_ostream &OS, const State &S) const {
	RegStatePrinter RegStatePrinter(BC);
	OS << "pacret-state<";
	OS << "NonAutClobRegs: ";
	RegStatePrinter.print(OS, S.NonAutClobRegs);
	OS << ", ";
	printLastInsts(OS, S.LastInstWritingReg);
	OS << ">";
	}

	class PacRetAnalysis
	: public DataflowAnalysis<PacRetAnalysis, State, /Backward=/false,
	PacStatePrinter> {
	using Parent =
	DataflowAnalysis<PacRetAnalysis, State, false, PacStatePrinter>;
	friend Parent;

	public:
	PacRetAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId,
	const std::vector<MCPhysReg> &RegsToTrackInstsFor)
	: Parent(BF, AllocId), NumRegs(BF.getBinaryContext().MRI->getNumRegs()),
	RegsToTrackInstsFor(RegsToTrackInstsFor),
	TrackingLastInsts(!RegsToTrackInstsFor.empty()),
	Reg2StateIdx(RegsToTrackInstsFor.empty()
	? 0
	: *llvm::max_element(RegsToTrackInstsFor) + 1,
	-1) {
	for (unsigned I = 0; I < RegsToTrackInstsFor.size(); ++I)
	Reg2StateIdx[RegsToTrackInstsFor[I]] = I;
	}
	virtual ~PacRetAnalysis() {}

	protected:
	const unsigned NumRegs;
	/// RegToTrackInstsFor is the set of registers for which the dataflow analysis
	/// must compute which the last set of instructions writing to it are.
	const std::vector<MCPhysReg> RegsToTrackInstsFor;
	const bool TrackingLastInsts;
	/// Reg2StateIdx maps Register to the index in the vector used in State to
	/// track which instructions last wrote to this register.
	std::vector<uint16_t> Reg2StateIdx;

	SmallPtrSet<const MCInst *, 4> &lastWritingInsts(State &S,
	MCPhysReg Reg) const {
	assert(Reg < Reg2StateIdx.size());
	assert(isTrackingReg(Reg));
	return S.LastInstWritingReg[Reg2StateIdx[Reg]];
	}
	const SmallPtrSet<const MCInst *, 4> &lastWritingInsts(const State &S,
	MCPhysReg Reg) const {
	assert(Reg < Reg2StateIdx.size());
	assert(isTrackingReg(Reg));
	return S.LastInstWritingReg[Reg2StateIdx[Reg]];
	}

	bool isTrackingReg(MCPhysReg Reg) const {
	return llvm::is_contained(RegsToTrackInstsFor, Reg);
	}

	void preflight() {}

	State getStartingStateAtBB(const BinaryBasicBlock &BB) {
	return State(NumRegs, RegsToTrackInstsFor.size());
	}

	State getStartingStateAtPoint(const MCInst &Point) {
	return State(NumRegs, RegsToTrackInstsFor.size());
	}

	void doConfluence(State &StateOut, const State &StateIn) {
	PacStatePrinter P(BC);
	LLVM_DEBUG({
	dbgs() << " PacRetAnalysis::Confluence(\n";
	dbgs() << " State 1: ";
	P.print(dbgs(), StateOut);
	dbgs() << "\n";
	dbgs() << " State 2: ";
	P.print(dbgs(), StateIn);
	dbgs() << ")\n";
	});

	StateOut \|= StateIn;

	LLVM_DEBUG({
	dbgs() << " merged state: ";
	P.print(dbgs(), StateOut);
	dbgs() << "\n";
	});
	}

	State computeNext(const MCInst &Point, const State &Cur) {
	PacStatePrinter P(BC);
	LLVM_DEBUG({
	dbgs() << " PacRetAnalysis::ComputeNext(";
	BC.InstPrinter->printInst(&const_cast<MCInst &>(Point), 0, "", *BC.STI,
	dbgs());
	dbgs() << ", ";
	P.print(dbgs(), Cur);
	dbgs() << ")\n";
	});

	State Next = Cur;
	BitVector Written = BitVector(NumRegs, false);
	// Assume a call can clobber all registers, including callee-saved
	// registers. There's a good chance that callee-saved registers will be
	// saved on the stack at some point during execution of the callee.
	// Therefore they should also be considered as potentially modified by an
	// attacker/written to.
	// Also, not all functions may respect the AAPCS ABI rules about
	// caller/callee-saved registers.
	if (BC.MIB->isCall(Point))
	Written.set();
	else
	// FIXME: `getWrittenRegs` only sets the register directly written in the
	// instruction, and the smaller aliasing registers. It does not set the
	// larger aliasing registers. To also set the larger aliasing registers,
	// we'd have to call `getClobberedRegs`.
	// It is unclear if there is any test case which shows a different
	// behaviour between using `getWrittenRegs` vs `getClobberedRegs`. We'd
	// first would like to see such a test case before making a decision
	// on whether using `getClobberedRegs` below would be better.
	// Also see the discussion on this at
	// https://github.com/llvm/llvm-project/pull/122304#discussion_r1939511909
	BC.MIB->getWrittenRegs(Point, Written);
	Next.NonAutClobRegs \|= Written;
	// Keep track of this instruction if it writes to any of the registers we
	// need to track that for:
	for (MCPhysReg Reg : RegsToTrackInstsFor)
	if (Written[Reg])
	lastWritingInsts(Next, Reg) = {&Point};

	ErrorOr<MCPhysReg> AutReg = BC.MIB->getAuthenticatedReg(Point);
	if (AutReg && *AutReg != BC.MIB->getNoRegister()) {
	// FIXME: should we use `OnlySmaller=false` below? See similar
	// FIXME about `getWrittenRegs` above and further discussion about this
	// at
	// https://github.com/llvm/llvm-project/pull/122304#discussion_r1939515516
	Next.NonAutClobRegs.reset(
	BC.MIB->getAliases(AutReg, /OnlySmaller=*/true));
	if (TrackingLastInsts && isTrackingReg(*AutReg))
	lastWritingInsts(Next, *AutReg).clear();
	}

	LLVM_DEBUG({
	dbgs() << " .. result: (";
	P.print(dbgs(), Next);
	dbgs() << ")\n";
	});

	return Next;
	}

	StringRef getAnnotationName() const { return StringRef("PacRetAnalysis"); }

	public:
	std::vector<MCInstReference>
	getLastClobberingInsts(const MCInst Ret, BinaryFunction &BF,
	const BitVector &UsedDirtyRegs) const {
	if (!TrackingLastInsts)
	return {};
	auto MaybeState = getStateAt(Ret);
	if (!MaybeState)
	llvm_unreachable("Expected State to be present");
	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 (const MCInst *Inst : lastWritingInsts(S, TrackedReg))
	LastWritingInsts.insert(Inst);
	}
	std::vector<MCInstReference> Result;
	for (const MCInst *Inst : LastWritingInsts) {
	MCInstInBBReference Ref = MCInstInBBReference::get(Inst, BF);
	assert(Ref.BB != nullptr && "Expected Inst to be found");
	Result.push_back(MCInstReference(Ref));
	}
	return Result;
	}
	};

	FunctionAnalysisResult
	Analysis::computeDfState(PacRetAnalysis &PRA, BinaryFunction &BF,
	MCPlusBuilder::AllocatorIdTy 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<GenDiag>(
	MCInstInBBReference(&BB, I),
	"Warning: pac-ret analysis could not analyze this return "
	"instruction"));
	continue;
	}
	MCPhysReg RetReg = *MaybeRetReg;
	LLVM_DEBUG({
	dbgs() << " Found RET inst: ";
	BC.printInstruction(dbgs(), Inst);
	dbgs() << " RetReg: " << BC.MRI->getName(RetReg)
	<< "; authenticatesReg: "
	<< BC.MIB->isAuthenticationOfReg(Inst, RetReg) << "\n";
	});
	if (BC.MIB->isAuthenticationOfReg(Inst, RetReg))
	break;
	BitVector UsedDirtyRegs = PRA.getStateAt(Inst)->NonAutClobRegs;
	LLVM_DEBUG({
	dbgs() << " NonAutClobRegs at Ret: ";
	RegStatePrinter RSP(BC);
	RSP.print(dbgs(), UsedDirtyRegs);
	dbgs() << "\n";
	});
	UsedDirtyRegs &= BC.MIB->getAliases(RetReg, /OnlySmaller=/true);
	LLVM_DEBUG({
	dbgs() << " Intersection with RetReg: ";
	RegStatePrinter RSP(BC);
	RSP.print(dbgs(), UsedDirtyRegs);
	dbgs() << "\n";
	});
	if (UsedDirtyRegs.any()) {
	// This return instruction needs to be reported
	Result.Diagnostics.push_back(std::make_shared<Gadget>(
	MCInstInBBReference(&BB, I),
	PRA.getLastClobberingInsts(Inst, BF, UsedDirtyRegs)));
	for (MCPhysReg RetRegWithGadget : UsedDirtyRegs.set_bits())
	Result.RegistersAffected.insert(RetRegWithGadget);
	}
	}
	}
	}
	return Result;
	}

	void Analysis::runOnFunction(BinaryFunction &BF,
	MCPlusBuilder::AllocatorIdTy AllocatorId) {
	LLVM_DEBUG({
	dbgs() << "Analyzing in function " << BF.getPrintName() << ", AllocatorId "
	<< AllocatorId << "\n";
	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);
	}

	// `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;
	}
	}
	}

	static void printBB(const BinaryContext &BC, const BinaryBasicBlock *BB,
	size_t StartIndex = 0, size_t EndIndex = -1) {
	if (EndIndex == (size_t)-1)
	EndIndex = BB->size() - 1;
	const BinaryFunction *BF = BB->getFunction();
	for (unsigned I = StartIndex; I <= EndIndex; ++I) {
	// FIXME: this assumes all instructions are 4 bytes in size. This is true
	// for AArch64, but it might be good to extract this function so it can be
	// used elsewhere and for other targets too.
	uint64_t Address = BB->getOffset() + BF->getAddress() + 4 * I;
	const MCInst &Inst = BB->getInstructionAtIndex(I);
	if (BC.MIB->isCFI(Inst))
	continue;
	BC.printInstruction(outs(), Inst, Address, BF);
	}
	}

	static void reportFoundGadgetInSingleBBSingleOverwInst(
	raw_ostream &OS, const BinaryContext &BC, const MCInstReference OverwInst,
	const MCInstReference RetInst) {
	BinaryBasicBlock *BB = RetInst.getBasicBlock();
	assert(OverwInst.ParentKind == MCInstReference::BasicBlockParent);
	assert(RetInst.ParentKind == MCInstReference::BasicBlockParent);
	MCInstInBBReference OverwInstBB = OverwInst.U.BBRef;
	if (BB == OverwInstBB.BB) {
	// overwriting inst and ret instruction are in the same basic block.
	assert(OverwInstBB.BBIndex < RetInst.U.BBRef.BBIndex);
	OS << " This happens in the following basic block:\n";
	printBB(BC, BB);
	}
	}

	void Gadget::generateReport(raw_ostream &OS, const BinaryContext &BC) const {
	GenDiag(RetInst, "non-protected ret found").generateReport(OS, BC);

	BinaryFunction *BF = RetInst.getFunction();
	OS << " The " << OverwritingRetRegInst.size()
	<< " instructions that write to the return register after any "
	"authentication are:\n";
	// Sort by address to ensure output is deterministic.
	std::vector<MCInstReference> ORRI = OverwritingRetRegInst;
	llvm::sort(ORRI, [](const MCInstReference &A, const MCInstReference &B) {
	return A.getAddress() < B.getAddress();
	});
	for (unsigned I = 0; I < ORRI.size(); ++I) {
	MCInstReference InstRef = ORRI[I];
	OS << " " << (I + 1) << ". ";
	BC.printInstruction(OS, InstRef, InstRef.getAddress(), BF);
	};
	LLVM_DEBUG({
	dbgs() << " .. OverWritingRetRegInst:\n";
	for (MCInstReference Ref : OverwritingRetRegInst) {
	dbgs() << " " << Ref << "\n";
	}
	});
	if (OverwritingRetRegInst.size() == 1) {
	const MCInstReference OverwInst = OverwritingRetRegInst[0];
	assert(OverwInst.ParentKind == MCInstReference::BasicBlockParent);
	reportFoundGadgetInSingleBBSingleOverwInst(OS, BC, OverwInst, RetInst);
	}
	}

	void GenDiag::generateReport(raw_ostream &OS, const BinaryContext &BC) const {
	BinaryFunction *BF = RetInst.getFunction();
	BinaryBasicBlock *BB = RetInst.getBasicBlock();

	OS << "\nGS-PACRET: " << Diag.Text;
	OS << " in function " << BF->getPrintName();
	if (BB)
	OS << ", basic block " << BB->getName();
	OS << ", at address " << llvm::format("%x", RetInst.getAddress()) << "\n";
	OS << " The return instruction is ";
	BC.printInstruction(OS, RetInst, RetInst.getAddress(), BF);
	}

	Error Analysis::runOnFunctions(BinaryContext &BC) {
	ParallelUtilities::WorkFuncWithAllocTy WorkFun =
	[&](BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocatorId) {
	runOnFunction(BF, AllocatorId);
	};

	ParallelUtilities::PredicateTy SkipFunc = [&](const BinaryFunction &BF) {
	return false;
	};

	ParallelUtilities::runOnEachFunctionWithUniqueAllocId(
	BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
	SkipFunc, "NonPacProtectedRetAnalysis");

	for (BinaryFunction *BF : BC.getAllBinaryFunctions())
	if (AnalysisResults.count(BF) > 0) {
	for (const std::shared_ptr<Annotation> &A :
	AnalysisResults[BF].Diagnostics)
	A->generateReport(outs(), BC);
	}
	return Error::success();
	}

	} // namespace NonPacProtectedRetAnalysis
	} // namespace bolt
	} // namespace llvm