Skip to content

Commit 837cc16

Browse files
Jmc18134JamieCunliffe
Jmc18134
authored and
JamieCunliffe
committed
Add codegen option for branch protection and pointer authentication on AArch64
The branch-protection codegen option enables the use of hint-space pointer authentication code for AArch64 targets
1 parent 2446a21 commit 837cc16

File tree

12 files changed

+266
-7
lines changed

12 files changed

+266
-7
lines changed

Diff for: ‎compiler/rustc_codegen_llvm/src/attributes.rs

+53-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use rustc_hir::def_id::DefId;
99
use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
1010
use rustc_middle::ty::layout::HasTyCtxt;
1111
use rustc_middle::ty::{self, TyCtxt};
12-
use rustc_session::config::OptLevel;
12+
use rustc_session::config::{BranchProtection, OptLevel, PAuthKey};
1313
use rustc_session::Session;
1414
use rustc_target::spec::abi::Abi;
1515
use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
@@ -203,6 +203,58 @@ pub fn non_lazy_bind(sess: &Session, llfn: &'ll Value) {
203203
}
204204
}
205205

206+
pub fn set_branch_protection(sess: &Session, llfn: &'ll Value) {
207+
// Setting PAC/BTI function attributes is only necessary for LLVM 11 and earlier.
208+
// For LLVM 12 and greater, module-level metadata attributes are set in
209+
// `compiler/rustc_codegen_llvm/src/context.rs`.
210+
if llvm_util::get_version() >= (12, 0, 0) {
211+
return;
212+
}
213+
214+
let BranchProtection { bti, pac_ret: pac } = sess.opts.cg.branch_protection;
215+
216+
if bti {
217+
llvm::AddFunctionAttrString(
218+
llfn,
219+
llvm::AttributePlace::Function,
220+
cstr!("branch-target-enforcement"),
221+
);
222+
}
223+
224+
if let Some(pac_opts) = pac {
225+
if pac_opts.leaf {
226+
llvm::AddFunctionAttrStringValue(
227+
llfn,
228+
llvm::AttributePlace::Function,
229+
cstr!("sign-return-address"),
230+
cstr!("non-leaf"),
231+
);
232+
} else {
233+
llvm::AddFunctionAttrStringValue(
234+
llfn,
235+
llvm::AttributePlace::Function,
236+
cstr!("sign-return-address"),
237+
cstr!("all"),
238+
);
239+
}
240+
241+
match pac_opts.key {
242+
PAuthKey::A => llvm::AddFunctionAttrStringValue(
243+
llfn,
244+
llvm::AttributePlace::Function,
245+
cstr!("sign-return-address-key"),
246+
cstr!("a_key"),
247+
),
248+
PAuthKey::B => llvm::AddFunctionAttrStringValue(
249+
llfn,
250+
llvm::AttributePlace::Function,
251+
cstr!("sign-return-address-key"),
252+
cstr!("b_key"),
253+
),
254+
}
255+
}
256+
}
257+
206258
pub(crate) fn default_optimisation_attrs(sess: &Session, llfn: &'ll Value) {
207259
match sess.opts.optimize {
208260
OptLevel::Size => {

Diff for: ‎compiler/rustc_codegen_llvm/src/context.rs

+33-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use rustc_middle::ty::layout::{
2121
};
2222
use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
2323
use rustc_middle::{bug, span_bug};
24-
use rustc_session::config::{CFGuard, CrateType, DebugInfo};
24+
use rustc_session::config::{BranchProtection, CFGuard, CrateType, DebugInfo, PAuthKey};
2525
use rustc_session::Session;
2626
use rustc_span::source_map::Span;
2727
use rustc_span::symbol::Symbol;
@@ -242,6 +242,38 @@ pub unsafe fn create_module(
242242
}
243243
}
244244

245+
if sess.target.arch == "aarch64" {
246+
let BranchProtection { bti, pac_ret: pac } = sess.opts.cg.branch_protection;
247+
248+
llvm::LLVMRustAddModuleFlag(
249+
llmod,
250+
"branch-target-enforcement\0".as_ptr().cast(),
251+
bti.into(),
252+
);
253+
254+
if let Some(pac_opts) = pac {
255+
llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address\0".as_ptr().cast(), 1);
256+
llvm::LLVMRustAddModuleFlag(
257+
llmod,
258+
"sign-return-address-all\0".as_ptr().cast(),
259+
pac_opts.leaf.into(),
260+
);
261+
llvm::LLVMRustAddModuleFlag(
262+
llmod,
263+
"sign-return-address-with-bkey\0".as_ptr().cast(),
264+
if pac_opts.key == PAuthKey::A { 0 } else { 1 },
265+
);
266+
} else {
267+
llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address\0".as_ptr().cast(), 0);
268+
llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address-all\0".as_ptr().cast(), 0);
269+
llvm::LLVMRustAddModuleFlag(
270+
llmod,
271+
"sign-return-address-with-bkey\0".as_ptr().cast(),
272+
0,
273+
);
274+
}
275+
}
276+
245277
llmod
246278
}
247279

Diff for: ‎compiler/rustc_codegen_llvm/src/declare.rs

+5
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,13 @@ fn declare_raw_fn(
4545
llvm::Attribute::NoRedZone.apply_llfn(Function, llfn);
4646
}
4747

48+
if cx.tcx.sess.target.arch == "aarch64" {
49+
attributes::set_branch_protection(cx.tcx.sess, llfn);
50+
}
51+
4852
attributes::default_optimisation_attrs(cx.tcx.sess, llfn);
4953
attributes::non_lazy_bind(cx.sess(), llfn);
54+
5055
llfn
5156
}
5257

Diff for: ‎compiler/rustc_interface/src/tests.rs

+7-2
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@ use rustc_session::config::{build_configuration, build_session_options, to_crate
88
use rustc_session::config::{
99
rustc_optgroups, ErrorOutputType, ExternLocation, LocationDetail, Options, Passes,
1010
};
11-
use rustc_session::config::{CFGuard, ExternEntry, LinkerPluginLto, LtoCli, SwitchWithOptPath};
1211
use rustc_session::config::{
13-
Externs, OutputType, OutputTypes, SymbolManglingVersion, WasiExecModel,
12+
BranchProtection, Externs, OutputType, OutputTypes, PAuthKey, PacRet, SymbolManglingVersion,
13+
WasiExecModel,
1414
};
15+
use rustc_session::config::{CFGuard, ExternEntry, LinkerPluginLto, LtoCli, SwitchWithOptPath};
1516
use rustc_session::lint::Level;
1617
use rustc_session::search_paths::SearchPath;
1718
use rustc_session::utils::{CanonicalizedPath, NativeLib, NativeLibKind};
@@ -566,6 +567,10 @@ fn test_codegen_options_tracking_hash() {
566567

567568
// Make sure that changing a [TRACKED] option changes the hash.
568569
// This list is in alphabetical order.
570+
tracked!(
571+
branch_protection,
572+
BranchProtection { bti: true, pac_ret: Some(PacRet { leaf: true, key: PAuthKey::B }) }
573+
);
569574
tracked!(code_model, Some(CodeModel::Large));
570575
tracked!(control_flow_guard, CFGuard::Checks);
571576
tracked!(debug_assertions, Some(true));

Diff for: ‎compiler/rustc_session/src/config.rs

+28-3
Original file line numberDiff line numberDiff line change
@@ -842,6 +842,30 @@ impl Passes {
842842
}
843843
}
844844

845+
#[derive(Clone, Copy, Hash, Debug, PartialEq)]
846+
pub enum PAuthKey {
847+
A,
848+
B,
849+
}
850+
851+
#[derive(Clone, Copy, Hash, Debug, PartialEq)]
852+
pub struct PacRet {
853+
pub leaf: bool,
854+
pub key: PAuthKey,
855+
}
856+
857+
#[derive(Clone, Copy, Hash, Debug, PartialEq)]
858+
pub struct BranchProtection {
859+
pub bti: bool,
860+
pub pac_ret: Option<PacRet>,
861+
}
862+
863+
impl Default for BranchProtection {
864+
fn default() -> Self {
865+
BranchProtection { bti: false, pac_ret: None }
866+
}
867+
}
868+
845869
pub const fn default_lib_output() -> CrateType {
846870
CrateType::Rlib
847871
}
@@ -2487,9 +2511,9 @@ impl PpMode {
24872511
crate mod dep_tracking {
24882512
use super::LdImpl;
24892513
use super::{
2490-
CFGuard, CrateType, DebugInfo, ErrorOutputType, InstrumentCoverage, LinkerPluginLto,
2491-
LocationDetail, LtoCli, OptLevel, OutputType, OutputTypes, Passes, SourceFileHashAlgorithm,
2492-
SwitchWithOptPath, SymbolManglingVersion, TrimmedDefPaths,
2514+
BranchProtection, CFGuard, CrateType, DebugInfo, ErrorOutputType, InstrumentCoverage,
2515+
LinkerPluginLto, LocationDetail, LtoCli, OptLevel, OutputType, OutputTypes, Passes,
2516+
SourceFileHashAlgorithm, SwitchWithOptPath, SymbolManglingVersion, TrimmedDefPaths,
24932517
};
24942518
use crate::lint;
24952519
use crate::options::WasiExecModel;
@@ -2583,6 +2607,7 @@ crate mod dep_tracking {
25832607
OutputType,
25842608
RealFileName,
25852609
LocationDetail,
2610+
BranchProtection,
25862611
);
25872612

25882613
impl<T1, T2> DepTrackingHash for (T1, T2)

Diff for: ‎compiler/rustc_session/src/options.rs

+31
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,8 @@ mod desc {
389389
pub const parse_gcc_ld: &str = "one of: no value, `lld`";
390390
pub const parse_stack_protector: &str =
391391
"one of (`none` (default), `basic`, `strong`, or `all`)";
392+
pub const parse_branch_protection: &str =
393+
"a `+` separated combination of `bti`, `b-key`, `pac-ret`, or `leaf`";
392394
}
393395

394396
mod parse {
@@ -929,6 +931,33 @@ mod parse {
929931
}
930932
true
931933
}
934+
935+
crate fn parse_branch_protection(slot: &mut BranchProtection, v: Option<&str>) -> bool {
936+
match v {
937+
Some(s) => {
938+
for opt in s.split('+') {
939+
match opt {
940+
"bti" => slot.bti = true,
941+
"pac-ret" if slot.pac_ret.is_none() => {
942+
slot.pac_ret = Some(PacRet { leaf: false, key: PAuthKey::A })
943+
}
944+
"leaf" => match slot.pac_ret.as_mut() {
945+
Some(pac) => pac.leaf = true,
946+
_ => return false,
947+
},
948+
"b-key" => match slot.pac_ret.as_mut() {
949+
Some(pac) => pac.key = PAuthKey::B,
950+
_ => return false,
951+
},
952+
_ => return false,
953+
};
954+
}
955+
}
956+
957+
_ => return false,
958+
}
959+
true
960+
}
932961
}
933962

934963
options! {
@@ -942,6 +971,8 @@ options! {
942971

943972
ar: String = (String::new(), parse_string, [UNTRACKED],
944973
"this option is deprecated and does nothing"),
974+
branch_protection: BranchProtection = (BranchProtection::default(), parse_branch_protection, [TRACKED],
975+
"set options for branch target identification and pointer authentication on AArch64"),
945976
code_model: Option<CodeModel> = (None, parse_code_model, [TRACKED],
946977
"choose the code model to use (`rustc --print code-models` for details)"),
947978
codegen_units: Option<usize> = (None, parse_opt_number, [UNTRACKED],

Diff for: ‎src/doc/rustc/src/codegen-options/index.md

+23
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,29 @@ a version of this list for your exact compiler by running `rustc -C help`.
77

88
This option is deprecated and does nothing.
99

10+
## branch-protection
11+
12+
This option lets you enable branch authentication instructions on AArch64.
13+
This option is ignored for non-AArch64 architectures.
14+
It takes some combination of the following values, separated by a `+`.
15+
16+
- `pac-ret` - Enable pointer authentication for non-leaf functions.
17+
- `leaf` - Enable pointer authentication for all functions, including leaf functions.
18+
- `b-key` - Sign return addresses with key B, instead of the default key A.
19+
- `bti` - Enable branch target identification.
20+
21+
`leaf` and `b-key` are only valid if `pac-ret` was previously specified.
22+
For example, `-C branch-protection=bti+pac-ret+leaf` is valid, but
23+
`-C branch-protection=bti+leaf+pac-ret` is not.
24+
25+
Repeated values are ignored.
26+
For example, `-C branch-protection=pac-ret+leaf+pac-ret` is equivalent to
27+
`-C branch-protection=pac-ret+leaf`.
28+
29+
Rust's standard library does not ship with BTI or pointer authentication enabled by default. \
30+
In Cargo projects the standard library can be recompiled with pointer authentication using the nightly
31+
[build-std](https://doc.rust-lang.org/nightly/cargo/reference/unstable.html#build-std) feature.
32+
1033
## code-model
1134

1235
This option lets you choose which code model to use. \

Diff for: ‎src/test/assembly/aarch64-pointer-auth.rs

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
// Test that PAC instructions are emitted when branch-protection is specified.
2+
3+
// min-llvm-version: 10.0.1
4+
// assembly-output: emit-asm
5+
// compile-flags: --target aarch64-unknown-linux-gnu
6+
// compile-flags: -C branch-protection=pac-ret+leaf
7+
// needs-llvm-components: aarch64
8+
9+
#![feature(no_core, lang_items)]
10+
#![no_std]
11+
#![no_core]
12+
#![crate_type = "lib"]
13+
14+
#[lang = "sized"]
15+
trait Sized {}
16+
17+
// CHECK: hint #25
18+
// CHECK: hint #29
19+
#[no_mangle]
20+
pub fn test() -> u8 {
21+
42
22+
}

Diff for: ‎src/test/codegen/branch-protection.rs

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
// Test that the correct module flags are emitted with different branch protection flags.
2+
3+
// revisions: bti pac-ret leaf b-key
4+
// min-llvm-version: 12.0.0
5+
// needs-llvm-components: aarch64
6+
// [bti] compile-flags: -C branch-protection=bti
7+
// [pac-ret] compile-flags: -C branch-protection=pac-ret
8+
// [leaf] compile-flags: -C branch-protection=pac-ret+leaf
9+
// [b-key] compile-flags: -C branch-protection=pac-ret+b-key
10+
// compile-flags: --target aarch64-unknown-linux-gnu
11+
12+
#![crate_type = "lib"]
13+
#![feature(no_core, lang_items)]
14+
#![no_core]
15+
16+
#[lang="sized"]
17+
trait Sized { }
18+
19+
// A basic test function.
20+
pub fn test() {
21+
}
22+
23+
// bti: !"branch-target-enforcement", i32 1
24+
// bti: !"sign-return-address", i32 0
25+
// bti: !"sign-return-address-all", i32 0
26+
// bti: !"sign-return-address-with-bkey", i32 0
27+
28+
// pac-ret: !"branch-target-enforcement", i32 0
29+
// pac-ret: !"sign-return-address", i32 1
30+
// pac-ret: !"sign-return-address-all", i32 0
31+
// pac-ret: !"sign-return-address-with-bkey", i32 0
32+
33+
// leaf: !"branch-target-enforcement", i32 0
34+
// leaf: !"sign-return-address", i32 1
35+
// leaf: !"sign-return-address-all", i32 1
36+
// leaf: !"sign-return-address-with-bkey", i32 0
37+
38+
// b-key: !"branch-target-enforcement", i32 0
39+
// b-key: !"sign-return-address", i32 1
40+
// b-key: !"sign-return-address-all", i32 0
41+
// b-key: !"sign-return-address-with-bkey", i32 1

Diff for: ‎src/test/run-make-fulldeps/pointer-auth-link-with-c/Makefile

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
-include ../tools.mk
2+
3+
# only-aarch64
4+
5+
all:
6+
$(COMPILE_OBJ) $(TMPDIR)/test.o test.c
7+
$(AR) rcs $(TMPDIR)/libtest.a $(TMPDIR)/test.o
8+
$(RUSTC) -C branch-protection=bti+pac-ret+leaf test.rs
9+
$(call RUN,test)
10+
11+
$(COMPILE_OBJ) $(TMPDIR)/test.o test.c -mbranch-protection=bti+pac-ret+leaf
12+
$(AR) rcs $(TMPDIR)/libtest.a $(TMPDIR)/test.o
13+
$(RUSTC) -C branch-protection=bti+pac-ret+leaf test.rs
14+
$(call RUN,test)

Diff for: ‎src/test/run-make-fulldeps/pointer-auth-link-with-c/test.c

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
int foo() { return 0; }

Diff for: ‎src/test/run-make-fulldeps/pointer-auth-link-with-c/test.rs

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#[link(name = "test")]
2+
extern "C" {
3+
fn foo() -> i32;
4+
}
5+
6+
fn main() {
7+
unsafe {foo();}
8+
}

0 commit comments

Comments
 (0)