[AArch64][v8.5A] Restrict indirect tail calls to use x16/17 only when using BTI

[olista01]

Summary

When branch target identification is enabled, all indirectly-callable
functions start with a BTI C instruction. this instruction can only be
the target of certain indirect branches (direct branches and
fall-through are not affected):

• A BLR instruction, in either a protected or unprotected page.
• A BR instruction in a protected page, using x16 or x17.
• A BR instruction in an unprotected page, using any register.

Without BTI, we can use any non call-preserved register to hold the
address for an indirect tail call. However, when BTI is enabled, then
the code being compiled might be loaded into a BTI-protected page, where
only x16 and x17 can be used for indirect tail calls.

Legacy code withiout this restriction can still indirectly tail-call
BTI-protected functions, because they will be loaded into an unprotected
page, so any register is allowed.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.td

6641

I don't think the Requires clauses here actually affect anything: there's no assembly syntax (or indeed MC-level feature) and the patterns are separate.

[olista01]

Remove unnecessary Requires on pseudo-instructions.

[t.p.northover]

Thanks. This looks fine now.

[efriedma]

Would it also make sense to enforce the corresponding restriction, that non-tail-call indirect branches shouldn't use x16/x17?

lib/Target/AArch64/AArch64AsmPrinter.cpp

case AArch64::DBG_VALUE: {
     }
     return;
   }
 
   // Tail calls use pseudo instructions so they have the proper code-gen
   // attributes (isCall, isReturn, etc.). We lower them to the real
   // instruction here.
   case AArch64::TCRETURNri:
+  case AArch64::TCRETURNriBTI:
   case AArch64::TCRETURNriALL: {
     MCInst TmpInst;
     TmpInst.setOpcode(AArch64::BR);
     TmpInst.addOperand(MCOperand::createReg(MI->getOperand(0).getReg()));
     EmitToStreamer(*OutStreamer, TmpInst);
     return;
   }
   case AArch64::TCRETURNdi: {

lib/Target/AArch64/AArch64FrameLowering.cpp

void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
   DebugLoc DL;
   bool IsTailCallReturn = false;
   if (MBB.end() != MBBI) {
     DL = MBBI->getDebugLoc();
     unsigned RetOpcode = MBBI->getOpcode();
     IsTailCallReturn = RetOpcode == AArch64::TCRETURNdi ||
-      RetOpcode == AArch64::TCRETURNri;
+                       RetOpcode == AArch64::TCRETURNri ||
+                       RetOpcode == AArch64::TCRETURNriBTI;
   }
   int NumBytes = MFI.getStackSize();
   const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
 
   // All calls are tail calls in GHC calling conv, and functions have no
   // prologue/epilogue.
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;

lib/Target/AArch64/AArch64InstrInfo.td

 // We could compute these on a per-module basis but doing so requires accessing
 // the Function object through the <Target>Subtarget and objections were raised
 // to that (see post-commit review comments for r301750).
 let RecomputePerFunction = 1 in {
   def ForCodeSize   : Predicate<"MF->getFunction().optForSize()">;
   def NotForCodeSize   : Predicate<"!MF->getFunction().optForSize()">;
   // Avoid generating STRQro if it is slow, unless we're optimizing for code size.
   def UseSTRQro : Predicate<"!Subtarget->isSTRQroSlow() || MF->getFunction().optForSize()">;
+
+  def UseBTI : Predicate<[{ MF->getFunction().hasFnAttribute("branch-target-enforcement") }]>;
+  def NotUseBTI : Predicate<[{ !MF->getFunction().hasFnAttribute("branch-target-enforcement") }]>;
 }
 
 include "AArch64InstrFormats.td"
 include "SVEInstrFormats.td"
 
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

 } // Predicates = [IsLE]
 
 // Tail call return handling. These are all compiler pseudo-instructions,
 // so no encoding information or anything like that.
 let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Uses = [SP] in {
   def TCRETURNdi : Pseudo<(outs), (ins i64imm:$dst, i32imm:$FPDiff), []>,
                    Sched<[WriteBrReg]>;
   def TCRETURNri : Pseudo<(outs), (ins tcGPR64:$dst, i32imm:$FPDiff), []>,
-                   Sched<[WriteBrReg]>;
+                   Sched<[WriteBrReg]>,
+                   Requires<[NotUseBTI]>;
   // Indirect tail-call with any register allowed, used by MachineOutliner when
   // this is proven safe.
   def TCRETURNriALL : Pseudo<(outs), (ins GPR64:$dst, i32imm:$FPDiff), []>,
                       Sched<[WriteBrReg]>;
+  // Indirect tail-call limited to only use registers (x16 and x17) which are
+  // allowed to tail-call a "BTI c" instruction.
+  def TCRETURNriBTI : Pseudo<(outs), (ins rtcGPR64:$dst, i32imm:$FPDiff), []>,
+                      Sched<[WriteBrReg]>,
+                      Requires<[UseBTI]>;
 }
 
 def : Pat<(AArch64tcret tcGPR64:$dst, (i32 timm:$FPDiff)),
-          (TCRETURNri tcGPR64:$dst, imm:$FPDiff)>;
+          (TCRETURNri tcGPR64:$dst, imm:$FPDiff)>,
+      Requires<[NotUseBTI]>;
+def : Pat<(AArch64tcret rtcGPR64:$dst, (i32 timm:$FPDiff)),
+          (TCRETURNriBTI rtcGPR64:$dst, imm:$FPDiff)>,
+      Requires<[UseBTI]>;
 def : Pat<(AArch64tcret tglobaladdr:$dst, (i32 timm:$FPDiff)),
           (TCRETURNdi texternalsym:$dst, imm:$FPDiff)>;
 def : Pat<(AArch64tcret texternalsym:$dst, (i32 timm:$FPDiff)),
           (TCRETURNdi texternalsym:$dst, imm:$FPDiff)>;
 
 include "AArch64InstrAtomics.td"
 include "AArch64SVEInstrInfo.td"

lib/Target/AArch64/AArch64RegisterInfo.td

 
 // For tail calls, we can't use callee-saved registers, as they are restored
 // to the saved value before the tail call, which would clobber a call address.
 // This is for indirect tail calls to store the address of the destination.
 def tcGPR64 : RegisterClass<"AArch64", [i64], 64, (sub GPR64common, X19, X20, X21,
                                                      X22, X23, X24, X25, X26,
                                                      X27, X28, FP, LR)>;
 
+// Restricted set of tail call registers, for use when branch target
+// enforcement is enabled. These are the only registers which can be used to
+// indirectly branch (not call) to the "BTI c" instruction at the start of a
+// BTI-protected function.
+def rtcGPR64 : RegisterClass<"AArch64", [i64], 64, (add X16, X17)>;
+
 // GPR register classes for post increment amount of vector load/store that
 // has alternate printing when Rm=31 and prints a constant immediate value
 // equal to the total number of bytes transferred.
 
 // FIXME: TableGen *should* be able to do these itself now. There appears to be
 // a bug in counting how many operands a Post-indexed MCInst should have which
 // means the aliases don't trigger.
 def GPR64pi1  : RegisterOperand<GPR64, "printPostIncOperand<1>">;

test/CodeGen/AArch64/branch-target-enforcement-indirect-calls.ll

+; RUN: llc -mtriple aarch64--none-eabi -mattr=+bti < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-arm-none-eabi"
+
+; When BTI is enabled, all indirect tail-calls must use x16 or x17 (the intra
+; procedure call scratch registers) to hold the address, as these instructions
+; are allowed to target the "BTI c" instruction at the start of the target
+; function. The alternative to this would be to start functions with "BTI jc",
+; which increases the number of potential ways they could be called, and
+; weakens the security protections.
+
+define void @bti_disabled(void ()* %p) {
+entry:
+  tail call void %p()
+; CHECK: br x0
+  ret void
+}
+
+define void @bti_enabled(void ()* %p) "branch-target-enforcement" {
+entry:
+  tail call void %p()
+; CHECK: br {{x16|x17}}
+  ret void
+}