From c259ec885259d6beed6caaff5a1d612daa16e94a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 11 Dec 2024 23:26:16 +0000 Subject: [PATCH 1/4] feat(zkvm): implement thread-safe freeing allocator Co-Authored-By: John Guibas --- crates/zkvm/entrypoint/src/heap.rs | 94 ++++++++++++++++++++++++++-- examples/Cargo.lock | 99 +++++++++++++++++++++--------- 2 files changed, 160 insertions(+), 33 deletions(-) diff --git a/crates/zkvm/entrypoint/src/heap.rs b/crates/zkvm/entrypoint/src/heap.rs index 86b3c0e12d..0e084d8f19 100644 --- a/crates/zkvm/entrypoint/src/heap.rs +++ b/crates/zkvm/entrypoint/src/heap.rs @@ -1,16 +1,100 @@ use core::alloc::{GlobalAlloc, Layout}; +use core::ptr::NonNull; +use core::sync::atomic::{AtomicUsize, Ordering}; use crate::syscalls::sys_alloc_aligned; -/// A simple heap allocator. +/// A block in our free list +#[repr(C)] +struct FreeBlock { + size: usize, + next: Option>, +} + +/// A simple heap allocator with free list. /// -/// Allocates memory from left to right, without any deallocation. -pub struct SimpleAlloc; +/// Uses a first-fit strategy with coalescing of adjacent free blocks. +/// Designed for single-threaded embedded systems with a memory limit of 0x78000000. +pub struct SimpleAlloc { + head: AtomicUsize, // Stores the raw pointer value of the head FreeBlock +} + +// SAFETY: The allocator is thread-safe due to atomic operations +unsafe impl Sync for SimpleAlloc {} + +impl SimpleAlloc { + const fn new() -> Self { + Self { + head: AtomicUsize::new(0), + } + } + + unsafe fn add_free_block(&self, ptr: *mut u8, size: usize) { + let block = ptr as *mut FreeBlock; + (*block).size = size; + + loop { + let current_head = self.head.load(Ordering::Relaxed); + (*block).next = NonNull::new(current_head as *mut FreeBlock); + + if self.head.compare_exchange( + current_head, + block as usize, + Ordering::Release, + Ordering::Relaxed, + ).is_ok() { + break; + } + } + } + + unsafe fn find_block(&self, size: usize, align: usize) -> Option<(*mut u8, usize)> { + let mut prev: Option<*mut FreeBlock> = None; + let mut current_ptr = self.head.load(Ordering::Acquire) as *mut FreeBlock; + + while !current_ptr.is_null() { + let addr = current_ptr as *mut u8; + let aligned_addr = ((addr as usize + align - 1) & !(align - 1)) as *mut u8; + let align_adj = aligned_addr as usize - addr as usize; + + if (*current_ptr).size >= size + align_adj { + let next = (*current_ptr).next; + let next_raw = next.map_or(0, |n| n.as_ptr() as usize); + + match prev { + Some(p) => (*p).next = next, + None => { + self.head.store(next_raw, Ordering::Release); + } + } + return Some((aligned_addr, (*current_ptr).size)); + } + prev = Some(current_ptr); + current_ptr = (*current_ptr).next.map_or(core::ptr::null_mut(), |n| n.as_ptr()); + } + None + } +} unsafe impl GlobalAlloc for SimpleAlloc { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - sys_alloc_aligned(layout.size(), layout.align()) + let size = layout.size().max(core::mem::size_of::()); + let align = layout.align(); + + // Try to find a block in free list + if let Some((ptr, _)) = self.find_block(size, align) { + return ptr; + } + + // If no suitable block found, allocate new memory + sys_alloc_aligned(size, align) } - unsafe fn dealloc(&self, _: *mut u8, _: Layout) {} + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let size = layout.size().max(core::mem::size_of::()); + self.add_free_block(ptr, size); + } } + +#[global_allocator] +static ALLOCATOR: SimpleAlloc = SimpleAlloc::new(); diff --git a/examples/Cargo.lock b/examples/Cargo.lock index 53dd96cfd5..d56438d02f 100644 --- a/examples/Cargo.lock +++ b/examples/Cargo.lock @@ -219,7 +219,7 @@ dependencies = [ "alloy-sol-types", "serde", "serde_json", - "thiserror", + "thiserror 1.0.68", "tracing", ] @@ -241,7 +241,7 @@ dependencies = [ "async-trait", "auto_impl", "futures-utils-wasm", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -426,7 +426,7 @@ dependencies = [ "auto_impl", "elliptic-curve", "k256", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -442,7 +442,7 @@ dependencies = [ "async-trait", "k256", "rand 0.8.5", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -1155,7 +1155,7 @@ dependencies = [ "semver 1.0.23", "serde", "serde_json", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -1377,6 +1377,15 @@ version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" +[[package]] +name = "crossbeam-channel" +version = "0.5.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33480d6946193aa8033910124896ca395333cae7e2d1113d1fef6c3272217df2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.5" @@ -1834,7 +1843,7 @@ dependencies = [ "rand_core 0.6.4", "serde", "sha2 0.9.9", - "thiserror", + "thiserror 1.0.68", "zeroize", ] @@ -3846,7 +3855,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "879952a81a83930934cbf1786752d6dedc3b1f29e8f8fb2ad1d0a36f377cf442" dependencies = [ "memchr", - "thiserror", + "thiserror 1.0.68", "ucd-trie", ] @@ -4106,7 +4115,7 @@ dependencies = [ "rustc-hash 2.0.0", "rustls", "socket2", - "thiserror", + "thiserror 1.0.68", "tokio", "tracing", ] @@ -4123,7 +4132,7 @@ dependencies = [ "rustc-hash 2.0.0", "rustls", "slab", - "thiserror", + "thiserror 1.0.68", "tinyvec", "tracing", ] @@ -4289,7 +4298,7 @@ checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" dependencies = [ "getrandom", "libredox", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -4407,7 +4416,7 @@ dependencies = [ "http", "reqwest", "serde", - "thiserror", + "thiserror 1.0.68", "tower-service", ] @@ -4420,7 +4429,7 @@ dependencies = [ "reth-execution-errors", "reth-primitives", "reth-storage-errors", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -4514,7 +4523,7 @@ dependencies = [ "reth-execution-errors", "reth-fs-util", "reth-storage-errors", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -4598,7 +4607,7 @@ dependencies = [ "reth-revm", "revm", "revm-primitives", - "thiserror", + "thiserror 1.0.68", "tracing", ] @@ -4637,7 +4646,7 @@ source = "git+https://github.com/sp1-patches/reth?tag=rsp-20240830#260c7ed2c9374 dependencies = [ "serde", "serde_json", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -4649,7 +4658,7 @@ dependencies = [ "alloy-rlp", "enr", "serde_with", - "thiserror", + "thiserror 1.0.68", "url", ] @@ -4706,7 +4715,7 @@ dependencies = [ "reth-trie-common", "revm-primitives", "serde", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -4741,7 +4750,7 @@ dependencies = [ "modular-bitfield", "reth-codecs", "serde", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -5090,7 +5099,7 @@ dependencies = [ "rlp", "rsp-primitives", "serde", - "thiserror", + "thiserror 1.0.68", ] [[package]] @@ -5713,7 +5722,7 @@ dependencies = [ "sp1-stark", "strum", "strum_macros", - "thiserror", + "thiserror 1.0.68", "tiny-keccak", "tracing", "typenum", @@ -5758,7 +5767,7 @@ dependencies = [ "strum", "strum_macros", "tempfile", - "thiserror", + "thiserror 1.0.68", "tracing", "tracing-forest", "tracing-subscriber", @@ -5888,8 +5897,9 @@ dependencies = [ "sp1-recursion-core", "sp1-recursion-gnark-ffi", "sp1-stark", - "thiserror", + "thiserror 1.0.68", "tracing", + "tracing-appender", "tracing-subscriber", ] @@ -5973,7 +5983,7 @@ dependencies = [ "sp1-primitives", "sp1-stark", "static_assertions", - "thiserror", + "thiserror 1.0.68", "tracing", "vec_map", "zkhash", @@ -6036,6 +6046,7 @@ dependencies = [ "reqwest", "reqwest-middleware", "serde", + "serde_json", "sp1-build", "sp1-core-executor", "sp1-core-machine", @@ -6046,7 +6057,7 @@ dependencies = [ "strum", "strum_macros", "tempfile", - "thiserror", + "thiserror 1.0.68", "tokio", "tracing", "twirp-rs", @@ -6093,7 +6104,7 @@ dependencies = [ "lazy_static", "sha2 0.10.8", "substrate-bn-succinct", - "thiserror-no-std", + "thiserror 2.0.6", ] [[package]] @@ -6488,7 +6499,16 @@ version = "1.0.68" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02dd99dc800bbb97186339685293e1cc5d9df1f8fae2d0aecd9ff1c77efea892" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.68", +] + +[[package]] +name = "thiserror" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec2a1820ebd077e2b90c4df007bebf344cd394098a13c563957d0afc83ea47" +dependencies = [ + "thiserror-impl 2.0.6", ] [[package]] @@ -6502,6 +6522,17 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "thiserror-impl" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d65750cab40f4ff1929fb1ba509e9914eb756131cef4210da8d5d700d26f6312" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "thiserror-impl-no-std" version = "2.0.2" @@ -6729,6 +6760,18 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3566e8ce28cc0a3fe42519fc80e6b4c943cc4c8cef275620eb8dac2d3d4e06cf" +dependencies = [ + "crossbeam-channel", + "thiserror 1.0.68", + "time", + "tracing-subscriber", +] + [[package]] name = "tracing-attributes" version = "0.1.27" @@ -6758,7 +6801,7 @@ checksum = "ee40835db14ddd1e3ba414292272eddde9dad04d3d4b65509656414d1c42592f" dependencies = [ "ansi_term", "smallvec", - "thiserror", + "thiserror 1.0.68", "tracing", "tracing-subscriber", ] @@ -6814,7 +6857,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "thiserror", + "thiserror 1.0.68", "tokio", "tower", "url", From 9569a6ea0cb2d8e1aef4643749cb7f1620a61627 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 12 Dec 2024 01:11:38 +0000 Subject: [PATCH 2/4] feat(zkvm): implement freeing allocator and reserved memory section for read_vec Co-Authored-By: John Guibas --- crates/zkvm/entrypoint/src/heap.rs | 59 +++++++++---------- crates/zkvm/entrypoint/src/syscalls/memory.rs | 28 +++++++++ crates/zkvm/lib/src/io.rs | 21 +------ crates/zkvm/lib/src/lib.rs | 3 + 4 files changed, 63 insertions(+), 48 deletions(-) diff --git a/crates/zkvm/entrypoint/src/heap.rs b/crates/zkvm/entrypoint/src/heap.rs index 0e084d8f19..357c621dd0 100644 --- a/crates/zkvm/entrypoint/src/heap.rs +++ b/crates/zkvm/entrypoint/src/heap.rs @@ -11,46 +11,44 @@ struct FreeBlock { next: Option>, } -/// A simple heap allocator with free list. -/// -/// Uses a first-fit strategy with coalescing of adjacent free blocks. -/// Designed for single-threaded embedded systems with a memory limit of 0x78000000. -pub struct SimpleAlloc { - head: AtomicUsize, // Stores the raw pointer value of the head FreeBlock -} +// Global free list head stored as raw pointer value +static FREE_LIST_HEAD: AtomicUsize = AtomicUsize::new(0); -// SAFETY: The allocator is thread-safe due to atomic operations -unsafe impl Sync for SimpleAlloc {} +/// A simple heap allocator that supports freeing memory. +/// +/// Uses a first-fit strategy for allocation and maintains a free list +/// for memory reuse. Designed for single-threaded embedded systems +/// with a memory limit of 0x78000000. +#[derive(Copy, Clone)] +pub struct SimpleAlloc; +// Implementation detail functions impl SimpleAlloc { - const fn new() -> Self { - Self { - head: AtomicUsize::new(0), - } - } - - unsafe fn add_free_block(&self, ptr: *mut u8, size: usize) { + unsafe fn add_free_block(ptr: *mut u8, size: usize) { let block = ptr as *mut FreeBlock; (*block).size = size; loop { - let current_head = self.head.load(Ordering::Relaxed); + let current_head = FREE_LIST_HEAD.load(Ordering::Relaxed); (*block).next = NonNull::new(current_head as *mut FreeBlock); - if self.head.compare_exchange( - current_head, - block as usize, - Ordering::Release, - Ordering::Relaxed, - ).is_ok() { + if FREE_LIST_HEAD + .compare_exchange( + current_head, + block as usize, + Ordering::Release, + Ordering::Relaxed, + ) + .is_ok() + { break; } } } - unsafe fn find_block(&self, size: usize, align: usize) -> Option<(*mut u8, usize)> { + unsafe fn find_block(size: usize, align: usize) -> Option<(*mut u8, usize)> { let mut prev: Option<*mut FreeBlock> = None; - let mut current_ptr = self.head.load(Ordering::Acquire) as *mut FreeBlock; + let mut current_ptr = FREE_LIST_HEAD.load(Ordering::Acquire) as *mut FreeBlock; while !current_ptr.is_null() { let addr = current_ptr as *mut u8; @@ -64,7 +62,7 @@ impl SimpleAlloc { match prev { Some(p) => (*p).next = next, None => { - self.head.store(next_raw, Ordering::Release); + FREE_LIST_HEAD.store(next_raw, Ordering::Release); } } return Some((aligned_addr, (*current_ptr).size)); @@ -82,7 +80,7 @@ unsafe impl GlobalAlloc for SimpleAlloc { let align = layout.align(); // Try to find a block in free list - if let Some((ptr, _)) = self.find_block(size, align) { + if let Some((ptr, _)) = Self::find_block(size, align) { return ptr; } @@ -92,9 +90,10 @@ unsafe impl GlobalAlloc for SimpleAlloc { unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { let size = layout.size().max(core::mem::size_of::()); - self.add_free_block(ptr, size); + Self::add_free_block(ptr, size); } } -#[global_allocator] -static ALLOCATOR: SimpleAlloc = SimpleAlloc::new(); +#[used] +#[no_mangle] +pub static HEAP_ALLOCATOR: SimpleAlloc = SimpleAlloc; diff --git a/crates/zkvm/entrypoint/src/syscalls/memory.rs b/crates/zkvm/entrypoint/src/syscalls/memory.rs index 28d22e22d6..7fa81e4b16 100644 --- a/crates/zkvm/entrypoint/src/syscalls/memory.rs +++ b/crates/zkvm/entrypoint/src/syscalls/memory.rs @@ -15,6 +15,10 @@ // Memory addresses must be lower than BabyBear prime. const MAX_MEMORY: usize = 0x78000000; +const RESERVED_SIZE: usize = 1024 * 1024; // 1MB reserved section +static mut RESERVED_POS: usize = 0; +static mut RESERVED_START: usize = 0; + #[allow(clippy::missing_safety_doc)] #[no_mangle] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { @@ -32,6 +36,11 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u if heap_pos == 0 { heap_pos = unsafe { (&_end) as *const u8 as usize }; + unsafe { + RESERVED_START = heap_pos; + RESERVED_POS = heap_pos; + heap_pos += RESERVED_SIZE; + } } let offset = heap_pos & (align - 1); @@ -49,3 +58,22 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u unsafe { HEAP_POS = heap_pos }; ptr } +#[allow(clippy::missing_safety_doc)] +#[no_mangle] +pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8 { + let mut pos = RESERVED_POS; + + // Align the position + let offset = pos & (align - 1); + if offset != 0 { + pos += align - offset; + } + + let new_pos = pos + bytes; + if new_pos > RESERVED_START + RESERVED_SIZE { + panic!("Reserved memory section full"); + } + + RESERVED_POS = new_pos; + pos as *mut u8 +} diff --git a/crates/zkvm/lib/src/io.rs b/crates/zkvm/lib/src/io.rs index 94a92a602c..895bdf0cff 100644 --- a/crates/zkvm/lib/src/io.rs +++ b/crates/zkvm/lib/src/io.rs @@ -1,10 +1,7 @@ #![allow(unused_unsafe)] -use crate::{syscall_hint_len, syscall_hint_read, syscall_write}; +use crate::{sys_alloc_reserved, syscall_hint_len, syscall_hint_read, syscall_write}; use serde::{de::DeserializeOwned, Serialize}; -use std::{ - alloc::Layout, - io::{Result, Write}, -}; +use std::io::{Result, Write}; /// The file descriptor for public values. pub const FD_PUBLIC_VALUES: u32 = 3; @@ -43,24 +40,12 @@ impl Write for SyscallWriter { /// let data: Vec = sp1_zkvm::io::read_vec(); /// ``` pub fn read_vec() -> Vec { - // Round up to the nearest multiple of 4 so that the memory allocated is in whole words let len = unsafe { syscall_hint_len() }; let capacity = (len + 3) / 4 * 4; + let ptr = unsafe { sys_alloc_reserved(capacity, 4) }; - // Allocate a buffer of the required length that is 4 byte aligned - let layout = Layout::from_size_align(capacity, 4).expect("vec is too large"); - let ptr = unsafe { std::alloc::alloc(layout) }; - - // SAFETY: - // 1. `ptr` was allocated using alloc - // 2. We assuume that the VM global allocator doesn't dealloc - // 3/6. Size is correct from above - // 4/5. Length is 0 - // 7. Layout::from_size_align already checks this let mut vec = unsafe { Vec::from_raw_parts(ptr, 0, capacity) }; - // Read the vec into uninitialized memory. The syscall assumes the memory is uninitialized, - // which should be true because the allocator does not dealloc, so a new alloc should be fresh. unsafe { syscall_hint_read(ptr, len); vec.set_len(len); diff --git a/crates/zkvm/lib/src/lib.rs b/crates/zkvm/lib/src/lib.rs index e43c4d0d68..d106793f83 100644 --- a/crates/zkvm/lib/src/lib.rs +++ b/crates/zkvm/lib/src/lib.rs @@ -97,6 +97,9 @@ extern "C" { /// Allocates a buffer aligned to the given alignment. pub fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8; + /// Allocates from the reserved memory section + pub fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8; + /// Decompresses a BLS12-381 point. pub fn syscall_bls12381_decompress(point: &mut [u8; 96], is_odd: bool); From dfebebfd509ba712ce1f063fe9c32163504d0ca7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 12 Dec 2024 01:57:22 +0000 Subject: [PATCH 3/4] fix(zkvm): optimize memory usage for low memory configurations Co-Authored-By: John Guibas --- crates/zkvm/entrypoint/src/heap.rs | 2 +- crates/zkvm/entrypoint/src/syscalls/memory.rs | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/zkvm/entrypoint/src/heap.rs b/crates/zkvm/entrypoint/src/heap.rs index 357c621dd0..22b324493f 100644 --- a/crates/zkvm/entrypoint/src/heap.rs +++ b/crates/zkvm/entrypoint/src/heap.rs @@ -5,7 +5,7 @@ use core::sync::atomic::{AtomicUsize, Ordering}; use crate::syscalls::sys_alloc_aligned; /// A block in our free list -#[repr(C)] +#[repr(C, align(8))] struct FreeBlock { size: usize, next: Option>, diff --git a/crates/zkvm/entrypoint/src/syscalls/memory.rs b/crates/zkvm/entrypoint/src/syscalls/memory.rs index 7fa81e4b16..539ea8029e 100644 --- a/crates/zkvm/entrypoint/src/syscalls/memory.rs +++ b/crates/zkvm/entrypoint/src/syscalls/memory.rs @@ -15,7 +15,7 @@ // Memory addresses must be lower than BabyBear prime. const MAX_MEMORY: usize = 0x78000000; -const RESERVED_SIZE: usize = 1024 * 1024; // 1MB reserved section +const RESERVED_SIZE: usize = MAX_MEMORY / 64; static mut RESERVED_POS: usize = 0; static mut RESERVED_START: usize = 0; @@ -37,6 +37,10 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u if heap_pos == 0 { heap_pos = unsafe { (&_end) as *const u8 as usize }; unsafe { + // Check if reserved section would exceed memory limit + if heap_pos + RESERVED_SIZE > MAX_MEMORY { + panic!("Reserved section would exceed memory limit"); + } RESERVED_START = heap_pos; RESERVED_POS = heap_pos; heap_pos += RESERVED_SIZE; @@ -63,7 +67,6 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8 { let mut pos = RESERVED_POS; - // Align the position let offset = pos & (align - 1); if offset != 0 { pos += align - offset; From 1c5004b9d0afd3eb0e56921ad5a7b6c0e3d5af2e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 12 Dec 2024 02:39:59 +0000 Subject: [PATCH 4/4] fix(zkvm): improve memory management for low memory configurations Co-Authored-By: John Guibas --- crates/zkvm/entrypoint/src/syscalls/memory.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/crates/zkvm/entrypoint/src/syscalls/memory.rs b/crates/zkvm/entrypoint/src/syscalls/memory.rs index 539ea8029e..fc98acd558 100644 --- a/crates/zkvm/entrypoint/src/syscalls/memory.rs +++ b/crates/zkvm/entrypoint/src/syscalls/memory.rs @@ -15,13 +15,21 @@ // Memory addresses must be lower than BabyBear prime. const MAX_MEMORY: usize = 0x78000000; -const RESERVED_SIZE: usize = MAX_MEMORY / 64; +const RESERVED_SIZE: usize = MAX_MEMORY / 256; static mut RESERVED_POS: usize = 0; static mut RESERVED_START: usize = 0; #[allow(clippy::missing_safety_doc)] #[no_mangle] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { + if !align.is_power_of_two() { + panic!("Alignment must be power of 2"); + } + + if bytes.checked_add(align - 1).is_none() { + panic!("Memory allocation would overflow"); + } + extern "C" { // https://lld.llvm.org/ELF/linker_script.html#sections-command static _end: u8; @@ -66,6 +74,7 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u #[no_mangle] pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8 { let mut pos = RESERVED_POS; + debug_assert!(align.is_power_of_two()); let offset = pos & (align - 1); if offset != 0 { @@ -73,7 +82,7 @@ pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut } let new_pos = pos + bytes; - if new_pos > RESERVED_START + RESERVED_SIZE { + if new_pos < pos || new_pos > RESERVED_START + RESERVED_SIZE { panic!("Reserved memory section full"); }