New signature of get_fn_addr, teach Rust how to sign fn pointers

Allow PAC metadata to be passed to `get_fn_addr` and related API
changes.

Author: jchlanda

compiler/rustc_codegen_gcc/src/common.rs

@@ -2,7 +2,8 @@ use gccjit::{LValue, RValue, ToRValue, Type};
 use rustc_abi::Primitive::Pointer;
 use rustc_abi::{self as abi, HasDataLayout};
 use rustc_codegen_ssa::traits::{
-    BaseTypeCodegenMethods, ConstCodegenMethods, MiscCodegenMethods, StaticCodegenMethods,
+    BaseTypeCodegenMethods, ConstCodegenMethods, MiscCodegenMethods, PacMetadata,
+    StaticCodegenMethods,
 };
 use rustc_middle::mir::Mutability;
 use rustc_middle::mir::interpret::{GlobalAlloc, PointerArithmetic, Scalar};
@@ -225,7 +226,13 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
         None
     }
 
-    fn scalar_to_backend(&self, cv: Scalar, layout: abi::Scalar, ty: Type<'gcc>) -> RValue<'gcc> {
+    fn scalar_to_backend_with_pac(
+        &self,
+        cv: Scalar,
+        layout: abi::Scalar,
+        ty: Type<'gcc>,
+        _pac: Option<PacMetadata>,
+    ) -> RValue<'gcc> {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {
             Scalar::Int(int) => {
@@ -274,7 +281,7 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
                         }
                         value
                     }
-                    GlobalAlloc::Function { instance, .. } => self.get_fn_addr(instance),
+                    GlobalAlloc::Function { instance, .. } => self.get_fn_addr(instance, None),
                     GlobalAlloc::VTable(ty, dyn_ty) => {
                         let alloc = self
                             .tcx

compiler/rustc_codegen_gcc/src/context.rs

@@ -7,7 +7,9 @@ use gccjit::{
 use rustc_abi::{Align, HasDataLayout, PointeeInfo, Size, TargetDataLayout, VariantIdx};
 use rustc_codegen_ssa::base::wants_msvc_seh;
 use rustc_codegen_ssa::errors as ssa_errors;
-use rustc_codegen_ssa::traits::{BackendTypes, BaseTypeCodegenMethods, MiscCodegenMethods};
+use rustc_codegen_ssa::traits::{
+    BackendTypes, BaseTypeCodegenMethods, MiscCodegenMethods, PacMetadata,
+};
 use rustc_data_structures::base_n::{ALPHANUMERIC_ONLY, ToBaseN};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::mir::interpret::Allocation;
@@ -403,7 +405,7 @@ impl<'gcc, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         get_fn(self, instance)
     }
 
-    fn get_fn_addr(&self, instance: Instance<'tcx>) -> RValue<'gcc> {
+    fn get_fn_addr(&self, instance: Instance<'tcx>, _pac: Option<PacMetadata>) -> RValue<'gcc> {
         let func_name = self.tcx.symbol_name(instance).name;
 
         let func = if let Some(variable) = self.get_declared_value(func_name) {

compiler/rustc_codegen_llvm/src/base.rs

@@ -25,12 +25,13 @@ use rustc_middle::mir::mono::Visibility;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::{DebugInfo, Offload};
 use rustc_span::Symbol;
-use rustc_target::spec::SanitizerSet;
+use rustc_target::spec::{Env, SanitizerSet};
 
 use super::ModuleLlvm;
 use crate::attributes;
 use crate::builder::Builder;
 use crate::builder::gpu_offload::OffloadGlobals;
+use crate::common::pauth_fn_attrs;
 use crate::context::CodegenCx;
 use crate::llvm::{self, Value};
 

compiler/rustc_codegen_llvm/src/common.rs

@@ -4,23 +4,81 @@ use std::borrow::Borrow;
 
 use libc::{c_char, c_uint};
 use rustc_abi::Primitive::Pointer;
-use rustc_abi::{self as abi, HasDataLayout as _};
+use rustc_abi::{self as abi, ExternAbi, HasDataLayout as _};
 use rustc_ast::Mutability;
 use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_hashes::Hash128;
+use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
 use rustc_middle::bug;
 use rustc_middle::mir::interpret::{GlobalAlloc, PointerArithmetic, Scalar};
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{Instance, TyCtxt};
 use rustc_session::cstore::DllImport;
+use rustc_target::spec::Env;
 use tracing::debug;
 
-use crate::consts::const_alloc_to_llvm;
+use crate::consts::{IsInitOrFini, IsStatic, const_alloc_to_llvm};
 pub(crate) use crate::context::CodegenCx;
 use crate::context::{GenericCx, SCx};
-use crate::llvm::{self, BasicBlock, ConstantInt, FALSE, TRUE, ToLlvmBool, Type, Value};
+use crate::llvm::{
+    self, BasicBlock, ConstantInt, FALSE, TRUE, ToLlvmBool, Type, Value, const_ptr_auth,
+};
+
+#[inline]
+pub(crate) fn pauth_fn_attrs() -> &'static [&'static str] {
+    // FIXME(jchlanda) This is not an exhaustive list of all `pauthtest`-related attributes, but
+    // only those currently supported. The list is expected to grow as additional functionality is
+    // implemented, particularly for C++ interoperability.
+    &[
+        "aarch64-jump-table-hardening",
+        "ptrauth-indirect-gotos",
+        "ptrauth-calls",
+        "ptrauth-returns",
+        "ptrauth-auth-traps",
+    ]
+}
+
+pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
+    cx: &CodegenCx<'ll, '_>,
+    instance: Instance<'tcx>,
+    llfn: &'ll llvm::Value,
+    pac: PacMetadata,
+) -> &'ll llvm::Value {
+    if cx.sess().target.env != Env::Pauthtest {
+        return llfn;
+    }
+
+    // Only free functions or methods
+    let def_id = instance.def_id();
+    if !matches!(cx.tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+        return llfn;
+    }
+    // Only C ABI
+    let abi = cx.tcx.fn_sig(def_id).skip_binder().abi();
+    if !matches!(abi, ExternAbi::C { .. }) {
+        return llfn;
+    }
+    // Ignore LLVM intrinsics
+    if llvm::get_value_name(llfn).starts_with(b"llvm.") {
+        return llfn;
+    }
+    if Some(def_id) == cx.tcx.lang_items().eh_personality() {
+        return llfn;
+    }
+
+    let addr_diversity = match pac.addr_diversity {
+        AddressDiversity::None => None,
+        AddressDiversity::Real => Some(llfn),
+        AddressDiversity::Synthetic(val) => {
+            let llval = cx.const_u64(val);
+            let llty = cx.val_ty(llfn);
+            Some(unsafe { llvm::LLVMConstIntToPtr(llval, llty) })
+        }
+    };
+    const_ptr_auth(llfn, pac.key, pac.disc, addr_diversity)
+}
 
 /*
 * A note on nomenclature of linking: "extern", "foreign", and "upcall".
@@ -264,7 +322,13 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
         })
     }
 
-    fn scalar_to_backend(&self, cv: Scalar, layout: abi::Scalar, llty: &'ll Type) -> &'ll Value {
+    fn scalar_to_backend_with_pac(
+        &self,
+        cv: Scalar,
+        layout: abi::Scalar,
+        llty: &'ll Type,
+        pac: Option<PacMetadata>,
+    ) -> &'ll Value {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {
             Scalar::Int(int) => {
@@ -293,8 +357,12 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                                 self.const_bitcast(llval, llty)
                             };
                         } else {
-                            let init =
-                                const_alloc_to_llvm(self, alloc.inner(), /*static*/ false);
+                            let init = const_alloc_to_llvm(
+                                self,
+                                alloc.inner(),
+                                IsStatic::No,
+                                IsInitOrFini::No,
+                            );
                             let alloc = alloc.inner();
                             let value = match alloc.mutability {
                                 Mutability::Mut => self.static_addr_of_mut(init, alloc.align, None),
@@ -315,7 +383,7 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                             value
                         }
                     }
-                    GlobalAlloc::Function { instance, .. } => self.get_fn_addr(instance),
+                    GlobalAlloc::Function { instance, .. } => self.get_fn_addr(instance, pac),
                     GlobalAlloc::VTable(ty, dyn_ty) => {
                         let alloc = self
                             .tcx
@@ -326,7 +394,12 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                                 }),
                             )))
                             .unwrap_memory();
-                        let init = const_alloc_to_llvm(self, alloc.inner(), /*static*/ false);
+                        let init = const_alloc_to_llvm(
+                            self,
+                            alloc.inner(),
+                            IsStatic::No,
+                            IsInitOrFini::No,
+                        );
                         self.static_addr_of_impl(init, alloc.inner().align, None)
                     }
                     GlobalAlloc::Static(def_id) => {

compiler/rustc_codegen_llvm/src/consts.rs

@@ -1,6 +1,6 @@
 use std::ops::Range;
 
-use rustc_abi::{Align, HasDataLayout, Primitive, Scalar, Size, WrappingRange};
+use rustc_abi::{Align, ExternAbi, HasDataLayout, Primitive, Scalar, Size, WrappingRange};
 use rustc_codegen_ssa::common;
 use rustc_codegen_ssa::traits::*;
 use rustc_hir::LangItem;
@@ -17,19 +17,30 @@ use rustc_middle::ty::layout::{HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{self, Instance};
 use rustc_middle::{bug, span_bug};
 use rustc_span::Symbol;
-use rustc_target::spec::Arch;
+use rustc_target::spec::{Arch, Env};
 use tracing::{debug, instrument, trace};
 
 use crate::common::CodegenCx;
 use crate::errors::SymbolAlreadyDefined;
-use crate::llvm::{self, Type, Value};
+use crate::llvm::{self, Type, Value, const_ptr_auth};
 use crate::type_of::LayoutLlvmExt;
 use crate::{base, debuginfo};
 
+/// Indicates whether a value originates from a `static`.
+pub(crate) enum IsStatic {
+    Yes,
+    No,
+}
+/// Indicates whether a symbol is part of `.init_array` or `.fini_array`.
+pub(crate) enum IsInitOrFini {
+    Yes,
+    No,
+}
 pub(crate) fn const_alloc_to_llvm<'ll>(
     cx: &CodegenCx<'ll, '_>,
     alloc: &Allocation,
-    is_static: bool,
+    is_static: IsStatic,
+    is_init_fini: IsInitOrFini,
 ) -> &'ll Value {
     // We expect that callers of const_alloc_to_llvm will instead directly codegen a pointer or
     // integer for any &ZST where the ZST is a constant (i.e. not a static). We should never be
@@ -775,7 +786,7 @@ impl<'ll> StaticCodegenMethods for CodegenCx<'ll, '_> {
     fn static_addr_of(&self, alloc: ConstAllocation<'_>, kind: Option<&str>) -> &'ll Value {
         // FIXME: should we cache `const_alloc_to_llvm` to avoid repeating this for the
         // same `ConstAllocation`?
-        let cv = const_alloc_to_llvm(self, alloc.inner(), /*static*/ false);
+        let cv = const_alloc_to_llvm(self, alloc.inner(), IsStatic::No, IsInitOrFini::No);
 
         let gv = self.static_addr_of_impl(cv, alloc.inner().align, kind);
         // static_addr_of_impl returns the bare global variable, which might not be in the default

compiler/rustc_codegen_llvm/src/context.rs

@@ -832,8 +832,12 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         get_fn(self, instance)
     }
 
-    fn get_fn_addr(&self, instance: Instance<'tcx>) -> &'ll Value {
-        get_fn(self, instance)
+    fn get_fn_addr(&self, instance: Instance<'tcx>, pac: Option<PacMetadata>) -> &'ll Value {
+        let llfn = get_fn(self, instance);
+        match pac {
+            Some(pac) => common::maybe_sign_fn_ptr(self, instance, llfn, pac),
+            None => llfn,
+        }
     }
 
     fn eh_personality(&self) -> &'ll Value {
@@ -875,13 +879,16 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
 
         let tcx = self.tcx;
         let llfn = match tcx.lang_items().eh_personality() {
-            Some(def_id) if name.is_none() => self.get_fn_addr(ty::Instance::expect_resolve(
-                tcx,
-                self.typing_env(),
-                def_id,
-                ty::List::empty(),
-                DUMMY_SP,
-            )),
+            Some(def_id) if name.is_none() => self.get_fn_addr(
+                ty::Instance::expect_resolve(
+                    tcx,
+                    self.typing_env(),
+                    def_id,
+                    ty::List::empty(),
+                    DUMMY_SP,
+                ),
+                Some(PacMetadata::default()),
+            ),
             _ => {
                 let name = name.unwrap_or("rust_eh_personality");
                 if let Some(llfn) = self.get_declared_value(name) {

compiler/rustc_codegen_ssa/src/base.rs

@@ -485,7 +485,7 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         return None;
     }
 
-    let main_llfn = cx.get_fn_addr(instance);
+    let main_llfn = cx.get_fn_addr(instance, Some(PacMetadata::default()));
 
     let entry_fn = create_entry_fn::<Bx>(cx, main_llfn, main_def_id, entry_type);
     return Some(entry_fn);
@@ -545,7 +545,7 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
                 cx.tcx().mk_args(&[main_ret_ty.into()]),
                 DUMMY_SP,
             );
-            let start_fn = cx.get_fn_addr(start_instance);
+            let start_fn = cx.get_fn_addr(start_instance, Some(PacMetadata::default()));
 
             let i8_ty = cx.type_i8();
             let arg_sigpipe = bx.const_u8(sigpipe);

compiler/rustc_codegen_ssa/src/common.rs

@@ -117,7 +117,11 @@ pub(crate) fn build_langcall<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let tcx = bx.tcx();
     let def_id = tcx.require_lang_item(li, span);
     let instance = ty::Instance::mono(tcx, def_id);
-    (bx.fn_abi_of_instance(instance, ty::List::empty()), bx.get_fn_addr(instance), instance)
+    (
+        bx.fn_abi_of_instance(instance, ty::List::empty()),
+        bx.get_fn_addr(instance, Some(PacMetadata::default())),
+        instance,
+    )
 }
 
 pub(crate) fn shift_mask_val<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -657,7 +657,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
             _ => (
                 false,
-                bx.get_fn_addr(drop_fn),
+                bx.get_fn_addr(drop_fn, Some(PacMetadata::default())),
                 bx.fn_abi_of_instance(drop_fn, ty::List::empty()),
                 drop_fn,
             ),
@@ -1032,7 +1032,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         )
                         .unwrap();
 
-                        (None, Some(bx.get_fn_addr(instance)))
+                        (None, Some(bx.get_fn_addr(instance, Some(PacMetadata::default()))))
                     }
                     _ => (Some(instance), None),
                 }
@@ -1338,7 +1338,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
 
         let fn_ptr = match (instance, llfn) {
-            (Some(instance), None) => bx.get_fn_addr(instance),
+            (Some(instance), None) => bx.get_fn_addr(instance, Some(PacMetadata::default())),
             (_, Some(llfn)) => llfn,
             _ => span_bug!(fn_span, "no instance or llfn for call"),
         };

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -415,7 +415,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                     args,
                                 )
                                 .unwrap();
-                                OperandValue::Immediate(bx.get_fn_addr(instance))
+                                OperandValue::Immediate(bx.get_fn_addr(instance,
+                                        Some(PacMetadata::default()),
+                                        ))
                             }
                             _ => bug!("{} cannot be reified to a fn ptr", operand.layout.ty),
                         }
@@ -429,7 +431,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                     args,
                                     ty::ClosureKind::FnOnce,
                                 );
-                                OperandValue::Immediate(bx.cx().get_fn_addr(instance))
+                                OperandValue::Immediate(bx.cx().get_fn_addr(instance,
+                                        Some(PacMetadata::default()),
+                                        ))
                             }
                             _ => bug!("{} cannot be cast to a fn ptr", operand.layout.ty),
                         }
@@ -628,7 +632,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         def: ty::InstanceKind::ThreadLocalShim(def_id),
                         args: ty::GenericArgs::empty(),
                     };
-                    let fn_ptr = bx.get_fn_addr(instance);
+                    let fn_ptr = bx.get_fn_addr(instance, Some(PacMetadata::default()));
                     let fn_abi = bx.fn_abi_of_instance(instance, ty::List::empty());
                     let fn_ty = bx.fn_decl_backend_type(fn_abi);
                     let fn_attrs = if bx.tcx().def_kind(instance.def_id()).has_codegen_attrs() {