Merge branch 'master' into jdetter/login-and-resume-gameplay

Author: jdetter

crates/cli/src/subcommands/subscribe.rs

@@ -9,10 +9,12 @@ use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
 use spacetimedb_lib::de::serde::{DeserializeWrapper, SeedWrapper};
 use spacetimedb_lib::ser::serde::SerializeWrapper;
+use std::io;
 use std::time::Duration;
+use thiserror::Error;
 use tokio::io::AsyncWriteExt;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
-use tokio_tungstenite::tungstenite::Message as WsMessage;
+use tokio_tungstenite::tungstenite::{Error as WsError, Message as WsMessage};
 
 use crate::api::ClientApi;
 use crate::common_args;
@@ -155,35 +157,88 @@ pub async fn exec(config: Config, args: &ArgMatches) -> Result<(), anyhow::Error
     if let Some(auth_header) = api.con.auth_header.to_header() {
         req.headers_mut().insert(header::AUTHORIZATION, auth_header);
     }
-    let (mut ws, _) = tokio_tungstenite::connect_async(req).await?;
+    let mut ws = tokio_tungstenite::connect_async(req).await.map(|(ws, _)| ws)?;
 
     let task = async {
         subscribe(&mut ws, queries.cloned().map(Into::into).collect()).await?;
         await_initial_update(&mut ws, print_initial_update.then_some(&module_def)).await?;
         consume_transaction_updates(&mut ws, num, &module_def).await
     };
 
-    let needs_shutdown = if let Some(timeout) = timeout {
+    let res = if let Some(timeout) = timeout {
         let timeout = Duration::from_secs(timeout.into());
         match tokio::time::timeout(timeout, task).await {
-            Ok(res) => res?,
-            Err(_elapsed) => true,
+            Ok(res) => res,
+            Err(_elapsed) => {
+                eprintln!("timed out after {}s", timeout.as_secs());
+                Ok(())
+            }
         }
     } else {
-        task.await?
+        task.await
     };
 
-    if needs_shutdown {
-        ws.close(None).await?;
-    }
+    // Close the connection gracefully.
+    // This will return an error if the server already closed,
+    // or the connection is in a bad state.
+    // The error (if any) relevant to the user is already stored in `res`,
+    // so we can ignore errors here -- graceful close is basically a
+    // courtesy to the server.
+    let _ = ws.close(None).await;
+    // The server closing the connection is not considered an error,
+    // but any other error is.
+    res.or_else(|e| {
+        if e.is_server_closed_connection() {
+            Ok(())
+        } else {
+            Err(e)
+        }
+    })
+    .map_err(anyhow::Error::from)
+}
 
-    Ok(())
+#[derive(Debug, Error)]
+enum Error {
+    #[error("error sending subscription queries")]
+    Subscribe {
+        #[source]
+        source: WsError,
+    },
+    #[error("protocol error: {details}")]
+    Protocol { details: &'static str },
+    #[error("websocket error: {source}")]
+    Websocket {
+        #[source]
+        source: WsError,
+    },
+    #[error("encountered failed transaction: {reason}")]
+    TransactionFailure { reason: Box<str> },
+    #[error("error formatting response: {source:#}")]
+    Reformat {
+        #[source]
+        source: anyhow::Error,
+    },
+    #[error(transparent)]
+    Serde(#[from] serde_json::Error),
+    #[error(transparent)]
+    Io(#[from] io::Error),
+}
+
+impl Error {
+    fn is_server_closed_connection(&self) -> bool {
+        matches!(
+            self,
+            Self::Websocket {
+                source: WsError::ConnectionClosed
+            }
+        )
+    }
 }
 
 /// Send the subscribe message.
-async fn subscribe<S>(ws: &mut S, query_strings: Box<[Box<str>]>) -> Result<(), S::Error>
+async fn subscribe<S>(ws: &mut S, query_strings: Box<[Box<str>]>) -> Result<(), Error>
 where
-    S: Sink<WsMessage> + Unpin,
+    S: Sink<WsMessage, Error = WsError> + Unpin,
 {
     let msg = serde_json::to_string(&SerializeWrapper::new(ws::ClientMessage::<()>::Subscribe(
         ws::Subscribe {
@@ -192,35 +247,39 @@ where
         },
     )))
     .unwrap();
-    ws.send(msg.into()).await
+    ws.send(msg.into()).await.map_err(|source| Error::Subscribe { source })
 }
 
 /// Await the initial [`ServerMessage::SubscriptionUpdate`].
 /// If `module_def` is `Some`, print a JSON representation to stdout.
-async fn await_initial_update<S>(ws: &mut S, module_def: Option<&RawModuleDefV9>) -> anyhow::Result<()>
+async fn await_initial_update<S>(ws: &mut S, module_def: Option<&RawModuleDefV9>) -> Result<(), Error>
 where
-    S: TryStream<Ok = WsMessage> + Unpin,
-    S::Error: std::error::Error + Send + Sync + 'static,
+    S: TryStream<Ok = WsMessage, Error = WsError> + Unpin,
 {
     const RECV_TX_UPDATE: &str = "protocol error: received transaction update before initial subscription update";
 
-    while let Some(msg) = ws.try_next().await? {
+    while let Some(msg) = ws.try_next().await.map_err(|source| Error::Websocket { source })? {
         let Some(msg) = parse_msg_json(&msg) else { continue };
         match msg {
             ws::ServerMessage::InitialSubscription(sub) => {
                 if let Some(module_def) = module_def {
-                    let formatted = reformat_update(&sub.database_update, module_def)?;
-                    let output = serde_json::to_string(&formatted)? + "\n";
+                    let output = format_output_json(&sub.database_update, module_def)?;
                     tokio::io::stdout().write_all(output.as_bytes()).await?
                 }
                 break;
             }
-            ws::ServerMessage::TransactionUpdate(ws::TransactionUpdate { status, .. }) => anyhow::bail!(match status {
-                ws::UpdateStatus::Failed(msg) => msg,
-                _ => RECV_TX_UPDATE.into(),
-            }),
+            ws::ServerMessage::TransactionUpdate(ws::TransactionUpdate { status, .. }) => {
+                return Err(match status {
+                    ws::UpdateStatus::Failed(msg) => Error::TransactionFailure { reason: msg },
+                    _ => Error::Protocol {
+                        details: RECV_TX_UPDATE,
+                    },
+                })
+            }
             ws::ServerMessage::TransactionUpdateLight(ws::TransactionUpdateLight { .. }) => {
-                anyhow::bail!(RECV_TX_UPDATE)
+                return Err(Error::Protocol {
+                    details: RECV_TX_UPDATE,
+                })
             }
             _ => continue,
         }
@@ -231,41 +290,47 @@ where
 
 /// Print `num` [`ServerMessage::TransactionUpdate`] messages as JSON.
 /// If `num` is `None`, keep going indefinitely.
-async fn consume_transaction_updates<S>(
-    ws: &mut S,
-    num: Option<u32>,
-    module_def: &RawModuleDefV9,
-) -> anyhow::Result<bool>
+async fn consume_transaction_updates<S>(ws: &mut S, num: Option<u32>, module_def: &RawModuleDefV9) -> Result<(), Error>
 where
-    S: TryStream<Ok = WsMessage> + Unpin,
-    S::Error: std::error::Error + Send + Sync + 'static,
+    S: TryStream<Ok = WsMessage, Error = WsError> + Unpin,
 {
     let mut stdout = tokio::io::stdout();
     let mut num_received = 0;
     loop {
         if num.is_some_and(|n| num_received >= n) {
-            break Ok(true);
+            return Ok(());
         }
-        let Some(msg) = ws.try_next().await? else {
+        let Some(msg) = ws.try_next().await.map_err(|source| Error::Websocket { source })? else {
             eprintln!("disconnected by server");
-            break Ok(false);
+            return Err(Error::Websocket {
+                source: WsError::ConnectionClosed,
+            });
         };
 
         let Some(msg) = parse_msg_json(&msg) else { continue };
         match msg {
             ws::ServerMessage::InitialSubscription(_) => {
-                anyhow::bail!("protocol error: received a second initial subscription update")
+                return Err(Error::Protocol {
+                    details: "received a second initial subscription update",
+                })
             }
             ws::ServerMessage::TransactionUpdateLight(ws::TransactionUpdateLight { update, .. })
             | ws::ServerMessage::TransactionUpdate(ws::TransactionUpdate {
                 status: ws::UpdateStatus::Committed(update),
                 ..
             }) => {
-                let output = serde_json::to_string(&reformat_update(&update, module_def)?)? + "\n";
+                let output = format_output_json(&update, module_def)?;
                 stdout.write_all(output.as_bytes()).await?;
                 num_received += 1;
             }
             _ => continue,
         }
     }
 }
+
+fn format_output_json(msg: &ws::DatabaseUpdate<JsonFormat>, schema: &RawModuleDefV9) -> Result<String, Error> {
+    let formatted = reformat_update(msg, schema).map_err(|source| Error::Reformat { source })?;
+    let output = serde_json::to_string(&formatted)? + "\n";
+
+    Ok(output)
+}

crates/core/src/error.rs

@@ -82,8 +82,6 @@ pub enum DatabaseError {
     DatabasedOpened(PathBuf, anyhow::Error),
 }
 
-// FIXME: reduce type size
-#[expect(clippy::large_enum_variant)]
 #[derive(Error, Debug, EnumAsInner)]
 pub enum DBError {
     #[error("LibError: {0}")]

crates/core/src/host/module_host.rs

@@ -539,8 +539,6 @@ pub enum InitDatabaseError {
     Other(anyhow::Error),
 }
 
-// FIXME: reduce type size
-#[expect(clippy::large_enum_variant)]
 #[derive(thiserror::Error, Debug)]
 pub enum ClientConnectedError {
     #[error(transparent)]

crates/core/src/sql/execute.rs

@@ -1126,6 +1126,11 @@ pub(crate) mod tests {
         Ok(())
     }
 
+    /// Test we are protected against stack overflows when:
+    /// 1. The query is too large (too many characters)
+    /// 2. The AST is too deep
+    ///
+    /// Exercise the limit [`recursion::MAX_RECURSION_EXPR`]
     #[test]
     fn test_large_query_no_panic() -> ResultTest<()> {
         let db = TestDB::durable()?;
@@ -1138,16 +1143,43 @@ 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 1..total {
+                let fragment = format!("x = {x} or ");
+                sql.push_str(&fragment.repeat((total - 1) as usize));
             }
+            sql.push_str("(y = 0)");
+            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())
+        );
+
+        let sql = build_query(41); // This causes stack overflow without the limit
+        assert_eq!(run(&db, ',', &sql), Err("Recursion limit exceeded".to_string()));
+
+        let sql = build_query(40); // The max we can with the current limit
+        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 push 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(())
     }
 

crates/expr/src/errors.rs

@@ -122,8 +122,6 @@ pub struct DuplicateName(pub String);
 #[error("`filter!` does not support column projections; Must return table rows")]
 pub struct FilterReturnType;
 
-// FIXME: reduce type size
-#[expect(clippy::large_enum_variant)]
 #[derive(Error, Debug)]
 pub enum TypingError {
     #[error(transparent)]

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;
@@ -78,8 +79,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> {
+// These types determine the size of each stack frame during type checking.
+// Changing their sizes will require updating the recursion limit to avoid stack overflows.
+const _: () = assert!(size_of::<TypingResult<Expr>>() == 64);
+const _: () = assert!(size_of::<SqlExpr>() == 40);
+
+fn _type_expr(vars: &Relvars, expr: SqlExpr, expected: Option<&AlgebraicType>, depth: 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 +124,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 + 1)?;
+            let b = _type_expr(vars, *b, Some(&AlgebraicType::Bool), depth + 1)?;
             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 + 1)?;
+            let a = _type_expr(vars, *a, Some(b.ty()), depth + 1)?;
             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 + 1)?;
+            let b = _type_expr(vars, *b, Some(a.ty()), depth + 1)?;
             if !op_supports_type(op, a.ty()) {
                 return Err(InvalidOp::new(op, a.ty()).into());
             }
@@ -144,6 +151,11 @@ pub(crate) fn type_expr(vars: &Relvars, expr: SqlExpr, expected: Option<&Algebra
     }
 }
 
+/// Type check and lower a [SqlExpr] into a logical [Expr].
+pub(crate) fn type_expr(vars: &Relvars, expr: SqlExpr, expected: Option<&AlgebraicType>) -> TypingResult<Expr> {
+    _type_expr(vars, expr, expected, 0)
+}
+
 /// Is this type compatible with this binary operator?
 fn op_supports_type(_op: BinOp, t: &AlgebraicType) -> bool {
     t.is_bool()