fix matcher logic; clean up logging; add Postgres SSL matching support

Author: whummer

paradedb/localstack_paradedb/extension.py

@@ -1,13 +1,10 @@
 import os
 import socket
-import logging
 
 from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
 from werkzeug.datastructures import Headers
 from localstack import config
 
-LOG = logging.getLogger(__name__)
-
 # Environment variables for configuration
 ENV_POSTGRES_USER = "PARADEDB_POSTGRES_USER"
 ENV_POSTGRES_PASSWORD = "PARADEDB_POSTGRES_PASSWORD"
@@ -65,11 +62,26 @@ def tcp_connection_matcher(self, data: bytes) -> bool:
         """
         Identify PostgreSQL/ParadeDB connections by protocol handshake.
 
-        PostgreSQL startup message format:
+        PostgreSQL can start with either:
+        1. SSL request: protocol code 80877103 (0x04D2162F)
+        2. Startup message: protocol version 3.0 (0x00030000)
+
+        Both use the same format:
         - 4 bytes: message length
-        - 4 bytes: protocol version (3.0 = 0x00030000)
+        - 4 bytes: protocol version/code
         """
-        return len(data) >= 8 and data[4:8] == b"\x00\x03\x00\x00"
+        if len(data) < 8:
+            return False
+
+        # Check for SSL request (80877103 = 0x04D2162F)
+        if data[4:8] == b"\x04\xd2\x16\x2f":
+            return True
+
+        # Check for protocol version 3.0 (0x00030000)
+        if data[4:8] == b"\x00\x03\x00\x00":
+            return True
+
+        return False
 
     def should_proxy_request(self, headers: Headers) -> bool:
         """

utils/localstack_extensions/utils/docker.py

@@ -170,8 +170,7 @@ def _setup_tcp_protocol_routing(self):
         )
 
         LOG.info(
-            f"Registered TCP extension {self.name} -> "
-            f"{self.container_host}:{target_port} on gateway"
+            f"Registered TCP extension {self.name} -> {self.container_host}:{target_port} on gateway"
         )
 
     @abstractmethod
@@ -220,8 +219,6 @@ def start_container(self) -> None:
             self._remove_container()
             raise
 
-        LOG.debug("Successfully started extension container %s", self.container_name)
-
     def _default_health_check(self) -> None:
         """Default health check: HTTP GET request to the main port."""
         response = requests.get(f"http://{self.container_host}:{self.main_port}/")

utils/localstack_extensions/utils/h2_proxy.py

@@ -29,16 +29,29 @@ def __init__(self, port: int, host: str = "localhost"):
         self.host = host
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self._socket.connect((self.host, self.port))
+        self._closed = False
 
     def receive_loop(self, callback):
-        while data := self._socket.recv(self.buffer_size):
+        while data := self.recv(self.buffer_size):
             callback(data)
 
+    def recv(self, length):
+        try:
+            return self._socket.recv(length)
+        except OSError as e:
+            if self._closed:
+                return None
+            else:
+                raise e
+
     def send(self, data):
         self._socket.sendall(data)
 
     def close(self):
+        if self._closed:
+            return
         LOG.debug(f"Closing connection to upstream HTTP2 server on port {self.port}")
+        self._closed = True
         try:
             self._socket.shutdown(socket.SHUT_RDWR)
             self._socket.close()
@@ -93,7 +106,6 @@ def __init__(self, http_response_stream):
             )
 
         def received_from_backend(self, data):
-            LOG.debug(f"Received {len(data)} bytes from backend")
             self.http_response_stream.write(data)
 
         def received_from_http2_client(self, data, default_handler: Callable):
@@ -113,9 +125,6 @@ def received_from_http2_client(self, data, default_handler: Callable):
 
                         if should_proxy_request(headers):
                             self.state = ForwardingState.FORWARDING
-                            LOG.debug(
-                                f"Forwarding {len(buffered_data)} bytes to backend"
-                            )
                             self.backend.send(buffered_data)
                         else:
                             self.state = ForwardingState.PASSTHROUGH

utils/localstack_extensions/utils/tcp_protocol_router.py

@@ -12,8 +12,10 @@
 from twisted.web.http import HTTPChannel
 
 from localstack.utils.patch import patch
+from localstack import config
 
 LOG = logging.getLogger(__name__)
+LOG.setLevel(logging.DEBUG if config.DEBUG else logging.INFO)
 
 # Global registry of extensions with TCP matchers
 # List of tuples: (extension_name, matcher_func, backend_host, backend_port)
@@ -31,14 +33,19 @@ def connectionMade(self):
         # Set up peer relationship
         server.set_tcp_peer(self)
 
+        # Unregister any existing producer on server transport (HTTPChannel may have one)
+        try:
+            server.transport.unregisterProducer()
+        except Exception:
+            pass  # No producer was registered, which is fine
+
         # Enable flow control
         self.transport.registerProducer(server.transport, True)
         server.transport.registerProducer(self.transport, True)
 
         # Send buffered data from detection phase
         if hasattr(self.factory, "initial_data"):
             initial_data = self.factory.initial_data
-            LOG.debug(f"Sending {len(initial_data)} buffered bytes to backend")
             self.transport.write(initial_data)
             del self.factory.initial_data
 
@@ -61,12 +68,8 @@ def patch_gateway_for_tcp_routing():
     global _gateway_patched
 
     if _gateway_patched:
-        LOG.debug("Gateway already patched for TCP routing")
         return
 
-    LOG.debug("Patching LocalStack gateway for TCP protocol detection")
-    peek_bytes_length = 32
-
     # Patch HTTPChannel to use our protocol-detecting version
     @patch(HTTPChannel.__init__)
     def _patched_init(fn, self, *args, **kwargs):
@@ -76,7 +79,6 @@ def _patched_init(fn, self, *args, **kwargs):
         self._detection_buffer = []
         self._detecting = True
         self._tcp_peer = None
-        self._detection_buffer_size = peek_bytes_length
 
     @patch(HTTPChannel.dataReceived)
     def _patched_dataReceived(fn, self, data):
@@ -102,10 +104,6 @@ def _patched_dataReceived(fn, self, data):
             for ext_name, matcher, backend_host, backend_port in _tcp_extensions:
                 try:
                     if matcher(buffered_data):
-                        LOG.info(
-                            f"Extension {ext_name} claimed connection, routing to "
-                            f"{backend_host}:{backend_port}"
-                        )
                         # Switch to TCP proxy mode
                         self._detecting = False
                         self.transport.pauseProducing()
@@ -123,14 +121,11 @@ def _patched_dataReceived(fn, self, data):
                     continue
 
             # No extension claimed the connection
-            buffer_size = getattr(self, "_detection_buffer_size", peek_bytes_length)
-            if len(buffered_data) >= buffer_size:
-                LOG.debug("No TCP extension matched, using HTTP handler")
-                self._detecting = False
-                # Feed buffered data to HTTP handler
-                for chunk in self._detection_buffer:
-                    fn(self, chunk)
-                self._detection_buffer = []
+            self._detecting = False
+            # Feed buffered data to HTTP handler
+            for chunk in self._detection_buffer:
+                fn(self, chunk)
+            self._detection_buffer = []
 
     @patch(HTTPChannel.connectionLost)
     def _patched_connectionLost(fn, self, reason):
@@ -150,7 +145,6 @@ def set_tcp_peer(self, peer):
     HTTPChannel.set_tcp_peer = set_tcp_peer
 
     _gateway_patched = True
-    LOG.info("Gateway patched successfully for TCP protocol routing")
 
 
 def register_tcp_extension(