Merge pull request #171 from influxdata/crepererum/more-evil

test: extend evil test suite

Author: crepererum

guests/evil/src/common.rs

@@ -0,0 +1,21 @@
+//! Methods used by multiple payloads.
+
+use std::sync::Arc;
+
+use datafusion_common::Result as DataFusionResult;
+use datafusion_expr::ScalarUDFImpl;
+
+/// Return empty root file system.
+///
+/// This always returns [`None`] because the example does not need any files.
+pub(crate) fn root_empty() -> Option<Vec<u8>> {
+    None
+}
+
+/// Returns empty list of UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs_empty(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    Ok(vec![])
+}

guests/evil/src/env.rs

@@ -0,0 +1,121 @@
+//! Payload that tries to read environment variables.
+use std::{hash::Hash, sync::Arc};
+
+use arrow::datatypes::DataType;
+use datafusion_common::{Result as DataFusionResult, ScalarValue};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+
+/// UDF that produces a string.
+struct StringUdf {
+    /// Name.
+    name: &'static str,
+
+    /// String producer.
+    effect: Box<dyn Fn() -> Option<String> + Send + Sync>,
+
+    /// Signature of the UDF.
+    ///
+    /// We store this here because [`ScalarUDFImpl::signature`] requires us to return a reference.
+    signature: Signature,
+}
+
+impl StringUdf {
+    /// Create new UDF.
+    fn new<F>(name: &'static str, effect: F) -> Self
+    where
+        F: Fn() -> Option<String> + Send + Sync + 'static,
+    {
+        Self {
+            name,
+            effect: Box::new(effect),
+            signature: Signature::uniform(0, vec![], Volatility::Immutable),
+        }
+    }
+}
+
+impl std::fmt::Debug for StringUdf {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Self {
+            name,
+            effect: _,
+            signature,
+        } = self;
+
+        f.debug_struct("StringUdf")
+            .field("name", name)
+            .field("effect", &"<EFFECT>")
+            .field("signature", signature)
+            .finish()
+    }
+}
+
+impl PartialEq<Self> for StringUdf {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+    }
+}
+
+impl Eq for StringUdf {}
+
+impl Hash for StringUdf {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+    }
+}
+
+impl ScalarUDFImpl for StringUdf {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Ok(DataType::Utf8)
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Ok(ColumnarValue::Scalar(ScalarValue::Utf8((self.effect)())))
+    }
+}
+
+/// 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(StringUdf::new("args", || {
+            let vars = std::env::args().collect::<Vec<_>>();
+            if vars.is_empty() {
+                None
+            } else {
+                Some(vars.join(","))
+            }
+        })),
+        Arc::new(StringUdf::new("current_dir", || {
+            let d = std::env::current_dir().unwrap();
+            Some(d.display().to_string())
+        })),
+        Arc::new(StringUdf::new("current_exe", || {
+            let e = std::env::current_exe().unwrap_err();
+            Some(e.to_string())
+        })),
+        Arc::new(StringUdf::new("env", || {
+            let vars = std::env::vars()
+                .map(|(k, v)| format!("{k}:{v}"))
+                .collect::<Vec<_>>();
+            if vars.is_empty() {
+                None
+            } else {
+                Some(vars.join(","))
+            }
+        })),
+    ])
+}