[WIP] pac ty disc

compiler/rustc_codegen_gcc/src/builder.rs

@@ -1843,6 +1843,17 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
     fn fptosi_sat(&mut self, val: RValue<'gcc>, dest_ty: Type<'gcc>) -> RValue<'gcc> {
         self.fptoint_sat(true, val, dest_ty)
     }
+
+    fn ptrauth_resign(
+        &mut self,
+        _value: Self::Value,
+        _old_key: u32,
+        _old_discriminator: u64,
+        _new_key: u32,
+        _new_discriminator: u64,
+    ) -> Self::Value {
+        bug!("Resigning of pointers not implemented");
+    }
 }
 
 impl<'a, 'gcc, 'tcx> Builder<'a, 'gcc, 'tcx> {

compiler/rustc_codegen_gcc/src/common.rs

@@ -247,7 +247,7 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
         cv: Scalar,
         layout: abi::Scalar,
         ty: Type<'gcc>,
-        _schema: Option<&PointerAuthSchema>,
+        _schema: Option<PointerAuthSchema>,
     ) -> RValue<'gcc> {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {

compiler/rustc_codegen_gcc/src/context.rs

@@ -401,7 +401,7 @@ impl<'gcc, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
     fn get_fn_addr(
         &self,
         instance: Instance<'tcx>,
-        _pointer_auth_schema: Option<&PointerAuthSchema>,
+        _pointer_auth_schema: Option<PointerAuthSchema>,
     ) -> RValue<'gcc> {
         let func_name = self.tcx.symbol_name(instance).name;
 

compiler/rustc_codegen_gcc/src/int.rs

@@ -375,6 +375,7 @@ impl<'a, 'gcc, 'tcx> Builder<'a, 'gcc, 'tcx> {
             fixed_count: 3,
             conv: CanonAbi::C,
             can_unwind: false,
+            ptrauth_type_discriminator: 0,
         };
         fn_abi.adjust_for_foreign_abi(self.cx, ExternAbi::C { unwind: false });
 

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1542,6 +1542,30 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         let cold_inline = llvm::AttributeKind::Cold.create_attr(self.llcx);
         attributes::apply_to_callsite(llret, llvm::AttributePlace::Function, &[cold_inline]);
     }
+
+    fn ptrauth_resign(
+        &mut self,
+        value: &'ll Value,
+        old_key: u32,
+        old_discriminator: u64,
+        new_key: u32,
+        new_discriminator: u64,
+    ) -> &'ll Value {
+        let ptr_as_int = self.ptrtoint(value, self.type_i64());
+        let resigned_int = self.call_intrinsic(
+            "llvm.ptrauth.resign",
+            &[],
+            &[
+                ptr_as_int,
+                self.const_i32(old_key as i32),
+                self.const_i64(old_discriminator as i64),
+                self.const_i32(new_key as i32),
+                self.const_i64(new_discriminator as i64),
+            ],
+        );
+
+        self.inttoptr(resigned_int, self.val_ty(value))
+    }
 }
 
 impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {
@@ -2059,8 +2083,14 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         // bundles.
         // Once this is resolved, we should analyze each call and skip direct calls. See the
         // discussion in the rust-lang issue: <https://github.com/rust-lang/rust/issues/152532>
-        let key: u32 = 0;
-        let discriminator: u64 = 0;
+
+        let key: u32 = self.sess().pointer_authentication_fn_ptr_key().unwrap() as u32;
+        let discriminator = if self.sess().pointer_authentication_fn_ptr_type_discrimination() {
+            fn_abi?.ptrauth_type_discriminator
+        } else {
+            0
+        };
+
         Some(llvm::OperandBundleBox::new(
             "ptrauth",
             &[self.const_u32(key), self.const_u64(discriminator)],

compiler/rustc_codegen_llvm/src/common.rs

@@ -30,11 +30,9 @@ pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
     cx: &CodegenCx<'ll, '_>,
     instance: Instance<'tcx>,
     llfn: &'ll llvm::Value,
-    schema: &PointerAuthSchema,
+    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();
@@ -317,7 +315,7 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
         cv: Scalar,
         layout: abi::Scalar,
         llty: &'ll Type,
-        schema: Option<&PointerAuthSchema>,
+        schema: Option<PointerAuthSchema>,
     ) -> &'ll Value {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {
@@ -352,6 +350,7 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                                 alloc.inner(),
                                 IsStatic::No,
                                 IsInitOrFini::No,
+                                None,
                             );
                             let alloc = alloc.inner();
                             let value = match alloc.mutability {
@@ -389,6 +388,7 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                             alloc.inner(),
                             IsStatic::No,
                             IsInitOrFini::No,
+                            None,
                         );
                         self.static_addr_of_impl(init, alloc.inner().align, None)
                     }

compiler/rustc_codegen_llvm/src/consts.rs

@@ -3,6 +3,7 @@ use std::ops::Range;
 use rustc_abi::{Align, ExternAbi, HasDataLayout, Primitive, Scalar, Size, WrappingRange};
 use rustc_codegen_ssa::common;
 use rustc_codegen_ssa::traits::*;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::LangItem;
 use rustc_hir::attrs::Linkage;
 use rustc_hir::def::DefKind;
@@ -13,8 +14,9 @@ use rustc_middle::mir::interpret::{
     read_target_uint,
 };
 use rustc_middle::mono::MonoItem;
+use rustc_middle::ptrauth::compute_fn_ptr_type_discriminator_for;
 use rustc_middle::ty::layout::{HasTypingEnv, LayoutOf};
-use rustc_middle::ty::{self, Instance};
+use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
 use rustc_span::Symbol;
 use rustc_target::spec::Arch;
@@ -37,11 +39,152 @@ pub(crate) enum IsInitOrFini {
     Yes,
     No,
 }
+
+/// Maps offsets within a static allocation to the function pointer
+/// discriminator that should be applied when authenticating the relocation
+/// emitted at that offset.
+///
+/// Offsets are relative to the beginning of the allocation.
+pub(crate) struct FnPtrDiscriminatorAtOffset {
+    map: FxHashMap<Size, u64>,
+}
+
+/// Recursively walks a type layout and records the offsets of all extern "C"
+/// function pointer fields together with their computed type discriminators.
+///
+/// Traversal currently supports:
+/// - direct function pointers
+/// - transparent wrappers
+/// - structs
+/// - tuples
+/// - arrays
+///
+/// Offsets are accumulated relative to the containing object.
+fn collect_fn_ptr_discriminators<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    typing_env: ty::TypingEnv<'tcx>,
+    ty: Ty<'tcx>,
+) -> FnPtrDiscriminatorAtOffset {
+    let mut map = FxHashMap::default();
+
+    collect_fn_ptr_discriminators_inner(tcx, typing_env, ty, Size::ZERO, &mut map);
+
+    FnPtrDiscriminatorAtOffset { map }
+}
+
+fn collect_fn_ptr_discriminators_inner<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    typing_env: ty::TypingEnv<'tcx>,
+    ty: Ty<'tcx>,
+    base_offset: Size,
+    map: &mut FxHashMap<Size, u64>,
+) {
+    // Direct function pointer.
+    if let Some(disc) = compute_fn_ptr_type_discriminator_for(tcx, ty) {
+        map.insert(base_offset, disc.into());
+
+        return;
+    }
+
+    match ty.kind() {
+        ty::Adt(def, args) if def.repr().transparent() => {
+            let variant = def.non_enum_variant();
+
+            let Some((_, field)) = variant.fields.iter_enumerated().next() else {
+                return;
+            };
+
+            let field_ty = tcx.normalize_erasing_regions(typing_env, field.ty(tcx, args));
+
+            collect_fn_ptr_discriminators_inner(tcx, typing_env, field_ty, base_offset, map);
+        }
+        ty::Adt(def, args) if def.is_struct() => {
+            let Ok(layout) = tcx.layout_of(typing_env.as_query_input(ty)) else {
+                return;
+            };
+
+            let variant = def.non_enum_variant();
+
+            for (idx, field_def) in variant.fields.iter_enumerated() {
+                let field_ty = tcx.normalize_erasing_regions(typing_env, field_def.ty(tcx, args));
+
+                let field_offset = layout.fields.offset(idx.into());
+
+                collect_fn_ptr_discriminators_inner(
+                    tcx,
+                    typing_env,
+                    field_ty,
+                    base_offset + field_offset,
+                    map,
+                );
+            }
+        }
+        ty::Tuple(fields) => {
+            let Ok(layout) = tcx.layout_of(typing_env.as_query_input(ty)) else {
+                return;
+            };
+
+            for (idx, field_ty) in fields.iter().enumerate() {
+                let field_offset = layout.fields.offset(idx);
+
+                collect_fn_ptr_discriminators_inner(
+                    tcx,
+                    typing_env,
+                    field_ty,
+                    base_offset + field_offset,
+                    map,
+                );
+            }
+        }
+        ty::Array(elem_ty, len) => {
+            let count = match len.try_to_target_usize(tcx) {
+                Some(v) => v,
+                None => return,
+            };
+
+            let Ok(elem_layout) = tcx.layout_of(typing_env.as_query_input(*elem_ty)) else {
+                return;
+            };
+
+            let stride = elem_layout.size;
+
+            // Collect discriminator of one element, so we don't have to recompute it for all the
+            // elements in the array.
+            let mut elem_map = FxHashMap::default();
+
+            collect_fn_ptr_discriminators_inner(
+                tcx,
+                typing_env,
+                *elem_ty,
+                Size::ZERO,
+                &mut elem_map,
+            );
+
+            // SAFETY: We immediately collect into a Vec and sort by offset.
+            // The HashMap iteration order is irrelevant and must not affect determinism.
+            #[allow(rustc::potential_query_instability)]
+            let mut entries: Vec<(Size, u64)> = elem_map.into_iter().collect();
+            entries.sort_unstable_by_key(|(offset, _)| *offset);
+
+            // Replicate for every array slot.
+            for i in 0..count {
+                let elem_base = base_offset + stride * i;
+
+                for (inner_offset, discr) in entries.iter().copied() {
+                    map.insert(elem_base + inner_offset, discr);
+                }
+            }
+        }
+        _ => {}
+    }
+}
+
 pub(crate) fn const_alloc_to_llvm<'ll>(
     cx: &CodegenCx<'ll, '_>,
     alloc: &Allocation,
     is_static: IsStatic,
     is_init_fini: IsInitOrFini,
+    fn_ptr_discriminators: Option<&FnPtrDiscriminatorAtOffset>,
 ) -> &'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
@@ -121,14 +264,24 @@ pub(crate) fn const_alloc_to_llvm<'ll>(
             as u64;
 
         let address_space = cx.tcx.global_alloc(prov.alloc_id()).address_space(cx);
-        let schema = if cx.sess().pointer_authentication() {
+        let mut schema = if cx.sess().pointer_authentication() {
             match is_init_fini {
                 IsInitOrFini::Yes => cx.sess().pointer_authentication_init_fini(),
                 IsInitOrFini::No => cx.sess().pointer_authentication_functions(),
             }
         } else {
             None
         };
+        let discr = fn_ptr_discriminators
+            .as_ref()
+            .and_then(|m| m.map.get(&Size::from_bytes(offset as u64)));
+
+        // Init/fini entries must not participate in function pointer type discrimination.
+        if let (Some(schema), Some(discr)) = (schema.as_mut(), discr)
+            && is_init_fini == IsInitOrFini::No
+        {
+            schema.constant_discriminator = *discr as u16;
+        }
         llvals.push(cx.scalar_to_backend_with_pac(
             InterpScalar::from_pointer(Pointer::new(prov, Size::from_bytes(ptr_offset)), &cx.tcx),
             Scalar::Initialized {
@@ -160,6 +313,15 @@ fn codegen_static_initializer<'ll, 'tcx>(
     cx: &CodegenCx<'ll, 'tcx>,
     def_id: DefId,
 ) -> Result<(&'ll Value, ConstAllocation<'tcx>), ErrorHandled> {
+    let fn_ptr_discriminators = if cx.sess().pointer_authentication_fn_ptr_type_discrimination() {
+        let instance = Instance::mono(cx.tcx, def_id);
+        let ty = instance.ty(cx.tcx, cx.typing_env());
+
+        Some(collect_fn_ptr_discriminators(cx.tcx, cx.typing_env(), ty))
+    } else {
+        None
+    };
+
     let alloc = cx.tcx.eval_static_initializer(def_id)?;
     let attrs = cx.tcx.codegen_fn_attrs(def_id);
     // FIXME(jchlanda) Decide if this could be better served by `ctor` crate. See the discussion
@@ -175,7 +337,16 @@ fn codegen_static_initializer<'ll, 'tcx>(
             }
         })
         .unwrap_or(IsInitOrFini::No);
-    Ok((const_alloc_to_llvm(cx, alloc.inner(), IsStatic::Yes, is_in_init_fini), alloc))
+    Ok((
+        const_alloc_to_llvm(
+            cx,
+            alloc.inner(),
+            IsStatic::Yes,
+            is_in_init_fini,
+            fn_ptr_discriminators.as_ref(),
+        ),
+        alloc,
+    ))
 }
 
 fn set_global_alignment<'ll>(cx: &CodegenCx<'ll, '_>, gv: &'ll Value, mut align: Align) {
@@ -837,7 +1008,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(), IsStatic::No, IsInitOrFini::No);
+        let cv = const_alloc_to_llvm(self, alloc.inner(), IsStatic::No, IsInitOrFini::No, None);
 
         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