Add pauth attributes and first IR tests

The set of supported attributes is:
function
* "aarch64-jump-table-hardening"
* "ptrauth-auth-traps"
* "ptrauth-calls"
* "ptrauth-indirect-gotos"
* "ptrauth-returns"
module
* "ptrauth-elf-got"
* "ptrauth-sign-personality"

Author: jchlanda

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -9,9 +9,10 @@ use rustc_middle::ty::{self, TyCtxt};
 use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
 use rustc_span::sym;
 use rustc_symbol_mangling::mangle_internal_symbol;
-use rustc_target::spec::{Arch, FramePointer, SanitizerSet, StackProbeType, StackProtector};
+use rustc_target::spec::{Arch, Env, FramePointer, SanitizerSet, StackProbeType, StackProtector};
 use smallvec::SmallVec;
 
+use crate::common::pauth_fn_attrs;
 use crate::context::SimpleCx;
 use crate::errors::{PackedStackBackchainNeedsSoftfloat, SanitizerMemtagRequiresMte};
 use crate::llvm::AttributePlace::Function;
@@ -605,6 +606,12 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
         }
     }
 
+    if sess.target.env == Env::Pauthtest {
+        for &ptrauth_attr in pauth_fn_attrs() {
+            to_add.push(llvm::CreateAttrString(cx.llcx, ptrauth_attr));
+        }
+    }
+
     to_add.extend(target_features_attr(cx, tcx, function_features));
 
     attributes::apply_to_llfn(llfn, Function, &to_add);

compiler/rustc_codegen_llvm/src/base.rs

@@ -125,7 +125,14 @@ pub(crate) fn compile_codegen_unit(
             if let Some(entry) =
                 maybe_create_entry_wrapper::<Builder<'_, '_, '_>>(&cx, cx.codegen_unit)
             {
-                let attrs = attributes::sanitize_attrs(&cx, tcx, SanitizerFnAttrs::default());
+                let mut attrs = attributes::sanitize_attrs(&cx, tcx, SanitizerFnAttrs::default());
+                // For pauthtest make sure that the ptrauth-* attributes are also attached to the
+                // entry wrapper.
+                if cx.sess().target.env == Env::Pauthtest {
+                    for &ptrauth_attr in pauth_fn_attrs() {
+                        attrs.push(llvm::CreateAttrString(cx.llcx, ptrauth_attr));
+                    }
+                }
                 attributes::apply_to_llfn(entry, llvm::AttributePlace::Function, &attrs);
             }
 
@@ -142,6 +149,19 @@ pub(crate) fn compile_codegen_unit(
                 cx.add_objc_module_flags();
             }
 
+            if cx.sess().target.env == Env::Pauthtest {
+                // FIXME(jchlanda): In LLVM/Clang, there also `aarch64-elf-pauthabi-platform` and
+                // `aarch64-elf-pauthabi-version` module flags, ending up in PAuth core info
+                // emitted in a special section in resulting ELF. This is used by static and
+                // dynamic linker to check binary compatibility. By default, absence of that info
+                // is treated as being compatible with anything. While not needed now, it will
+                // become relevant when C++ interop is implemented.
+                if cx.sess().opts.unstable_opts.pauth_enable_elf_got {
+                    cx.add_ptrauth_elf_got_flag();
+                }
+                cx.add_ptrauth_sign_personality_flag();
+            }
+
             // Finalize code coverage by injecting the coverage map. Note, the coverage map will
             // also be added to the `llvm.compiler.used` variable, created next.
             if cx.sess().instrument_coverage() {

compiler/rustc_codegen_llvm/src/context.rs

@@ -688,6 +688,24 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         }
     }
 
+    pub(crate) fn add_ptrauth_elf_got_flag(&self) {
+        llvm::add_module_flag_u32(
+            self.llmod,
+            llvm::ModuleFlagMergeBehavior::Max,
+            "ptrauth-elf-got",
+            1,
+        );
+    }
+
+    pub(crate) fn add_ptrauth_sign_personality_flag(&self) {
+        llvm::add_module_flag_u32(
+            self.llmod,
+            llvm::ModuleFlagMergeBehavior::Max,
+            "ptrauth-sign-personality",
+            1,
+        );
+    }
+
     // We do our best here to match what Clang does when compiling Objective-C natively.
     // See Clang's `CGObjCCommonMac::EmitImageInfo`:
     // https://github.com/llvm/llvm-project/blob/llvmorg-20.1.8/clang/lib/CodeGen/CGObjCMac.cpp#L5085

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -23,7 +23,7 @@ use rustc_session::config::CrateType;
 use rustc_span::{Span, Symbol, sym};
 use rustc_symbol_mangling::{mangle_internal_symbol, symbol_name_for_instance_in_crate};
 use rustc_target::callconv::PassMode;
-use rustc_target::spec::Os;
+use rustc_target::spec::{Env, Os};
 use tracing::debug;
 
 use crate::abi::FnAbiLlvmExt;
@@ -32,12 +32,13 @@ use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::{
     OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload,
 };
+use crate::common::pauth_fn_attrs;
 use crate::context::CodegenCx;
 use crate::declare::declare_raw_fn;
 use crate::errors::{
     AutoDiffWithoutEnable, AutoDiffWithoutLto, OffloadWithoutEnable, OffloadWithoutFatLTO,
 };
-use crate::llvm::{self, Type, Value};
+use crate::llvm::{self, Attribute, AttributePlace, Type, Value};
 use crate::type_of::LayoutLlvmExt;
 use crate::va_arg::emit_va_arg;
 
@@ -1312,6 +1313,13 @@ fn get_rust_try_fn<'a, 'll, 'tcx>(
         ExternAbi::Rust,
     ));
     let rust_try = gen_fn(cx, "__rust_try", rust_fn_sig, codegen);
+    if cx.sess().target.env == Env::Pauthtest {
+        let attrs: Vec<&Attribute> =
+            pauth_fn_attrs().iter().map(|name| llvm::CreateAttrString(cx.llcx, name)).collect();
+        let (_ty, rust_try_fn) = rust_try;
+        crate::attributes::apply_to_llfn(rust_try_fn, AttributePlace::Function, &attrs);
+    }
+
     cx.rust_try_fn.set(Some(rust_try));
     rust_try
 }

tests/codegen-llvm/pauth-attr-special-funcs.rs

@@ -0,0 +1,31 @@
+//@ only-aarch64-unknown-linux-pauthtest
+//@ compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=0
+// Make sure that compiler generated functions (main wrapper and __rust_try) also have ptrauth
+// attributes set correctly. Rustc only generates __rust_try at O0, so use that opt level for the
+// test.
+
+//@ needs-llvm-components: aarch64
+
+use std::panic;
+
+// CHECK: define {{.*}} @__rust_try{{.*}} [[ATTR_TRY:#[0-9]+]]
+// CHECK: define {{.*}} @main{{.*}} [[ATTR_MAIN:#[0-9]+]]
+
+// CHECK: attributes [[ATTR_TRY]] = { {{.*}}"aarch64-jump-table-hardening"
+// CHECK-DAG: "ptrauth-auth-traps"
+// CHECK-DAG: "ptrauth-calls"
+// CHECK-DAG: "ptrauth-indirect-gotos"
+// CHECK-DAG: "ptrauth-returns"
+
+// CHECK: attributes [[ATTR_MAIN]] = { {{.*}}"aarch64-jump-table-hardening"
+// CHECK-DAG: "ptrauth-auth-traps"
+// CHECK-DAG: "ptrauth-calls"
+// CHECK-DAG: "ptrauth-indirect-gotos"
+// CHECK-DAG: "ptrauth-returns"
+fn main() {
+    let _ = panic::catch_unwind(|| {
+        panic!("BOOM");
+    });
+}
+
+// CHECK: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}

tests/codegen-llvm/pauth-extern-c-direct-indirect-call.rs

@@ -0,0 +1,87 @@
+// ignore-tidy-linelength
+//@ only-aarch64-unknown-linux-pauthtest
+//@ revisions: O0_PAUTH O3_PAUTH
+
+//@ [O0_PAUTH] needs-llvm-components: aarch64
+//@ [O0_PAUTH] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=0
+//@ [O3_PAUTH] needs-llvm-components: aarch64
+//@ [O3_PAUTH] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=3
+
+// Make sure that direct extern "C" calls are not handled by pointer authentication operand bundle
+// logic.
+use std::ffi::c_void;
+use std::hint::black_box;
+
+extern "C" {
+    fn rand() -> i32;
+    fn add(a: i32, b: i32) -> i32;
+    fn sub(a: i32, b: i32) -> i32;
+
+    // Corresponds to: void *woof;
+    static mut woof: *mut c_void;
+    fn direct_function_taking_void_arg(data: *mut c_void);
+    fn direct_no_arg();
+    fn direct_function_taking_fp_arg(func: unsafe extern "C" fn());
+}
+
+type CFnPtr = unsafe extern "C" fn(i32, i32) -> i32;
+
+// CHECK-LABE: test_indirect_call
+#[inline(never)]
+fn test_indirect_call() {
+    let fp_add: CFnPtr = black_box(add);
+    let fp_sub: CFnPtr = black_box(sub);
+
+    unsafe {
+        // O0_PAUTH: call {{.*}}i32 %fp_add({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32 %fp_add({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        let _id1 = fp_add(7, 4);
+        // O0_PAUTH: call {{.*}}i32 %fp_sub({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32 %fp_sub({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        let _id2 = fp_sub(10, 6);
+    }
+
+    // Also test calling via conditional pointer
+    unsafe {
+        // O0_PAUTH: call {{.*}}i32 ptrauth (ptr @rand, i32 0)({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32  @rand() #
+        let use_add = rand() % 2 == 0;
+        // O0_PAUTH: store ptr ptrauth (ptr @sub, i32 0), ptr %[[FP_O0:[a-zA-Z0-9_.]+]]
+        // O0_PAUTH: store ptr ptrauth (ptr @add, i32 0), ptr %[[FP_O0]]{{.*}}
+        // O0_PAUTH: %[[LOAD_FP_O0:[a-zA-Z0-9_.]+]] = load ptr, ptr %[[FP_O0]]{{.*}}
+        // O3_PAUTH: %[[FP_O3:.*]] = select i1 %{{.*}}, ptr ptrauth (ptr @add, i32 0), ptr ptrauth (ptr @sub, i32 0)
+        let fp: CFnPtr = if use_add { add } else { sub };
+        // O0_PAUTH: call i32 %[[LOAD_FP_O0]](i32 1, i32 2) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32 %[[FP_O3]](i32 noundef 1, i32 noundef 2) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        let _id3 = fp(1, 2);
+    }
+
+    unsafe {
+        direct_function_taking_fp_arg(direct_no_arg);
+    }
+}
+
+// CHECK-LABE: test_direct_call
+#[inline(never)]
+fn test_direct_call() {
+    unsafe {
+        // O0_PAUTH: call {{.*}}i32 ptrauth (ptr @add, i32 0)({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32 @add(i32 {{.*}}2, i32 {{.*}}3) #
+        let _d1 = add(2, 3);
+        // O0_PAUTH: call {{.*}}i32 ptrauth (ptr @sub, i32 0)({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: call {{.*}}i32 @sub(i32 {{.*}}5, i32 {{.*}}1) #
+        let _d2 = sub(5, 1);
+
+        // O0_PAUTH: call {{.*}}void ptrauth (ptr @direct_function_taking_void_arg, i32 0)({{.*}}) #{{[0-9]+}} [ "ptrauth"(i32 0, i64 0) ]
+        // O3_PAUTH: {{(tail )?}}call void @direct_function_taking_void_arg(ptr noundef %{{.*}}) #
+        direct_function_taking_void_arg(woof);
+    }
+}
+
+fn main() {
+    test_indirect_call();
+    test_direct_call();
+}
+
+// O0_PAUTH: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}
+// O3_PAUTH: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}

tests/codegen-llvm/pauth-extern-c.rs

@@ -0,0 +1,71 @@
+// ignore-tidy-linelength
+//@ only-aarch64-unknown-linux-pauthtest
+//@ revisions: O0_PAUTH O3_PAUTH O0_PAUTH-ELF-GOT O3_PAUTH-ELF-GOT O0_NO_PAUTH O3_NO_PAUTH
+
+//@ [O0_PAUTH] needs-llvm-components: aarch64
+//@ [O0_PAUTH] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=0
+//@ [O3_PAUTH] needs-llvm-components: aarch64
+//@ [O3_PAUTH] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=3
+//@ [O0_PAUTH-ELF-GOT] needs-llvm-components: aarch64
+//@ [O0_PAUTH-ELF-GOT] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=0 -Z pauth_enable_elf_got
+//@ [O3_PAUTH-ELF-GOT] needs-llvm-components: aarch64
+//@ [O3_PAUTH-ELF-GOT] compile-flags: --target=aarch64-unknown-linux-pauthtest -C opt-level=3 -Z pauth_enable_elf_got
+//@ [O0_NO_PAUTH] needs-llvm-components: aarch64
+//@ [O0_NO_PAUTH] compile-flags: --target=aarch64-unknown-linux-gnu -C opt-level=0
+//@ [O3_NO_PAUTH] needs-llvm-components: aarch64
+//@ [O3_NO_PAUTH] compile-flags: --target=aarch64-unknown-linux-gnu -C opt-level=3
+
+use std::hint::black_box;
+
+type FnPtr = unsafe extern "C" fn(i32, i32) -> i32;
+
+// O0_NO_PAUTH-NOT: "ptrauth"(i32
+// O3_NO_PAUTH-NOT: "ptrauth"(i32
+
+// O0_PAUTH: define {{.*}}pauth_extern_c4main
+// O3_PAUTH: define {{.*}}pauth_extern_c4main
+fn main() {
+    // O0_PAUTH: ptr ptrauth (ptr @add_from_c, i32 0)
+    // O3_PAUTH: ptr ptrauth (ptr @add_from_c, i32 0)
+    let add_ptr: FnPtr = black_box(add_from_c);
+    // O0_PAUTH: call i32 @{{.*}}pauth_extern_c7call_it{{.*}}(ptr {{.*}})
+    let _sum = call_it(add_ptr, 5, 7);
+    assert!(12 == _sum);
+}
+
+// O0_PAUTH: define {{.*}}pauth_extern_c7call_it{{.*}} #[[ATTR_O0_1:[0-9]+]]
+// O3_PAUTH: define {{.*}}pauth_extern_c7call_it{{.*}} #[[ATTR_O3_1:[0-9]+]]
+#[inline(never)]
+fn call_it(fn_ptr: FnPtr, arg_1: i32, arg_2: i32) -> i32 {
+    // O0_PAUTH: call i32 %fn_ptr(i32 %arg_1, i32 %arg_2) {{.*}} [ "ptrauth"(i32 0, i64 0) ]
+    // O3_PAUTH: call {{.*}} i32 %fn_ptr(i32 {{.*}}, i32 {{.*}}) {{.*}} [ "ptrauth"(i32 0, i64 0) ]
+    unsafe { fn_ptr(arg_1, arg_2) }
+}
+
+extern "C" {
+    fn add_from_c(a: i32, b: i32) -> i32;
+}
+
+// O0_PAUTH-CHECK: attributes #[[ATTR_O0_1]] = { {{.*}}"aarch64-jump-table-hardening"
+// O0_PAUTH-CHECK-DAG: "ptrauth-auth-traps"
+// O0_PAUTH-CHECK-DAG: "ptrauth-calls"
+// O0_PAUTH-CHECK-DAG: "ptrauth-indirect-gotos"
+// O0_PAUTH-CHECK-DAG: "ptrauth-returns"
+
+// O3_PAUTH-CHECK: attributes #[[ATTR_O3_1]] = { {{.*}}"aarch64-jump-table-hardening"
+// O3_PAUTH-CHECK-DAG: "ptrauth-auth-traps"
+// O3_PAUTH-CHECK-DAG: "ptrauth-calls"
+// O3_PAUTH-CHECK-DAG: "ptrauth-indirect-gotos"
+// O3_PAUTH-CHECK-DAG: "ptrauth-returns"
+
+// O0_PAUTH-ELF-GOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O0_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O0_PAUTH: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}
+// O3_PAUTH-ELF-GOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O3_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O3_PAUTH: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}
+
+// O0_NO_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O0_NO_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}
+// O3_NO_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-elf-got", i32 1}
+// O3_NO_PAUTH-NOT: !{{[0-9]+}} = !{i32 7, !"ptrauth-sign-personality", i32 1}