Accept any sequence of header pairs, not just list

CHANGELOG.rst

@@ -4,6 +4,7 @@ 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.
 

src/wsproto/handshake.py

@@ -81,15 +81,22 @@ 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:
             msg = "Cannot initiate an upgrade connection when acting as the client"
             raise LocalProtocolError(
                 msg,
             )
-        upgrade_request = h11.Request(method=b"GET", target=path, headers=headers)
+        # h11 walks any sequence of pairs at runtime, but its stubs only
+        # accept a list, so materialise the sequence here instead of
+        # silencing the type checker with a cast.
+        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 +211,7 @@ def _process_connection_request(
         subprotocols: list[str] = []
         upgrade = b""
         version = None
-        headers: Headers = []
+        collected: list[tuple[bytes, bytes]] = []
         for name, value in event.headers:
             name = name.lower()
             if name == b"connection":
@@ -224,7 +231,8 @@ def _process_connection_request(
                 version = value
             elif name == b"upgrade":
                 upgrade = value
-            headers.append((name, value))
+            collected.append((name, value))
+            headers: Headers = collected
         if connection_tokens is None or not any(
             token.lower() == "upgrade" for token in connection_tokens
         ):
@@ -299,7 +307,7 @@ def _accept(self, event: AcceptConnection) -> bytes:
 
         response = h11.InformationalResponse(
             status_code=101,
-            headers=headers + event.extra_headers,
+            headers=[*headers, *event.extra_headers],
             reason=b"Switching Protocols",
         )
         self._connection = Connection(
@@ -381,7 +389,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, *request.extra_headers],
         )
         return self._h11_connection.send(upgrade) or b""
 
@@ -397,7 +405,7 @@ def _establish_client_connection(
         accepts: list[str] = []
         subprotocol = None
         upgrade = b""
-        headers: Headers = []
+        collected: list[tuple[bytes, bytes]] = []
         for name, value in event.headers:
             name = name.lower()
             if name == b"connection":
@@ -415,7 +423,8 @@ def _establish_client_connection(
             if name == b"upgrade":
                 upgrade = value
                 continue  # Skip appending to headers
-            headers.append((name, value))
+            collected.append((name, value))
+            headers: Headers = collected
 
         if connection_tokens is None or not any(
             token.lower() == "upgrade" for token in connection_tokens

src/wsproto/typing.py

@@ -1,3 +1,5 @@
 from __future__ import annotations
 
-Headers = list[tuple[bytes, bytes]]
+from collections.abc import Sequence
+
+Headers = Sequence[tuple[bytes, bytes]]

tests/test_client.py

@@ -59,7 +59,22 @@ def test_connection_request_additional_headers() -> None:
     headers = normed_header_dict(request.headers)
     assert headers[b"x-foo"] == b"Bar"
     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)
@@ -142,7 +157,7 @@ def test_connection_send_state() -> None:
 
 def _make_handshake(
     response_status: int,
-    response_headers: Headers,
+    response_headers: list[tuple[bytes, bytes]],
     subprotocols: Optional[List[str]] = None,
     extensions: Optional[List[Extension]] = None,
     auto_accept_key: bool = True,
@@ -205,14 +220,14 @@ def test_handshake_extra_accept_headers() -> None:
 @pytest.mark.parametrize("extra_headers", [[], [(b"connection", b"Keep-Alive")]])
 def test_handshake_response_broken_connection_header(extra_headers: Headers) -> None:
     with pytest.raises(RemoteProtocolError) as excinfo:
-        _make_handshake(101, [(b"upgrade", b"websocket")] + extra_headers)
+        _make_handshake(101, [(b"upgrade", b"websocket"), *extra_headers])
     assert str(excinfo.value) == "Missing header, 'Connection: Upgrade'"
 
 
 @pytest.mark.parametrize("extra_headers", [[], [(b"upgrade", b"h2")]])
 def test_handshake_response_broken_upgrade_header(extra_headers: Headers) -> None:
     with pytest.raises(RemoteProtocolError) as excinfo:
-        _make_handshake(101, [(b"connection", b"Upgrade")] + extra_headers)
+        _make_handshake(101, [(b"connection", b"Upgrade"), *extra_headers])
     assert str(excinfo.value) == "Missing header, 'Upgrade: websocket'"
 
 

tests/test_server.py

@@ -24,7 +24,7 @@ def _make_connection_request(request_headers: Headers, method: str = "GET") -> R
     client = h11.Connection(h11.CLIENT)
     server = WSConnection(SERVER)
     server.receive_data(
-        client.send(h11.Request(method=method, target="/", headers=request_headers)),
+        client.send(h11.Request(method=method, target="/", headers=list(request_headers))),
     )
     event = next(server.events())
     assert isinstance(event, Request)
@@ -47,14 +47,14 @@ def test_connection_request() -> None:
     assert event.host == "localhost"
     assert event.subprotocols == []
     assert event.target == "/"
-    headers = normed_header_dict(event.extra_headers)
-    assert b"host" not in headers
-    assert b"sec-websocket-extensions" not in headers
-    assert b"sec-websocket-protocol" not in headers
-    assert headers[b"connection"] == b"Keep-Alive, Upgrade"
-    assert headers[b"sec-websocket-version"] == b"13"
-    assert headers[b"upgrade"] == b"websocket"
-    assert headers[b"x-foo"] == b"bar"
+    header_dict = normed_header_dict(event.extra_headers)
+    assert b"host" not in header_dict
+    assert b"sec-websocket-extensions" not in header_dict
+    assert b"sec-websocket-protocol" not in header_dict
+    assert header_dict[b"connection"] == b"Keep-Alive, Upgrade"
+    assert header_dict[b"sec-websocket-version"] == b"13"
+    assert header_dict[b"upgrade"] == b"websocket"
+    assert header_dict[b"x-foo"] == b"bar"
 
 
 def test_connection_request_bad_method() -> None:
@@ -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,23 +142,36 @@ 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)
 
     assert event.extensions == []
     assert event.host == "localhost"
     assert event.subprotocols == []
     assert event.target == "/"
-    headers = normed_header_dict(event.extra_headers)
-    assert b"host" not in headers
-    assert b"sec-websocket-extensions" not in headers
-    assert b"sec-websocket-protocol" not in headers
-    assert headers[b"connection"] == b"Keep-Alive, Upgrade"
-    assert headers[b"sec-websocket-version"] == b"13"
-    assert headers[b"upgrade"] == b"websocket"
-    assert headers[b"x-foo"] == b"bar"
+    header_dict = normed_header_dict(event.extra_headers)
+    assert b"host" not in header_dict
+    assert b"sec-websocket-extensions" not in header_dict
+    assert b"sec-websocket-protocol" not in header_dict
+    assert header_dict[b"connection"] == b"Keep-Alive, Upgrade"
+    assert header_dict[b"sec-websocket-version"] == b"13"
+    assert header_dict[b"upgrade"] == b"websocket"
+    assert header_dict[b"x-foo"] == b"bar"
 
 
 def _make_handshake(
@@ -181,8 +194,8 @@ def _make_handshake(
                     (b"Upgrade", b"websocket"),
                     (b"Sec-WebSocket-Version", b"13"),
                     (b"Sec-WebSocket-Key", nonce),
+                    *request_headers,
                 ]
-                + request_headers,
             ),
         ),
     )
@@ -219,7 +232,19 @@ def test_handshake_extra_headers() -> None:
         (b"sec-websocket-accept", generate_accept_token(nonce)),
         (b"upgrade", b"websocket"),
         (b"x-foo", b"bar"),
-    ]
+    ]        
+
+
+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"])