[BOLT] Gadget scanner: Detect address materialization and arithmetics

In addition to authenticated pointers, consider the contents of a
register safe if it was
* written by PC-relative address computation
* updated by an arithmetic instruction whose input address is safe

Author: atrosinenko

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -587,6 +587,22 @@ class MCPlusBuilder {
     return getNoRegister();
   }
 
+  virtual MCPhysReg getSafelyMaterializedAddressReg(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return getNoRegister();
+  }
+
+  /// Analyzes if this instruction can safely perform address arithmetics.
+  ///
+  /// If the first element of the returned pair is no-register, this instruction
+  /// is considered unknown. Otherwise, (output, input) pair is returned,
+  /// so that output is as trusted as input is.
+  virtual std::pair<MCPhysReg, MCPhysReg>
+  analyzeSafeAddressArithmetics(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return std::make_pair(getNoRegister(), getNoRegister());
+  }
+
   virtual bool isTerminator(const MCInst &Inst) const;
 
   virtual bool isNoop(const MCInst &Inst) const {

bolt/lib/Passes/PAuthGadgetScanner.cpp

@@ -325,18 +325,7 @@ class PacRetAnalysis
     });
   }
 
-  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 getClobberedRegs(const MCInst &Point) const {
     BitVector Clobbered(NumRegs, false);
     // Assume a call can clobber all registers, including callee-saved
     // registers. There's a good chance that callee-saved registers will be
@@ -349,24 +338,81 @@ class PacRetAnalysis
       Clobbered.set();
     else
       BC.MIB->getClobberedRegs(Point, Clobbered);
+    return Clobbered;
+  }
+
+  // Returns all registers that can be treated as if they are written by an
+  // authentication instruction.
+  SmallVector<MCPhysReg> getAuthenticatedRegs(const MCInst &Point,
+                                              const State &Cur) const {
+    SmallVector<MCPhysReg> Regs;
+    const MCPhysReg NoReg = BC.MIB->getNoRegister();
+
+    // A signed pointer can be authenticated, or
+    ErrorOr<MCPhysReg> AutReg = BC.MIB->getAuthenticatedReg(Point);
+    if (AutReg && *AutReg != NoReg)
+      Regs.push_back(*AutReg);
+
+    // ... a safe address can be materialized, or
+    MCPhysReg NewAddrReg = BC.MIB->getSafelyMaterializedAddressReg(Point);
+    if (NewAddrReg != NoReg)
+      Regs.push_back(NewAddrReg);
+
+    // ... address can be updated in a safe manner, producing the result
+    // which is as trusted as the input address.
+    MCPhysReg ArithResult, ArithSrc;
+    std::tie(ArithResult, ArithSrc) =
+        BC.MIB->analyzeSafeAddressArithmetics(Point);
+    if (ArithResult != NoReg && Cur.SafeToDerefRegs[ArithSrc])
+      Regs.push_back(ArithResult);
+
+    return Regs;
+  }
+
+  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";
+    });
+
+    // First, compute various properties of the instruction, taking the state
+    // before its execution into account, if necessary.
+
+    BitVector Clobbered = getClobberedRegs(Point);
+    // Compute the set of registers that can be considered as written by
+    // an authentication instruction. This includes operations that are
+    // *strictly better* than authentication, such as materializing a
+    // PC-relative constant.
+    SmallVector<MCPhysReg> AuthenticatedOrBetter =
+        getAuthenticatedRegs(Point, Cur);
+
+    // Then, compute the state after this instruction is executed.
+    State Next = Cur;
+
     Next.SafeToDerefRegs.reset(Clobbered);
     // Keep track of this instruction if it writes to any of the registers we
     // need to track that for:
     for (MCPhysReg Reg : RegsToTrackInstsFor.getRegisters())
       if (Clobbered[Reg])
         lastWritingInsts(Next, Reg) = {&Point};
 
-    ErrorOr<MCPhysReg> AutReg = BC.MIB->getAuthenticatedReg(Point);
-    if (AutReg && *AutReg != BC.MIB->getNoRegister()) {
-      // The sub-registers of *AutReg are also trusted now, but not its
-      // super-registers (as they retain untrusted register units).
-      BitVector AuthenticatedSubregs =
-          BC.MIB->getAliases(*AutReg, /*OnlySmaller=*/true);
-      for (MCPhysReg Reg : AuthenticatedSubregs.set_bits()) {
-        Next.SafeToDerefRegs.set(Reg);
-        if (RegsToTrackInstsFor.isTracked(Reg))
-          lastWritingInsts(Next, Reg).clear();
-      }
+    // After accounting for clobbered registers in general, override the state
+    // according to authentication and other *special cases* of clobbering.
+
+    // The sub-registers of each authenticated register are also trusted now,
+    // but not their super-registers (as they retain untrusted register units).
+    BitVector AuthenticatedSubregs(NumRegs);
+    for (MCPhysReg AutReg : AuthenticatedOrBetter)
+      AuthenticatedSubregs |= BC.MIB->getAliases(AutReg, /*OnlySmaller=*/true);
+    for (MCPhysReg Reg : AuthenticatedSubregs.set_bits()) {
+      Next.SafeToDerefRegs.set(Reg);
+      if (RegsToTrackInstsFor.isTracked(Reg))
+        lastWritingInsts(Next, Reg).clear();
     }
 
     LLVM_DEBUG({

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -319,6 +319,36 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
+  MCPhysReg getSafelyMaterializedAddressReg(const MCInst &Inst) const override {
+    switch (Inst.getOpcode()) {
+    case AArch64::ADR:
+    case AArch64::ADRP:
+      return Inst.getOperand(0).getReg();
+    default:
+      return getNoRegister();
+    }
+  }
+
+  std::pair<MCPhysReg, MCPhysReg>
+  analyzeSafeAddressArithmetics(const MCInst &Inst) const override {
+    switch (Inst.getOpcode()) {
+    default:
+      return std::make_pair(getNoRegister(), getNoRegister());
+    case AArch64::ADDXri:
+    case AArch64::SUBXri:
+      return std::make_pair(Inst.getOperand(0).getReg(),
+                            Inst.getOperand(1).getReg());
+    case AArch64::ORRXrs:
+      // "mov Xd, Xm" is equivalent to "orr Xd, XZR, Xm, lsl #0"
+      if (Inst.getOperand(1).getReg() != AArch64::XZR ||
+          Inst.getOperand(3).getImm() != 0)
+        return std::make_pair(getNoRegister(), getNoRegister());
+
+      return std::make_pair(Inst.getOperand(0).getReg(),
+                            Inst.getOperand(2).getReg());
+    }
+  }
+
   bool isADRP(const MCInst &Inst) const override {
     return Inst.getOpcode() == AArch64::ADRP;
   }

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

@@ -141,24 +141,9 @@ f_nonx30_ret_ok:
         stp     x29, x30, [sp, #-16]!
         mov     x29, sp
         bl      g
-        add     x0, x0, #3
         ldp     x29, x30, [sp], #16
-        // FIXME: Should the scanner understand that an authenticated register (below x30,
-        //        after the autiasp instruction), is OK to be moved to another register
-        //        and then that register being used to return?
-        //        This respects that pac-ret hardening intent, but the scanner currently
-        //        will produce a false positive for this.
-        //        Is it worthwhile to make the scanner more complex for this case?
-        //        So far, scanning many millions of instructions across a linux distro,
-        //        I haven't encountered such an example.
-        //        The ".if 0" block below tests this case and currently fails.
-.if 0
         autiasp
         mov     x16, x30
-.else
-        mov     x16, x30
-        autia   x16, sp
-.endif
 // CHECK-NOT: function f_nonx30_ret_ok
         ret     x16
         .size f_nonx30_ret_ok, .-f_nonx30_ret_ok