Auto merge of rust-lang#123052 - maurer:addr-taken, r=<try>

CFI: Support function pointers for trait methods

Adds support for both CFI and KCFI for function pointers to trait methods by attaching both concrete and abstract types to functions.

KCFI does this through generation of a `ReifyShim` on any function pointer for a method that could go into a vtable, and keeping this separate from `ReifyShim`s that are *intended* for vtable us by setting a `ReifyReason` on them.

CFI does this by setting both the concrete and abstract type on every instance.

This should land after rust-lang#123024 or a similar PR, as it diverges the implementation of CFI vs KCFI.

r? `@compiler-errors`

compiler/rustc_codegen_llvm/src/declare.rs

@@ -18,6 +18,7 @@ use crate::llvm;
 use crate::llvm::AttributePlace::Function;
 use crate::type_::Type;
 use crate::value::Value;
+use itertools::Itertools;
 use rustc_codegen_ssa::traits::TypeMembershipMethods;
 use rustc_middle::ty::{Instance, Ty};
 use rustc_symbol_mangling::typeid::{
@@ -143,18 +144,18 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             if let Some(instance) = instance {
                 let typeid = typeid_for_instance(self.tcx, &instance, TypeIdOptions::empty());
                 self.set_type_metadata(llfn, typeid);
-                let typeid =
-                    typeid_for_instance(self.tcx, &instance, TypeIdOptions::GENERALIZE_POINTERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid =
-                    typeid_for_instance(self.tcx, &instance, TypeIdOptions::NORMALIZE_INTEGERS);
-                self.add_type_metadata(llfn, typeid);
-                let typeid = typeid_for_instance(
-                    self.tcx,
-                    &instance,
-                    TypeIdOptions::GENERALIZE_POINTERS | TypeIdOptions::NORMALIZE_INTEGERS,
-                );
-                self.add_type_metadata(llfn, typeid);
+                for options in [
+                    TypeIdOptions::GENERALIZE_POINTERS,
+                    TypeIdOptions::NORMALIZE_INTEGERS,
+                    TypeIdOptions::CONCRETE,
+                ]
+                .into_iter()
+                .powerset()
+                .map(TypeIdOptions::from_iter)
+                {
+                    let typeid = typeid_for_instance(self.tcx, &instance, options);
+                    self.add_type_metadata(llfn, typeid);
+                }
             } else {
                 let typeid = typeid_for_fnabi(self.tcx, fn_abi, TypeIdOptions::empty());
                 self.set_type_metadata(llfn, typeid);

compiler/rustc_middle/src/mir/visit.rs

@@ -341,7 +341,7 @@ macro_rules! make_mir_visitor {
 
                         ty::InstanceDef::Intrinsic(_def_id) |
                         ty::InstanceDef::VTableShim(_def_id) |
-                        ty::InstanceDef::ReifyShim(_def_id) |
+                        ty::InstanceDef::ReifyShim(_def_id, _) |
                         ty::InstanceDef::Virtual(_def_id, _) |
                         ty::InstanceDef::ThreadLocalShim(_def_id) |
                         ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |

compiler/rustc_middle/src/ty/instance.rs

@@ -31,6 +31,13 @@ pub struct Instance<'tcx> {
     pub args: GenericArgsRef<'tcx>,
 }
 
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+#[derive(TyEncodable, TyDecodable, HashStable)]
+pub enum ReifyReason {
+    FnPtr,
+    Vtable,
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 #[derive(TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable, Lift)]
 pub enum InstanceDef<'tcx> {
@@ -67,7 +74,13 @@ pub enum InstanceDef<'tcx> {
     /// Because this is a required part of the function's ABI but can't be tracked
     /// as a property of the function pointer, we use a single "caller location"
     /// (the definition of the function itself).
-    ReifyShim(DefId),
+    ///
+    /// The second field encodes *why* this shim was created. This allows distinguishing between
+    /// a `ReifyShim` that appears in a vtable vs one that appears as a function pointer.
+    ///
+    /// This field will only be populated if we are compiling in a mode that needs these shims
+    /// to be separable, currently only when KCFI is enabled.
+    ReifyShim(DefId, Option<ReifyReason>),
 
     /// `<fn() as FnTrait>::call_*` (generated `FnTrait` implementation for `fn()` pointers).
     ///
@@ -194,7 +207,7 @@ impl<'tcx> InstanceDef<'tcx> {
         match self {
             InstanceDef::Item(def_id)
             | InstanceDef::VTableShim(def_id)
-            | InstanceDef::ReifyShim(def_id)
+            | InstanceDef::ReifyShim(def_id, _)
             | InstanceDef::FnPtrShim(def_id, _)
             | InstanceDef::Virtual(def_id, _)
             | InstanceDef::Intrinsic(def_id)
@@ -354,7 +367,9 @@ fn fmt_instance(
     match instance.def {
         InstanceDef::Item(_) => Ok(()),
         InstanceDef::VTableShim(_) => write!(f, " - shim(vtable)"),
-        InstanceDef::ReifyShim(_) => write!(f, " - shim(reify)"),
+        InstanceDef::ReifyShim(_, None) => write!(f, " - shim(reify)"),
+        InstanceDef::ReifyShim(_, Some(ReifyReason::FnPtr)) => write!(f, " - shim(reify-fnptr)"),
+        InstanceDef::ReifyShim(_, Some(ReifyReason::Vtable)) => write!(f, " - shim(reify-vtable)"),
         InstanceDef::ThreadLocalShim(_) => write!(f, " - shim(tls)"),
         InstanceDef::Intrinsic(_) => write!(f, " - intrinsic"),
         InstanceDef::Virtual(_, num) => write!(f, " - virtual#{num}"),
@@ -476,15 +491,24 @@ impl<'tcx> Instance<'tcx> {
         debug!("resolve(def_id={:?}, args={:?})", def_id, args);
         // Use either `resolve_closure` or `resolve_for_vtable`
         assert!(!tcx.is_closure_like(def_id), "Called `resolve_for_fn_ptr` on closure: {def_id:?}");
+        let reason = tcx.sess.is_sanitizer_kcfi_enabled().then_some(ReifyReason::FnPtr);
         Instance::resolve(tcx, param_env, def_id, args).ok().flatten().map(|mut resolved| {
             match resolved.def {
                 InstanceDef::Item(def) if resolved.def.requires_caller_location(tcx) => {
                     debug!(" => fn pointer created for function with #[track_caller]");
-                    resolved.def = InstanceDef::ReifyShim(def);
+                    resolved.def = InstanceDef::ReifyShim(def, reason);
                 }
                 InstanceDef::Virtual(def_id, _) => {
                     debug!(" => fn pointer created for virtual call");
-                    resolved.def = InstanceDef::ReifyShim(def_id);
+                    resolved.def = InstanceDef::ReifyShim(def_id, reason);
+                }
+                // FIXME(maurer) only shim it if it is a vtable-safe function
+                _ if tcx.sess.is_sanitizer_kcfi_enabled()
+                    && tcx.associated_item(def_id).trait_item_def_id.is_some() =>
+                {
+                    // If this function could also go in a vtable, we need to `ReifyShim` it with
+                    // KCFI because it can only attach one type per function.
+                    resolved.def = InstanceDef::ReifyShim(resolved.def_id(), reason)
                 }
                 _ => {}
             }
@@ -508,6 +532,7 @@ impl<'tcx> Instance<'tcx> {
             debug!(" => associated item with unsizeable self: Self");
             Some(Instance { def: InstanceDef::VTableShim(def_id), args })
         } else {
+            let reason = tcx.sess.is_sanitizer_kcfi_enabled().then_some(ReifyReason::Vtable);
             Instance::resolve(tcx, param_env, def_id, args).ok().flatten().map(|mut resolved| {
                 match resolved.def {
                     InstanceDef::Item(def) => {
@@ -544,18 +569,18 @@ impl<'tcx> Instance<'tcx> {
                                 // Create a shim for the `FnOnce/FnMut/Fn` method we are calling
                                 // - unlike functions, invoking a closure always goes through a
                                 // trait.
-                                resolved = Instance { def: InstanceDef::ReifyShim(def_id), args };
+                                resolved = Instance { def: InstanceDef::ReifyShim(def_id, reason), args };
                             } else {
                                 debug!(
                                     " => vtable fn pointer created for function with #[track_caller]: {:?}", def
                                 );
-                                resolved.def = InstanceDef::ReifyShim(def);
+                                resolved.def = InstanceDef::ReifyShim(def, reason);
                             }
                         }
                     }
                     InstanceDef::Virtual(def_id, _) => {
                         debug!(" => vtable fn pointer created for virtual call");
-                        resolved.def = InstanceDef::ReifyShim(def_id);
+                        resolved.def = InstanceDef::ReifyShim(def_id, reason)
                     }
                     _ => {}
                 }

compiler/rustc_middle/src/ty/mod.rs

@@ -89,7 +89,7 @@ pub use self::context::{
     tls, CtxtInterners, DeducedParamAttrs, Feed, FreeRegionInfo, GlobalCtxt, Lift, TyCtxt,
     TyCtxtFeed,
 };
-pub use self::instance::{Instance, InstanceDef, ShortInstance, UnusedGenericParams};
+pub use self::instance::{Instance, InstanceDef, ReifyReason, ShortInstance, UnusedGenericParams};
 pub use self::list::List;
 pub use self::parameterized::ParameterizedOverTcx;
 pub use self::predicate::{

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -448,6 +448,7 @@ TrivialTypeTraversalAndLiftImpls! {
     crate::ty::ClosureKind,
     crate::ty::ParamConst,
     crate::ty::ParamTy,
+    crate::ty::instance::ReifyReason,
     interpret::AllocId,
     interpret::CtfeProvenance,
     interpret::Scalar,

compiler/rustc_mir_transform/src/inline.rs

@@ -324,7 +324,7 @@ impl<'tcx> Inliner<'tcx> {
             // do not need to catch this here, we can wait until the inliner decides to continue
             // inlining a second time.
             InstanceDef::VTableShim(_)
-            | InstanceDef::ReifyShim(_)
+            | InstanceDef::ReifyShim(..)
             | InstanceDef::FnPtrShim(..)
             | InstanceDef::ClosureOnceShim { .. }
             | InstanceDef::ConstructCoroutineInClosureShim { .. }

compiler/rustc_mir_transform/src/inline/cycle.rs

@@ -84,7 +84,7 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
                 // again, a function item can end up getting inlined. Thus we'll be able to cause
                 // a cycle that way
                 InstanceDef::VTableShim(_)
-                | InstanceDef::ReifyShim(_)
+                | InstanceDef::ReifyShim(..)
                 | InstanceDef::FnPtrShim(..)
                 | InstanceDef::ClosureOnceShim { .. }
                 | InstanceDef::ConstructCoroutineInClosureShim { .. }

compiler/rustc_mir_transform/src/shim.rs

@@ -55,7 +55,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
         // a virtual call, or a direct call to a function for which
         // indirect calls must be codegen'd differently than direct ones
         // (such as `#[track_caller]`).
-        ty::InstanceDef::ReifyShim(def_id) => {
+        ty::InstanceDef::ReifyShim(def_id, _) => {
             build_call_shim(tcx, instance, None, CallKind::Direct(def_id))
         }
         ty::InstanceDef::ClosureOnceShim { call_once: _, track_caller: _ } => {

compiler/rustc_symbol_mangling/src/legacy.rs

@@ -2,7 +2,7 @@ use rustc_data_structures::stable_hasher::{Hash64, HashStable, StableHasher};
 use rustc_hir::def_id::CrateNum;
 use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData};
 use rustc_middle::ty::print::{PrettyPrinter, Print, PrintError, Printer};
-use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, Instance, ReifyReason, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::ty::{GenericArg, GenericArgKind};
 
 use std::fmt::{self, Write};
@@ -71,8 +71,14 @@ pub(super) fn mangle<'tcx>(
         ty::InstanceDef::VTableShim(..) => {
             printer.write_str("{{vtable-shim}}").unwrap();
         }
-        ty::InstanceDef::ReifyShim(..) => {
-            printer.write_str("{{reify-shim}}").unwrap();
+        ty::InstanceDef::ReifyShim(_, reason) => {
+            printer.write_str("{{reify-shim").unwrap();
+            match reason {
+                Some(ReifyReason::FnPtr) => printer.write_str("-fnptr").unwrap(),
+                Some(ReifyReason::Vtable) => printer.write_str("-vtable").unwrap(),
+                None => (),
+            }
+            printer.write_str("}}").unwrap();
         }
         // FIXME(async_closures): This shouldn't be needed when we fix
         // `Instance::ty`/`Instance::def_id`.

compiler/rustc_symbol_mangling/src/typeid.rs

@@ -4,7 +4,7 @@
 /// For more information about LLVM CFI and cross-language LLVM CFI support for the Rust compiler,
 /// see design document in the tracking issue #89653.
 use bitflags::bitflags;
-use rustc_middle::ty::{Instance, Ty, TyCtxt};
+use rustc_middle::ty::{Instance, InstanceDef, ReifyReason, Ty, TyCtxt};
 use rustc_target::abi::call::FnAbi;
 use std::hash::Hasher;
 use twox_hash::XxHash64;
@@ -16,6 +16,8 @@ bitflags! {
         const GENERALIZE_POINTERS = 1;
         const GENERALIZE_REPR_C = 2;
         const NORMALIZE_INTEGERS = 4;
+        // Only relevant for `typeid_for_instance`
+        const CONCRETE = 8;
     }
 }
 
@@ -56,8 +58,13 @@ pub fn kcfi_typeid_for_fnabi<'tcx>(
 pub fn kcfi_typeid_for_instance<'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: &Instance<'tcx>,
-    options: TypeIdOptions,
+    mut options: TypeIdOptions,
 ) -> u32 {
+    // If we receive a `ReifyShim` intended to produce a function pointer, we need to remain
+    // concrete - abstraction is for vtables.
+    if matches!(instance.def, InstanceDef::ReifyShim(_, Some(ReifyReason::FnPtr))) {
+        options |= TypeIdOptions::CONCRETE
+    }
     // A KCFI type metadata identifier is a 32-bit constant produced by taking the lower half of the
     // xxHash64 of the type metadata identifier. (See llvm/llvm-project@cff5bef.)
     let mut hash: XxHash64 = Default::default();

compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs

@@ -1116,7 +1116,8 @@ pub fn typeid_for_instance<'tcx>(
         instance.args = strip_receiver_auto(tcx, instance.args)
     }
 
-    if let Some(impl_id) = tcx.impl_of_method(instance.def_id())
+    if !options.contains(TypeIdOptions::CONCRETE)
+        && let Some(impl_id) = tcx.impl_of_method(instance.def_id())
         && let Some(trait_ref) = tcx.impl_trait_ref(impl_id)
     {
         let impl_method = tcx.associated_item(instance.def_id());

compiler/rustc_symbol_mangling/src/v0.rs

@@ -8,8 +8,8 @@ use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData};
 use rustc_middle::ty::layout::IntegerExt;
 use rustc_middle::ty::print::{Print, PrintError, Printer};
 use rustc_middle::ty::{
-    self, EarlyBinder, FloatTy, Instance, IntTy, Ty, TyCtxt, TypeVisitable, TypeVisitableExt,
-    UintTy,
+    self, EarlyBinder, FloatTy, Instance, IntTy, ReifyReason, Ty, TyCtxt, TypeVisitable,
+    TypeVisitableExt, UintTy,
 };
 use rustc_middle::ty::{GenericArg, GenericArgKind};
 use rustc_span::symbol::kw;
@@ -44,7 +44,9 @@ pub(super) fn mangle<'tcx>(
     let shim_kind = match instance.def {
         ty::InstanceDef::ThreadLocalShim(_) => Some("tls"),
         ty::InstanceDef::VTableShim(_) => Some("vtable"),
-        ty::InstanceDef::ReifyShim(_) => Some("reify"),
+        ty::InstanceDef::ReifyShim(_, None) => Some("reify"),
+        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::FnPtr)) => Some("reify-fnptr"),
+        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::Vtable)) => Some("reify-vtable"),
 
         ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
         | ty::InstanceDef::CoroutineKindShim { .. } => Some("fn_once"),

compiler/rustc_target/src/spec/targets/aarch64_unknown_linux_gnu.rs

@@ -19,6 +19,7 @@ pub fn target() -> Target {
             stack_probes: StackProbeType::Inline,
             supported_sanitizers: SanitizerSet::ADDRESS
                 | SanitizerSet::CFI
+                | SanitizerSet::KCFI
                 | SanitizerSet::LEAK
                 | SanitizerSet::MEMORY
                 | SanitizerSet::MEMTAG

compiler/rustc_target/src/spec/targets/x86_64_unknown_linux_gnu.rs

@@ -10,6 +10,7 @@ pub fn target() -> Target {
     base.static_position_independent_executables = true;
     base.supported_sanitizers = SanitizerSet::ADDRESS
         | SanitizerSet::CFI
+        | SanitizerSet::KCFI
         | SanitizerSet::DATAFLOW
         | SanitizerSet::LEAK
         | SanitizerSet::MEMORY

src/tools/compiletest/src/header.rs

@@ -826,6 +826,7 @@ const KNOWN_DIRECTIVE_NAMES: &[&str] = &[
     "needs-sanitizer-cfi",
     "needs-sanitizer-dataflow",
     "needs-sanitizer-hwaddress",
+    "needs-sanitizer-kcfi",
     "needs-sanitizer-leak",
     "needs-sanitizer-memory",
     "needs-sanitizer-memtag",

tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs

@@ -5,11 +5,16 @@
 //
 // This checks that the reified function pointer will have the expected alias set at its call-site.
 
-//@ needs-sanitizer-cfi
+//@ revisions: cfi kcfi
 // FIXME(#122848) Remove only-linux once OSX CFI binaries work
 //@ only-linux
-//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
-//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C opt-level=0
+//@ [cfi] compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C lto
+//@ [cfi] compile-flags: -C prefer-dynamic=off
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
 //@ run-pass
 
 pub fn main() {

tests/ui/sanitizer/cfi-complex-receiver.rs

@@ -2,11 +2,16 @@
 // * Arc<dyn Foo> as for custom receivers
 // * &dyn Bar<T=Baz> for type constraints
 
-//@ needs-sanitizer-cfi
+//@ revisions: cfi kcfi
 // FIXME(#122848) Remove only-linux once OSX CFI binaries work
 //@ only-linux
-//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
-//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C opt-level=0
+//@ [cfi] compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C lto
+//@ [cfi] compile-flags: -C prefer-dynamic=off
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
 //@ run-pass
 
 use std::sync::Arc;

tests/ui/sanitizer/cfi-self-ref.rs

@@ -1,10 +1,15 @@
 // Check that encoding self-referential types works with #[repr(transparent)]
 
-//@ needs-sanitizer-cfi
+//@ revisions: cfi kcfi
 // FIXME(#122848) Remove only-linux once OSX CFI binaries work
 //@ only-linux
-//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
-//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C opt-level=0
+//@ [cfi] compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C lto
+//@ [cfi] compile-flags: -C prefer-dynamic=off
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
 //@ run-pass
 
 use std::marker::PhantomData;

tests/ui/sanitizer/cfi-virtual-auto.rs

@@ -1,10 +1,15 @@
 // Tests that calling a trait object method on a trait object with additional auto traits works.
 
-//@ needs-sanitizer-cfi
+//@ revisions: cfi kcfi
 // FIXME(#122848) Remove only-linux once OSX CFI binaries work
 //@ only-linux
-//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
-//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C opt-level=0
+//@ [cfi] compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C lto
+//@ [cfi] compile-flags: -C prefer-dynamic=off
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
 //@ run-pass
 
 trait Foo {