Run handler only when opening handshake succeeds.

When process_request or process_response returned a HTTP response with a
status code other than 101, the connection handler used to start, which
was incorrect.

Fix #1419.

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

@@ -202,11 +202,10 @@ async def handshake(
 
         # 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.
+        # the handshake encounters an error. It isn't set 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,13 +368,19 @@ async def conn_handler(self, connection: ServerConnection) -> None:
                     connection.close_transport()
                     return
 
-            try:
-                await self.handler(connection)
-            except Exception:
-                connection.logger.error("connection handler failed", exc_info=True)
-                await connection.close(CloseCode.INTERNAL_ERROR)
+            if connection.protocol.state is OPEN:
+                try:
+                    connection.start_keepalive()
+                    await self.handler(connection)
+                except Exception:
+                    connection.logger.error("connection handler failed", exc_info=True)
+                    await connection.close(CloseCode.INTERNAL_ERROR)
+                else:
+                    await connection.close()
             else:
-                await connection.close()
+                # process_request or process_response sent a non-101 response.
+                connection.close_transport()
+                return
 
         except TimeoutError:
             # When the opening handshake times out, there's nothing to log.

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
@@ -167,7 +167,8 @@ def handshake(
 
         # 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.
+        # the handshake encounters an error. It isn't set 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
@@ -562,20 +563,23 @@ def protocol_select_subprotocol(
             except TimeoutError:
                 connection.close_socket()
                 connection.recv_events_thread.join()
-                return
             except Exception:
                 connection.logger.error("opening handshake failed", exc_info=True)
                 connection.close_socket()
                 connection.recv_events_thread.join()
-                return
 
-            try:
-                handler(connection)
-            except Exception:
-                connection.logger.error("connection handler failed", exc_info=True)
-                connection.close(CloseCode.INTERNAL_ERROR)
+            if connection.protocol.state is OPEN:
+                try:
+                    handler(connection)
+                except Exception:
+                    connection.logger.error("connection handler failed", exc_info=True)
+                    connection.close(CloseCode.INTERNAL_ERROR)
+                else:
+                    connection.close()
             else:
-                connection.close()
+                # process_request or process_response sent a non-101 response.
+                connection.close_socket()
+                connection.recv_events_thread.join()
 
         except Exception:  # pragma: no cover
             # Don't leak sockets on unexpected errors.

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 = True
+
         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 = True
+
         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 = True
+
         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."""