Stop if process_request/response fails the handshake.

Fix #1419.

Author: aaugustin

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.