From fd3d33424f2e7aaf0ae69d5bdceb2a542eefdd31 Mon Sep 17 00:00:00 2001 From: ulleo Date: Thu, 12 Nov 2026 18:30:19 +0800 Subject: [PATCH] fix: retain SQL table extraction when skipping whitelist checks --- backend/README.md | 4 ++-- backend/apps/chat/task/llm.py | 29 +++++++++++++------------ backend/common/core/config.py | 2 +- backend/tests/test_table_check.py | 35 ++++++++++++++++++++++++------- docs/agents/security.md | 4 ++-- 5 files changed, 46 insertions(+), 28 deletions(-) diff --git a/backend/README.md b/backend/README.md index 0b2efd9b..21a2011f 100644 --- a/backend/README.md +++ b/backend/README.md @@ -4,7 +4,7 @@ `SQLBOT_SKIP_TABLE_CHECK` 默认为 `false`。不配置时，生成的 SQL 仍通过真实表名提取和本轮允许表集合检查。 -仅在需要跳过这组检查时设置： +仅在需要跳过表白名单检查时设置： ```dotenv SQLBOT_SKIP_TABLE_CHECK=true @@ -12,4 +12,4 @@ SQLBOT_SKIP_TABLE_CHECK=true 通过进程环境变量传入；本地开发也可写入仓库根目录 `.env`，安装版可写入安装目录的 `conf/sqlbot.conf`。Docker Compose 部署可在 SQLBot 服务的 `environment` 中设置。修改后需重启服务或重建容器，所有问数入口共用该启动配置，不提供配置页面或运行时修改。 -开启后跳过表名提取、无法提取表名时的拒绝及表白名单检查，可能允许查询数据源账号可读取的未勾选表。SQL 只读、危险函数、元数据查询控制、行权限处理及行数限制仍按原有流程执行。 +开启后仅跳过表白名单检查，可能允许查询数据源账号可读取的未勾选表。真实表名提取始终执行，无法提取表名时仍拒绝 SQL。SQL 只读、危险函数、元数据查询控制、行权限处理及行数限制仍按原有流程执行。 diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 245f3dc9..a205b142 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -1578,21 +1578,20 @@ def run_task(self, in_chat: bool = True, stream: bool = True, sql_operate = OperationEnum.GENERATE_SQL sql, tables = self.check_sql(session=_session, res=full_sql_text, operate=sql_operate) - if not settings.SQLBOT_SKIP_TABLE_CHECK: - # 表名安全检查：用 sqlglot 解析真实 SQL，不信任 AI 返回的 tables - actual_tables = extract_tables_from_sql(sql, ds_type=self.ds.type) - if not actual_tables: - raise SingleMessageError( - "SQL parsing failed: unable to extract table names. " - "This may indicate an unsupported SQL syntax or a security issue." - ) - allowed_tables = set(self.table_name_list) - unauthorized_tables = actual_tables - allowed_tables - if unauthorized_tables: - raise SingleMessageError( - f"SQL contains unauthorized tables: {', '.join(unauthorized_tables)}. " - f"Allowed tables: {', '.join(allowed_tables)}" - ) + # 表名安全检查：用 sqlglot 解析真实 SQL，不信任 AI 返回的 tables + actual_tables = extract_tables_from_sql(sql, ds_type=self.ds.type) + if not actual_tables: + raise SingleMessageError( + "SQL parsing failed: unable to extract table names. " + "This may indicate an unsupported SQL syntax or a security issue." + ) + allowed_tables = set(self.table_name_list) + unauthorized_tables = actual_tables - allowed_tables + if unauthorized_tables and not settings.SQLBOT_SKIP_TABLE_CHECK: + raise SingleMessageError( + f"SQL contains unauthorized tables: {', '.join(unauthorized_tables)}. " + f"Allowed tables: {', '.join(allowed_tables)}" + ) if ((not self.current_assistant or is_page_embedded) and is_normal_user( self.current_user)) or use_dynamic_ds: diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 8305c2c8..7eef381b 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -117,7 +117,7 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str: # 默认关闭，防止通过元数据查询泄露数据库结构 SQLBOT_ALLOW_METADATA_QUERIES: bool = False - # 启动配置：显式开启时跳过生成 SQL 的表名提取和白名单检查 + # 启动配置：显式开启时跳过生成 SQL 的表白名单检查，仍须成功提取表名 SQLBOT_SKIP_TABLE_CHECK: bool = False PARSE_REASONING_BLOCK_ENABLED: bool = True diff --git a/backend/tests/test_table_check.py b/backend/tests/test_table_check.py index 1adb7f2a..e4d03c5c 100644 --- a/backend/tests/test_table_check.py +++ b/backend/tests/test_table_check.py @@ -92,8 +92,6 @@ def test_default_allows_known_tables_and_cte_aliases(task, sql): "SELECT id FROM sales UNION ALL SELECT id FROM secret_payroll", "unauthorized tables: secret_payroll", ), - ("SELECT 1", "unable to extract table names"), - ("SELECT (", "unable to extract table names"), ], ) def test_default_rejects_invalid_or_unknown_tables_before_save_and_execute( @@ -107,16 +105,37 @@ def test_default_rejects_invalid_or_unknown_tables_before_save_and_execute( task.execute_sql.assert_not_called() -@pytest.mark.parametrize("sql", ["SELECT * FROM secret_payroll", "SELECT 1"]) -def test_explicit_skip_bypasses_extraction_and_keeps_remaining_flow( +@pytest.mark.parametrize("skip", [False, True]) +@pytest.mark.parametrize("sql", ["SELECT 1", "SELECT ("]) +def test_missing_tables_are_rejected_even_when_whitelist_check_is_skipped( + task, monkeypatch, sql, skip +): + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", skip) + result = run(task, sql) + assert result["success"] is False + assert "unable to extract table names" in result["message"] + task.save_error.assert_called_once() + task.generate_filter.assert_not_called() + task.check_save_sql.assert_not_called() + task.execute_sql.assert_not_called() + + +@pytest.mark.parametrize( + "sql", + [ + "SELECT * FROM secret_payroll", + "SELECT id FROM sales UNION ALL SELECT id FROM secret_payroll", + ], +) +def test_explicit_skip_only_bypasses_whitelist_check_and_keeps_remaining_flow( task, monkeypatch, sql ): monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", True) - extract = Mock(side_effect=AssertionError("table extraction must be skipped")) + extract = Mock(wraps=llm.extract_tables_from_sql) monkeypatch.setattr(llm, "extract_tables_from_sql", extract) result = run(task, sql) assert result["success"] is True - extract.assert_not_called() + extract.assert_called_once_with(sql, ds_type="pg") task.generate_filter.assert_called_once() task.check_save_sql.assert_called_once() task.execute_sql.assert_called_once_with(sql=sql) @@ -127,8 +146,8 @@ def test_explicit_skip_bypasses_extraction_and_keeps_remaining_flow( "sql", [ "DELETE FROM sales", - "SELECT pg_read_file('/etc/passwd')", - "SHOW TABLES", + "SELECT pg_read_file('/etc/passwd') FROM sales", + "DESCRIBE sales", ], ) def test_skip_keeps_database_execution_guards(task, monkeypatch, sql): diff --git a/docs/agents/security.md b/docs/agents/security.md index 81396a3e..2cc0343b 100644 --- a/docs/agents/security.md +++ b/docs/agents/security.md @@ -32,8 +32,8 @@ 1. 只把允许的表和字段元数据提供给模型； 2. 解析模型返回 JSON 并提取 SQL； -3. 用 `sqlglot` 提取真实表名并排除 CTE 别名； -4. 默认对照允许表集合拒绝未授权表；部署显式设置 `SQLBOT_SKIP_TABLE_CHECK=true` 时，跳过第 3、4 步及无法提取表名时的拒绝； +3. 用 `sqlglot` 提取真实表名并排除 CTE 别名，无法提取表名时始终拒绝； +4. 默认对照允许表集合拒绝未授权表；部署显式设置 `SQLBOT_SKIP_TABLE_CHECK=true` 时，仅跳过本步； 5. 按用户、助手和页面嵌入场景应用行权限或动态 SQL 变换； 6. 校验并保存最终 SQL； 7. 执行最终 SQL 并保存结果。