CFI: Support function pointers for trait methods

compiler/rustc_codegen_llvm/src/declare.rs

@@ -147,7 +147,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
                 for options in [
                     TypeIdOptions::GENERALIZE_POINTERS,
                     TypeIdOptions::NORMALIZE_INTEGERS,
-                    TypeIdOptions::NO_SELF_TYPE_ERASURE,
+                    TypeIdOptions::ERASE_SELF_TYPE,
                 ]
                 .into_iter()
                 .powerset()
@@ -173,7 +173,9 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 
         if self.tcx.sess.is_sanitizer_kcfi_enabled() {
             // LLVM KCFI does not support multiple !kcfi_type attachments
-            let mut options = TypeIdOptions::empty();
+            // Default to erasing the self type. If we need the concrete type, there will be a
+            // hint in the instance.
+            let mut options = TypeIdOptions::ERASE_SELF_TYPE;
             if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
                 options.insert(TypeIdOptions::GENERALIZE_POINTERS);
             }

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,28 @@ pub struct Instance<'tcx> {
     pub args: GenericArgsRef<'tcx>,
 }
 
+/// Describes why a `ReifyShim` was created. This is needed to distingish a ReifyShim created to
+/// adjust for things like `#[track_caller]` in a vtable from a `ReifyShim` created to produce a
+/// function pointer from a vtable entry.
+/// Currently, this is only used when KCFI is enabled, as only KCFI needs to treat those two
+/// `ReifyShim`s differently.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+#[derive(TyEncodable, TyDecodable, HashStable)]
+pub enum ReifyReason {
+    /// The `ReifyShim` was created to produce a function pointer. This happens when:
+    /// * A vtable entry is directly converted to a function call (e.g. creating a fn ptr from a
+    ///   method on a `dyn` object).
+    /// * A function with `#[track_caller]` is converted to a function pointer
+    /// * If KCFI is enabled, creating a function pointer from a method on an object-safe trait.
+    /// This includes the case of converting `::call`-like methods on closure-likes to function
+    /// pointers.
+    FnPtr,
+    /// This `ReifyShim` was created to populate a vtable. Currently, this happens when a
+    /// `#[track_caller]` mismatch occurs between the implementation of a method and the method.
+    /// This includes the case of `::call`-like methods in closure-likes' vtables.
+    Vtable,
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 #[derive(TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable, Lift)]
 pub enum InstanceDef<'tcx> {
@@ -67,7 +89,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 +222,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 +382,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 +506,34 @@ 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);
+                }
+                // Reify `Trait::method` implementations if KCFI is enabled
+                // FIXME(maurer) only reify 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)
+                }
+                // Reify `::call`-like method implementations if KCFI is enabled
+                _ if tcx.sess.is_sanitizer_kcfi_enabled()
+                    && tcx.is_closure_like(resolved.def_id()) =>
+                {
+                    // Reroute through a reify via the *unresolved* instance. The resolved one can't
+                    // be directly reified because it's closure-like. The reify can handle the
+                    // unresolved instance.
+                    resolved = Instance { def: InstanceDef::ReifyShim(def_id, reason), args }
                 }
                 _ => {}
             }
@@ -508,6 +557,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 +594,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

@@ -88,7 +88,7 @@ pub use self::context::{
     tls, CtxtInterners, CurrentGcx, 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

@@ -449,6 +449,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;
@@ -24,9 +24,9 @@ bitflags! {
         /// `-fsanitize-cfi-icall-experimental-normalize-integers` option for cross-language LLVM
         /// CFI and  KCFI support.
         const NORMALIZE_INTEGERS = 4;
-        /// Do not perform self type erasure for attaching a secondary type id to methods with their
-        /// concrete self so they can be used as function pointers.
-        const NO_SELF_TYPE_ERASURE = 8;
+        /// Generalize the instance by erasing the concrete `Self` type where possible.
+        /// Only has an effect on `{kcfi_,}typeid_for_instance`.
+        const ERASE_SELF_TYPE = 8;
     }
 }
 
@@ -67,8 +67,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.remove(TypeIdOptions::ERASE_SELF_TYPE);
+    }
     // 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

@@ -1172,7 +1172,7 @@ pub fn typeid_for_instance<'tcx>(
         instance.args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
     }
 
-    if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE) {
+    if options.contains(EncodeTyOptions::ERASE_SELF_TYPE) {
         if let Some(impl_id) = tcx.impl_of_method(instance.def_id())
             && let Some(trait_ref) = tcx.impl_trait_ref(impl_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"),

tests/codegen/sanitizer/cfi/emit-type-metadata-id-itanium-cxx-abi-method-secondary-typeid.rs

@@ -18,5 +18,5 @@ impl Trait1 for Type1 {
 }
 
 
-// CHECK: ![[TYPE1]] = !{i64 0, !"_ZTSFvu3refIu3dynIu{{[0-9]+}}NtC{{[[:print:]]+}}_{{[[:print:]]+}}6Trait1u6regionEEE"}
-// CHECK: ![[TYPE2]] = !{i64 0, !"_ZTSFvu3refIu{{[0-9]+}}NtC{{[[:print:]]+}}_{{[[:print:]]+}}5Type1EE"}
+// CHECK: ![[TYPE1]] = !{i64 0, !"_ZTSFvu3refIu{{[0-9]+}}NtC{{[[:print:]]+}}_{{[[:print:]]+}}5Type1EE"}
+// CHECK: ![[TYPE2]] = !{i64 0, !"_ZTSFvu3refIu3dynIu{{[0-9]+}}NtC{{[[:print:]]+}}_{{[[:print:]]+}}6Trait1u6regionEEE"}

tests/ui/sanitizer/cfi-closures.rs

@@ -15,7 +15,6 @@
 
 #![feature(fn_traits)]
 #![feature(unboxed_closures)]
-#![feature(cfg_sanitize)]
 
 fn foo<'a, T>() -> Box<dyn Fn(&'a T) -> &'a T> {
     Box::new(|x| x)
@@ -72,9 +71,6 @@ fn use_closure<C>(call: extern "rust-call" fn(&C, ()) -> i32, f: &C) -> i32 {
 }
 
 #[test]
-// FIXME after KCFI reify support is added, remove this
-// It will appear to work if you test locally, set -C opt-level=0 to see it fail.
-#[cfg_attr(sanitize = "kcfi", ignore)]
 fn closure_addr_taken() {
     let x = 3i32;
     let f = || x;

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

@@ -1,11 +1,41 @@
 // Verifies that casting a method to a function pointer works.
-//
-// FIXME(#122848): Remove only-linux when fixed.
+
+//@ revisions: cfi kcfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
 //@ only-linux
-//@ needs-sanitizer-cfi
-//@ compile-flags: -Clto -Copt-level=0 -Cprefer-dynamic=off -Ctarget-feature=-crt-static -Zsanitizer=cfi
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C target-feature=-crt-static
+//@ [cfi] compile-flags: -C opt-level=0 -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
+//@ [kcfi] compile-flags: -C panic=abort -C prefer-dynamic=off
 //@ run-pass
 
+trait Foo {
+    fn foo(&self);
+    fn bar(&self);
+}
+
+struct S;
+
+impl Foo for S {
+    fn foo(&self) {}
+    #[track_caller]
+    fn bar(&self) {}
+}
+
+struct S2 {
+    f: fn(&S)
+}
+
+impl S2 {
+    fn foo(&self, s: &S) {
+        (self.f)(s)
+    }
+}
+
 trait Trait1 {
     fn foo(&self);
 }
@@ -20,4 +50,8 @@ fn main() {
     let type1 = Type1 {};
     let f = <Type1 as Trait1>::foo;
     f(&type1);
+    // Check again with different optimization barriers
+    S2 { f: <S as Foo>::foo }.foo(&S);
+    // Check mismatched #[track_caller]
+    S2 { f: <S as Foo>::bar }.foo(&S)
 }