coverage++

Author: Kriechi

src/wsproto/__init__.py

@@ -36,10 +36,9 @@ def state(self) -> ConnectionState:
         :returns: Connection state
         :rtype: wsproto.connection.ConnectionState
         """
-        if self.connection is None:  # noqa
+        if self.connection is None:
             return self.handshake.state
-        else:
-            return self.connection.state
+        return self.connection.state
 
     def initiate_upgrade_connection(self, headers: Headers, path: str) -> None:
         self.handshake.initiate_upgrade_connection(headers, path)

src/wsproto/connection.py

@@ -132,6 +132,8 @@ def receive_data(self, data: Optional[bytes]) -> None:
             self._proto.receive_bytes(data)
         elif self.state is ConnectionState.CLOSED:
             raise LocalProtocolError("Connection already closed.")
+        else:
+            pass  # pragma: no cover
 
     def events(self) -> Generator[Event, None, None]:
         """
@@ -179,5 +181,8 @@ def events(self) -> Generator[Event, None, None]:
                         frame_finished=frame.frame_finished,
                         message_finished=frame.message_finished,
                     )
+
+                else:
+                    pass  # pragma: no cover
         except ParseFailed as exc:
             yield CloseConnection(code=exc.code, reason=str(exc))

src/wsproto/frame_protocol.py

@@ -12,7 +12,7 @@
 from typing import Generator, List, NamedTuple, Optional, Tuple, TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
-    from .extensions import Extension  # noqa
+    from .extensions import Extension  # pragma: no cover
 
 
 _XOR_TABLE = [bytes(a ^ b for a in range(256)) for b in range(256)]
@@ -23,24 +23,23 @@ def __init__(self, masking_key: bytes) -> None:
         self._masking_key = masking_key
 
     def process(self, data: bytes) -> bytes:
-        if data:  # pylint:disable=no-else-return
+        if data:
             data_array = bytearray(data)
             a, b, c, d = (_XOR_TABLE[n] for n in self._masking_key)
             data_array[::4] = data_array[::4].translate(a)
             data_array[1::4] = data_array[1::4].translate(b)
             data_array[2::4] = data_array[2::4].translate(c)
             data_array[3::4] = data_array[3::4].translate(d)
 
-            # Rotate the maksing key so that the next usage continues
+            # Rotate the masking key so that the next usage continues
             # with the next key element, rather than restarting.
             key_rotation = len(data) % 4
             self._masking_key = (
                 self._masking_key[key_rotation:] + self._masking_key[:key_rotation]
             )
 
             return bytes(data_array)
-        else:
-            return data
+        return data
 
 
 class XorMaskerNull:
@@ -75,7 +74,7 @@ class Opcode(IntEnum):
     RFC 6455, Section 5.2 - Base Framing Protocol
     """
 
-    #: Contiuation frame
+    #: Continuation frame
     CONTINUATION = 0x0
 
     #: Text message

src/wsproto/handshake.py

@@ -224,7 +224,7 @@ def _process_connection_request(  # noqa: MC0001
                 "Missing header, 'Sec-WebSocket-Version'",
                 event_hint=RejectConnection(
                     headers=[(b"Sec-WebSocket-Version", WEBSOCKET_VERSION)],
-                    status_code=426,
+                    status_code=426 if version else 400,
                 ),
             )
         if key is None:
@@ -235,10 +235,6 @@ def _process_connection_request(  # noqa: MC0001
             raise RemoteProtocolError(
                 "Missing header, 'Upgrade: WebSocket'", event_hint=RejectConnection()
             )
-        if version is None:
-            raise RemoteProtocolError(
-                "Missing header, 'Sec-WebSocket-Version'", event_hint=RejectConnection()
-            )
         if host is None:
             raise RemoteProtocolError(
                 "Missing header, 'Host'", event_hint=RejectConnection()

test/test_client.py

@@ -1,12 +1,10 @@
-# These tests test the behaviours expected of wsproto in when the
-# connection is a client.
 from typing import List, Optional
 
 import h11
 import pytest
 
 from wsproto import WSConnection
-from wsproto.connection import CLIENT
+from wsproto.connection import CLIENT, ConnectionState
 from wsproto.events import (
     AcceptConnection,
     Event,
@@ -18,6 +16,7 @@
 from wsproto.typing import Headers
 from wsproto.utilities import (
     generate_accept_token,
+    LocalProtocolError,
     normed_header_dict,
     RemoteProtocolError,
 )
@@ -102,6 +101,42 @@ def test_connection_request_subprotocols() -> None:
     assert headers[b"sec-websocket-protocol"] == b"one, two"
 
 
+def test_connection_send_state() -> None:
+    client = WSConnection(CLIENT)
+    assert client.state is ConnectionState.CONNECTING
+
+    server = h11.Connection(h11.SERVER)
+    server.receive_data(
+        client.send(
+            Request(
+                host="localhost",
+                target="/",
+            )
+        )
+    )
+    headers = normed_header_dict(server.next_event().headers)
+    response = h11.InformationalResponse(
+        status_code=101,
+        headers=[
+            (b"connection", b"Upgrade"),
+            (b"upgrade", b"WebSocket"),
+            (
+                b"Sec-WebSocket-Accept",
+                generate_accept_token(headers[b"sec-websocket-key"]),
+            ),
+        ],
+    )
+    client.receive_data(server.send(response))
+    assert len(list(client.events())) == 1
+    assert client.state is ConnectionState.OPEN  # type: ignore # https://github.com/python/mypy/issues/9005
+
+    with pytest.raises(LocalProtocolError) as excinfo:
+        client.send(Request(host="localhost", target="/"))
+
+    client.receive_data(b"foobar")
+    assert len(list(client.events())) == 1
+
+
 def _make_handshake(
     response_status: int,
     response_headers: Headers,
@@ -110,6 +145,8 @@ def _make_handshake(
     auto_accept_key: bool = True,
 ) -> List[Event]:
     client = WSConnection(CLIENT)
+    assert client.state is ConnectionState.CONNECTING
+
     server = h11.Connection(h11.SERVER)
     server.receive_data(
         client.send(
@@ -134,6 +171,7 @@ def _make_handshake(
         status_code=response_status, headers=response_headers
     )
     client.receive_data(server.send(response))
+    assert client.state is not ConnectionState.CONNECTING
 
     return list(client.events())
 

test/test_connection.py

@@ -68,6 +68,9 @@ def test_closure(client_sends: bool, code: CloseReason, reason: str) -> None:
     assert remote.state is ConnectionState.CLOSED  # type: ignore[comparison-overlap]
     assert local.state is ConnectionState.CLOSED
 
+    with pytest.raises(LocalProtocolError):
+        local.receive_data(b"foobar")
+
 
 def test_abnormal_closure() -> None:
     client = Connection(CLIENT)

test/test_frame_protocol.py

@@ -1193,6 +1193,7 @@ def test_data_we_have_no_idea_what_to_do_with(self) -> None:
 
 def test_xor_mask_simple() -> None:
     masker = fp.XorMaskerSimple(b"1234")
+    assert masker.process(b"") == b""
     assert masker.process(b"some very long data for masking by websocket") == (
         b"B]^Q\x11DVFH\x12_[_U\x13PPFR\x14W]A\x14\\S@_X\\T\x14SK\x13CTP@[RYV@"
     )

test/test_handshake.py

@@ -3,7 +3,7 @@
 from wsproto.connection import CLIENT, ConnectionState, SERVER
 from wsproto.events import AcceptConnection, Ping, RejectConnection, Request
 from wsproto.handshake import H11Handshake
-from wsproto.utilities import LocalProtocolError
+from wsproto.utilities import LocalProtocolError, RemoteProtocolError
 
 
 def test_successful_handshake() -> None:
@@ -19,19 +19,22 @@ def test_successful_handshake() -> None:
     assert client.state is ConnectionState.OPEN
     assert server.state is ConnectionState.OPEN
 
+    assert repr(client) == "H11Handshake(client=True, state=ConnectionState.OPEN)"
+    assert repr(server) == "H11Handshake(client=False, state=ConnectionState.OPEN)"
 
-def test_rejected_handshake() -> None:
-    client = H11Handshake(CLIENT)
-    server = H11Handshake(SERVER)
 
-    server.receive_data(client.send(Request(host="localhost", target="/")))
-    assert isinstance(next(server.events()), Request)
-
-    client.receive_data(server.send(RejectConnection()))
-    assert isinstance(next(client.events()), RejectConnection)
-
-    assert client.state is ConnectionState.CLOSED
-    assert server.state is ConnectionState.CLOSED
+@pytest.mark.parametrize("http", [b"HTTP/1.0", b"HTTP/1.1"])
+def test_rejected_handshake(http: bytes) -> None:
+    server = H11Handshake(SERVER)
+    with pytest.raises(RemoteProtocolError):
+        server.receive_data(
+            b"GET / " + http + b"\r\n"
+            b"Upgrade: WebSocket\r\n"
+            b"Connection: Upgrade\r\n"
+            b"Sec-WebSocket-Key: VQr8cvwwZ1fEk62PDq8J3A==\r\n"
+            b"Sec-WebSocket-Version: 13\r\n"
+            b"\r\n"
+        )
 
 
 def test_initiate_upgrade_as_client() -> None:

test/test_server.py

@@ -1,5 +1,3 @@
-# These tests test the behaviours expected of wsproto in when the
-# connection is a server.
 from typing import cast, List, Optional, Tuple
 
 import h11