From e47089c63a92bca4391a5fbb197e6b0872867458 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:18 +0000 Subject: [PATCH 01/12] Bump version to 3.0.0 --- Cargo.lock | 2 +- Cargo.toml | 2 +- fuzz/Cargo.lock | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bca4e24..f8923f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,7 +4,7 @@ version = 4 [[package]] name = "picoalloc" -version = "2.1.0" +version = "3.0.0" dependencies = [ "polkavm-derive", ] diff --git a/Cargo.toml b/Cargo.toml index ce249ca..d210db1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ members = [ [package] name = "picoalloc" -version = "2.1.0" +version = "3.0.0" edition = "2021" authors = ["Jan Bujak "] repository = "https://github.com/koute/picoalloc" diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index 0c9ee1f..e483782 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -66,7 +66,7 @@ checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" [[package]] name = "picoalloc" -version = "2.1.0" +version = "3.0.0" [[package]] name = "picoalloc-fuzz" From 110b3b3997dc212716bbc3b5a1aac3f64bdc90ad Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:19 +0000 Subject: [PATCH 02/12] Prevent an out-of-bounds write if we can't expand to full address space --- src/allocator.rs | 58 ++++++++++++++++++++++---------- src/env.rs | 6 ++-- src/env/corevm.rs | 2 +- src/env/linux.rs | 2 +- src/env/polkavm.rs | 10 ++---- src/lib.rs | 84 +++++++++++++++++++++++++++++++++++++++++----- 6 files changed, 123 insertions(+), 39 deletions(-) diff --git a/src/allocator.rs b/src/allocator.rs index 3883bcc..809cacf 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -552,6 +552,7 @@ struct FreeChunkHeader { } const HEADER_SIZE: Size = Size::from_bytes_usize(core::mem::size_of::()).unwrap(); +const FREE_CHUNK_HEADER_SIZE: Size = Size::from_bytes_usize(core::mem::size_of::()).unwrap(); const _: () = { assert!(core::mem::size_of::() <= ALLOCATION_GRANULARITY as usize); @@ -590,24 +591,41 @@ impl Allocator { } #[inline(always)] - fn initialize(&mut self) { + fn initialize(&mut self) -> bool { if self.base_address.is_null() { - self.initialize_impl(); + self.initialize_impl() + } else { + true } } #[inline(never)] #[cold] - fn initialize_impl(&mut self) { - self.base_address = unsafe { self.env.allocate_address_space(self.total_space) }; - paranoid_assert_eq!(self.base_address.addr() % ALLOCATION_GRANULARITY as usize, 0); + fn initialize_impl(&mut self) -> bool { + let base_address = unsafe { self.env.allocate_address_space(self.total_space) }; + if base_address.is_null() { + return false; + } - let bin = Self::size_to_bin_round_down(self.total_space); - self.free_lists_with_unallocated_memory.set(bin); + paranoid_assert_eq!(base_address.addr() % ALLOCATION_GRANULARITY as usize, 0); - let chunk = self.base_address.cast::(); + let chunk = base_address.cast::(); paranoid_assert!(chunk.is_aligned()); + let is_ok = unsafe { self.env.expand_memory_until(base_address, FREE_CHUNK_HEADER_SIZE) }; + if !is_ok { + unsafe { + self.env.free_address_space(base_address, self.total_space); + } + + return false; + } + + self.base_address = base_address; + + let bin = Self::size_to_bin_round_down(self.total_space); + self.free_lists_with_unallocated_memory.set(bin); + let chunk_header = FreeChunkHeader { prev_chunk_size: Size(0), size: ChunkSize::new_unallocated(self.total_space), @@ -619,6 +637,8 @@ impl Allocator { chunk.write(chunk_header); *get_mut_unchecked(&mut self.first_in_free_list, bin.index()) = Pointer::from_pointer(chunk); } + + true } #[inline] @@ -727,7 +747,9 @@ impl Allocator { return None; } - self.initialize(); + if !self.initialize() { + return None; + } let min_size = requested_size.checked_add(HEADER_SIZE)?.checked_add(align.unchecked_sub(Size(1)))?; if min_size.0 > MAX_ALLOCATION_SIZE.0 { @@ -769,21 +791,21 @@ impl Allocator { paranoid_assert!(header_offset >= chunk_offset); - if !unsafe { - self.env.expand_memory_until( - self.base_address - .byte_add(data_offset.unchecked_add(requested_size).bytes() as usize), - ) - } { - return None; - } - let free_space_lhs = header_offset.unchecked_sub(chunk_offset); let free_space_rhs = chunk_size .unchecked_sub(requested_size) .unchecked_sub(free_space_lhs) .unchecked_sub(HEADER_SIZE); + let mut end_offset = data_offset.unchecked_add(requested_size); + if !free_space_rhs.is_empty() { + end_offset = end_offset.unchecked_add(FREE_CHUNK_HEADER_SIZE); + } + + if !unsafe { self.env.expand_memory_until(self.base_address, end_offset) } { + return None; + } + unsafe { let mut prev_chunk_size = chunk.get_unchecked(self.base_address).prev_chunk_size; self.unregister_free_space_first_chunk(chunk, bin); diff --git a/src/env.rs b/src/env.rs index fbdf7a1..385d3c4 100644 --- a/src/env.rs +++ b/src/env.rs @@ -23,7 +23,7 @@ pub fn abort() -> ! { pub trait Env { unsafe fn allocate_address_space(&mut self, size: Size) -> *mut u8; - unsafe fn expand_memory_until(&mut self, end: *mut u8) -> bool; + unsafe fn expand_memory_until(&mut self, base: *mut u8, size: Size) -> bool; unsafe fn free_address_space(&mut self, base: *mut u8, size: Size); } @@ -38,8 +38,8 @@ impl Env for ArrayPointer { self.0.cast() } - unsafe fn expand_memory_until(&mut self, end: *mut u8) -> bool { - (end.addr() - self.0.addr()) <= SIZE + unsafe fn expand_memory_until(&mut self, _base: *mut u8, size: Size) -> bool { + size <= const { Size::from_bytes_usize(SIZE).unwrap() } } unsafe fn free_address_space(&mut self, _base: *mut u8, _size: Size) {} diff --git a/src/env/corevm.rs b/src/env/corevm.rs index d6462ad..4637c1e 100644 --- a/src/env/corevm.rs +++ b/src/env/corevm.rs @@ -15,7 +15,7 @@ impl Env for System { } #[inline] - unsafe fn expand_memory_until(&mut self, _end: *mut u8) -> bool { + unsafe fn expand_memory_until(&mut self, _base: *mut u8, _size: Size) -> bool { true } diff --git a/src/env/linux.rs b/src/env/linux.rs index 5c73d8a..c620184 100644 --- a/src/env/linux.rs +++ b/src/env/linux.rs @@ -68,7 +68,7 @@ impl Env for System { } #[inline] - unsafe fn expand_memory_until(&mut self, _end: *mut u8) -> bool { + unsafe fn expand_memory_until(&mut self, _base: *mut u8, _size: Size) -> bool { true } diff --git a/src/env/polkavm.rs b/src/env/polkavm.rs index bc562a9..878b0ca 100644 --- a/src/env/polkavm.rs +++ b/src/env/polkavm.rs @@ -28,13 +28,9 @@ impl Env for System { } #[inline] - unsafe fn expand_memory_until(&mut self, end: *mut u8) -> bool { - let heap_end = sbrk(0); - if heap_end.addr() >= end.addr() { - return true; - } - - !sbrk(end.addr() - heap_end.addr()).is_null() + unsafe fn expand_memory_until(&mut self, base: *mut u8, size: Size) -> bool { + let bytes = sbrk(0).addr() - base.addr(); + bytes >= size.bytes() as usize || !sbrk(size.bytes() as usize - bytes).is_null() } #[inline] diff --git a/src/lib.rs b/src/lib.rs index 881ec8a..bf9bc70 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -93,14 +93,18 @@ fn test_many_small_allocations(env: E, count: usize) { extern crate alloc; let mut allocator = Allocator::new(env, Size::from_bytes_usize(32 * 1024 * 1024).unwrap()); let mut allocations = alloc::vec::Vec::new(); - for _ in 0..count { - allocations.push( - allocator - .alloc(Size::from_bytes_usize(1).unwrap(), Size::from_bytes_usize(1).unwrap()) - .unwrap(), - ); + for nth in 0..count { + let Some(pointer) = allocator.alloc(Size::from_bytes_usize(1).unwrap(), Size::from_bytes_usize(1).unwrap()) else { + panic!("allocation {nth} failed!"); + }; + + allocations.push(pointer); } + assert!(allocator + .alloc(Size::from_bytes_usize(1).unwrap(), Size::from_bytes_usize(1).unwrap()) + .is_none()); + unsafe { allocator.free(allocations.pop().unwrap()); allocator.free(allocations.swap_remove(0)); @@ -116,11 +120,73 @@ fn test_many_small_allocations(env: E, count: usize) { #[cfg(any(all(target_arch = "x86_64", target_os = "linux"), target_env = "polkavm"))] #[test] fn test_many_small_allocations_native() { - test_many_small_allocations(System, 10000); + test_many_small_allocations(System, 524288); } #[test] fn test_many_small_allocations_buffer() { - let mut buffer = Array([0_u8; 1024 * 16]); - test_many_small_allocations(ArrayPointer(&mut buffer), 256); + #[repr(C)] + struct Storage { + buffer: Array<{ 1024 * 16 }>, + sentinel: [u8; 64], + } + let mut storage = Storage { + buffer: Array([0_u8; 1024 * 16]), + sentinel: [0b10101010; 64], + }; + + test_many_small_allocations(ArrayPointer(&mut storage.buffer), 255); + unsafe { + for offset in 0..storage.sentinel.len() { + // Make sure there were no out-of-bounds writes. + assert_eq!(core::ptr::read_volatile(storage.sentinel.as_ptr().add(offset)), 0b10101010); + } + } +} + +#[test] +fn test_boundary() { + #[repr(align(32))] + pub struct TestEnv { + buffer: [u8; SIZE], + } + + impl Drop for TestEnv { + fn drop(&mut self) { + // Make sure there were no out-of-bounds writes. + assert!(self.buffer[LIMIT..].iter().all(|&byte| byte == 0b10101010)); + } + } + + impl Env for TestEnv { + unsafe fn allocate_address_space(&mut self, _size: Size) -> *mut u8 { + self.buffer.as_mut_ptr() + } + + unsafe fn expand_memory_until(&mut self, _base: *mut u8, size: Size) -> bool { + size <= const { Size::from_bytes_usize(LIMIT).unwrap() } + } + + unsafe fn free_address_space(&mut self, _base: *mut u8, _size: Size) {} + } + + impl Default for TestEnv { + fn default() -> Self { + Self { + buffer: [0b10101010; SIZE], + } + } + } + + let env = TestEnv::<256, 64>::default(); + let mut alloc = Allocator::new(env, Size::from_bytes_usize(64).unwrap()); + let p = alloc + .alloc(Size::from_bytes_usize(1).unwrap(), Size::from_bytes_usize(32).unwrap()) + .unwrap(); + unsafe { + alloc.free(p); + } + assert!(alloc + .alloc(Size::from_bytes_usize(1).unwrap(), Size::from_bytes_usize(33).unwrap()) + .is_none()); } From 853bc0f7862bb8632d56633d911a5b034a8850f8 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:20 +0000 Subject: [PATCH 03/12] Only test the system allocator on linux/x86_64 --- src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index bf9bc70..4b0fb29 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -76,7 +76,7 @@ fn test_allocator(env: E) { } } -#[cfg(any(all(target_arch = "x86_64", target_os = "linux"), target_env = "polkavm"))] +#[cfg(all(target_arch = "x86_64", target_os = "linux"))] #[test] fn test_allocator_system() { test_allocator(System); @@ -117,7 +117,7 @@ fn test_many_small_allocations(env: E, count: usize) { } } -#[cfg(any(all(target_arch = "x86_64", target_os = "linux"), target_env = "polkavm"))] +#[cfg(all(target_arch = "x86_64", target_os = "linux"))] #[test] fn test_many_small_allocations_native() { test_many_small_allocations(System, 524288); From f11da24bdf49550a888d5a5da58c497920653e2a Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:20 +0000 Subject: [PATCH 04/12] Also run the tests with the global allocator enabled --- ci/jobs/build-and-test.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ci/jobs/build-and-test.sh b/ci/jobs/build-and-test.sh index 582518e..db68e8c 100755 --- a/ci/jobs/build-and-test.sh +++ b/ci/jobs/build-and-test.sh @@ -13,6 +13,9 @@ cargo test --all --release echo ">> cargo test (paranoid)" cargo test --features paranoid +echo ">> cargo test (paranoid, global allocator)" +cargo test --features paranoid,global_allocator_rust + echo ">> cargo build (native)" cargo build -p picoalloc_native --release From 461a7eff302b25147467351a9bba3bcc8b9e2c78 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:21 +0000 Subject: [PATCH 05/12] Unexport `System` For certain targets having multiple allocators using this is unsound. --- fuzz/fuzz_targets/allocator.rs | 6 +++--- src/env/corevm.rs | 3 ++- src/env/linux.rs | 4 ++-- src/env/polkavm.rs | 3 ++- src/global_allocator_libc.rs | 2 +- src/lib.rs | 15 ++++++++++----- 6 files changed, 20 insertions(+), 13 deletions(-) diff --git a/fuzz/fuzz_targets/allocator.rs b/fuzz/fuzz_targets/allocator.rs index 5dec1fb..3a071f2 100644 --- a/fuzz/fuzz_targets/allocator.rs +++ b/fuzz/fuzz_targets/allocator.rs @@ -10,7 +10,7 @@ enum Op { Free { index: usize }, } -use picoalloc::{Allocator, System, Size}; +use picoalloc::{Allocator, Size, UnsafeSystem}; fn fill_slice(seed: u128, slice: &mut [u8]) { let mut rng = oorandom::Rand64::new(seed); @@ -27,7 +27,7 @@ fn fill_slice(seed: u128, slice: &mut [u8]) { } fuzz_target!(|ops: Vec| { - let mut allocator = Allocator::new(System, Size::from_bytes_usize(32 * 1024 * 1024).unwrap()); + let mut allocator = Allocator::new(UnsafeSystem, Size::from_bytes_usize(32 * 1024 * 1024).unwrap()); let mut allocations: Vec<(*mut u8, Vec)> = vec![]; let mut alive_addresses = BTreeSet::new(); @@ -42,7 +42,7 @@ fuzz_target!(|ops: Vec| { assert_eq!(pointer.addr() % align, 0); - let usable_size = unsafe { Allocator::::usable_size(pointer) }; + let usable_size = unsafe { Allocator::::usable_size(pointer) }; assert!(usable_size >= size); let data = { diff --git a/src/env/corevm.rs b/src/env/corevm.rs index 4637c1e..36a176f 100644 --- a/src/env/corevm.rs +++ b/src/env/corevm.rs @@ -1,4 +1,5 @@ -use crate::{Env, Size, System}; +use crate::env::System; +use crate::{Env, Size}; #[polkavm_derive::polkavm_import] extern "C" { diff --git a/src/env/linux.rs b/src/env/linux.rs index c620184..6114fa0 100644 --- a/src/env/linux.rs +++ b/src/env/linux.rs @@ -1,5 +1,5 @@ -use crate::env::abort; -use crate::{Env, Size, System}; +use crate::env::{abort, System}; +use crate::{Env, Size}; #[inline] fn abort_on_fail(result: usize) -> usize { diff --git a/src/env/polkavm.rs b/src/env/polkavm.rs index 878b0ca..703cddb 100644 --- a/src/env/polkavm.rs +++ b/src/env/polkavm.rs @@ -1,4 +1,5 @@ -use crate::{Env, Size, System}; +use crate::env::System; +use crate::{Env, Size}; #[inline] fn sbrk(size: usize) -> *mut u8 { diff --git a/src/global_allocator_libc.rs b/src/global_allocator_libc.rs index 6014cc9..61fea20 100644 --- a/src/global_allocator_libc.rs +++ b/src/global_allocator_libc.rs @@ -143,5 +143,5 @@ pub unsafe extern "C" fn realloc(pointer: *mut c_void, size: usize) -> *mut c_vo #[no_mangle] pub unsafe extern "C" fn malloc_usable_size(pointer: *mut c_void) -> usize { - Allocator::::usable_size(pointer.cast()) + Allocator::::usable_size(pointer.cast()) } diff --git a/src/lib.rs b/src/lib.rs index 4b0fb29..cbcf9c6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,16 +13,21 @@ mod global_allocator_libc; #[cfg(any(feature = "global_allocator_rust", feature = "global_allocator_libc"))] #[cfg_attr(feature = "global_allocator_rust", global_allocator)] -pub(crate) static GLOBAL_ALLOCATOR: Mutex> = - Mutex::new(Allocator::new(System, Size::from_bytes_usize(1024 * 1024 * 1024).unwrap())); +pub(crate) static GLOBAL_ALLOCATOR: Mutex> = Mutex::new(Allocator::new( + crate::env::System, + Size::from_bytes_usize(1024 * 1024 * 1024).unwrap(), +)); pub use crate::allocator::{Allocator, Size}; -pub use crate::env::{Array, ArrayPointer, Env, System}; +pub use crate::env::{Array, ArrayPointer, Env}; pub use crate::mutex::Mutex; #[doc(hidden)] pub use crate::env::abort; +#[doc(hidden)] +pub use crate::env::System as UnsafeSystem; + #[cfg(test)] fn test_allocator(env: E) { let mut allocator = Allocator::new(env, Size::from_bytes_usize(8 * 1024 * 1024).unwrap()); @@ -79,7 +84,7 @@ fn test_allocator(env: E) { #[cfg(all(target_arch = "x86_64", target_os = "linux"))] #[test] fn test_allocator_system() { - test_allocator(System); + test_allocator(crate::env::System); } #[test] @@ -120,7 +125,7 @@ fn test_many_small_allocations(env: E, count: usize) { #[cfg(all(target_arch = "x86_64", target_os = "linux"))] #[test] fn test_many_small_allocations_native() { - test_many_small_allocations(System, 524288); + test_many_small_allocations(crate::env::System, 524288); } #[test] From 37043c3d5ddb83d108b9cf9ff7492fd0b1b400f5 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:21 +0000 Subject: [PATCH 06/12] Remove the PolkaVM-specific special case in `Mutex` Since this API is public it's unsound to do this. --- src/mutex.rs | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/mutex.rs b/src/mutex.rs index 1ab7ead..2da3e6c 100644 --- a/src/mutex.rs +++ b/src/mutex.rs @@ -3,7 +3,6 @@ use core::ops::{Deref, DerefMut}; pub struct Mutex { value: UnsafeCell, - #[cfg(not(target_env = "polkavm"))] flag: core::sync::atomic::AtomicBool, } @@ -20,22 +19,18 @@ impl Mutex { pub const fn new(value: T) -> Self { Mutex { value: UnsafeCell::new(value), - #[cfg(not(target_env = "polkavm"))] flag: core::sync::atomic::AtomicBool::new(false), } } #[inline] pub fn lock(&self) -> MutexGuard { - #[cfg(not(target_env = "polkavm"))] - { - use core::sync::atomic::Ordering; - while self - .flag - .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed) - .is_err() - {} - } + use core::sync::atomic::Ordering; + while self + .flag + .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_err() + {} MutexGuard(self) } @@ -44,7 +39,6 @@ impl Mutex { impl Drop for MutexGuard<'_, T> { #[inline] fn drop(&mut self) { - #[cfg(not(target_env = "polkavm"))] self.0.flag.store(false, core::sync::atomic::Ordering::Release); } } From ab3ed3113c1ca77cefe9f0665b6586d633d8af26 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:22 +0000 Subject: [PATCH 07/12] Add WASM support --- .github/workflows/rust.yml | 6 ++++++ ci/jobs/build-and-test-wasm.sh | 16 +++++++++++++++ ci/run-all-tests.sh | 1 + src/env.rs | 3 +++ src/env/wasm.rs | 36 ++++++++++++++++++++++++++++++++++ 5 files changed, 62 insertions(+) create mode 100755 ci/jobs/build-and-test-wasm.sh create mode 100644 src/env/wasm.rs diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 8f677a4..d6aa979 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -16,8 +16,14 @@ jobs: - uses: actions/checkout@v4 - name: uname -a run: uname -a + - name: Install wasmtime-cli + run: cargo install --version 33.0.0 wasmtime-cli + - name: Install target wasm32-wasip2 + run: rustup target add wasm32-wasip2 - name: Build and test run: ./ci/jobs/build-and-test.sh + - name: Build and test (WASM) + run: ./ci/jobs/build-and-test-wasm.sh rustfmt: runs-on: ubuntu-24.04 steps: diff --git a/ci/jobs/build-and-test-wasm.sh b/ci/jobs/build-and-test-wasm.sh new file mode 100755 index 0000000..ede30be --- /dev/null +++ b/ci/jobs/build-and-test-wasm.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +set -euo pipefail +cd -- "$(dirname -- "${BASH_SOURCE[0]}")" +cd ../.. + +cargo check --features paranoid,global_allocator_rust --target=wasm32-wasip2 + +echo ">> cargo test (debug, WASM)" +wasmtime $(cargo test --features paranoid --target=wasm32-wasip2 --no-run --message-format=json | grep -oE '"executable":"[^"]+"' | cut -d ":" -f 2 | grep -oE '[^"]+') + +echo ">> cargo test (release, WASM)" +wasmtime $(cargo test --features paranoid --release --target=wasm32-wasip2 --no-run --message-format=json | grep -oE '"executable":"[^"]+"' | cut -d ":" -f 2 | grep -oE '[^"]+') + +echo ">> cargo test (debug, global allocator, WASM)" +wasmtime $(cargo test --features paranoid,global_allocator_rust --target=wasm32-wasip2 --no-run --message-format=json | grep -oE '"executable":"[^"]+"' | cut -d ":" -f 2 | grep -oE '[^"]+') diff --git a/ci/run-all-tests.sh b/ci/run-all-tests.sh index bfe5700..d128b94 100755 --- a/ci/run-all-tests.sh +++ b/ci/run-all-tests.sh @@ -5,6 +5,7 @@ cd -- "$(dirname -- "${BASH_SOURCE[0]}")" cd .. ./ci/jobs/build-and-test.sh +./ci/jobs/build-and-test-wasm.sh case "$OSTYPE" in linux*) diff --git a/src/env.rs b/src/env.rs index 385d3c4..9fc7856 100644 --- a/src/env.rs +++ b/src/env.rs @@ -55,3 +55,6 @@ mod polkavm; #[cfg(all(target_env = "polkavm", feature = "corevm"))] mod corevm; + +#[cfg(target_family = "wasm")] +mod wasm; diff --git a/src/env/wasm.rs b/src/env/wasm.rs new file mode 100644 index 0000000..dc02c6d --- /dev/null +++ b/src/env/wasm.rs @@ -0,0 +1,36 @@ +use crate::env::System; +use crate::{Env, Size}; + +extern "C" { + static mut __heap_base: u8; +} + +#[inline] +fn heap_base() -> *mut u8 { + let pointer = (&raw mut __heap_base); + pointer.with_addr(pointer.addr().next_multiple_of(32)) +} + +const PAGE_SIZE: usize = 64 * 1024; + +impl Env for System { + #[inline] + unsafe fn allocate_address_space(&mut self, _size: Size) -> *mut u8 { + heap_base() + } + + #[inline] + unsafe fn expand_memory_until(&mut self, base: *mut u8, size: Size) -> bool { + let current_size = core::arch::wasm32::memory_size(0) * PAGE_SIZE; + let required_size = base.addr() + size.bytes() as usize; + if current_size >= required_size { + return true; + } + + let delta = (required_size - current_size).next_multiple_of(PAGE_SIZE); + core::arch::wasm32::memory_grow(0, delta / PAGE_SIZE) != usize::MAX + } + + #[inline] + unsafe fn free_address_space(&mut self, _base: *mut u8, _size: Size) {} +} From c63ce2d9a148d503254a58a949674d433634d552 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:22 +0000 Subject: [PATCH 08/12] Move `GlobalAlloc` impl to allocator.rs --- src/allocator.rs | 22 ++++++++++++++++++++++ src/global_allocator.rs | 21 --------------------- src/lib.rs | 3 --- 3 files changed, 22 insertions(+), 24 deletions(-) delete mode 100644 src/global_allocator.rs diff --git a/src/allocator.rs b/src/allocator.rs index 809cacf..05b3faf 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -954,3 +954,25 @@ impl Allocator { unsafe { &*pointer.byte_sub(HEADER_SIZE.bytes() as usize).cast::() } } } + +#[cfg(feature = "alloc")] +extern crate alloc; + +#[cfg(feature = "alloc")] +unsafe impl alloc::alloc::GlobalAlloc for crate::Mutex> { + unsafe fn alloc(&self, layout: alloc::alloc::Layout) -> *mut u8 { + let Some(align) = Size::from_bytes_usize(layout.align()) else { + return core::ptr::null_mut(); + }; + + let Some(size) = Size::from_bytes_usize(layout.size()) else { + return core::ptr::null_mut(); + }; + + self.lock().alloc(align, size).unwrap_or(core::ptr::null_mut()) + } + + unsafe fn dealloc(&self, pointer: *mut u8, _layout: alloc::alloc::Layout) { + self.lock().free(pointer); + } +} diff --git a/src/global_allocator.rs b/src/global_allocator.rs deleted file mode 100644 index 7badd12..0000000 --- a/src/global_allocator.rs +++ /dev/null @@ -1,21 +0,0 @@ -extern crate alloc; - -use crate::{Allocator, Env, Mutex, Size}; - -unsafe impl alloc::alloc::GlobalAlloc for Mutex> { - unsafe fn alloc(&self, layout: alloc::alloc::Layout) -> *mut u8 { - let Some(align) = Size::from_bytes_usize(layout.align()) else { - return core::ptr::null_mut(); - }; - - let Some(size) = Size::from_bytes_usize(layout.size()) else { - return core::ptr::null_mut(); - }; - - self.lock().alloc(align, size).unwrap_or(core::ptr::null_mut()) - } - - unsafe fn dealloc(&self, pointer: *mut u8, _layout: alloc::alloc::Layout) { - self.lock().free(pointer); - } -} diff --git a/src/lib.rs b/src/lib.rs index cbcf9c6..5ce56b5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,9 +5,6 @@ mod allocator; mod env; mod mutex; -#[cfg(feature = "alloc")] -mod global_allocator; - #[cfg(feature = "global_allocator_libc")] mod global_allocator_libc; From 5a704b4f1f703c0403a5fb67ab1c903511ee2d00 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:23 +0000 Subject: [PATCH 09/12] Remove the `alloc` feature --- Cargo.toml | 3 +-- src/allocator.rs | 10 +++------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index d210db1..c37d0c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,9 +30,8 @@ panic = "abort" [features] default = [] -alloc = [] global_allocator_libc = [] -global_allocator_rust = ["alloc"] +global_allocator_rust = [] paranoid = ["strict_provenance"] strict_provenance = [] corevm = ["dep:polkavm-derive"] diff --git a/src/allocator.rs b/src/allocator.rs index 05b3faf..e67d571 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -955,12 +955,8 @@ impl Allocator { } } -#[cfg(feature = "alloc")] -extern crate alloc; - -#[cfg(feature = "alloc")] -unsafe impl alloc::alloc::GlobalAlloc for crate::Mutex> { - unsafe fn alloc(&self, layout: alloc::alloc::Layout) -> *mut u8 { +unsafe impl core::alloc::GlobalAlloc for crate::Mutex> { + unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 { let Some(align) = Size::from_bytes_usize(layout.align()) else { return core::ptr::null_mut(); }; @@ -972,7 +968,7 @@ unsafe impl alloc::alloc::GlobalAlloc for crate::Mutex Date: Tue, 17 Jun 2025 08:29:23 +0000 Subject: [PATCH 10/12] Track expanded memory size and add more paranoid asserts --- src/allocator.rs | 63 +++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 57 insertions(+), 6 deletions(-) diff --git a/src/allocator.rs b/src/allocator.rs index e67d571..eea8511 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -562,6 +562,7 @@ const _: () = { pub struct Allocator { total_space: Size, + allocated_space: Size, base_address: *mut u8, free_lists_with_unallocated_memory: BitMask, first_in_free_list: [Pointer; BIN_CONFIG.bin_count as usize], @@ -583,6 +584,7 @@ impl Allocator { pub const fn new(env: E, total_space: Size) -> Self { Allocator { total_space, + allocated_space: const { Size::from_bytes_usize(0).unwrap() }, base_address: core::ptr::null_mut(), free_lists_with_unallocated_memory: BitMask::new(), first_in_free_list: [Pointer::NULL; BIN_CONFIG.bin_count as usize], @@ -622,6 +624,8 @@ impl Allocator { } self.base_address = base_address; + self.allocated_space = FREE_CHUNK_HEADER_SIZE; + self.paranoid_check_access(Pointer::from_pointer(chunk)); let bin = Self::size_to_bin_round_down(self.total_space); self.free_lists_with_unallocated_memory.set(bin); @@ -638,6 +642,7 @@ impl Allocator { *get_mut_unchecked(&mut self.first_in_free_list, bin.index()) = Pointer::from_pointer(chunk); } + self.paranoid_check_chunk(Pointer::from_pointer(chunk).cast()); true } @@ -659,6 +664,7 @@ impl Allocator { #[inline(always)] fn unregister_free_space_first_chunk(&mut self, chunk: Pointer, bin: BitIndex) { + self.paranoid_check_access(chunk); paranoid_assert_eq!(self.first_in_free_list[bin.index()], chunk); unsafe { @@ -677,6 +683,8 @@ impl Allocator { #[inline(always)] fn unregister_free_space(&mut self, chunk: Pointer, bin: BitIndex) { + self.paranoid_check_access(chunk); + if unsafe { *get_unchecked(&self.first_in_free_list, bin.index()) } == chunk { self.unregister_free_space_first_chunk(chunk, bin); } else { @@ -702,7 +710,7 @@ impl Allocator { return prev_chunk_size; } - paranoid_assert!(!chunk.is_null()); + self.paranoid_check_access(chunk); let bin = Self::size_to_bin_round_down(size); unsafe { @@ -728,7 +736,8 @@ impl Allocator { #[inline(always)] fn register_allocation(&mut self, chunk: Pointer, prev_chunk_size: Size, size: Size) { - paranoid_assert!(!chunk.is_null()); + self.paranoid_check_access(chunk); + unsafe { chunk.write_no_drop( self.base_address, @@ -769,7 +778,6 @@ impl Allocator { })?; let chunk = unsafe { *get_unchecked(&self.first_in_free_list, bin.index()) }; - paranoid_assert!(!chunk.is_null()); self.paranoid_check_chunk(chunk.cast::()); let chunk_size = unsafe { chunk.get_unchecked(self.base_address).size }; @@ -802,8 +810,12 @@ impl Allocator { end_offset = end_offset.unchecked_add(FREE_CHUNK_HEADER_SIZE); } - if !unsafe { self.env.expand_memory_until(self.base_address, end_offset) } { - return None; + if self.allocated_space < end_offset { + if !unsafe { self.env.expand_memory_until(self.base_address, end_offset) } { + return None; + } + + self.allocated_space = end_offset; } unsafe { @@ -823,6 +835,7 @@ impl Allocator { let final_chunk = next_chunk.unchecked_add(free_space_rhs); if final_chunk.cast() < Pointer::from_pointer(self.base_address).unchecked_add(self.total_space) { + self.paranoid_check_access(final_chunk); final_chunk.get_mut_unchecked(self.base_address).prev_chunk_size = prev_chunk_size; } @@ -838,21 +851,54 @@ impl Allocator { let data: Pointer = allocation_chunk.unchecked_add(HEADER_SIZE).cast(); paranoid_assert_eq!(data.address() % align.bytes() as Address, 0); - Some(data.raw_pointer_mut(self.base_address)) + let output = data.raw_pointer_mut(self.base_address); + self.paranoid_check_chunk(Pointer::from_pointer(output).unchecked_sub(HEADER_SIZE).cast::()); + + Some(output) + } + + #[cfg(any(debug_assertions, test, feature = "paranoid"))] + #[inline(never)] + #[track_caller] + fn paranoid_check_access(&self, pointer: Pointer) { + paranoid_assert!(!pointer.is_null()); + paranoid_assert!( + pointer + .address() + .wrapping_sub(self.base_address.addr() as Address) + .wrapping_add(core::mem::size_of::() as Address) + <= Address::from(self.allocated_space.bytes()) + ); } + #[cfg(not(any(debug_assertions, test, feature = "paranoid")))] + #[inline(always)] + fn paranoid_check_access(&self, _pointer: Pointer) {} + #[cfg(any(debug_assertions, test, feature = "paranoid"))] #[inline(never)] + #[track_caller] fn paranoid_check_chunk(&self, chunk: Pointer) { paranoid_assert!(!chunk.is_null()); + paranoid_assert!(!self.base_address.is_null()); + paranoid_assert!(!self.total_space.is_empty()); unsafe { let base_address = Pointer::from_pointer(self.base_address); + paranoid_assert!(chunk.cast() >= base_address); + let end_of_address_space = base_address.unchecked_add(self.total_space); if chunk.cast() == end_of_address_space { return; } + paranoid_assert!(chunk.address() < end_of_address_space.address()); + paranoid_assert!( + chunk.address().wrapping_add(core::mem::size_of::() as Address) <= end_of_address_space.address() + ); + + self.paranoid_check_access(chunk); + let prev_chunk_size = chunk.get_unchecked(self.base_address).prev_chunk_size; if prev_chunk_size.is_empty() { paranoid_assert_eq!(chunk.cast(), base_address); @@ -900,6 +946,8 @@ impl Allocator { // Try to merge with the previous free chunk. if !Size::from_pointer_and_base_unchecked(chunk, Pointer::from_pointer_mut(self.base_address)).is_empty() { let prev_chunk = chunk.unchecked_sub(prev_chunk_size); + self.paranoid_check_access(prev_chunk); + let prev_size = unsafe { prev_chunk.get_unchecked(self.base_address).size }; paranoid_assert_eq!(prev_size.size(), prev_chunk_size); @@ -917,6 +965,8 @@ impl Allocator { { let next_chunk = chunk.unchecked_add(size); if next_chunk.cast() < end_of_address_space { + self.paranoid_check_access(next_chunk); + let next_size = unsafe { next_chunk.get_unchecked(self.base_address).size }; if !next_size.is_allocated() { let next_size = next_size.size(); @@ -931,6 +981,7 @@ impl Allocator { let next_chunk = chunk.unchecked_add(size); if next_chunk.cast() < end_of_address_space { + self.paranoid_check_access(next_chunk); unsafe { next_chunk.get_mut_unchecked(self.base_address).prev_chunk_size = size; }; From 2d29e27736c6f7cf0c9fe9f5671719a6f3f168cd Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:29:25 +0000 Subject: [PATCH 11/12] Make sure the heap base is aligned under PolkaVM --- src/env/polkavm.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/env/polkavm.rs b/src/env/polkavm.rs index 703cddb..b857f85 100644 --- a/src/env/polkavm.rs +++ b/src/env/polkavm.rs @@ -19,12 +19,22 @@ impl Env for System { #[inline] unsafe fn allocate_address_space(&mut self, _size: Size) -> *mut u8 { unsafe { - let mut output; + let mut pointer: *mut u8; core::arch::asm!( ".insn r 0xb, 3, 0, {dst}, zero, zero", - dst = out(reg) output, + dst = out(reg) pointer, ); - output + + let aligned_pointer = pointer.with_addr(pointer.addr().next_multiple_of(32)); + let sbrk_pointer = sbrk(0); + if sbrk_pointer.addr() < aligned_pointer.addr() { + let bytes = aligned_pointer.addr() - sbrk_pointer.addr(); + if sbrk(bytes).is_null() { + return core::ptr::null_mut(); + } + } + + aligned_pointer } } From 29beb9a9ec082d5d6f893c5b26005c2618003886 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 17 Jun 2025 08:35:56 +0000 Subject: [PATCH 12/12] Do not use `Option::or_else` for better codegen --- src/allocator.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/allocator.rs b/src/allocator.rs index eea8511..6f374e8 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -769,13 +769,17 @@ impl Allocator { // // First calculate the minimum bin to fit this allocation; round up in case the size doesn't match the bin size exactly. // If this doesn't work then try rounding down and see if maybe we can find an oversized region in the previous bin. - let bin = self + let mut bin = self .free_lists_with_unallocated_memory - .find_first(Self::size_to_bin_round_up(min_size)) - .or_else(|| { - self.free_lists_with_unallocated_memory - .find_first(Self::size_to_bin_round_down(min_size)) - })?; + .find_first(Self::size_to_bin_round_up(min_size)); + + if bin.is_none() { + bin = self + .free_lists_with_unallocated_memory + .find_first(Self::size_to_bin_round_down(min_size)); + } + + let bin = bin?; let chunk = unsafe { *get_unchecked(&self.first_in_free_list, bin.index()) }; self.paranoid_check_chunk(chunk.cast::());