Stricter WASM SQLite query filter to avoid specific sqlean SELECT queries to run on read-only connections.

crates/wasm-runtime-host/src/sqlite.rs

@@ -3,7 +3,7 @@ use http::Uri;
 use http_body_util::{BodyExt, combinators::UnsyncBoxBody};
 use rusqlite::Transaction;
 use self_cell::{MutBorrow, self_cell};
-use sqlite3_parser::ast::Stmt;
+use sqlite3_parser::ast::{Expr, OneSelect, ResultColumn, Select, Stmt};
 use tokio::time::Duration;
 use trailbase_schema::parse::parse_into_statement;
 use trailbase_schema::sqlite::unquote_expr;
@@ -127,12 +127,14 @@ async fn handle_sqlite_query(
       conn.detach(&unquote_expr(&name)).map_err(sqlite_err)?;
       Rows::empty()
     }
-    Ok(Some(Stmt::Select(_))) => match read(conn, request).await {
-      Ok(rows) => rows,
-      Err(err) => {
-        return Ok(SqliteResponse::Error(err));
+    Ok(Some(Stmt::Select(select))) if is_readonly_select(&select) => {
+      match read(conn, request).await {
+        Ok(rows) => rows,
+        Err(err) => {
+          return Ok(SqliteResponse::Error(err));
+        }
       }
-    },
+    }
     _ => match write(conn, request).await {
       Ok(rows) => rows,
       Err(err) => {
@@ -232,6 +234,54 @@ pub fn sqlite_err<E: std::error::Error>(err: E) -> String {
   return err.to_string();
 }
 
+#[allow(clippy::single_match)]
+#[inline]
+fn is_readonly_select(select: &Select) -> bool {
+  fn is_readonly_one_select(select: &OneSelect) -> bool {
+    return match select {
+      OneSelect::Select { columns, .. } => {
+        for column in columns {
+          if let ResultColumn::Expr(Expr::FunctionCall { name, .. }, _) = column {
+            // Filter out SQLean's "define" which is clearly mutating and will
+            // leave connections in an inconsistent state.
+            //
+            // QUESTION: Should we do more, e.g. error and reject the query? It's likely not
+            // enough to just relegate this to the write connection.
+            match name.0.as_bytes() {
+              b"define" | b"undefine" | b"define_free" => {
+                return false;
+              }
+              _ => {}
+            }
+          }
+        }
+
+        return true;
+      }
+      OneSelect::Values(_) => true,
+    };
+  }
+
+  if let Some(ref with) = select.with {
+    for cte in &with.ctes {
+      if !is_readonly_select(&cte.select) {
+        return false;
+      }
+    }
+  }
+
+  let body = &select.body;
+  if let Some(ref compounds) = body.compounds {
+    for compound in compounds {
+      if !is_readonly_one_select(&compound.select) {
+        return false;
+      }
+    }
+  }
+
+  return is_readonly_one_select(&body.select);
+}
+
 // #[inline]
 // fn empty<E>() -> BoxBody<Bytes, E> {
 //   BoxBody::new(http_body_util::Empty::new().map_err(|_| unreachable!()))
@@ -389,4 +439,20 @@ mod tests {
       "Got: {err:?}"
     );
   }
+
+  fn parse_select(s: &str) -> Box<Select> {
+    let stmt = parse_into_statement(s).unwrap().unwrap();
+    if let Stmt::Select(select) = stmt {
+      return select;
+    }
+    panic!("Expected SELECT, got: {stmt:?}");
+  }
+
+  #[test]
+  fn readonly_select_filter_test() {
+    assert!(is_readonly_select(&parse_select("SELECT * FROM test;")));
+
+    let select = parse_select("SELECT define('sumn', ':n * (:n + 1) / 2');");
+    assert!(!is_readonly_select(&select), "{select:?}");
+  }
 }