[llvm] cca40aa - [AArch64][v8.5A] Add BTI to all function starts

[Pablo Barrio via llvm-commits]

Author: Pablo Barrio
Date: 2021-04-14T15:24:01+01:00
New Revision: cca40aa8d8aa732a226c8978e53cd47e7b7c76ec

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

LOG: [AArch64][v8.5A] Add BTI to all function starts

The existing BTI placement pass avoids inserting "BTI c" when the
function has local linkage and is only directly called. However,
even in this case, there is a (small) chance that the linker later
adds a hunk with an indirect call to the function, e.g. if the
function is placed in a separate section and moved far away from
its callers. Make sure to add BTI for these functions too.

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

Added: 
    

Modified: 
    llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
    llvm/test/CodeGen/AArch64/branch-target-enforcement.mir
    llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index d3b5166585c3..6506180bee56 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -64,7 +64,6 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(
       dbgs() << "********** AArch64 Branch Targets  **********\n"
              << "********** Function: " << MF.getName() << '\n');
-  const Function &F = MF.getFunction();
 
   // LLVM does not consider basic blocks which are the targets of jump tables
   // to be address-taken (the address can't escape anywhere else), but they are
@@ -78,13 +77,16 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   bool MadeChange = false;
   for (MachineBasicBlock &MBB : MF) {
     bool CouldCall = false, CouldJump = false;
-    // If the function is address-taken or externally-visible, it could be
-    // indirectly called. PLT entries and tail-calls use BR, but when they are
+    // Even in cases where a function has internal linkage and is only called
+    // directly in its translation unit, it can still be called indirectly if
+    // the linker decides to add a thunk to it for whatever reason (say, for
+    // example, if it is finally placed far from its call site and a BL is not
+    // long-range enough). PLT entries and tail-calls use BR, but when they are
     // are in guarded pages should all use x16 or x17 to hold the called
     // address, so we don't need to set CouldJump here. BR instructions in
     // non-guarded pages (which might be non-BTI-aware code) are allowed to
     // branch to a "BTI c" using any register.
-    if (&MBB == &*MF.begin() && (F.hasAddressTaken() || !F.hasLocalLinkage()))
+    if (&MBB == &*MF.begin())
       CouldCall = true;
 
     // If the block itself is address-taken, it could be indirectly branched

diff  --git a/llvm/test/CodeGen/AArch64/branch-target-enforcement.mir b/llvm/test/CodeGen/AArch64/branch-target-enforcement.mir
index f34fb2b84bcc..b8cb7750ff2c 100644
--- a/llvm/test/CodeGen/AArch64/branch-target-enforcement.mir
+++ b/llvm/test/CodeGen/AArch64/branch-target-enforcement.mir
@@ -116,12 +116,15 @@ body:             |
     RET undef $lr, implicit killed $w0
 
 ---
-# Internal function, not address-taken in this module, so no BTI needed.
+# Internal function, not address-taken in this module, however the compiler
+# cannot 100% ensure that later parts of the toolchain won't add indirect
+# jumps. E.g. a linker adding a thunk to extend the range of a direct jump.
+# Therefore, even this case needs a BTI.
 name:            simple_internal
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: simple_internal
-    ; CHECK-NOT: HINT
+    ; CHECK: HINT 34
     ; CHECK: RET
     $w0 = ORRWrs $wzr, $wzr, 0
     RET undef $lr, implicit killed $w0

diff  --git a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
index 01c3d2b0666a..5a89c5148b1e 100644
--- a/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
+++ b/llvm/test/CodeGen/AArch64/patchable-function-entry-bti.ll
@@ -48,19 +48,20 @@ define void @f2_1() "patchable-function-entry"="1" "patchable-function-prefix"="
 }
 
 ;; -fpatchable-function-entry=1 -mbranch-protection=bti
-;; For M=0, don't create .Lpatch0 if the initial instruction is not BTI,
-;; even if other basic blocks may have BTI.
+;; We add BTI c even when the function has internal linkage
 define internal void @f1i(i64 %v) "patchable-function-entry"="1" "branch-target-enforcement"="true" {
 ; CHECK-LABEL: f1i:
 ; CHECK-NEXT: .Lfunc_begin3:
 ; CHECK:      // %bb.0:
+; CHECK-NEXT:  hint #34
+; CHECK-NEXT:  .Lpatch1:
 ; CHECK-NEXT:  nop
 ;; Other basic blocks have BTI, but they don't affect our decision to not create .Lpatch0
 ; CHECK:      .LBB{{.+}} // %sw.bb1
 ; CHECK-NEXT:  hint #36
 ; CHECK:      .section __patchable_function_entries,"awo", at progbits,f1i{{$}}
 ; CHECK-NEXT: .p2align 3
-; CHECK-NEXT: .xword .Lfunc_begin3
+; CHECK-NEXT: .xword .Lpatch1
 entry:
   switch i64 %v, label %sw.bb0 [
     i64 1, label %sw.bb1