use #[naked] for __rust_probestack

Author: folkertdev

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -5,6 +5,7 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFunctionEntry};
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
+use rustc_symbol_mangling::mangle_internal_symbol;
 use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
 use smallvec::SmallVec;
 
@@ -256,11 +257,11 @@ fn probestack_attr<'ll>(cx: &CodegenCx<'ll, '_>) -> Option<&'ll Attribute> {
         StackProbeType::Inline => "inline-asm",
         // Flag our internal `__rust_probestack` function as the stack probe symbol.
         // This is defined in the `compiler-builtins` crate for each architecture.
-        StackProbeType::Call => "__rust_probestack",
+        StackProbeType::Call => &mangle_internal_symbol(cx.tcx, "__rust_probestack"),
         // Pick from the two above based on the LLVM version.
         StackProbeType::InlineOrCall { min_llvm_version_for_inline } => {
             if llvm_util::get_version() < min_llvm_version_for_inline {
-                "__rust_probestack"
+                &mangle_internal_symbol(cx.tcx, "__rust_probestack")
             } else {
                 "inline-asm"
             }

library/compiler-builtins/compiler-builtins/src/lib.rs

@@ -8,6 +8,7 @@
 #![feature(linkage)]
 #![feature(naked_functions)]
 #![feature(repr_simd)]
+#![feature(rustc_attrs)]
 #![cfg_attr(f16_enabled, feature(f16))]
 #![cfg_attr(f128_enabled, feature(f128))]
 #![no_builtins]

library/compiler-builtins/compiler-builtins/src/probestack.rs

@@ -49,82 +49,11 @@
 // We only define stack probing for these architectures today.
 #![cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 
-// SAFETY: defined in this module.
-// FIXME(extern_custom): the ABI is not correct.
-unsafe extern "C" {
-    pub fn __rust_probestack();
-}
-
-// A wrapper for our implementation of __rust_probestack, which allows us to
-// keep the assembly inline while controlling all CFI directives in the assembly
-// emitted for the function.
-//
-// This is the ELF version.
-#[cfg(not(any(target_vendor = "apple", target_os = "uefi")))]
-macro_rules! define_rust_probestack {
-    ($body: expr) => {
-        concat!(
-            "
-            .pushsection .text.__rust_probestack
-            .globl __rust_probestack
-            .type  __rust_probestack, @function
-            .hidden __rust_probestack
-        __rust_probestack:
-            ",
-            $body,
-            "
-            .size __rust_probestack, . - __rust_probestack
-            .popsection
-            "
-        )
-    };
-}
-
-#[cfg(all(target_os = "uefi", target_arch = "x86_64"))]
-macro_rules! define_rust_probestack {
-    ($body: expr) => {
-        concat!(
-            "
-            .globl __rust_probestack
-        __rust_probestack:
-            ",
-            $body
-        )
-    };
-}
-
-// Same as above, but for Mach-O. Note that the triple underscore
-// is deliberate
-#[cfg(target_vendor = "apple")]
-macro_rules! define_rust_probestack {
-    ($body: expr) => {
-        concat!(
-            "
-            .globl ___rust_probestack
-        ___rust_probestack:
-            ",
-            $body
-        )
-    };
-}
-
-// In UEFI x86 arch, triple underscore is deliberate.
-#[cfg(all(target_os = "uefi", target_arch = "x86"))]
-macro_rules! define_rust_probestack {
-    ($body: expr) => {
-        concat!(
-            "
-            .globl ___rust_probestack
-        ___rust_probestack:
-            ",
-            $body
-        )
-    };
-}
-
 // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
 // ensuring that if any pages are unmapped we'll make a page fault.
 //
+// This function is unsafe because it uses a custom ABI, it does not actually match `extern "C"`.
+//
 // The ABI here is that the stack frame size is located in `%rax`. Upon
 // return we're not supposed to modify `%rsp` or `%rax`.
 //
@@ -133,8 +62,10 @@ macro_rules! define_rust_probestack {
     target_arch = "x86_64",
     not(all(target_env = "sgx", target_vendor = "fortanix"))
 ))]
-core::arch::global_asm!(
-    define_rust_probestack!(
+#[unsafe(naked)]
+#[rustc_std_internal_symbol]
+pub unsafe extern "C" fn __rust_probestack() {
+    core::arch::naked_asm!(
         "
     .cfi_startproc
     pushq  %rbp
@@ -184,10 +115,10 @@ core::arch::global_asm!(
     .cfi_adjust_cfa_offset -8
     ret
     .cfi_endproc
-    "
-    ),
-    options(att_syntax)
-);
+    ",
+        options(att_syntax)
+    )
+}
 
 // This function is the same as above, except that some instructions are
 // [manually patched for LVI].
@@ -197,8 +128,10 @@ core::arch::global_asm!(
     target_arch = "x86_64",
     all(target_env = "sgx", target_vendor = "fortanix")
 ))]
-core::arch::global_asm!(
-    define_rust_probestack!(
+#[unsafe(naked)]
+#[no_mangle]
+pub unsafe extern "C" fn __rust_probestack() {
+    core::arch::naked_asm!(
         "
     .cfi_startproc
     pushq  %rbp
@@ -250,19 +183,23 @@ core::arch::global_asm!(
     lfence
     jmp *%r11
     .cfi_endproc
-    "
-    ),
-    options(att_syntax)
-);
+    ",
+        options(att_syntax)
+    )
+}
 
 #[cfg(all(target_arch = "x86", not(target_os = "uefi")))]
 // This is the same as x86_64 above, only translated for 32-bit sizes. Note
 // that on Unix we're expected to restore everything as it was, this
 // function basically can't tamper with anything.
 //
+// This function is unsafe because it uses a custom ABI, it does not actually match `extern "C"`.
+//
 // The ABI here is the same as x86_64, except everything is 32-bits large.
-core::arch::global_asm!(
-    define_rust_probestack!(
+#[unsafe(naked)]
+#[rustc_std_internal_symbol]
+pub unsafe extern "C" fn __rust_probestack() {
+    core::arch::naked_asm!(
         "
     .cfi_startproc
     push   %ebp
@@ -293,24 +230,28 @@ core::arch::global_asm!(
     .cfi_adjust_cfa_offset -4
     ret
     .cfi_endproc
-    "
-    ),
-    options(att_syntax)
-);
+    ",
+        options(att_syntax)
+    )
+}
 
 #[cfg(all(target_arch = "x86", target_os = "uefi"))]
 // UEFI target is windows like target. LLVM will do _chkstk things like windows.
 // probestack function will also do things like _chkstk in MSVC.
 // So we need to sub %ax %sp in probestack when arch is x86.
 //
+// This function is unsafe because it uses a custom ABI, it does not actually match `extern "C"`.
+//
 // REF: Rust commit(74e80468347)
 // rust\src\llvm-project\llvm\lib\Target\X86\X86FrameLowering.cpp: 805
 // Comments in LLVM:
 //   MSVC x32's _chkstk and cygwin/mingw's _alloca adjust %esp themselves.
 //   MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust %rsp
 //   themselves.
-core::arch::global_asm!(
-    define_rust_probestack!(
+#[unsafe(naked)]
+#[rustc_std_internal_symbol]
+pub unsafe extern "C" fn __rust_probestack() {
+    core::arch::naked_asm!(
         "
     .cfi_startproc
     push   %ebp
@@ -346,7 +287,7 @@ core::arch::global_asm!(
     .cfi_adjust_cfa_offset -4
     ret
     .cfi_endproc
-    "
-    ),
-    options(att_syntax)
-);
+    ",
+        options(att_syntax)
+    )
+}