[ARM] Implement setjmp BTI placement for PACBTI-M

[stuij]

Summary

This patch intends to guard indirect branches performed by longjmp
by inserting BTI instructions after calls to setjmp.

Calls with 'returns-twice' are lowered to a new pseudo-instruction
named t2CALL_BTI that is later expanded to a bundle of {tBL,t2BTI}.

This patch is part of a series that adds support for the PACBTI-M extension of
the Armv8.1-M architecture, as detailed here:

https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/armv8-1-m-pointer-authentication-and-branch-target-identification-extension

The PACBTI-M specification can be found in the Armv8-M Architecture Reference
Manual:

https://developer.arm.com/documentation/ddi0553/latest

The following people contributed to this patch:

• Alexandros Lamprineas
• Ties Stuij

[chill]

I would suggest to add a command line option to control placement of BTI after setjmp/etc.
The default should be as it is now, but IMHO we should provide means for people who *know* how their
runtime is implemented[1] to avoid placing the extra BTI instructions.

[1] e.g. longjmp uses PAC, in which case destination address is already authenticated, or
longjmp uses a non-BTI setting instruction to jump, in which case the BTI does nothing
other than take space and provide opportunities to jump from elsewhere.

[labrinea]

We only perform this transformation when armlib is not in use.
That information is passed on by the module flags.

This was the downstream strategy. Like Momchill pointed out, it'd be preferable to provide a way for the user to control the codegen for 'returns-twice' calls.

[stuij]

addressed review comments, namely adding a commandline argument, -mno-bti-at-return-twice, to not place a bti instruction after a setjmp call.

[stuij]

slight rewording

[DavidSpickett]

llvm/lib/Target/ARM/ARMInstrThumb2.td

5745

Should this require IsMClass instead/also? Though I wasn't able to get anything weird to happen when using an A profile triple so maybe I'm missing a check elsewhere that means you'd never get to this point with A profile Arm.

For example this A profile triple:

$ ./bin/clang --target=thumbv8-arm-none-eabi /tmp/test.c -o /tmp/test.o -o - -S -mbranch-protection=bti -mthumb

Doesn't put anything after a call to setjmp, nop or otherwise, but I can't place where that decision is made.

[chill]

llvm/lib/Target/ARM/ARMInstrThumb2.td

5745

The decision is made in ARMMachineFunctionInfo

https://github.com/llvm/llvm-project/blob/a02af37560ff5aa22dcef5735ef25eaf58eaaf64/llvm/lib/Target/ARM/ARMMachineFunctionInfo.cpp#L18

[DavidSpickett]

llvm/lib/Target/ARM/ARMInstrThumb2.td

5745

Never mind, I figured it out as per usual just after leaving the comment.

static bool GetBranchTargetEnforcement(MachineFunction &MF) {
  const auto &Subtarget = MF.getSubtarget<ARMSubtarget>();
  if (!Subtarget.isMClass() || !Subtarget.hasV7Ops())
    return false;

This returns false for the A profile which means that GuardWithBTI is false so we don't add a BTI. Maybe one could craft some IR example that got around that but doesn't seem a likely issue.

[DavidSpickett]

llvm/lib/Target/ARM/ARMInstrThumb2.td

5745

The decision is made in ARMMachineFunctionInfo

Left my comment just as you did. Thanks! I see how it works now.

clang/docs/ClangCommandLineReference.rst

 .. option:: -mtp=<arg>
 
 Thread pointer access method (AArch32/AArch64 only)
 
 .. option:: -munaligned-access, -mno-unaligned-access
 
 Allow memory accesses to be unaligned (AArch32/AArch64 only)
 
+.. option:: -mno-bti-at-return-twice
+
+Do not add a BTI instruction after a setjmp or other return-twice construct (Arm
+only)
+
 Hexagon
 -------
 .. option:: -mieee-rnd-near
 
 .. option:: -mmemops, -mno-memops
 
 Enable generation of memop instructions
 

clang/include/clang/Driver/Options.td

def mfix_cortex_a53_835769 : Flag<["-"], "mfix-cortex-a53-835769">,
   Group<m_aarch64_Features_Group>,
   HelpText<"Workaround Cortex-A53 erratum 835769 (AArch64 only)">;
 def mno_fix_cortex_a53_835769 : Flag<["-"], "mno-fix-cortex-a53-835769">,
   Group<m_aarch64_Features_Group>,
   HelpText<"Don't workaround Cortex-A53 erratum 835769 (AArch64 only)">;
 def mmark_bti_property : Flag<["-"], "mmark-bti-property">,
   Group<m_aarch64_Features_Group>,
   HelpText<"Add .note.gnu.property with BTI to assembly files (AArch64 only)">;
+def mno_bti_at_return_twice : Flag<["-"], "mno-bti-at-return-twice">,
+  Group<m_arm_Features_Group>,
+  HelpText<"Do not add a BTI instruction after a setjmp or other"
+           " return-twice construct (Arm only)">;
+
 foreach i = {1-31} in
   def ffixed_x#i : Flag<["-"], "ffixed-x"#i>, Group<m_Group>,
     HelpText<"Reserve the x"#i#" register (AArch64/RISC-V only)">;
 
 foreach i = {8-15,18} in
   def fcall_saved_x#i : Flag<["-"], "fcall-saved-x"#i>, Group<m_aarch64_Features_Group>,
     HelpText<"Make the x"#i#" register call-saved (AArch64 only)">;
 

clang/lib/Driver/ToolChains/Arch/ARM.cpp

if (EnableRetBr)
       Features.push_back("+harden-sls-retbr");
     if (EnableBlr)
       Features.push_back("+harden-sls-blr");
     if (DisableComdat) {
       Features.push_back("+harden-sls-nocomdat");
     }
   }
 
+  if (Args.getLastArg(options::OPT_mno_bti_at_return_twice))
+    Features.push_back("+no-bti-at-return-twice");
 }
 
 std::string arm::getARMArch(StringRef Arch, const llvm::Triple &Triple) {
   std::string MArch;
   if (!Arch.empty())
     MArch = std::string(Arch);
   else
     MArch = std::string(Triple.getArchName());

clang/test/Driver/arm-bti-return-twice.c

+// RUN: %clang -target arm-arm-none-eabi -march=armv8-m.main -mbranch-protection=bti \
+// RUN: -mno-bti-at-return-twice -### %s 2>&1 | FileCheck %s --check-prefix=FEAT
+// RUN: %clang -target arm-arm-none-eabi -march=armv8-m.main -mbranch-protection=bti \
+// RUN: -### %s 2>&1 | FileCheck %s --check-prefix=NOFEAT
+
+// FEAT: "+no-bti-at-return-twice"
+// NOFEAT-NOT: "+no-bti-at-return-twice"

llvm/lib/Target/ARM/ARM.td

def FeatureFixCMSE_CVE_2021_35465 : SubtargetFeature<"fix-cmse-cve-2021-35465",
                                         "FixCMSE_CVE_2021_35465", "true",
                                         "Mitigate against the cve-2021-35465 "
                                         "security vulnurability">;
 
 def FeaturePACBTI         : SubtargetFeature<"pacbti", "HasPACBTI", "true",
                                              "Enable Pointer Authentication and Branch "
                                              "Target Identification">;
 
+def FeatureNoBTIAtReturnTwice : SubtargetFeature<"no-bti-at-return-twice",
+                                                 "NoBTIAtReturnTwice", "true",
+                                                 "Don't place a BTI instruction "
+                                                 "after a return-twice">;
+
 //===----------------------------------------------------------------------===//
 // ARM architecture class
 //
 
 // A-series ISA
 def FeatureAClass : SubtargetFeature<"aclass", "ARMProcClass", "AClass",
                                      "Is application profile ('A' series)">;
 

llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp

case ARM::BL_PUSHLR: {
         MIB = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(ARM::BL));
       }
       MIB.cloneMemRefs(MI);
       for (const MachineOperand &MO : llvm::drop_begin(MI.operands()))
         MIB.add(MO);
       MI.eraseFromParent();
       return true;
     }
+    case ARM::t2CALL_BTI: {
+      MachineFunction &MF = *MI.getMF();
+      MachineInstrBuilder MIB =
+          BuildMI(MF, MI.getDebugLoc(), TII->get(ARM::tBL));
+      MIB.cloneMemRefs(MI);
+      for (unsigned i = 0; i < MI.getNumOperands(); ++i)
+        MIB.add(MI.getOperand(i));
+      if (MI.isCandidateForCallSiteEntry())
+        MF.moveCallSiteInfo(&MI, MIB.getInstr());
+      MIBundleBuilder Bundler(MBB, MI);
+      Bundler.append(MIB);
+      Bundler.append(BuildMI(MF, MI.getDebugLoc(), TII->get(ARM::t2BTI)));
+      finalizeBundle(MBB, Bundler.begin(), Bundler.end());
+      MI.eraseFromParent();
+      return true;
+    }
     case ARM::LOADDUAL:
     case ARM::STOREDUAL: {
       Register PairReg = MI.getOperand(0).getReg();
 
       MachineInstrBuilder MIB =
           BuildMI(MBB, MBBI, MI.getDebugLoc(),
                   TII->get(Opcode == ARM::LOADDUAL ? ARM::LDRD : ARM::STRD))
               .addReg(TRI->getSubReg(PairReg, ARM::gsub_0),

llvm/lib/Target/ARM/ARMISelLowering.h

enum NodeType : unsigned {
 
     // Add pseudo op to model memcpy for struct byval.
     COPY_STRUCT_BYVAL,
 
     CALL,        // Function call.
     CALL_PRED,   // Function call that's predicable.
     CALL_NOLINK, // Function call with branch not branch-and-link.
     tSECALL,     // CMSE non-secure function call.
+    t2CALL_BTI,  // Thumb function call followed by BTI instruction.
     BRCOND,      // Conditional branch.
     BR_JT,       // Jumptable branch.
     BR2_JT,      // Jumptable branch (2 level - jumptable entry is a jump).
     RET_FLAG,    // Return with a flag operand.
     SERET_FLAG,  // CMSE Entry function return with a flag operand.
     INTRET_FLAG, // Interrupt return with an LR-offset and a flag operand.
 
     PIC_ADD, // Add with a PC operand and a PIC label.

llvm/lib/Target/ARM/ARMISelLowering.cpp

case ARMISD::FIRST_NUMBER:
     MAKE_CASE(ARMISD::Wrapper)
     MAKE_CASE(ARMISD::WrapperPIC)
     MAKE_CASE(ARMISD::WrapperJT)
     MAKE_CASE(ARMISD::COPY_STRUCT_BYVAL)
     MAKE_CASE(ARMISD::CALL)
     MAKE_CASE(ARMISD::CALL_PRED)
     MAKE_CASE(ARMISD::CALL_NOLINK)
     MAKE_CASE(ARMISD::tSECALL)
+    MAKE_CASE(ARMISD::t2CALL_BTI)
     MAKE_CASE(ARMISD::BRCOND)
     MAKE_CASE(ARMISD::BR_JT)
     MAKE_CASE(ARMISD::BR2_JT)
     MAKE_CASE(ARMISD::RET_FLAG)
     MAKE_CASE(ARMISD::SERET_FLAG)
     MAKE_CASE(ARMISD::INTRET_FLAG)
     MAKE_CASE(ARMISD::PIC_ADD)
     MAKE_CASE(ARMISD::CMP)
ARMTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   MachineFunction &MF = DAG.getMachineFunction();
   ARMFunctionInfo *AFI = MF.getInfo<ARMFunctionInfo>();
   MachineFunction::CallSiteInfo CSInfo;
   bool isStructRet = (Outs.empty()) ? false : Outs[0].Flags.isSRet();
   bool isThisReturn = false;
   bool isCmseNSCall   = false;
   bool isSibCall = false;
   bool PreferIndirect = false;
+  bool GuardWithBTI = false;
+
+  // Lower 'returns_twice' calls to a pseudo-instruction.
+  if (CLI.CB && CLI.CB->getAttributes().hasFnAttr(Attribute::ReturnsTwice) &&
+      !Subtarget->getNoBTIAtReturnTwice())
+    GuardWithBTI = AFI->branchTargetEnforcement();
 
   // Determine whether this is a non-secure function call.
   if (CLI.CB && CLI.CB->getAttributes().hasFnAttr("cmse_nonsecure_call"))
     isCmseNSCall = true;
 
   // Disable tail calls if they're not supported.
   if (!Subtarget->supportsTailCall())
     isTailCall = false;
if (isStructRet) {
           dl.getDebugLoc());
       DAG.getContext()->diagnose(Diag);
     }
   }
 
   // FIXME: handle tail calls differently.
   unsigned CallOpc;
   if (Subtarget->isThumb()) {
-    if (isCmseNSCall)
+    if (GuardWithBTI)
+      CallOpc = ARMISD::t2CALL_BTI;
+    else if (isCmseNSCall)
       CallOpc = ARMISD::tSECALL;
     else if ((!isDirect || isARMFunc) && !Subtarget->hasV5TOps())
       CallOpc = ARMISD::CALL_NOLINK;
     else
       CallOpc = ARMISD::CALL;
   } else {
     if (!isDirect && !Subtarget->hasV5TOps())
       CallOpc = ARMISD::CALL_NOLINK;

llvm/lib/Target/ARM/ARMInstrThumb2.td

 }
 
 def t2PAC    : PACBTIHintSpaceDefInst<"pac", 0b00011101>;
 def t2PACBTI : PACBTIHintSpaceDefInst<"pacbti", 0b00001101>;
 def t2BTI    : PACBTIHintSpaceNoOpsInst<"bti", 0b00001111>;
 def t2AUT    : PACBTIHintSpaceUseInst<"aut", 0b00101101> {
   let hasSideEffects = 1;
 }
+
+def ARMt2CallBTI : SDNode<"ARMISD::t2CALL_BTI", SDT_ARMcall,
+                   [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue, SDNPVariadic]>;
+
+def t2CALL_BTI : PseudoInst<(outs), (ins pred:$p, thumb_bl_target:$func),
+                 IIC_Br, [(ARMt2CallBTI tglobaladdr:$func)]>,
+                 Requires<[IsThumb2]>, Sched<[WriteBrL]>;

llvm/lib/Target/ARM/ARMSubtarget.h

protected:
   Triple TargetTriple;
 
   /// SchedModel - Processor specific instruction costs.
   MCSchedModel SchedModel;
 
   /// Selected instruction itineraries (one entry per itinerary class.)
   InstrItineraryData InstrItins;
 
+  /// NoBTIAtReturnTwice - Don't place a BTI instruction after
+  /// return-twice constructs (setjmp)
+  bool NoBTIAtReturnTwice = false;
+
   /// Options passed via command line that could influence the target
   const TargetOptions &Options;
 
   const ARMBaseTargetMachine &TM;
 
 public:
   /// This constructor initializes the data members to match that
   /// of the specified triple.
bool ignoreCSRForAllocationOrder(const MachineFunction &MF,
                                    unsigned PhysReg) const override;
   unsigned getGPRAllocationOrder(const MachineFunction &MF) const;
 
   bool fixCMSE_CVE_2021_35465() const { return FixCMSE_CVE_2021_35465; }
 
   bool hardenSlsRetBr() const { return HardenSlsRetBr; }
   bool hardenSlsBlr() const { return HardenSlsBlr; }
   bool hardenSlsNoComdat() const { return HardenSlsNoComdat; }
+
+  bool getNoBTIAtReturnTwice() const { return NoBTIAtReturnTwice; }
 };
 
 } // end namespace llvm
 
 #endif  // LLVM_LIB_TARGET_ARM_ARMSUBTARGET_H

llvm/test/CodeGen/ARM/setjmp-bti-basic.ll

+; RUN: llc -mtriple=thumbv8.1m.main-arm-none-eabi < %s | FileCheck %s --check-prefix=BTI
+; RUN: llc -mtriple=thumbv8.1m.main-arm-none-eabi -mattr=+no-bti-at-return-twice < %s | \
+; RUN: FileCheck %s --check-prefix=NOBTI
+
+; C source
+; --------
+; jmp_buf buf;
+;
+; extern void bar(int x);
+;
+; int foo(int x) {
+;   if (setjmp(buf))
+;     x = 0;
+;   else
+;     bar(x);
+;   return x;
+; }
+
+@buf = global [20 x i64] zeroinitializer, align 8
+
+define i32 @foo(i32 %x) {
+; BTI-LABEL: foo:
+; BTI:       bl setjmp
+; BTI-NEXT:  bti
+; NOBTI-LABEL: foo:
+; NOBTI:       bl setjmp
+; NOBTI-NOT:   bti
+
+entry:
+  %call = call i32 @setjmp(i64* getelementptr inbounds ([20 x i64], [20 x i64]* @buf, i32 0, i32 0)) #0
+  %tobool.not = icmp eq i32 %call, 0
+  br i1 %tobool.not, label %if.else, label %if.end
+
+if.else:                                          ; preds = %entry
+  call void @bar(i32 %x)
+  br label %if.end
+
+if.end:                                           ; preds = %entry, %if.else
+  %x.addr.0 = phi i32 [ %x, %if.else ], [ 0, %entry ]
+  ret i32 %x.addr.0
+}
+
+declare void @bar(i32)
+declare i32 @setjmp(i64*) #0
+
+attributes #0 = { returns_twice }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"branch-target-enforcement", i32 1}

llvm/test/CodeGen/ARM/setjmp-bti-outliner.ll

+; RUN: llc -mtriple=thumbv8.1m.main-arm-none-eabi -enable-machine-outliner < %s | \
+; RUN: FileCheck %s --check-prefix=BTI
+; RUN: llc -mtriple=thumbv8.1m.main-arm-none-eabi -enable-machine-outliner -mattr=+no-bti-at-return-twice < %s | FileCheck %s --check-prefix=NOBTI
+
+; C source
+; --------
+; jmp_buf buf;
+;
+; extern void h(int a, int b, int *c);
+;
+; int f(int a, int b, int c, int d) {
+;   if (setjmp(buf) != 0)
+;     return -1;
+;   h(a, b, &a);
+;   return 2 + a * (a + b) / (c + d);
+; }
+;
+; int g(int a, int b, int c, int d) {
+;   if (setjmp(buf) != 0)
+;     return -1;
+;   h(a, b, &a);
+;   return 1 + a * (a + b) / (c + d);
+; }
+
+@buf = global [20 x i64] zeroinitializer, align 8
+
+define i32 @f(i32 %a, i32 %b, i32 %c, i32 %d) {
+; BTI-LABEL: f:
+; BTI:       bl OUTLINED_FUNCTION_0
+; BTI-NEXT:  bti
+; NOBTI-LABEL: f:
+; NOBTI:       bl OUTLINED_FUNCTION_0
+; NOBTI-NEXT:   cbz r0, .LBB0_2
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %call = call i32 @setjmp(i64* getelementptr inbounds ([20 x i64], [20 x i64]* @buf, i32 0, i32 0)) #0
+  %cmp.not = icmp eq i32 %call, 0
+  br i1 %cmp.not, label %if.end, label %return
+
+if.end:                                           ; preds = %entry
+  call void @h(i32 %a, i32 %b, i32* nonnull %a.addr)
+  %0 = load i32, i32* %a.addr, align 4
+  %add = add nsw i32 %0, %b
+  %mul = mul nsw i32 %add, %0
+  %add1 = add nsw i32 %d, %c
+  %div = sdiv i32 %mul, %add1
+  %add2 = add nsw i32 %div, 2
+  br label %return
+
+return:                                           ; preds = %entry, %if.end
+  %retval.0 = phi i32 [ %add2, %if.end ], [ -1, %entry ]
+  ret i32 %retval.0
+}
+
+define i32 @g(i32 %a, i32 %b, i32 %c, i32 %d) {
+; BTI-LABEL: g:
+; BTI:       bl OUTLINED_FUNCTION_0
+; BTI-NEXT:  bti
+; NOBTI-LABEL: g:
+; NOBTI:       bl OUTLINED_FUNCTION_0
+; NOBTI-NEXT:  cbz  r0, .LBB1_2
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %call = call i32 @setjmp(i64* getelementptr inbounds ([20 x i64], [20 x i64]* @buf, i32 0, i32 0)) #0
+  %cmp.not = icmp eq i32 %call, 0
+  br i1 %cmp.not, label %if.end, label %return
+
+if.end:                                           ; preds = %entry
+  call void @h(i32 %a, i32 %b, i32* nonnull %a.addr)
+  %0 = load i32, i32* %a.addr, align 4
+  %add = add nsw i32 %0, %b
+  %mul = mul nsw i32 %add, %0
+  %add1 = add nsw i32 %d, %c
+  %div = sdiv i32 %mul, %add1
+  %add2 = add nsw i32 %div, 1
+  br label %return
+
+return:                                           ; preds = %entry, %if.end
+  %retval.0 = phi i32 [ %add2, %if.end ], [ -1, %entry ]
+  ret i32 %retval.0
+}
+
+declare void @h(i32, i32, i32*)
+declare i32 @setjmp(i64*) #0
+
+attributes #0 = { returns_twice }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"branch-target-enforcement", i32 1}