Merge pull request #219 from influxdata/crepererum/sandbox-udf-edition

feat: sandbox lock-down -- "set of UDFs" edition

Author: crepererum

guests/evil/src/complex/mod.rs

@@ -0,0 +1,42 @@
+//! Overly complex data emitted by the payload.
+
+use arrow::datatypes::DataType;
+use datafusion_common::{Result as DataFusionResult, ScalarValue};
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+
+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)]
+struct NamedUdf(String);
+
+impl ScalarUDFImpl for NamedUdf {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.0
+    }
+
+    fn signature(&self) -> &Signature {
+        static S: Signature = Signature {
+            type_signature: TypeSignature::Uniform(0, vec![]),
+            volatility: Volatility::Immutable,
+        };
+
+        &S
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Ok(DataType::Null)
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Ok(ColumnarValue::Scalar(ScalarValue::Null))
+    }
+}

guests/evil/src/complex/udf_long_name.rs

@@ -0,0 +1,23 @@
+//! Long UDF names.
+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();
+    let name = std::iter::repeat_n('x', limit + 1).collect::<String>();
+
+    Ok(vec![
+        // Emit twice. This shouldn't trigger the "duplicate names" detection though because the length MUST be
+        // checked BEFORE potentially hashing the names.
+        Arc::new(NamedUdf(name.clone())),
+        Arc::new(NamedUdf(name)),
+    ])
+}

guests/evil/src/complex/udfs_duplicate_names.rs

@@ -0,0 +1,19 @@
+//! Duplicate UDF names.
+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>>> {
+    Ok(vec![
+        Arc::new(NamedUdf("foo".to_owned())),
+        Arc::new(NamedUdf("bar".to_owned())),
+        Arc::new(NamedUdf("foo".to_owned())),
+    ])
+}

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

@@ -8,6 +8,7 @@ use datafusion_expr::ScalarUDFImpl;
 use datafusion_udf_wasm_guest::export;
 
 mod common;
+mod complex;
 mod env;
 mod fs;
 mod net;
@@ -35,6 +36,18 @@ impl Evil {
     /// Get evil, multiplexed by env.
     fn get() -> Self {
         match std::env::var("EVIL").expect("evil specified").as_str() {
+            "complex::udf_long_name" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::udf_long_name::udfs),
+            },
+            "complex::udfs_duplicate_names" => Self {
+                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

@@ -2,7 +2,14 @@
 //!
 //!
 //! [DataFusion]: https://datafusion.apache.org/
-use std::{any::Any, collections::BTreeMap, hash::Hash, ops::DerefMut, sync::Arc, time::Duration};
+use std::{
+    any::Any,
+    collections::{BTreeMap, HashSet},
+    hash::Hash,
+    ops::DerefMut,
+    sync::Arc,
+    time::Duration,
+};
 
 use ::http::HeaderName;
 use arrow::datatypes::DataType;
@@ -33,7 +40,7 @@ use wasmtime_wasi_http::{
 
 use crate::{
     bindings::exports::datafusion_udf_wasm::udf::types as wit_types,
-    conversion::limits::{CheckedInto, TrustedDataLimits},
+    conversion::limits::{CheckedInto, ComplexityToken, TrustedDataLimits},
     error::{DataFusionResultExt, WasmToDataFusionResultExt, WitDataFusionResultExt},
     http::{HttpRequestValidator, RejectAllHttpRequests},
     limiter::{Limiter, StaticResourceLimits},
@@ -267,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>,
 }
@@ -293,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(),
         }
     }
@@ -371,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);
@@ -571,10 +595,18 @@ 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));
 
         let mut udfs = Vec::with_capacity(udf_resources.len());
+        let mut names_seen = HashSet::with_capacity(udf_resources.len());
         for resource in udf_resources {
             let mut store_guard = store.lock().await;
             let store2: &mut Store<WasmStateImpl> = &mut store_guard;
@@ -587,6 +619,13 @@ impl WasmScalarUdf {
                     "call ScalarUdf::name",
                     Some(&store_guard.data().stderr.contents()),
                 )?;
+            ComplexityToken::new(permissions.trusted_data_limits.clone())?
+                .check_identifier(&name)?;
+            if !names_seen.insert(name.clone()) {
+                return Err(DataFusionError::External(
+                    format!("non-unique UDF name: '{name}'").into(),
+                ));
+            }
 
             let store2: &mut Store<WasmStateImpl> = &mut store_guard;
             let signature: Signature = bindings

host/tests/integration_tests/evil/complex.rs

@@ -0,0 +1,47 @@
+use datafusion_udf_wasm_host::{WasmPermissions, conversion::limits::TrustedDataLimits};
+
+use crate::integration_tests::evil::test_utils::{try_scalar_udfs, try_scalar_udfs_with_env};
+
+#[tokio::test]
+async fn test_udf_long_name() {
+    let err = try_scalar_udfs_with_env(
+        "complex::udf_long_name",
+        &[(
+            "limit",
+            &TrustedDataLimits::default()
+                .max_identifier_length
+                .to_string(),
+        )],
+    )
+    .await
+    .unwrap_err();
+
+    insta::assert_snapshot!(
+        err,
+        @"Resources exhausted: identifier length: got=51, limit=50");
+}
+
+#[tokio::test]
+async fn test_udfs_duplicate_names() {
+    let err = try_scalar_udfs("complex::udfs_duplicate_names")
+        .await
+        .unwrap_err();
+
+    insta::assert_snapshot!(
+        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");
+}

host/tests/integration_tests/evil/mod.rs

@@ -1,3 +1,4 @@
+mod complex;
 mod env;
 mod fs;
 mod net;