[clang] 0fbb174 - [ARM] Implement setjmp BTI placement for PACBTI-M

Ties Stuij via cfe-commits cfe-commits at lists.llvm.org
Mon Dec 6 03:07:32 PST 2021


Author: Ties Stuij
Date: 2021-12-06T11:07:10Z
New Revision: 0fbb17458a01a6b388fc67661ffb92969503e977

URL: https://github.com/llvm/llvm-project/commit/0fbb17458a01a6b388fc67661ffb92969503e977
DIFF: https://github.com/llvm/llvm-project/commit/0fbb17458a01a6b388fc67661ffb92969503e977.diff

LOG: [ARM] Implement setjmp BTI placement for PACBTI-M

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

Reviewed By: labrinea

Differential Revision: https://reviews.llvm.org/D112427

Added: 
    clang/test/Driver/arm-bti-return-twice.c
    llvm/test/CodeGen/ARM/setjmp-bti-basic.ll
    llvm/test/CodeGen/ARM/setjmp-bti-outliner.ll

Modified: 
    clang/docs/ClangCommandLineReference.rst
    clang/include/clang/Driver/Options.td
    clang/lib/Driver/ToolChains/Arch/ARM.cpp
    llvm/lib/Target/ARM/ARM.td
    llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp
    llvm/lib/Target/ARM/ARMISelLowering.cpp
    llvm/lib/Target/ARM/ARMISelLowering.h
    llvm/lib/Target/ARM/ARMInstrThumb2.td
    llvm/lib/Target/ARM/ARMSubtarget.h

Removed: 
    


################################################################################
diff  --git a/clang/docs/ClangCommandLineReference.rst b/clang/docs/ClangCommandLineReference.rst
index 8d4ffa6a38fc4..97807009fd911 100644
--- a/clang/docs/ClangCommandLineReference.rst
+++ b/clang/docs/ClangCommandLineReference.rst
@@ -3260,6 +3260,11 @@ Thread pointer access method (AArch32/AArch64 only)
 
 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

diff  --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 4e6dd20503446..7c257e4a474fa 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -3338,6 +3338,11 @@ def mno_fix_cortex_a53_835769 : Flag<["-"], "mno-fix-cortex-a53-835769">,
 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)">;

diff  --git a/clang/lib/Driver/ToolChains/Arch/ARM.cpp b/clang/lib/Driver/ToolChains/Arch/ARM.cpp
index 8d5c64d975502..e03bed0a6de64 100644
--- a/clang/lib/Driver/ToolChains/Arch/ARM.cpp
+++ b/clang/lib/Driver/ToolChains/Arch/ARM.cpp
@@ -875,6 +875,8 @@ void arm::getARMTargetFeatures(const Driver &D, const llvm::Triple &Triple,
     }
   }
 
+  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) {

diff  --git a/clang/test/Driver/arm-bti-return-twice.c b/clang/test/Driver/arm-bti-return-twice.c
new file mode 100644
index 0000000000000..c5cd385a34e20
--- /dev/null
+++ b/clang/test/Driver/arm-bti-return-twice.c
@@ -0,0 +1,7 @@
+// 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"

diff  --git a/llvm/lib/Target/ARM/ARM.td b/llvm/lib/Target/ARM/ARM.td
index e03dd597eb650..8173fe4036a85 100644
--- a/llvm/lib/Target/ARM/ARM.td
+++ b/llvm/lib/Target/ARM/ARM.td
@@ -446,6 +446,11 @@ 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
 //

diff  --git a/llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp b/llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp
index 7a35f252b22ae..d6b1444c35e48 100644
--- a/llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/ARM/ARMExpandPseudoInsts.cpp
@@ -3073,6 +3073,22 @@ bool ARMExpandPseudo::ExpandMI(MachineBasicBlock &MBB,
       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();

diff  --git a/llvm/lib/Target/ARM/ARMISelLowering.cpp b/llvm/lib/Target/ARM/ARMISelLowering.cpp
index 9a4e64285caa8..add5b1cf6dfe3 100644
--- a/llvm/lib/Target/ARM/ARMISelLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMISelLowering.cpp
@@ -1658,6 +1658,7 @@ const char *ARMTargetLowering::getTargetNodeName(unsigned Opcode) const {
     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)
@@ -2321,6 +2322,12 @@ ARMTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   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"))
@@ -2726,7 +2733,9 @@ ARMTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   // FIXME: handle tail calls 
diff erently.
   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;

diff  --git a/llvm/lib/Target/ARM/ARMISelLowering.h b/llvm/lib/Target/ARM/ARMISelLowering.h
index e3b422358caee..1c5f8389f57cd 100644
--- a/llvm/lib/Target/ARM/ARMISelLowering.h
+++ b/llvm/lib/Target/ARM/ARMISelLowering.h
@@ -69,6 +69,7 @@ class VectorType;
     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).

diff  --git a/llvm/lib/Target/ARM/ARMInstrThumb2.td b/llvm/lib/Target/ARM/ARMInstrThumb2.td
index 4471317f4ea41..6e8e61ca2b8e4 100644
--- a/llvm/lib/Target/ARM/ARMInstrThumb2.td
+++ b/llvm/lib/Target/ARM/ARMInstrThumb2.td
@@ -5736,3 +5736,10 @@ 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]>;

diff  --git a/llvm/lib/Target/ARM/ARMSubtarget.h b/llvm/lib/Target/ARM/ARMSubtarget.h
index d51a888c951fe..cb46a6b7dc26b 100644
--- a/llvm/lib/Target/ARM/ARMSubtarget.h
+++ b/llvm/lib/Target/ARM/ARMSubtarget.h
@@ -534,6 +534,10 @@ class ARMSubtarget : public ARMGenSubtargetInfo {
   /// 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;
 
@@ -948,6 +952,8 @@ class ARMSubtarget : public ARMGenSubtargetInfo {
   bool hardenSlsRetBr() const { return HardenSlsRetBr; }
   bool hardenSlsBlr() const { return HardenSlsBlr; }
   bool hardenSlsNoComdat() const { return HardenSlsNoComdat; }
+
+  bool getNoBTIAtReturnTwice() const { return NoBTIAtReturnTwice; }
 };
 
 } // end namespace llvm

diff  --git a/llvm/test/CodeGen/ARM/setjmp-bti-basic.ll b/llvm/test/CodeGen/ARM/setjmp-bti-basic.ll
new file mode 100644
index 0000000000000..e18f87d82e6c6
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/setjmp-bti-basic.ll
@@ -0,0 +1,50 @@
+; 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;
+; }
+
+ at 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}

diff  --git a/llvm/test/CodeGen/ARM/setjmp-bti-outliner.ll b/llvm/test/CodeGen/ARM/setjmp-bti-outliner.ll
new file mode 100644
index 0000000000000..20fd1768d6077
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/setjmp-bti-outliner.ll
@@ -0,0 +1,92 @@
+; 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);
+; }
+
+ at 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}


        


More information about the cfe-commits mailing list