fix -Zsanitizer=kcfi on #[naked] functions

Author: folkertdev

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)
-                {
+                // A shim used by `-Zsanitizer=kcfi`. Its body is a call to the real function.
+                // The shim should be codegened like a normal function.
+                let is_reify_shim = matches!(instance.def, InstanceKind::ReifyShim(..));
+
+                let flags = cx.tcx().codegen_fn_attrs(instance.def_id()).flags;
+                if flags.contains(CodegenFnAttrFlags::NAKED) && !is_reify_shim {
                     naked_asm::codegen_naked_asm::<Bx::CodegenCx>(cx, instance, item_data);
                 } else {
                     base::codegen_instance::<Bx>(cx, instance);

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