Repository: aaugustin/websockets
Pull request #1274 (2022-12-10 20:15:36)

Refactor create_protocol to represent the real type
`[Any]` means one argument of `Any` type, `...` means "ignore". This is relevant since `WebSocketServerProtocol` and similar in fact have multiple arguments. Ideally a `Protocol` with a `__call__` method included should be used, but that's up for discussion since that would add maintainability issues as the signature would be duplicated in multiple places.

**Or**, `Type[WebSocketServerProtocol]` *may* work too.


Patch:
diff --git a/SECURITY.md b/SECURITY.md
index 82024b4..175b20c 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -1,5 +1,12 @@
-# Security policy
+# Security
+
+## Policy
 
 Only the latest version receives security updates.
 
-Please report vulnerabilities [via Tidelift](https://tidelift.com/security).
+## Contact information
+
+Please report security vulnerabilities to the
+[Tidelift security team](https://tidelift.com/security).
+
+Tidelift will coordinate the fix and disclosure.
diff --git a/docs/faq/common.rst b/docs/faq/common.rst
index dff64f6..66e2730 100644
--- a/docs/faq/common.rst
+++ b/docs/faq/common.rst
@@ -125,11 +125,11 @@ You can bind additional arguments to the protocol factory with
     import websockets
 
     class MyServerProtocol(websockets.WebSocketServerProtocol):
-        def __init__(self, extra_argument, *args, **kwargs):
+        def __init__(self, *args, extra_argument=None, **kwargs):
             super().__init__(*args, **kwargs)
             # do something with extra_argument
 
-    create_protocol = functools.partial(MyServerProtocol, extra_argument='spam')
+    create_protocol = functools.partial(MyServerProtocol, extra_argument=42)
     start_server = websockets.serve(..., create_protocol=create_protocol)
 
 This example was for a server. The same pattern applies on a client.
diff --git a/docs/faq/server.rst b/docs/faq/server.rst
index 4e4622d..02a248c 100644
--- a/docs/faq/server.rst
+++ b/docs/faq/server.rst
@@ -185,7 +185,7 @@ You can bind additional arguments to the connection handler with
     async def handler(websocket, extra_argument):
         ...
 
-    bound_handler = functools.partial(handler, extra_argument='spam')
+    bound_handler = functools.partial(handler, extra_argument=42)
     start_server = websockets.serve(bound_handler, ...)
 
 Another way to achieve this result is to define the ``handler`` coroutine in
diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst
index aab9afd..34e1f90 100644
--- a/docs/project/changelog.rst
+++ b/docs/project/changelog.rst
@@ -37,7 +37,7 @@ Backwards-incompatible changes
     :class: caution
 
     Aliases provide compatibility for all previously public APIs according to
-    the `backwards-compatibility policy`_
+    the `backwards-compatibility policy`_.
 
     * The ``connection`` module was renamed to ``protocol``.
 
@@ -67,6 +67,9 @@ New features
 
 * Made it possible to close a server without closing existing connections.
 
+* Added :attr:`~protocol.ServerProtocol.select_subprotocol` to customize
+  negotiation of subprotocols in the Sans-I/O layer.
+
 10.4
 ----
 
diff --git a/docs/reference/server.rst b/docs/reference/server.rst
index 50ef4ee..08bfe2f 100644
--- a/docs/reference/server.rst
+++ b/docs/reference/server.rst
@@ -123,6 +123,8 @@ Sans-I/O
 
     .. automethod:: accept
 
+    .. automethod:: select_subprotocol
+
     .. automethod:: reject
 
     .. automethod:: send_response
diff --git a/src/websockets/client.py b/src/websockets/client.py
index ad4c1a1..bfc8080 100644
--- a/src/websockets/client.py
+++ b/src/websockets/client.py
@@ -316,11 +316,17 @@ class ClientProtocol(Protocol):
 
     def parse(self) -> Generator[None, None, None]:
         if self.state is CONNECTING:
-            response = yield from Response.parse(
-                self.reader.read_line,
-                self.reader.read_exact,
-                self.reader.read_to_eof,
-            )
+            try:
+                response = yield from Response.parse(
+                    self.reader.read_line,
+                    self.reader.read_exact,
+                    self.reader.read_to_eof,
+                )
+            except Exception as exc:
+                self.handshake_exc = exc
+                self.parser = self.discard()
+                next(self.parser)  # start coroutine
+                yield
 
             if self.debug:
                 code, phrase = response.status_code, response.reason_phrase
@@ -334,14 +340,15 @@ class ClientProtocol(Protocol):
                 self.process_response(response)
             except InvalidHandshake as exc:
                 response._exception = exc
+                self.events.append(response)
                 self.handshake_exc = exc
                 self.parser = self.discard()
                 next(self.parser)  # start coroutine
-            else:
-                assert self.state is CONNECTING
-                self.state = OPEN
-            finally:
-                self.events.append(response)
+                yield
+
+            assert self.state is CONNECTING
+            self.state = OPEN
+            self.events.append(response)
 
         yield from super().parse()
 
diff --git a/src/websockets/legacy/auth.py b/src/websockets/legacy/auth.py
index ac24c17..def36a3 100644
--- a/src/websockets/legacy/auth.py
+++ b/src/websockets/legacy/auth.py
@@ -118,7 +118,7 @@ def basic_auth_protocol_factory(
     realm: Optional[str] = None,
     credentials: Optional[Union[Credentials, Iterable[Credentials]]] = None,
     check_credentials: Optional[Callable[[str, str], Awaitable[bool]]] = None,
-    create_protocol: Optional[Callable[[Any], BasicAuthWebSocketServerProtocol]] = None,
+    create_protocol: Optional[Callable[..., BasicAuthWebSocketServerProtocol]] = None,
 ) -> Callable[[Any], BasicAuthWebSocketServerProtocol]:
     """
     Protocol factory that enforces HTTP Basic Auth.
diff --git a/src/websockets/legacy/client.py b/src/websockets/legacy/client.py
index 9b953df..4981ac9 100644
--- a/src/websockets/legacy/client.py
+++ b/src/websockets/legacy/client.py
@@ -438,7 +438,7 @@ class Connect:
         self,
         uri: str,
         *,
-        create_protocol: Optional[Callable[[Any], WebSocketClientProtocol]] = None,
+        create_protocol: Optional[Callable[..., WebSocketClientProtocol]] = None,
         logger: Optional[LoggerLike] = None,
         compression: Optional[str] = "deflate",
         origin: Optional[Origin] = None,
diff --git a/src/websockets/legacy/http.py b/src/websockets/legacy/http.py
index d9e44cc..7cc3db8 100644
--- a/src/websockets/legacy/http.py
+++ b/src/websockets/legacy/http.py
@@ -10,8 +10,8 @@ from ..exceptions import SecurityError
 
 __all__ = ["read_request", "read_response"]
 
-MAX_HEADERS = 256
-MAX_LINE = 4110
+MAX_HEADERS = 128
+MAX_LINE = 8192
 
 
 def d(value: bytes) -> str:
diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py
index eabeb8e..6f8833a 100644
--- a/src/websockets/legacy/server.py
+++ b/src/websockets/legacy/server.py
@@ -520,31 +520,29 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
         server_subprotocols: Sequence[Subprotocol],
     ) -> Optional[Subprotocol]:
         """
-        Pick a subprotocol among those offered by the client.
+        Pick a subprotocol among those supported by the client and the server.
 
-        If several subprotocols are supported by the client and the server,
-        the default implementation selects the preferred subprotocol by
-        giving equal value to the priorities of the client and the server.
-        If no subprotocol is supported by the client and the server, it
-        proceeds without a subprotocol.
+        If several subprotocols are available, select the preferred subprotocol
+        by giving equal weight to the preferences of the client and the server.
 
-        This is unlikely to be the most useful implementation in practice.
-        Many servers providing a subprotocol will require that the client
-        uses that subprotocol. Such rules can be implemented in a subclass.
+        If no subprotocol is available, proceed without a subprotocol.
 
-        You may also override this method with the ``select_subprotocol``
-        argument of :func:`serve` and :class:`WebSocketServerProtocol`.
+        You may provide a ``select_subprotocol`` argument to :func:`serve` or
+        :class:`WebSocketServerProtocol` to override this logic. For example,
+        you could reject the handshake if the client doesn't support a
+        particular subprotocol, rather than accept the handshake without that
+        subprotocol.
 
         Args:
             client_subprotocols: list of subprotocols offered by the client.
             server_subprotocols: list of subprotocols available on the server.
 
         Returns:
-            Optional[Subprotocol]: Selected subprotocol.
+            Optional[Subprotocol]: Selected subprotocol, if a common subprotocol
+            was found.
 
             :obj:`None` to continue without a subprotocol.
 
-
         """
         if self._select_subprotocol is not None:
             return self._select_subprotocol(client_subprotocols, server_subprotocols)
@@ -991,7 +989,7 @@ class Serve:
         host: Optional[Union[str, Sequence[str]]] = None,
         port: Optional[int] = None,
         *,
-        create_protocol: Optional[Callable[[Any], WebSocketServerProtocol]] = None,
+        create_protocol: Optional[Callable[..., WebSocketServerProtocol]] = None,
         logger: Optional[LoggerLike] = None,
         compression: Optional[str] = "deflate",
         origins: Optional[Sequence[Optional[Origin]]] = None,
diff --git a/src/websockets/server.py b/src/websockets/server.py
index cc5ba79..214b38b 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -5,7 +5,7 @@ import binascii
 import email.utils
 import http
 import warnings
-from typing import Any, Generator, List, Optional, Sequence, Tuple, cast
+from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, cast
 
 from .datastructures import Headers, MultipleValuesError
 from .exceptions import (
@@ -58,6 +58,10 @@ class ServerProtocol(Protocol):
             should be tried.
         subprotocols: list of supported subprotocols, in order of decreasing
             preference.
+        select_subprotocol: callback for selecting a subprotocol among
+            those supported by the client and the server. It has the same
+            signature as the :meth:`select_subprotocol` method, including a
+            :class:`ServerProtocol` instance as first argument.
         state: initial state of the WebSocket connection.
         max_size: maximum size of incoming messages in bytes;
             :obj:`None` to disable the limit.
@@ -73,6 +77,7 @@ class ServerProtocol(Protocol):
         origins: Optional[Sequence[Optional[Origin]]] = None,
         extensions: Optional[Sequence[ServerExtensionFactory]] = None,
         subprotocols: Optional[Sequence[Subprotocol]] = None,
+        select_subprotocol: Optional[SelectSubprotocol] = None,
         state: State = CONNECTING,
         max_size: Optional[int] = 2**20,
         logger: Optional[LoggerLike] = None,
@@ -86,6 +91,14 @@ class ServerProtocol(Protocol):
         self.origins = origins
         self.available_extensions = extensions
         self.available_subprotocols = subprotocols
+        if select_subprotocol is not None:
+            # Bind select_subprotocol then shadow self.select_subprotocol.
+            # Use setattr to work around https://github.com/python/mypy/issues/2427.
+            setattr(
+                self,
+                "select_subprotocol",
+                select_subprotocol.__get__(self, self.__class__),
+            )
 
     def accept(self, request: Request) -> Response:
         """
@@ -96,7 +109,7 @@ class ServerProtocol(Protocol):
 
         You must send the handshake response with :meth:`send_response`.
 
-        You can modify it before sending it, for example to add HTTP headers.
+        You may modify it before sending it, for example to add HTTP headers.
 
         Args:
             request: WebSocket handshake request event received from the client.
@@ -175,7 +188,8 @@ class ServerProtocol(Protocol):
         return Response(101, "Switching Protocols", headers)
 
     def process_request(
-        self, request: Request
+        self,
+        request: Request,
     ) -> Tuple[str, Optional[str], Optional[str]]:
         """
         Check a handshake request and negotiate extensions and subprotocol.
@@ -273,6 +287,7 @@ class ServerProtocol(Protocol):
            Optional[Origin]: origin, if it is acceptable.
 
         Raises:
+            InvalidHandshake: if the Origin header is invalid.
             InvalidOrigin: if the origin isn't acceptable.
 
         """
@@ -323,7 +338,7 @@ class ServerProtocol(Protocol):
             HTTP response header and list of accepted extensions.
 
         Raises:
-            InvalidHandshake: to abort the handshake with an HTTP 400 error.
+            InvalidHandshake: if the Sec-WebSocket-Extensions header is invalid.
 
         """
         response_header_value: Optional[str] = None
@@ -383,60 +398,79 @@ class ServerProtocol(Protocol):
            also the value of the ``Sec-WebSocket-Protocol`` response header.
 
         Raises:
-            InvalidHandshake: to abort the handshake with an HTTP 400 error.
+            InvalidHandshake: if the Sec-WebSocket-Subprotocol header is invalid.
 
         """
-        subprotocol: Optional[Subprotocol] = None
-
-        header_values = headers.get_all("Sec-WebSocket-Protocol")
-
-        if header_values and self.available_subprotocols:
-
-            parsed_header_values: List[Subprotocol] = sum(
-                [parse_subprotocol(header_value) for header_value in header_values], []
-            )
-
-            subprotocol = self.select_subprotocol(
-                parsed_header_values, self.available_subprotocols
-            )
+        subprotocols: Sequence[Subprotocol] = sum(
+            [
+                parse_subprotocol(header_value)
+                for header_value in headers.get_all("Sec-WebSocket-Protocol")
+            ],
+            [],
+        )
 
-        return subprotocol
+        return self.select_subprotocol(subprotocols)
 
     def select_subprotocol(
         self,
-        client_subprotocols: Sequence[Subprotocol],
-        server_subprotocols: Sequence[Subprotocol],
+        subprotocols: Sequence[Subprotocol],
     ) -> Optional[Subprotocol]:
         """
         Pick a subprotocol among those offered by the client.
 
-        If several subprotocols are supported by the client and the server,
-        the default implementation selects the preferred subprotocols by
-        giving equal value to the priorities of the client and the server.
+        If several subprotocols are supported by both the client and the server,
+        pick the first one in the list declared the server.
+
+        If the server doesn't support any subprotocols, continue without a
+        subprotocol, regardless of what the client offers.
 
-        If no common subprotocol is supported by the client and the server, it
-        proceeds without a subprotocol.
+        If the server supports at least one subprotocol and the client doesn't
+        offer any, abort the handshake with an HTTP 400 error.
 
-        This is unlikely to be the most useful implementation in practice, as
-        many servers providing a subprotocol will require that the client uses
-        that subprotocol.
+        You provide a ``select_subprotocol`` argument to :class:`ServerProtocol`
+        to override this logic. For example, you could accept the connection
+        even if client doesn't offer a subprotocol, rather than reject it.
+
+        Here's how to negotiate the ``chat`` subprotocol if the client supports
+        it and continue without a subprotocol otherwise::
+
+            def select_subprotocol(protocol, subprotocols):
+                if "chat" in subprotocols:
+                    return "chat"
 
         Args:
-            client_subprotocols: list of subprotocols offered by the client.
-            server_subprotocols: list of subprotocols available on the server.
+            subprotocols: list of subprotocols offered by the client.
 
         Returns:
-            Optional[Subprotocol]: Subprotocol, if a common subprotocol was
-            found.
+            Optional[Subprotocol]: Selected subprotocol, if a common subprotocol
+            was found.
+
+            :obj:`None` to continue without a subprotocol.
+
+        Raises:
+            NegotiationError: custom implementations may raise this exception
+                to abort the handshake with an HTTP 400 error.
 
         """
-        subprotocols = set(client_subprotocols) & set(server_subprotocols)
-        if not subprotocols:
+        # Server doesn't offer any subprotocols.
+        if not self.available_subprotocols:  # None or empty list
             return None
-        priority = lambda p: (
-            client_subprotocols.index(p) + server_subprotocols.index(p)
+
+        # Server offers at least one subprotocol but client doesn't offer any.
+        if not subprotocols:
+            raise NegotiationError("missing subprotocol")
+
+        # Server and client both offer subprotocols. Look for a shared one.
+        proposed_subprotocols = set(subprotocols)
+        for subprotocol in self.available_subprotocols:
+            if subprotocol in proposed_subprotocols:
+                return subprotocol
+
+        # No common subprotocol was found.
+        raise NegotiationError(
+            "invalid subprotocol; expected one of "
+            + ", ".join(self.available_subprotocols)
         )
-        return sorted(subprotocols, key=priority)[0]
 
     def reject(
         self,
@@ -507,7 +541,16 @@ class ServerProtocol(Protocol):
 
     def parse(self) -> Generator[None, None, None]:
         if self.state is CONNECTING:
-            request = yield from Request.parse(self.reader.read_line)
+            try:
+                request = yield from Request.parse(
+                    self.reader.read_line,
+                )
+            except Exception as exc:
+                self.handshake_exc = exc
+                self.send_eof()
+                self.parser = self.discard()
+                next(self.parser)  # start coroutine
+                yield
 
             if self.debug:
                 self.logger.debug("< GET %s HTTP/1.1", request.path)
@@ -519,6 +562,12 @@ class ServerProtocol(Protocol):
         yield from super().parse()
 
 
+SelectSubprotocol = Callable[
+    [ServerProtocol, Sequence[Subprotocol]],
+    Optional[Subprotocol],
+]
+
+
 class ServerConnection(ServerProtocol):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         warnings.warn(


Test patch:
diff --git a/tests/legacy/test_http.py b/tests/legacy/test_http.py
index 5c9adc9..15d53e0 100644
--- a/tests/legacy/test_http.py
+++ b/tests/legacy/test_http.py
@@ -119,13 +119,13 @@ class HTTPAsyncTests(AsyncioTestCase):
             await read_headers(self.stream)
 
     async def test_headers_limit(self):
-        self.stream.feed_data(b"foo: bar\r\n" * 257 + b"\r\n")
+        self.stream.feed_data(b"foo: bar\r\n" * 129 + b"\r\n")
         with self.assertRaises(SecurityError):
             await read_headers(self.stream)
 
     async def test_line_limit(self):
-        # Header line contains 5 + 4104 + 2 = 4111 bytes.
-        self.stream.feed_data(b"foo: " + b"a" * 4104 + b"\r\n\r\n")
+        # Header line contains 5 + 8186 + 2 = 8193 bytes.
+        self.stream.feed_data(b"foo: " + b"a" * 8186 + b"\r\n\r\n")
         with self.assertRaises(SecurityError):
             await read_headers(self.stream)
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 718219b..c83c870 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -218,6 +218,32 @@ class AcceptRejectTests(unittest.TestCase):
         )
         self.assertEqual(response.body, b"Sorry folks.\n")
 
+    def test_no_response(self):
+        with unittest.mock.patch("websockets.client.generate_key", return_value=KEY):
+            client = ClientProtocol(parse_uri("ws://example.com/test"))
+        client.connect()
+        client.receive_eof()
+        self.assertEqual(client.events_received(), [])
+
+    def test_partial_response(self):
+        with unittest.mock.patch("websockets.client.generate_key", return_value=KEY):
+            client = ClientProtocol(parse_uri("ws://example.com/test"))
+        client.connect()
+        client.receive_data(b"HTTP/1.1 101 Switching Protocols\r\n")
+        client.receive_eof()
+        self.assertEqual(client.events_received(), [])
+
+    def test_random_response(self):
+        with unittest.mock.patch("websockets.client.generate_key", return_value=KEY):
+            client = ClientProtocol(parse_uri("ws://example.com/test"))
+        client.connect()
+        client.receive_data(b"220 smtp.invalid\r\n")
+        client.receive_data(b"250 Hello relay.invalid\r\n")
+        client.receive_data(b"250 Ok\r\n")
+        client.receive_data(b"250 Ok\r\n")
+        client.receive_eof()
+        self.assertEqual(client.events_received(), [])
+
     def make_accept_response(self, client):
         request = client.connect()
         return Response(
diff --git a/tests/test_server.py b/tests/test_server.py
index 2645dc8..ecf3d4c 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -4,7 +4,12 @@ import unittest
 import unittest.mock
 
 from websockets.datastructures import Headers
-from websockets.exceptions import InvalidHeader, InvalidOrigin, InvalidUpgrade
+from websockets.exceptions import (
+    InvalidHeader,
+    InvalidOrigin,
+    InvalidUpgrade,
+    NegotiationError,
+)
 from websockets.frames import OP_TEXT, Frame
 from websockets.http11 import Request, Response
 from websockets.protocol import CONNECTING, OPEN
@@ -67,6 +72,24 @@ class ConnectTests(unittest.TestCase):
             ),
         )
 
+    def test_no_request(self):
+        server = ServerProtocol()
+        server.receive_eof()
+        self.assertEqual(server.events_received(), [])
+
+    def test_partial_request(self):
+        server = ServerProtocol()
+        server.receive_data(b"GET /test HTTP/1.1\r\n")
+        server.receive_eof()
+        self.assertEqual(server.events_received(), [])
+
+    def test_random_request(self):
+        server = ServerProtocol()
+        server.receive_data(b"HELO relay.invalid\r\n")
+        server.receive_data(b"MAIL FROM: <alice@invalid>\r\n")
+        server.receive_data(b"RCPT TO: <bob@invalid>\r\n")
+        self.assertEqual(server.events_received(), [])
+
 
 class AcceptRejectTests(unittest.TestCase):
     def make_request(self):
@@ -544,9 +567,12 @@ class AcceptRejectTests(unittest.TestCase):
         request = self.make_request()
         response = server.accept(request)
 
-        self.assertEqual(response.status_code, 101)
-        self.assertNotIn("Sec-WebSocket-Protocol", response.headers)
-        self.assertIsNone(server.subprotocol)
+        self.assertEqual(response.status_code, 400)
+        with self.assertRaisesRegex(
+            NegotiationError,
+            r"missing subprotocol",
+        ):
+            raise server.handshake_exc
 
     def test_subprotocol(self):
         server = ServerProtocol(subprotocols=["chat"])
@@ -571,8 +597,8 @@ class AcceptRejectTests(unittest.TestCase):
     def test_multiple_subprotocols(self):
         server = ServerProtocol(subprotocols=["superchat", "chat"])
         request = self.make_request()
-        request.headers["Sec-WebSocket-Protocol"] = "superchat"
         request.headers["Sec-WebSocket-Protocol"] = "chat"
+        request.headers["Sec-WebSocket-Protocol"] = "superchat"
         response = server.accept(request)
 
         self.assertEqual(response.status_code, 101)
@@ -595,6 +621,34 @@ class AcceptRejectTests(unittest.TestCase):
         request.headers["Sec-WebSocket-Protocol"] = "otherchat"
         response = server.accept(request)
 
+        self.assertEqual(response.status_code, 400)
+        with self.assertRaisesRegex(
+            NegotiationError,
+            r"invalid subprotocol; expected one of superchat, chat",
+        ):
+            raise server.handshake_exc
+
+    @staticmethod
+    def optional_chat(protocol, subprotocols):
+        if "chat" in subprotocols:
+            return "chat"
+
+    def test_select_subprotocol(self):
+        server = ServerProtocol(select_subprotocol=self.optional_chat)
+        request = self.make_request()
+        request.headers["Sec-WebSocket-Protocol"] = "chat"
+        response = server.accept(request)
+
+        self.assertEqual(response.status_code, 101)
+        self.assertEqual(response.headers["Sec-WebSocket-Protocol"], "chat")
+        self.assertEqual(server.subprotocol, "chat")
+
+    def test_select_no_subprotocol(self):
+        server = ServerProtocol(select_subprotocol=self.optional_chat)
+        request = self.make_request()
+        request.headers["Sec-WebSocket-Protocol"] = "otherchat"
+        response = server.accept(request)
+
         self.assertEqual(response.status_code, 101)
         self.assertNotIn("Sec-WebSocket-Protocol", response.headers)
         self.assertIsNone(server.subprotocol)