From ea7625f4266e69498ea5b9287a72e543a514a837 Mon Sep 17 00:00:00 2001 From: Nikita Popov Date: Wed, 24 Jul 2024 16:47:56 +0200 Subject: [PATCH] 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. --- compiler/rustc_codegen_llvm/src/attributes.rs | 31 +++++++++++-- tests/codegen/branch-protection-old-llvm.rs | 45 +++++++++++++++++++ tests/codegen/branch-protection.rs | 10 +++++ 3 files changed, 83 insertions(+), 3 deletions(-) create mode 100644 tests/codegen/branch-protection-old-llvm.rs diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs index 3877460fcdb0d..ff01a2088812d 100644 --- a/compiler/rustc_codegen_llvm/src/attributes.rs +++ b/compiler/rustc_codegen_llvm/src/attributes.rs @@ -4,7 +4,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; @@ -407,8 +407,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 0000000000000..bb3c7a4b70c3c --- /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 a29ec67d578b8..2f5ff9e98c22d 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