[PAC] Function pointer type discrimination for transmutes

Implement pointer authentication domain handling for function pointer
transmutes. When function pointer type discrimination is enabled,
transmuting between function pointer types with different authentication
domains now re-signs the pointer using the appropriate discriminator.

Author: jchlanda

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -1,8 +1,11 @@
 use std::assert_matches;
 
 use itertools::Itertools as _;
-use rustc_abi::{self as abi, BackendRepr, FIRST_VARIANT};
+use rustc_abi::{self as abi, BackendRepr, ExternAbi, FIRST_VARIANT};
 use rustc_index::IndexVec;
+use rustc_middle::ptrauth::{
+    ptrauth_clone_discriminated_schema_for, ptrauth_compute_fn_ptr_type_discriminator_for,
+};
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Mutability, Ty, TyCtxt};
@@ -17,6 +20,13 @@ use crate::common::{IntPredicate, TypeKind};
 use crate::traits::*;
 use crate::{MemFlags, base};
 
+/// Type metadata used when applying pointer authentication semantics during
+/// transmute lowering.
+struct TransmuteInfo<'tcx> {
+    src_ty: Ty<'tcx>,
+    dst_ty: Ty<'tcx>,
+}
+
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     fn try_codegen_const_aggregate_as_immediate(
         &mut self,
@@ -91,6 +101,126 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         true
     }
 
+    /// Applies pointer-authentication-aware semantic transmute, that is ensuring that when a
+    /// function pointer is transmuted between two types that map to different type
+    /// discriminators, the resulting pointer is re-signed appropriately.
+    ///
+    /// Only SSA `OperandValue::Immediate` values are eligible for this path.
+    fn ptrauth_codegen_transmute_operand(
+        &mut self,
+        bx: &mut Bx,
+        operand: OperandRef<'tcx, Bx::Value>,
+        cast: TyAndLayout<'tcx>,
+    ) -> OperandValue<Bx::Value> {
+        let val = self.codegen_transmute_operand(bx, operand, cast);
+
+        let OperandValue::Immediate(ptr) = val else {
+            return val;
+        };
+
+        let info = TransmuteInfo { src_ty: operand.layout.ty, dst_ty: cast.ty };
+
+        OperandValue::Immediate(self.resign_transmuted_fn_ptr(bx, ptr, info))
+    }
+
+    /// Applies pointer-authentication type discriminator correction for a function pointer value
+    /// being transmuted between two types.
+    ///
+    /// If the source and destination types have different type discriminator values, the pointer
+    /// must be resigned using `llvm.ptrauth.resign` intrinsic.
+    ///
+    /// A discriminator value of `0` is used to represent non-function-pointer
+    /// or "raw pointer" values:
+    /// ```text
+    ///   static mut CPTR: *const u8 = 0 as *const u8;
+    ///   ... = mem::transmute::<*const u8, unsafe extern "C" fn()>(CPTR);
+    /// ```
+    /// where the source has no type discriminator.
+    fn resign_transmuted_fn_ptr(
+        &mut self,
+        bx: &mut Bx,
+        val: Bx::Value,
+        info: TransmuteInfo<'tcx>,
+    ) -> Bx::Value {
+        // Resigning can only happen in the context of function pointer type discrimination.
+        assert!(self.cx.tcx().sess.pointer_authentication_fn_ptr_type_discrimination());
+
+        let tcx = bx.tcx();
+
+        let src_disc = ptrauth_compute_fn_ptr_type_discriminator_for(tcx, info.src_ty).unwrap_or(0);
+        let dst_disc = ptrauth_compute_fn_ptr_type_discriminator_for(tcx, info.dst_ty).unwrap_or(0);
+
+        if src_disc == dst_disc {
+            return val;
+        }
+
+        debug!("resign_transmuted_fn_ptr\t{:#x} -> {:#x}", src_disc, dst_disc);
+
+        let key = self.cx.tcx().sess.pointer_authentication_fn_ptr_key().unwrap() as u32;
+        bx.ptrauth_resign(val, key, src_disc.into(), key, dst_disc.into())
+    }
+
+    /// Walks through `#[repr(transparent)]` wrappers to find an underlying
+    /// function pointer or function definition.
+    ///
+    /// Returns the corresponding layout if one is found, otherwise `None`.
+    fn transparent_fn_ptr_layout(
+        &self,
+        mut layout: TyAndLayout<'tcx>,
+    ) -> Option<TyAndLayout<'tcx>> {
+        loop {
+            match layout.ty.kind() {
+                ty::FnPtr(..) | ty::FnDef(..) => return Some(layout),
+
+                ty::Adt(def, _) if def.repr().transparent() => {
+                    layout = layout.field(self.cx, 0);
+                }
+
+                _ => return None,
+            }
+        }
+    }
+
+    /// Applies pointer-authentication type discriminator change during a transmute into a
+    /// memory-backed place.
+    ///
+    /// Unlike the operand version, this path handles values stored in memory
+    /// and therefore must unwrap #[repr(transparent)] wrapper types so that pointer
+    /// authentication is based on the underlying function pointer type.
+    ///
+    /// Only immediate values are subject to ptrauth adjustment; other
+    /// representations are passed through unchanged.
+    fn ptrauth_codegen_transmute_place(
+        &mut self,
+        bx: &mut Bx,
+        src: OperandRef<'tcx, Bx::Value>,
+        dst: PlaceRef<'tcx, Bx::Value>,
+    ) {
+        debug!(
+            "ptrauth_codegen_transmute_place\tsrc={:?}, dst={:?}",
+            src.layout.ty.kind(),
+            dst.layout.ty.kind()
+        );
+
+        let info = TransmuteInfo {
+            src_ty: self.transparent_fn_ptr_layout(src.layout).map_or(src.layout.ty, |l| l.ty),
+            dst_ty: self.transparent_fn_ptr_layout(dst.layout).map_or(dst.layout.ty, |l| l.ty),
+        };
+
+        let dest = dst.val.with_type(src.layout);
+
+        let val = match src.val {
+            OperandValue::Immediate(v) => {
+                let v = self.resign_transmuted_fn_ptr(bx, v, info);
+                OperandValue::Immediate(v)
+            }
+            other => other,
+        };
+
+        OperandRef { val, layout: src.layout, move_annotation: None }
+            .store_with_annotation(bx, dest);
+    }
+
     #[instrument(level = "trace", skip(self, bx))]
     pub(crate) fn codegen_rvalue(
         &mut self,
@@ -179,8 +309,25 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::Rvalue::Cast(
                 mir::CastKind::Transmute | mir::CastKind::Subtype,
                 ref operand,
-                _ty,
+                ty,
             ) => {
+                if self.cx.tcx().sess.pointer_authentication_fn_ptr_type_discrimination() {
+                    let src_ty = operand.ty(self.mir, self.cx.tcx());
+                    let dst_ty = self.monomorphize(ty);
+
+                    if src_ty.is_fn_ptr() || dst_ty.is_fn_ptr() {
+                        let op = self.codegen_operand(bx, operand);
+                        let cast = bx.cx().layout_of(dst_ty);
+
+                        let val = self.ptrauth_codegen_transmute_operand(bx, op, cast);
+
+                        OperandRef { val, layout: cast, move_annotation: None }
+                            .store_with_annotation(bx, dest);
+
+                        return;
+                    }
+                }
+
                 let src = self.codegen_operand(bx, operand);
                 self.codegen_transmute(bx, src, dest);
             }
@@ -302,7 +449,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             // Since in this path we have a place anyway, we can store or copy to it,
             // making sure we use the destination place's alignment even if the
             // source would normally have a higher one.
-            src.store_with_annotation(bx, dst.val.with_type(src.layout));
+
+            if self.cx.tcx().sess.pointer_authentication_fn_ptr_type_discrimination() {
+                self.ptrauth_codegen_transmute_place(bx, src, dst);
+            } else {
+                src.store_with_annotation(bx, dst.val.with_type(src.layout));
+            }
         }
     }
 
@@ -316,6 +468,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         operand: OperandRef<'tcx, Bx::Value>,
         cast: TyAndLayout<'tcx>,
     ) -> OperandValue<Bx::Value> {
+        debug!(
+            "codegen_transmute_operand\t
+                from_ty={:?} to_ty={:?} from_layout={:?} to_layout={:?} is fnptr=({}, {})",
+            operand.layout.ty,
+            cast.ty,
+            operand.layout.backend_repr,
+            cast.backend_repr,
+            operand.layout.ty.is_fn_ptr(),
+            cast.ty.is_fn_ptr()
+        );
         if let abi::BackendRepr::Memory { .. } = cast.backend_repr
             && !cast.is_zst()
         {
@@ -502,12 +664,18 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                     args.no_bound_vars().unwrap(),
                                 )
                                 .unwrap();
-                                OperandValue::Immediate(
-                                    bx.get_fn_addr(
-                                        instance,
-                                        bx.sess().pointer_authentication_functions(),
-                                    ),
+
+                            let schema = if bx.sess().pointer_authentication_fn_ptr_type_discrimination() {
+                                ptrauth_clone_discriminated_schema_for(
+                                    bx.tcx(),
+                                    bx.sess().pointer_authentication_functions(),
+                                    operand.layout.ty,
                                 )
+                            } else {
+                                bx.sess().pointer_authentication_functions().clone()
+                            };
+
+                                OperandValue::Immediate(bx.get_fn_addr(instance, schema))
                             }
                             _ => bug!("{} cannot be reified to a fn ptr", operand.layout.ty),
                         }
@@ -521,10 +689,20 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                     args,
                                     ty::ClosureKind::FnOnce,
                                 );
+                                assert!(
+                                    !matches!(
+                                        bx.cx().tcx().fn_sig(instance.def_id()).skip_binder().abi(),
+                                        ExternAbi::C { .. } | ExternAbi::System { .. }
+                                    )
+                                );
                                 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(),
+                                        /* ptrauth_schema */ None,
                                     ),
                                 )
                             }
@@ -601,7 +779,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         })
                     }
                     mir::CastKind::Transmute | mir::CastKind::Subtype => {
-                        self.codegen_transmute_operand(bx, operand, cast)
+                        if self.cx.tcx().sess.pointer_authentication_fn_ptr_type_discrimination() {
+                            self.ptrauth_codegen_transmute_operand(bx, operand, cast)
+                        } else {
+                            self.codegen_transmute_operand(bx, operand, cast)
+                        }
                     }
                 };
                 OperandRef { val, layout: cast, move_annotation: None }
@@ -746,8 +928,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());
+                    // needs_thread_local_shim implies Windows/MSVC, for which pointer
+                    // authentication is not yet supported.
+                    assert!(!self.cx.tcx().sess.pointer_authentication());
+                    let fn_ptr = bx.get_fn_addr(instance, /* ptrauth_schema */ 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() {