feat: ensure that guest only returns a limited set of UDFs

Author: crepererum

guests/evil/src/complex/mod.rs

@@ -8,6 +8,7 @@ use datafusion_expr::{
 
 pub(crate) mod udf_long_name;
 pub(crate) mod udfs_duplicate_names;
+pub(crate) mod udfs_many;
 
 /// UDF with a name
 #[derive(Debug, PartialEq, Eq, Hash)]

guests/evil/src/complex/udfs_many.rs

@@ -0,0 +1,19 @@
+//! Create many UDFs.
+use std::sync::Arc;
+
+use datafusion_common::Result as DataFusionResult;
+use datafusion_expr::ScalarUDFImpl;
+
+use crate::complex::NamedUdf;
+
+/// Returns our evil UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+
+    Ok((0..=limit)
+        .map(|i| Arc::new(NamedUdf(i.to_string())) as _)
+        .collect())
+}

guests/evil/src/lib.rs

@@ -43,6 +43,10 @@ impl Evil {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(complex::udfs_duplicate_names::udfs),
             },
+            "complex::udfs_many" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::udfs_many::udfs),
+            },
             "env" => Self {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(env::udfs),

host/src/lib.rs

@@ -274,6 +274,9 @@ pub struct WasmPermissions {
     /// Trusted data limits.
     trusted_data_limits: TrustedDataLimits,
 
+    /// Maximum number of UDFs.
+    max_udfs: usize,
+
     /// Environment variables.
     envs: BTreeMap<String, String>,
 }
@@ -300,6 +303,7 @@ impl Default for WasmPermissions {
             stderr_bytes: 1024, // 1KB
             resource_limits: StaticResourceLimits::default(),
             trusted_data_limits: TrustedDataLimits::default(),
+            max_udfs: 20,
             envs: BTreeMap::default(),
         }
     }
@@ -378,6 +382,19 @@ impl WasmPermissions {
         }
     }
 
+    /// Get the maximum number of UDFs that a payload/guest can produce.
+    pub fn max_udfs(&self) -> usize {
+        self.max_udfs
+    }
+
+    /// Set the maximum number of UDFs that a payload/guest can produce.
+    pub fn with_max_udfs(self, limit: usize) -> Self {
+        Self {
+            max_udfs: limit,
+            ..self
+        }
+    }
+
     /// Add environment variable.
     pub fn with_env(mut self, key: String, value: String) -> Self {
         self.envs.insert(key, value);
@@ -578,6 +595,13 @@ impl WasmScalarUdf {
             )?
             .convert_err(permissions.trusted_data_limits.clone())
             .context("scalar_udfs")?;
+        if udf_resources.len() > permissions.max_udfs {
+            return Err(DataFusionError::ResourcesExhausted(format!(
+                "guest returned too many UDFs: got={}, limit={}",
+                udf_resources.len(),
+                permissions.max_udfs,
+            )));
+        }
 
         let store = Arc::new(Mutex::new(store));
 

host/tests/integration_tests/evil/complex.rs

@@ -1,4 +1,4 @@
-use datafusion_udf_wasm_host::conversion::limits::TrustedDataLimits;
+use datafusion_udf_wasm_host::{WasmPermissions, conversion::limits::TrustedDataLimits};
 
 use crate::integration_tests::evil::test_utils::{try_scalar_udfs, try_scalar_udfs_with_env};
 
@@ -31,3 +31,17 @@ async fn test_udfs_duplicate_names() {
         err,
         @"External error: non-unique UDF name: 'foo'");
 }
+
+#[tokio::test]
+async fn test_udfs_many() {
+    let err = try_scalar_udfs_with_env(
+        "complex::udfs_many",
+        &[("limit", &WasmPermissions::default().max_udfs().to_string())],
+    )
+    .await
+    .unwrap_err();
+
+    insta::assert_snapshot!(
+        err,
+        @"Resources exhausted: guest returned too many UDFs: got=21, limit=20");
+}