Impose a 64MiB ArrayBuffer memory limit per isolate (#37851)

V8's ArrayBuffer allocates outside of the V8 heap and doesn't abide the normal V8 heap limit - instead the default allocator directly calls `malloc`/`calloc`/`free`. This is very sad and means that JS can crash the entire process with an OOM.
We can install our own allocator instead that tracks the amount of memory handed out, and put a limit of 64MiB.

GitOrigin-RevId: 47f744add45dbda5e588a499877d3c7a3338f393

Author: goffrie

Cargo.lock

@@ -793,6 +793,10 @@ pub static ISOLATE_MAX_USER_HEAP_SIZE: LazyLock<usize> =
 pub static ISOLATE_MAX_HEAP_EXTRA_SIZE: LazyLock<usize> =
     LazyLock::new(|| env_config("ISOLATE_MAX_HEAP_EXTRA_SIZE", 1 << 25));
 
+/// Set a separate 64MB limit on ArrayBuffer allocations.
+pub static ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE: LazyLock<usize> =
+    LazyLock::new(|| env_config("ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE", 1 << 26));
+
 /// Chunk sizes: 1, 2, 3, ..., MAX_DYNAMIC_SMART_CHUNK_SIZE incrementing by 1.
 /// These chunk sizes allow small (common) batches to be handled in a single
 /// chunk, while limiting the size of a chunk (don't overload the db), and

crates/common/src/knobs.rs

@@ -45,6 +45,7 @@ http-body-util = { workspace = true }
 humansize = { workspace = true }
 itertools = { workspace = true }
 keybroker = { path = "../keybroker" }
+libc = { workspace = true }
 maplit = { workspace = true, optional = true }
 metrics = { path = "../metrics" }
 mime = { workspace = true }

crates/isolate/Cargo.toml

@@ -0,0 +1,108 @@
+use std::{
+    alloc::{
+        alloc,
+        alloc_zeroed,
+        dealloc,
+        Layout,
+    },
+    ffi::c_void,
+    mem,
+    ptr,
+    sync::{
+        atomic::{
+            AtomicUsize,
+            Ordering,
+        },
+        Arc,
+    },
+};
+
+use deno_core::v8::{
+    new_rust_allocator,
+    Allocator,
+    RustAllocatorVtable,
+    UniqueRef,
+};
+
+pub struct ArrayBufferMemoryLimit {
+    available: AtomicUsize,
+    limit: usize,
+}
+
+impl ArrayBufferMemoryLimit {
+    /// Returns the amount of memory used by ArrayBuffers.
+    /// This can be an overestimate since it includes buffers that would be GC'd
+    /// but haven't yet.
+    pub fn used(&self) -> usize {
+        self.limit
+            .saturating_sub(self.available.load(Ordering::Relaxed))
+    }
+
+    fn consume(&self, amount: usize) -> bool {
+        let mut limit = self.available.load(Ordering::Relaxed);
+        loop {
+            if limit < amount {
+                crate::metrics::log_array_buffer_oom();
+                return false;
+            }
+            match self.available.compare_exchange_weak(
+                limit,
+                limit - amount,
+                Ordering::Relaxed,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => return true,
+                Err(v) => limit = v,
+            }
+        }
+    }
+}
+
+const ALIGNMENT: usize = mem::align_of::<libc::max_align_t>();
+
+unsafe extern "C" fn allocate(handle: &ArrayBufferMemoryLimit, len: usize) -> *mut c_void {
+    if handle.consume(len) {
+        alloc_zeroed(Layout::from_size_align_unchecked(len, ALIGNMENT)).cast()
+    } else {
+        ptr::null_mut()
+    }
+}
+unsafe extern "C" fn allocate_uninitialized(
+    handle: &ArrayBufferMemoryLimit,
+    len: usize,
+) -> *mut c_void {
+    if handle.consume(len) {
+        alloc(Layout::from_size_align_unchecked(len, ALIGNMENT)).cast()
+    } else {
+        ptr::null_mut()
+    }
+}
+unsafe extern "C" fn free(handle: &ArrayBufferMemoryLimit, data: *mut c_void, len: usize) {
+    handle.available.fetch_add(len, Ordering::Relaxed);
+    dealloc(
+        data.cast(),
+        Layout::from_size_align_unchecked(len, ALIGNMENT),
+    );
+}
+unsafe extern "C" fn drop(handle: *const ArrayBufferMemoryLimit) {
+    Arc::from_raw(handle);
+}
+
+pub fn limited_array_buffer_allocator(
+    limit: usize,
+) -> (Arc<ArrayBufferMemoryLimit>, UniqueRef<Allocator>) {
+    unsafe {
+        let limit = Arc::new(ArrayBufferMemoryLimit {
+            available: AtomicUsize::new(limit),
+            limit,
+        });
+        const VTABLE: RustAllocatorVtable<ArrayBufferMemoryLimit> = RustAllocatorVtable {
+            allocate,
+            allocate_uninitialized,
+            free,
+            drop,
+        };
+        let allocator = new_rust_allocator(Arc::into_raw(limit.clone()), &VTABLE);
+        (limit, allocator)
+    }
+}

crates/isolate/src/array_buffer_allocator.rs

@@ -11,6 +11,7 @@ use std::{
 use anyhow::{
     anyhow,
     bail,
+    Context as _,
 };
 use async_recursion::async_recursion;
 use common::{
@@ -36,6 +37,7 @@ use serde_json::Value as JsonValue;
 use value::heap_size::HeapSize;
 
 use crate::{
+    array_buffer_allocator::ArrayBufferMemoryLimit,
     bundled_js::system_udf_file,
     environment::{
         IsolateEnvironment,
@@ -200,13 +202,18 @@ impl<'a, 'b: 'a, RT: Runtime, E: IsolateEnvironment<RT>> ExecutionScope<'a, 'b,
 
     pub fn record_heap_stats(&mut self) -> anyhow::Result<()> {
         let stats = self.get_heap_statistics();
+        let array_buffer_size = self
+            .get_slot::<Arc<ArrayBufferMemoryLimit>>()
+            .context("missing ArrayBufferMemoryLimit?")?
+            .used();
         self.with_state_mut(|state| {
             let blobs_heap_size = state.blob_parts.heap_size();
             let streams_heap_size = state.streams.heap_size() + state.stream_listeners.heap_size();
             state.environment.record_heap_stats(IsolateHeapStats::new(
                 stats,
                 blobs_heap_size,
                 streams_heap_size,
+                array_buffer_size,
             ));
         })
     }

crates/isolate/src/execution_scope.rs

@@ -10,6 +10,7 @@ use anyhow::Context as _;
 use common::{
     knobs::{
         FUNRUN_INITIAL_PERMIT_TIMEOUT,
+        ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE,
         ISOLATE_MAX_USER_HEAP_SIZE,
     },
     runtime::Runtime,
@@ -31,6 +32,7 @@ use humansize::{
 use value::heap_size::WithHeapSize;
 
 use crate::{
+    array_buffer_allocator::ArrayBufferMemoryLimit,
     concurrency_limiter::ConcurrencyLimiter,
     environment::IsolateEnvironment,
     helpers::pump_message_loop,
@@ -64,6 +66,7 @@ pub struct Isolate<RT: Runtime> {
     // we reclaim after removing the callback.
     heap_ctx_ptr: *mut HeapContext,
     limiter: ConcurrencyLimiter,
+    array_buffer_memory_limit: Arc<ArrayBufferMemoryLimit>,
 
     created: tokio::time::Instant,
 }
@@ -124,13 +127,15 @@ pub struct IsolateHeapStats {
 
     pub blobs_heap_size: usize,
     pub streams_heap_size: usize,
+    pub array_buffer_size: usize,
 }
 
 impl IsolateHeapStats {
     pub fn new(
         stats: v8::HeapStatistics,
         blobs_heap_size: usize,
         streams_heap_size: usize,
+        array_buffer_size: usize,
     ) -> Self {
         Self {
             v8_total_heap_size: stats.total_heap_size(),
@@ -142,6 +147,7 @@ impl IsolateHeapStats {
             environment_heap_size: 0,
             blobs_heap_size,
             streams_heap_size,
+            array_buffer_size,
         }
     }
 
@@ -153,7 +159,13 @@ impl IsolateHeapStats {
 impl<RT: Runtime> Isolate<RT> {
     pub fn new(rt: RT, max_user_timeout: Option<Duration>, limiter: ConcurrencyLimiter) -> Self {
         let _timer = create_isolate_timer();
-        let mut v8_isolate = crate::udf_runtime::create_isolate_with_udf_runtime();
+        let (array_buffer_memory_limit, array_buffer_allocator) =
+            crate::array_buffer_allocator::limited_array_buffer_allocator(
+                *ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE,
+            );
+        let mut v8_isolate = crate::udf_runtime::create_isolate_with_udf_runtime(
+            v8::CreateParams::default().array_buffer_allocator(array_buffer_allocator),
+        );
 
         // Tells V8 to capture current stack trace when uncaught exception occurs and
         // report it to the message listeners. The option is off by default.
@@ -189,6 +201,7 @@ impl<RT: Runtime> Isolate<RT> {
         );
 
         assert!(v8_isolate.set_slot(handle.clone()));
+        assert!(v8_isolate.set_slot(array_buffer_memory_limit.clone()));
 
         Self {
             created: rt.monotonic_now(),
@@ -198,6 +211,7 @@ impl<RT: Runtime> Isolate<RT> {
             heap_ctx_ptr,
             max_user_timeout,
             limiter,
+            array_buffer_memory_limit,
         }
     }
 
@@ -235,7 +249,7 @@ impl<RT: Runtime> Isolate<RT> {
     // Heap stats for an isolate that has no associated state or environment.
     pub fn heap_stats(&mut self) -> IsolateHeapStats {
         let stats = self.v8_isolate.get_heap_statistics();
-        IsolateHeapStats::new(stats, 0, 0)
+        IsolateHeapStats::new(stats, 0, 0, self.array_buffer_memory_limit.used())
     }
 
     pub fn check_isolate_clean(&mut self) -> Result<(), IsolateNotClean> {

crates/isolate/src/isolate.rs

@@ -8,7 +8,8 @@ pub struct Thread {
 
 impl Thread {
     pub fn new() -> Self {
-        let mut isolate = crate::udf_runtime::create_isolate_with_udf_runtime();
+        let mut isolate =
+            crate::udf_runtime::create_isolate_with_udf_runtime(v8::CreateParams::default());
 
         // Tells V8 to capture current stack trace when uncaught exception occurs and
         // report it to the message listeners. The option is off by default.

crates/isolate/src/isolate2/thread.rs

@@ -9,6 +9,7 @@
 #![feature(impl_trait_in_assoc_type)]
 #![feature(round_char_boundary)]
 
+mod array_buffer_allocator;
 pub mod bundled_js;
 pub mod client;
 mod concurrency_limiter;

crates/isolate/src/lib.rs

@@ -356,6 +356,14 @@ pub fn log_system_timeout() {
     log_counter(&UDF_SYSTEM_TIMEOUT_TOTAL, 1);
 }
 
+register_convex_counter!(
+    ARRAY_BUFFER_OOM_TOTAL,
+    "Number of times that isolates hit the ArrayBuffer memory limit"
+);
+pub fn log_array_buffer_oom() {
+    log_counter(&ARRAY_BUFFER_OOM_TOTAL, 1);
+}
+
 register_convex_counter!(
     RECREATE_ISOLATE_TOTAL,
     "Number of times an isolate is recreated",
@@ -500,6 +508,10 @@ register_convex_gauge!(
     ISOLATE_TOTAL_MALLOCED_MEMORY_BYTES,
     "Total isolate malloc'd memory across all isolates"
 );
+register_convex_gauge!(
+    ISOLATE_TOTAL_ARRAY_BUFFER_MEMORY_BYTES,
+    "Total isolate ArrayBuffer-allocated memory across all isolates"
+);
 
 pub fn log_aggregated_heap_stats(stats: &IsolateHeapStats) {
     log_gauge(
@@ -526,6 +538,10 @@ pub fn log_aggregated_heap_stats(stats: &IsolateHeapStats) {
         &ISOLATE_TOTAL_MALLOCED_MEMORY_BYTES,
         stats.v8_malloced_memory as f64,
     );
+    log_gauge(
+        &ISOLATE_TOTAL_ARRAY_BUFFER_MEMORY_BYTES,
+        stats.array_buffer_size as f64,
+    );
 }
 
 register_convex_histogram!(UDF_FETCH_SECONDS, "Duration of UDF fetch", &STATUS_LABEL);

crates/isolate/src/metrics.rs

@@ -14,6 +14,7 @@ use common::{
         ComponentPath,
         PublicFunctionPath,
     },
+    knobs::ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE,
     log_lines::{
         LogLines,
         TRUNCATED_LINE_SUFFIX,
@@ -972,3 +973,53 @@ async fn test_subfunction_depth(rt: TestRuntime) -> anyhow::Result<()> {
     assert_eq!(result, ConvexValue::Float64(8.0));
     Ok(())
 }
+
+#[convex_macro::test_runtime]
+async fn test_array_buffer_size_limit(rt: TestRuntime) -> anyhow::Result<()> {
+    let t = UdfTest::default(rt).await?;
+    // Allocating an ArrayBuffer over the size limit fails
+    let e = t
+        .query_js_error(
+            "adversarial:allocateArrayBuffers",
+            assert_obj!(
+                "size" => (*ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE + 1) as f64,
+                "count" => 1.0,
+                "retain" => 1.0,
+            ),
+        )
+        .await?;
+    assert_contains(&e, "RangeError: Array buffer allocation failed");
+    // However, we can repeatedly allocate under the size limit and succeed as
+    // long as the old ones can be GC'd
+    t.query(
+        "adversarial:allocateArrayBuffers",
+        assert_obj!(
+            "size" => *ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE as f64,
+            "count" => 10.0,
+            "retain" => 0.0,
+        ),
+    )
+    .await?;
+    t.query(
+        "adversarial:allocateArrayBuffers",
+        assert_obj!(
+            "size" => (*ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE / 2) as f64,
+            "count" => 10.0,
+            "retain" => 1.0,
+        ),
+    )
+    .await?;
+    // ... but retaining too many ArrayBuffers will fail
+    let e = t
+        .query_js_error(
+            "adversarial:allocateArrayBuffers",
+            assert_obj!(
+                "size" => (*ISOLATE_MAX_ARRAY_BUFFER_TOTAL_SIZE / 2) as f64,
+                "count" => 10.0,
+                "retain" => 2.0,
+            ),
+        )
+        .await?;
+    assert_contains(&e, "RangeError: Array buffer allocation failed");
+    Ok(())
+}

crates/isolate/src/tests/adversarial.rs

@@ -40,12 +40,12 @@ const INITIAL_HEAP_SIZE: usize = 1 << 16;
 
 /// Creates a new V8 isolate from the saved snapshot. Contexts created in this
 /// isolate will have the UDF runtime already loaded.
-pub(crate) fn create_isolate_with_udf_runtime() -> v8::OwnedIsolate {
+pub(crate) fn create_isolate_with_udf_runtime(create_params: v8::CreateParams) -> v8::OwnedIsolate {
     let snapshot = BASE_SNAPSHOT
         .get()
         .expect("udf_runtime::initialize not called");
     v8::Isolate::new(
-        v8::CreateParams::default()
+        create_params
             .heap_limits(
                 INITIAL_HEAP_SIZE,
                 *ISOLATE_MAX_USER_HEAP_SIZE + *ISOLATE_MAX_HEAP_EXTRA_SIZE,