From 222a9d07558dbfa5798d73bfd65f1f3b646e2ee4 Mon Sep 17 00:00:00 2001 From: Ram Rachum Date: Sat, 13 Jun 2020 23:00:34 +0300 Subject: [PATCH 1/2] Fix exception causes in handshake.py --- src/websockets/handshake.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/websockets/handshake.py b/src/websockets/handshake.py index 646b6dba4..0c9cc69e8 100644 --- a/src/websockets/handshake.py +++ b/src/websockets/handshake.py @@ -94,28 +94,28 @@ def check_request(headers: Headers) -> str: try: s_w_key = headers["Sec-WebSocket-Key"] - except KeyError: - raise InvalidHeader("Sec-WebSocket-Key") - except MultipleValuesError: + except KeyError as exc: + raise InvalidHeader("Sec-WebSocket-Key") from exc + except MultipleValuesError as exc: raise InvalidHeader( "Sec-WebSocket-Key", "more than one Sec-WebSocket-Key header found" - ) + ) from exc try: raw_key = base64.b64decode(s_w_key.encode(), validate=True) - except binascii.Error: - raise InvalidHeaderValue("Sec-WebSocket-Key", s_w_key) + except binascii.Error as exc: + raise InvalidHeaderValue("Sec-WebSocket-Key", s_w_key) from exc if len(raw_key) != 16: raise InvalidHeaderValue("Sec-WebSocket-Key", s_w_key) try: s_w_version = headers["Sec-WebSocket-Version"] - except KeyError: - raise InvalidHeader("Sec-WebSocket-Version") - except MultipleValuesError: + except KeyError as exc: + raise InvalidHeader("Sec-WebSocket-Version") from exc + except MultipleValuesError as exc: raise InvalidHeader( "Sec-WebSocket-Version", "more than one Sec-WebSocket-Version header found" - ) + ) from exc if s_w_version != "13": raise InvalidHeaderValue("Sec-WebSocket-Version", s_w_version) @@ -171,12 +171,12 @@ def check_response(headers: Headers, key: str) -> None: try: s_w_accept = headers["Sec-WebSocket-Accept"] - except KeyError: - raise InvalidHeader("Sec-WebSocket-Accept") - except MultipleValuesError: + except KeyError as exc: + raise InvalidHeader("Sec-WebSocket-Accept") from exc + except MultipleValuesError as exc: raise InvalidHeader( "Sec-WebSocket-Accept", "more than one Sec-WebSocket-Accept header found" - ) + ) from exc if s_w_accept != accept(key): raise InvalidHeaderValue("Sec-WebSocket-Accept", s_w_accept) From ba9e901db00bc8f9dd959a1ff2845fb9a7b753f6 Mon Sep 17 00:00:00 2001 From: Ram Rachum Date: Sat, 13 Jun 2020 23:09:20 +0300 Subject: [PATCH 2/2] Refactor header exception logic into `get_safe` method --- src/websockets/handshake.py | 31 +++++-------------------------- src/websockets/http.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 26 deletions(-) diff --git a/src/websockets/handshake.py b/src/websockets/handshake.py index 0c9cc69e8..cc6d1973e 100644 --- a/src/websockets/handshake.py +++ b/src/websockets/handshake.py @@ -31,9 +31,9 @@ import random from typing import List -from .exceptions import InvalidHeader, InvalidHeaderValue, InvalidUpgrade +from .exceptions import InvalidHeaderValue, InvalidUpgrade from .headers import ConnectionOption, UpgradeProtocol, parse_connection, parse_upgrade -from .http import Headers, MultipleValuesError +from .http import Headers __all__ = ["build_request", "check_request", "build_response", "check_response"] @@ -92,14 +92,7 @@ def check_request(headers: Headers) -> str: if not (len(upgrade) == 1 and upgrade[0].lower() == "websocket"): raise InvalidUpgrade("Upgrade", ", ".join(upgrade)) - try: - s_w_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", "more than one Sec-WebSocket-Key header found" - ) from exc + s_w_key = headers.get_safe("Sec-WebSocket-Key") try: raw_key = base64.b64decode(s_w_key.encode(), validate=True) @@ -108,14 +101,7 @@ def check_request(headers: Headers) -> str: if len(raw_key) != 16: raise InvalidHeaderValue("Sec-WebSocket-Key", s_w_key) - try: - s_w_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", "more than one Sec-WebSocket-Version header found" - ) from exc + s_w_version = headers.get_safe("Sec-WebSocket-Version") if s_w_version != "13": raise InvalidHeaderValue("Sec-WebSocket-Version", s_w_version) @@ -169,14 +155,7 @@ def check_response(headers: Headers, key: str) -> None: if not (len(upgrade) == 1 and upgrade[0].lower() == "websocket"): raise InvalidUpgrade("Upgrade", ", ".join(upgrade)) - 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", "more than one Sec-WebSocket-Accept header found" - ) from exc + s_w_accept = headers.get_safe("Sec-WebSocket-Accept") if s_w_accept != accept(key): raise InvalidHeaderValue("Sec-WebSocket-Accept", s_w_accept) diff --git a/src/websockets/http.py b/src/websockets/http.py index f87bfb76a..6e4400924 100644 --- a/src/websockets/http.py +++ b/src/websockets/http.py @@ -292,6 +292,16 @@ def copy(self) -> "Headers": copy._list = self._list.copy() return copy + def get_safe(self, header: str) -> str: + from .exceptions import InvalidHeader + + try: + return self[header] + except KeyError as exc: + raise InvalidHeader(header) from exc + except MultipleValuesError as exc: + raise InvalidHeader(header, f"more than one {header} header found") from exc + # Collection methods def __contains__(self, key: object) -> bool: