Skip to content

use #[naked] for __rust_probestack #897

New issue

Have a question about this project? No Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

No Sign up for GitHub

By clicking “No Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? No Sign in to your account

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

use #[naked] for __rust_probestack #897

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
260c31d
use `#[naked]` for `__rust_probestack`
folkertdev May 1, 2025
a548e1f
merge the sgx/fortanix `__rust_probestack` into the general `x86_64` one
folkertdev May 1, 2025
b1934ae
add note about why this function is unsafe
folkertdev May 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions compiler-builtins/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
215 changes: 54 additions & 161 deletions compiler-builtins/src/probestack.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,154 +49,39 @@
// We only define stack probing for these architectures today.
#![cfg(any(target_arch = "x86_64", target_arch = "x86"))]

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`.
//
// Any changes to this function should be replicated to the SGX version below.
#[cfg(all(
target_arch = "x86_64",
not(all(target_env = "sgx", target_vendor = "fortanix"))
))]
core::arch::global_asm!(
define_rust_probestack!(
"
.cfi_startproc
pushq %rbp
.cfi_adjust_cfa_offset 8
.cfi_offset %rbp, -16
movq %rsp, %rbp
.cfi_def_cfa_register %rbp

mov %rax,%r11 // duplicate %rax as we're clobbering %r11

// Main loop, taken in one page increments. We're decrementing rsp by
// a page each time until there's less than a page remaining. We're
// guaranteed that this function isn't called unless there's more than a
// page needed.
//
// Note that we're also testing against `8(%rsp)` to account for the 8
// bytes pushed on the stack orginally with our return address. Using
// `8(%rsp)` simulates us testing the stack pointer in the caller's
// context.

// It's usually called when %rax >= 0x1000, but that's not always true.
// Dynamic stack allocation, which is needed to implement unsized
// rvalues, triggers stackprobe even if %rax < 0x1000.
// Thus we have to check %r11 first to avoid segfault.
cmp $0x1000,%r11
jna 3f
2:
sub $0x1000,%rsp
test %rsp,8(%rsp)
sub $0x1000,%r11
cmp $0x1000,%r11
ja 2b

3:
// Finish up the last remaining stack space requested, getting the last
// bits out of r11
sub %r11,%rsp
test %rsp,8(%rsp)

// Restore the stack pointer to what it previously was when entering
// this function. The caller will readjust the stack pointer after we
// return.
add %rax,%rsp

leave
.cfi_def_cfa_register %rsp
.cfi_adjust_cfa_offset -8
ret
.cfi_endproc
"
),
options(att_syntax)
);
#[cfg(target_arch = "x86_64")]
#[unsafe(naked)]
#[rustc_std_internal_symbol]
pub unsafe extern "C" fn __rust_probestack() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you mention in the doc comment that the custom ABI is the reason this is unsafe?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you mean in the module doc comment? We can add some specific details (about what registers are used etc) if we add a doc comment to the function itself.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was referring to the comments above each of the functions, I guess they're not actually doc comments though they could be. Something like (hence the `unsafe`, since the ABI does not actually match `extern "C"`) in the paragraph describing ABI is sufficient, unless you want to add more. Just to have a note that it isn't supposed to be extern "C".

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That reminded me of a few other cases we have where naked functions are extern "C" but that doesn't really make sense. Seems like a somewhat common issue, I wrote up rust-lang/rust#140566 as a possible way to make this more clear (not for this pr of course).

#[cfg(not(all(target_env = "sgx", target_vendor = "fortanix")))]
macro_rules! ret {
() => {
"ret"
};
}

#[cfg(all(target_env = "sgx", target_vendor = "fortanix"))]
macro_rules! ret {
// for this target, [manually patch for LVI].
//
// [manually patch for LVI]: https://software.intel.com/security-software-guidance/insights/deep-dive-load-value-injection#specialinstructions
() => {
"
pop %r11
lfence
jmp *%r11
"
};
}

// This function is the same as above, except that some instructions are
// [manually patched for LVI].
//
// [manually patched for LVI]: https://software.intel.com/security-software-guidance/insights/deep-dive-load-value-injection#specialinstructions
#[cfg(all(
target_arch = "x86_64",
all(target_env = "sgx", target_vendor = "fortanix")
))]
core::arch::global_asm!(
define_rust_probestack!(
core::arch::naked_asm!(
"
.cfi_startproc
pushq %rbp
Expand Down Expand Up @@ -244,23 +129,27 @@ core::arch::global_asm!(
leave
.cfi_def_cfa_register %rsp
.cfi_adjust_cfa_offset -8
pop %r11
lfence
jmp *%r11
",
ret!(),
"
.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
Expand Down Expand Up @@ -291,24 +180,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
Expand Down Expand Up @@ -344,7 +237,7 @@ core::arch::global_asm!(
.cfi_adjust_cfa_offset -4
ret
.cfi_endproc
"
),
options(att_syntax)
);
",
options(att_syntax)
)
}
Loading