feat: limit memory consumption

Closes #14.

Author: crepererum

Cargo.lock

@@ -45,6 +45,10 @@ impl Evil {
                 root: Box::new(root::invalid_entry::root),
                 udfs: Box::new(common::udfs_empty),
             },
+            "root::large_file" => Self {
+                root: Box::new(root::large_file::root),
+                udfs: Box::new(common::udfs_empty),
+            },
             "root::many_files" => Self {
                 root: Box::new(root::many_files::root),
                 udfs: Box::new(common::udfs_empty),
@@ -57,6 +61,10 @@ impl Evil {
                 root: Box::new(root::path_long::root),
                 udfs: Box::new(common::udfs_empty),
             },
+            "root::sparse" => Self {
+                root: Box::new(root::sparse::root),
+                udfs: Box::new(common::udfs_empty),
+            },
             "root::unsupported_entry" => Self {
                 root: Box::new(root::unsupported_entry::root),
                 udfs: Box::new(common::udfs_empty),

guests/evil/src/lib.rs

@@ -0,0 +1,19 @@
+//! Evil payloads that creates a LARGE file.
+
+/// Return root file system.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn root() -> Option<Vec<u8>> {
+    let mut ar = tar::Builder::new(Vec::new());
+
+    let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+    let data = vec![0u8; limit];
+
+    let mut header = tar::Header::new_gnu();
+    header.set_path("foo").unwrap();
+    header.set_size(0);
+    header.set_cksum();
+
+    ar.append(&header, data.as_slice()).unwrap();
+
+    Some(ar.into_inner().unwrap())
+}

guests/evil/src/root/large_file.rs

@@ -1,6 +1,8 @@
 //! Evil things concerning the root file system.
 pub(crate) mod invalid_entry;
+pub(crate) mod large_file;
 pub(crate) mod many_files;
 pub(crate) mod not_tar;
 pub(crate) mod path_long;
+pub(crate) mod sparse;
 pub(crate) mod unsupported_entry;

guests/evil/src/root/mod.rs

@@ -0,0 +1,9 @@
+//! Evil payloads that returns a HUGE sparse tar file (i.e. the file itself is small, but the sparse payload is big).
+
+/// Return root file system.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn root() -> Option<Vec<u8>> {
+    // Sparse TARs are really hard to get right using the public APIs in `tar`, hence we just include a test file
+    // directly.
+    Some(include_bytes!("sparse-large.tar").to_vec())
+}

guests/evil/src/root/sparse-large.tar

@@ -109,6 +109,21 @@ fn ackermann(m: u128, n: u128) -> u128 {
 #[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
 pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
     Ok(vec![
+        Arc::new(SideEffect::new("alloc", || {
+            let data = vec![1u8; 1_000_000_000];
+            // side effect to prevent optimization
+            for x in data {
+                println!("{x}");
+            }
+        })),
+        Arc::new(SideEffect::new("alloc_try", || {
+            let mut data = Vec::<u8>::with_capacity(0);
+            data.try_reserve(1_000_000_000).unwrap();
+            // side effect to prevent optimization
+            for x in data {
+                println!("{x}");
+            }
+        })),
         Arc::new(SideEffect::new("fillstderr", || {
             let s: String = std::iter::repeat_n('x', 10_000).collect();
             for _ in 0..10_000 {
@@ -144,6 +159,7 @@ pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImp
             }
         })),
         Arc::new(SideEffect::new("panic", || panic!("foo"))),
+        Arc::new(SideEffect::new("pass", || ())),
         Arc::new(SideEffect::new("stackoverflow", || {
             // simulate a side-effect via I/O so the compiler cannot optimize this method away
             println!("{}", ackermann(10, 10));

guests/evil/src/root/sparse.rs

@@ -8,6 +8,7 @@ license.workspace = true
 anyhow.workspace = true
 arrow.workspace = true
 datafusion-common.workspace = true
+datafusion-execution.workspace = true
 datafusion-expr.workspace = true
 datafusion-udf-wasm-arrow2bytes.workspace = true
 http.workspace = true

guests/evil/src/runtime.rs

@@ -7,6 +7,7 @@ use std::{any::Any, collections::BTreeMap, hash::Hash, ops::DerefMut, sync::Arc,
 use ::http::HeaderName;
 use arrow::datatypes::DataType;
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_execution::memory_pool::MemoryPool;
 use datafusion_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
     async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl},
@@ -34,6 +35,7 @@ use crate::{
     bindings::exports::datafusion_udf_wasm::udf::types as wit_types,
     error::{DataFusionResultExt, WasmToDataFusionResultExt},
     http::{HttpRequestValidator, RejectAllHttpRequests},
+    limiter::{Limiter, StaticResourceLimits},
     linker::link,
     tokio_helpers::async_in_sync_context,
     vfs::{VfsCtxView, VfsLimits, VfsState, VfsView},
@@ -51,6 +53,7 @@ mod bindings;
 mod conversion;
 pub mod error;
 pub mod http;
+pub mod limiter;
 mod linker;
 mod tokio_helpers;
 pub mod vfs;
@@ -62,6 +65,9 @@ struct WasmStateImpl {
     /// This filesystem is provided to the payload in memory with read-write support.
     vfs_state: VfsState,
 
+    /// Resource limiter.
+    limiter: Limiter,
+
     /// A limited buffer for stderr.
     ///
     /// This is especially useful for when the payload crashes.
@@ -87,6 +93,7 @@ impl std::fmt::Debug for WasmStateImpl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let Self {
             vfs_state,
+            limiter,
             stderr,
             wasi_ctx: _,
             wasi_http_ctx: _,
@@ -96,6 +103,7 @@ impl std::fmt::Debug for WasmStateImpl {
         } = self;
         f.debug_struct("WasmStateImpl")
             .field("vfs_state", vfs_state)
+            .field("limiter", limiter)
             .field("stderr", stderr)
             .field("wasi_ctx", &"<WASI_CTX>")
             .field("resource_table", resource_table)
@@ -249,6 +257,12 @@ pub struct WasmPermissions {
     /// Virtual file system limits.
     vfs: VfsLimits,
 
+    /// Limit of the stored stderr data.
+    stderr_bytes: usize,
+
+    /// Static resource limits.
+    resource_limits: StaticResourceLimits,
+
     /// Environment variables.
     envs: BTreeMap<String, String>,
 }
@@ -272,6 +286,8 @@ impl Default for WasmPermissions {
                 .floor() as _,
             http: Arc::new(RejectAllHttpRequests),
             vfs: VfsLimits::default(),
+            stderr_bytes: 1024, // 1KB
+            resource_limits: StaticResourceLimits::default(),
             envs: BTreeMap::default(),
         }
     }
@@ -316,6 +332,24 @@ impl WasmPermissions {
         }
     }
 
+    /// Limit of the stored stderr data.
+    pub fn with_stderr_bytes(self, limit: usize) -> Self {
+        Self {
+            stderr_bytes: limit,
+            ..self
+        }
+    }
+
+    /// Set static resource limits.
+    ///
+    /// Note that this does NOT limit the overall memory consumption of the payload. This will be done via [`MemoryPool`].
+    pub fn with_resource_limits(self, limits: StaticResourceLimits) -> Self {
+        Self {
+            resource_limits: limits,
+            ..self
+        }
+    }
+
     /// Set virtual filesystem limits.
     pub fn with_vfs_limits(self, limits: VfsLimits) -> Self {
         Self {
@@ -406,6 +440,7 @@ impl WasmScalarUdf {
         component: &WasmComponentPrecompiled,
         permissions: &WasmPermissions,
         io_rt: Handle,
+        memory_pool: &Arc<dyn MemoryPool>,
         source: String,
     ) -> DataFusionResult<Vec<Self>> {
         let WasmComponentPrecompiled { compiled_component } = component;
@@ -446,31 +481,35 @@ impl WasmScalarUdf {
         let component_res = unsafe { Component::deserialize(&engine, compiled_component) };
         let component = component_res.context("create WASM component", None)?;
 
+        // resource/mem limiter
+        let mut limiter = Limiter::new(permissions.resource_limits.clone(), memory_pool);
+
         // Create in-memory VFS
         let vfs_state = VfsState::new(permissions.vfs.clone());
 
         // set up WASI p2 context
-        let stderr = MemoryOutputPipe::new(1024);
+        limiter.grow(permissions.stderr_bytes)?;
+        let stderr = MemoryOutputPipe::new(permissions.stderr_bytes);
         let mut wasi_ctx_builder = WasiCtx::builder();
         wasi_ctx_builder.stderr(stderr.clone());
         permissions.envs.iter().for_each(|(k, v)| {
             wasi_ctx_builder.env(k, v);
         });
 
+        // configure store
+        // NOTE: Do that BEFORE linking so that memory limits are checked for the initial allocation of the WASM
+        //       component as well.
         let state = WasmStateImpl {
             vfs_state,
+            limiter,
             stderr,
             wasi_ctx: wasi_ctx_builder.build(),
             wasi_http_ctx: WasiHttpCtx::new(),
             resource_table: ResourceTable::new(),
             http_validator: Arc::clone(&permissions.http),
             io_rt,
         };
-        let (bindings, mut store) = link(&engine, &component, state)
-            .await
-            .context("link WASM components", None)?;
-
-        // configure store
+        let mut store = Store::new(&engine, state);
         store.epoch_deadline_callback(|_| {
             Ok(UpdateDeadline::YieldCustom(
                 // increment deadline epoch by one step
@@ -482,6 +521,11 @@ impl WasmScalarUdf {
                 Box::pin(tokio::task::consume_budget()),
             ))
         });
+        store.limiter(|state| &mut state.limiter);
+
+        let bindings = link(&engine, &component, &mut store)
+            .await
+            .context("link WASM components", None)?;
 
         // Populate VFS from tar archive
         let root_data = bindings
@@ -493,11 +537,12 @@ impl WasmScalarUdf {
                 Some(&store.data().stderr.contents()),
             )?;
         if let Some(root_data) = root_data {
-            store
-                .data_mut()
+            let state = store.data_mut();
+
+            state
                 .vfs_state
-                .populate_from_tar(&root_data)
-                .map_err(DataFusionError::IoError)?;
+                .populate_from_tar(&root_data, &mut state.limiter)
+                .map_err(|e| DataFusionError::IoError(e).context("populate root FS from TAR"))?;
         }
 
         let udf_resources = bindings