Run handler only when opening handshake succeeds.

When process_request() or process_response() returned a HTTP response
without calling accept() or reject() and with a status code other than
101, the connection handler used to start, which was incorrect.

Fix #1419.

Also move start_keepalive() outside of handshake() and bring it together
with starting the connection handler, which is more logical.

Author: aaugustin

docs/project/changelog.rst

@@ -68,6 +68,13 @@ Improvements
   Previously, :exc:`RuntimeError` was raised. For backwards compatibility,
   :exc:`~exceptions.ConcurrencyError` is a subclass of :exc:`RuntimeError`.
 
+Bug fixes
+.........
+
+* The new :mod:`asyncio` and :mod:`threading` implementations of servers don't
+  start the connection handler anymore when ``process_request`` or
+  ``process_response`` returns a HTTP response.
+
 13.0.1
 ------
 

src/websockets/asyncio/client.py

@@ -98,9 +98,7 @@ async def handshake(
         # before receiving a response, when the response cannot be parsed, or
         # when the response fails the handshake.
 
-        if self.protocol.handshake_exc is None:
-            self.start_keepalive()
-        else:
+        if self.protocol.handshake_exc is not None:
             raise self.protocol.handshake_exc
 
     def process_event(self, event: Event) -> None:
@@ -465,6 +463,7 @@ async def __await_impl__(self) -> ClientConnection:
                             raise uri_or_exc from exc
 
                     else:
+                        self.connection.start_keepalive()
                         return self.connection
                 else:
                     raise SecurityError(f"more than {MAX_REDIRECTS} redirects")

src/websockets/asyncio/server.py

@@ -201,12 +201,11 @@ async def handshake(
                 self.protocol.send_response(self.response)
 
         # self.protocol.handshake_exc is always set when the connection is lost
-        # before receiving a request, when the request cannot be parsed, or when
-        # the response fails the handshake.
+        # before receiving a request, when the request cannot be parsed, when
+        # the handshake encounters an error, or when process_request or
+        # process_response sends a HTTP response that rejects the handshake.
 
-        if self.protocol.handshake_exc is None:
-            self.start_keepalive()
-        else:
+        if self.protocol.handshake_exc is not None:
             raise self.protocol.handshake_exc
 
     def process_event(self, event: Event) -> None:
@@ -369,7 +368,9 @@ async def conn_handler(self, connection: ServerConnection) -> None:
                     connection.close_transport()
                     return
 
+            assert connection.protocol.state is OPEN
             try:
+                connection.start_keepalive()
                 await self.handler(connection)
             except Exception:
                 connection.logger.error("connection handler failed", exc_info=True)

src/websockets/server.py

@@ -204,7 +204,6 @@ def accept(self, request: Request) -> Response:
         if protocol_header is not None:
             headers["Sec-WebSocket-Protocol"] = protocol_header
 
-        self.logger.info("connection open")
         return Response(101, "Switching Protocols", headers)
 
     def process_request(
@@ -515,14 +514,7 @@ def reject(self, status: StatusLike, text: str) -> Response:
                 ("Content-Type", "text/plain; charset=utf-8"),
             ]
         )
-        response = Response(status.value, status.phrase, headers, body)
-        # When reject() is called from accept(), handshake_exc is already set.
-        # If a user calls reject(), set handshake_exc to guarantee invariant:
-        # "handshake_exc is None if and only if opening handshake succeeded."
-        if self.handshake_exc is None:
-            self.handshake_exc = InvalidStatus(response)
-        self.logger.info("connection rejected (%d %s)", status.value, status.phrase)
-        return response
+        return Response(status.value, status.phrase, headers, body)
 
     def send_response(self, response: Response) -> None:
         """
@@ -545,7 +537,20 @@ def send_response(self, response: Response) -> None:
         if response.status_code == 101:
             assert self.state is CONNECTING
             self.state = OPEN
+            self.logger.info("connection open")
+
         else:
+            # handshake_exc may be already set if accept() encountered an error.
+            # If the connection isn't open, set handshake_exc to guarantee that
+            # handshake_exc is None if and only if opening handshake succeeded.
+            if self.handshake_exc is None:
+                self.handshake_exc = InvalidStatus(response)
+            self.logger.info(
+                "connection rejected (%d %s)",
+                response.status_code,
+                response.reason_phrase,
+            )
+
             self.send_eof()
             self.parser = self.discard()
             next(self.parser)  # start coroutine

src/websockets/sync/server.py

@@ -23,7 +23,7 @@
     validate_subprotocols,
 )
 from ..http11 import SERVER, Request, Response
-from ..protocol import CONNECTING, Event
+from ..protocol import CONNECTING, OPEN, Event
 from ..server import ServerProtocol
 from ..typing import LoggerLike, Origin, StatusLike, Subprotocol
 from .connection import Connection
@@ -166,8 +166,9 @@ def handshake(
                 self.protocol.send_response(self.response)
 
         # self.protocol.handshake_exc is always set when the connection is lost
-        # before receiving a request, when the request cannot be parsed, or when
-        # the response fails the handshake.
+        # before receiving a request, when the request cannot be parsed, when
+        # the handshake encounters an error, or when process_request or
+        # process_response sends a HTTP response that rejects the handshake.
 
         if self.protocol.handshake_exc is not None:
             raise self.protocol.handshake_exc
@@ -569,6 +570,7 @@ def protocol_select_subprotocol(
                 connection.recv_events_thread.join()
                 return
 
+            assert connection.protocol.state is OPEN
             try:
                 handler(connection)
             except Exception:

tests/asyncio/test_server.py

@@ -142,32 +142,46 @@ async def process_request(ws, request):
     async def test_process_request_returns_response(self):
         """Server aborts handshake if process_request returns a response."""
 
+        handler_ran = False
+
         def process_request(ws, request):
             return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden")
 
-        async with serve(*args, process_request=process_request) as server:
+        async def handler(ws):
+            nonlocal handler_ran
+            handler_ran = True
+
+        async with serve(handler, *args[1:], process_request=process_request) as server:
             with self.assertRaises(InvalidStatus) as raised:
                 async with connect(get_uri(server)):
                     self.fail("did not raise")
             self.assertEqual(
                 str(raised.exception),
                 "server rejected WebSocket connection: HTTP 403",
             )
+            self.assertFalse(handler_ran)
 
     async def test_async_process_request_returns_response(self):
         """Server aborts handshake if async process_request returns a response."""
 
+        handler_ran = False
+
         async def process_request(ws, request):
             return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden")
 
-        async with serve(*args, process_request=process_request) as server:
+        async def handler(ws):
+            nonlocal handler_ran
+            handler_ran = True
+
+        async with serve(handler, *args[1:], process_request=process_request) as server:
             with self.assertRaises(InvalidStatus) as raised:
                 async with connect(get_uri(server)):
                     self.fail("did not raise")
             self.assertEqual(
                 str(raised.exception),
                 "server rejected WebSocket connection: HTTP 403",
             )
+            self.assertFalse(handler_ran)
 
     async def test_process_request_raises_exception(self):
         """Server returns an error if process_request raises an exception."""

tests/sync/test_server.py

@@ -130,9 +130,15 @@ def process_request(ws, request):
     def test_process_request_returns_response(self):
         """Server aborts handshake if process_request returns a response."""
 
+        handler_ran = False
+
         def process_request(ws, request):
             return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden")
 
+        async def handler(ws):
+            nonlocal handler_ran
+            handler_ran = True
+
         with run_server(process_request=process_request) as server:
             with self.assertRaises(InvalidStatus) as raised:
                 with connect(get_uri(server)):
@@ -141,6 +147,7 @@ def process_request(ws, request):
                 str(raised.exception),
                 "server rejected WebSocket connection: HTTP 403",
             )
+            self.assertFalse(handler_ran)
 
     def test_process_request_raises_exception(self):
         """Server returns an error if process_request raises an exception."""