Security-harden the guest's memory allocator

ceno_rt/src/allocator.rs

@@ -4,26 +4,34 @@
 use core::alloc::{GlobalAlloc, Layout};
 
 struct SimpleAllocator {
-    next_alloc: usize,
+    next_alloc: *mut u8,
 }
 
+extern "C" {
+    // The address of this variable is the start of the heap (growing upwards).
+    static mut _sheap: u8;
+}
+
+#[global_allocator]
+static mut HEAP: SimpleAllocator = SimpleAllocator {
+    next_alloc: &raw mut _sheap,
+};
+
 unsafe impl GlobalAlloc for SimpleAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         // SAFETY: Single threaded, so nothing else can touch this while we're working.
         let mut heap_pos = HEAP.next_alloc;
 
         let align = layout.align();
         // `Layout` contract forbids making a `Layout` with align=0, or align not power of 2.
-        // So we can safely use subtraction and a mask to ensure alignment without worrying about UB.
-        let offset = heap_pos & (align - 1);
-        if offset != 0 {
-            heap_pos = heap_pos.strict_add(align.strict_sub(offset));
-        }
+        core::hint::assert_unchecked(align.is_power_of_two());
+        core::hint::assert_unchecked(align != 0);
+        heap_pos = heap_pos.add(heap_pos.align_offset(align));
 
-        let ptr = heap_pos as *mut u8;
+        let ptr = heap_pos;
         // Panic on overflow.  We don't want to wrap around, and overwrite stack etc.
         // (We could also return a null pointer, but only malicious programs would ever hit this.)
-        heap_pos = heap_pos.strict_add(layout.size());
+        heap_pos = heap_pos.add(layout.size());
 
         HEAP.next_alloc = heap_pos;
         ptr
@@ -36,14 +44,3 @@ unsafe impl GlobalAlloc for SimpleAllocator {
     /// Never deallocate.
     unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
 }
-
-#[global_allocator]
-// We initialize `next_alloc` to 0xFFFF_FFFF to indicate that the heap has not been initialized.
-// The value is chosen to make any premature allocation fail.
-static mut HEAP: SimpleAllocator = SimpleAllocator {
-    next_alloc: 0xFFFF_FFFF,
-};
-
-pub unsafe fn init_heap() {
-    HEAP.next_alloc = core::ptr::from_ref::<u8>(&crate::_sheap).cast::<u8>() as usize;
-}

ceno_rt/src/lib.rs

@@ -79,19 +79,14 @@ macro_rules! entry {
 /// _start_rust is called by the assembly entry point and it calls the Rust main().
 #[no_mangle]
 unsafe extern "C" fn _start_rust() -> ! {
-    allocator::init_heap();
-    {
-        extern "C" {
-            fn bespoke_entrypoint();
-        }
-        bespoke_entrypoint();
+    extern "C" {
+        fn bespoke_entrypoint();
     }
+    bespoke_entrypoint();
     halt(0)
 }
 
 extern "C" {
     // The address of this variable is the start of the stack (growing downwards).
     static _stack_start: u8;
-    // The address of this variable is the start of the heap (growing upwards).
-    static _sheap: u8;
 }