diff --git a/basm-std/src/platform/codegen.rs b/basm-std/src/platform/codegen.rs index d6e361b4..b880ce88 100644 --- a/basm-std/src/platform/codegen.rs +++ b/basm-std/src/platform/codegen.rs @@ -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")] @@ -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, @@ -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 diff --git a/basm-std/src/platform/os/linux.rs b/basm-std/src/platform/os/linux.rs index 0012e5af..48ea23cb 100644 --- a/basm-std/src/platform/os/linux.rs +++ b/basm-std/src/platform/os/linux.rs @@ -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)] @@ -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)] diff --git a/basm-std/src/platform/os/macos.rs b/basm-std/src/platform/os/macos.rs index 7ab5026e..d6b9bfc8 100644 --- a/basm-std/src/platform/os/macos.rs +++ b/basm-std/src/platform/os/macos.rs @@ -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, @@ -40,17 +40,23 @@ pub mod syscall { static mut DLMALLOC: dlmalloc::Dlmalloc = 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, @@ -58,55 +64,59 @@ unsafe fn dlmalloc_realloc( 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); + } }