Merge pull request #225 from supabase/custom-array-buffer-allocator

feat: enforce limits on array buffer allocations

Author: laktek

crates/base/src/deno_runtime.rs

@@ -24,6 +24,7 @@ use event_worker::js_interceptors::sb_events_js_interceptors;
 use event_worker::sb_user_event_worker;
 use sb_core::cache::CacheSetting;
 use sb_core::cert::ValueRootCertStoreProvider;
+use sb_core::external_memory::custom_allocator;
 use sb_core::http_start::sb_core_http;
 use sb_core::net::sb_core_net;
 use sb_core::permissions::{sb_core_permissions, Permissions};
@@ -265,19 +266,20 @@ impl DenoRuntime {
             sb_core_runtime::init_ops(Some(main_module_url.clone())),
         ];
 
+        let mut create_params = None;
+        if conf.is_user_worker() {
+            let memory_limit =
+                mib_to_bytes(conf.as_user_worker().unwrap().memory_limit_mb) as usize;
+            create_params = Some(
+                deno_core::v8::CreateParams::default()
+                    .heap_limits(mib_to_bytes(0) as usize, memory_limit)
+                    .array_buffer_allocator(custom_allocator(memory_limit)),
+            )
+        };
         let runtime_options = RuntimeOptions {
             extensions,
             is_main: true,
-            create_params: {
-                if conf.is_user_worker() {
-                    Some(deno_core::v8::CreateParams::default().heap_limits(
-                        mib_to_bytes(0) as usize,
-                        mib_to_bytes(conf.as_user_worker().unwrap().memory_limit_mb) as usize,
-                    ))
-                } else {
-                    None
-                }
-            },
+            create_params,
             get_error_class_fn: Some(&get_error_class_name),
             shared_array_buffer_store: None,
             compiled_wasm_module_store: Default::default(),
@@ -903,4 +905,27 @@ mod test {
             _ => panic!("Invalid Result"),
         };
     }
+
+    #[tokio::test]
+    async fn test_array_buffer_allocation_below_limit() {
+        let user_rt = create_basic_user_runtime("./test_cases/array_buffers", 20, 1000).await;
+        let (_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStream>();
+        let result = user_rt.run(unix_stream_rx).await;
+        assert!(result.is_ok(), "expected no errors");
+    }
+
+    #[tokio::test]
+    async fn test_array_buffer_allocation_above_limit() {
+        let user_rt = create_basic_user_runtime("./test_cases/array_buffers", 15, 1000).await;
+        let (_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStream>();
+        let result = user_rt.run(unix_stream_rx).await;
+        match result {
+            Err(err) => {
+                assert!(err
+                    .to_string()
+                    .contains("RangeError: Array buffer allocation failed"));
+            }
+            _ => panic!("Invalid Result"),
+        };
+    }
 }

crates/base/test_cases/array_buffers/index.ts

@@ -0,0 +1 @@
+const arr = new ArrayBuffer(17 * 1024 * 1024);

crates/sb_core/external_memory.rs

@@ -0,0 +1,82 @@
+use deno_core::v8;
+use deno_core::v8::UniqueRef;
+use std::ffi::c_void;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+pub struct CustomAllocator {
+    max: usize,
+    count: AtomicUsize,
+}
+
+#[allow(clippy::unnecessary_cast)]
+unsafe extern "C" fn allocate(allocator: &CustomAllocator, n: usize) -> *mut c_void {
+    allocator.count.fetch_add(n, Ordering::SeqCst);
+    let count_loaded = allocator.count.load(Ordering::SeqCst);
+    if count_loaded > allocator.max {
+        return std::ptr::null::<*mut [u8]>() as *mut c_void;
+    }
+
+    Box::into_raw(vec![0u8; n].into_boxed_slice()) as *mut [u8] as *mut c_void
+}
+
+#[allow(clippy::unnecessary_cast)]
+#[allow(clippy::uninit_vec)]
+unsafe extern "C" fn allocate_uninitialized(allocator: &CustomAllocator, n: usize) -> *mut c_void {
+    allocator.count.fetch_add(n, Ordering::SeqCst);
+    let count_loaded = allocator.count.load(Ordering::SeqCst);
+    if count_loaded > allocator.max {
+        return std::ptr::null::<*mut [u8]>() as *mut c_void;
+    }
+
+    let mut store = Vec::with_capacity(n);
+    store.set_len(n);
+    Box::into_raw(store.into_boxed_slice()) as *mut [u8] as *mut c_void
+}
+
+unsafe extern "C" fn free(allocator: &CustomAllocator, data: *mut c_void, n: usize) {
+    allocator.count.fetch_sub(n, Ordering::SeqCst);
+    let _ = Box::from_raw(std::slice::from_raw_parts_mut(data as *mut u8, n));
+}
+
+#[allow(clippy::unnecessary_cast)]
+unsafe extern "C" fn reallocate(
+    allocator: &CustomAllocator,
+    prev: *mut c_void,
+    oldlen: usize,
+    newlen: usize,
+) -> *mut c_void {
+    allocator
+        .count
+        .fetch_add(newlen.wrapping_sub(oldlen), Ordering::SeqCst);
+    let count_loaded = allocator.count.load(Ordering::SeqCst);
+    if count_loaded > allocator.max {
+        return std::ptr::null::<*mut [u8]>() as *mut c_void;
+    }
+
+    let old_store = Box::from_raw(std::slice::from_raw_parts_mut(prev as *mut u8, oldlen));
+    let mut new_store = Vec::with_capacity(newlen);
+    let copy_len = oldlen.min(newlen);
+    new_store.extend_from_slice(&old_store[..copy_len]);
+    new_store.resize(newlen, 0u8);
+    Box::into_raw(new_store.into_boxed_slice()) as *mut [u8] as *mut c_void
+}
+
+unsafe extern "C" fn drop(allocator: *const CustomAllocator) {
+    Arc::from_raw(allocator);
+}
+
+pub fn custom_allocator(max: usize) -> UniqueRef<deno_core::v8::Allocator> {
+    let vtable: &'static v8::RustAllocatorVtable<CustomAllocator> = &v8::RustAllocatorVtable {
+        allocate,
+        allocate_uninitialized,
+        free,
+        reallocate,
+        drop,
+    };
+    let allocator = Arc::new(CustomAllocator {
+        count: AtomicUsize::new(0),
+        max,
+    });
+    unsafe { v8::new_rust_allocator(Arc::into_raw(allocator), vtable) }
+}

crates/sb_core/lib.rs

@@ -3,6 +3,7 @@ pub mod cache;
 pub mod cert;
 pub mod emit;
 pub mod errors_rt;
+pub mod external_memory;
 pub mod file_fetcher;
 pub mod http_start;
 pub mod net;