Strengthen alloc arithmetic against edge cases (#1720)

If a contract were, for some inexplicable reason, to manually and
_unsafely_ construct a malformed `Layout` and pass it into the bump
allocator, they could cause it to misbehave. While this represents no
risk to any normal code, or indeed any safe client code at all, it turns
out that just eliminating the entire space of potential misbehavior from
the allocator is both easy and makes it more-obviously more-correct. So
we do that here.

Author: graydon

soroban-sdk/src/alloc.rs

@@ -1,12 +1,28 @@
 // This code is adapted from https://github.com/wenyuzhao/bump-allocator-rs
+//
+// We've altered it to work entirely with usize values internally and only cast
+// back to an exposed-provenance pointer when returning from alloc. This gives
+// us a richer checked-arithmetic API we can use to trap overflows internally,
+// and also avoids some potential UB issues with pointer provenance. Since the
+// provenance of __heap_base is a 1-byte value anyway, and all the rest of the
+// wasm heap is considered to have exposed provenance, we think this is the best
+// we can do. Writing allocators is tricky!
+//
+// NB: technically these alterations only handle corner cases that cannot be hit
+// using safe client code. Safe clients pass `Layout` structs that always meet
+// additional size and alignment constraints. But hardening the code to tolerate
+// even _unsafe_ inputs -- malformed `Layout` inputs one can only create by
+// calling unsafe methods -- is not only easy to do, it makes the code simpler
+// and more readable, so we went ahead and did it.
 
+use crate::unwrap::UnwrapOptimized;
 use core::alloc::{GlobalAlloc, Layout};
 
-pub static mut LOCAL_ALLOCATOR: BumpPointerLocal = BumpPointerLocal::new();
+static mut LOCAL_ALLOCATOR: BumpPointerLocal = BumpPointerLocal::new();
 
-pub struct BumpPointerLocal {
-    cursor: *mut u8,
-    limit: *mut u8,
+struct BumpPointerLocal {
+    cursor: usize,
+    limit: usize,
 }
 
 impl BumpPointerLocal {
@@ -16,24 +32,20 @@ impl BumpPointerLocal {
 
     pub const fn new() -> Self {
         Self {
-            cursor: 0 as _,
-            limit: 0 as _,
+            cursor: 0,
+            limit: 0,
         }
     }
 
-    #[inline(always)]
-    fn align_allocation(cursor: *mut u8, align: usize) -> *mut u8 {
-        let mask = align - 1;
-        (((cursor as usize) + mask) & !mask) as _
-    }
-
     #[inline(always)]
     fn maybe_init_inline(&mut self) {
-        if self.limit as usize == 0 {
+        if self.limit == 0 {
             // This is a slight over-estimate and ideally we would use __heap_base
             // but that seems not to be easy to access and in any case it is just a
             // convention whereas this is more guaranteed by the wasm spec to work.
-            self.cursor = (core::arch::wasm32::memory_size(Self::MEM) * Self::PAGE_SIZE) as _;
+            self.cursor = core::arch::wasm32::memory_size(Self::MEM)
+                .checked_mul(Self::PAGE_SIZE)
+                .unwrap_optimized();
             self.limit = self.cursor;
         }
     }
@@ -43,11 +55,15 @@ impl BumpPointerLocal {
         self.maybe_init_inline()
     }
 
+    // Allocate `bytes` bytes with `align` alignment.
     #[inline(always)]
-    pub fn alloc(&mut self, bytes: usize, align: usize) -> *mut u8 {
+    fn alloc(&mut self, bytes: usize, align: usize) -> usize {
         self.maybe_init();
-        let start = Self::align_allocation(self.cursor, align);
-        let new_cursor = unsafe { start.add(bytes) };
+        let start = self
+            .cursor
+            .checked_next_multiple_of(align)
+            .unwrap_optimized();
+        let new_cursor = start.checked_add(bytes).unwrap_optimized();
         if new_cursor <= self.limit {
             self.cursor = new_cursor;
             start
@@ -57,17 +73,18 @@ impl BumpPointerLocal {
     }
 
     #[inline(always)]
-    fn alloc_slow_inline(&mut self, bytes: usize, align: usize) -> *mut u8 {
-        let pages = (bytes + Self::PAGE_SIZE - 1) / Self::PAGE_SIZE;
+    fn alloc_slow_inline(&mut self, bytes: usize, align: usize) -> usize {
+        let pages = bytes.div_ceil(Self::PAGE_SIZE);
         if core::arch::wasm32::memory_grow(Self::MEM, pages) == usize::MAX {
-            panic!()
+            core::arch::wasm32::unreachable();
         }
-        self.limit = unsafe { self.limit.add(pages * Self::PAGE_SIZE) };
+        let bytes_grown = pages.checked_mul(Self::PAGE_SIZE).unwrap_optimized();
+        self.limit = self.limit.checked_add(bytes_grown).unwrap_optimized();
         self.alloc(bytes, align)
     }
 
     #[inline(never)]
-    fn alloc_slow(&mut self, bytes: usize, align: usize) -> *mut u8 {
+    fn alloc_slow(&mut self, bytes: usize, align: usize) -> usize {
         self.alloc_slow_inline(bytes, align)
     }
 }
@@ -80,7 +97,7 @@ unsafe impl GlobalAlloc for BumpPointer {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         let (bytes, align) = (layout.size(), layout.align());
         let ptr = LOCAL_ALLOCATOR.alloc(bytes, align);
-        ptr
+        core::ptr::with_exposed_provenance_mut(ptr)
     }
 
     #[inline(always)]