[AArch64] Avoid lowering setjmp call to CALL_BTI if harden-sls-blr is enabled

[pzheng]

Summary

Commit c3b9819 enabled inserting "bti j" after call to setjmp through lowering
setjmp calls to AArch64ISD::CALL_BTI which is later pattern matched to pseudo
instruction BLR_BTI. However, the lowering to BLR_BTI is infeasible if SLS BLR
mitigation (harden-sls-blr) is enabled because the pattern
Requires<[NoSLSBLRMitigation]>. Therefore, when harden-sls-blr is enabled,
ISel crashes due to the "can not select AArch64ISD::CALL_BTI" error. This patch
fixes this corner case by avoiding lowering setjmp call to CLL_BTI if
harden-sls-blr is enabled.

[DavidSpickett]

Not crashing seems reasonable :)

I don't know what the interaction of BTI and the sls-blr is though. Does not emitting the bti actually break things if you have sls-blr and branch target enforcement enabled? Perhaps the __llvm_slsblr_thunk_x would make that all work because it would contain the required btis.

I defer to @kristof.beyls on that one.

[DavidSpickett]

This change would lead to a situation where you have sls-blr and bti enabled but you don't get a bti after the setjmp call, which is going to be a hard failure once you return from the setjmp.

Talking to Kristof, we don't think there's any reason that CALL_BTI can't coexist with the sls mitigation. This could be done by following what AArch64call does https://github.com/llvm/llvm-project/blob/6e1ebb916e467f26d3c0cb0819770cd67f956cc3/llvm/lib/Target/AArch64/AArch64InstrInfo.td#L2567.

Since this was my change originally, I will work on a patch to do that.

[pzheng]

Thanks for the feedback, @DavidSpickett! Our longjmp implementation does not use br, so this patch is safe for us. Having said that, I do understand your concern over not generating a bti after setjmp when both sls-blr and branch target enforcement are enabled.

Please kindly keep me in the loop once you have a proper fix for this issue. I am abandoning this patch.

[DavidSpickett]

Posted https://reviews.llvm.org/D143915 instead.

llvm/lib/Target/AArch64/AArch64FastISel.cpp

bool AArch64FastISel::fastLowerCall(CallLoweringInfo &CLI) {
   MCSymbol *Symbol = CLI.Symbol;
 
   if (!Callee && !Symbol)
     return false;
 
   // Allow SelectionDAG isel to handle calls to functions like setjmp that need
   // a bti instruction following the call.
   if (CLI.CB && CLI.CB->hasFnAttr(Attribute::ReturnsTwice) &&
-      !Subtarget->noBTIAtReturnTwice() &&
+      !Subtarget->noBTIAtReturnTwice() && !Subtarget->hardenSlsBlr() &&
       MF->getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
     return false;
 
   // Allow SelectionDAG isel to handle indirect calls with KCFI checks.
   if (CLI.CB && CLI.CB->isIndirectCall() &&
       CLI.CB->getOperandBundle(LLVMContext::OB_kcfi))
     return false;
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

assert(!IsTailCall &&
            "tail calls cannot be marked with clang.arc.attachedcall");
     CallOpc = AArch64ISD::CALL_RVMARKER;
 
     // Add a target global address for the retainRV/claimRV runtime function
     // just before the call target.
     Function *ARCFn = *objcarc::getAttachedARCFunction(CLI.CB);
     auto GA = DAG.getTargetGlobalAddress(ARCFn, DL, PtrVT);
     Ops.insert(Ops.begin() + 1, GA);
-  } else if (GuardWithBTI)
+  } else if (GuardWithBTI && !Subtarget->hardenSlsBlr())
     CallOpc = AArch64ISD::CALL_BTI;
 
   // Returns a chain and a flag for retval copy to use.
   Chain = DAG.getNode(CallOpc, DL, NodeTys, Ops);
 
   if (IsCFICall)
     Chain.getNode()->setCFIType(CLI.CFIType->getZExtValue());
 

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

bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   // be expanded to the call, directly followed by a special marker sequence and
   // a call to an ObjC library function.
   if (Info.CB && objcarc::hasAttachedCallOpBundle(Info.CB))
     Opc = AArch64::BLR_RVMARKER;
   // A call to a returns twice function like setjmp must be followed by a bti
   // instruction.
   else if (Info.CB &&
            Info.CB->getAttributes().hasFnAttr(Attribute::ReturnsTwice) &&
-           !Subtarget.noBTIAtReturnTwice() &&
+           !Subtarget.noBTIAtReturnTwice() && !Subtarget.hardenSlsBlr() &&
            MF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
     Opc = AArch64::BLR_BTI;
   else
     Opc = getCallOpcode(MF, Info.Callee.isReg(), false);
 
   auto MIB = MIRBuilder.buildInstrNoInsert(Opc);
   unsigned CalleeOpNo = 0;
 

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=+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
+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+harden-sls-blr < %s | \
+; RUN: FileCheck %s --check-prefix=SLS
+; RUN: llc -mtriple=aarch64-none-linux-gnu -global-isel -mattr=+harden-sls-blr < %s | \
+; RUN: FileCheck %s --check-prefix=SLS
+; RUN: llc -mtriple=aarch64-none-linux-gnu -fast-isel -mattr=+harden-sls-blr < %s | \
+; RUN: FileCheck %s --check-prefix=SLS
 
 ; C source
 ; --------
 ; extern int setjmp(ptr);
 ; extern void notsetjmp(void);
 ;
 ; void bbb(void) {
 ;   setjmp(0);

 
 ; 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
+
+; SLS-LABEL: bbb:
+; SLS:     bl setjmp
+; SLS-NOT: hint #36
+; SLS:     bl __llvm_slsblr_thunk_x{{[0-9]+}}
+; SLS-NOT: hint #36
+; SLS:     bl notsetjmp
+; SLS-NOT: hint #36
 entry:
   %fnptr = alloca ptr, align 8
   %call = call i32 @setjmp(ptr noundef null) #0
   store ptr @setjmp, ptr %fnptr, align 8
   %0 = load ptr, ptr %fnptr, align 8
   %call1 = call i32 %0(ptr noundef null) #0
   call void @notsetjmp()
   ret void