diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fe46003..8258108 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,8 @@ Release History 1.4.0 (unreleased) ------------------ +- Accept any sequence of ``(name, value)`` header pairs, matching ``h11``. + - Require ``event_hint`` when constructing ``RemoteProtocolError``. This is an API-breaking change. diff --git a/src/wsproto/handshake.py b/src/wsproto/handshake.py index ab390b8..c9c3284 100644 --- a/src/wsproto/handshake.py +++ b/src/wsproto/handshake.py @@ -81,7 +81,7 @@ def initiate_upgrade_connection( This should be used if the request has already be received and parsed. - :param list headers: HTTP headers represented as a list of 2-tuples. + :param headers: HTTP headers as a sequence of ``(name, value)`` pairs. :param str path: A URL path. """ if self.client: @@ -89,7 +89,9 @@ def initiate_upgrade_connection( raise LocalProtocolError( msg, ) - upgrade_request = h11.Request(method=b"GET", target=path, headers=headers) + upgrade_request = h11.Request( + method=b"GET", target=path, headers=list(headers), + ) h11_client = h11.Connection(h11.CLIENT) self.receive_data(h11_client.send(upgrade_request)) @@ -204,7 +206,7 @@ def _process_connection_request( subprotocols: list[str] = [] upgrade = b"" version = None - headers: Headers = [] + headers: list[tuple[bytes, bytes]] = [] for name, value in event.headers: name = name.lower() if name == b"connection": @@ -299,7 +301,7 @@ def _accept(self, event: AcceptConnection) -> bytes: response = h11.InformationalResponse( status_code=101, - headers=headers + event.extra_headers, + headers=headers + list(event.extra_headers), reason=b"Switching Protocols", ) self._connection = Connection( @@ -381,7 +383,7 @@ def _initiate_connection(self, request: Request) -> bytes: upgrade = h11.Request( method=b"GET", target=request.target.encode("ascii"), - headers=headers + request.extra_headers, + headers=headers + list(request.extra_headers), ) return self._h11_connection.send(upgrade) or b"" @@ -397,7 +399,7 @@ def _establish_client_connection( accepts: list[str] = [] subprotocol = None upgrade = b"" - headers: Headers = [] + headers: list[tuple[bytes, bytes]] = [] for name, value in event.headers: name = name.lower() if name == b"connection": diff --git a/src/wsproto/typing.py b/src/wsproto/typing.py index 0063d21..1674b92 100644 --- a/src/wsproto/typing.py +++ b/src/wsproto/typing.py @@ -1,3 +1,3 @@ -from __future__ import annotations +from collections.abc import Sequence -Headers = list[tuple[bytes, bytes]] +Headers = Sequence[tuple[bytes, bytes]] diff --git a/tests/test_client.py b/tests/test_client.py index d1bbdbc..6ab9f1c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -61,6 +61,21 @@ def test_connection_request_additional_headers() -> None: assert headers[b"x-bar"] == b"Foo" +def test_connection_request_tuple_extra_headers() -> None: + # Headers accepts any sequence of pairs, not only list. + request = _make_connection_request( + Request( + host="localhost", + target="/", + extra_headers=((b"X-Foo", b"Bar"), (b"X-Bar", b"Foo")), + ), + ) + + headers = normed_header_dict(request.headers) + assert headers[b"x-foo"] == b"Bar" + assert headers[b"x-bar"] == b"Foo" + + def test_connection_request_simple_extension() -> None: extension = FakeExtension(offer_response=True) request = _make_connection_request( diff --git a/tests/test_server.py b/tests/test_server.py index 10e3e7f..d760fc4 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -131,9 +131,9 @@ def test_connection_request_key_header() -> None: assert str(excinfo.value) == "Missing header, 'Sec-WebSocket-Key'" -def test_upgrade_request() -> None: - server = WSConnection(SERVER) - server.initiate_upgrade_connection( +@pytest.mark.parametrize( + "headers", + [ [ (b"Host", b"localhost"), (b"Connection", b"Keep-Alive, Upgrade"), @@ -142,8 +142,21 @@ def test_upgrade_request() -> None: (b"Sec-WebSocket-Key", generate_nonce()), (b"X-Foo", b"bar"), ], - "/", - ) + # Sequence that is not a list (the point of Headers = Sequence[...]). + ( + (b"Host", b"localhost"), + (b"Connection", b"Keep-Alive, Upgrade"), + (b"Upgrade", b"websocket"), + (b"Sec-WebSocket-Version", b"13"), + (b"Sec-WebSocket-Key", generate_nonce()), + (b"X-Foo", b"bar"), + ), + ], + ids=["list", "tuple"], +) +def test_upgrade_request(headers: Headers) -> None: + server = WSConnection(SERVER) + server.initiate_upgrade_connection(headers, "/") event = next(server.events()) event = cast("Request", event) @@ -222,6 +235,18 @@ def test_handshake_extra_headers() -> None: ] +def test_handshake_tuple_extra_headers() -> None: + response, nonce = _make_handshake([], accept_headers=((b"X-Foo", b"bar"),)) + + assert response.status_code == 101 + assert sorted(response.headers) == [ + (b"connection", b"Upgrade"), + (b"sec-websocket-accept", generate_accept_token(nonce)), + (b"upgrade", b"websocket"), + (b"x-foo", b"bar"), + ] + + @pytest.mark.parametrize("accept_subprotocol", ["one", "two"]) def test_handshake_with_subprotocol(accept_subprotocol: str) -> None: response, _ = _make_handshake(