Add a recursion limit to the evaluation of type_expr & parse_expr

mamcx · mamcx · commit 1ae2a88bb2db · 2025-07-24T11:44:22.000-05:00
diff --git a/crates/core/src/sql/execute.rs b/crates/core/src/sql/execute.rs
@@ -1126,6 +1126,9 @@ pub(crate) mod tests {
         Ok(())
     }
 
+    /// Test we are protected against recursion when:
+    /// 1. The query is too large
+    /// 2. The AST is too deep
     #[test]
     fn test_large_query_no_panic() -> ResultTest<()> {
         let db = TestDB::durable()?;
@@ -1138,16 +1141,46 @@ pub(crate) mod tests {
             )
             .unwrap();
 
-        let mut query = "select * from test where ".to_string();
-        for x in 0..1_000 {
-            for y in 0..1_000 {
-                let fragment = format!("((x = {x}) and y = {y}) or");
-                query.push_str(&fragment);
+        let build_query = |total| {
+            let mut sql = "select * from test where ".to_string();
+            for x in 0..total {
+                for y in 0..total {
+                    let fragment = format!("((x = {x}) and (y = {y})) or ");
+                    sql.push_str(&fragment);
+                }
             }
+            sql.push_str("((x = 1000) and (y = 1000))");
+            sql
+        };
+        let run = |db: &RelationalDB, sep: char, sql_text: &str| {
+            run_for_testing(db, sql_text).map_err(|e| e.to_string().split(sep).next().unwrap_or_default().to_string())
+        };
+        let sql = build_query(1_000);
+        assert_eq!(
+            run(&db, ':', &sql),
+            Err("SQL query exceeds maximum allowed length".to_string())
+        );
+
+        // Exercise the limit [recursion::MAX_RECURSION_EXPR] && [recursion::MAX_RECURSION_TYP_EXPR]
+        let sql = build_query(8);
+        assert_eq!(run(&db, ',', &sql), Err("Recursion limit exceeded".to_string()));
+
+        let sql = build_query(7);
+        assert!(run(&db, ',', &sql).is_ok(), "Expected query to run without panic");
+
+        // Check no overflow with lot of joins
+        let mut sql = "SELECT test.* FROM test ".to_string();
+        // We could pust up to 700 joins without overflow as long we don't have any conditions,
+        // but here execution become too slow.
+        // TODO: Move this test to the `Plan`
+        for i in 0..200 {
+            sql.push_str(&format!("JOIN test AS m{i} ON test.x = m{i}.y "));
         }
-        query.push_str("((x = 1000) and (y = 1000))");
 
-        assert!(run_for_testing(&db, &query).is_err());
+        assert!(
+            run(&db, ',', &sql).is_ok(),
+            "Query with many joins and conditions should not overflow"
+        );
         Ok(())
     }
 
diff --git a/crates/expr/src/check.rs b/crates/expr/src/check.rs
@@ -99,7 +99,7 @@ pub trait TypeChecker {
                     vars.insert(rhs.alias.clone(), rhs.schema.clone());
 
                     if let Some(on) = on {
-                        if let Expr::BinOp(BinOp::Eq, a, b) = type_expr(vars, on, Some(&AlgebraicType::Bool))? {
+                        if let Expr::BinOp(BinOp::Eq, a, b) = type_expr(vars, on, Some(&AlgebraicType::Bool), &mut 0)? {
                             if let (Expr::Field(a), Expr::Field(b)) = (*a, *b) {
                                 join = RelExpr::EqJoin(LeftDeepJoin { lhs, rhs }, a, b);
                                 continue;
diff --git a/crates/expr/src/lib.rs b/crates/expr/src/lib.rs
@@ -19,6 +19,7 @@ use spacetimedb_sats::algebraic_type::fmt::fmt_algebraic_type;
 use spacetimedb_sats::algebraic_value::ser::ValueSerializer;
 use spacetimedb_schema::schema::ColumnSchema;
 use spacetimedb_sql_parser::ast::{self, BinOp, ProjectElem, SqlExpr, SqlIdent, SqlLiteral};
+use spacetimedb_sql_parser::parser::recursion;
 
 pub mod check;
 pub mod errors;
@@ -30,7 +31,7 @@ pub mod statement;
 pub(crate) fn type_select(input: RelExpr, expr: SqlExpr, vars: &Relvars) -> TypingResult<RelExpr> {
     Ok(RelExpr::Select(
         Box::new(input),
-        type_expr(vars, expr, Some(&AlgebraicType::Bool))?,
+        type_expr(vars, expr, Some(&AlgebraicType::Bool), &mut 0)?,
     ))
 }
 
@@ -68,7 +69,7 @@ pub(crate) fn type_proj(input: RelExpr, proj: ast::Project, vars: &Relvars) -> T
                     return Err(DuplicateName(alias.into_string()).into());
                 }
 
-                if let Expr::Field(p) = type_expr(vars, expr.into(), None)? {
+                if let Expr::Field(p) = type_expr(vars, expr.into(), None, &mut 0)? {
                     projections.push((alias, p));
                 }
             }
@@ -79,7 +80,14 @@ pub(crate) fn type_proj(input: RelExpr, proj: ast::Project, vars: &Relvars) -> T
 }
 
 /// Type check and lower a [SqlExpr] into a logical [Expr].
-pub(crate) fn type_expr(vars: &Relvars, expr: SqlExpr, expected: Option<&AlgebraicType>) -> TypingResult<Expr> {
+pub(crate) fn type_expr(
+    vars: &Relvars,
+    expr: SqlExpr,
+    expected: Option<&AlgebraicType>,
+    depth: &mut usize,
+) -> TypingResult<Expr> {
+    recursion::guard(depth, recursion::MAX_RECURSION_TYP_EXPR, "expr::type_expr")?;
+
     match (expr, expected) {
         (SqlExpr::Lit(SqlLiteral::Bool(v)), None | Some(AlgebraicType::Bool)) => Ok(Expr::bool(v)),
         (SqlExpr::Lit(SqlLiteral::Bool(_)), Some(ty)) => Err(UnexpectedType::new(&AlgebraicType::Bool, ty).into()),
@@ -117,21 +125,21 @@ pub(crate) fn type_expr(vars: &Relvars, expr: SqlExpr, expected: Option<&Algebra
             }))
         }
         (SqlExpr::Log(a, b, op), None | Some(AlgebraicType::Bool)) => {
-            let a = type_expr(vars, *a, Some(&AlgebraicType::Bool))?;
-            let b = type_expr(vars, *b, Some(&AlgebraicType::Bool))?;
+            let a = type_expr(vars, *a, Some(&AlgebraicType::Bool), depth)?;
+            let b = type_expr(vars, *b, Some(&AlgebraicType::Bool), depth)?;
             Ok(Expr::LogOp(op, Box::new(a), Box::new(b)))
         }
         (SqlExpr::Bin(a, b, op), None | Some(AlgebraicType::Bool)) if matches!(&*a, SqlExpr::Lit(_)) => {
-            let b = type_expr(vars, *b, None)?;
-            let a = type_expr(vars, *a, Some(b.ty()))?;
+            let b = type_expr(vars, *b, None, depth)?;
+            let a = type_expr(vars, *a, Some(b.ty()), depth)?;
             if !op_supports_type(op, a.ty()) {
                 return Err(InvalidOp::new(op, a.ty()).into());
             }
             Ok(Expr::BinOp(op, Box::new(a), Box::new(b)))
         }
         (SqlExpr::Bin(a, b, op), None | Some(AlgebraicType::Bool)) => {
-            let a = type_expr(vars, *a, None)?;
-            let b = type_expr(vars, *b, Some(a.ty()))?;
+            let a = type_expr(vars, *a, None, depth)?;
+            let b = type_expr(vars, *b, Some(a.ty()), depth)?;
             if !op_supports_type(op, a.ty()) {
                 return Err(InvalidOp::new(op, a.ty()).into());
             }
diff --git a/crates/expr/src/statement.rs b/crates/expr/src/statement.rs
@@ -162,7 +162,7 @@ pub fn type_delete(delete: SqlDelete, tx: &impl SchemaView) -> TypingResult<Tabl
     let mut vars = Relvars::default();
     vars.insert(table_name.clone(), from.clone());
     let expr = filter
-        .map(|expr| type_expr(&vars, expr, Some(&AlgebraicType::Bool)))
+        .map(|expr| type_expr(&vars, expr, Some(&AlgebraicType::Bool), &mut 0))
         .transpose()?;
     Ok(TableDelete {
         table: from,
@@ -216,7 +216,7 @@ pub fn type_update(update: SqlUpdate, tx: &impl SchemaView) -> TypingResult<Tabl
     vars.insert(table_name.clone(), schema.clone());
     let values = values.into_boxed_slice();
     let filter = filter
-        .map(|expr| type_expr(&vars, expr, Some(&AlgebraicType::Bool)))
+        .map(|expr| type_expr(&vars, expr, Some(&AlgebraicType::Bool), &mut 0))
         .transpose()?;
     Ok(TableUpdate {
         table: schema,
diff --git a/crates/sql-parser/src/parser/errors.rs b/crates/sql-parser/src/parser/errors.rs
@@ -93,6 +93,13 @@ pub enum SqlRequired {
     JoinAlias,
 }
 
+#[derive(Error, Debug)]
+#[error("Recursion limit exceeded, `{message}` limit: {limit}")]
+pub struct RecursionError {
+    pub(crate) limit: usize,
+    pub(crate) message: String,
+}
+
 #[derive(Error, Debug)]
 pub enum SqlParseError {
     #[error(transparent)]
@@ -103,4 +110,6 @@ pub enum SqlParseError {
     SqlRequired(#[from] SqlRequired),
     #[error(transparent)]
     ParserError(#[from] ParserError),
+    #[error(transparent)]
+    Recursion(#[from] RecursionError),
 }
diff --git a/crates/sql-parser/src/parser/mod.rs b/crates/sql-parser/src/parser/mod.rs
@@ -10,6 +10,7 @@ use crate::ast::{
 };
 
 pub mod errors;
+pub mod recursion;
 pub mod sql;
 pub mod sub;
 
@@ -61,11 +62,14 @@ trait RelParser {
                 Ok(SqlJoin {
                     var,
                     alias,
-                    on: Some(parse_expr(Expr::BinaryOp {
-                        left,
-                        op: BinaryOperator::Eq,
-                        right,
-                    })?),
+                    on: Some(parse_expr(
+                        Expr::BinaryOp {
+                            left,
+                            op: BinaryOperator::Eq,
+                            right,
+                        },
+                        &mut 0,
+                    )?),
                 })
             }
             _ => Err(SqlUnsupported::JoinType.into()),
@@ -204,15 +208,16 @@ pub(crate) fn parse_proj(expr: Expr) -> SqlParseResult<ProjectExpr> {
 }
 
 /// Parse a scalar expression
-pub(crate) fn parse_expr(expr: Expr) -> SqlParseResult<SqlExpr> {
+pub(crate) fn parse_expr(expr: Expr, depth: &mut usize) -> SqlParseResult<SqlExpr> {
     fn signed_num(sign: impl Into<String>, expr: Expr) -> Result<SqlExpr, SqlUnsupported> {
         match expr {
             Expr::Value(Value::Number(n, _)) => Ok(SqlExpr::Lit(SqlLiteral::Num((sign.into() + &n).into_boxed_str()))),
             expr => Err(SqlUnsupported::Expr(expr)),
         }
     }
+    recursion::guard(depth, recursion::MAX_RECURSION_EXPR, "sql-parser::parse_expr")?;
     match expr {
-        Expr::Nested(expr) => parse_expr(*expr),
+        Expr::Nested(expr) => parse_expr(*expr, depth),
         Expr::Value(Value::Placeholder(param)) if &param == ":sender" => Ok(SqlExpr::Param(Parameter::Sender)),
         Expr::Value(v) => Ok(SqlExpr::Lit(parse_literal(v)?)),
         Expr::UnaryOp {
@@ -238,31 +243,31 @@ pub(crate) fn parse_expr(expr: Expr) -> SqlParseResult<SqlExpr> {
             op: BinaryOperator::And,
             right,
         } => {
-            let l = parse_expr(*left)?;
-            let r = parse_expr(*right)?;
+            let l = parse_expr(*left, depth)?;
+            let r = parse_expr(*right, depth)?;
             Ok(SqlExpr::Log(Box::new(l), Box::new(r), LogOp::And))
         }
         Expr::BinaryOp {
             left,
             op: BinaryOperator::Or,
             right,
         } => {
-            let l = parse_expr(*left)?;
-            let r = parse_expr(*right)?;
+            let l = parse_expr(*left, depth)?;
+            let r = parse_expr(*right, depth)?;
             Ok(SqlExpr::Log(Box::new(l), Box::new(r), LogOp::Or))
         }
         Expr::BinaryOp { left, op, right } => {
-            let l = parse_expr(*left)?;
-            let r = parse_expr(*right)?;
+            let l = parse_expr(*left, depth)?;
+            let r = parse_expr(*right, depth)?;
             Ok(SqlExpr::Bin(Box::new(l), Box::new(r), parse_binop(op)?))
         }
         _ => Err(SqlUnsupported::Expr(expr).into()),
     }
 }
 
 /// Parse an optional scalar expression
-pub(crate) fn parse_expr_opt(opt: Option<Expr>) -> SqlParseResult<Option<SqlExpr>> {
-    opt.map(parse_expr).transpose()
+pub(crate) fn parse_expr_opt(opt: Option<Expr>, depth: &mut usize) -> SqlParseResult<Option<SqlExpr>> {
+    opt.map(|expr| parse_expr(expr, depth)).transpose()
 }
 
 /// Parse a scalar binary operator
diff --git a/crates/sql-parser/src/parser/recursion.rs b/crates/sql-parser/src/parser/recursion.rs
@@ -0,0 +1,34 @@
+//! A utility for guarding against excessive recursion depth in the SQL parser.
+//!
+//! Different parts of the parser may have different recursion limits.
+//!
+//! Removing one could allow the others to be higher, but depending on how the `SQL` is structured, it could lead to a `stack overflow`
+//! if is not guarded against, so is incorrect to assume that a limit is sufficient for the next part of the parser.
+use crate::parser::errors::{RecursionError, SqlParseError};
+use std::fmt::Display;
+
+/// A conservative limit for recursion depth on `parse_expr`.
+pub const MAX_RECURSION_EXPR: usize = 700;
+/// A conservative limit for recursion depth on `type_expr`.
+pub const MAX_RECURSION_TYP_EXPR: usize = 5_000;
+
+/// A utility for guarding against excessive recursion depth.
+///
+/// **Usage:**
+/// ```
+/// use spacetimedb_sql_parser::parser::recursion;
+/// let mut depth = 0;
+/// assert!(recursion::guard(&mut depth, 10, "test").is_ok());
+/// ```
+pub fn guard(depth: &mut usize, limit: usize, msg: impl Display) -> Result<(), SqlParseError> {
+    *depth += 1;
+    if *depth > limit {
+        Err(RecursionError {
+            limit,
+            message: msg.to_string(),
+        }
+        .into())
+    } else {
+        Ok(())
+    }
+}
diff --git a/crates/sql-parser/src/parser/sql.rs b/crates/sql-parser/src/parser/sql.rs
@@ -202,7 +202,7 @@ fn parse_statement(stmt: Statement) -> SqlParseResult<SqlAst> {
         } if joins.is_empty() && with_hints.is_empty() && partitions.is_empty() => Ok(SqlAst::Update(SqlUpdate {
             table: parse_ident(name)?,
             assignments: parse_assignments(assignments)?,
-            filter: parse_expr_opt(selection)?,
+            filter: parse_expr_opt(selection, &mut 0)?,
         })),
         Statement::Delete {
             tables,
@@ -297,7 +297,7 @@ fn parse_delete(mut from: Vec<TableWithJoins>, selection: Option<Expr>) -> SqlPa
                 joins,
             } if joins.is_empty() && with_hints.is_empty() && partitions.is_empty() => Ok(SqlDelete {
                 table: parse_ident(name)?,
-                filter: parse_expr_opt(selection)?,
+                filter: parse_expr_opt(selection, &mut 0)?,
             }),
             t => Err(SqlUnsupported::DeleteTable(t).into()),
         }
@@ -395,7 +395,7 @@ fn parse_select(select: Select, limit: Option<Box<str>>) -> SqlParseResult<SqlSe
             Ok(SqlSelect {
                 project: parse_projection(projection)?,
                 from: SqlParser::parse_from(from)?,
-                filter: parse_expr_opt(selection)?,
+                filter: parse_expr_opt(selection, &mut 0)?,
                 limit,
             })
         }
diff --git a/crates/sql-parser/src/parser/sub.rs b/crates/sql-parser/src/parser/sub.rs
@@ -142,7 +142,7 @@ fn parse_select(select: Select) -> SqlParseResult<SqlSelect> {
         {
             Ok(SqlSelect {
                 from: SubParser::parse_from(from)?,
-                filter: parse_expr_opt(selection)?,
+                filter: parse_expr_opt(selection, &mut 0)?,
                 project: parse_projection(projection)?,
             })
         }

Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ fn parse_select(select: Select) -> SqlParseResult<SqlSelect> {`
`142`	`142`	`{`
`143`	`143`	`Ok(SqlSelect {`
`144`	`144`	`from: SubParser::parse_from(from)?,`
`145`		`- filter: parse_expr_opt(selection)?,`
	`145`	`+ filter: parse_expr_opt(selection, &mut 0)?,`
`146`	`146`	`project: parse_projection(projection)?,`
`147`	`147`	`})`
`148`	`148`	`}`