mirror of
https://github.com/pgadmin-org/pgadmin4.git
synced 2026-08-19 01:15:04 -05:00
fix(llm): reject multi-statement and non-read-only AI assistant queries
The AI Assistant's execute_sql_query tool runs LLM-generated SQL inside a BEGIN TRANSACTION READ ONLY wrapper. However, the LLM-supplied query was sent to psycopg as-is, so a multi-statement payload beginning with COMMIT, END, ROLLBACK, or ABORT terminated the read-only transaction and ran subsequent statements in autocommit mode. With ordinary write privileges this allowed unauthorised data modification; for a superuser or pg_execute_server_program role this chained to remote code execution on the database host via COPY ... TO PROGRAM. Validate the LLM-supplied query before any connection work happens: * The input must parse to exactly one non-empty/non-comment statement. * The leading real token (after stripping leading whitespace, comments, and punctuation) must be one of SELECT, WITH, EXPLAIN, SHOW, VALUES, TABLE. Everything else -- DML, DDL, CALL, COPY, DO, SET/RESET, the transaction-control verbs, and the rest -- is rejected up front. PostgreSQL's READ ONLY mode continues to enforce the remaining cases (data-modifying CTEs, EXPLAIN ANALYZE on writes, volatile side effects) at runtime, so the validator is the load-bearing check for multi- statement / top-level escapes and READ ONLY is the backstop for the rest. Add a 60-scenario regression suite under web/pgadmin/llm/tests/test_database_tool_security.py covering the original PoC payloads (COMMIT/END/ROLLBACK/ABORT/SET/BEGIN), multi- statement masked by comments, the full allow- and deny-list of leading keywords, dollar-quoted literals containing semicolons, and degenerate inputs (empty, whitespace-only, comment-only, quoted identifier). Reported by: Isaac Chen <isaac9503@gmail.com> Reviewed by: Kundan Sable <kundan.sable@enterprisedb.com>
This commit is contained in:
Dave Page
Ashesh Vashi
@@ -0,0 +1,327 @@
|
||||
##########################################################################
|
||||
#
|
||||
# pgAdmin 4 - PostgreSQL Tools
|
||||
#
|
||||
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
|
||||
# This software is released under the PostgreSQL Licence
|
||||
#
|
||||
##########################################################################
|
||||
|
||||
"""Security regression tests for execute_sql_query input validation.
|
||||
|
||||
These tests pin the multi-statement / leading-keyword guard added in
|
||||
``_validate_readonly_query``. The validator is the load-bearing
|
||||
defense against an attacker escaping the ``BEGIN TRANSACTION READ
|
||||
ONLY`` wrapper by emitting a transaction-control statement (COMMIT,
|
||||
END, ROLLBACK, ABORT) followed by writes -- the original Isaac Chen
|
||||
report from 2026-06-08.
|
||||
|
||||
The tests are pure unit tests (no DB, no Flask client) -- they exercise
|
||||
the validator directly. Anything that reaches the connection layer is
|
||||
already too late.
|
||||
"""
|
||||
|
||||
from pgadmin.utils.route import BaseTestGenerator
|
||||
from pgadmin.llm.tools.database import (
|
||||
DatabaseToolError,
|
||||
_validate_readonly_query,
|
||||
)
|
||||
|
||||
|
||||
class ValidateReadonlyQueryAcceptTestCase(BaseTestGenerator):
|
||||
"""Queries that MUST be accepted by the validator.
|
||||
|
||||
These cover the six allowlisted leading keywords plus a handful of
|
||||
sqlparse edge cases (leading paren, leading block comment, dollar-
|
||||
quoted literal containing ``;``, trailing comment-only "statement")
|
||||
where over-eager rejection would break legitimate LLM usage.
|
||||
"""
|
||||
|
||||
scenarios = [
|
||||
('Plain SELECT', dict(
|
||||
query='SELECT 1',
|
||||
)),
|
||||
('SELECT with semicolon', dict(
|
||||
query='SELECT 1;',
|
||||
)),
|
||||
('Leading paren SELECT', dict(
|
||||
query='(SELECT 1)',
|
||||
)),
|
||||
('Leading block comment then SELECT', dict(
|
||||
query='/* hint */ SELECT 1',
|
||||
)),
|
||||
('SELECT followed by line comment', dict(
|
||||
query='SELECT 1 -- trailing',
|
||||
)),
|
||||
('SELECT followed by comment-only second statement', dict(
|
||||
query='SELECT 1; -- trailing',
|
||||
)),
|
||||
('WITH plain CTE', dict(
|
||||
query='WITH x AS (SELECT 1) SELECT * FROM x',
|
||||
)),
|
||||
('WITH RECURSIVE', dict(
|
||||
query=(
|
||||
'WITH RECURSIVE t(n) AS ('
|
||||
' SELECT 1 UNION ALL SELECT n+1 FROM t WHERE n < 3'
|
||||
') SELECT * FROM t'
|
||||
),
|
||||
)),
|
||||
('EXPLAIN SELECT', dict(
|
||||
query='EXPLAIN SELECT 1',
|
||||
)),
|
||||
('EXPLAIN ANALYZE SELECT', dict(
|
||||
query='EXPLAIN ANALYZE SELECT 1',
|
||||
)),
|
||||
('EXPLAIN with options', dict(
|
||||
query='EXPLAIN (ANALYZE, BUFFERS) SELECT 1',
|
||||
)),
|
||||
('SHOW GUC', dict(
|
||||
query='SHOW search_path',
|
||||
)),
|
||||
('VALUES list', dict(
|
||||
query='VALUES (1), (2), (3)',
|
||||
)),
|
||||
('TABLE statement', dict(
|
||||
query='TABLE pg_catalog.pg_class',
|
||||
)),
|
||||
('Lowercase keyword', dict(
|
||||
query='select 1',
|
||||
)),
|
||||
('Mixed case keyword', dict(
|
||||
query='SeLeCt 1',
|
||||
)),
|
||||
('Dollar-quoted string containing semicolon', dict(
|
||||
# Single SELECT whose literal contains ; -- the validator
|
||||
# must NOT mistake this for a multi-statement payload.
|
||||
query="SELECT $$;$$",
|
||||
)),
|
||||
('Tagged dollar quote containing semicolon', dict(
|
||||
query="SELECT $tag$;$tag$",
|
||||
)),
|
||||
('Standard string with doubled quote', dict(
|
||||
query="SELECT 'a''b'",
|
||||
)),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
# Should not raise.
|
||||
_validate_readonly_query(self.query)
|
||||
|
||||
|
||||
class ValidateReadonlyQueryRejectTestCase(BaseTestGenerator):
|
||||
"""Queries that MUST be rejected by the validator.
|
||||
|
||||
The first block reproduces the Isaac Chen PoC payloads (multi-
|
||||
statement with a leading transaction-control keyword to close the
|
||||
wrapping READ ONLY transaction). The remaining scenarios pin the
|
||||
leading-keyword allowlist against every common write / state-
|
||||
changing top-level statement.
|
||||
"""
|
||||
|
||||
scenarios = [
|
||||
# --- Isaac Chen PoC family: transaction control + multi-stmt ---
|
||||
('PoC: COMMIT then COPY TO PROGRAM', dict(
|
||||
query=(
|
||||
"COMMIT; "
|
||||
"COPY (SELECT 1) TO PROGRAM 'id > /tmp/id.txt 2>&1'; "
|
||||
"SELECT 1"
|
||||
),
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: COMMIT then DELETE', dict(
|
||||
query='COMMIT; DELETE FROM t; SELECT 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: END then DELETE', dict(
|
||||
query='END; DELETE FROM t; SELECT 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: ROLLBACK then INSERT', dict(
|
||||
query="ROLLBACK; INSERT INTO t VALUES (1); SELECT 1",
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: ABORT then UPDATE', dict(
|
||||
query='ABORT; UPDATE t SET c = 1; SELECT 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: SET then DELETE', dict(
|
||||
query=(
|
||||
"SET role superuser; "
|
||||
"DELETE FROM t; "
|
||||
"SELECT 1"
|
||||
),
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: BEGIN then SELECT', dict(
|
||||
query='BEGIN; SELECT 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('PoC: trailing COMMIT after SELECT', dict(
|
||||
query='SELECT 1; COMMIT',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
|
||||
# --- Multi-statement: comments don't smuggle past the count ---
|
||||
('Multi-statement masked by block comment', dict(
|
||||
query='SELECT 1; /* hide */ DROP TABLE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Multi-statement masked by line comment', dict(
|
||||
query='SELECT 1; -- ignore\nDROP TABLE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
|
||||
# --- Disallowed leading keywords (single-statement) ---
|
||||
('Leading UPDATE', dict(
|
||||
query='UPDATE t SET c = 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading DELETE', dict(
|
||||
query='DELETE FROM t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading INSERT', dict(
|
||||
query='INSERT INTO t VALUES (1)',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading MERGE', dict(
|
||||
query=(
|
||||
'MERGE INTO t USING s ON t.id = s.id '
|
||||
'WHEN MATCHED THEN UPDATE SET c = s.c'
|
||||
),
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading CALL', dict(
|
||||
query="CALL bad_proc()",
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading COPY TO PROGRAM', dict(
|
||||
query="COPY (SELECT 1) TO PROGRAM 'id'",
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading DO block', dict(
|
||||
query="DO $$ BEGIN PERFORM 1; END $$",
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading SET', dict(
|
||||
query='SET role superuser',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading RESET', dict(
|
||||
query='RESET role',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading CREATE', dict(
|
||||
query='CREATE TABLE t (c int)',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading DROP', dict(
|
||||
query='DROP TABLE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading ALTER', dict(
|
||||
query='ALTER TABLE t ADD COLUMN c int',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading TRUNCATE', dict(
|
||||
query='TRUNCATE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading LOCK', dict(
|
||||
query='LOCK TABLE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading GRANT', dict(
|
||||
query='GRANT ALL ON t TO public',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading REVOKE', dict(
|
||||
query='REVOKE ALL ON t FROM public',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading NOTIFY', dict(
|
||||
query="NOTIFY chan, 'msg'",
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading LISTEN', dict(
|
||||
query='LISTEN chan',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading PREPARE', dict(
|
||||
query='PREPARE p AS SELECT 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading EXECUTE', dict(
|
||||
query='EXECUTE p',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading REFRESH MATERIALIZED VIEW', dict(
|
||||
query='REFRESH MATERIALIZED VIEW mv',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading VACUUM', dict(
|
||||
query='VACUUM t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading ANALYZE (standalone)', dict(
|
||||
# 'ANALYZE foo' is a maintenance command, not EXPLAIN
|
||||
# ANALYZE -- must be rejected because the validator looks
|
||||
# at the *first* keyword only.
|
||||
query='ANALYZE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading CHECKPOINT', dict(
|
||||
query='CHECKPOINT',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading CLUSTER', dict(
|
||||
query='CLUSTER t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Leading REINDEX', dict(
|
||||
query='REINDEX TABLE t',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
|
||||
# --- Empty / degenerate inputs ---
|
||||
('Empty string', dict(
|
||||
query='',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Whitespace only', dict(
|
||||
query=' \n\t ',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Only semicolons', dict(
|
||||
query=';;;',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Only a block comment', dict(
|
||||
query='/* nothing here */',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
('Quoted identifier "SELECT" is not the keyword', dict(
|
||||
# PostgreSQL: references a column/table named SELECT.
|
||||
# The validator must NOT treat the quoted identifier as
|
||||
# the allowlisted keyword.
|
||||
query='"SELECT" 1',
|
||||
expected_code='INVALID_QUERY',
|
||||
)),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
try:
|
||||
_validate_readonly_query(self.query)
|
||||
except DatabaseToolError as e:
|
||||
self.assertEqual(e.code, self.expected_code)
|
||||
return
|
||||
self.fail(
|
||||
f"Validator accepted query that should have been "
|
||||
f"rejected: {self.query!r}"
|
||||
)
|
||||
@@ -19,6 +19,7 @@ Uses pgAdmin's SQL template infrastructure for version-aware queries.
|
||||
import secrets
|
||||
from typing import Optional
|
||||
|
||||
import sqlparse
|
||||
from flask import render_template
|
||||
|
||||
from pgadmin.utils.driver import get_driver
|
||||
@@ -38,6 +39,23 @@ INDEXES_TEMPLATE_PATH = 'indexes/sql'
|
||||
LLM_APP_NAME_PREFIX = 'pgAdmin 4 - LLM'
|
||||
|
||||
|
||||
# Statement keywords permitted at the start of an LLM-supplied query.
|
||||
# The BEGIN TRANSACTION READ ONLY wrapper around the query is only
|
||||
# effective if the LLM cannot exit that transaction. A multi-statement
|
||||
# payload such as `COMMIT; <write>; SELECT 1` would otherwise terminate
|
||||
# the read-only transaction (via COMMIT/END/ROLLBACK/ABORT) and run
|
||||
# subsequent statements in autocommit mode. The allowlist below, in
|
||||
# combination with the single-statement check, ensures the LLM cannot
|
||||
# emit transaction-control statements, SET/RESET, DML/DDL, CALL, COPY,
|
||||
# or anything else that could escape or weaken the read-only sandbox.
|
||||
# EXPLAIN ANALYZE on a SELECT remains supported; EXPLAIN ANALYZE on a
|
||||
# write statement is blocked by PostgreSQL itself inside the read-only
|
||||
# transaction.
|
||||
_ALLOWED_LEADING_KEYWORDS = frozenset({
|
||||
'SELECT', 'WITH', 'EXPLAIN', 'SHOW', 'VALUES', 'TABLE',
|
||||
})
|
||||
|
||||
|
||||
class DatabaseToolError(Exception):
|
||||
"""Exception raised when a database tool operation fails."""
|
||||
|
||||
@@ -127,6 +145,96 @@ def _connect_readonly(
|
||||
return False, str(e)
|
||||
|
||||
|
||||
def _first_real_keyword(statement) -> str:
|
||||
"""
|
||||
Return the first non-trivial leaf token of a parsed statement.
|
||||
|
||||
Whitespace, comments, and punctuation are skipped so that a leading
|
||||
open-parenthesis (e.g. ``(SELECT 1) UNION (SELECT 2)``) or a leading
|
||||
comment block does not mask the real leading keyword. The result is
|
||||
upper-cased; an empty string is returned for an empty statement.
|
||||
"""
|
||||
for tok in statement.flatten():
|
||||
if tok.is_whitespace:
|
||||
continue
|
||||
ttype = str(tok.ttype) if tok.ttype is not None else ''
|
||||
if 'Comment' in ttype or 'Punctuation' in ttype:
|
||||
continue
|
||||
return (tok.normalized or '').upper()
|
||||
return ''
|
||||
|
||||
|
||||
def _validate_readonly_query(query: str) -> None:
|
||||
"""
|
||||
Ensure an LLM-supplied query is a single, read-only statement.
|
||||
|
||||
This is the load-bearing defense against an attacker escaping the
|
||||
read-only transaction wrapper by emitting multiple statements. The
|
||||
PostgreSQL simple-query protocol cheerfully runs a whole
|
||||
semicolon-separated batch in one round-trip, so a payload such as
|
||||
``COMMIT; <write>; SELECT 1`` would terminate the read-only
|
||||
transaction (via the leading ``COMMIT``) and execute the remaining
|
||||
statements in autocommit mode. The trailing ``ROLLBACK`` would then
|
||||
be a no-op.
|
||||
|
||||
Validation rules:
|
||||
|
||||
* The input must contain exactly one non-empty statement.
|
||||
* The leading keyword must be in :data:`_ALLOWED_LEADING_KEYWORDS`.
|
||||
|
||||
PostgreSQL is left to enforce the rest -- ``EXPLAIN ANALYZE`` on a
|
||||
write statement, data-modifying CTEs, and volatile function side
|
||||
effects are all rejected at runtime by the read-only transaction.
|
||||
|
||||
Args:
|
||||
query: SQL query supplied by the LLM tool call.
|
||||
|
||||
Raises:
|
||||
DatabaseToolError: If the query is empty, contains more than
|
||||
one statement, or is not a read-only statement.
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
raise DatabaseToolError(
|
||||
"Query is empty",
|
||||
code="INVALID_QUERY"
|
||||
)
|
||||
|
||||
try:
|
||||
parsed = sqlparse.parse(query)
|
||||
except Exception as e:
|
||||
raise DatabaseToolError(
|
||||
f"Failed to parse query: {e}",
|
||||
code="INVALID_QUERY"
|
||||
)
|
||||
|
||||
statements = [
|
||||
s for s in parsed
|
||||
if s.token_first(skip_cm=True, skip_ws=True) is not None
|
||||
]
|
||||
|
||||
if not statements:
|
||||
raise DatabaseToolError(
|
||||
"Query contains no SQL statement",
|
||||
code="INVALID_QUERY"
|
||||
)
|
||||
|
||||
if len(statements) > 1:
|
||||
raise DatabaseToolError(
|
||||
"Only a single SQL statement is allowed; multi-statement "
|
||||
"queries are rejected",
|
||||
code="INVALID_QUERY"
|
||||
)
|
||||
|
||||
keyword = _first_real_keyword(statements[0])
|
||||
if keyword not in _ALLOWED_LEADING_KEYWORDS:
|
||||
raise DatabaseToolError(
|
||||
f"Statement type '{keyword or 'UNKNOWN'}' is not permitted; "
|
||||
"only read-only statements (SELECT, WITH, EXPLAIN, SHOW, "
|
||||
"VALUES, TABLE) are allowed",
|
||||
code="INVALID_QUERY"
|
||||
)
|
||||
|
||||
|
||||
def _execute_readonly_query(conn, query: str) -> dict:
|
||||
"""
|
||||
Execute a query in a read-only transaction.
|
||||
@@ -231,9 +339,16 @@ def execute_readonly_query(
|
||||
- truncated: True if results were limited
|
||||
|
||||
Raises:
|
||||
DatabaseToolError: If the query fails or connection
|
||||
cannot be established
|
||||
DatabaseToolError: If the query is rejected by validation, the
|
||||
query fails at runtime, or a connection cannot be
|
||||
established.
|
||||
"""
|
||||
# Validate the LLM-supplied query before allocating a connection.
|
||||
# The BEGIN TRANSACTION READ ONLY wrapper below is only effective
|
||||
# if the query is a single read-only statement; see
|
||||
# _validate_readonly_query for the threat model and rules.
|
||||
_validate_readonly_query(query)
|
||||
|
||||
# Generate unique connection ID for this LLM query
|
||||
conn_id = f"llm_{secrets.choice(range(1, 9999999))}"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user