[PAC] ABI-compatibility rules for discriminator calculation (4.5/8)

compiler/rustc_attr_ir/src/data_structures.rs

@@ -590,6 +590,12 @@ pub enum RustcDumpLayoutKind {
     Size,
 }
 
+#[derive(Clone, Debug, StableHash, Encodable, Decodable, PrintAttribute)]
+pub enum RustcDumpPtrauthDiscriminatorKind {
+    Encoding,
+    Hash,
+}
+
 #[derive(Clone, Debug, StableHash, Encodable, Decodable, PrintAttribute, PartialEq, Eq)]
 pub enum RustcMirKind {
     PeekMaybeInit,
@@ -1252,6 +1258,9 @@ pub enum AttributeKind {
     /// Represents the [`rustc_dump_object_lifetime_defaults`](./attribute.rustc_dump_object_lifetime_defaults.html) attribute.
     RustcDumpObjectLifetimeDefaults,
 
+    /// Represents the [`rustc_dump_ptrauth_discriminator`](./attribute.:rustc_dump_ptrauth_discriminator.html) attribute.
+    RustcDumpPtrauthDiscriminator(ThinVec<RustcDumpPtrauthDiscriminatorKind>),
+
     /// Represents the [`rustc_dump_symbol_name`](./attribute.rustc_dump_symbol_name.html) attribute.
     RustcDumpSymbolName(Span),
 

compiler/rustc_attr_ir/src/encode_cross_crate.rs

@@ -140,6 +140,7 @@ impl AttributeKind {
             RustcDumpItemBounds => No,
             RustcDumpLayout(..) => No,
             RustcDumpObjectLifetimeDefaults => No,
+            RustcDumpPtrauthDiscriminator(..) => No,
             RustcDumpSymbolName(..) => Yes,
             RustcDumpUserArgs => No,
             RustcDumpVariances => No,
@@ -355,6 +356,7 @@ impl AttributeKind {
             RustcDumpItemBounds => false,
             RustcDumpLayout(..) => false,
             RustcDumpObjectLifetimeDefaults => false,
+            RustcDumpPtrauthDiscriminator(..) => false,
             RustcDumpSymbolName(..) => false,
             RustcDumpUserArgs => false,
             RustcDumpVariances => false,

compiler/rustc_attr_parsing/src/attributes/rustc_dump.rs

@@ -1,5 +1,5 @@
 use rustc_attr_ir::target::{AssocCtxt, MethodKind, Target};
-use rustc_attr_ir::{AttributeKind, RustcDumpLayoutKind};
+use rustc_attr_ir::{AttributeKind, RustcDumpLayoutKind, RustcDumpPtrauthDiscriminatorKind};
 use rustc_feature::AttributeStability;
 use rustc_span::{Span, Symbol, sym};
 
@@ -178,6 +178,56 @@ impl CombineAttributeParser for RustcDumpLayoutParser {
     }
 }
 
+pub(crate) struct RustcDumpPtrauthDiscriminatorParser;
+
+impl CombineAttributeParser for RustcDumpPtrauthDiscriminatorParser {
+    const PATH: &[Symbol] = &[sym::rustc_dump_ptrauth_discriminator];
+
+    type Item = RustcDumpPtrauthDiscriminatorKind;
+
+    const CONVERT: ConvertFn<Self::Item> =
+        |items, _| AttributeKind::RustcDumpPtrauthDiscriminator(items);
+
+    const ALLOWED_TARGETS: AllowedTargets<'_> = AllowedTargets::AllowList(&[Allow(Target::Fn)]);
+
+    const TEMPLATE: AttributeTemplate = template!(List: &["encoding", "hash"]);
+    const STABILITY: AttributeStability = unstable!(rustc_attrs);
+
+    fn extend(
+        cx: &mut AcceptContext<'_, '_>,
+        args: &ArgParser,
+    ) -> impl IntoIterator<Item = Self::Item> {
+        let Some(items) = cx.expect_list(args, cx.attr_span) else {
+            return vec![];
+        };
+
+        let mut result = Vec::new();
+        for item in items.mixed() {
+            let Some(arg) = item.meta_item_no_args() else {
+                cx.adcx().expected_not_literal(item.span());
+                continue;
+            };
+            let Some(ident) = arg.ident() else {
+                cx.adcx().expected_identifier(arg.span());
+                return vec![];
+            };
+            let kind = match ident.name {
+                sym::ptrauth_encoding => RustcDumpPtrauthDiscriminatorKind::Encoding,
+                sym::ptrauth_hash => RustcDumpPtrauthDiscriminatorKind::Hash,
+                _ => {
+                    cx.adcx().expected_specific_argument(
+                        ident.span,
+                        &[sym::ptrauth_encoding, sym::ptrauth_hash],
+                    );
+                    continue;
+                }
+            };
+            result.push(kind);
+        }
+        result
+    }
+}
+
 pub(crate) struct RustcDumpObjectLifetimeDefaultsParser;
 
 impl NoArgsAttributeParser for RustcDumpObjectLifetimeDefaultsParser {

compiler/rustc_attr_parsing/src/context.rs

@@ -194,6 +194,7 @@ attribute_parsers!(
         Combine<RustcAllowConstFnUnstableParser>,
         Combine<RustcCleanParser>,
         Combine<RustcDumpLayoutParser>,
+        Combine<RustcDumpPtrauthDiscriminatorParser>,
         Combine<RustcMirParser>,
         Combine<RustcThenThisWouldNeedParser>,
         Combine<TargetFeatureParser>,

compiler/rustc_codegen_gcc/src/builder.rs

@@ -1943,6 +1943,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

@@ -1,9 +1,10 @@
 use gccjit::{GlobalKind, LValue, RValue, ToRValue, Type};
 use rustc_abi::Primitive::Pointer;
-use rustc_abi::{self as abi, HasDataLayout};
+use rustc_abi::{self as abi, HasDataLayout, Size};
 use rustc_codegen_ssa::traits::{
     BaseTypeCodegenMethods, ConstCodegenMethods, MiscCodegenMethods, StaticCodegenMethods,
 };
+use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::mir::Mutability;
 use rustc_middle::mir::interpret::{GlobalAlloc, PointerArithmetic, Scalar};
 use rustc_middle::ty::layout::LayoutOf;
@@ -116,7 +117,7 @@ impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
             Mutability::Mut => {
                 self.static_addr_of_mut(const_alloc_to_gcc(self, alloc), alloc.inner().align, None)
             }
-            _ => self.static_addr_of(alloc, None),
+            _ => self.static_addr_of(alloc, None, None),
         };
         if !self.sess().fewer_names() {
             // FIXME(antoyo): set value name.
@@ -332,7 +333,8 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
         cv: Scalar,
         layout: abi::Scalar,
         ty: Type<'gcc>,
-        _schema: Option<&PointerAuthSchema>,
+        _ptrauth_schema: Option<PointerAuthSchema>,
+        _ptrauth_discriminators: Option<&FxHashMap<Size, u64>>,
     ) -> RValue<'gcc> {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {

compiler/rustc_codegen_gcc/src/consts.rs

@@ -8,6 +8,7 @@ use rustc_attr_ir::Linkage;
 use rustc_codegen_ssa::traits::{
     BaseTypeCodegenMethods, ConstCodegenMethods, StaticCodegenMethods,
 };
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_log::tracing::trace;
@@ -61,7 +62,12 @@ fn set_global_alignment<'gcc, 'tcx>(
 }
 
 impl<'gcc, 'tcx> StaticCodegenMethods for CodegenCx<'gcc, 'tcx> {
-    fn static_addr_of(&self, alloc: ConstAllocation<'_>, kind: Option<&str>) -> RValue<'gcc> {
+    fn static_addr_of(
+        &self,
+        alloc: ConstAllocation<'_>,
+        kind: Option<&str>,
+        _ptrauth_discriminators: Option<&FxHashMap<Size, u64>>,
+    ) -> RValue<'gcc> {
         let cv = const_alloc_to_gcc(self, alloc);
         let align = alloc.inner().align;
 

compiler/rustc_codegen_gcc/src/context.rs

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

compiler/rustc_codegen_gcc/src/int.rs

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

compiler/rustc_codegen_llvm/src/asm.rs

@@ -166,7 +166,8 @@ impl<'ll, 'tcx> AsmBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     ConstScalar::Ptr(ptr, _) => {
                         let (prov, _) = ptr.prov_and_relative_offset();
                         let global_alloc = self.tcx.global_alloc(prov.alloc_id());
-                        let value = self.cx.alloc_to_backend(global_alloc, false, None).unwrap();
+                        let value =
+                            self.cx.alloc_to_backend(global_alloc, false, None, None).unwrap();
                         inputs.push(value);
                         op_idx.insert(idx, constraints.len());
                         constraints.push("s".to_string());
@@ -456,8 +457,9 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
                                 ConstScalar::Ptr(ptr, _) => {
                                     let (prov, offset) = ptr.prov_and_relative_offset();
                                     let global_alloc = self.tcx.global_alloc(prov.alloc_id());
-                                    let llval =
-                                        self.alloc_to_backend(global_alloc, true, None).unwrap();
+                                    let llval = self
+                                        .alloc_to_backend(global_alloc, true, None, None)
+                                        .unwrap();
 
                                     self.add_compiler_used_global(llval);
                                     let symbol = llvm::build_string(|s| unsafe {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1567,6 +1567,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, '_> {
@@ -2213,8 +2237,13 @@ 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;
+        // If sess().pointer_authentication_fn_ptr_type_discrimination() is enabled, this contains
+        // the function pointer type discriminator; otherwise, it is None. LLVM expects a u64 here,
+        // so use 0 when no discriminator is present.
+        let discriminator = fn_abi?.ptrauth_discriminator.unwrap_or(0);
+
         Some(llvm::OperandBundleBox::new(
             "ptrauth",
             &[self.const_u32(key), self.const_u64(discriminator)],

compiler/rustc_codegen_llvm/src/common.rs

@@ -4,11 +4,12 @@ use std::borrow::Borrow;
 
 use libc::{c_char, c_uint};
 use rustc_abi::Primitive::Pointer;
-use rustc_abi::{self as abi, ExternAbi, HasDataLayout as _};
+use rustc_abi::{self as abi, ExternAbi, HasDataLayout as _, Size};
 use rustc_ast::Mutability;
 use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::*;
 use rustc_crate_store::DllImport;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hash::{StableHash, StableHasher};
 use rustc_hashes::Hash128;
 use rustc_hir::def::DefKind;
@@ -30,11 +31,9 @@ pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
     cx: &CodegenCx<'ll, '_>,
     instance: Instance<'tcx>,
     llfn: &'ll llvm::Value,
-    schema: &PointerAuthSchema,
+    ptrauth_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();
@@ -54,7 +53,7 @@ pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
         return llfn;
     }
 
-    let addr_diversity = match schema.is_address_discriminated {
+    let addr_diversity = match ptrauth_schema.is_address_discriminated {
         PointerAuthAddressDiscriminator::HardwareAddress(true) => Some(llfn),
         PointerAuthAddressDiscriminator::HardwareAddress(false) => None,
         PointerAuthAddressDiscriminator::Synthetic(val) => {
@@ -63,7 +62,12 @@ pub(crate) fn maybe_sign_fn_ptr<'ll, 'tcx>(
             Some(unsafe { llvm::LLVMConstIntToPtr(llval, llty) })
         }
     };
-    const_ptr_auth(llfn, schema.key as u32, schema.constant_discriminator as u64, addr_diversity)
+    const_ptr_auth(
+        llfn,
+        ptrauth_schema.key as u32,
+        ptrauth_schema.constant_discriminator as u64,
+        addr_diversity,
+    )
 }
 
 /*
@@ -179,11 +183,12 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         &self,
         global_alloc: GlobalAlloc<'tcx>,
         need_symbol_name: bool,
-        schema: Option<&PointerAuthSchema>,
+        ptrauth_schema: Option<PointerAuthSchema>,
+        ptrauth_discriminators: Option<&FxHashMap<Size, u64>>,
     ) -> Result<&'ll Value, u64> {
         let alloc = match global_alloc {
             GlobalAlloc::Function { instance, .. } => {
-                return Ok(self.get_fn_addr(instance, schema));
+                return Ok(self.get_fn_addr(instance, ptrauth_schema));
             }
             GlobalAlloc::Static(def_id) => {
                 assert!(self.tcx.is_static(def_id));
@@ -226,7 +231,13 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             }
         };
 
-        let init = const_alloc_to_llvm(self, alloc.inner(), IsStatic::No, IsInitOrFini::No);
+        let init = const_alloc_to_llvm(
+            self,
+            alloc.inner(),
+            IsStatic::No,
+            IsInitOrFini::No,
+            ptrauth_discriminators,
+        );
         let alloc = alloc.inner();
 
         if need_symbol_name {
@@ -405,7 +416,8 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
         cv: Scalar,
         layout: abi::Scalar,
         llty: &'ll Type,
-        schema: Option<&PointerAuthSchema>,
+        ptrauth_schema: Option<PointerAuthSchema>,
+        ptrauth_discriminators: Option<&FxHashMap<Size, u64>>,
     ) -> &'ll Value {
         let bitsize = if layout.is_bool() { 1 } else { layout.size(self).bits() };
         match cv {
@@ -422,7 +434,12 @@ impl<'ll, 'tcx> ConstCodegenMethods for CodegenCx<'ll, 'tcx> {
                 let (prov, offset) = ptr.prov_and_relative_offset();
                 let global_alloc = self.tcx.global_alloc(prov.alloc_id());
                 let base_addr_space = global_alloc.address_space(self);
-                let base_addr = match self.alloc_to_backend(global_alloc, false, schema) {
+                let base_addr = match self.alloc_to_backend(
+                    global_alloc,
+                    false,
+                    ptrauth_schema,
+                    ptrauth_discriminators,
+                ) {
                     Ok(base_addr) => base_addr,
                     Err(base_addr) => {
                         let val = base_addr.wrapping_add(offset.bytes());