[BOLT] Gadget scanner: detect untrusted LR before tail call

Implement the detection of tail calls performed with untrusted link
register, which violates the assumption made on entry to every function.

Unlike other pauth gadgets, this one involves some amount of guessing
which branch instructions should be checked as tail calls.

Author: atrosinenko

bolt/lib/Passes/PAuthGadgetScanner.cpp

@@ -655,8 +655,9 @@ class DataflowSrcSafetyAnalysis
 //
 // Then, a function can be split into a number of disjoint contiguous sequences
 // of instructions without labels in between. These sequences can be processed
-// the same way basic blocks are processed by data-flow analysis, assuming
-// pessimistically that all registers are unsafe at the start of each sequence.
+// the same way basic blocks are processed by data-flow analysis, with the same
+// pessimistic estimation of the initial state at the start of each sequence
+// (except the first instruction of the function).
 class CFGUnawareSrcSafetyAnalysis : public SrcSafetyAnalysis {
   BinaryFunction &BF;
   MCPlusBuilder::AllocatorIdTy AllocId;
@@ -667,6 +668,30 @@ class CFGUnawareSrcSafetyAnalysis : public SrcSafetyAnalysis {
       BC.MIB->removeAnnotation(I.second, StateAnnotationIndex);
   }
 
+  /// Compute a reasonably pessimistic estimation of the register state when
+  /// the previous instruction is not known for sure. Take the set of registers
+  /// which are trusted at function entry and remove all registers that can be
+  /// clobbered inside this function.
+  SrcState computePessimisticState(BinaryFunction &BF) {
+    BitVector ClobberedRegs(NumRegs);
+    for (auto &I : BF.instrs()) {
+      MCInst &Inst = I.second;
+      BC.MIB->getClobberedRegs(Inst, ClobberedRegs);
+
+      // If this is a call instruction, no register is safe anymore, unless
+      // it is a tail call. Ignore tail calls for the purpose of estimating the
+      // worst-case scenario, assuming no instructions are executed in the
+      // caller after this point anyway.
+      if (BC.MIB->isCall(Inst) && !BC.MIB->isTailCall(Inst))
+        ClobberedRegs.set();
+    }
+
+    SrcState S = createEntryState();
+    S.SafeToDerefRegs.reset(ClobberedRegs);
+    S.TrustedRegs.reset(ClobberedRegs);
+    return S;
+  }
+
 public:
   CFGUnawareSrcSafetyAnalysis(BinaryFunction &BF,
                               MCPlusBuilder::AllocatorIdTy AllocId,
@@ -677,6 +702,7 @@ class CFGUnawareSrcSafetyAnalysis : public SrcSafetyAnalysis {
   }
 
   void run() override {
+    const SrcState DefaultState = computePessimisticState(BF);
     SrcState S = createEntryState();
     for (auto &I : BF.instrs()) {
       MCInst &Inst = I.second;
@@ -691,7 +717,7 @@ class CFGUnawareSrcSafetyAnalysis : public SrcSafetyAnalysis {
         LLVM_DEBUG({
           traceInst(BC, "Due to label, resetting the state before", Inst);
         });
-        S = createUnsafeState();
+        S = DefaultState;
       }
 
       // Check if we need to remove an old annotation (this is the case if
@@ -1226,6 +1252,83 @@ shouldReportReturnGadget(const BinaryContext &BC, const MCInstReference &Inst,
   return make_report(RetKind, Inst, *RetReg);
 }
 
+/// While BOLT already marks some of the branch instructions as tail calls,
+/// this function tries to improve the coverage by including less obvious cases
+/// when it is possible to do without introducing too many false positives.
+static bool shouldAnalyzeTailCallInst(const BinaryContext &BC,
+                                      const BinaryFunction &BF,
+                                      const MCInstReference &Inst) {
+  // Some BC.MIB->isXYZ(Inst) methods simply delegate to MCInstrDesc::isXYZ()
+  // (such as isBranch at the time of writing this comment), some don't (such
+  // as isCall). For that reason, call MCInstrDesc's methods explicitly when
+  // it is important.
+  const MCInstrDesc &Desc =
+      BC.MII->get(static_cast<const MCInst &>(Inst).getOpcode());
+  // Tail call should be a branch (but not necessarily an indirect one).
+  if (!Desc.isBranch())
+    return false;
+
+  // Always analyze the branches already marked as tail calls by BOLT.
+  if (BC.MIB->isTailCall(Inst))
+    return true;
+
+  // Try to also check the branches marked as "UNKNOWN CONTROL FLOW" - the
+  // below is a simplified condition from BinaryContext::printInstruction.
+  bool IsUnknownControlFlow =
+      BC.MIB->isIndirectBranch(Inst) && !BC.MIB->getJumpTable(Inst);
+
+  if (BF.hasCFG() && IsUnknownControlFlow)
+    return true;
+
+  return false;
+}
+
+static std::optional<BriefReport<MCPhysReg>>
+shouldReportUnsafeTailCall(const BinaryContext &BC, const BinaryFunction &BF,
+                           const MCInstReference &Inst, const SrcState &S) {
+  static const GadgetKind UntrustedLRKind(
+      "untrusted link register found before tail call");
+
+  if (!shouldAnalyzeTailCallInst(BC, BF, Inst))
+    return std::nullopt;
+
+  // Not only the set of registers returned by getTrustedLiveInRegs() can be
+  // seen as a reasonable target-independent _approximation_ of "the LR", these
+  // are *exactly* those registers used by SrcSafetyAnalysis to initialize the
+  // set of trusted registers on function entry.
+  // Thus, this function basically checks that the precondition expected to be
+  // imposed by a function call instruction (which is hardcoded into the target-
+  // specific getTrustedLiveInRegs() function) is also respected on tail calls.
+  SmallVector<MCPhysReg> RegsToCheck = BC.MIB->getTrustedLiveInRegs();
+  LLVM_DEBUG({
+    traceInst(BC, "Found tail call inst", Inst);
+    traceRegMask(BC, "Trusted regs", S.TrustedRegs);
+  });
+
+  // In musl on AArch64, the _start function sets LR to zero and calls the next
+  // stage initialization function at the end, something along these lines:
+  //
+  //   _start:
+  //     mov     x30, #0
+  //     ; ... other initialization ...
+  //     b       _start_c ; performs "exit" system call at some point
+  //
+  // As this would produce a false positive for every executable linked with
+  // such libc, ignore tail calls performed by ELF entry function.
+  if (BC.StartFunctionAddress &&
+      *BC.StartFunctionAddress == Inst.getFunction()->getAddress()) {
+    LLVM_DEBUG({ dbgs() << "  Skipping tail call in ELF entry function.\n"; });
+    return std::nullopt;
+  }
+
+  // Returns at most one report per instruction - this is probably OK...
+  for (auto Reg : RegsToCheck)
+    if (!S.TrustedRegs[Reg])
+      return make_report(UntrustedLRKind, Inst, Reg);
+
+  return std::nullopt;
+}
+
 static std::optional<BriefReport<MCPhysReg>>
 shouldReportCallGadget(const BinaryContext &BC, const MCInstReference &Inst,
                        const SrcState &S) {
@@ -1366,6 +1469,9 @@ void FunctionAnalysis::findUnsafeUses(
     if (PacRetGadgetsOnly)
       return;
 
+    if (auto Report = shouldReportUnsafeTailCall(BC, BF, Inst, S))
+      Reports.push_back(*Report);
+
     if (auto Report = shouldReportCallGadget(BC, Inst, S))
       Reports.push_back(*Report);
     if (auto Report = shouldReportSigningOracle(BC, Inst, S))

bolt/test/binary-analysis/AArch64/gs-pacret-autiasp.s

@@ -229,20 +229,33 @@ f_unreachable_instruction:
         ret
         .size f_unreachable_instruction, .-f_unreachable_instruction
 
-// Expected false positive: without CFG, the state is reset to all-unsafe
-// after an unconditional branch.
-
-        .globl  state_is_reset_after_indirect_branch_nocfg
-        .type   state_is_reset_after_indirect_branch_nocfg,@function
-state_is_reset_after_indirect_branch_nocfg:
-// CHECK-LABEL: GS-PAUTH: non-protected ret found in function state_is_reset_after_indirect_branch_nocfg, at address
-// CHECK-NEXT:  The instruction is     {{[0-9a-f]+}}:         ret
+// Without CFG, the state is reset at labels, assuming every register that can
+// be clobbered in the function was actually clobbered.
+
+        .globl  lr_untouched_nocfg
+        .type   lr_untouched_nocfg,@function
+lr_untouched_nocfg:
+// CHECK-NOT: lr_untouched_nocfg
+        adr     x2, 1f
+        br      x2
+1:
+        ret
+        .size lr_untouched_nocfg, .-lr_untouched_nocfg
+
+        .globl  lr_clobbered_nocfg
+        .type   lr_clobbered_nocfg,@function
+lr_clobbered_nocfg:
+// CHECK-LABEL: GS-PAUTH: non-protected ret found in function lr_clobbered_nocfg, at address
+// CHECK-NEXT:  The instruction is     {{[0-9a-f]+}}:      ret
 // CHECK-NEXT:  The 0 instructions that write to the affected registers after any authentication are:
         adr     x2, 1f
         br      x2
 1:
+        b       2f
+        bl      g   // never executed, but affects the expected worst-case scenario
+2:
         ret
-        .size state_is_reset_after_indirect_branch_nocfg, .-state_is_reset_after_indirect_branch_nocfg
+        .size lr_clobbered_nocfg, .-lr_clobbered_nocfg
 
 /// Now do a basic sanity check on every different Authentication instruction:
 

bolt/test/binary-analysis/AArch64/gs-pauth-debug-output.s

@@ -199,8 +199,8 @@ nocfg:
 // CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(  br      x0, src-state<SafeToDerefRegs: LR W0 W30 X0 W0_HI W30_HI , TrustedRegs: LR W0 W30 X0 W0_HI W30_HI , Insts: >)
 // CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: LR W0 W30 X0 W0_HI W30_HI , TrustedRegs: LR W0 W30 X0 W0_HI W30_HI , Insts: >)
 // CHECK-NEXT:   Due to label, resetting the state before:     00000000:       ret # Offset: 8
-// CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(  ret     x30, src-state<SafeToDerefRegs: , TrustedRegs: , Insts: >)
-// CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: , TrustedRegs: , Insts: >)
+// CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(  ret     x30, src-state<SafeToDerefRegs: LR W30 W30_HI , TrustedRegs: LR W30 W30_HI , Insts: >)
+// CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: LR W30 W30_HI , TrustedRegs: LR W30 W30_HI , Insts: >)
 // CHECK-NEXT: After src register safety analysis:
 // CHECK-NEXT: Binary Function "nocfg"  {
 // CHECK-NEXT:   Number      : 3
@@ -224,32 +224,6 @@ nocfg:
 // CHECK-NEXT:   Found RET inst:     00000000:         ret # Offset: 8 # CFGUnawareSrcSafetyAnalysis: src-state<SafeToDerefRegs: BitVector, TrustedRegs: BitVector, Insts: >
 // CHECK-NEXT:     RetReg: LR
 // CHECK-NEXT:     SafeToDerefRegs:
-// CHECK-EMPTY:
-// CHECK-NEXT: Running detailed src register safety analysis...
-// CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(   adr     x0, __ENTRY_nocfg@0x[[ENTRY_ADDR]], src-state<SafeToDerefRegs: LR W30 W30_HI , TrustedRegs: LR W30 W30_HI , Insts: [0]()>)
-// CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: LR W0 W30 X0 W0_HI W30_HI , TrustedRegs: LR W0 W30 X0 W0_HI W30_HI , Insts: [0]()>)
-// CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(  br      x0, src-state<SafeToDerefRegs: LR W0 W30 X0 W0_HI W30_HI , TrustedRegs: LR W0 W30 X0 W0_HI W30_HI , Insts: [0]()>)
-// CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: LR W0 W30 X0 W0_HI W30_HI , TrustedRegs: LR W0 W30 X0 W0_HI W30_HI , Insts: [0]()>)
-// CHECK-NEXT:   Due to label, resetting the state before:     00000000:       ret # Offset: 8
-// CHECK-NEXT:   SrcSafetyAnalysis::ComputeNext(  ret     x30, src-state<SafeToDerefRegs: , TrustedRegs: , Insts: [0]()>)
-// CHECK-NEXT:     .. result: (src-state<SafeToDerefRegs: , TrustedRegs: , Insts: [0]()>)
-// CHECK-NEXT: After detailed src register safety analysis:
-// CHECK-NEXT: Binary Function "nocfg"  {
-// CHECK-NEXT:   Number      : 3
-// ...
-// CHECK:        Secondary Entry Points : __ENTRY_nocfg@0x[[ENTRY_ADDR]]
-// CHECK-NEXT: }
-// CHECK-NEXT: .{{[A-Za-z0-9]+}}:
-// CHECK-NEXT:     00000000:   adr     x0, __ENTRY_nocfg@0x[[ENTRY_ADDR]] # CFGUnawareSrcSafetyAnalysis: src-state<SafeToDerefRegs: BitVector, TrustedRegs: BitVector, Insts: [0]()>
-// CHECK-NEXT:     00000004:   br      x0 # UNKNOWN CONTROL FLOW # Offset: 4 # CFGUnawareSrcSafetyAnalysis: src-state<SafeToDerefRegs: BitVector, TrustedRegs: BitVector, Insts: [0]()>
-// CHECK-NEXT: __ENTRY_nocfg@0x[[ENTRY_ADDR]] (Entry Point):
-// CHECK-NEXT: .{{[A-Za-z0-9]+}}:
-// CHECK-NEXT:     00000008:   ret # Offset: 8 # CFGUnawareSrcSafetyAnalysis: src-state<SafeToDerefRegs: BitVector, TrustedRegs: BitVector, Insts: [0]()>
-// CHECK-NEXT: DWARF CFI Instructions:
-// CHECK-NEXT:     <empty>
-// CHECK-NEXT: End of Function "nocfg"
-// CHECK-EMPTY:
-// CHECK-NEXT:   Attaching clobbering info to:     00000000:   ret # Offset: 8 # CFGUnawareSrcSafetyAnalysis: src-state<SafeToDerefRegs: BitVector, TrustedRegs: BitVector, Insts: [0]()>
 
         .globl  auth_oracle
         .type   auth_oracle,@function