Rollup merge of #143293 - folkertdev:naked-function-kcfi, r=compiler-errors

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

fixes https://github.com/rust-lang/rust/issues/143266

With `-Zsanitizer=kcfi`, indirect calls happen via generated intermediate shim that forwards the call. The generated shim preserves the attributes of the original, including `#[unsafe(naked)]`. The shim is not a naked function though, and violates its invariants (like having a body that consists of a single `naked_asm!` call).

My fix here is to match on the `InstanceKind`, and only use `codegen_naked_asm` when the instance is not a `ReifyShim`. That does beg the question whether there are other `InstanceKind`s that could come up. As far as I can tell the answer is no: calling via `dyn` seems to work find, and `#[track_caller]` is disallowed in combination with `#[naked]`.

r? codegen
````@rustbot```` label +A-naked
cc ````@maurer```` ````@rcvalle````

19 files changed, 67 insertions, 34 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/driver/aot.rs b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
index 442151fe32d..8ec3599b63d 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/aot.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
@@ -530,8 +530,8 @@ 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)
-                {
+                let flags = tcx.codegen_instance_attrs(instance.def).flags;
+                if 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_cranelift/src/driver/jit.rs b/compiler/rustc_codegen_cranelift/src/driver/jit.rs
index b1f185b551c..b3497503bf0 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/jit.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/jit.rs
@@ -127,7 +127,7 @@ fn codegen_and_compile_fn<'tcx>(
     module: &mut dyn Module,
     instance: Instance<'tcx>,
 ) {
-    if tcx.codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED) {
+    if tcx.codegen_instance_attrs(instance.def).flags.contains(CodegenFnAttrFlags::NAKED) {
         tcx.dcx()
             .span_fatal(tcx.def_span(instance.def_id()), "Naked asm is not supported in JIT mode");
     }
diff --git a/compiler/rustc_codegen_cranelift/src/driver/mod.rs b/compiler/rustc_codegen_cranelift/src/driver/mod.rs
index ffd47cace38..8f83c30b598 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/mod.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/mod.rs
@@ -35,7 +35,7 @@ fn predefine_mono_items<'tcx>(
                         is_compiler_builtins,
                     );
                     let is_naked = tcx
-                        .codegen_fn_attrs(instance.def_id())
+                        .codegen_instance_attrs(instance.def)
                         .flags
                         .contains(CodegenFnAttrFlags::NAKED);
                     module
diff --git a/compiler/rustc_codegen_gcc/src/attributes.rs b/compiler/rustc_codegen_gcc/src/attributes.rs
index bf0927dc590..7a1ae6ca9c8 100644
--- a/compiler/rustc_codegen_gcc/src/attributes.rs
+++ b/compiler/rustc_codegen_gcc/src/attributes.rs
@@ -87,7 +87,7 @@ pub fn from_fn_attrs<'gcc, 'tcx>(
     #[cfg_attr(not(feature = "master"), allow(unused_variables))] func: Function<'gcc>,
     instance: ty::Instance<'tcx>,
 ) {
-    let codegen_fn_attrs = cx.tcx.codegen_fn_attrs(instance.def_id());
+    let codegen_fn_attrs = cx.tcx.codegen_instance_attrs(instance.def);
 
     #[cfg(feature = "master")]
     {
diff --git a/compiler/rustc_codegen_gcc/src/callee.rs b/compiler/rustc_codegen_gcc/src/callee.rs
index 189ac7cd779..e7ca95af594 100644
--- a/compiler/rustc_codegen_gcc/src/callee.rs
+++ b/compiler/rustc_codegen_gcc/src/callee.rs
@@ -105,7 +105,7 @@ pub fn get_fn<'gcc, 'tcx>(cx: &CodegenCx<'gcc, 'tcx>, instance: Instance<'tcx>)
             let is_hidden = if is_generic {
                 // This is a monomorphization of a generic function.
                 if !(cx.tcx.sess.opts.share_generics()
-                    || tcx.codegen_fn_attrs(instance_def_id).inline
+                    || tcx.codegen_instance_attrs(instance.def).inline
                         == rustc_attr_data_structures::InlineAttr::Never)
                 {
                     // When not sharing generics, all instances are in the same
diff --git a/compiler/rustc_codegen_gcc/src/mono_item.rs b/compiler/rustc_codegen_gcc/src/mono_item.rs
index 539e3ac8507..51f35cbdee4 100644
--- a/compiler/rustc_codegen_gcc/src/mono_item.rs
+++ b/compiler/rustc_codegen_gcc/src/mono_item.rs
@@ -53,7 +53,7 @@ impl<'gcc, 'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
         self.linkage.set(base::linkage_to_gcc(linkage));
         let decl = self.declare_fn(symbol_name, fn_abi);
-        //let attrs = self.tcx.codegen_fn_attrs(instance.def_id());
+        //let attrs = self.tcx.codegen_instance_attrs(instance.def);
 
         attributes::from_fn_attrs(self, decl, instance);
 
diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index 1ea5a062254..c32f11b27f3 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -344,7 +344,7 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     llfn: &'ll Value,
     instance: ty::Instance<'tcx>,
 ) {
-    let codegen_fn_attrs = cx.tcx.codegen_fn_attrs(instance.def_id());
+    let codegen_fn_attrs = cx.tcx.codegen_instance_attrs(instance.def);
 
     let mut to_add = SmallVec::<[_; 16]>::new();
 
diff --git a/compiler/rustc_codegen_llvm/src/callee.rs b/compiler/rustc_codegen_llvm/src/callee.rs
index 6d68eca60af..5a3dd90ab24 100644
--- a/compiler/rustc_codegen_llvm/src/callee.rs
+++ b/compiler/rustc_codegen_llvm/src/callee.rs
@@ -102,7 +102,7 @@ pub(crate) fn get_fn<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, instance: Instance<'t
         let is_hidden = if is_generic {
             // This is a monomorphization of a generic function.
             if !(cx.tcx.sess.opts.share_generics()
-                || tcx.codegen_fn_attrs(instance_def_id).inline
+                || tcx.codegen_instance_attrs(instance.def).inline
                     == rustc_attr_data_structures::InlineAttr::Never)
             {
                 // When not sharing generics, all instances are in the same
diff --git a/compiler/rustc_codegen_llvm/src/mono_item.rs b/compiler/rustc_codegen_llvm/src/mono_item.rs
index 8f70270f203..f9edaded60d 100644
--- a/compiler/rustc_codegen_llvm/src/mono_item.rs
+++ b/compiler/rustc_codegen_llvm/src/mono_item.rs
@@ -55,8 +55,8 @@ impl<'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
         let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
         let lldecl = self.declare_fn(symbol_name, fn_abi, Some(instance));
         llvm::set_linkage(lldecl, base::linkage_to_llvm(linkage));
-        let attrs = self.tcx.codegen_fn_attrs(instance.def_id());
-        base::set_link_section(lldecl, attrs);
+        let attrs = self.tcx.codegen_instance_attrs(instance.def);
+        base::set_link_section(lldecl, &attrs);
         if (linkage == Linkage::LinkOnceODR || linkage == Linkage::WeakODR)
             && self.tcx.sess.target.supports_comdat()
         {
diff --git a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
index 025f5fb54f4..b8f635ab781 100644
--- a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
@@ -356,7 +356,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             LocalRef::Operand(operand) => {
                 // Don't spill operands onto the stack in naked functions.
                 // See: https://github.com/rust-lang/rust/issues/42779
-                let attrs = bx.tcx().codegen_fn_attrs(self.instance.def_id());
+                let attrs = bx.tcx().codegen_instance_attrs(self.instance.def);
                 if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
                     return;
                 }
diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index fa69820d5d2..50d0f910744 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -390,9 +390,8 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     let mut num_untupled = None;
 
-    let codegen_fn_attrs = bx.tcx().codegen_fn_attrs(fx.instance.def_id());
-    let naked = codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED);
-    if naked {
+    let codegen_fn_attrs = bx.tcx().codegen_instance_attrs(fx.instance.def);
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
         return vec![];
     }
 
diff --git a/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs b/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
index beaf8950978..42e435cf0a3 100644
--- a/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
@@ -128,7 +128,7 @@ fn prefix_and_suffix<'tcx>(
     let is_arm = tcx.sess.target.arch == "arm";
     let is_thumb = tcx.sess.unstable_target_features.contains(&sym::thumb_mode);
 
-    let attrs = tcx.codegen_fn_attrs(instance.def_id());
+    let attrs = tcx.codegen_instance_attrs(instance.def);
     let link_section = attrs.link_section.map(|symbol| symbol.as_str().to_string());
 
     // If no alignment is specified, an alignment of 4 bytes is used.
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index e90463aacc8..2587e89417a 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -611,11 +611,19 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     let fn_abi = bx.fn_abi_of_instance(instance, ty::List::empty());
                     let fn_ty = bx.fn_decl_backend_type(fn_abi);
                     let fn_attrs = if bx.tcx().def_kind(instance.def_id()).has_codegen_attrs() {
-                        Some(bx.tcx().codegen_fn_attrs(instance.def_id()))
+                        Some(bx.tcx().codegen_instance_attrs(instance.def))
                     } else {
                         None
                     };
-                    bx.call(fn_ty, fn_attrs, Some(fn_abi), fn_ptr, &[], None, Some(instance))
+                    bx.call(
+                        fn_ty,
+                        fn_attrs.as_deref(),
+                        Some(fn_abi),
+                        fn_ptr,
+                        &[],
+                        None,
+                        Some(instance),
+                    )
                 } else {
                     bx.get_static(def_id)
                 };
diff --git a/compiler/rustc_codegen_ssa/src/mono_item.rs b/compiler/rustc_codegen_ssa/src/mono_item.rs
index 7b4268abe4b..b9040c330fb 100644
--- a/compiler/rustc_codegen_ssa/src/mono_item.rs
+++ b/compiler/rustc_codegen_ssa/src/mono_item.rs
@@ -41,12 +41,8 @@ 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)
-                {
+                let flags = cx.tcx().codegen_instance_attrs(instance.def).flags;
+                if flags.contains(CodegenFnAttrFlags::NAKED) {
                     naked_asm::codegen_naked_asm::<Bx::CodegenCx>(cx, instance, item_data);
                 } else {
                     base::codegen_instance::<Bx>(cx, instance);
@@ -75,7 +71,7 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
                 cx.predefine_static(def_id, linkage, visibility, symbol_name);
             }
             MonoItem::Fn(instance) => {
-                let attrs = cx.tcx().codegen_fn_attrs(instance.def_id());
+                let attrs = cx.tcx().codegen_instance_attrs(instance.def);
 
                 if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
                     // do not define this function; it will become a global assembly block
diff --git a/compiler/rustc_const_eval/src/interpret/memory.rs b/compiler/rustc_const_eval/src/interpret/memory.rs
index b2cbbc21430..47228de5213 100644
--- a/compiler/rustc_const_eval/src/interpret/memory.rs
+++ b/compiler/rustc_const_eval/src/interpret/memory.rs
@@ -936,7 +936,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         if let Some(fn_val) = self.get_fn_alloc(id) {
             let align = match fn_val {
                 FnVal::Instance(instance) => {
-                    self.tcx.codegen_fn_attrs(instance.def_id()).alignment.unwrap_or(Align::ONE)
+                    self.tcx.codegen_instance_attrs(instance.def).alignment.unwrap_or(Align::ONE)
                 }
                 // Machine-specific extra functions currently do not support alignment restrictions.
                 FnVal::Other(_) => Align::ONE,
diff --git a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
index 2f16d385efb..6eae3b51e29 100644
--- a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
+++ b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use rustc_abi::Align;
 use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
 use rustc_attr_data_structures::{InlineAttr, InstructionSetAttr, OptimizeAttr};
@@ -6,6 +8,26 @@ use rustc_span::Symbol;
 use rustc_target::spec::SanitizerSet;
 
 use crate::mir::mono::Linkage;
+use crate::ty::{InstanceKind, TyCtxt};
+
+impl<'tcx> TyCtxt<'tcx> {
+    pub fn codegen_instance_attrs(
+        self,
+        instance_kind: InstanceKind<'_>,
+    ) -> Cow<'tcx, CodegenFnAttrs> {
+        let mut attrs = Cow::Borrowed(self.codegen_fn_attrs(instance_kind.def_id()));
+
+        // Drop the `#[naked]` attribute on non-item `InstanceKind`s, like the shims that
+        // are generated for indirect function calls.
+        if !matches!(instance_kind, InstanceKind::Item(_)) {
+            if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
+                attrs.to_mut().flags.remove(CodegenFnAttrFlags::NAKED);
+            }
+        }
+
+        attrs
+    }
+}
 
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, Debug)]
 pub struct CodegenFnAttrs {
diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
index 47ba850d50d..2d7ddd105bd 100644
--- a/compiler/rustc_middle/src/mir/mono.rs
+++ b/compiler/rustc_middle/src/mir/mono.rs
@@ -152,7 +152,7 @@ impl<'tcx> MonoItem<'tcx> {
         // If the function is #[naked] or contains any other attribute that requires exactly-once
         // instantiation:
         // We emit an unused_attributes lint for this case, which should be kept in sync if possible.
-        let codegen_fn_attrs = tcx.codegen_fn_attrs(instance.def_id());
+        let codegen_fn_attrs = tcx.codegen_instance_attrs(instance.def);
         if codegen_fn_attrs.contains_extern_indicator()
             || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED)
         {
@@ -219,7 +219,7 @@ impl<'tcx> MonoItem<'tcx> {
         // functions the same as those that unconditionally get LocalCopy codegen. It's only when
         // we get here that we can at least not codegen a #[inline(never)] generic function in all
         // of our CGUs.
-        if let InlineAttr::Never = tcx.codegen_fn_attrs(instance.def_id()).inline
+        if let InlineAttr::Never = codegen_fn_attrs.inline
             && self.is_generic_fn()
         {
             return InstantiationMode::GloballyShared { may_conflict: true };
@@ -234,14 +234,13 @@ impl<'tcx> MonoItem<'tcx> {
     }
 
     pub fn explicit_linkage(&self, tcx: TyCtxt<'tcx>) -> Option<Linkage> {
-        let def_id = match *self {
-            MonoItem::Fn(ref instance) => instance.def_id(),
-            MonoItem::Static(def_id) => def_id,
+        let instance_kind = match *self {
+            MonoItem::Fn(ref instance) => instance.def,
+            MonoItem::Static(def_id) => InstanceKind::Item(def_id),
             MonoItem::GlobalAsm(..) => return None,
         };
 
-        let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
-        codegen_fn_attrs.linkage
+        tcx.codegen_instance_attrs(instance_kind).linkage
     }
 
     /// Returns `true` if this instance is instantiable - whether it has no unsatisfied
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 935cc889574..3c1be069c9e 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1505,6 +1505,15 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    /// Returns the `CodegenFnAttrs` for the item at `def_id`.
+    ///
+    /// If possible, use `tcx.codegen_instance_attrs` instead. That function takes the
+    /// instance kind into account.
+    ///
+    /// For example, the `#[naked]` attribute should be applied for `InstanceKind::Item`,
+    /// but should not be applied if the instance kind is `InstanceKind::ReifyShim`.
+    /// Using this query would include the attribute regardless of the actual instance
+    /// kind at the call site.
     query codegen_fn_attrs(def_id: DefId) -> &'tcx CodegenFnAttrs {
         desc { |tcx| "computing codegen attributes of `{}`", tcx.def_path_str(def_id) }
         arena_cache
diff --git a/compiler/rustc_symbol_mangling/src/lib.rs b/compiler/rustc_symbol_mangling/src/lib.rs
index a9bf5eae445..6bcb7f6e093 100644
--- a/compiler/rustc_symbol_mangling/src/lib.rs
+++ b/compiler/rustc_symbol_mangling/src/lib.rs
@@ -180,7 +180,7 @@ fn compute_symbol_name<'tcx>(
 
     // FIXME(eddyb) Precompute a custom symbol name based on attributes.
     let attrs = if tcx.def_kind(def_id).has_codegen_attrs() {
-        tcx.codegen_fn_attrs(def_id)
+        &tcx.codegen_instance_attrs(instance.def)
     } else {
         CodegenFnAttrs::EMPTY
     };