mirror of
https://github.com/pgadmin-org/pgadmin4.git
synced 2026-09-03 20:52:57 -05:00
fix(llm): close lexer-differential bypass in AI Assistant read-only guard
sqlparse's string-literal lexing can disagree with PostgreSQL's: under
standard_conforming_strings = on (the default), a backslash before a
quote is an ordinary character to PostgreSQL but sqlparse treats it as
escaping the quote, so a payload like
SELECT '\';COMMIT;CREATE TABLE pwn(x int);SELECT 1 --'
passes _validate_readonly_query as a single SELECT while PostgreSQL
executes it as four statements -- the smuggled COMMIT ends the wrapping
BEGIN TRANSACTION READ ONLY and the trailing ROLLBACK is a no-op,
reintroducing the write/RCE bypass the bf47924444 fix was meant to
close (reported by Kai Aizen / SnailSploit).
Run the LLM-supplied query with prepare=True, forcing psycopg3's
extended query protocol. PostgreSQL's own Parse step -- not a
client-side approximation of it -- rejects any text containing more
than one statement, independent of how it's lexed. Threaded through
execute_2darray as an opt-in parameter (default None) so no other
caller's behavior changes. Also set SESSION CHARACTERISTICS AS
TRANSACTION READ ONLY as defense-in-depth against a smuggled COMMIT.
prepare=True alone is not sufficient: psycopg3's PrepareManager.get()
returns Prepare.NO -- checked before it even inspects the prepare
argument -- whenever the connection's prepare_threshold is None, which
is pgAdmin's per-server default ("Prepare threshold" is blank unless an
administrator sets it). On a default-configured server the extended
protocol never actually engaged, so the bypass stayed live. Force
prepare_threshold=0 on the LLM's single-use connection in
_connect_readonly() so the extended protocol -- and PostgreSQL's
single-statement Parse-step guarantee -- is unconditional on this
connection, without touching the server-wide setting or any other
session/caller.
Adds regression coverage at three levels: unit tests pinning that
prepare=True is always passed and that prepare_threshold=0 is forced
on the connection; and an end-to-end test driving the real
/sqleditor/nlq/chat/<trans_id>/stream route (LLM client mocked, real
tool-dispatch path) with the exact smuggled-COMMIT payload, confirming
the guarantee holds from the HTTP entry point down to the driver call.
This commit is contained in:
@@ -10,21 +10,32 @@
|
||||
"""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.
|
||||
``_validate_readonly_query``. It rejects the original Isaac Chen
|
||||
report from 2026-06-08 (a transaction-control statement such as COMMIT,
|
||||
END, ROLLBACK, ABORT followed by writes) whenever sqlparse and
|
||||
PostgreSQL agree on where the statement boundary is.
|
||||
|
||||
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.
|
||||
``_validate_readonly_query`` is a fast pre-filter, not the security
|
||||
boundary -- see ``ValidateReadonlyQueryLexerDifferentialTestCase``
|
||||
below and the module docstring in ``database.py`` for why sqlparse
|
||||
cannot be the load-bearing check, and ``ExecuteReadonlyQueryProtocolTestCase``
|
||||
for the protocol-level enforcement that actually is (Kai Aizen /
|
||||
SnailSploit report, 2026-07-23, on the bf4792444446 fix).
|
||||
|
||||
Most of these 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 for validator, which is exactly
|
||||
why the protocol-level fix exists.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from pgadmin.utils.route import BaseTestGenerator
|
||||
from pgadmin.llm.tools.database import (
|
||||
DatabaseToolError,
|
||||
_validate_readonly_query,
|
||||
_connect_readonly,
|
||||
execute_readonly_query,
|
||||
)
|
||||
|
||||
|
||||
@@ -404,3 +415,170 @@ class ValidateReadonlyQueryRejectTestCase(BaseTestGenerator):
|
||||
f"Validator accepted query that should have been "
|
||||
f"rejected: {self.query!r}"
|
||||
)
|
||||
|
||||
|
||||
class ValidateReadonlyQueryLexerDifferentialTestCase(BaseTestGenerator):
|
||||
"""Pins a KNOWN, ACCEPTED limitation of the sqlparse-based validator.
|
||||
|
||||
sqlparse's string-literal lexing does not always match PostgreSQL's.
|
||||
Under standard_conforming_strings = on (the default), a backslash
|
||||
inside a '...'-quoted literal is an ordinary character to PostgreSQL
|
||||
but sqlparse treats it as escaping the following quote, so it keeps
|
||||
reading past what PostgreSQL considers the end of the string. That
|
||||
lets a payload like the one below smuggle ``;COMMIT;<write>;`` past
|
||||
the "exactly one statement" check disguised as a single string
|
||||
literal, even though PostgreSQL executes it as four statements --
|
||||
the report from Kai Aizen / SnailSploit (2026-07-23) against the
|
||||
bf4792444446 fix.
|
||||
|
||||
This is NOT something ``_validate_readonly_query`` can be made to
|
||||
catch in general (it would require re-implementing PostgreSQL's
|
||||
lexer, including its dependence on server-side GUCs like
|
||||
standard_conforming_strings). The test below asserts the validator
|
||||
accepts the payload -- documenting that this is expected -- and
|
||||
exists only to point at where the real protection lives: see
|
||||
ExecuteReadonlyQueryProtocolTestCase, which pins that the query is
|
||||
always executed with prepare=True. That forces PostgreSQL's own
|
||||
Parse step -- not a client-side approximation of it -- to reject
|
||||
any text containing more than one statement, regardless of how it
|
||||
is lexed.
|
||||
"""
|
||||
|
||||
scenarios = [
|
||||
('Backslash-quote smuggled COMMIT + DDL', dict(
|
||||
query=(
|
||||
"SELECT '\\';COMMIT;CREATE TABLE pwn(x int);"
|
||||
"SELECT 1 --'"
|
||||
),
|
||||
)),
|
||||
('Backslash-quote smuggled COMMIT + COPY TO PROGRAM', dict(
|
||||
query=(
|
||||
"SELECT '\\';COMMIT;COPY (SELECT 1) TO PROGRAM 'id';"
|
||||
"SELECT 1 --'"
|
||||
),
|
||||
)),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
# Documents current, accepted behavior -- must NOT raise.
|
||||
# If this ever starts raising, _validate_readonly_query has
|
||||
# changed in a way that may be worth understanding, but the
|
||||
# query is still safe only because of prepare=True downstream;
|
||||
# don't mistake a change here for the security fix itself.
|
||||
keyword = _validate_readonly_query(self.query)
|
||||
self.assertIsNone(keyword)
|
||||
|
||||
|
||||
class ExecuteReadonlyQueryProtocolTestCase(BaseTestGenerator):
|
||||
"""Pins that execute_readonly_query always runs the LLM's query with
|
||||
prepare=True.
|
||||
|
||||
This is the actual load-bearing defense against statement-boundary
|
||||
smuggling (see ValidateReadonlyQueryLexerDifferentialTestCase):
|
||||
prepare=True forces psycopg3's extended query protocol, whose Parse
|
||||
step is answered by PostgreSQL's own parser and rejects a query
|
||||
string containing more than one SQL statement -- independent of
|
||||
sqlparse, and independent of any particular payload shape. Verified
|
||||
end-to-end against a live PostgreSQL 16 instance during the
|
||||
SnailSploit report triage (2026-07-23): the smuggled-COMMIT payload
|
||||
that _validate_readonly_query accepts is rejected by PostgreSQL's
|
||||
Parse step with "cannot insert multiple commands into a prepared
|
||||
statement" once prepare=True is set, both before and after the
|
||||
max_rows LIMIT-wrapping applied to SELECT queries.
|
||||
|
||||
This test mocks the connection layer (no DB) and only pins the
|
||||
wiring: that prepare=True is passed on every call, not just for
|
||||
inputs that look suspicious. A future refactor that drops the
|
||||
keyword argument, or only sets it conditionally, must fail this
|
||||
test.
|
||||
"""
|
||||
|
||||
scenarios = [
|
||||
('Benign SELECT', dict(
|
||||
query='SELECT 1',
|
||||
)),
|
||||
('Backslash-quote smuggled COMMIT + DDL', dict(
|
||||
query=(
|
||||
"SELECT '\\';COMMIT;CREATE TABLE pwn(x int);"
|
||||
"SELECT 1 --'"
|
||||
),
|
||||
)),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
mock_manager = MagicMock()
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.execute_void.return_value = (True, None)
|
||||
mock_conn.execute_2darray.return_value = (
|
||||
True, {'columns': [], 'rows': []}
|
||||
)
|
||||
|
||||
with patch(
|
||||
'pgadmin.llm.tools.database._get_connection',
|
||||
return_value=(mock_manager, mock_conn)
|
||||
), patch(
|
||||
'pgadmin.llm.tools.database._connect_readonly',
|
||||
return_value=(True, None)
|
||||
):
|
||||
execute_readonly_query(sid=1, did=1, query=self.query)
|
||||
|
||||
mock_conn.execute_2darray.assert_called_once()
|
||||
_args, kwargs = mock_conn.execute_2darray.call_args
|
||||
self.assertTrue(
|
||||
kwargs.get('prepare') is True,
|
||||
"execute_readonly_query must run the LLM-supplied query "
|
||||
"with prepare=True so PostgreSQL's own Parse step -- not "
|
||||
"the sqlparse pre-filter -- enforces the single-statement "
|
||||
"guarantee."
|
||||
)
|
||||
|
||||
|
||||
class ConnectReadonlyForcesPrepareThresholdTestCase(BaseTestGenerator):
|
||||
"""Pins that _connect_readonly forces conn.conn.prepare_threshold = 0.
|
||||
|
||||
prepare=True on the execute_2darray call is necessary but NOT
|
||||
sufficient: psycopg3's PrepareManager.get() returns Prepare.NO --
|
||||
silently falling back to the multi-statement-capable simple query
|
||||
protocol -- whenever the connection's prepare_threshold is None,
|
||||
*before* it even inspects the prepare argument. pgAdmin's per-server
|
||||
"Prepare threshold" field defaults to blank/None, so on a default
|
||||
server the prepare=True guarantee never engages and the
|
||||
CVE-2026-12045 smuggled-COMMIT bypass stays live (Kai Aizen /
|
||||
SnailSploit follow-up, verified live against PostgreSQL on
|
||||
2026-07-24).
|
||||
|
||||
_connect_readonly() therefore forces prepare_threshold = 0 on the
|
||||
LLM's single-use connection. This test pins that override so a
|
||||
future refactor that drops it -- reopening the bypass -- fails here,
|
||||
without needing a live server.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
mock_manager = MagicMock()
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.connected.return_value = True
|
||||
mock_conn.execute_void.return_value = (True, None)
|
||||
# Emulate pgAdmin's default: extended protocol disabled.
|
||||
mock_conn.conn.prepare_threshold = None
|
||||
|
||||
status, msg = _connect_readonly(
|
||||
mock_manager, mock_conn, 'llm_test_conn')
|
||||
|
||||
self.assertTrue(status, msg)
|
||||
self.assertEqual(
|
||||
mock_conn.conn.prepare_threshold, 0,
|
||||
"_connect_readonly must force prepare_threshold=0 on the "
|
||||
"LLM connection; otherwise psycopg3 ignores prepare=True "
|
||||
"(when the server's Prepare threshold is blank/None, the "
|
||||
"default) and falls back to the multi-statement-capable "
|
||||
"simple query protocol, reopening CVE-2026-12045."
|
||||
)
|
||||
|
||||
@@ -44,10 +44,21 @@ LLM_APP_NAME_PREFIX = 'pgAdmin 4 - LLM'
|
||||
# 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.
|
||||
# subsequent statements in autocommit mode.
|
||||
#
|
||||
# The allowlist below rejects transaction-control statements, SET/RESET,
|
||||
# DML/DDL, CALL, COPY, and anything else that could weaken the read-only
|
||||
# sandbox, *for queries sqlparse correctly recognises as a single
|
||||
# statement*. It is a fast, user-friendly pre-filter, not the boundary
|
||||
# enforcement: sqlparse's lexer can disagree with PostgreSQL's own
|
||||
# parser about where one statement ends and another begins (e.g. string
|
||||
# literal escaping under standard_conforming_strings), so a payload can
|
||||
# look like one safe statement to sqlparse while PostgreSQL executes it
|
||||
# as several. The actual guarantee that only one statement reaches the
|
||||
# server is enforced at the protocol level -- see _execute_readonly_query,
|
||||
# which runs the query with prepare=True so PostgreSQL's own Parse step
|
||||
# (not a client-side approximation of it) rejects multi-statement text.
|
||||
#
|
||||
# EXPLAIN ANALYZE on a SELECT remains supported; EXPLAIN ANALYZE on a
|
||||
# write statement is blocked by PostgreSQL itself inside the read-only
|
||||
# transaction.
|
||||
@@ -126,6 +137,24 @@ def _connect_readonly(
|
||||
if not status:
|
||||
return False, msg
|
||||
|
||||
# Force the extended query protocol on this connection,
|
||||
# regardless of the server's configured "Prepare threshold"
|
||||
# (which defaults to blank/None, meaning "never prepare"). The
|
||||
# per-call prepare=True passed to execute_2darray() in
|
||||
# _execute_readonly_query() is silently ignored by psycopg3
|
||||
# whenever conn.prepare_threshold is None -- PrepareManager.get()
|
||||
# checks that before it looks at the prepare argument at all, and
|
||||
# falls back to the simple query protocol, which is exactly the
|
||||
# multi-statement-capable protocol this defense must avoid.
|
||||
# Setting the threshold to 0 ("always prepare") on this
|
||||
# single-use, per-query connection (see conn_id generation in
|
||||
# execute_readonly_query) makes the extended protocol -- and
|
||||
# therefore PostgreSQL's single-statement Parse-step guarantee --
|
||||
# unconditional here, without touching the server-wide setting or
|
||||
# any other connection.
|
||||
if getattr(conn, 'conn', None) is not None:
|
||||
conn.conn.prepare_threshold = 0
|
||||
|
||||
# Set application name via SQL - this is thread-safe and doesn't
|
||||
# require environment variables. The name will be visible in
|
||||
# pg_stat_activity to identify LLM connections.
|
||||
@@ -139,6 +168,16 @@ def _connect_readonly(
|
||||
# Non-fatal - connection still works without custom app name
|
||||
pass
|
||||
|
||||
# Defense-in-depth: make READ ONLY the session default, not just a
|
||||
# property of the transaction started by BEGIN TRANSACTION READ
|
||||
# ONLY. If a statement ever managed to end that transaction early
|
||||
# (e.g. a smuggled COMMIT), the next transaction on this connection
|
||||
# -- implicit or explicit -- would still be read-only rather than
|
||||
# falling back to a writable default.
|
||||
conn.execute_void(
|
||||
"SET SESSION CHARACTERISTICS AS TRANSACTION READ ONLY"
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
except Exception as e:
|
||||
@@ -166,20 +205,28 @@ def _first_real_keyword(statement) -> str:
|
||||
|
||||
def _validate_readonly_query(query: str) -> None:
|
||||
"""
|
||||
Ensure an LLM-supplied query is a single, read-only statement.
|
||||
Reject LLM-supplied queries that are obviously not 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.
|
||||
This is a fast pre-filter, not the security boundary: it rejects
|
||||
plainly-disallowed statement types and multi-statement input *as
|
||||
sqlparse's lexer sees it*. sqlparse's notion of where a string
|
||||
literal (and therefore a statement) ends can disagree with
|
||||
PostgreSQL's own parser -- e.g. a backslash before a quote is an
|
||||
escape to sqlparse but an ordinary character to PostgreSQL when
|
||||
standard_conforming_strings is on (the default) -- so a payload can
|
||||
pass this check as a single SELECT while PostgreSQL would actually
|
||||
execute it as several statements, including a COMMIT that ends the
|
||||
read-only transaction. That class of attack is closed at the
|
||||
protocol level, not here: see _execute_readonly_query, which
|
||||
executes the query with prepare=True so PostgreSQL's own Parse step
|
||||
-- the actual authority on statement boundaries -- rejects any text
|
||||
containing more than one statement.
|
||||
|
||||
Validation rules:
|
||||
|
||||
* The input must contain exactly one non-empty statement.
|
||||
* The input must contain exactly one non-empty statement (as far as
|
||||
sqlparse can tell).
|
||||
* The leading keyword must be in :data:`_ALLOWED_LEADING_KEYWORDS`.
|
||||
|
||||
PostgreSQL is left to enforce the rest -- ``EXPLAIN ANALYZE`` on a
|
||||
@@ -242,6 +289,19 @@ def _execute_readonly_query(conn, query: str) -> dict:
|
||||
The query is wrapped in a read-only transaction to ensure
|
||||
no data modifications can occur.
|
||||
|
||||
The query is executed with prepare=True, and the connection's
|
||||
prepare_threshold is forced to 0 in _connect_readonly(), which
|
||||
together force it through PostgreSQL's extended query protocol.
|
||||
(prepare=True alone is not enough: psycopg3 ignores it whenever
|
||||
prepare_threshold is None -- pgAdmin's blank-by-default server
|
||||
setting -- and falls back to the multi-statement-capable simple
|
||||
query protocol.) The server's Parse step accepts only a single SQL
|
||||
statement in the extended protocol, so this is what actually
|
||||
guarantees the LLM cannot smuggle a second statement (e.g. a COMMIT
|
||||
to end the read-only transaction early) past
|
||||
_validate_readonly_query's sqlparse-based check -- see the note on
|
||||
_validate_readonly_query for why that check alone is not sufficient.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query: SQL query to execute
|
||||
@@ -266,8 +326,10 @@ def _execute_readonly_query(conn, query: str) -> dict:
|
||||
)
|
||||
|
||||
try:
|
||||
# Execute the actual query
|
||||
status, result = conn.execute_2darray(query)
|
||||
# Execute the actual query. prepare=True, combined with the
|
||||
# prepare_threshold=0 forced in _connect_readonly(), is
|
||||
# load-bearing -- see the docstring above.
|
||||
status, result = conn.execute_2darray(query, prepare=True)
|
||||
|
||||
if not status:
|
||||
raise DatabaseToolError(
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
##########################################################################
|
||||
#
|
||||
# pgAdmin 4 - PostgreSQL Tools
|
||||
#
|
||||
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
|
||||
# This software is released under the PostgreSQL Licence
|
||||
#
|
||||
##########################################################################
|
||||
|
||||
"""End-to-end regression test for the CVE-2026-12045 bypass fix, driven
|
||||
through the REAL /sqleditor/nlq/chat/<trans_id>/stream HTTP route.
|
||||
|
||||
The existing NLQ tests in test_nlq_chat.py all patch
|
||||
``pgadmin.llm.chat.chat_with_database_stream`` directly, which stubs out
|
||||
the entire tool-call loop and never touches ``execute_tool`` /
|
||||
``execute_readonly_query`` in ``pgadmin.llm.tools.database``. That leaves
|
||||
a gap: nothing exercises the actual production code path a real AI
|
||||
Assistant request takes -- Flask route -> chat_with_database_stream's
|
||||
tool-dispatch loop -> execute_tool -> execute_readonly_query ->
|
||||
conn.execute_2darray -- to confirm prepare=True is still wired through
|
||||
when the "chosen" query originates from an LLM tool call rather than a
|
||||
direct Python call (as in test_database_tool_security.py's
|
||||
ExecuteReadonlyQueryProtocolTestCase).
|
||||
|
||||
This test closes that gap without needing a real LLM API key: it patches
|
||||
only ``pgadmin.llm.chat.get_llm_client`` (the same patch point the
|
||||
project's own EXPLAIN-analysis and NLQ tests already use for other
|
||||
scenarios) to return a fake client whose ``chat_stream()`` simulates the
|
||||
model choosing to call ``execute_sql_query`` with the smuggled-COMMIT
|
||||
payload from the Kai Aizen / SnailSploit report (2026-07-23). Everything
|
||||
below that -- chat.py's real tool-dispatch loop, execute_tool,
|
||||
execute_readonly_query -- runs for real. Only the DB connection object
|
||||
itself is mocked (as in ExecuteReadonlyQueryProtocolTestCase), since a
|
||||
real PostgreSQL connection isn't available in this test environment.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from pgadmin.utils.route import BaseTestGenerator
|
||||
from pgadmin.llm.models import LLMResponse, ToolCall, StopReason
|
||||
|
||||
|
||||
BYPASS_PAYLOAD = (
|
||||
"SELECT '\\';COMMIT;CREATE TABLE pwn_marker(x int);SELECT 1 --'"
|
||||
)
|
||||
|
||||
|
||||
def _make_mock_llm_client(tool_query, final_text='Here is what I found.'):
|
||||
"""A fake LLM client whose chat_stream() first requests a tool call
|
||||
with `tool_query`, then (once the tool result is fed back) ends the
|
||||
turn with `final_text`. Mirrors the shape client.chat_stream() is
|
||||
expected to yield: str chunks and/or a terminal LLMResponse.
|
||||
"""
|
||||
state = {'iteration': 0}
|
||||
|
||||
def chat_stream_side_effect(*_args, **_kwargs):
|
||||
state['iteration'] += 1
|
||||
if state['iteration'] == 1:
|
||||
yield LLMResponse(
|
||||
content='',
|
||||
tool_calls=[ToolCall(
|
||||
id='tc-1',
|
||||
name='execute_sql_query',
|
||||
arguments={'query': tool_query},
|
||||
)],
|
||||
stop_reason=StopReason.TOOL_USE,
|
||||
)
|
||||
else:
|
||||
yield final_text
|
||||
yield LLMResponse(
|
||||
content=final_text,
|
||||
stop_reason=StopReason.END_TURN,
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat_stream.side_effect = chat_stream_side_effect
|
||||
return mock_client
|
||||
|
||||
|
||||
class NLQToolDispatchProtocolTestCase(BaseTestGenerator):
|
||||
"""Pins that a real request through /sqleditor/nlq/chat/.../stream,
|
||||
when the model chooses to call execute_sql_query with the
|
||||
backslash-quote smuggled-COMMIT payload, still reaches
|
||||
conn.execute_2darray with prepare=True -- through the full
|
||||
production dispatch path, not just a direct Python call.
|
||||
|
||||
The DB connection is mocked (no real PostgreSQL available here), so
|
||||
this does not re-prove that PostgreSQL itself rejects the payload
|
||||
-- that is already covered by ExecuteReadonlyQueryProtocolTestCase
|
||||
(mock-based) and by manual verification against a live server. This
|
||||
test's job is narrower and complementary: prove the HTTP route
|
||||
really drives execute_readonly_query with this exact attacker-
|
||||
controlled string and that prepare=True is not lost or made
|
||||
conditional anywhere along the way.
|
||||
"""
|
||||
|
||||
scenarios = [
|
||||
('Tool call with bypass payload reaches DB layer with '
|
||||
'prepare=True', dict(
|
||||
mock_execute_2darray_result=(True, {
|
||||
'columns': [], 'rows': []
|
||||
}),
|
||||
)),
|
||||
('Tool call failure (simulating real Postgres Parse-step '
|
||||
'rejection) is surfaced, not a crash', dict(
|
||||
mock_execute_2darray_result=(
|
||||
False,
|
||||
'cannot insert multiple commands into a prepared '
|
||||
'statement'
|
||||
),
|
||||
)),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def runTest(self):
|
||||
trans_id = 12345
|
||||
|
||||
mock_trans_obj = MagicMock()
|
||||
mock_trans_obj.sid = 1
|
||||
mock_trans_obj.did = 1
|
||||
|
||||
mock_session_conn = MagicMock()
|
||||
mock_session_conn.connected.return_value = True
|
||||
|
||||
mock_session = {'sid': 1, 'did': 1}
|
||||
|
||||
# The connection object *inside* pgadmin.llm.tools.database --
|
||||
# this is what execute_2darray(prepare=True) is called on.
|
||||
mock_llm_conn = MagicMock()
|
||||
mock_llm_conn.execute_void.return_value = (True, None)
|
||||
mock_llm_conn.execute_2darray.return_value = \
|
||||
self.mock_execute_2darray_result
|
||||
|
||||
mock_client = _make_mock_llm_client(BYPASS_PAYLOAD)
|
||||
|
||||
patches = [
|
||||
patch('pgadmin.llm.utils.is_llm_enabled', return_value=True),
|
||||
patch('pgadmin.llm.chat.is_llm_available', return_value=True),
|
||||
patch(
|
||||
'pgadmin.llm.chat.get_llm_client',
|
||||
return_value=mock_client
|
||||
),
|
||||
patch(
|
||||
'pgadmin.tools.sqleditor.check_transaction_status',
|
||||
return_value=(
|
||||
True, None, mock_session_conn, mock_trans_obj,
|
||||
mock_session
|
||||
)
|
||||
),
|
||||
patch(
|
||||
'pgadmin.llm.tools.database._get_connection',
|
||||
return_value=(MagicMock(), mock_llm_conn)
|
||||
),
|
||||
patch(
|
||||
'pgadmin.llm.tools.database._connect_readonly',
|
||||
return_value=(True, None)
|
||||
),
|
||||
patch(
|
||||
'pgadmin.authenticate.mfa.utils.mfa_required',
|
||||
lambda f: f
|
||||
),
|
||||
]
|
||||
|
||||
for p in patches:
|
||||
p.start()
|
||||
|
||||
try:
|
||||
response = self.tester.post(
|
||||
f'/sqleditor/nlq/chat/{trans_id}/stream',
|
||||
data=json.dumps({'message': 'please run my query'}),
|
||||
content_type='application/json',
|
||||
follow_redirects=True
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertIn('text/event-stream', response.content_type)
|
||||
|
||||
# Consume the stream so the generator (and therefore the
|
||||
# real tool-dispatch loop) actually executes.
|
||||
raw = response.data.decode('utf-8')
|
||||
|
||||
# The route must not blow up with an unhandled exception --
|
||||
# either a 'complete' event (tool succeeded from the
|
||||
# model's perspective) or the loop simply continuing to a
|
||||
# final answer after a tool-error message is acceptable;
|
||||
# what must NOT happen is the request dying before
|
||||
# execute_2darray is ever reached.
|
||||
events = []
|
||||
for line in raw.split('\n'):
|
||||
if line.startswith('data: '):
|
||||
try:
|
||||
events.append(json.loads(line[6:]))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
event_types = [e.get('type') for e in events]
|
||||
self.assertIn('complete', event_types)
|
||||
|
||||
# The actual security-critical assertion: regardless of
|
||||
# whether the simulated DB call "succeeded" or "failed",
|
||||
# it was invoked with the attacker's exact payload and
|
||||
# prepare=True -- proving the fix is wired all the way
|
||||
# through the real HTTP route, not bypassed by some
|
||||
# earlier short-circuit.
|
||||
mock_llm_conn.execute_2darray.assert_called_once()
|
||||
call_args, call_kwargs = \
|
||||
mock_llm_conn.execute_2darray.call_args
|
||||
# execute_readonly_query wraps SELECT-prefixed queries with
|
||||
# a LIMIT subquery before executing (see database.py) -- the
|
||||
# attacker's payload must still be present verbatim inside
|
||||
# that wrapper, and it's the wrapped string that actually
|
||||
# gets sent with prepare=True.
|
||||
self.assertIn(BYPASS_PAYLOAD, call_args[0])
|
||||
self.assertTrue(
|
||||
call_kwargs.get('prepare') is True,
|
||||
"The real /sqleditor/nlq/chat/.../stream route must "
|
||||
"still invoke execute_2darray(query, prepare=True) "
|
||||
"for an LLM-chosen query, even when that query is the "
|
||||
"backslash-quote smuggled-COMMIT bypass payload."
|
||||
)
|
||||
finally:
|
||||
for p in patches:
|
||||
p.stop()
|
||||
|
||||
def tearDown(self):
|
||||
pass
|
||||
@@ -197,7 +197,7 @@ class BaseConnection(metaclass=ABCMeta):
|
||||
|
||||
@abstractmethod
|
||||
def execute_2darray(self, query, params=None,
|
||||
formatted_exception_msg=False):
|
||||
formatted_exception_msg=False, prepare=None):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -830,7 +830,7 @@ WHERE db.datname = current_database()""")
|
||||
25,
|
||||
'Psycopg3 Cursor: {0}'.format(str(e)))
|
||||
|
||||
def __internal_blocking_execute(self, cur, query, params):
|
||||
def __internal_blocking_execute(self, cur, query, params, prepare=None):
|
||||
"""
|
||||
This function executes the query using cursor's execute function,
|
||||
but in case of asynchronous connection we need to wait for the
|
||||
@@ -841,10 +841,17 @@ WHERE db.datname = current_database()""")
|
||||
cur: Cursor object
|
||||
query: SQL query to run.
|
||||
params: Extra parameters
|
||||
prepare: force the query through PostgreSQL's PREPARE step
|
||||
(extended query protocol) when True. Unlike the simple
|
||||
query protocol, the server rejects more than one SQL
|
||||
statement in a single Parse message, so this is used by
|
||||
callers that must guarantee a single statement is
|
||||
executed regardless of how a client-side SQL lexer
|
||||
would have classified the text.
|
||||
"""
|
||||
|
||||
query = query.encode(self.python_encoding)
|
||||
cur.execute(query, params)
|
||||
cur.execute(query, params, prepare=prepare)
|
||||
|
||||
def execute_on_server_as_csv(self, records=2000):
|
||||
"""
|
||||
@@ -1246,7 +1253,7 @@ WHERE db.datname = current_database()""")
|
||||
)
|
||||
|
||||
def execute_2darray(self, query, params=None,
|
||||
formatted_exception_msg=False):
|
||||
formatted_exception_msg=False, prepare=None):
|
||||
status, cur = self.__cursor()
|
||||
self.row_count = 0
|
||||
|
||||
@@ -1270,14 +1277,16 @@ WHERE db.datname = current_database()""")
|
||||
)
|
||||
)
|
||||
try:
|
||||
self.__internal_blocking_execute(cur, query, params)
|
||||
self.__internal_blocking_execute(
|
||||
cur, query, params, prepare=prepare
|
||||
)
|
||||
except psycopg.Error as pe:
|
||||
cur.close_cursor()
|
||||
if not self.connected() and self.auto_reconnect and \
|
||||
not self.reconnecting:
|
||||
return self.__attempt_execution_reconnect(
|
||||
self.execute_2darray, query, params,
|
||||
formatted_exception_msg
|
||||
formatted_exception_msg, prepare
|
||||
)
|
||||
errmsg = self._formatted_exception_msg(pe, formatted_exception_msg)
|
||||
current_app.logger.error(
|
||||
|
||||
Reference in New Issue
Block a user