fix connection handling for non-TypeDB HTTP2 requests

Author: whummer

typedb/Makefile

@@ -2,6 +2,7 @@ VENV_BIN = python3 -m venv
 VENV_DIR ?= .venv
 VENV_ACTIVATE = $(VENV_DIR)/bin/activate
 VENV_RUN = . $(VENV_ACTIVATE)
+TEST_PATH ?= tests
 
 usage:          ## Shows usage for this Makefile
 	@cat Makefile | grep -E '^[a-zA-Z_-]+:.*?## .*$$' | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-15s\033[0m %s\n", $$1, $$2}'
@@ -39,7 +40,7 @@ lint:		    ## Run ruff to lint the codebase
 	$(VENV_RUN); python -m ruff check --output-format=full .
 
 test:		    ## Run integration tests (requires LocalStack running with the Extension installed)
-	$(VENV_RUN); pytest tests $(PYTEST_ARGS)
+	$(VENV_RUN); pytest $(PYTEST_ARGS) $(TEST_PATH)
 
 clean-dist: clean
 	rm -rf dist/

typedb/localstack_typedb/utils/h2_proxy.py

@@ -70,14 +70,18 @@ class ForwardingBuffer:
         to the backend, or leave it to the default handler.
         """
 
+        backend: TcpForwarder
+        buffer: list
+        proxying: bool | None
+
         def __init__(self, http_response_stream):
             self.http_response_stream = http_response_stream
             LOG.debug(
                 f"Starting TCP forwarder to port {target_port} for new HTTP2 connection"
             )
             self.backend = TcpForwarder(target_port, host=target_host)
             self.buffer = []
-            self.proxying = False
+            self.proxying = None
             reactor.getThreadPool().callInThread(
                 self.backend.receive_loop, self.received_from_backend
             )
@@ -86,23 +90,36 @@ 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):
+        def received_from_http2_client(self, data, default_handler: Callable):
+            if self.proxying is False:
+                # Note: Return here only if `proxying` is `False` (a value of `None` indicates
+                # that the headers have not fully been received yet)
+                return default_handler(data)
+
             if self.proxying:
                 assert not self.buffer
                 # Keep sending data to the backend for the lifetime of this connection
                 self.backend.send(data)
-            else:
-                self.buffer.append(data)
-                if headers := get_headers_from_data_stream(self.buffer):
-                    self.proxying = should_proxy_request(headers)
-                    # Now we know what to do with the buffer
-                    buffered_data = b"".join(self.buffer)
-                    self.buffer = []
-                    if self.proxying:
-                        LOG.debug(f"Forwarding {len(buffered_data)} bytes to backend")
-                        self.backend.send(buffered_data)
-                    else:
-                        return default_handler(buffered_data)
+                return
+
+            self.buffer.append(data)
+
+            if not (headers := get_headers_from_data_stream(self.buffer)):
+                # If no headers received yet, then return (method will be called again for next chunk of data)
+                return
+
+            self.proxying = should_proxy_request(headers)
+
+            buffered_data = b"".join(self.buffer)
+            self.buffer = []
+
+            if not self.proxying:
+                # if this is not a target request, then call the default handler
+                default_handler(buffered_data)
+                return
+
+            LOG.debug(f"Forwarding {len(buffered_data)} bytes to backend")
+            self.backend.send(buffered_data)
 
         def close(self):
             self.backend.close()

typedb/tests/test_extension.py

@@ -1,4 +1,5 @@
 import requests
+import httpx
 from localstack.utils.strings import short_uid
 from localstack_typedb.utils.h2_proxy import (
     get_frames_from_http2_stream,
@@ -75,6 +76,30 @@ def test_connect_to_db_via_grpc_endpoint():
             assert len(results) == 2
 
 
+def test_connect_to_h2_endpoint_non_typedb():
+    url = "https://s3.localhost.localstack.cloud:4566/"
+
+    # make an HTTP/2 request to the LocalStack health endpoint
+    with httpx.Client(http2=True, verify=False, trust_env=False) as client:
+        health_url = f"{url}/_localstack/health"
+        response = client.get(health_url)
+
+    assert response.status_code == 200
+    assert response.http_version == "HTTP/2"
+    assert '"services":' in response.text
+
+    # make an HTTP/2 request to a LocalStack endpoint outside the extension (S3 list buckets)
+    headers = {
+        "Authorization": "AWS4-HMAC-SHA256 Credential=000000000000/20250101/us-east-1/s3/aws4_request, ..."
+    }
+    with httpx.Client(http2=True, verify=False, trust_env=False) as client:
+        response = client.get(url, headers=headers)
+
+    assert response.status_code == 200
+    assert response.http_version == "HTTP/2"
+    assert "<ListAllMyBucketsResult" in response.text
+
+
 def test_get_frames_from_http2_stream():
     # note: the data below is a dump taken from a browser request made against the emulator
     data = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x18\x04\x00\x00\x00\x00\x00\x00\x01\x00\x01\x00\x00\x00\x02\x00\x00\x00\x00\x00\x04\x00\x02\x00\x00\x00\x05\x00\x00@\x00\x00\x00\x04\x08\x00\x00\x00\x00\x00\x00\xbf\x00\x01"