fix `-Zsanitizer=kcfi` on `#[naked]` functions

And more broadly only codegen `InstanceKind::Item` using the naked
function codegen code. Other instance kinds should follow the normal
path.

4 files changed, 91 insertions, 8 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/driver/aot.rs b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
index 442151fe32d..727f2760c0f 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/aot.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
@@ -530,8 +530,12 @@ fn codegen_cgu_content(
     for (mono_item, item_data) in mono_items {
         match mono_item {
             MonoItem::Fn(instance) => {
-                if tcx.codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED)
-                {
+                // Other `InstanceKind`s (e.g. `ReifyShim` generated by indirect calls) should be
+                // codegened like a normal function.
+                let is_item_instance = matches!(instance.def, InstanceKind::Item(_));
+
+                let flags = tcx.codegen_fn_attrs(instance.def_id()).flags;
+                if is_item_instance && flags.contains(CodegenFnAttrFlags::NAKED) {
                     rustc_codegen_ssa::mir::naked_asm::codegen_naked_asm(
                         &mut GlobalAsmContext { tcx, global_asm: &mut cx.global_asm },
                         instance,
diff --git a/compiler/rustc_codegen_ssa/src/mono_item.rs b/compiler/rustc_codegen_ssa/src/mono_item.rs
index 7b4268abe4b..fa0ef04b2e6 100644
--- a/compiler/rustc_codegen_ssa/src/mono_item.rs
+++ b/compiler/rustc_codegen_ssa/src/mono_item.rs
@@ -1,5 +1,6 @@
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
+use rustc_middle::ty::InstanceKind;
 use rustc_middle::ty::layout::HasTyCtxt;
 use tracing::debug;
 
@@ -41,12 +42,12 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
                 base::codegen_global_asm(cx, item_id);
             }
             MonoItem::Fn(instance) => {
-                if cx
-                    .tcx()
-                    .codegen_fn_attrs(instance.def_id())
-                    .flags
-                    .contains(CodegenFnAttrFlags::NAKED)
-                {
+                // Other `InstanceKind`s (e.g. `ReifyShim` generated by indirect calls) should be
+                // codegened like a normal function.
+                let is_item_instance = matches!(instance.def, InstanceKind::Item(_));
+
+                let flags = cx.tcx().codegen_fn_attrs(instance.def_id()).flags;
+                if is_item_instance && flags.contains(CodegenFnAttrFlags::NAKED) {
                     naked_asm::codegen_naked_asm::<Bx::CodegenCx>(cx, instance, item_data);
                 } else {
                     base::codegen_instance::<Bx>(cx, instance);
diff --git a/tests/codegen/sanitizer/kcfi/naked-function.rs b/tests/codegen/sanitizer/kcfi/naked-function.rs
new file mode 100644
index 00000000000..92431fc7d63
--- /dev/null
+++ b/tests/codegen/sanitizer/kcfi/naked-function.rs
@@ -0,0 +1,47 @@
+//@ 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]
+
+extern crate minicore;
+use minicore::*;
+
+struct Thing;
+trait MyTrait {
+    #[unsafe(naked)]
+    extern "C" fn my_naked_function() {
+        // the real function is defined
+        // CHECK: .globl
+        // CHECK-SAME: my_naked_function
+        naked_asm!("ret")
+    }
+}
+impl MyTrait for Thing {}
+
+// CHECK-LABEL: main
+#[unsafe(no_mangle)]
+pub fn main() {
+    // Trick the compiler into generating an indirect call.
+    const F: extern "C" fn() = Thing::my_naked_function;
+
+    // main calls the shim function
+    // CHECK: call
+    // CHECK-SAME: my_naked_function
+    // CHECK-SAME: reify.shim.fnptr
+    (F)();
+}
+
+// the shim calls the real function
+// CHECK: define
+// CHECK-SAME: my_naked_function
+// CHECK-SAME: reify.shim.fnptr
+
+// CHECK: declare !kcfi_type
+// CHECK-SAME: my_naked_function
diff --git a/tests/ui/asm/naked-function-shim.rs b/tests/ui/asm/naked-function-shim.rs
new file mode 100644
index 00000000000..4694d0cd963
--- /dev/null
+++ b/tests/ui/asm/naked-function-shim.rs
@@ -0,0 +1,31 @@
+// The indirect call will generate a shim that then calls the actual function. Test that
+// this is handled correctly. See also https://github.com/rust-lang/rust/issues/143266.
+
+//@ build-pass
+//@ 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
+
+#![feature(no_core, lang_items)]
+#![crate_type = "lib"]
+#![no_core]
+
+extern crate minicore;
+use minicore::*;
+
+trait MyTrait {
+    #[unsafe(naked)]
+    extern "C" fn foo(&self) {
+        naked_asm!("ret")
+    }
+}
+
+impl MyTrait for i32 {}
+
+fn main() {
+    let x: extern "C" fn(&_) = <dyn MyTrait as MyTrait>::foo;
+    x(&1);
+}