Merge pull request #220 from influxdata/crepererum/sandbox-error-edition

test: ensure error conversion is guarded

Author: crepererum

guests/evil/src/complex/error.rs

@@ -0,0 +1,85 @@
+//! Badness related to errors returned by the UDF.
+
+use std::sync::Arc;
+
+use arrow::datatypes::DataType;
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+
+use crate::common::DynBox;
+
+/// UDF that returns an error when invoked.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct ErrorUDF {
+    /// Name.
+    name: &'static str,
+
+    /// The error generator.
+    err: DynBox<dyn Fn() -> DataFusionError + Send + Sync>,
+}
+
+impl ErrorUDF {
+    /// Create new UDF.
+    fn new<F>(name: &'static str, f: F) -> Self
+    where
+        F: Fn() -> DataFusionError + Send + Sync + 'static,
+    {
+        Self {
+            name,
+            err: DynBox(Box::new(f)),
+        }
+    }
+}
+
+impl ScalarUDFImpl for ErrorUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    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> {
+        Err((self.err)())
+    }
+}
+
+/// 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(ErrorUDF::new("long_ctx", || {
+            let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+            DataFusionError::Execution("foo".to_owned())
+                .context(std::iter::repeat_n('x', limit + 1).collect::<String>())
+        })),
+        Arc::new(ErrorUDF::new("long_msg", || {
+            let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+            DataFusionError::Execution(std::iter::repeat_n('x', limit + 1).collect())
+        })),
+        Arc::new(ErrorUDF::new("nested_ctx", || {
+            let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+            (0..=limit).fold(DataFusionError::Execution("foo".to_owned()), |err, _| {
+                err.context("bar")
+            })
+        })),
+    ])
+}

guests/evil/src/complex/mod.rs

@@ -6,6 +6,7 @@ use datafusion_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
 };
 
+pub(crate) mod error;
 pub(crate) mod udf_long_name;
 pub(crate) mod udfs_duplicate_names;
 pub(crate) mod udfs_many;

guests/evil/src/lib.rs

@@ -36,6 +36,10 @@ impl Evil {
     /// Get evil, multiplexed by env.
     fn get() -> Self {
         match std::env::var("EVIL").expect("evil specified").as_str() {
+            "complex::error" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::error::udfs),
+            },
             "complex::udf_long_name" => Self {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(complex::udf_long_name::udfs),

host/tests/integration_tests/evil/complex.rs

@@ -1,7 +1,62 @@
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{DataFusionError, config::ConfigOptions};
+use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, async_udf::AsyncScalarUDFImpl};
 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_err_long_ctx() {
+    let err = run_err_udf(
+        "long_ctx",
+        TrustedDataLimits::default().max_aux_string_length,
+    )
+    .await;
+
+    insta::assert_snapshot!(
+        err,
+        @r"
+    convert error from WASI
+    caused by
+    Resources exhausted: auxiliary string length: got=10001, limit=10000
+    ",
+    );
+}
+
+#[tokio::test]
+async fn test_err_long_msg() {
+    let err = run_err_udf(
+        "long_msg",
+        TrustedDataLimits::default().max_aux_string_length,
+    )
+    .await;
+
+    insta::assert_snapshot!(
+        err,
+        @r"
+    convert error from WASI
+    caused by
+    Resources exhausted: auxiliary string length: got=10001, limit=10000
+    ",
+    );
+}
+
+#[tokio::test]
+async fn test_err_nested_ctx() {
+    let err = run_err_udf("nested_ctx", TrustedDataLimits::default().max_depth as _).await;
+
+    insta::assert_snapshot!(
+        err,
+        @r"
+    convert error from WASI
+    caused by
+    Resources exhausted: data structure depth: limit=10
+    ",
+    );
+}
+
 #[tokio::test]
 async fn test_udf_long_name() {
     let err = try_scalar_udfs_with_env(
@@ -45,3 +100,23 @@ async fn test_udfs_many() {
         err,
         @"Resources exhausted: guest returned too many UDFs: got=21, limit=20");
 }
+
+/// Test UDF related to Error handling.
+async fn run_err_udf(name: &'static str, limit: usize) -> DataFusionError {
+    let udf = try_scalar_udfs_with_env("complex::error", &[("limit", &limit.to_string())])
+        .await
+        .unwrap()
+        .into_iter()
+        .find(|udf| udf.name() == name)
+        .unwrap();
+
+    udf.invoke_async_with_args(ScalarFunctionArgs {
+        args: vec![],
+        arg_fields: vec![],
+        number_rows: 1,
+        return_field: Arc::new(Field::new("r", DataType::Null, true)),
+        config_options: Arc::new(ConfigOptions::default()),
+    })
+    .await
+    .unwrap_err()
+}