fix header casing for Twisted WebSockets (#26)

Author: bentsku

pyproject.toml

@@ -7,7 +7,7 @@ name = "rolo"
 authors = [
     { name = "LocalStack Contributors", email = "info@localstack.cloud" }
 ]
-version = "0.7.4"
+version = "0.7.5"
 description = "A Python framework for building HTTP-based server applications"
 dependencies = [
     "requests>=2.20",

rolo/testing/pytest.py

@@ -17,7 +17,7 @@
 from rolo.gateway.asgi import AsgiGateway
 from rolo.gateway.wsgi import WsgiGateway
 from rolo.routing import handler_dispatcher
-from rolo.serving.twisted import TwistedGateway
+from rolo.serving.twisted import HeaderPreservingHTTPChannel, TwistedGateway
 from rolo.websocket.adapter import WebSocketListener
 
 if typing.TYPE_CHECKING:
@@ -237,6 +237,12 @@ def _create(gateway):
 
 @pytest.fixture
 def serve_twisted_websocket_listener(twisted_reactor, serve_twisted_tcp_server):
+    """
+    This fixture creates a Twisted Site, without the need to serve a fully-fledged rolo Gateway.
+    This is inspired by `rolo.serving.twisted.TwistedGateway`, but directly uses `WebsocketResourceDecorator` to
+    pass the `WebSocketListener` instead of `gateway.accept` to the `websocketListener` parameter.
+    It allows us to test the low-level behavior of WebSockets without being dependent on the Gateway implementation.
+    """
     from twisted.web.server import Site
 
     from rolo.serving.twisted import HeaderPreservingWSGIResource, WebsocketResourceDecorator
@@ -250,6 +256,7 @@ def _create(websocket_listener: WebSocketListener):
                 websocketListener=websocket_listener,
             )
         )
+        site.protocol = HeaderPreservingHTTPChannel.protocol_factory
         return serve_twisted_tcp_server(site)
 
     return _create

rolo/websocket/request.py

@@ -3,7 +3,7 @@
 
 from werkzeug import Response
 from werkzeug._internal import _wsgi_decoding_dance
-from werkzeug.datastructures import EnvironHeaders, Headers
+from werkzeug.datastructures import EnvironHeaders, Headers, MultiDict
 from werkzeug.sansio.request import Request as _SansIORequest
 from werkzeug.wsgi import _get_server
 
@@ -115,6 +115,17 @@ def __init__(self, environ: WebSocketEnvironment):
 
         :param environ: the WebSocketEnvironment
         """
+        raw_headers = environ.get("rolo.headers")
+        if raw_headers:
+            # restores raw headers from the server scope, to have proper casing or dashes. This can depend on server
+            # behavior, but we want a unified way to keep header casing/formatting.
+            # This is similar to what we do in wsgi.py
+            headers = Headers(
+                MultiDict([(k.decode("latin-1"), v.decode("latin-1")) for (k, v) in raw_headers])
+            )
+        else:
+            headers = Headers(EnvironHeaders(environ))
+
         # copied from werkzeug.wrappers.request
         super().__init__(
             method=environ.get("REQUEST_METHOD", "WEBSOCKET"),
@@ -123,7 +134,7 @@ def __init__(self, environ: WebSocketEnvironment):
             root_path=_wsgi_decoding_dance(environ.get("SCRIPT_NAME") or ""),
             path=_wsgi_decoding_dance(environ.get("PATH_INFO") or ""),
             query_string=environ.get("QUERY_STRING", "").encode("latin1"),
-            headers=Headers(EnvironHeaders(environ)),
+            headers=headers,
             remote_addr=environ.get("REMOTE_ADDR"),
         )
         self.environ = environ

tests/websocket/test_websockets.py

@@ -92,7 +92,8 @@ def echo_headers(request: WebSocketRequest):
 
     client = websocket.WebSocket()
     client.connect(
-        server.url.replace("http://", "ws://"), header=["Authorization: Basic let-me-in"]
+        server.url.replace("http://", "ws://"),
+        header=["Authorization: Basic let-me-in", "CasedHeader: hello"],
     )
 
     assert client.handshake_response.status == 101
@@ -101,6 +102,7 @@ def echo_headers(request: WebSocketRequest):
     headers = json.loads(doc)
     assert headers["Connection"] == "Upgrade"
     assert headers["Authorization"] == "Basic let-me-in"
+    assert headers["CasedHeader"] == "hello"
 
 
 def test_websocket_reject(serve_websocket_listener):
@@ -174,11 +176,15 @@ def _handler(request: WebSocketRequest, request_args: dict):
         with request.accept() as ws:
             ws.send("foo")
             ws.send(f"id={request_args['id']}")
+            ws.send(json.dumps(dict(request.headers)))
 
     router.add("/foo/<id>", _handler)
 
     server = serve_websocket_listener(WebSocketRequest.listener(router.dispatch))
     client = websocket.WebSocket()
-    client.connect(server.url.replace("http://", "ws://") + "/foo/bar")
+    client.connect(
+        server.url.replace("http://", "ws://") + "/foo/bar", header=["CasedHeader: hello"]
+    )
     assert client.recv() == "foo"
     assert client.recv() == "id=bar"
+    assert "CasedHeader" in json.loads(client.recv())