[AArch64] Avoid incompatibility between SLSBLR mitigation and BTI codegen.

[kristof.beyls]

Summary

A "BTI c" instruction only allows jumping/calling to using a BLR* instruction.
However, the SLSBLR mitigation changes a BLR to a BR to implement the
function call. Therefore, a "BTI c" check that passed before could
trigger after the BLR->BL change done by the SLSBLR mitigation.
However, if the register used in BR is X16 or X17, this trigger will not
fire (see ArmARM for further details).

Therefore, this patch simply changes the function stubs for the SLSBLR
mitigation from
__llvm_slsblr_thunk_x<N>:

br x<N>
SpeculationBarrier

to
__llvm_slsblr_thunk_x<N>:

mov x16, x<N>
br  x16
SpeculationBarrier

[ostannard]

I have the same concern as in D81402: will this still work if the call to the thunk goes through a long-branch veneer or PLT, which could clobber x16 and x17?

llvm/lib/Target/AArch64/AArch64FastISel.cpp

3274 ↗

(On Diff #269248)

Do we need to conditionally emit BLRCallX16X17 here too?

3307 ↗

(On Diff #269248)

Again, might we need to emit BLRCallX16X17 here?

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1129 ↗

(On Diff #269248)

Why is this hard-coded to BLRCallX16X17, never BLRCall?

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6562 ↗

(On Diff #269248)

Would it be possible for a BLRCallX16X17 to reach here?

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

779 ↗

(On Diff #269248)

Same as fast-sel, do we need to emit BLRCallX16X17 here?

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

2874 ↗

(On Diff #269248)

Same as fast-sel, do we need to emit BLRCallX16X17 here?

[kristof.beyls]

Based on the review feedback that using X16 or X17 could be problematic since a linker can clobber those registers on a call, this patch goes for a different approach to achieve compatibility with BTI-enabled callees.
It simply moves the register to be called always to be in X16, and then does BR X16.
It seems it's simplest to do that unconditionally, as it may not always be easy to find out if the call target has been BTI-enabled.

[ostannard]

LGTM.

I think we could make this conditional based on whether the caller has BTI enabled, because:

• If the caller has BTI enabled, then we correctly use X16 for a BTI callee, or wastefully (but still with correct behaviour) use X16 for a non-BTI callee.
• If the caller has BTI disabled, then it must be in a page with BTI disabled, and so is able to use BR with any register, even if the destination is BTI-protected.

This relies on the caller and thunk being allocated in memory with the same BTI state, which I expect to be a safe assumption.

This patch still looks good by itself, the above would just be an optimisation and can be done later.

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
   MachineBasicBlock *Entry = &MF.front();
   Entry->clear();
 
   //  These thunks need to consist of the following instructions:
   //  __llvm_slsblr_thunk_xN:
   //      BR xN
   //      barrierInsts
   Entry->addLiveIn(ThunkReg);
-  BuildMI(Entry, DebugLoc(), TII->get(AArch64::BR)).addReg(ThunkReg);
+  // MOV X16, ThunkReg == ORR X16, XZR, ThunkReg, LSL #0
+  BuildMI(Entry, DebugLoc(), TII->get(AArch64::ORRXrs), AArch64::X16)
+      .addReg(AArch64::XZR)
+      .addReg(ThunkReg)
+      .addImm(0);
+  BuildMI(Entry, DebugLoc(), TII->get(AArch64::BR)).addReg(AArch64::X16);
   // Make sure the thunks do not make use of the SB extension in case there is
   // a function somewhere that will call to it that for some reason disabled
   // the SB extension locally on that function, even though it's enabled for
   // the module otherwise. Therefore set AlwaysUseISBSDB to true.
   insertSpeculationBarrier(&MF.getSubtarget<AArch64Subtarget>(), *Entry,
                            Entry->end(), DebugLoc(), true /*AlwaysUseISBDSB*/);
 }
 

llvm/test/CodeGen/AArch64/speculation-hardening-sls.ll

; CHECK
   %call2 = call i64 %6(i8* %2, i8* %8, i64 0, i64 0, i64 0, i64 0, i64 0, i64 0)
   %add = add nsw i64 %call1, %call
   %add1 = add nsw i64 %call2, %add
   ret i64 %add1
 ; CHECK: .Lfunc_end
 }
 
 ; HARDEN-label: __llvm_slsblr_thunk_x0:
-; HARDEN:    br x0
+; HARDEN:    mov x16, x0
+; HARDEN:    br x16
 ; ISBDSB-NEXT: dsb sy
 ; ISBDSB-NEXT: isb
 ; SB-NEXT:     dsb sy
 ; SB-NEXT:     isb
 ; HARDEN-NEXT: .Lfunc_end
 ; HARDEN-label: __llvm_slsblr_thunk_x19:
-; HARDEN:    br x19
+; HARDEN:    mov x16, x19
+; HARDEN:    br x16
 ; ISBDSB-NEXT: dsb sy
 ; ISBDSB-NEXT: isb
 ; SB-NEXT:     dsb sy
 ; SB-NEXT:     isb
 ; HARDEN-NEXT: .Lfunc_end