CFI: Support function pointers for trait methods

Adds support for both CFI and KCFI for attaching concrete and abstract
types to functions. KCFI does this through generation of `ReifyShim` on
any function pointer that could go in a vtable, and checking the
`ReifyReason` when emitting the instance. CFI does this by attaching
both the concrete and abstract type to every instance.

TypeID codegen tests are switched to be anchored on the left rather than
the right in order to allow emission of additional type attachments.

Fixes #115953

4 files changed, 63 insertions, 10 deletions

diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs
index e5625c8a5c6..4a7720b38f8 100644
--- a/compiler/rustc_middle/src/ty/instance.rs
+++ b/compiler/rustc_middle/src/ty/instance.rs
@@ -517,6 +517,24 @@ impl<'tcx> Instance<'tcx> {
                     debug!(" => fn pointer created for virtual call");
                     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 }
+                }
                 _ => {}
             }
 
diff --git a/compiler/rustc_symbol_mangling/src/typeid.rs b/compiler/rustc_symbol_mangling/src/typeid.rs
index 11657b3b303..862ba285db8 100644
--- a/compiler/rustc_symbol_mangling/src/typeid.rs
+++ b/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;
@@ -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();
diff --git a/tests/ui/sanitizer/cfi-closures.rs b/tests/ui/sanitizer/cfi-closures.rs
index 54f1cc035bc..f3d9be35716 100644
--- a/tests/ui/sanitizer/cfi-closures.rs
+++ b/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;
diff --git a/tests/ui/sanitizer/cfi-method-fn-ptr-cast.rs b/tests/ui/sanitizer/cfi-method-fn-ptr-cast.rs
index 273b8785fae..8f79de11748 100644
--- a/tests/ui/sanitizer/cfi-method-fn-ptr-cast.rs
+++ b/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)
 }