Improve exception handling during handshake.

Also refactor tests for Sans-I/O client and server.
python-websockets · Sep 21, 2024 · 20739e0 · 20739e0
1 parent 206624a
commit 20739e0
Show file tree

Hide file tree

Showing 5 changed files with 799 additions and 616 deletions.
diff --git a/src/websockets/client.py b/src/websockets/client.py
@@ -175,10 +175,10 @@ def process_response(self, response: Response) -> None:
 
         try:
             s_w_accept = headers["Sec-WebSocket-Accept"]
-        except KeyError as exc:
-            raise InvalidHeader("Sec-WebSocket-Accept") from exc
-        except MultipleValuesError as exc:
-            raise InvalidHeader("Sec-WebSocket-Accept", "multiple values") from exc
+        except KeyError:
+            raise InvalidHeader("Sec-WebSocket-Accept") from None
+        except MultipleValuesError:
+            raise InvalidHeader("Sec-WebSocket-Accept", "multiple values") from None
 
         if s_w_accept != accept_key(self.key):
             raise InvalidHeaderValue("Sec-WebSocket-Accept", s_w_accept)

diff --git a/src/websockets/server.py b/src/websockets/server.py
@@ -253,10 +253,10 @@ def process_request(
 
         try:
             key = headers["Sec-WebSocket-Key"]
-        except KeyError as exc:
-            raise InvalidHeader("Sec-WebSocket-Key") from exc
-        except MultipleValuesError as exc:
-            raise InvalidHeader("Sec-WebSocket-Key", "multiple values") from exc
+        except KeyError:
+            raise InvalidHeader("Sec-WebSocket-Key") from None
+        except MultipleValuesError:
+            raise InvalidHeader("Sec-WebSocket-Key", "multiple values") from None
 
         try:
             raw_key = base64.b64decode(key.encode(), validate=True)
@@ -267,10 +267,10 @@ def process_request(
 
         try:
             version = headers["Sec-WebSocket-Version"]
-        except KeyError as exc:
-            raise InvalidHeader("Sec-WebSocket-Version") from exc
-        except MultipleValuesError as exc:
-            raise InvalidHeader("Sec-WebSocket-Version", "multiple values") from exc
+        except KeyError:
+            raise InvalidHeader("Sec-WebSocket-Version") from None
+        except MultipleValuesError:
+            raise InvalidHeader("Sec-WebSocket-Version", "multiple values") from None
 
         if version != "13":
             raise InvalidHeaderValue("Sec-WebSocket-Version", version)
@@ -308,8 +308,8 @@ def process_origin(self, headers: Headers) -> Origin | None:
         # per https://datatracker.ietf.org/doc/html/rfc6454#section-7.3.
         try:
             origin = headers.get("Origin")
-        except MultipleValuesError as exc:
-            raise InvalidHeader("Origin", "multiple values") from exc
+        except MultipleValuesError:
+            raise InvalidHeader("Origin", "multiple values") from None
         if origin is not None:
             origin = cast(Origin, origin)
         if self.origins is not None:
@@ -503,7 +503,7 @@ def reject(self, status: StatusLike, text: str) -> Response:
             HTTP response to send to the client.
 
         """
-        # If a user passes an int instead of a HTTPStatus, fix it automatically.
+        # If status is an int instead of an HTTPStatus, fix it automatically.
         status = http.HTTPStatus(status)
         body = text.encode()
         headers = Headers(