Skip to content

Commit

Permalink
Fix unsafe
Browse files Browse the repository at this point in the history
  • Loading branch information
@byeongkeunahn
byeongkeunahn committed Sep 7, 2024
1 parent 551b7b0 commit 19964cd
Show file tree
Hide file tree
Showing 3 changed files with 154 additions and 134 deletions.
192 changes: 101 additions & 91 deletions basm-std/src/platform/codegen.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,36 +56,38 @@ compile_error!("AArch64 (aarch64-apple-darwin) is only supported for local execu
#[no_mangle]
#[naked]
pub unsafe extern "win64" fn _basm_start() -> ! {
// AMD64 System V ABI requires RSP to be aligned
// on the 16-byte boundary BEFORE `call` instruction.
// However, when called as the entrypoint by the Linux OS,
// RSP will be 16-byte aligned AFTER `call` instruction.
asm!(
"clc", // CF=0 (running without loader) / CF=1 (running with loader)
"mov rbx, rcx", // Save PLATFORM_DATA table
"jnc 2f",
"test rbx, rbx",
"jz 2f",
"jmp 3f",
"2:",
"sub rsp, 72", // 16 + 72 + 8 = 96 = 16*6 -> stack alignment preserved
"push 3", // env_flags = 3 (ENV_FLAGS_LINUX_STYLE_CHKSTK | ENV_FLAGS_NATIVE)
"push 2", // env_id = 2 (ENV_ID_LINUX)
"lea rbx, [rsp]", // rbx = PLATFORM_DATA table
"3:",
"push rcx", // short form of "sub rsp, 8"
"lea rdi, [rip + __ehdr_start]",
"lea rsi, [rip + _DYNAMIC]",
"mov QWORD PTR [rbx + 32], rdi", // overwrite ptr_alloc_rwx with in-memory ImageBase
"call {0}",
"mov rdi, rbx",
"call {1}",
"pop rcx", // short form of "add rsp, 8"
"ret",
sym loader::amd64_elf::relocate,
sym _start_rust,
options(noreturn)
);
unsafe {
// AMD64 System V ABI requires RSP to be aligned
// on the 16-byte boundary BEFORE `call` instruction.
// However, when called as the entrypoint by the Linux OS,
// RSP will be 16-byte aligned AFTER `call` instruction.
asm!(
"clc", // CF=0 (running without loader) / CF=1 (running with loader)
"mov rbx, rcx", // Save PLATFORM_DATA table
"jnc 2f",
"test rbx, rbx",
"jz 2f",
"jmp 3f",
"2:",
"sub rsp, 72", // 16 + 72 + 8 = 96 = 16*6 -> stack alignment preserved
"push 3", // env_flags = 3 (ENV_FLAGS_LINUX_STYLE_CHKSTK | ENV_FLAGS_NATIVE)
"push 2", // env_id = 2 (ENV_ID_LINUX)
"lea rbx, [rsp]", // rbx = PLATFORM_DATA table
"3:",
"push rcx", // short form of "sub rsp, 8"
"lea rdi, [rip + __ehdr_start]",
"lea rsi, [rip + _DYNAMIC]",
"mov QWORD PTR [rbx + 32], rdi", // overwrite ptr_alloc_rwx with in-memory ImageBase
"call {0}",
"mov rdi, rbx",
"call {1}",
"pop rcx", // short form of "add rsp, 8"
"ret",
sym loader::amd64_elf::relocate,
sym _start_rust,
options(noreturn)
);
}
}

#[cfg(target_os = "windows")]
Expand Down Expand Up @@ -177,72 +179,78 @@ pub unsafe extern "win64" fn _basm_start() -> ! {
}

#[cfg(target_arch = "x86")]
#[no_mangle]
#[unsafe(no_mangle)]
#[naked]
#[link_section = ".data"]
unsafe extern "cdecl" fn _get_start_offset() -> ! {
asm!("lea eax, [_basm_start]", "ret", options(noreturn));
unsafe {
asm!("lea eax, [_basm_start]", "ret", options(noreturn));
}
}

#[cfg(target_arch = "x86")]
#[no_mangle]
#[unsafe(no_mangle)]
#[naked]
#[link_section = ".data"]
unsafe extern "cdecl" fn _get_dynamic_section_offset() -> ! {
asm!("lea eax, [_DYNAMIC]", "ret", options(noreturn));
unsafe {
asm!("lea eax, [_DYNAMIC]", "ret", options(noreturn));
}
}

#[cfg(target_arch = "x86")]
#[no_mangle]
#[unsafe(no_mangle)]
#[naked]
pub unsafe extern "cdecl" fn _basm_start() -> ! {
// i386 System V ABI requires ESP to be aligned
// on the 16-byte boundary BEFORE `call` instruction
asm!(
"clc", // CF=0 (running without loader) / CF=1 (running with loader)
"jc 2f",
"sub esp, 44", // 44 = 40 + 4; PLATFORM_DATA ptr (4 bytes, pushed later) + PLATFORM_DATA (40 (+ 16 = 56 bytes)) + alignment (4 bytes wasted)
"push 0", // zero upper dword
"push 3", // env_flags = 3 (ENV_FLAGS_LINUX_STYLE_CHKSTK | ENV_FLAGS_NATIVE)
"push 0", // zero upper dword
"push 2", // env_id = 2 (ENV_ID_LINUX)
"mov edx, esp", // edx = PLATFORM_DATA table
"jmp 3f",
"2:",
"mov edx, DWORD PTR [esp + 4]", // edx = PLATFORM_DATA table
"push ebp",
"mov ebp, esp",
"and esp, 0xFFFFFFF0",
"sub esp, 12",
"3:",
"call 4f",
"4:",
"pop ecx", // ecx = _basm_start + 36 (obtained by counting the opcode size in bytes)
"push edx", // [esp + 0] = PLATFORM_DATA table
"call {2}", // eax = offset of _basm_start from the image base
"sub ecx, eax",
"sub ecx, 36", // ecx = the in-memory image base (i.e., __ehdr_start)
"call {3}", // eax = offset of _DYNAMIC table from the image base
"add eax, ecx", // eax = _DYNAMIC table
"sub esp, 8", // For stack alignment
"push eax",
"push ecx",
"call {0}",
"add esp, 16",
"call {1}",
"mov esp, ebp",
"pop ebp",
"ret",
sym loader::i686_elf::relocate,
sym _start_rust,
sym _get_start_offset,
sym _get_dynamic_section_offset,
options(noreturn)
);
unsafe {
// i386 System V ABI requires ESP to be aligned
// on the 16-byte boundary BEFORE `call` instruction
asm!(
"clc", // CF=0 (running without loader) / CF=1 (running with loader)
"jc 2f",
"sub esp, 44", // 44 = 40 + 4; PLATFORM_DATA ptr (4 bytes, pushed later) + PLATFORM_DATA (40 (+ 16 = 56 bytes)) + alignment (4 bytes wasted)
"push 0", // zero upper dword
"push 3", // env_flags = 3 (ENV_FLAGS_LINUX_STYLE_CHKSTK | ENV_FLAGS_NATIVE)
"push 0", // zero upper dword
"push 2", // env_id = 2 (ENV_ID_LINUX)
"mov edx, esp", // edx = PLATFORM_DATA table
"jmp 3f",
"2:",
"mov edx, DWORD PTR [esp + 4]", // edx = PLATFORM_DATA table
"push ebp",
"mov ebp, esp",
"and esp, 0xFFFFFFF0",
"sub esp, 12",
"3:",
"call 4f",
"4:",
"pop ecx", // ecx = _basm_start + 36 (obtained by counting the opcode size in bytes)
"push edx", // [esp + 0] = PLATFORM_DATA table
"call {2}", // eax = offset of _basm_start from the image base
"sub ecx, eax",
"sub ecx, 36", // ecx = the in-memory image base (i.e., __ehdr_start)
"call {3}", // eax = offset of _DYNAMIC table from the image base
"add eax, ecx", // eax = _DYNAMIC table
"sub esp, 8", // For stack alignment
"push eax",
"push ecx",
"call {0}",
"add esp, 16",
"call {1}",
"mov esp, ebp",
"pop ebp",
"ret",
sym loader::i686_elf::relocate,
sym _start_rust,
sym _get_start_offset,
sym _get_dynamic_section_offset,
options(noreturn)
);
}
}

#[cfg(target_arch = "wasm32")]
#[no_mangle]
#[unsafe(no_mangle)]
pub extern "C" fn _basm_start() {
let mut pd = platform::services::PlatformData {
env_id: platform::services::ENV_ID_WASM,
Expand All @@ -252,21 +260,23 @@ pub extern "C" fn _basm_start() {
}

#[cfg(target_arch = "aarch64")]
#[no_mangle]
#[unsafe(no_mangle)]
#[naked]
#[repr(align(8))]
pub unsafe extern "C" fn _basm_start() -> ! {
asm!(
"sub sp, sp, #96",
"mov x0, #4", // 4 = ENV_ID_MACOS
"str x0, [sp, #(8 * 0)]",
"mov x0, #2", // 2 = ENV_FLAGS_NATIVE
"str x0, [sp, #(8 * 1)]",
"mov x0, sp",
"bl {0}",
sym _start_rust,
options(noreturn)
)
unsafe {
asm!(
"sub sp, sp, #96",
"mov x0, #4", // 4 = ENV_ID_MACOS
"str x0, [sp, #(8 * 0)]",
"mov x0, #2", // 2 = ENV_FLAGS_NATIVE
"str x0, [sp, #(8 * 1)]",
"mov x0, sp",
"bl {0}",
sym _start_rust,
options(noreturn)
)
}
}

/* We prevent inlining solution::main, since if the user allocates
Expand Down
4 changes: 2 additions & 2 deletions basm-std/src/platform/os/linux.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ pub mod syscall {
}
#[cfg(not(target_arch = "x86_64"))]
pub unsafe fn syscall1(call_id: usize, arg0: usize) -> usize {
syscall(call_id, arg0, 0, 0, 0, 0, 0)
unsafe { syscall(call_id, arg0, 0, 0, 0, 0, 0) }
}
#[cfg(target_arch = "x86_64")]
#[inline(always)]
Expand All @@ -97,7 +97,7 @@ pub mod syscall {
arg1: usize,
arg2: usize,
) -> usize {
syscall(call_id, arg0, arg1, arg2, 0, 0, 0)
unsafe { syscall(call_id, arg0, arg1, arg2, 0, 0, 0) }
}
#[cfg(target_arch = "x86_64")]
#[inline(always)]
Expand Down
92 changes: 51 additions & 41 deletions basm-std/src/platform/os/macos.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pub mod syscall {
}

#[link(name = "System", kind = "dylib")]
extern "C" {
unsafe extern "C" {
pub fn mmap(
addr: *const u8,
len: usize,
Expand All @@ -40,73 +40,83 @@ pub mod syscall {
static mut DLMALLOC: dlmalloc::Dlmalloc<dlmalloc_macos::System> =
dlmalloc::Dlmalloc::new(dlmalloc_macos::System::new());
unsafe fn dlmalloc_alloc(size: usize, align: usize) -> *mut u8 {
DLMALLOC.memalign(align, size)
unsafe {
DLMALLOC.memalign(align, size)
}
}
unsafe fn dlmalloc_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
let ptr = DLMALLOC.memalign(align, size);
if !ptr.is_null() && DLMALLOC.calloc_must_clear(ptr) {
core::ptr::write_bytes(ptr, 0, size);
unsafe {
let ptr = DLMALLOC.memalign(align, size);
if !ptr.is_null() && DLMALLOC.calloc_must_clear(ptr) {
core::ptr::write_bytes(ptr, 0, size);
}
ptr
}
ptr
}
unsafe fn dlmalloc_dealloc(ptr: *mut u8, _size: usize, _align: usize) {
DLMALLOC.free(ptr);
unsafe {
DLMALLOC.free(ptr);
}
}
unsafe fn dlmalloc_realloc(
ptr: *mut u8,
old_size: usize,
old_align: usize,
new_size: usize,
) -> *mut u8 {
if old_align <= DLMALLOC.malloc_alignment() {
DLMALLOC.realloc(ptr, new_size)
} else {
let ptr_new = DLMALLOC.memalign(old_align, new_size);
if !ptr_new.is_null() {
core::ptr::copy_nonoverlapping(ptr, ptr_new, core::cmp::min(old_size, new_size));
DLMALLOC.free(ptr);
unsafe {
if old_align <= DLMALLOC.malloc_alignment() {
DLMALLOC.realloc(ptr, new_size)
} else {
let ptr_new = DLMALLOC.memalign(old_align, new_size);
if !ptr_new.is_null() {
core::ptr::copy_nonoverlapping(ptr, ptr_new, core::cmp::min(old_size, new_size));
DLMALLOC.free(ptr);
}
ptr_new
}
ptr_new
}
}

mod services_override {
#[inline(always)]
pub unsafe extern "C" fn svc_read_stdio(fd: usize, buf: *mut u8, count: usize) -> usize {
super::syscall::read(fd, buf, count)
unsafe { super::syscall::read(fd, buf, count) }
}
#[inline(always)]
pub unsafe extern "C" fn svc_write_stdio(fd: usize, buf: *const u8, count: usize) -> usize {
super::syscall::write(fd, buf, count)
unsafe { super::syscall::write(fd, buf, count) }
}
}

pub unsafe fn init() {
/* Ensure stack size is at least 256 MiB, when running locally
* (online judges usually have their stack sizes set large).
* For Windows, this is set as a linker option in the build script.
* However, on Linux, the linker option only marks this value
* in an ELF section, which must be interpreted and applied
* by the runtime startup code (e.g., glibc).
* Thus, instead of parsing the ELF section, we just invoke
* the kernel APIs directly. */
let pd = services::platform_data();
if pd.env_flags & services::ENV_FLAGS_NATIVE != 0 {
let mut rlim: syscall::RLimit = Default::default();
let ret = syscall::getrlimit(syscall::RLIMIT_STACK, &mut rlim as *mut syscall::RLimit);
if ret == 0 && rlim.rlim_cur < 256 * 1024 * 1024 {
rlim.rlim_cur = 256 * 1024 * 1024;
syscall::setrlimit(syscall::RLIMIT_STACK, &rlim as *const syscall::RLimit);
unsafe {
/* Ensure stack size is at least 256 MiB, when running locally
* (online judges usually have their stack sizes set large).
* For Windows, this is set as a linker option in the build script.
* However, on Linux, the linker option only marks this value
* in an ELF section, which must be interpreted and applied
* by the runtime startup code (e.g., glibc).
* Thus, instead of parsing the ELF section, we just invoke
* the kernel APIs directly. */
let pd = services::platform_data();
if pd.env_flags & services::ENV_FLAGS_NATIVE != 0 {
let mut rlim: syscall::RLimit = Default::default();
let ret = syscall::getrlimit(syscall::RLIMIT_STACK, &mut rlim as *mut syscall::RLimit);
if ret == 0 && rlim.rlim_cur < 256 * 1024 * 1024 {
rlim.rlim_cur = 256 * 1024 * 1024;
syscall::setrlimit(syscall::RLIMIT_STACK, &rlim as *const syscall::RLimit);
}
}
}

allocator::install_malloc_impl(
dlmalloc_alloc,
dlmalloc_alloc_zeroed,
dlmalloc_dealloc,
dlmalloc_realloc,
);
allocator::install_malloc_impl(
dlmalloc_alloc,
dlmalloc_alloc_zeroed,
dlmalloc_dealloc,
dlmalloc_realloc,
);

services::install_single_service(5, services_override::svc_read_stdio as usize);
services::install_single_service(6, services_override::svc_write_stdio as usize);
services::install_single_service(5, services_override::svc_read_stdio as usize);
services::install_single_service(6, services_override::svc_write_stdio as usize);
}
}

0 comments on commit 19964cd

Please sign in to comment.