feat: ensure that UDF names are unique

Author: crepererum

guests/evil/src/complex/mod.rs

@@ -0,0 +1,3 @@
+//! Overly complex data emitted by the payload.
+
+pub(crate) mod udfs_duplicate_names;

guests/evil/src/complex/udfs_duplicate_names.rs

@@ -0,0 +1,51 @@
+//! 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,
+};
+
+/// 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))
+    }
+}
+
+/// 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")),
+    ])
+}

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;
@@ -34,6 +35,10 @@ impl Evil {
     /// Get evil, multiplexed by env.
     fn get() -> Self {
         match std::env::var("EVIL").expect("evil specified").as_str() {
+            "complex::udfs_duplicate_names" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::udfs_duplicate_names::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;
@@ -575,6 +582,7 @@ impl WasmScalarUdf {
         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 +595,11 @@ impl WasmScalarUdf {
                     "call ScalarUdf::name",
                     Some(&store_guard.data().stderr.contents()),
                 )?;
+            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,12 @@
+use crate::integration_tests::evil::test_utils::try_scalar_udfs;
+
+#[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'");
+}

host/tests/integration_tests/evil/mod.rs

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