Skip to content

Commit

Permalink
Set branch protection function attributes
Browse files Browse the repository at this point in the history
Since LLVM 19, it is necessary to set not only module flags, but
also function attributes for branch protection on aarch64. See
llvm/llvm-project@e15d67c
for the relevant LLVM change.
  • Loading branch information
nikic committed Jul 24, 2024
1 parent d24930c commit ea7625f
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 3 deletions.
31 changes: 28 additions & 3 deletions compiler/rustc_codegen_llvm/src/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions tests/codegen/branch-protection-old-llvm.rs
Original file line number Diff line number Diff line change
@@ -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
10 changes: 10 additions & 0 deletions tests/codegen/branch-protection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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
Expand Down

0 comments on commit ea7625f

Please sign in to comment.