kcfi: only reify trait methods when dyn-compatible

compiler/rustc_middle/src/ty/instance.rs

@@ -618,11 +618,11 @@ impl<'tcx> Instance<'tcx> {
                         // be directly reified because it's closure-like. The reify can handle the
                         // unresolved instance.
                         resolved = Instance { def: InstanceKind::ReifyShim(def_id, reason), args }
-                    // Reify `Trait::method` implementations
-                    // FIXME(maurer) only reify it if it is a vtable-safe function
+                    // Reify `Trait::method` implementations if the trait is dyn-compatible.
                     } else if let Some(assoc) = tcx.opt_associated_item(def_id)
                         && let AssocContainer::Trait | AssocContainer::TraitImpl(Ok(_)) =
                             assoc.container
+                        && tcx.is_dyn_compatible(assoc.container_id(tcx))
                     {
                         // 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.

tests/codegen-llvm/sanitizer/kcfi/fn-ptr-reify-shim.rs

@@ -0,0 +1,74 @@
+//@ add-core-stubs
+//@ revisions: aarch64 x86_64
+//@ [aarch64] compile-flags: --target aarch64-unknown-none
+//@ [aarch64] needs-llvm-components: aarch64
+//@ [x86_64] compile-flags: --target x86_64-unknown-none
+//@ [x86_64] needs-llvm-components: x86
+//@ compile-flags: -Ctarget-feature=-crt-static -Zsanitizer=kcfi -Cno-prepopulate-passes -Copt-level=0
+
+#![feature(no_core, lang_items)]
+#![crate_type = "lib"]
+#![no_core]
+
+// A `ReifyShim` should only be created when the trait is dyn-compatible.
+
+extern crate minicore;
+use minicore::*;
+
+trait DynCompatible {
+    fn dyn_name(&self) -> &'static str;
+
+    fn dyn_name_default(&self) -> &'static str {
+        let _ = self;
+        "dyn_default"
+    }
+}
+
+// Not dyn-compatible because the `Self: Sized` bound is missing.
+trait NotDynCompatible {
+    fn not_dyn_name() -> &'static str;
+
+    fn not_dyn_name_default() -> &'static str {
+        "not_dyn_default"
+    }
+}
+
+struct S;
+
+impl DynCompatible for S {
+    fn dyn_name(&self) -> &'static str {
+        "dyn_compatible"
+    }
+}
+
+impl NotDynCompatible for S {
+    fn not_dyn_name() -> &'static str {
+        "not_dyn_compatible"
+    }
+}
+
+#[no_mangle]
+pub fn main() {
+    let s = S;
+
+    // `DynCompatible` is indeed dyn-compatible.
+    let _: &dyn DynCompatible = &s;
+
+    // CHECK: call <fn_ptr_reify_shim::S as fn_ptr_reify_shim::DynCompatible>::dyn_name{{.*}}reify.shim.fnptr
+    let dyn_name = S::dyn_name as fn(&S) -> &str;
+    let _unused = dyn_name(&s);
+
+    // CHECK: call fn_ptr_reify_shim::DynCompatible::dyn_name_default{{.*}}reify.shim.fnptr
+    let dyn_name_default = S::dyn_name_default as fn(&S) -> &str;
+    let _unused = dyn_name_default(&s);
+
+    // Check using $ (end-of-line) that these calls do not contain `reify.shim.fnptr`.
+
+    // CHECK: call <fn_ptr_reify_shim::S as fn_ptr_reify_shim::NotDynCompatible>::not_dyn_name{{$}}
+    let not_dyn_name = S::not_dyn_name as fn() -> &'static str;
+    let _unused = not_dyn_name();
+
+    // CHECK: call fn_ptr_reify_shim::NotDynCompatible::not_dyn_name_default{{$}}
+    let not_dyn_name_default = S::not_dyn_name_default as fn() -> &'static str;
+    let _unused = not_dyn_name_default();
+}

tests/codegen-llvm/sanitizer/kcfi/naked-function.rs

@@ -15,8 +15,9 @@ use minicore::*;
 
 struct Thing;
 trait MyTrait {
+    // NOTE: this test assumes that this trait is dyn-compatible.
     #[unsafe(naked)]
-    extern "C" fn my_naked_function() {
+    extern "C" fn my_naked_function(&self) {
         // the real function is defined
         // CHECK: .globl
         // CHECK-SAME: my_naked_function
@@ -34,13 +35,13 @@ impl MyTrait for Thing {}
 #[unsafe(no_mangle)]
 pub fn main() {
     // Trick the compiler into generating an indirect call.
-    const F: extern "C" fn() = Thing::my_naked_function;
+    const F: extern "C" fn(&Thing) = Thing::my_naked_function;
 
     // main calls the shim function
     // CHECK: call void
     // CHECK-SAME: my_naked_function
     // CHECK-SAME: reify.shim.fnptr
-    (F)();
+    (F)(&Thing);
 }
 
 // CHECK: declare !kcfi_type

tests/ui/sanitizer/kcfi-c-variadic.rs

@@ -0,0 +1,20 @@
+//@ needs-sanitizer-kcfi
+//@ no-prefer-dynamic
+//@ compile-flags: -Zsanitizer=kcfi -Cpanic=abort -Cunsafe-allow-abi-mismatch=sanitizer
+//@ ignore-backends: gcc
+//@ run-pass
+
+#![feature(c_variadic)]
+
+trait Trait {
+    unsafe extern "C" fn foo(x: i32, y: i32, mut ap: ...) -> i32 {
+        x + y + ap.arg::<i32>() + ap.arg::<i32>()
+    }
+}
+
+impl Trait for i32 {}
+
+fn main() {
+    let f = i32::foo as unsafe extern "C" fn(i32, i32, ...) -> i32;
+    assert_eq!(unsafe { f(1, 2, 3, 4) }, 1 + 2 + 3 + 4);
+}