diff --git a/evoagentx/tools/database_postgresql.py b/evoagentx/tools/database_postgresql.py index 3bfddd2c..a0fe0dc7 100644 --- a/evoagentx/tools/database_postgresql.py +++ b/evoagentx/tools/database_postgresql.py @@ -4,6 +4,7 @@ from pathlib import Path import psycopg2 import psycopg2.extras +import psycopg2.sql import re from .database_base import DatabaseBase, DatabaseType, QueryType, DatabaseConnection @@ -893,7 +894,14 @@ def get_collection_info(self, collection_name: str) -> Dict[str, Any]: return self.format_error_result("PostgreSQL server not available") with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute(f"SELECT COUNT(*) as row_count FROM {collection_name}") + # Table names are identifiers and cannot be passed as bound + # parameters; use psycopg2.sql.Identifier to quote them + # safely instead of interpolating them into the SQL string. + cur.execute( + psycopg2.sql.SQL("SELECT COUNT(*) as row_count FROM {}").format( + psycopg2.sql.Identifier(collection_name) + ) + ) row_count = cur.fetchone()["row_count"] cur.execute("SELECT column_name, data_type FROM information_schema.columns WHERE table_name = %s", (collection_name,)) columns = cur.fetchall() @@ -904,6 +912,13 @@ def get_collection_info(self, collection_name: str) -> Dict[str, Any]: } return self.format_query_result(info, QueryType.SELECT) except Exception as e: + # Rollback on error so an aborted transaction does not poison + # subsequent queries (same convention as execute_query). + try: + if self.conn: + self.conn.rollback() + except Exception as rollback_error: + logger.warning(f"Error during rollback: {str(rollback_error)}") return self.format_error_result(str(e)) def get_schema(self, collection_name: str = None) -> Dict[str, Any]: diff --git a/tests/src/tools/test_database_postgresql.py b/tests/src/tools/test_database_postgresql.py new file mode 100644 index 00000000..faf99c75 --- /dev/null +++ b/tests/src/tools/test_database_postgresql.py @@ -0,0 +1,118 @@ +import os +from unittest.mock import MagicMock + +import pytest + +try: + import psycopg2 + import psycopg2.sql + + HAS_PSYCOPG2 = True +except ImportError: # pragma: no cover + HAS_PSYCOPG2 = False + +from evoagentx.tools.database_postgresql import PostgreSQLDatabase + +PAYLOAD = "information_schema.tables WHERE table_name='users'--" + + +def _make_db(collection_name="public"): + """Build a PostgreSQLDatabase whose remote init is fully mocked.""" + db = PostgreSQLDatabase.__new__(PostgreSQLDatabase) + db.conn = MagicMock() + db.file_based_mode = False + db._is_initialized = True + db.is_local_database = False + db.local_path = None + db.auto_save = False + db.database_name = collection_name + return db + + +def _executed_statements(db): + """Return every first positional arg passed to cursor.execute.""" + cursor = db.conn.cursor.return_value.__enter__.return_value + return [call.args[0] for call in cursor.execute.call_args_list] + + +@pytest.mark.skipif(not HAS_PSYCOPG2, reason="psycopg2 not installed") +class TestGetCollectionInfoSqlInjection: + def test_identifier_is_parameterized_not_interpolated(self): + """Regression test for issue #287: the table name must reach the + driver as a psycopg2.sql.Identifier, never interpolated into a + plain SQL string.""" + db = _make_db() + db.get_collection_info(PAYLOAD) + + statements = _executed_statements(db) + assert statements, "expected a statement to be executed" + for stmt in statements: + # No raw interpolated SQL may contain the payload verbatim. + assert not (isinstance(stmt, str) and PAYLOAD in stmt) + count_stmt = statements[0] + assert isinstance(count_stmt, psycopg2.sql.Composed) + # The payload is carried as a quoted identifier inside the Composed. + assert any( + isinstance(part, psycopg2.sql.Identifier) + and getattr(part, "strings", None) == (PAYLOAD,) + for part in count_stmt.seq + ) + + def test_injection_is_never_executed_as_sql(self): + """The whole payload is one quoted identifier: it cannot split the + statement, so the injected WHERE clause never runs.""" + db = _make_db() + db.get_collection_info(PAYLOAD) + + count_stmt = _executed_statements(db)[0] + literal = "".join( + part.string for part in count_stmt.seq if isinstance(part, psycopg2.sql.SQL) + ) + # The literal SQL offers no surface: the payload only rides along as + # the quoted identifier, never as executable statement text. + assert literal == "SELECT COUNT(*) as row_count FROM " + assert PAYLOAD not in literal + + def test_rollback_on_error_keeps_connection_usable(self): + """A rejected name must roll back so the aborted transaction does + not poison subsequent queries (same convention as execute_query).""" + db = _make_db() + cursor = db.conn.cursor.return_value.__enter__.return_value + cursor.execute.side_effect = psycopg2.errors.UndefinedTable( + 'relation "x" does not exist' + ) + + result = db.get_collection_info("no_such_table") + + assert result["success"] is False + db.conn.rollback.assert_called_once() + + +@pytest.mark.live_postgres +class TestGetCollectionInfoLive: + """Opt-in live check against a disposable PostgreSQL instance: + EVOAGENTX_TEST_PG_DSN=postgresql://postgres:postgres@localhost:5434/postgres \ + pytest tests/src/tools/test_database_postgresql.py -m live_postgres + """ + + @pytest.mark.skipif( + not os.environ.get("EVOAGENTX_TEST_PG_DSN"), + reason="EVOAGENTX_TEST_PG_DSN not set", + ) + def test_injected_payload_returns_error_and_legit_query_works(self): + dsn = os.environ["EVOAGENTX_TEST_PG_DSN"] + db = PostgreSQLDatabase(connection_string=dsn) + assert not db.file_based_mode, "live test requires a real connection" + + with db.conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS users") + cur.execute("CREATE TABLE users (id SERIAL PRIMARY KEY, name TEXT)") + cur.execute("INSERT INTO users (name) VALUES ('alice'), ('bob')") + db.conn.commit() + + injected = db.get_collection_info(PAYLOAD) + assert injected["success"] is False + + legit = db.get_collection_info("users") + assert legit["success"] is True + assert legit["data"]["row_count"] == 2