yggdrasil/api/database/
sql_console.rs

1#![allow(clippy::unused_unit, deprecated)]
2
3//! SQL 控制台执行（全读写 + 4 道护栏）。
4//!
5//! 护栏：
6//! 1. 高危语句闸门：`DROP DATABASE`/`DROP SCHEMA`（字符串预检）绝禁；
7//!    `DROP`/`TRUNCATE`/`ALTER` 需 `confirm_dangerous`。
8//! 2. 无 WHERE 拦截：`UPDATE`/`DELETE` 无 `selection` 拒绝。
9//! 3. 查询超时上限：复用 `STATEMENT_TIMEOUT_SECS`（pool 层已注入 GUC）。
10//! 4. 前端二次确认（前端实现）。
11//!
12//! 默认禁止多语句（`allow_multi` 放开）。
13
14use crate::api::error::ApiError;
15use leptos::prelude::*;
16use serde::{Deserialize, Serialize};
17
18// admin 鉴权 + DB 查询仅在 server 构建里被 server function 体引用。
19#[cfg(feature = "server")]
20use crate::api::auth::get_current_admin_user;
21#[cfg(feature = "server")]
22use crate::api::database::read_only::{executed_statements, is_read_only_statement};
23#[cfg(feature = "server")]
24use crate::api::error::AppError;
25#[cfg(feature = "server")]
26use crate::db::pool::get_conn;
27
28#[derive(Deserialize, Serialize, Debug, Clone, Copy)]
29pub struct ExecuteSqlOpts {
30    /// 是否允许多语句（`;` 分隔），默认 false。
31    pub allow_multi: bool,
32    /// 是否勾选「我了解后果」（放开 DROP/TRUNCATE/ALTER 等高危）。
33    pub confirm_dangerous: bool,
34    /// 是否带 EXPLAIN 执行计划。
35    pub with_explain: bool,
36}
37
38#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq)]
39pub struct SqlResult {
40    pub columns: Vec<String>,
41    /// 每格用 JSON 表示（text/int/timestamp/bool/null）。
42    pub rows: Vec<Vec<serde_json::Value>>,
43    pub affected_rows: u64,
44    pub elapsed_ms: u64,
45    /// 语句类型（来自 AST，如 "Select"/"Update"/"CreateTable"）。
46    pub statement_type: String,
47    pub explain: Option<String>,
48    /// SQL 已解析但执行失败时的脱敏/截断错误文本（admin-only 控制台可见）。
49    pub error: Option<String>,
50    /// 是否因 500 行上限截断。
51    pub truncated: bool,
52}
53
54// 以下常量/枚举仅被 server function 体引用（WASM 构建里 server fn 体被 cfg 剥掉，
55// 故这些符号也需 gate，否则非 server 构建会报 dead_code）。
56
57/// 结果行数上限（超出截断 + 提示）。
58#[cfg(feature = "server")]
59const MAX_ROWS: usize = 500;
60
61/// 绝对禁止的语句关键词（字符串预检，sqlparser 无 ObjectType::Database/Schema）。
62/// 命中即拒，不可放行。
63///
64/// 注意：这里存的是「关键字序列」，由 [`is_absolutely_forbidden`] 做 token 级
65/// 匹配——而非原始 `contains()` 子串。这样 `DROP   DATABASE`（多空格）、
66/// `DROP\tDATABASE`、`DROP\nDATABASE` 等绕过单空格子串的写法都能命中。
67/// 关键背景：sqlparser 的 PostgreSqlDialect 无法解析 DROP/CREATE DATABASE，
68/// 这类语句【没有 AST 兜底】，字符串预检是唯一防线，故必须 token 级鲁棒。
69#[cfg(feature = "server")]
70const ABSOLUTELY_FORBIDDEN: &[&[&str]] = &[
71    &["drop", "database"],
72    &["drop", "schema"],
73    &["create", "database"],
74];
75
76/// 护栏 1 的字符串预检（token 序列匹配）。
77///
78/// 把 SQL 按空白拆成 token（小写、去逗号/分号尾缀），再扫描是否出现
79/// `ABSOLUTELY_FORBIDDEN` 里任一连续关键字序列。返回命中的序列描述供报错。
80///
81/// 用 token 序列而非 `contains()` 是为了拦多空格/制表符/换行绕过——
82/// `drop   database` 在子串匹配下漏，在 token 序列下命中。
83#[cfg(feature = "server")]
84fn is_absolutely_forbidden(sql: &str) -> Option<&'static str> {
85    // 规范化：小写 + 按任意空白拆分 + 去掉 token 尾部的 , ; ( )
86    let lowered = sql.to_lowercase();
87    let tokens: Vec<&str> = lowered
88        .split_whitespace()
89        .map(|t| t.trim_end_matches([',', ';', '(', ')']))
90        .collect();
91    for forbidden in ABSOLUTELY_FORBIDDEN {
92        // 在 token 流里滑窗查连续序列
93        for window in tokens.windows(forbidden.len()) {
94            if window == *forbidden {
95                return Some(match forbidden {
96                    ["drop", "database"] => "DROP DATABASE",
97                    ["drop", "schema"] => "DROP SCHEMA",
98                    ["create", "database"] => "CREATE DATABASE",
99                    _ => "未知高危操作",
100                });
101            }
102        }
103    }
104    None
105}
106
107/// 护栏检查返回值。
108#[cfg(feature = "server")]
109#[derive(Debug)]
110enum GuardResult {
111    Allowed,
112    /// 需 confirm_dangerous 才放行。
113    NeedsConfirm,
114    /// 不可放行（附带原因）。
115    Forbidden(String),
116}
117
118/// 护栏 1+2：sqlparser 解析后遍历 AST，检查高危语句与无 WHERE 的 UPDATE/DELETE。
119///
120/// 逐条检查语句本身及其内部会被执行的语句（`EXPLAIN ANALYZE` 内层、WITH 中的
121/// INSERT/UPDATE/DELETE），防止包一层就绕过护栏。
122#[cfg(feature = "server")]
123fn check_guards(asts: &[sqlparser::ast::Statement], confirm_dangerous: bool) -> GuardResult {
124    use sqlparser::ast::{ObjectType, Statement};
125
126    for stmt in asts.iter().flat_map(executed_statements) {
127        match stmt {
128            // 护栏 1（绝禁）：DROP SCHEMA 和 DROP DATABASE 永远禁止。
129            // 这里在 AST 层结构性禁止，防 SQL 注释/空白绕过字符串预检。
130            Statement::Drop {
131                object_type: ObjectType::Schema,
132                ..
133            } => {
134                return GuardResult::Forbidden("禁止 DROP SCHEMA".to_string());
135            }
136            Statement::Drop {
137                object_type: ObjectType::Database,
138                ..
139            } => {
140                return GuardResult::Forbidden("禁止 DROP DATABASE".to_string());
141            }
142            // 护栏 1（绝禁）：CREATE DATABASE 永远禁止。sqlparser 能解析它，
143            // 故在 AST 层补上结构禁止——这是字符串预检之外的第二道防线，
144            // 防多空格/注释绕过 is_absolutely_forbidden 的 token 匹配。
145            Statement::CreateDatabase { .. } => {
146                return GuardResult::Forbidden("禁止 CREATE DATABASE".to_string());
147            }
148            // 护栏 1：需确认的高危语句（DROP TABLE/VIEW/INDEX 等、TRUNCATE、ALTER）
149            Statement::Drop { .. } | Statement::Truncate { .. } | Statement::AlterTable { .. } => {
150                if !confirm_dangerous {
151                    return GuardResult::NeedsConfirm;
152                }
153            }
154            // 护栏 2：UPDATE 无 WHERE
155            Statement::Update(sqlparser::ast::Update {
156                selection: None, ..
157            }) => {
158                return GuardResult::Forbidden(
159                    "UPDATE 缺少 WHERE 子句，将影响全表。请加 WHERE 条件。".to_string(),
160                );
161            }
162            // 护栏 2：DELETE 无 WHERE
163            Statement::Delete(sqlparser::ast::Delete {
164                selection: None, ..
165            }) => {
166                return GuardResult::Forbidden(
167                    "DELETE 缺少 WHERE 子句，将影响全表。请加 WHERE 条件。".to_string(),
168                );
169            }
170            _ => {}
171        }
172    }
173    GuardResult::Allowed
174}
175
176/// 提取语句类型名（AST 变体名，如 "Select"/"Insert"/"Update"）。
177#[cfg(feature = "server")]
178fn statement_type_name(stmt: &sqlparser::ast::Statement) -> String {
179    use sqlparser::ast::Statement;
180    let name = match stmt {
181        Statement::Query(_) => "Select",
182        Statement::Insert(_) => "Insert",
183        Statement::Update(_) => "Update",
184        Statement::Delete(_) => "Delete",
185        Statement::CreateTable { .. } => "CreateTable",
186        Statement::AlterTable { .. } => "AlterTable",
187        Statement::Drop { .. } => "Drop",
188        Statement::Truncate { .. } => "Truncate",
189        Statement::Explain { .. } => "Explain",
190        _ => "Other",
191    };
192    name.to_string()
193}
194
195/// 判断语句是否只读（无副作用的查询 / 非 ANALYZE 的 EXPLAIN / SHOW 系列）。
196///
197/// 查询与 EXPLAIN 的判定与数据导出共用 [`is_read_only_statement`]：
198/// `EXPLAIN ANALYZE <DML>`、data-modifying CTE、`SELECT INTO`、行锁查询都算写。
199#[cfg(feature = "server")]
200fn is_read_only(stmt: &sqlparser::ast::Statement) -> bool {
201    use sqlparser::ast::Statement;
202    is_read_only_statement(stmt)
203        || matches!(
204            stmt,
205            // SHOW 系列只读取元数据/配置参数，不会写入业务表。
206            Statement::ShowVariable { .. }
207                | Statement::ShowVariables { .. }
208                | Statement::ShowStatus { .. }
209                | Statement::ShowCreate { .. }
210                | Statement::ShowColumns { .. }
211                | Statement::ShowCatalogs { .. }
212                | Statement::ShowDatabases { .. }
213                | Statement::ShowProcessList { .. }
214                | Statement::ShowSchemas { .. }
215                | Statement::ShowCharset(_)
216                | Statement::ShowObjects(_)
217                | Statement::ShowTables { .. }
218                | Statement::ShowViews { .. }
219                | Statement::ShowFunctions { .. }
220                | Statement::ShowCollation { .. }
221        )
222}
223
224/// 判断一组语句中是否含有写操作（非只读语句）。
225///
226/// SQL 控制台执行写操作（INSERT/UPDATE/DELETE/TRUNCATE/ALTER/DROP 等）可能直改
227/// posts/comments/tags 等业务表，绕过 server function 的正常缓存失效路径。
228/// 抽成纯函数便于单测「哪些语句集合应触发兜底失效」。
229#[cfg(feature = "server")]
230fn writes_affect_cache(stmts: &[sqlparser::ast::Statement]) -> bool {
231    stmts.iter().any(|s| !is_read_only(s))
232}
233
234/// 把一列的值转成 JSON（按 PG 类型名分发）。
235#[cfg(feature = "server")]
236fn col_to_json(row: &tokio_postgres::Row, idx: usize) -> serde_json::Value {
237    use serde_json::json;
238    let ty = row
239        .columns()
240        .get(idx)
241        .map(|c| c.type_().name())
242        .unwrap_or("");
243    match ty {
244        "int2" => row
245            .try_get::<_, Option<i16>>(idx)
246            .ok()
247            .flatten()
248            .map(|v| json!(v))
249            .unwrap_or(serde_json::Value::Null),
250        "int4" => row
251            .try_get::<_, Option<i32>>(idx)
252            .ok()
253            .flatten()
254            .map(|v| json!(v))
255            .unwrap_or(serde_json::Value::Null),
256        "int8" => row
257            .try_get::<_, Option<i64>>(idx)
258            .ok()
259            .flatten()
260            .map(|v| json!(v))
261            .unwrap_or(serde_json::Value::Null),
262        "float4" => row
263            .try_get::<_, Option<f32>>(idx)
264            .ok()
265            .flatten()
266            .map(|v| json!(v))
267            .unwrap_or(serde_json::Value::Null),
268        "float8" => row
269            .try_get::<_, Option<f64>>(idx)
270            .ok()
271            .flatten()
272            .map(|v| json!(v))
273            .unwrap_or(serde_json::Value::Null),
274        "bool" => row
275            .try_get::<_, Option<bool>>(idx)
276            .ok()
277            .flatten()
278            .map(|v| json!(v))
279            .unwrap_or(serde_json::Value::Null),
280        // 其余（text/varchar/timestamp/jsonb/...）一律按字符串取，失败则 null
281        _ => row
282            .try_get::<_, Option<String>>(idx)
283            .ok()
284            .flatten()
285            .map(|v| json!(v))
286            .unwrap_or(serde_json::Value::Null),
287    }
288}
289
290/// 执行 SQL（全读写，管理员）。护栏见模块文档。
291#[server(ExecuteSql, endpoint = "execute_sql", input = leptos::server_fn::codec::Json)]
292pub async fn execute_sql(sql: String, opts: ExecuteSqlOpts) -> Result<SqlResult, ApiError> {
293    let _user = get_current_admin_user().await?;
294
295    #[cfg(feature = "server")]
296    {
297        use sqlparser::dialect::PostgreSqlDialect;
298        use sqlparser::parser::Parser;
299
300        // 护栏 1（绝禁）：字符串预检 DROP/CREATE DATABASE、DROP SCHEMA
301        // （token 序列匹配，防多空格绕过；这类语句无 AST 兜底）
302        if let Some(name) = is_absolutely_forbidden(&sql) {
303            return Err(AppError::BadRequest(format!("禁止的操作：{}", name)).into());
304        }
305
306        // 解析 SQL
307        let dialect = PostgreSqlDialect {};
308        let asts = Parser::parse_sql(&dialect, &sql)
309            .map_err(|e| AppError::BadRequest(format!("SQL 解析失败：{e}")))?;
310        if asts.is_empty() {
311            return Err(AppError::BadRequest("空的 SQL 语句".into()).into());
312        }
313
314        // 多语句检查（默认禁止）
315        if asts.len() > 1 && !opts.allow_multi {
316            return Err(AppError::BadRequest(
317                "检测到多条语句，请勾选「允许多语句」后再执行".into(),
318            )
319            .into());
320        }
321
322        // 护栏 1+2：AST 检查
323        match check_guards(&asts, opts.confirm_dangerous) {
324            GuardResult::Forbidden(msg) => {
325                return Err(AppError::BadRequest(msg).into());
326            }
327            GuardResult::NeedsConfirm => {
328                return Err(AppError::BadRequest(
329                    "高危操作（DROP/TRUNCATE/ALTER），需勾选「我了解后果」".into(),
330                )
331                .into());
332            }
333            GuardResult::Allowed => {}
334        }
335
336        let client = get_conn().await.map_err(AppError::db_conn)?;
337        let start = std::time::Instant::now;
338
339        // 逐条执行：每条用其 AST 重序列化的形式（stmt.to_string()），
340        // 而非原始整段 SQL——保证执行的语句与护栏检查的 AST 完全一致，
341        // 避免 allow_multi 时整段 SQL 被重复执行，也杜绝读/写分类与实际执行解耦。
342        let mut last_result = SqlResult::default();
343        for stmt in &asts {
344            // 重序列化单条语句为可执行 SQL（去掉末尾分号，避免与 execute 的隐式分号冲突）
345            let stmt_sql = stmt.to_string();
346            last_result = execute_one(&client, stmt, &stmt_sql, opts.with_explain, start).await?;
347            if last_result.error.is_some() {
348                break;
349            }
350        }
351        // 是否含写语句：抽成纯函数便于单测（见 writes_affect_cache）。
352        let has_write = writes_affect_cache(&asts);
353        // SQL 控制台写操作可能直改 posts/comments/tags，绕过了 server function 的
354        // 正常失效路径（delete_post/create_post 等已在内部失效）。此处全量兜底失效，
355        // 避免最长 10 分钟（TTL_SINGLE_POST=600s）内继续吐出陈旧数据（含已删除文章）。
356        if has_write {
357            crate::cache::invalidate_all_post_caches();
358            crate::cache::invalidate_search_results();
359            crate::cache::invalidate_all_comments();
360            crate::ssr_cache::invalidate_ssr_all_public();
361            crate::ssr_cache::bump_global_generation();
362        }
363        Ok(last_result)
364    }
365    #[cfg(not(feature = "server"))]
366    {
367        let _ = (sql, opts);
368        Ok(SqlResult::default())
369    }
370}
371
372/// 执行单条语句，返回结果。
373///
374/// `stmt_sql` 必须是 `stmt` 重序列化后的**单条**语句 SQL（不含其他语句），
375/// 保证护栏检查的 AST 与实际执行的语句一致。
376/// 把 admin-only SQL 执行错误限制在可控长度，避免异常文本撑大 server function 响应。
377#[cfg(feature = "server")]
378fn format_sql_error(error: impl std::fmt::Display) -> String {
379    error.to_string().chars().take(2_000).collect()
380}
381
382#[cfg(feature = "server")]
383async fn execute_one(
384    client: &deadpool_postgres::Object,
385    stmt: &sqlparser::ast::Statement,
386    stmt_sql: &str,
387    with_explain: bool,
388    start: impl Fn() -> std::time::Instant + Copy,
389) -> Result<SqlResult, ApiError> {
390    let statement_type = statement_type_name(stmt);
391    let read_only = is_read_only(stmt);
392    // 只有查询（SELECT/VALUES/WITH…）是表表达式，能加 EXPLAIN 前缀或包成子查询；
393    // SHOW / EXPLAIN 只能原样执行，`EXPLAIN EXPLAIN …`、`FROM (SHOW …)` 都是语法错误。
394    let is_query = matches!(stmt, sqlparser::ast::Statement::Query(_));
395    let stmt_sql = stmt_sql.trim_end_matches(';');
396
397    if with_explain && read_only && is_query {
398        // EXPLAIN 模式：包裹单条语句取执行计划，取首列文本拼接
399        let explain_sql = format!("EXPLAIN {stmt_sql}");
400        let rows = match client.query(&explain_sql, &[]).await {
401            Ok(rows) => rows,
402            Err(error) => {
403                return Ok(SqlResult {
404                    statement_type,
405                    error: Some(format_sql_error(error)),
406                    elapsed_ms: start().elapsed().as_millis() as u64,
407                    ..Default::default()
408                });
409            }
410        };
411        let explain = rows
412            .iter()
413            .filter_map(|r| r.try_get::<_, String>(0).ok())
414            .collect::<Vec<_>>()
415            .join("\n");
416        return Ok(SqlResult {
417            statement_type,
418            explain: Some(explain),
419            elapsed_ms: start().elapsed().as_millis() as u64,
420            ..Default::default()
421        });
422    }
423
424    if read_only {
425        // 只读：取结果集。列名从第一行取（空结果集时无列名，前端容错）。
426        // 查询包成子查询并追加 LIMIT MAX_ROWS+1，让 DB 只回传这么多行——避免
427        // `SELECT * FROM big_table` 把整张表物化进内存（仅 statement_timeout 兜底，M11）。
428        // +1 行供客户端检测截断：实际行数 > MAX_ROWS 时 truncated=true。
429        // 内层 SQL 已带 LIMIT 时外层 LIMIT 取更小值，结果正确；列名随子查询透传。
430        // SHOW / EXPLAIN 输出由服务端生成、行数有界，原样执行后同样按 MAX_ROWS 截断。
431        let limited_sql = if is_query {
432            format!("SELECT * FROM ({stmt_sql}) AS q LIMIT {}", MAX_ROWS + 1)
433        } else {
434            stmt_sql.to_string()
435        };
436        let rows = match client.query(&limited_sql, &[]).await {
437            Ok(rows) => rows,
438            Err(error) => {
439                return Ok(SqlResult {
440                    statement_type,
441                    error: Some(format_sql_error(error)),
442                    elapsed_ms: start().elapsed().as_millis() as u64,
443                    ..Default::default()
444                });
445            }
446        };
447        let columns: Vec<String> = rows
448            .first()
449            .map(|r| r.columns().iter().map(|c| c.name().to_string()).collect())
450            .unwrap_or_default();
451        let mut data: Vec<Vec<serde_json::Value>> = Vec::new();
452        let mut truncated = false;
453        for r in &rows {
454            if data.len() >= MAX_ROWS {
455                truncated = true;
456                break;
457            }
458            let row: Vec<serde_json::Value> = (0..r.len()).map(|i| col_to_json(r, i)).collect();
459            data.push(row);
460        }
461        Ok(SqlResult {
462            columns,
463            rows: data,
464            truncated,
465            statement_type,
466            elapsed_ms: start().elapsed().as_millis() as u64,
467            ..Default::default()
468        })
469    } else {
470        // 写操作：返回影响行数
471        let affected = match client.execute(stmt_sql, &[]).await {
472            Ok(affected) => affected,
473            Err(error) => {
474                return Ok(SqlResult {
475                    statement_type,
476                    error: Some(format_sql_error(error)),
477                    elapsed_ms: start().elapsed().as_millis() as u64,
478                    ..Default::default()
479                });
480            }
481        };
482        Ok(SqlResult {
483            affected_rows: affected,
484            statement_type,
485            elapsed_ms: start().elapsed().as_millis() as u64,
486            ..Default::default()
487        })
488    }
489}
490
491#[cfg(all(test, feature = "server"))]
492mod tests {
493    use super::*;
494
495    /// 用 PostgreSql 方言解析 SQL 为 AST 列表,供测试复用。
496    fn parse(sql: &str) -> Vec<sqlparser::ast::Statement> {
497        use sqlparser::dialect::PostgreSqlDialect;
498        use sqlparser::parser::Parser;
499        Parser::parse_sql(&PostgreSqlDialect {}, sql).unwrap_or_default()
500    }
501
502    // ── 护栏 1:绝对禁止(token 序列字符串预检层) ──────────────────
503    // is_absolutely_forbidden 做 token 级匹配,防多空格/制表符/换行绕过。
504    // 关键背景:sqlparser 无法解析 DROP/CREATE DATABASE,这一层是唯一防线。
505
506    #[test]
507    fn absolutely_forbidden_targets_database_and_schema() {
508        // 锁定关键字序列集合:任何改动都应是有意识的。
509        assert_eq!(
510            ABSOLUTELY_FORBIDDEN,
511            &[
512                &["drop", "database"],
513                &["drop", "schema"],
514                &["create", "database"],
515            ]
516        );
517    }
518
519    #[test]
520    fn precheck_catches_forbidden_regardless_of_case() {
521        for sql in [
522            "DROP DATABASE yggdrasil",
523            "drop schema public",
524            "CREATE DATABASE evil",
525            "Drop Database x",
526        ] {
527            assert!(is_absolutely_forbidden(sql).is_some(), "应拦截: {sql:?}");
528        }
529    }
530
531    #[test]
532    fn precheck_catches_multi_space_bypass() {
533        // token 序列匹配必须命中多空格/制表符/换行绕过——
534        // 这是改用 is_absolutely_forbidden(替代旧 contains)的核心动机:
535        // sqlparser 无法解析 DROP/CREATE DATABASE,这一层是唯一防线。
536        for sql in [
537            "DROP   DATABASE x",
538            "drop\tdatabase\tx",
539            "DROP\nDATABASE\nx",
540            "DROP\t\tDATABASE x;",
541        ] {
542            assert!(
543                is_absolutely_forbidden(sql).is_some(),
544                "多空格绕过应被拦截: {sql:?}"
545            );
546        }
547    }
548
549    #[test]
550    fn precheck_returns_canonical_name_for_error_message() {
551        assert_eq!(
552            is_absolutely_forbidden("DROP DATABASE x"),
553            Some("DROP DATABASE")
554        );
555        assert_eq!(
556            is_absolutely_forbidden("drop schema public"),
557            Some("DROP SCHEMA")
558        );
559        assert_eq!(
560            is_absolutely_forbidden("CREATE DATABASE evil"),
561            Some("CREATE DATABASE")
562        );
563    }
564
565    #[test]
566    fn create_database_is_blocked_at_both_string_and_ast_layers() {
567        // CREATE DATABASE 能被 sqlparser 解析(→ AST),故有双重防线:
568        // 1. is_absolutely_forbidden(字符串 token 预检)
569        // 2. check_guards 的 Statement::CreateDatabase 分支(AST 层绝禁)
570        // 本测试锁定两道防线都生效,且确认 CREATE DATABASE 确实可被解析
571        // (若未来它变得不可解析,字符串预检就成唯一防线,需另作加固)。
572        assert!(is_absolutely_forbidden("CREATE DATABASE x").is_some());
573        let asts = parse("CREATE DATABASE x");
574        assert!(!asts.is_empty(), "CREATE DATABASE 应可被 sqlparser 解析");
575        assert!(matches!(
576            check_guards(&asts, true),
577            GuardResult::Forbidden(_)
578        ));
579    }
580
581    #[test]
582    fn drop_database_is_guarded_by_both_precheck_and_ast_check() {
583        let asts = parse("DROP DATABASE x");
584        assert!(!asts.is_empty(), "DROP DATABASE 应可被 sqlparser 解析");
585        assert!(matches!(
586            check_guards(&asts, true),
587            GuardResult::Forbidden(_)
588        ));
589        assert!(is_absolutely_forbidden("DROP DATABASE x").is_some());
590    }
591
592    #[test]
593    fn precheck_ignores_benign_statements() {
594        for sql in [
595            "SELECT * FROM users",
596            "DROP TABLE old_logs",
597            "CREATE TABLE t (id int)",
598            "DELETE FROM t WHERE id = 1",
599            "CREATE INDEX idx ON t (col)",
600        ] {
601            assert!(is_absolutely_forbidden(sql).is_none(), "不应误拦: {sql:?}");
602        }
603    }
604
605    // ── 护栏 1:AST 层 DROP SCHEMA 绝禁(防注释/空白绕过字符串预检) ──
606
607    #[test]
608    fn guard_forbids_drop_schema_even_though_string_precheck_is_bypassable() {
609        // 即使字符串预检被某种方式绕过,AST 层仍绝禁 DROP SCHEMA。
610        let asts = parse("DROP SCHEMA public");
611        match check_guards(&asts, true) {
612            GuardResult::Forbidden(msg) => {
613                assert!(msg.contains("SCHEMA"), "DROP SCHEMA 应被禁止, 得到: {msg}")
614            }
615            other => panic!("DROP SCHEMA 应 Forbidden, 得到 {other:?}"),
616        }
617    }
618
619    #[test]
620    fn guard_drop_schema_ignores_confirm_flag() {
621        // 即便勾选了 confirm_dangerous,DROP SCHEMA 仍绝禁。
622        let asts = parse("DROP SCHEMA public");
623        assert!(matches!(
624            check_guards(&asts, true),
625            GuardResult::Forbidden(_)
626        ));
627    }
628
629    // ── 护栏 1:DROP/TRUNCATE/ALTER 需确认 ─────────────────────────
630
631    #[test]
632    fn guard_drop_table_needs_confirm() {
633        let asts = parse("DROP TABLE old_logs");
634        assert!(matches!(
635            check_guards(&asts, false),
636            GuardResult::NeedsConfirm
637        ));
638    }
639
640    #[test]
641    fn guard_drop_table_allowed_with_confirm() {
642        let asts = parse("DROP TABLE old_logs");
643        assert!(matches!(check_guards(&asts, true), GuardResult::Allowed));
644    }
645
646    #[test]
647    fn guard_truncate_needs_confirm() {
648        let asts = parse("TRUNCATE TABLE sessions");
649        assert!(matches!(
650            check_guards(&asts, false),
651            GuardResult::NeedsConfirm
652        ));
653    }
654
655    #[test]
656    fn guard_alter_table_needs_confirm() {
657        let asts = parse("ALTER TABLE posts ADD COLUMN foo text");
658        assert!(matches!(
659            check_guards(&asts, false),
660            GuardResult::NeedsConfirm
661        ));
662    }
663
664    #[test]
665    fn guard_alter_table_allowed_with_confirm() {
666        let asts = parse("ALTER TABLE posts ADD COLUMN foo text");
667        assert!(matches!(check_guards(&asts, true), GuardResult::Allowed));
668    }
669
670    // ── 护栏 2:UPDATE/DELETE 无 WHERE 绝禁 ────────────────────────
671
672    #[test]
673    fn guard_update_without_where_is_forbidden() {
674        let asts = parse("UPDATE posts SET title = 'x'");
675        match check_guards(&asts, true) {
676            GuardResult::Forbidden(msg) => assert!(msg.contains("WHERE")),
677            other => panic!("无 WHERE 的 UPDATE 应 Forbidden, 得到 {other:?}"),
678        }
679    }
680
681    #[test]
682    fn guard_delete_without_where_is_forbidden() {
683        let asts = parse("DELETE FROM posts");
684        match check_guards(&asts, true) {
685            GuardResult::Forbidden(msg) => assert!(msg.contains("WHERE")),
686            other => panic!("无 WHERE 的 DELETE 应 Forbidden, 得到 {other:?}"),
687        }
688    }
689
690    #[test]
691    fn guard_update_with_where_allowed() {
692        let asts = parse("UPDATE posts SET title = 'x' WHERE id = 1");
693        assert!(matches!(check_guards(&asts, false), GuardResult::Allowed));
694    }
695
696    #[test]
697    fn guard_delete_with_where_allowed() {
698        let asts = parse("DELETE FROM posts WHERE id = 1");
699        assert!(matches!(check_guards(&asts, false), GuardResult::Allowed));
700    }
701
702    #[test]
703    fn guard_update_with_where_not_rescued_by_confirm() {
704        // confirm_dangerous 不应放行无 WHERE 的 UPDATE/DELETE——这是不可协商的安全护栏。
705        let asts = parse("UPDATE posts SET title = 'x'");
706        assert!(matches!(
707            check_guards(&asts, true),
708            GuardResult::Forbidden(_)
709        ));
710    }
711
712    // ── 正常语句 ───────────────────────────────────────────────────
713
714    #[test]
715    fn guard_allows_select_insert_create() {
716        for sql in [
717            "SELECT * FROM posts",
718            "INSERT INTO posts (title) VALUES ('x')",
719            "CREATE TABLE t (id int)",
720        ] {
721            let asts = parse(sql);
722            assert!(
723                matches!(check_guards(&asts, false), GuardResult::Allowed),
724                "应放行: {sql}"
725            );
726        }
727    }
728
729    // ── 多语句:任一命中即拒 ───────────────────────────────────────
730
731    #[test]
732    fn guard_checks_all_statements_in_batch() {
733        // 第二条是无 WHERE 的 DELETE,即便第一条正常,整体也应被拦。
734        let asts = parse("SELECT 1; DELETE FROM posts");
735        assert!(matches!(
736            check_guards(&asts, false),
737            GuardResult::Forbidden(_)
738        ));
739    }
740
741    #[test]
742    fn guard_first_dangerous_short_circuits() {
743        // 第一条 DROP TABLE 未确认,应在 NeedsConfirm 处停下。
744        let asts = parse("DROP TABLE a; SELECT 1");
745        assert!(matches!(
746            check_guards(&asts, false),
747            GuardResult::NeedsConfirm
748        ));
749    }
750
751    // ── statement_type_name ───────────────────────────────────────
752
753    #[test]
754    fn statement_type_name_maps_variants() {
755        assert_eq!(statement_type_name(&parse("SELECT 1")[0]), "Select");
756        assert_eq!(
757            statement_type_name(&parse("INSERT INTO t (a) VALUES (1)")[0]),
758            "Insert"
759        );
760        assert_eq!(
761            statement_type_name(&parse("UPDATE t SET a = 1 WHERE id = 1")[0]),
762            "Update"
763        );
764        assert_eq!(
765            statement_type_name(&parse("DELETE FROM t WHERE id = 1")[0]),
766            "Delete"
767        );
768        assert_eq!(
769            statement_type_name(&parse("CREATE TABLE t (id int)")[0]),
770            "CreateTable"
771        );
772        assert_eq!(
773            statement_type_name(&parse("ALTER TABLE t ADD COLUMN x int")[0]),
774            "AlterTable"
775        );
776        assert_eq!(statement_type_name(&parse("DROP TABLE t")[0]), "Drop");
777        assert_eq!(statement_type_name(&parse("TRUNCATE t")[0]), "Truncate");
778    }
779
780    // ── is_read_only ──────────────────────────────────────────────
781
782    #[test]
783    fn is_read_only_classifies_correctly() {
784        assert!(is_read_only(&parse("SELECT 1")[0]));
785        assert!(is_read_only(&parse("EXPLAIN SELECT 1")[0]));
786        assert!(!is_read_only(&parse("UPDATE t SET a = 1 WHERE id = 1")[0]));
787        assert!(!is_read_only(&parse("DELETE FROM t WHERE id = 1")[0]));
788        assert!(!is_read_only(&parse("INSERT INTO t (a) VALUES (1)")[0]));
789    }
790
791    #[test]
792    fn is_read_only_keeps_show_statements_read_only() {
793        assert!(is_read_only(&parse("SHOW search_path")[0]));
794    }
795
796    #[test]
797    fn guard_forbids_unbounded_delete_hidden_behind_explain_analyze_or_cte() {
798        // EXPLAIN ANALYZE 与 data-modifying CTE 会真正执行内层 DML，
799        // 无 WHERE 护栏必须同样作用于内层语句。
800        for sql in [
801            "EXPLAIN ANALYZE DELETE FROM posts",
802            "EXPLAIN (ANALYZE) UPDATE posts SET title = 'x'",
803            "WITH d AS (DELETE FROM posts RETURNING *) SELECT * FROM d",
804            "WITH x AS (SELECT 1) DELETE FROM posts",
805        ] {
806            match check_guards(&parse(sql), true) {
807                GuardResult::Forbidden(msg) => assert!(msg.contains("WHERE"), "{sql}: {msg}"),
808                other => panic!("{sql} 应被 WHERE 护栏拦截, got {other:?}"),
809            }
810        }
811    }
812
813    // ── writes_affect_cache（SQL 控制台兜底失效开关） ─────────────
814    // 任何写语句（非只读）都应触发全量缓存失效，兜底绕过 server function 的直改 DB。
815
816    #[test]
817    fn writes_affect_cache_true_for_write_statements() {
818        for sql in [
819            "UPDATE posts SET deleted_at = NOW() WHERE id = 648",
820            "DELETE FROM posts WHERE id = 648",
821            "INSERT INTO posts (title) VALUES ('x')",
822            "TRUNCATE posts",
823            "ALTER TABLE posts ADD COLUMN x int",
824            "DROP TABLE posts",
825            "EXPLAIN ANALYZE DELETE FROM posts WHERE id = 648",
826            "EXPLAIN (ANALYZE) UPDATE posts SET title = 'x' WHERE id = 648",
827            "WITH d AS (DELETE FROM posts WHERE id = 648 RETURNING *) SELECT * FROM d",
828            "SELECT * INTO posts_copy FROM posts",
829        ] {
830            assert!(
831                writes_affect_cache(&parse(sql)),
832                "写语句应触发兜底失效：{sql:?}"
833            );
834        }
835    }
836
837    #[test]
838    fn writes_affect_cache_false_for_read_only_statements() {
839        for sql in [
840            "SELECT 1",
841            "EXPLAIN SELECT * FROM posts",
842            "SELECT id FROM posts WHERE id = 1",
843        ] {
844            assert!(
845                !writes_affect_cache(&parse(sql)),
846                "只读语句不应触发兜底失效：{sql:?}"
847            );
848        }
849    }
850
851    #[test]
852    fn writes_affect_cache_mixed_statements_flagged() {
853        // 多语句中混入任一写语句即应触发（与 allow_multi 语义一致）。
854        let stmts = parse("SELECT 1; UPDATE posts SET deleted_at = NULL WHERE id = 648");
855        assert!(writes_affect_cache(&stmts));
856    }
857    #[test]
858    fn format_sql_error_truncates_long_messages() {
859        let message = format_sql_error("x".repeat(2_500));
860        assert_eq!(message.chars().count(), 2_000);
861    }
862
863    #[test]
864    #[ignore = "requires DATABASE_URL"]
865    fn sql_console_database_runs_show_and_explain_statements() {
866        crate::db::TEST_DATABASE_RUNTIME.block_on(show_and_explain_impl());
867    }
868
869    /// SHOW / EXPLAIN 原样执行出结果；查询在两种模式下仍各自走 EXPLAIN 前缀与 LIMIT 包裹。
870    async fn show_and_explain_impl() {
871        let client = get_conn().await.unwrap();
872        let run = |sql: &'static str, with_explain: bool| {
873            let client = &client;
874            async move {
875                let stmt = parse(sql).remove(0);
876                execute_one(
877                    client,
878                    &stmt,
879                    &stmt.to_string(),
880                    with_explain,
881                    std::time::Instant::now,
882                )
883                .await
884                .unwrap()
885            }
886        };
887
888        for with_explain in [false, true] {
889            let show = run("SHOW search_path", with_explain).await;
890            assert_eq!(show.error, None);
891            assert_eq!(show.columns, ["search_path"]);
892            assert_eq!(show.rows.len(), 1);
893
894            let all = run("SHOW ALL;", with_explain).await;
895            assert_eq!(all.error, None);
896            assert_eq!(all.columns, ["name", "setting", "description"]);
897            assert!(all.rows.len() > 1);
898
899            let plan = run("EXPLAIN SELECT 1", with_explain).await;
900            assert_eq!(plan.error, None);
901            assert_eq!(plan.columns, ["QUERY PLAN"]);
902            assert_eq!(plan.explain, None);
903        }
904
905        let explained = run("SELECT 1 AS one", true).await;
906        assert_eq!(explained.error, None);
907        assert!(explained.explain.unwrap().contains("Result"));
908
909        let selected = run("SELECT generate_series(1, 600) AS n", false).await;
910        assert_eq!(selected.error, None);
911        assert_eq!(selected.columns, ["n"]);
912        assert_eq!(selected.rows.len(), MAX_ROWS);
913        assert!(selected.truncated);
914    }
915}
yggdrasil/api/database/sql_console.rs

yggdrasil/api/database/
sql_console.rs
