Skip to content

Commit 29fbe3e

Browse files
tmandrytmandry
committed
Add control flow information to __rust_probestack
1 parent 0df0cf5 commit 29fbe3e

File tree

2 files changed

+144
-89
lines changed

2 files changed

+144
-89
lines changed

‎src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#![cfg_attr(feature = "compiler-builtins", compiler_builtins)]
22
#![feature(abi_unadjusted)]
33
#![feature(asm)]
4+
#![feature(global_asm)]
45
#![feature(cfg_target_has_atomic)]
56
#![feature(compiler_builtins)]
67
#![feature(core_intrinsics)]

‎src/probestack.rs

+143-89
Original file line numberDiff line numberDiff line change
@@ -41,95 +41,149 @@
4141
//! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
4242
//! be more than welcome to accept such a change!
4343
44-
#![cfg(not(windows))] // Windows already has builtins to do this
45-
46-
#[naked]
47-
#[no_mangle]
48-
#[cfg(all(target_arch = "x86_64", not(feature = "mangled-names")))]
49-
pub unsafe extern "C" fn __rust_probestack() {
50-
// Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
51-
// ensuring that if any pages are unmapped we'll make a page fault.
52-
//
53-
// The ABI here is that the stack frame size is located in `%eax`. Upon
54-
// return we're not supposed to modify `%esp` or `%eax`.
55-
asm!("
56-
pushq %rbp
57-
movq %rsp, %rbp
58-
59-
mov %rax,%r11 // duplicate %rax as we're clobbering %r11
60-
61-
// Main loop, taken in one page increments. We're decrementing rsp by
62-
// a page each time until there's less than a page remaining. We're
63-
// guaranteed that this function isn't called unless there's more than a
64-
// page needed.
65-
//
66-
// Note that we're also testing against `8(%rsp)` to account for the 8
67-
// bytes pushed on the stack orginally with our return address. Using
68-
// `8(%rsp)` simulates us testing the stack pointer in the caller's
69-
// context.
70-
71-
// It's usually called when %rax >= 0x1000, but that's not always true.
72-
// Dynamic stack allocation, which is needed to implement unsized
73-
// rvalues, triggers stackprobe even if %rax < 0x1000.
74-
// Thus we have to check %r11 first to avoid segfault.
75-
cmp $$0x1000,%r11
76-
jna 3f
77-
2:
78-
sub $$0x1000,%rsp
79-
test %rsp,8(%rsp)
80-
sub $$0x1000,%r11
81-
cmp $$0x1000,%r11
82-
ja 2b
83-
84-
3:
85-
// Finish up the last remaining stack space requested, getting the last
86-
// bits out of r11
87-
sub %r11,%rsp
88-
test %rsp,8(%rsp)
89-
90-
// Restore the stack pointer to what it previously was when entering
91-
// this function. The caller will readjust the stack pointer after we
92-
// return.
93-
add %rax,%rsp
94-
95-
leave
96-
ret
97-
" ::: "memory" : "volatile");
98-
::core::intrinsics::unreachable();
44+
#![cfg(not(feature = "mangled-names"))]
45+
// Windows already has builtins to do this.
46+
#![cfg(not(windows))]
47+
// We only define stack probing for these architectures today.
48+
#![cfg(any(target_arch = "x86_64", target_arch = "x86"))]
49+
50+
extern "C" {
51+
pub fn __rust_probestack();
9952
}
10053

101-
#[naked]
102-
#[no_mangle]
103-
#[cfg(all(target_arch = "x86", not(feature = "mangled-names")))]
104-
pub unsafe extern "C" fn __rust_probestack() {
105-
// This is the same as x86_64 above, only translated for 32-bit sizes. Note
106-
// that on Unix we're expected to restore everything as it was, this
107-
// function basically can't tamper with anything.
108-
//
109-
// The ABI here is the same as x86_64, except everything is 32-bits large.
110-
asm!("
111-
push %ebp
112-
mov %esp, %ebp
113-
push %ecx
114-
mov %eax,%ecx
115-
116-
cmp $$0x1000,%ecx
117-
jna 3f
118-
2:
119-
sub $$0x1000,%esp
120-
test %esp,8(%esp)
121-
sub $$0x1000,%ecx
122-
cmp $$0x1000,%ecx
123-
ja 2b
124-
125-
3:
126-
sub %ecx,%esp
127-
test %esp,8(%esp)
128-
129-
add %eax,%esp
130-
pop %ecx
131-
leave
132-
ret
133-
" ::: "memory" : "volatile");
134-
::core::intrinsics::unreachable();
54+
// A wrapper for our implementation of __rust_probestack, which allows us to
55+
// keep the assembly inline while controlling all CFI directives in the assembly
56+
// emitted for the function.
57+
//
58+
// This is the ELF version.
59+
#[cfg(not(target_vendor = "apple"))]
60+
macro_rules! define_rust_probestack {
61+
($body: expr) => {
62+
concat!(
63+
"
64+
.pushsection .text.__rust_probestack
65+
.globl __rust_probestack
66+
.type __rust_probestack, @function
67+
__rust_probestack:
68+
",
69+
$body,
70+
"
71+
.size __rust_probestack, . - __rust_probestack
72+
.popsection
73+
"
74+
)
75+
};
76+
}
77+
78+
// Same as above, but for Mach-O.
79+
#[cfg(target_vendor = "apple")]
80+
macro_rules! define_rust_probestack {
81+
($body: expr) => {
82+
concat!(
83+
"
84+
.globl ___rust_probestack
85+
___rust_probestack:
86+
",
87+
$body
88+
)
89+
};
13590
}
91+
92+
// Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
93+
// ensuring that if any pages are unmapped we'll make a page fault.
94+
//
95+
// The ABI here is that the stack frame size is located in `%rax`. Upon
96+
// return we're not supposed to modify `%rsp` or `%rax`.
97+
#[cfg(target_arch = "x86_64")]
98+
global_asm!(define_rust_probestack!(
99+
"
100+
.cfi_startproc
101+
pushq %rbp
102+
.cfi_adjust_cfa_offset 8
103+
.cfi_offset %rbp, -16
104+
movq %rsp, %rbp
105+
.cfi_def_cfa_register %rbp
106+
107+
mov %rax,%r11 // duplicate %rax as we're clobbering %r11
108+
109+
// Main loop, taken in one page increments. We're decrementing rsp by
110+
// a page each time until there's less than a page remaining. We're
111+
// guaranteed that this function isn't called unless there's more than a
112+
// page needed.
113+
//
114+
// Note that we're also testing against `8(%rsp)` to account for the 8
115+
// bytes pushed on the stack orginally with our return address. Using
116+
// `8(%rsp)` simulates us testing the stack pointer in the caller's
117+
// context.
118+
119+
// It's usually called when %rax >= 0x1000, but that's not always true.
120+
// Dynamic stack allocation, which is needed to implement unsized
121+
// rvalues, triggers stackprobe even if %rax < 0x1000.
122+
// Thus we have to check %r11 first to avoid segfault.
123+
cmp $0x1000,%r11
124+
jna 3f
125+
2:
126+
sub $0x1000,%rsp
127+
test %rsp,8(%rsp)
128+
sub $0x1000,%r11
129+
cmp $0x1000,%r11
130+
ja 2b
131+
132+
3:
133+
// Finish up the last remaining stack space requested, getting the last
134+
// bits out of r11
135+
sub %r11,%rsp
136+
test %rsp,8(%rsp)
137+
138+
// Restore the stack pointer to what it previously was when entering
139+
// this function. The caller will readjust the stack pointer after we
140+
// return.
141+
add %rax,%rsp
142+
143+
leave
144+
.cfi_def_cfa_register %rsp
145+
.cfi_adjust_cfa_offset -8
146+
ret
147+
.cfi_endproc
148+
"
149+
));
150+
151+
#[cfg(target_arch = "x86")]
152+
// This is the same as x86_64 above, only translated for 32-bit sizes. Note
153+
// that on Unix we're expected to restore everything as it was, this
154+
// function basically can't tamper with anything.
155+
//
156+
// The ABI here is the same as x86_64, except everything is 32-bits large.
157+
global_asm!(define_rust_probestack!(
158+
"
159+
.cfi_startproc
160+
push %ebp
161+
.cfi_adjust_cfa_offset 4
162+
.cfi_offset %ebp, -8
163+
mov %esp, %ebp
164+
.cfi_def_cfa_register %ebp
165+
push %ecx
166+
mov %eax,%ecx
167+
168+
cmp $0x1000,%ecx
169+
jna 3f
170+
2:
171+
sub $0x1000,%esp
172+
test %esp,8(%esp)
173+
sub $0x1000,%ecx
174+
cmp $0x1000,%ecx
175+
ja 2b
176+
177+
3:
178+
sub %ecx,%esp
179+
test %esp,8(%esp)
180+
181+
add %eax,%esp
182+
pop %ecx
183+
leave
184+
.cfi_def_cfa_register %esp
185+
.cfi_adjust_cfa_offset -4
186+
ret
187+
.cfi_endproc
188+
"
189+
));

0 commit comments

Comments
 (0)