[llvm] [AArch64][BTI] Add BTI at EH entries. (PR #155308)

[Shashi Shankar via llvm-commits]

https://github.com/shashforge updated https://github.com/llvm/llvm-project/pull/155308

>From 8daa8a743a65554898bdf424b9d05ca7992eb7e9 Mon Sep 17 00:00:00 2001
From: Shashi Shankar <shashishankar1687 at gmail.com>
Date: Mon, 25 Aug 2025 23:13:19 +0200
Subject: [PATCH] [AArch64][BTI] Add BTI at EH entries: BTI j for Itanium EH,
 BTI c for WinEH

Signed-off-by: Shashi Shankar <shashishankar1687 at gmail.com>
---
 .../Target/AArch64/AArch64BranchTargets.cpp   | 27 ++++++++++++++----
 llvm/test/CodeGen/AArch64/bti-ehpad.ll        | 28 +++++++++++++++++++
 .../test/CodeGen/AArch64/wineh-bti-funclet.ll | 27 ++++++++++++++++++
 3 files changed, 77 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/bti-ehpad.ll
 create mode 100644 llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll

diff --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index 3436dc9ef4521..83abdaac3fec1 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -62,9 +62,8 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   if (!MF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
     return false;
 
-  LLVM_DEBUG(
-      dbgs() << "********** AArch64 Branch Targets  **********\n"
-             << "********** Function: " << MF.getName() << '\n');
+  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
@@ -103,6 +102,12 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
         JumpTableTargets.count(&MBB))
       CouldJump = true;
 
+    if (MBB.isEHPad()) {
+      if (HasWinCFI && (MBB.isEHFuncletEntry() || MBB.isCleanupFuncletEntry()))
+        CouldCall = true;
+      else
+        CouldJump = true;
+    }
     if (CouldCall || CouldJump) {
       addBTI(MBB, CouldCall, CouldJump, HasWinCFI);
       MadeChange = true;
@@ -129,8 +134,15 @@ void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
   assert(HintNum != 32 && "No target kinds!");
 
   auto MBBI = MBB.begin();
+  bool SawEHLabel = false;
 
-  // Skip the meta instructions, those will be removed anyway.
+  // If the block starts with EH_LABEL(s), skip them first and remember we saw
+  // one.
+  while (MBBI != MBB.end() && MBBI->isEHLabel()) {
+    ++MBBI;
+    SawEHLabel = true;
+  }
+  /// Skip the meta instructions, those will be removed anyway.
   for (; MBBI != MBB.end() &&
          (MBBI->isMetaInstruction() || MBBI->getOpcode() == AArch64::EMITBKEY);
        ++MBBI)
@@ -147,7 +159,12 @@ void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
     BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
             TII->get(AArch64::SEH_Nop));
   }
-  BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
+  // Insertion policy:
+  //  If the block started with an EH_LABEL, insert BTI after the label and
+  //  meta (at MBBI).Otherwise, keep legacy behavior: insert at block begin
+  //  (before CFI/SEH NOP), to preserve existing codegen/test expectations.
+  MachineBasicBlock::iterator InsertHere = SawEHLabel ? MBBI : MBB.begin();
+  BuildMI(MBB, InsertHere, MBB.findDebugLoc(InsertHere),
           TII->get(AArch64::HINT))
       .addImm(HintNum);
 }
diff --git a/llvm/test/CodeGen/AArch64/bti-ehpad.ll b/llvm/test/CodeGen/AArch64/bti-ehpad.ll
new file mode 100644
index 0000000000000..16d12810e2c2d
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/bti-ehpad.ll
@@ -0,0 +1,28 @@
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu %s -o - | FileCheck %s
+
+declare i32 @__gxx_personality_v0(...)
+declare void @may_throw()
+
+define void @test() #0 personality ptr @__gxx_personality_v0 {
+entry:
+  invoke void @may_throw()
+          to label %ret unwind label %lpad
+
+lpad:
+  landingpad { ptr, i32 } cleanup
+  ret void
+
+ret:
+  ret void
+}
+
+; Enforce BTI at codegen.
+attributes #0 = { "branch-target-enforcement"="true" "target-features"="+bti" }
+
+; CHECK-LABEL: test:
+; The assembler prints ".LBB...: // %lpad" and then an EH_LABEL line.
+; We require BTI j immediately *after* the EH_LABEL.
+; CHECK:      {{// *%lpad}}
+; CHECK-NEXT: {{.*EH_LABEL}}
+; CHECK-NEXT: {{(bti +j|hint +#36)}}
+
diff --git a/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll b/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll
new file mode 100644
index 0000000000000..75945919b3b7d
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll
@@ -0,0 +1,27 @@
+; RUN: llc -mtriple=aarch64-windows -mattr=+bti -o - %s | FileCheck %s
+
+declare i32 @__CxxFrameHandler3(...)
+declare void @may_throw()
+
+define dso_local void @"?w@@YAXXZ"() #0 personality ptr @__CxxFrameHandler3 {
+entry:
+  invoke void @may_throw()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:
+  %cs = catchswitch within none [label %catch] unwind to caller
+
+catch:
+  %cp = catchpad within %cs [ptr null, i32 0, ptr null]
+  call void @may_throw() ["funclet"(token %cp)]
+  catchret from %cp to label %try.cont
+
+try.cont:
+  ret void
+}
+
+attributes #0 = { "branch-target-enforcement"="true" }
+
+; CHECK-LABEL: "?w@@YAXXZ":
+; Either 'bti c' / 'hint #34' or a PAC prologue satisfies call-like entry.
+; CHECK: {{(hint +#34|bti +c|paciasp|pacibsp)}}