Fix: Extended Query Protocol Parse packets corrupted by interceptor (#18)

* fix: correctly parse Extended Query Protocol Parse packets with binary OID suffix
The Parse packet body (type 'P') has the format:
  statement_name\x00 + query\x00 + int16(param_count) + uint32[] OIDs
The old handler used data[1:-2] / data[-2:] which treated only the last
2 bytes as the param-count field, leaking binary OID bytes into the query
slice fed to _intercept_query().  When a parameter type OID contains a
byte >= 0x80 (e.g. jsonb OID 3802 = 0x00000EDA), the subsequent UTF-8
decode raised UnicodeDecodeError, causing the connection to drop or hang.
Fix: use find(b'\x00') to locate the true boundaries of statement name
and query text, so the binary suffix is never touched by the text decoder.
Adds a regression test using psycopg v3 (Extended Query Protocol) with
a jsonb+text parameterized INSERT through the proxy.

Co-authored-by: GitHub Copilot <copilot@github.com>

* PR comments

---------

Co-authored-by: GitHub Copilot <copilot@github.com>

Author: nik-localstack

postgresql_proxy/interceptors.py

@@ -55,14 +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":
-                # Statement that needs parsing.
-                # First byte of the body is some Statement flag. Ignore, don't lose
-                # Next is the query, same as above, ends with an b'\x00'
-                # Last 2 bytes are the number of parameters. Ignore, don't lose
-                statement = data[0:1]
-                query = self._intercept_query(data[1:-2], ic_queries)
-                params = data[-2:]
-                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
@@ -100,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)

requirements-test.txt

@@ -1,2 +1,3 @@
 pytest==9.0.3
 pytest-timeout==2.4.0
+psycopg[binary]==3.3.4

tests/test_proxy.py

@@ -1,3 +1,4 @@
+from collections.abc import Generator
 import contextlib
 import os
 import shutil
@@ -8,13 +9,25 @@
 import threading
 import time
 
+import psycopg
 import psycopg2
 import pytest
 
 from postgresql_proxy import config_schema as cfg
 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))
@@ -125,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},
@@ -135,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
     )
@@ -332,3 +354,90 @@ def test_psql_ssl_file_batch_stress_no_hang(postgres_settings, ssl_proxy_port):
                     "psql -f batch succeeded but expected marker missing "
                     f"(run={run_idx + 1}, {elapsed=:.2f}s) stdout_tail={out_tail}"
                 )
+
+
+def test_extended_query_protocol_parse_packet_with_high_oid_params_passes_through_proxy(
+    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.
+    """
+    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()
+
+    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
+):
+    """Parse packets with a non-empty statement name must also be relayed correctly.
+
+    The statement_name field precedes the query text in the Parse body.  The fix uses
+    find(b'\\x00') to locate boundaries, so named statements work the same as anonymous
+    ones (empty name).
+    """
+    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",
+        # Prepare after the first execution of the same query (i.e. on 2nd run).
+        prepare_threshold=1,
+    ) as conn:
+        with conn.cursor() as cur:
+            # Execute twice so psycopg can promote the query to a named statement.
+            for val in (1, 2):
+                cur.execute("SELECT %s::int + 1", (val,))
+                result = cur.fetchone()
+                assert result == (val + 1,)
+
+            # Verify psycopg created a named prepared statement in this session.
+            cur.execute(
+                "SELECT count(*) FROM pg_prepared_statements WHERE name LIKE '_pg3_%'"
+            )
+            prepared_count = cur.fetchone()
+            assert prepared_count is not None and prepared_count[0] >= 1