feat: check over overlong UDF names

Author: crepererum

guests/evil/src/complex/mod.rs

@@ -1,3 +1,41 @@
 //! 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;
+
+/// 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

@@ -1,51 +1,19 @@
 //! Duplicate UDF names.
 use std::sync::Arc;
 
-use arrow::datatypes::DataType;
-use datafusion_common::{Result as DataFusionResult, ScalarValue};
-use datafusion_expr::{
-    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
-};
+use datafusion_common::Result as DataFusionResult;
+use datafusion_expr::ScalarUDFImpl;
 
-/// UDF with a name
-#[derive(Debug, PartialEq, Eq, Hash)]
-struct NamedUdf(&'static str);
-
-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))
-    }
-}
+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")),
-        Arc::new(NamedUdf("bar")),
-        Arc::new(NamedUdf("foo")),
+        Arc::new(NamedUdf("foo".to_owned())),
+        Arc::new(NamedUdf("bar".to_owned())),
+        Arc::new(NamedUdf("foo".to_owned())),
     ])
 }

guests/evil/src/lib.rs

@@ -35,6 +35,10 @@ 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),

host/src/lib.rs

@@ -40,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},
@@ -595,6 +595,8 @@ 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(),

host/tests/integration_tests/evil/complex.rs

@@ -1,4 +1,25 @@
-use crate::integration_tests::evil::test_utils::try_scalar_udfs;
+use datafusion_udf_wasm_host::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() {