Only explain sampled Postgres statements that are a single SQL statement

postgres/changelog.d/24839.security

@@ -0,0 +1 @@
+Send the explain PREPARE for a parameterized query over the extended query protocol, so the server rejects sampled text holding more than one statement.

postgres/datadog_checks/postgres/explain_parameterized_queries.py

@@ -39,6 +39,14 @@
 )
 
 
+class ExtendedProtocolUnavailable(Exception):
+    """Raised when a sampled statement can't be prepared over the extended query protocol.
+
+    Preparing it over the simple protocol instead would execute every statement in the sampled text, so
+    the query goes unexplained rather than being sent in a way that could run a planted separator.
+    """
+
+
 def agent_check_getter(self):
     return self._check
 
@@ -83,6 +91,15 @@ def __init__(self, check, config, explain_function):
         self._check = check
         self._config = config
         self._explain_function = explain_function
+        # Checked once here rather than per statement: it only depends on the libpq this agent was built
+        # against. See _execute_prepare for why a pipeline is required to explain a parameterized query.
+        self._can_use_pipeline = psycopg.capabilities.has_pipeline(check=False)
+        if not self._can_use_pipeline:
+            logger.warning(
+                "Parameterized queries cannot be explained: this build links libpq %s, and pipeline mode "
+                "(libpq 14+) is required to send the prepared statement safely",
+                psycopg.pq.version(),
+            )
 
     @tracked_method(agent_check_getter=agent_check_getter)
     def explain_statement(
@@ -146,7 +163,7 @@ def _create_prepared_statement(
         # Returns None on success, or a (DBExplainError, err_msg) tuple when the query can't be prepared because
         # a parameter's type can't be resolved. Other unexpected errors are re-raised.
         try:
-            self._execute_query(
+            self._execute_prepare(
                 conn,
                 PREPARE_STATEMENT_QUERY.format(query_signature=query_signature, statement=statement),
             )
@@ -251,6 +268,18 @@ def _deallocate_prepared_statement(self, conn, query_signature):
                 e,
             )
 
+    def _execute_prepare(self, conn, query):
+        # The PREPARE is the one place sampled query text becomes SQL, so it goes out over the extended query
+        # protocol, where the server rejects a multi-command string ("cannot insert multiple commands into a
+        # prepared statement") rather than executing every statement in it as the monitoring user. psycopg
+        # leaves the simple protocol only for a query with parameters, one requesting binary results (which
+        # the client-side cursors this pool uses refuse) or one inside a pipeline, so a pipeline is the only
+        # route available here.
+        if not self._can_use_pipeline:
+            raise ExtendedProtocolUnavailable("cannot prepare a sampled statement without pipeline support (libpq 14+)")
+        with conn.pipeline():
+            self._execute_query(conn, query)
+
     def _execute_query(self, conn, query):
         with conn.cursor() as cursor:
             logger.debug('Executing query=[%s]', query)

postgres/datadog_checks/postgres/util.py

@@ -277,7 +277,7 @@ class DBExplainError(Enum):
     # not able to access the required function
     database_error = 'database_error'
 
-    # datatype mismatch occurs when return type is not json, for instance when multiple queries are explained
+    # datatype mismatch occurs when the return type of the EXPLAIN function is not json
     datatype_mismatch = 'datatype_mismatch'
 
     # this could be the result of a missing EXPLAIN function

postgres/tests/test_explain_parameterized_queries.py

@@ -8,6 +8,7 @@
 import pytest
 
 from datadog_checks.base.utils.db.sql import compute_sql_signature
+from datadog_checks.postgres.explain_parameterized_queries import ExtendedProtocolUnavailable
 from datadog_checks.postgres.util import DBExplainError
 from datadog_checks.postgres.version_utils import V12
 
@@ -38,6 +39,8 @@ def dbm_instance(pg_instance):
     "query,expected_explain_err_code",
     [
         ("SELECT * FROM pg_settings WHERE name = $1", DBExplainError.explained_with_prepared_statement),
+        # a single trailing statement terminator is legitimate and must still be explained
+        ("SELECT * FROM pg_settings WHERE name = $1;", DBExplainError.explained_with_prepared_statement),
         (
             "SELECT * FROM pg_settings WHERE name = $1 AND "
             "context = (SELECT context FROM pg_settings WHERE vartype = $2) AND source = $3",
@@ -95,6 +98,67 @@ def test_explain_parameterized_queries_generic_params(integration_check, dbm_ins
         )
 
 
+@pytest.mark.integration
+@pytest.mark.usefixtures("dd_environment")
+@requires_over_12
+def test_stacked_statements_are_rejected_by_the_server(integration_check, dbm_instance):
+    '''
+    VULN-92306: sampled text comes from pg_stat_activity, so whoever ran the query chose it, and the agent
+    runs the PREPARE as the monitoring user. The extended query protocol is what stops a planted separator
+    from executing, so this asserts the server refuses it and that nothing ran.
+    '''
+    check = integration_check(dbm_instance)
+    check.check(dbm_instance)
+
+    # the trailing comment supplies the $1 marker that routes the statement into the prepared statement path
+    query = "SELECT 1; CREATE TEMP TABLE dd_injection_marker(x int); --$1"
+    query_signature = compute_sql_signature(query)
+
+    plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
+        DB_NAME, query, query, query_signature
+    )
+
+    assert plan_dict is None
+    assert explain_err_code == DBExplainError.failed_to_explain_with_prepared_statement
+    assert err == "<class 'psycopg.errors.SyntaxError'>"
+
+    with check.db_pool.get_connection(DB_NAME) as conn:
+        rows = check.statement_samples._explain_parameterized_queries._execute_query_and_fetch_rows(
+            conn, "SELECT 1 FROM pg_class WHERE relname = 'dd_injection_marker'"
+        )
+    assert rows == [], "the agent executed the injected statement"
+
+
+@pytest.mark.unit
+def test_execute_prepare_uses_a_pipeline(integration_check, dbm_instance):
+    check = integration_check(dbm_instance)
+    epq = check.statement_samples._explain_parameterized_queries
+    conn = mock.MagicMock()
+
+    with mock.patch.object(epq, '_execute_query') as mock_execute:
+        epq._execute_prepare(conn, "PREPARE dd_test AS SELECT 1")
+
+    conn.pipeline.assert_called_once()
+    mock_execute.assert_called_once_with(conn, "PREPARE dd_test AS SELECT 1")
+
+
+@pytest.mark.unit
+def test_execute_prepare_fails_closed_without_pipeline_support(integration_check, dbm_instance):
+    """Without a pipeline the PREPARE would go out over the simple query protocol, which executes every
+    statement in the sampled text, so the query goes unexplained instead."""
+    check = integration_check(dbm_instance)
+    epq = check.statement_samples._explain_parameterized_queries
+    epq._can_use_pipeline = False
+    conn = mock.MagicMock()
+
+    with mock.patch.object(epq, '_execute_query') as mock_execute:
+        with pytest.raises(ExtendedProtocolUnavailable):
+            epq._execute_prepare(conn, "PREPARE dd_test AS SELECT 1")
+
+    mock_execute.assert_not_called()
+    conn.pipeline.assert_not_called()
+
+
 @pytest.mark.integration
 @pytest.mark.usefixtures("dd_environment")
 def test_explain_parameterized_queries_version_below_12(integration_check, dbm_instance):
@@ -293,7 +357,7 @@ def test_create_prepared_statement_exception(integration_check, dbm_instance, ex
     query = "SELECT * FROM pg_settings WHERE name = $1"
     query_signature = compute_sql_signature(query)
     with mock.patch(
-        'datadog_checks.postgres.explain_parameterized_queries.ExplainParameterizedQueries._execute_query',
+        'datadog_checks.postgres.explain_parameterized_queries.ExplainParameterizedQueries._execute_prepare',
         side_effect=exception_class,
     ):
         with pytest.raises(exception_class):
@@ -315,7 +379,7 @@ def test_create_prepared_statement_datatype_mismatch_maps_to_code(integration_ch
     check = integration_check(dbm_instance)
     epq = check.statement_samples._explain_parameterized_queries
 
-    with mock.patch.object(epq, '_execute_query', side_effect=psycopg.errors.DatatypeMismatch("type mismatch")):
+    with mock.patch.object(epq, '_execute_prepare', side_effect=psycopg.errors.DatatypeMismatch("type mismatch")):
         result = epq._create_prepared_statement(
             None,
             "SELECT id FROM t WHERE id = $1",
@@ -335,7 +399,7 @@ def test_create_prepared_statement_ambiguous_function_maps_to_code(integration_c
 
     with mock.patch.object(
         epq,
-        '_execute_query',
+        '_execute_prepare',
         side_effect=psycopg.errors.AmbiguousFunction("function unnest(unknown) is not unique"),
     ):
         result = epq._create_prepared_statement(