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

Author: Jmc18134

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -9,7 +9,7 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::ty::layout::HasTyCtxt;
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_session::config::OptLevel;
+use rustc_session::config::{BranchProtection, OptLevel, PAuthKey};
 use rustc_session::Session;
 use rustc_target::spec::abi::Abi;
 use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
@@ -203,6 +203,58 @@ pub fn non_lazy_bind(sess: &Session, llfn: &'ll Value) {
     }
 }
 
+pub fn set_branch_protection(sess: &Session, llfn: &'ll Value) {
+    // Setting PAC/BTI function attributes is only necessary for LLVM 11 and earlier.
+    // For LLVM 12 and greater, module-level metadata attributes are set in
+    // `compiler/rustc_codegen_llvm/src/context.rs`.
+    if llvm_util::get_version() >= (12, 0, 0) {
+        return;
+    }
+
+    let BranchProtection { bti, pac_ret: pac } = sess.opts.cg.branch_protection;
+
+    if bti {
+        llvm::AddFunctionAttrString(
+            llfn,
+            llvm::AttributePlace::Function,
+            cstr!("branch-target-enforcement"),
+        );
+    }
+
+    if let Some(pac_opts) = pac {
+        if pac_opts.leaf {
+            llvm::AddFunctionAttrStringValue(
+                llfn,
+                llvm::AttributePlace::Function,
+                cstr!("sign-return-address"),
+                cstr!("non-leaf"),
+            );
+        } else {
+            llvm::AddFunctionAttrStringValue(
+                llfn,
+                llvm::AttributePlace::Function,
+                cstr!("sign-return-address"),
+                cstr!("all"),
+            );
+        }
+
+        match pac_opts.key {
+            PAuthKey::A => llvm::AddFunctionAttrStringValue(
+                llfn,
+                llvm::AttributePlace::Function,
+                cstr!("sign-return-address-key"),
+                cstr!("a_key"),
+            ),
+            PAuthKey::B => llvm::AddFunctionAttrStringValue(
+                llfn,
+                llvm::AttributePlace::Function,
+                cstr!("sign-return-address-key"),
+                cstr!("b_key"),
+            ),
+        }
+    }
+}
+
 pub(crate) fn default_optimisation_attrs(sess: &Session, llfn: &'ll Value) {
     match sess.opts.optimize {
         OptLevel::Size => {

compiler/rustc_codegen_llvm/src/context.rs

@@ -21,7 +21,7 @@ use rustc_middle::ty::layout::{
 };
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
-use rustc_session::config::{CFGuard, CrateType, DebugInfo};
+use rustc_session::config::{BranchProtection, CFGuard, CrateType, DebugInfo, PAuthKey};
 use rustc_session::Session;
 use rustc_span::source_map::Span;
 use rustc_span::symbol::Symbol;
@@ -242,6 +242,38 @@ pub unsafe fn create_module(
         }
     }
 
+    if sess.target.arch == "aarch64" {
+        let BranchProtection { bti, pac_ret: pac } = sess.opts.cg.branch_protection;
+
+        llvm::LLVMRustAddModuleFlag(
+            llmod,
+            "branch-target-enforcement\0".as_ptr().cast(),
+            bti.into(),
+        );
+
+        if let Some(pac_opts) = pac {
+            llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address\0".as_ptr().cast(), 1);
+            llvm::LLVMRustAddModuleFlag(
+                llmod,
+                "sign-return-address-all\0".as_ptr().cast(),
+                pac_opts.leaf.into(),
+            );
+            llvm::LLVMRustAddModuleFlag(
+                llmod,
+                "sign-return-address-with-bkey\0".as_ptr().cast(),
+                if pac_opts.key == PAuthKey::A { 0 } else { 1 },
+            );
+        } else {
+            llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address\0".as_ptr().cast(), 0);
+            llvm::LLVMRustAddModuleFlag(llmod, "sign-return-address-all\0".as_ptr().cast(), 0);
+            llvm::LLVMRustAddModuleFlag(
+                llmod,
+                "sign-return-address-with-bkey\0".as_ptr().cast(),
+                0,
+            );
+        }
+    }
+
     llmod
 }
 

compiler/rustc_codegen_llvm/src/declare.rs

@@ -45,8 +45,13 @@ fn declare_raw_fn(
         llvm::Attribute::NoRedZone.apply_llfn(Function, llfn);
     }
 
+    if cx.tcx.sess.target.arch == "aarch64" {
+        attributes::set_branch_protection(cx.tcx.sess, llfn);
+    }
+
     attributes::default_optimisation_attrs(cx.tcx.sess, llfn);
     attributes::non_lazy_bind(cx.sess(), llfn);
+
     llfn
 }
 

compiler/rustc_interface/src/tests.rs

@@ -8,10 +8,11 @@ use rustc_session::config::{build_configuration, build_session_options, to_crate
 use rustc_session::config::{
     rustc_optgroups, ErrorOutputType, ExternLocation, LocationDetail, Options, Passes,
 };
-use rustc_session::config::{CFGuard, ExternEntry, LinkerPluginLto, LtoCli, SwitchWithOptPath};
 use rustc_session::config::{
-    Externs, OutputType, OutputTypes, SymbolManglingVersion, WasiExecModel,
+    BranchProtection, Externs, OutputType, OutputTypes, PAuthKey, PacRet, SymbolManglingVersion,
+    WasiExecModel,
 };
+use rustc_session::config::{CFGuard, ExternEntry, LinkerPluginLto, LtoCli, SwitchWithOptPath};
 use rustc_session::lint::Level;
 use rustc_session::search_paths::SearchPath;
 use rustc_session::utils::{CanonicalizedPath, NativeLib, NativeLibKind};
@@ -566,6 +567,10 @@ fn test_codegen_options_tracking_hash() {
 
     // Make sure that changing a [TRACKED] option changes the hash.
     // This list is in alphabetical order.
+    tracked!(
+        branch_protection,
+        BranchProtection { bti: true, pac_ret: Some(PacRet { leaf: true, key: PAuthKey::B }) }
+    );
     tracked!(code_model, Some(CodeModel::Large));
     tracked!(control_flow_guard, CFGuard::Checks);
     tracked!(debug_assertions, Some(true));

compiler/rustc_session/src/config.rs

@@ -842,6 +842,30 @@ impl Passes {
     }
 }
 
+#[derive(Clone, Copy, Hash, Debug, PartialEq)]
+pub enum PAuthKey {
+    A,
+    B,
+}
+
+#[derive(Clone, Copy, Hash, Debug, PartialEq)]
+pub struct PacRet {
+    pub leaf: bool,
+    pub key: PAuthKey,
+}
+
+#[derive(Clone, Copy, Hash, Debug, PartialEq)]
+pub struct BranchProtection {
+    pub bti: bool,
+    pub pac_ret: Option<PacRet>,
+}
+
+impl Default for BranchProtection {
+    fn default() -> Self {
+        BranchProtection { bti: false, pac_ret: None }
+    }
+}
+
 pub const fn default_lib_output() -> CrateType {
     CrateType::Rlib
 }
@@ -2487,9 +2511,9 @@ impl PpMode {
 crate mod dep_tracking {
     use super::LdImpl;
     use super::{
-        CFGuard, CrateType, DebugInfo, ErrorOutputType, InstrumentCoverage, LinkerPluginLto,
-        LocationDetail, LtoCli, OptLevel, OutputType, OutputTypes, Passes, SourceFileHashAlgorithm,
-        SwitchWithOptPath, SymbolManglingVersion, TrimmedDefPaths,
+        BranchProtection, CFGuard, CrateType, DebugInfo, ErrorOutputType, InstrumentCoverage,
+        LinkerPluginLto, LocationDetail, LtoCli, OptLevel, OutputType, OutputTypes, Passes,
+        SourceFileHashAlgorithm, SwitchWithOptPath, SymbolManglingVersion, TrimmedDefPaths,
     };
     use crate::lint;
     use crate::options::WasiExecModel;
@@ -2583,6 +2607,7 @@ crate mod dep_tracking {
         OutputType,
         RealFileName,
         LocationDetail,
+        BranchProtection,
     );
 
     impl<T1, T2> DepTrackingHash for (T1, T2)

compiler/rustc_session/src/options.rs

@@ -389,6 +389,8 @@ mod desc {
     pub const parse_gcc_ld: &str = "one of: no value, `lld`";
     pub const parse_stack_protector: &str =
         "one of (`none` (default), `basic`, `strong`, or `all`)";
+    pub const parse_branch_protection: &str =
+        "a `+` separated combination of `bti`, `b-key`, `pac-ret`, or `leaf`";
 }
 
 mod parse {
@@ -929,6 +931,33 @@ mod parse {
         }
         true
     }
+
+    crate fn parse_branch_protection(slot: &mut BranchProtection, v: Option<&str>) -> bool {
+        match v {
+            Some(s) => {
+                for opt in s.split('+') {
+                    match opt {
+                        "bti" => slot.bti = true,
+                        "pac-ret" if slot.pac_ret.is_none() => {
+                            slot.pac_ret = Some(PacRet { leaf: false, key: PAuthKey::A })
+                        }
+                        "leaf" => match slot.pac_ret.as_mut() {
+                            Some(pac) => pac.leaf = true,
+                            _ => return false,
+                        },
+                        "b-key" => match slot.pac_ret.as_mut() {
+                            Some(pac) => pac.key = PAuthKey::B,
+                            _ => return false,
+                        },
+                        _ => return false,
+                    };
+                }
+            }
+
+            _ => return false,
+        }
+        true
+    }
 }
 
 options! {
@@ -942,6 +971,8 @@ options! {
 
     ar: String = (String::new(), parse_string, [UNTRACKED],
         "this option is deprecated and does nothing"),
+    branch_protection: BranchProtection = (BranchProtection::default(), parse_branch_protection, [TRACKED],
+        "set options for branch target identification and pointer authentication on AArch64"),
     code_model: Option<CodeModel> = (None, parse_code_model, [TRACKED],
         "choose the code model to use (`rustc --print code-models` for details)"),
     codegen_units: Option<usize> = (None, parse_opt_number, [UNTRACKED],

src/doc/rustc/src/codegen-options/index.md

@@ -7,6 +7,29 @@ a version of this list for your exact compiler by running `rustc -C help`.
 
 This option is deprecated and does nothing.
 
+## branch-protection
+
+This option lets you enable branch authentication instructions on AArch64.
+This option is ignored for non-AArch64 architectures.
+It takes some combination of the following values, separated by a `+`.
+
+- `pac-ret` - Enable pointer authentication for non-leaf functions.
+- `leaf` - Enable pointer authentication for all functions, including leaf functions.
+- `b-key` - Sign return addresses with key B, instead of the default key A.
+- `bti` - Enable branch target identification.
+
+`leaf` and `b-key` are only valid if `pac-ret` was previously specified.
+For example, `-C branch-protection=bti+pac-ret+leaf` is valid, but
+`-C branch-protection=bti+leaf+pac-ret` is not.
+
+Repeated values are ignored.
+For example, `-C branch-protection=pac-ret+leaf+pac-ret` is equivalent to
+`-C branch-protection=pac-ret+leaf`.
+
+Rust's standard library does not ship with BTI or pointer authentication enabled by default. \
+In Cargo projects the standard library can be recompiled with pointer authentication using the nightly
+[build-std](https://doc.rust-lang.org/nightly/cargo/reference/unstable.html#build-std) feature.
+
 ## code-model
 
 This option lets you choose which code model to use. \

src/test/assembly/aarch64-pointer-auth.rs

@@ -0,0 +1,22 @@
+// Test that PAC instructions are emitted when branch-protection is specified.
+
+// min-llvm-version: 10.0.1
+// assembly-output: emit-asm
+// compile-flags: --target aarch64-unknown-linux-gnu
+// compile-flags: -C branch-protection=pac-ret+leaf
+// needs-llvm-components: aarch64
+
+#![feature(no_core, lang_items)]
+#![no_std]
+#![no_core]
+#![crate_type = "lib"]
+
+#[lang = "sized"]
+trait Sized {}
+
+// CHECK: hint #25
+// CHECK: hint #29
+#[no_mangle]
+pub fn test() -> u8 {
+    42
+}

src/test/codegen/branch-protection.rs

@@ -0,0 +1,41 @@
+// Test that the correct module flags are emitted with different branch protection flags.
+
+// revisions: bti pac-ret leaf b-key
+// min-llvm-version: 12.0.0
+// needs-llvm-components: aarch64
+// [bti] compile-flags: -C branch-protection=bti
+// [pac-ret] compile-flags: -C branch-protection=pac-ret
+// [leaf] compile-flags: -C branch-protection=pac-ret+leaf
+// [b-key] compile-flags: -C branch-protection=pac-ret+b-key
+// compile-flags: --target aarch64-unknown-linux-gnu
+
+#![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
+
+// pac-ret: !"branch-target-enforcement", i32 0
+// pac-ret: !"sign-return-address", i32 1
+// pac-ret: !"sign-return-address-all", i32 0
+// pac-ret: !"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
+
+// b-key: !"branch-target-enforcement", i32 0
+// b-key: !"sign-return-address", i32 1
+// b-key: !"sign-return-address-all", i32 0
+// b-key: !"sign-return-address-with-bkey", i32 1

src/test/run-make-fulldeps/pointer-auth-link-with-c/Makefile

@@ -0,0 +1,14 @@
+-include ../tools.mk
+
+# only-aarch64
+
+all:
+	$(COMPILE_OBJ) $(TMPDIR)/test.o test.c
+	$(AR) rcs $(TMPDIR)/libtest.a $(TMPDIR)/test.o
+	$(RUSTC) -C branch-protection=bti+pac-ret+leaf test.rs
+	$(call RUN,test)
+
+	$(COMPILE_OBJ) $(TMPDIR)/test.o test.c -mbranch-protection=bti+pac-ret+leaf
+	$(AR) rcs $(TMPDIR)/libtest.a $(TMPDIR)/test.o
+	$(RUSTC) -C branch-protection=bti+pac-ret+leaf test.rs
+	$(call RUN,test)

src/test/run-make-fulldeps/pointer-auth-link-with-c/test.c

@@ -0,0 +1 @@
+int foo() { return 0; }

src/test/run-make-fulldeps/pointer-auth-link-with-c/test.rs

@@ -0,0 +1,8 @@
+#[link(name = "test")]
+extern "C" {
+    fn foo() -> i32;
+}
+
+fn main() {
+    unsafe {foo();}
+}