Unify creation of discrimination data and fill in missing get_fn_addr call sites

Author: jchlanda

compiler/rustc_codegen_llvm/src/common.rs

@@ -32,9 +32,7 @@ pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
     llfn: &'ll llvm::Value,
     schema: PointerAuthSchema,
 ) -> &'ll llvm::Value {
-    if cx.tcx.sess.pointer_authentication_functions().is_none() {
-        return llfn;
-    }
+    assert!(cx.tcx.sess.pointer_authentication_functions().is_some());
 
     // Only free functions or methods
     let def_id = instance.def_id();

compiler/rustc_codegen_llvm/src/consts.rs

@@ -15,7 +15,7 @@ use rustc_middle::mir::interpret::{
 };
 use rustc_middle::mono::MonoItem;
 use rustc_middle::ptrauth::{
-    build_fn_ptr_type_discriminator_input, compute_fn_ptr_type_discriminator,
+    build_fn_ptr_type_discriminator_input_from_ty, compute_fn_ptr_type_discriminator,
 };
 use rustc_middle::ty::layout::{HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};

compiler/rustc_codegen_llvm/src/context.rs

@@ -973,6 +973,16 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
 
         let tcx = self.tcx;
         let llfn = match tcx.lang_items().eh_personality() {
+            // We intentionally do NOT apply pointer authentication (and/or function type
+            // discriminators to the EH personality function).
+            //
+            // Although `get_fn_addr` normally produces a signed function pointer for
+            // externally-callable functions, the EH personality is not an indirect call
+            // target in the SSA sense. Instead, it is a compile-time constant attached to
+            // the Function object (via LLVM's `setPersonalityFn`) and consumed only by
+            // exception handling metadata generation (landing pads / unwind tables).
+            // LLVM never loads or invokes the personality via a function pointer value;
+            // it is not part of the program's call graph or data flow.
             Some(def_id) if name.is_none() => self.get_fn_addr(
                 ty::Instance::expect_resolve(
                     tcx,
@@ -981,7 +991,7 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
                     ty::List::empty(),
                     DUMMY_SP,
                 ),
-                tcx.sess.pointer_authentication_functions(),
+                None,
             ),
             _ => {
                 let name = name.unwrap_or("rust_eh_personality");

compiler/rustc_codegen_ssa/src/base.rs

@@ -492,8 +492,11 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         // We want to create the wrapper only when the codegen unit is the primary one
         return None;
     }
-
-    let main_llfn = cx.get_fn_addr(instance, cx.sess().pointer_authentication_functions());
+    // No function pointer signing / type discriminator is needed here. Although `get_fn_addr` is
+    // used to obtain function pointers, both the user's `main` and `LangItem::Start` use the Rust
+    // ABI (currently pointer authentication is only supported for C/System ABI). The same applies
+    // to the logic in `create_entry_fn` further below.
+    let main_llfn = cx.get_fn_addr(instance, None);
 
     let entry_fn = create_entry_fn::<Bx>(cx, main_llfn, main_def_id, entry_type);
     return Some(entry_fn);
@@ -554,8 +557,8 @@ 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, cx.sess().pointer_authentication_functions());
+
+            let start_fn = cx.get_fn_addr(start_instance, None);
 
             let i8_ty = cx.type_i8();
             let arg_sigpipe = bx.const_u8(sigpipe);

compiler/rustc_codegen_ssa/src/common.rs

@@ -2,6 +2,9 @@
 
 use rustc_hir::LangItem;
 use rustc_hir::attrs::PeImportNameType;
+use rustc_middle::ptrauth::{
+    build_fn_ptr_type_discriminator_input_from_instance, compute_fn_ptr_type_discriminator,
+};
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{self, Instance, TyCtxt};
 use rustc_middle::{bug, mir, span_bug};
@@ -117,11 +120,19 @@ 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, tcx.sess.pointer_authentication_functions()),
-        instance,
-    )
+    let mut schema = bx.sess().pointer_authentication_functions().clone();
+
+    if let Some(ref mut s) = schema
+        && bx.sess().pointer_authentication_fn_ptr_type_discrimination()
+    {
+        // It is unlikely that any of LangItem will follow the extern C/System ABI, but it future
+        // proofs the implementation.
+        let disc_input = build_fn_ptr_type_discriminator_input_from_instance(tcx, instance);
+        let disc = compute_fn_ptr_type_discriminator(tcx, &disc_input) as u16;
+
+        s.constant_discriminator = disc;
+    }
+    (bx.fn_abi_of_instance(instance, ty::List::empty()), bx.get_fn_addr(instance, schema), instance)
 }
 
 pub(crate) fn shift_mask_val<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -11,6 +11,9 @@ use rustc_hir::attrs::AttributeKind;
 use rustc_hir::lang_items::LangItem;
 use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER;
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
+use rustc_middle::ptrauth::{
+    build_fn_ptr_type_discriminator_input_from_instance, compute_fn_ptr_type_discriminator,
+};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
 use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt};
@@ -684,12 +687,25 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     virtual_drop,
                 )
             }
-            _ => (
-                false,
-                bx.get_fn_addr(drop_fn, bx.sess().pointer_authentication_functions()),
-                bx.fn_abi_of_instance(drop_fn, ty::List::empty()),
-                drop_fn,
-            ),
+            _ => {
+                let mut schema = bx.sess().pointer_authentication_functions().clone();
+
+                if let Some(ref mut s) = schema
+                    && bx.sess().pointer_authentication_fn_ptr_type_discrimination()
+                {
+                    let disc_input =
+                        build_fn_ptr_type_discriminator_input_from_instance(bx.tcx(), drop_fn);
+                    let disc = compute_fn_ptr_type_discriminator(bx.tcx(), &disc_input) as u16;
+
+                    s.constant_discriminator = disc;
+                }
+                (
+                    false,
+                    bx.get_fn_addr(drop_fn, schema),
+                    bx.fn_abi_of_instance(drop_fn, ty::List::empty()),
+                    drop_fn,
+                )
+            }
         };
 
         // We generate a null check for the drop_fn. This saves a bunch of relocations being
@@ -1102,13 +1118,22 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         )
                         .unwrap();
 
-                        (
-                            None,
-                            Some(bx.get_fn_addr(
+                        let mut schema = bx.sess().pointer_authentication_functions().clone();
+
+                        if let Some(ref mut s) = schema
+                            && bx.sess().pointer_authentication_fn_ptr_type_discrimination()
+                        {
+                            let disc_input = build_fn_ptr_type_discriminator_input_from_instance(
+                                bx.tcx(),
                                 instance,
-                                bx.sess().pointer_authentication_functions(),
-                            )),
-                        )
+                            );
+                            let disc =
+                                compute_fn_ptr_type_discriminator(bx.tcx(), &disc_input) as u16;
+
+                            s.constant_discriminator = disc;
+                        };
+
+                        (None, Some(bx.get_fn_addr(instance, schema)))
                     }
                     _ => (Some(instance), None),
                 }
@@ -1422,7 +1447,19 @@ 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, bx.sess().pointer_authentication_functions())
+                let mut schema = bx.sess().pointer_authentication_functions().clone();
+
+                if let Some(ref mut s) = schema
+                    && bx.sess().pointer_authentication_fn_ptr_type_discrimination()
+                {
+                    let disc_input =
+                        build_fn_ptr_type_discriminator_input_from_instance(bx.tcx(), instance);
+                    let disc = compute_fn_ptr_type_discriminator(bx.tcx(), &disc_input) as u16;
+
+                    s.constant_discriminator = disc;
+                }
+
+                bx.get_fn_addr(instance, schema)
             }
             (_, Some(llfn)) => llfn,
             _ => span_bug!(fn_span, "no instance or llfn for call"),

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -1,7 +1,7 @@
 use itertools::Itertools as _;
 use rustc_abi::{self as abi, BackendRepr, FIRST_VARIANT};
 use rustc_middle::ptrauth::{
-    build_fn_ptr_type_discriminator_input, compute_fn_ptr_type_discriminator,
+    build_fn_ptr_type_discriminator_input_from_ty, compute_fn_ptr_type_discriminator,
 };
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
@@ -70,8 +70,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     ) -> Bx::Value {
         let tcx = bx.tcx();
 
-        let src_input = build_fn_ptr_type_discriminator_input(tcx, info.src_ty);
-        let dst_input = build_fn_ptr_type_discriminator_input(tcx, info.dst_ty);
+        let src_input = build_fn_ptr_type_discriminator_input_from_ty(tcx, info.src_ty);
+        let dst_input = build_fn_ptr_type_discriminator_input_from_ty(tcx, info.dst_ty);
 
         let src_disc = match src_input {
             Some(src) => compute_fn_ptr_type_discriminator(tcx, &src),
@@ -612,7 +612,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                                 if let Some(ref mut s) = schema {
                                     if bx.sess().pointer_authentication_fn_ptr_type_discrimination() {
-                                        if let Some(input) = build_fn_ptr_type_discriminator_input(
+                                        if let Some(input) = build_fn_ptr_type_discriminator_input_from_ty(
                                             bx.tcx(),
                                             operand.layout.ty,
                                         ) {
@@ -640,9 +640,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                     ty::ClosureKind::FnOnce,
                                 );
                                 OperandValue::Immediate(
+                                    // A closure coerced to a function pointer retains the Rust
+                                    // ABI. Pointer authentication only applies to extern
+                                    // "C"/System ABI function pointer, hence pass None to
+                                    // `get_fn_addr`.
                                     bx.cx().get_fn_addr(
                                         instance,
-                                        bx.sess().pointer_authentication_functions(),
+                                        None,
                                     ),
                                 )
                             }
@@ -868,8 +872,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         def: ty::InstanceKind::Shim(ty::ShimKind::ThreadLocal(def_id)),
                         args: ty::GenericArgs::empty(),
                     };
-                    let fn_ptr =
-                        bx.get_fn_addr(instance, bx.sess().pointer_authentication_functions());
+                    // This is the address of a compiler-generated TLS shim function. It is not an
+                    // externally visible function pointer and does not require function pointer
+                    // authentication signing.
+                    let fn_ptr = bx.get_fn_addr(instance, None);
                     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() {

compiler/rustc_middle/src/ptrauth/discriminator.rs

@@ -42,7 +42,9 @@ The computation is structured into three conceptual stages:
 
 - Public API
   - `FnPtrTypeDiscriminatorInput`
-  - `build_fn_ptr_type_discriminator_input`
+  - `build_fn_ptr_type_discriminator_input_from_instance`
+  - `build_fn_ptr_type_discriminator_input_from_sig`
+  - `build_fn_ptr_type_discriminator_input_from_ty`
   - `compute_fn_ptr_type_discriminator`
 
 - Signature extraction
@@ -68,7 +70,7 @@ function types only. It does NOT attempt to model full type system rules.
 */
 
 use rustc_abi::ExternAbi;
-use rustc_middle::ty::{self, Ty, TyCtxt, Unnormalized};
+use rustc_middle::ty::{self, Instance, Ty, TyCtxt, Unnormalized};
 use rustc_span::sym;
 
 use crate::ptrauth::llvm_siphash::llvm_pointer_auth_stable_siphash;
@@ -84,7 +86,7 @@ pub struct FnPtrTypeDiscriminatorInput<'tcx> {
 }
 
 impl<'tcx> FnPtrTypeDiscriminatorInput<'tcx> {
-    pub fn from_sig(sig: ty::FnSig<'tcx>) -> Self {
+    fn from_sig(sig: ty::FnSig<'tcx>) -> Self {
         FnPtrTypeDiscriminatorInput {
             inputs: sig.inputs(),
             output: sig.output(),
@@ -93,10 +95,7 @@ impl<'tcx> FnPtrTypeDiscriminatorInput<'tcx> {
         }
     }
 
-    pub fn from_sig_tys(
-        sig: ty::FnSigTys<TyCtxt<'tcx>>,
-        header: &ty::FnHeader<TyCtxt<'tcx>>,
-    ) -> Self {
+    fn from_sig_tys(sig: ty::FnSigTys<TyCtxt<'tcx>>, header: &ty::FnHeader<TyCtxt<'tcx>>) -> Self {
         FnPtrTypeDiscriminatorInput {
             inputs: sig.inputs(),
             output: sig.output(),
@@ -137,7 +136,7 @@ pub fn extract_fn_ptr_type<'tcx>(tcx: TyCtxt<'tcx>, mut ty: Ty<'tcx>) -> Option<
 ///
 /// FnDef is only accepted for convenience; the discriminator is still computed
 /// from the instantiated function signature.
-pub fn build_fn_ptr_type_discriminator_input<'tcx>(
+pub fn build_fn_ptr_type_discriminator_input_from_ty<'tcx>(
     tcx: TyCtxt<'tcx>,
     ty: Ty<'tcx>,
 ) -> Option<FnPtrTypeDiscriminatorInput<'tcx>> {
@@ -160,6 +159,37 @@ pub fn build_fn_ptr_type_discriminator_input<'tcx>(
     }
 }
 
+/// Builds type discrimination input from a monomorphized function instance.
+///
+/// The instance's signature is instantiated using its generic arguments and
+/// normalized before constructing the canonical discriminator input. Unlike
+/// `build_fn_ptr_type_discriminator_input_from_ty`, this function cannot fail
+/// because an `Instance` always represents a callable item with a well-defined
+/// function signature.
+pub fn build_fn_ptr_type_discriminator_input_from_instance<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+) -> FnPtrTypeDiscriminatorInput<'tcx> {
+    let sig = tcx
+        .instantiate_and_normalize_erasing_regions(
+            instance.args,
+            ty::TypingEnv::fully_monomorphized(),
+            tcx.fn_sig(instance.def_id()),
+        )
+        .skip_binder();
+
+    FnPtrTypeDiscriminatorInput::from_sig(sig)
+}
+
+/// Builds type discrimination input from a function signature.
+///
+/// The signature is assumed to already be instantiated and normalized.
+pub fn build_fn_ptr_type_discriminator_input_from_sig<'tcx>(
+    sig: ty::FnSig<'tcx>,
+) -> FnPtrTypeDiscriminatorInput<'tcx> {
+    FnPtrTypeDiscriminatorInput::from_sig(sig)
+}
+
 pub fn compute_fn_ptr_type_discriminator<'tcx>(
     tcx: TyCtxt<'tcx>,
     input: &FnPtrTypeDiscriminatorInput<'tcx>,

compiler/rustc_middle/src/ptrauth/mod.rs

@@ -2,6 +2,7 @@ pub mod discriminator;
 pub mod llvm_siphash;
 
 pub use discriminator::{
-    FnPtrTypeDiscriminatorInput, build_fn_ptr_type_discriminator_input,
+    FnPtrTypeDiscriminatorInput, build_fn_ptr_type_discriminator_input_from_instance,
+    build_fn_ptr_type_discriminator_input_from_sig, build_fn_ptr_type_discriminator_input_from_ty,
     compute_fn_ptr_type_discriminator,
 };

compiler/rustc_session/src/session.rs

@@ -96,7 +96,7 @@ pub enum PointerAuthARM8_3Key {
 }
 
 /// Forms of extra discrimination.
-#[derive(Clone, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum PointerAuthDiscrimination {
     /// No additional discrimination.
     None,
@@ -109,7 +109,7 @@ pub enum PointerAuthDiscrimination {
 }
 
 /// Types of address discrimination.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum PointerAuthAddressDiscriminator {
     /// Enable/disable hardware address discrimination.
     HardwareAddress(bool),
@@ -118,7 +118,7 @@ pub enum PointerAuthAddressDiscriminator {
     Synthetic(u64),
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct PointerAuthSchema {
     pub is_address_discriminated: PointerAuthAddressDiscriminator,
     pub discrimination_kind: PointerAuthDiscrimination,