Rollup merge of #128141 - nikic:aarch64-bti, r=DianQK,cuviper

Set branch protection function attributes

Since LLVM 19, it is necessary to set not only module flags, but also function attributes for branch protection on aarch64. See https://github.com/llvm/llvm-project/commit/e15d67cfc2e5775cc79281aa860f3ad3be628f39 for the relevant LLVM change.

Fixes https://github.com/rust-lang/rust/issues/127829.

3 files changed, 83 insertions, 3 deletions

diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index 4ef8cd7ce30..ad38814a68b 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -5,7 +5,7 @@ use rustc_codegen_ssa::traits::*;
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFunctionEntry};
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_session::config::{FunctionReturn, OptLevel};
+use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
 use rustc_span::symbol::sym;
 use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
 use smallvec::SmallVec;
@@ -405,8 +405,33 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         // And it is a module-level attribute, so the alternative is pulling naked functions into new LLVM modules.
         // Otherwise LLVM's "naked" functions come with endbr prefixes per https://github.com/rust-lang/rust/issues/98768
         to_add.push(AttributeKind::NoCfCheck.create_attr(cx.llcx));
-        // Need this for AArch64.
-        to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
+        if llvm_util::get_version() < (19, 0, 0) {
+            // Prior to LLVM 19, branch-target-enforcement was disabled by setting the attribute to
+            // the string "false". Now it is disabled by absence of the attribute.
+            to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
+        }
+    } else if llvm_util::get_version() >= (19, 0, 0) {
+        // For non-naked functions, set branch protection attributes on aarch64.
+        if let Some(BranchProtection { bti, pac_ret }) =
+            cx.sess().opts.unstable_opts.branch_protection
+        {
+            assert!(cx.sess().target.arch == "aarch64");
+            if bti {
+                to_add.push(llvm::CreateAttrString(cx.llcx, "branch-target-enforcement"));
+            }
+            if let Some(PacRet { leaf, key }) = pac_ret {
+                to_add.push(llvm::CreateAttrStringValue(
+                    cx.llcx,
+                    "sign-return-address",
+                    if leaf { "all" } else { "non-leaf" },
+                ));
+                to_add.push(llvm::CreateAttrStringValue(
+                    cx.llcx,
+                    "sign-return-address-key",
+                    if key == PAuthKey::A { "a_key" } else { "b_key" },
+                ));
+            }
+        }
     }
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR)
         || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR_ZEROED)
diff --git a/tests/codegen/branch-protection-old-llvm.rs b/tests/codegen/branch-protection-old-llvm.rs
new file mode 100644
index 00000000000..bb3c7a4b70c
--- /dev/null
+++ b/tests/codegen/branch-protection-old-llvm.rs
@@ -0,0 +1,45 @@
+// Test that the correct module flags are emitted with different branch protection flags.
+
+//@ revisions: BTI PACRET LEAF BKEY NONE
+//@ needs-llvm-components: aarch64
+//@ [BTI] compile-flags: -Z branch-protection=bti
+//@ [PACRET] compile-flags: -Z branch-protection=pac-ret
+//@ [LEAF] compile-flags: -Z branch-protection=pac-ret,leaf
+//@ [BKEY] compile-flags: -Z branch-protection=pac-ret,b-key
+//@ compile-flags: --target aarch64-unknown-linux-gnu
+//@ ignore-llvm-version: 19 - 99
+
+#![crate_type = "lib"]
+#![feature(no_core, lang_items)]
+#![no_core]
+
+#[lang = "sized"]
+trait Sized {}
+
+// A basic test function.
+pub fn test() {}
+
+// BTI: !"branch-target-enforcement", i32 1
+// BTI: !"sign-return-address", i32 0
+// BTI: !"sign-return-address-all", i32 0
+// BTI: !"sign-return-address-with-bkey", i32 0
+
+// PACRET: !"branch-target-enforcement", i32 0
+// PACRET: !"sign-return-address", i32 1
+// PACRET: !"sign-return-address-all", i32 0
+// PACRET: !"sign-return-address-with-bkey", i32 0
+
+// LEAF: !"branch-target-enforcement", i32 0
+// LEAF: !"sign-return-address", i32 1
+// LEAF: !"sign-return-address-all", i32 1
+// LEAF: !"sign-return-address-with-bkey", i32 0
+
+// BKEY: !"branch-target-enforcement", i32 0
+// BKEY: !"sign-return-address", i32 1
+// BKEY: !"sign-return-address-all", i32 0
+// BKEY: !"sign-return-address-with-bkey", i32 1
+
+// NONE-NOT: branch-target-enforcement
+// NONE-NOT: sign-return-address
+// NONE-NOT: sign-return-address-all
+// NONE-NOT: sign-return-address-with-bkey
diff --git a/tests/codegen/branch-protection.rs b/tests/codegen/branch-protection.rs
index a29ec67d578..2f5ff9e98c2 100644
--- a/tests/codegen/branch-protection.rs
+++ b/tests/codegen/branch-protection.rs
@@ -7,6 +7,7 @@
 //@ [LEAF] compile-flags: -Z branch-protection=pac-ret,leaf
 //@ [BKEY] compile-flags: -Z branch-protection=pac-ret,b-key
 //@ compile-flags: --target aarch64-unknown-linux-gnu
+//@ min-llvm-version: 19
 
 #![crate_type = "lib"]
 #![feature(no_core, lang_items)]
@@ -16,23 +17,32 @@
 trait Sized {}
 
 // A basic test function.
+// CHECK: @test(){{.*}} [[ATTR:#[0-9]+]] {
+#[no_mangle]
 pub fn test() {}
 
+// BTI: attributes [[ATTR]] = {{.*}} "branch-target-enforcement"
 // BTI: !"branch-target-enforcement", i32 1
 // BTI: !"sign-return-address", i32 0
 // BTI: !"sign-return-address-all", i32 0
 // BTI: !"sign-return-address-with-bkey", i32 0
 
+// PACRET: attributes [[ATTR]] = {{.*}} "sign-return-address"="non-leaf"
+// PACRET-SAME: "sign-return-address-key"="a_key"
 // PACRET: !"branch-target-enforcement", i32 0
 // PACRET: !"sign-return-address", i32 1
 // PACRET: !"sign-return-address-all", i32 0
 // PACRET: !"sign-return-address-with-bkey", i32 0
 
+// LEAF: attributes [[ATTR]] = {{.*}} "sign-return-address"="all"
+// LEAF-SAME: "sign-return-address-key"="a_key"
 // LEAF: !"branch-target-enforcement", i32 0
 // LEAF: !"sign-return-address", i32 1
 // LEAF: !"sign-return-address-all", i32 1
 // LEAF: !"sign-return-address-with-bkey", i32 0
 
+// BKEY: attributes [[ATTR]] = {{.*}} "sign-return-address"="non-leaf"
+// BKEY-SAME: "sign-return-address-key"="b_key"
 // BKEY: !"branch-target-enforcement", i32 0
 // BKEY: !"sign-return-address", i32 1
 // BKEY: !"sign-return-address-all", i32 0