[llvm][AArch64] Fix an interaction of SLS and BTI after a returns twice call

[DavidSpickett]

Summary

This fixes the combination of two things:

• Placing a BTI after calls to a returns twice function like setjmp. This allows the setjmp to return with a br instead of a ret.
• Straight line speculation mitigations that replace BLR with a BL to a thunk that does the mitigation, and then goes to the original target.

Originally I marked AArch64call_bti as requiring that SLS mitigation
be disabled. This caused a crash when you tried to codegen with both.
Since CALL_BTI tried to match with AArch64call_bti but could not.

This change does 2 things:

• Follow the pattern set by AArch64call and add 2 patterns for AArch64call_bti. One with no IP (interprocedural) registers, and one with. For SLS mitigation on and off respectively.
• Modify the sls hardening pass to iterate through bundled instructions, as the AArch64 KCFI pass does.

Since there is a 1:1 replacement of the BLR with a BL,
the bundle remains intact. This is checked with an MIR test.

The ir -> asm testing is updated to add runs with the sls
mitigation enabled.

[DavidSpickett]

I should note that the IR that describes an indirect call to a returns twice function will not be generated by current clang. It seems like the SROA pass drops that attribute when it replaces the value. Perhaps it did at one point do that, as I don't think I hand wrote the IR. I will see if that is a bug or not.

That doesn't prevent this patch going in though.

I asked my colleagues about the merits of pseudos that are expanded late, vs. having passes look inside bundles. There were a lot of opinions mostly toward bundles being a bad way to do this. However, it's what we have and the KCFI pass already handles something like this so I went with it. Let's teach 1 pass to handle a bundle rather than 2 how to handle a pseudo.

@pzheng Please check whatever build configuration you were hitting the crash in. I think I'm testing all of SelectionDag/FastIsel/GlobalIsel but not 100% sure.

llvm/test/CodeGen/AArch64/setjmp-bti.ll

4

And now I see I forgot a SelectionDAG test for this one.

[DavidSpickett]

llvm/test/CodeGen/AArch64/speculation-hardening-sls-blr-bti.mir

74

Note that the sls thunk is transparent such that when the setjmp returns it will hit this bti, not to somewhere in the thunk.

[DavidSpickett]

Add SelectionDAG test for bti+sls.

[DavidSpickett]

llvm/test/CodeGen/AArch64/setjmp-bti.ll

4

Added.

[pzheng]

Thanks for the fix, @DavidSpickett! I confirm that this patch fixes the crash we encountered.

[DavidSpickett]

I should note that the IR that describes an indirect call to a returns twice function will not be generated by current clang. It seems like the SROA pass drops that attribute when it replaces the value. Perhaps it did at one point do that, as I don't think I hand wrote the IR. I will see if that is a bug or not.

Filed https://github.com/llvm/llvm-project/issues/60732. It is a generic issue with SROA I think.

[kristof.beyls]

LGTM, thank you!

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 
 def : Pat<(AArch64call_rvmarker (i64 tglobaladdr:$rvfunc), GPR64:$Rn),
           (BLR_RVMARKER tglobaladdr:$rvfunc, GPR64:$Rn)>,
       Requires<[NoSLSBLRMitigation]>;
 
 def : Pat<(AArch64call_bti GPR64:$Rn),
           (BLR_BTI GPR64:$Rn)>,
       Requires<[NoSLSBLRMitigation]>;
+def : Pat<(AArch64call_bti GPR64noip:$Rn),
+          (BLR_BTI GPR64noip:$Rn)>,
+      Requires<[SLSBLRMitigation]>;
 
 let isBranch = 1, isTerminator = 1, isBarrier = 1, isIndirectBranch = 1 in {
 def BR  : BranchReg<0b0000, "br", [(brind GPR64:$Rn)]>;
 } // isBranch, isTerminator, isBarrier, isIndirectBranch
 
 // Create a separate pseudo-instruction for codegen to use so that we don't
 // flag lr as used in every function. It'll be restored before the RET by the
 // epilogue if it's legitimately used.

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

public:
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
   StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
 
 private:
   bool hardenReturnsAndBRs(MachineBasicBlock &MBB) const;
   bool hardenBLRs(MachineBasicBlock &MBB) const;
   MachineBasicBlock &ConvertBLRToBL(MachineBasicBlock &MBB,
-                                    MachineBasicBlock::iterator) const;
+                                    MachineBasicBlock::instr_iterator) const;
 };
 
 } // end anonymous namespace
 
 char AArch64SLSHardening::ID = 0;
 
 INITIALIZE_PASS(AArch64SLSHardening, "aarch64-sls-hardening",
                 AARCH64_SLS_HARDENING_NAME, false, false)
void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
   // 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*/);
 }
 
-MachineBasicBlock &
-AArch64SLSHardening::ConvertBLRToBL(MachineBasicBlock &MBB,
-                                    MachineBasicBlock::iterator MBBI) const {
+MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
+    MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator MBBI) const {
   // Transform a BLR to a BL as follows:
   // Before:
   //   |-----------------------------|
   //   |      ...                    |
   //   |  instI                      |
   //   |  BLR xN                     |
   //   |  instJ                      |
   //   |      ...                    |
MachineBasicBlock &AArch64SLSHardening::ConvertBLRToBL(
 
   return MBB;
 }
 
 bool AArch64SLSHardening::hardenBLRs(MachineBasicBlock &MBB) const {
   if (!ST->hardenSlsBlr())
     return false;
   bool Modified = false;
-  MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
-  MachineBasicBlock::iterator NextMBBI;
+  MachineBasicBlock::instr_iterator MBBI = MBB.instr_begin(),
+                                    E = MBB.instr_end();
+  MachineBasicBlock::instr_iterator NextMBBI;
   for (; MBBI != E; MBBI = NextMBBI) {
     MachineInstr &MI = *MBBI;
     NextMBBI = std::next(MBBI);
     if (isBLR(MI)) {
       ConvertBLRToBL(MBB, MBBI);
       Modified = true;
     }
   }

llvm/test/CodeGen/AArch64/setjmp-bti.ll

 ; RUN: llc -mtriple=aarch64-none-linux-gnu < %s | FileCheck %s --check-prefix=BTI
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -global-isel < %s | FileCheck %s --check-prefix=BTI
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -fast-isel < %s | FileCheck %s --check-prefix=BTI
+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+harden-sls-blr< %s | FileCheck %s --check-prefix=BTISLS
+; RUN: llc -mtriple=aarch64-none-linux-gnu -global-isel -mattr=+harden-sls-blr< %s | FileCheck %s --check-prefix=BTISLS
+; RUN: llc -mtriple=aarch64-none-linux-gnu -fast-isel   -mattr=+harden-sls-blr< %s | FileCheck %s --check-prefix=BTISLS
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+no-bti-at-return-twice < %s | \
 ; RUN: FileCheck %s --check-prefix=NOBTI
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -global-isel -mattr=+no-bti-at-return-twice < %s | \
 ; RUN: FileCheck %s --check-prefix=NOBTI
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -fast-isel -mattr=+no-bti-at-return-twice < %s | \
 ; RUN: FileCheck %s --check-prefix=NOBTI
 
 ; C source

 ; BTI-LABEL: bbb:
 ; BTI:       bl setjmp
 ; BTI-NEXT:  hint #36
 ; BTI:       blr x{{[0-9]+}}
 ; BTI-NEXT:  hint #36
 ; BTI:       bl notsetjmp
 ; BTI-NOT:   hint #36
 
+; BTISLS-LABEL: bbb:
+; BTISLS:       bl setjmp
+; BTISLS-NEXT:  hint #36
+; BTISLS:       bl __llvm_slsblr_thunk_x{{[0-9]+}}
+; BTISLS-NEXT:  hint #36
+; BTISLS:       bl notsetjmp
+; BTISLS-NOT:   hint #36
+
 ; NOBTI-LABEL: bbb:
 ; NOBTI:     bl setjmp
 ; NOBTI-NOT: hint #36
 ; NOBTI:     blr x{{[0-9]+}}
 ; NOBTI-NOT: hint #36
 ; NOBTI:     bl notsetjmp
 ; NOBTI-NOT: hint #36
 entry:

llvm/test/CodeGen/AArch64/speculation-hardening-sls-blr-bti.mir

+# RUN: llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu \
+# RUN:     -start-before aarch64-sls-hardening \
+# RUN:     -stop-after aarch64-sls-hardening -o - %s \
+# RUN:   | FileCheck %s --check-prefixes=CHECK
+
+# Check when the BLR SLS hardening encounters a BLR/BTI bundle, the BTI
+# instruction remains after the BLR is replaced with a BL.
+# These BLR/BTI bundles are produced when calling a returns_twice function
+# (like setjmp) indirectly.
+--- |
+  $__llvm_slsblr_thunk_x8 = comdat any
+
+  define dso_local void @fn() #0 {
+  entry:
+    %fnptr = alloca ptr, align 8
+    store ptr @setjmp, ptr %fnptr, align 8
+    %0 = load ptr, ptr %fnptr, align 8
+    %call1 = call i32 %0(ptr noundef null) #1
+    ret void
+  }
+
+  ; Function Attrs: returns_twice
+  declare i32 @setjmp(ptr noundef) #1
+
+  ; Function Attrs: naked nounwind
+  define linkonce_odr hidden void @__llvm_slsblr_thunk_x8() #2 comdat {
+  entry:
+    ret void
+  }
+
+  attributes #0 = { "target-features"="+harden-sls-blr" }
+  attributes #1 = { returns_twice }
+  attributes #2 = { naked nounwind }
+
+  !llvm.module.flags = !{!0}
+  !0 = !{i32 8, !"branch-target-enforcement", i32 1}
+...
+---
+name:            fn
+exposesReturnsTwice: true
+tracksRegLiveness: true
+fixedStack:      []
+stack:
+  - { id: 0, name: fnptr, type: default, offset: -8, size: 8, alignment: 8,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -8, debug-info-variable: '', debug-info-expression: '',
+      debug-info-location: '' }
+  - { id: 1, name: '', type: spill-slot, offset: -16, size: 8, alignment: 16,
+      stack-id: default, callee-saved-register: '$lr', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0.entry:
+    liveins: $lr
+
+    early-clobber $sp = frame-setup STRXpre killed $lr, $sp, -16 :: (store (s64) into %stack.1)
+    frame-setup CFI_INSTRUCTION def_cfa_offset 16
+    frame-setup CFI_INSTRUCTION offset $w30, -16
+    $x8 = ADRP target-flags(aarch64-page, aarch64-got) @setjmp
+    renamable $x8 = LDRXui killed $x8, target-flags(aarch64-pageoff, aarch64-got, aarch64-nc) @setjmp
+    STRXui renamable $x8, $sp, 1 :: (store (s64) into %ir.fnptr)
+    $x0 = ORRXrs $xzr, $xzr, 0
+    BUNDLE implicit-def $lr, implicit-def $w30, implicit killed $x8, implicit $sp {
+      BLR killed renamable $x8, implicit-def $lr, implicit $sp
+      HINT 36
+    }
+    ; CHECK:      BUNDLE implicit-def $lr, implicit-def $w30, implicit killed $x8, implicit $sp {
+    ; CHECK-NEXT:   BL <mcsymbol __llvm_slsblr_thunk_x8>, implicit-def $lr, implicit $sp, implicit killed $x8
+    ; CHECK-NEXT:   HINT 36
+    ; CHECK-NEXT: }
+    early-clobber $sp, $lr = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.1)
+    RET undef $lr
+...
+---
+name:            __llvm_slsblr_thunk_x8
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $x8
+
+    $x16 = ORRXrs $xzr, $x8, 0
+    BR $x16
+    SpeculationBarrierISBDSBEndBB
+...