From 018ac7847bb5a18ff68cf05b1d68f37fb32792d8 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Martin=20Storsj=C3=B6?= <martin@martin.st>
Date: Wed, 5 Oct 2022 13:49:20 +0300
Subject: [PATCH] [AArch64] Add SEH_Nop opcodes for BTI hints

These are harmless for the unwinder - the unwinder doesn't need to
handle them for being able to unwind correctly.

Only add the opcodes when the branch target is in a SEH prologue;
for jumptables e.g. within a function, we shouldn't add any SEH
opcodes.

Differential Revision: https://reviews.llvm.org/D135277
---
 llvm/lib/Target/AArch64/AArch64BranchTargets.cpp | 12 +++--
 llvm/test/CodeGen/AArch64/wineh-bti.ll           | 65 ++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/wineh-bti.ll

diff --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index d2acd1d..e2d1051 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -39,7 +39,8 @@ public:
   StringRef getPassName() const override { return AARCH64_BRANCH_TARGETS_NAME; }
 
 private:
-  void addBTI(MachineBasicBlock &MBB, bool CouldCall, bool CouldJump);
+  void addBTI(MachineBasicBlock &MBB, bool CouldCall, bool CouldJump,
+              bool NeedsWinCFI);
 };
 } // end anonymous namespace
 
@@ -75,6 +76,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
         JumpTableTargets.insert(MBB);
 
   bool MadeChange = false;
+  bool HasWinCFI = MF.hasWinCFI();
   for (MachineBasicBlock &MBB : MF) {
     bool CouldCall = false, CouldJump = false;
     // Even in cases where a function has internal linkage and is only called
@@ -95,7 +97,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
       CouldJump = true;
 
     if (CouldCall || CouldJump) {
-      addBTI(MBB, CouldCall, CouldJump);
+      addBTI(MBB, CouldCall, CouldJump, HasWinCFI);
       MadeChange = true;
     }
   }
@@ -104,7 +106,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
 }
 
 void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
-                                  bool CouldJump) {
+                                  bool CouldJump, bool HasWinCFI) {
   LLVM_DEBUG(dbgs() << "Adding BTI " << (CouldJump ? "j" : "")
                     << (CouldCall ? "c" : "") << " to " << MBB.getName()
                     << "\n");
@@ -134,6 +136,10 @@ void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
        MBBI->getOpcode() == AArch64::PACIBSP))
     return;
 
+  if (HasWinCFI && MBBI->getFlag(MachineInstr::FrameSetup)) {
+    BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
+            TII->get(AArch64::SEH_Nop));
+  }
   BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
           TII->get(AArch64::HINT))
       .addImm(HintNum);
diff --git a/llvm/test/CodeGen/AArch64/wineh-bti.ll b/llvm/test/CodeGen/AArch64/wineh-bti.ll
new file mode 100644
index 0000000..125b0cb7
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/wineh-bti.ll
@@ -0,0 +1,65 @@
+; RUN: llc < %s -mtriple=aarch64-windows | FileCheck %s
+
+define dso_local i32 @func(i32 %in) {
+entry:
+  call void asm sideeffect "", "~{x19}"()
+  switch i32 %in, label %def [
+    i32 0, label %lbl1
+    i32 1, label %lbl2
+    i32 2, label %lbl3
+    i32 4, label %lbl4
+  ]
+
+def:
+  ret i32 0
+
+lbl1:
+  call void asm sideeffect "", ""()
+  ret i32 1
+
+lbl2:
+  ret i32 2
+
+lbl3:
+  ret i32 4
+
+lbl4:
+  ret i32 8
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 8, !"branch-target-enforcement", i32 1}
+!1 = !{i32 8, !"sign-return-address", i32 1}
+
+; CHECK-LABEL: func:
+; CHECK-NEXT: .seh_proc func
+; CHECK-NEXT: // %bb.0:
+; CHECK-NEXT: hint #34
+; CHECK-NEXT: .seh_nop
+; CHECK-NEXT: str x19, [sp, #-16]!
+; CHECK-NEXT: .seh_save_reg_x x19, 16
+; CHECK-NEXT: .seh_endprologue
+
+; CHECK:      .LBB0_2:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #1
+
+; CHECK:      .LBB0_3:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #2
+
+; CHECK:      .LBB0_4:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #4
+
+; CHECK:      .LBB0_5:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #8
+
+; CHECK:      .seh_startepilogue
+; CHECK-NEXT: ldr x19, [sp], #16
+; CHECK-NEXT: .seh_save_reg_x x19, 16
+; CHECK-NEXT: .seh_endepilogue
+; CHECK-NEXT: ret
+; CHECK-NEXT: .seh_endfunclet
-- 
2.7.4