PR comments

Author: nik-localstack

postgresql_proxy/interceptors.py

@@ -55,20 +55,14 @@ def intercept(self, packet_type, data):
                 # Query, ends with b'\x00'
                 data = self._intercept_query(data, ic_queries)
             elif packet_type == b"P":
-                # Parse packet body:
-                #   statement_name\x00 + query\x00 + int16(param_count) + uint32[]
-                # Keep the binary suffix untouched (count + OID array).
-                statement_end = data.find(b"\x00")
-                if statement_end != -1:
-                    query_start = statement_end + 1
-                    query_end = data.find(b"\x00", query_start)
-                    if query_end != -1:
-                        statement = data[:query_start]
-                        query = self._intercept_query(
-                            data[query_start : query_end + 1], ic_queries
-                        )
-                        params = data[query_end + 1 :]
-                        data = statement + query + params
+                # Parse message (Extended Query Protocol). Body format:
+                #   statement_name\x00 + query\x00 + int16(param_count) + uint32[] OIDs
+                # maxsplit=2 keeps the binary suffix intact even when OID values contain \x00 bytes.
+                parts = data.split(b"\x00", 2)
+                if len(parts) == 3:
+                    statement, raw_query, params = parts
+                    query = self._intercept_query(raw_query, ic_queries)
+                    data = statement + b"\x00" + query + params
 
         if packet_type == b"":
             # Connection request / context. Ignore the first 4 bytes, keep it
@@ -106,8 +100,8 @@ def _intercept_context_data(self, data):
 
     def _intercept_query(self, query, interceptors):
         logging.getLogger("intercept").debug("intercepting query\n%s", query)
-        # Remove zero byte at the end
-        query = query[:-1].decode("utf-8")
+        # Remove null terminator
+        query = query.rstrip(b"\x00").decode("utf-8")
         for interceptor in interceptors:
             func = self._get_plugin_interceptor_function(interceptor)
             query = func(query, self.context)

tests/test_proxy.py

@@ -1,3 +1,4 @@
+from collections.abc import Generator
 import contextlib
 import os
 import shutil
@@ -16,6 +17,17 @@
 from postgresql_proxy.proxy import Proxy
 
 
+class _QuerySpy:
+    """Minimal query interceptor that records every query it sees."""
+
+    def __init__(self):
+        self.captured: list[str] = []
+
+    def capture(self, query: str, context) -> str:
+        self.captured.append(query)
+        return query
+
+
 def _get_free_tcp_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
         sock.bind(("127.0.0.1", 0))
@@ -126,8 +138,18 @@ def _temporary_server_cert_pair():
 
 
 @contextlib.contextmanager
-def _run_proxy(postgres_settings, ssl_context: ssl.SSLContext | None = None):
+def _run_proxy(
+    postgres_settings,
+    ssl_context: ssl.SSLContext | None = None,
+    *,
+    plugins: dict | None = None,
+    query_interceptors: list[dict] | None = None,
+) -> Generator[int, None, None]:
+    """Start a proxy in a background thread and yield its listening port."""
     proxy_port = _get_free_tcp_port()
+    commands_config: dict = {}
+    if query_interceptors:
+        commands_config["queries"] = query_interceptors
     instance = cfg.InstanceSettings(
         {
             "listen": {"name": "proxy", "host": "127.0.0.1", "port": proxy_port},
@@ -136,14 +158,13 @@ def _run_proxy(postgres_settings, ssl_context: ssl.SSLContext | None = None):
                 "host": postgres_settings["host"],
                 "port": postgres_settings["port"],
             },
-            # Keep interceptors active with default no-op behavior.
-            "intercept": {"commands": {}, "responses": {}},
+            "intercept": {"commands": commands_config, "responses": {}},
         }
     )
     if not hasattr(instance.intercept.responses, "parameter_status"):
         instance.intercept.responses.parameter_status = []
 
-    proxy = Proxy(instance, plugins={}, debug=True, ssl_context=ssl_context)
+    proxy = Proxy(instance, plugins=plugins or {}, debug=True, ssl_context=ssl_context)
     thread = threading.Thread(
         target=proxy.listen, kwargs={"max_connections": 32}, daemon=True
     )
@@ -336,39 +357,57 @@ def test_psql_ssl_file_batch_stress_no_hang(postgres_settings, ssl_proxy_port):
 
 
 def test_extended_query_protocol_parse_packet_with_high_oid_params_passes_through_proxy(
-    postgres_settings, plain_proxy_port
+    postgres_settings,
 ):
     """Regression: proxy must not corrupt Extended Query Protocol Parse packets.
 
     psycopg v3 sends Parse → Bind → Execute for parameterized queries.  The Parse body
     ends with binary uint32 OIDs; jsonb OID 3802 (0x00000EDA) contains 0xDA which is
     not valid UTF-8.  The old interceptor sliced the body incorrectly and crashed on
     decode, causing the connection to hang or drop.
+
+    A _QuerySpy is wired into the proxy to verify the interceptor receives the correct
+    SQL text — not corrupted bytes from the binary OID suffix.
     """
-    with psycopg.connect(
-        host="127.0.0.1",
-        port=plain_proxy_port,
-        user=postgres_settings["user"],
-        password=postgres_settings["password"],
-        dbname=postgres_settings["dbname"],
-        sslmode="disable",
-    ) as conn:
-        with conn.cursor() as cur:
-            cur.execute(
-                "DROP TABLE IF EXISTS _test_jsonb_proxy_params;"
-                "CREATE TABLE _test_jsonb_proxy_params "
-                "(id serial PRIMARY KEY, data jsonb, label text);"
-            )
+    spy = _QuerySpy()
+    with _run_proxy(
+        postgres_settings,
+        plugins={"spy": spy},
+        query_interceptors=[{"plugin": "spy", "function": "capture"}],
+    ) as proxy_port:
+        with psycopg.connect(
+            host="127.0.0.1",
+            port=proxy_port,
+            user=postgres_settings["user"],
+            password=postgres_settings["password"],
+            dbname=postgres_settings["dbname"],
+            sslmode="disable",
+        ) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "DROP TABLE IF EXISTS _test_jsonb_proxy_params;"
+                    "CREATE TABLE _test_jsonb_proxy_params "
+                    "(id serial PRIMARY KEY, data jsonb, label text);"
+                )
 
-            cur.execute(
-                "INSERT INTO _test_jsonb_proxy_params (data, label) "
-                "VALUES (%s, %s) RETURNING id",
-                (psycopg.types.json.Jsonb({"key": "value"}), "hello"),
-            )
-            row = cur.fetchone()
+                cur.execute(
+                    "INSERT INTO _test_jsonb_proxy_params (data, label) "
+                    "VALUES (%s, %s) RETURNING id",
+                    (psycopg.types.json.Jsonb({"key": "value"}), "hello"),
+                )
+                row = cur.fetchone()
 
     assert row is not None and row[0] >= 1
 
+    # Verify the interceptor received clean SQL — no binary OID bytes leaked in.
+    insert_queries = [
+        q for q in spy.captured if "INSERT INTO _test_jsonb_proxy_params" in q
+    ]
+    assert insert_queries
+    assert all("\x00" not in q for q in insert_queries), (
+        "null byte leaked into intercepted query"
+    )
+
 
 def test_extended_query_protocol_named_prepared_statement_passes_through_proxy(
     postgres_settings, plain_proxy_port