fix(proto): check plan integer conversions across usize boundaries

datafusion/common/src/utils/mod.rs

@@ -25,7 +25,9 @@ pub mod proxy;
 pub mod string_utils;
 
 use crate::assert_or_internal_err;
-use crate::error::{_exec_datafusion_err, _exec_err, _internal_datafusion_err};
+use crate::error::{
+    _exec_datafusion_err, _exec_err, _internal_datafusion_err, _plan_datafusion_err,
+};
 use crate::{Result, ScalarValue};
 use arrow::array::{
     Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
@@ -1142,6 +1144,32 @@ pub fn combine_limit(
     (combined_skip, combined_fetch)
 }
 
+/// Converts a wire integer to `usize`, rejecting out-of-range values.
+/// `context` and `field` identify the value in the error message.
+pub fn usize_from_wire<T>(value: T, context: &str, field: &str) -> Result<usize>
+where
+    T: TryInto<usize> + std::fmt::Display + Copy,
+{
+    value.try_into().map_err(|_| {
+        _plan_datafusion_err!(
+            "{context}: {field} wire value {value} is out of range for usize"
+        )
+    })
+}
+
+/// Converts a `usize` to a wire integer, rejecting out-of-range values.
+pub fn usize_to_wire<T: TryFrom<usize>>(
+    value: usize,
+    context: &str,
+    field: &str,
+) -> Result<T> {
+    T::try_from(value).map_err(|_| {
+        _plan_datafusion_err!(
+            "{context}: {field} value {value} is out of range for the plan wire format"
+        )
+    })
+}
+
 /// Returns the estimated number of threads available for parallel execution.
 ///
 /// This is a wrapper around `std::thread::available_parallelism`, providing a default value
@@ -1482,6 +1510,30 @@ mod tests {
     #[cfg(feature = "sql")]
     use sqlparser::ast::Ident;
 
+    #[test]
+    fn test_usize_wire_conversions() {
+        assert_eq!(usize_from_wire(42_u64, "SomeExec", "fetch").unwrap(), 42);
+        let err = usize_from_wire(-1_i64, "SomeExec", "skip").unwrap_err();
+        assert_eq!(
+            err.strip_backtrace(),
+            "Error during planning: SomeExec: skip wire value -1 is out of range for usize"
+        );
+
+        assert_eq!(usize_to_wire::<u32>(42, "SomeExec", "fetch").unwrap(), 42);
+        let err = usize_to_wire::<u8>(256, "SomeExec", "fetch").unwrap_err();
+        assert_eq!(
+            err.strip_backtrace(),
+            "Error during planning: SomeExec: fetch value 256 is out of range for the plan wire format"
+        );
+
+        let max = usize_from_wire(u64::MAX, "SomeExec", "fetch");
+        if usize::BITS >= u64::BITS {
+            assert_eq!(max.unwrap(), usize::MAX);
+        } else {
+            assert!(max.is_err());
+        }
+    }
+
     #[test]
     fn test_bisect_linear_left_and_right() -> Result<()> {
         let arrays: Vec<ArrayRef> = vec![

datafusion/datasource/src/file_scan_config/proto.rs

@@ -39,6 +39,7 @@ use std::sync::Arc;
 
 use arrow::datatypes::Schema;
 use datafusion_common::parsers::CompressionTypeVariant;
+use datafusion_common::utils::{usize_from_wire, usize_to_wire};
 use datafusion_common::{DataFusionError, Result, internal_datafusion_err};
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs};
@@ -128,7 +129,11 @@ impl FileScanConfig {
         Ok(protobuf::FileScanExecConf {
             file_groups,
             statistics: Some((&self.statistics()).into()),
-            limit: self.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }),
+            limit: self
+                .limit
+                .map(|limit| usize_to_wire::<u32>(limit, "FileScanConfig", "limit"))
+                .transpose()?
+                .map(|limit| protobuf::ScanLimit { limit }),
             projection: vec![],
             schema: Some((&schema).try_into()?),
             table_partition_cols: self
@@ -245,14 +250,28 @@ impl FileScanConfig {
             file_source
         };
 
+        let limit = conf
+            .limit
+            .as_ref()
+            .map(|limit| usize_from_wire(limit.limit, "FileScanConfig", "limit"))
+            .transpose()?;
+        let batch_size = conf
+            .batch_size
+            .map(|size| usize_from_wire(size, "FileScanConfig", "batch_size"))
+            .transpose()?;
+        if batch_size == Some(0) {
+            return datafusion_common::plan_err!(
+                "FileScanConfig: batch_size must be greater than 0"
+            );
+        }
         let config_builder = FileScanConfigBuilder::new(object_store_url, file_source)
             .with_file_groups(file_groups)
             .with_constraints(constraints)
             .with_statistics(statistics)
-            .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize))
+            .with_limit(limit)
             .with_output_ordering(output_ordering)
             .with_output_partitioning(output_partitioning)
-            .with_batch_size(conf.batch_size.map(|s| s as usize))
+            .with_batch_size(batch_size)
             .with_file_compression_type(file_compression_type);
         Ok(config_builder.build())
     }

datafusion/datasource/src/memory.rs

@@ -280,6 +280,7 @@ impl DataSource for MemorySourceConfig {
         &self,
         ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
     ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_common::utils::usize_to_wire;
         use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto;
         use datafusion_proto_models::protobuf;
 
@@ -315,7 +316,12 @@ impl DataSource for MemorySourceConfig {
                         projection,
                         sort_information,
                         show_sizes: self.show_sizes,
-                        fetch: self.fetch.map(|f| f as u32),
+                        fetch: self
+                            .fetch
+                            .map(|fetch| {
+                                usize_to_wire(fetch, "MemoryScanExecNode", "fetch")
+                            })
+                            .transpose()?,
                     },
                 ),
             ),
@@ -682,6 +688,7 @@ impl MemorySourceConfig {
         ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
     ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
         use datafusion_common::internal_datafusion_err;
+        use datafusion_common::utils::usize_from_wire;
         use datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto;
         use datafusion_proto_models::protobuf;
 
@@ -718,8 +725,12 @@ impl MemorySourceConfig {
             sort_information.extend(LexOrdering::new(sort_exprs));
         }
 
+        let fetch = scan
+            .fetch
+            .map(|fetch| usize_from_wire(fetch, "MemoryScanExecNode", "fetch"))
+            .transpose()?;
         let source = Self::try_new(&partitions, schema, projection)?
-            .with_limit(scan.fetch.map(|f| f as usize))
+            .with_limit(fetch)
             .with_show_sizes(scan.show_sizes)
             .try_with_sort_information(sort_information)?;
 

datafusion/expr/src/logical_plan/plan.rs

@@ -65,7 +65,7 @@ use datafusion_common::{
     FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result,
     ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions,
     aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err,
-    internal_err, plan_err, validate_range_split_points,
+    internal_err, plan_datafusion_err, plan_err, validate_range_split_points,
 };
 use indexmap::IndexSet;
 use itertools::Itertools as _;
@@ -3782,7 +3782,12 @@ impl Limit {
                     // `skip = NULL` is equivalent to `skip = 0`
                     let s = s.unwrap_or(0);
                     if s >= 0 {
-                        Ok(SkipType::Literal(s as usize))
+                        let s = usize::try_from(s).map_err(|_| {
+                            plan_datafusion_err!(
+                                "OFFSET value {s} cannot be represented as usize"
+                            )
+                        })?;
+                        Ok(SkipType::Literal(s))
                     } else {
                         plan_err!("OFFSET must be >=0, '{}' was provided", s)
                     }
@@ -3800,7 +3805,12 @@ impl Limit {
             Some(expr) => match *expr {
                 Expr::Literal(ScalarValue::Int64(Some(s)), _) => {
                     if s >= 0 {
-                        Ok(FetchType::Literal(Some(s as usize)))
+                        let s = usize::try_from(s).map_err(|_| {
+                            plan_datafusion_err!(
+                                "LIMIT value {s} cannot be represented as usize"
+                            )
+                        })?;
+                        Ok(FetchType::Literal(Some(s)))
                     } else {
                         plan_err!("LIMIT must be >= 0, '{}' was provided", s)
                     }
@@ -5031,6 +5041,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn limit_literals_use_checked_usize_conversion() -> Result<()> {
+        let value = i64::from(u32::MAX) + 1;
+        let input = Arc::new(LogicalPlanBuilder::empty(false).build()?);
+        let limit = Limit {
+            skip: Some(Box::new(lit(value))),
+            fetch: Some(Box::new(lit(value))),
+            input,
+        };
+
+        if usize::BITS < 64 {
+            assert!(limit.get_skip_type().is_err());
+            assert!(limit.get_fetch_type().is_err());
+        } else {
+            let expected = usize::try_from(value).unwrap();
+            let SkipType::Literal(skip) = limit.get_skip_type()? else {
+                panic!("expected literal skip")
+            };
+            let FetchType::Literal(Some(fetch)) = limit.get_fetch_type()? else {
+                panic!("expected literal fetch")
+            };
+            assert_eq!(skip, expected);
+            assert_eq!(fetch, expected);
+        }
+        Ok(())
+    }
+
     fn employee_schema() -> Schema {
         Schema::new(vec![
             Field::new("id", DataType::Int32, false),

datafusion/functions-table/src/generate_series.rs

@@ -291,6 +291,9 @@ impl GenerateSeriesTable {
         &self,
         batch_size: usize,
     ) -> Result<Arc<RwLock<dyn LazyBatchGenerator>>> {
+        if batch_size == 0 {
+            return plan_err!("GenerateSeriesTable: batch_size must be greater than 0");
+        }
         let generator: Arc<RwLock<dyn LazyBatchGenerator>> = match &self.args {
             GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })),
             GenSeriesArgs::Int64Args {
@@ -877,7 +880,26 @@ mod generate_series_tests {
     use datafusion_common::Result;
     use datafusion_physical_plan::memory::LazyBatchGenerator;
 
-    use crate::generate_series::GenericSeriesState;
+    use crate::generate_series::{
+        GenSeriesArgs, GenerateSeriesTable, GenericSeriesState,
+    };
+
+    #[test]
+    fn generate_series_rejects_zero_batch_size() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
+        let table = GenerateSeriesTable::new(
+            schema,
+            GenSeriesArgs::Int64Args {
+                start: 1,
+                end: 2,
+                step: 1,
+                include_end: true,
+                name: "generate_series",
+            },
+        );
+
+        assert!(table.as_generator(0).is_err());
+    }
 
     #[test]
     fn test_generic_series_state_reset() -> Result<()> {