diff --git a/src/websockets/handshake.py b/src/websockets/handshake.py index 646b6dba4..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,30 +92,16 @@ 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: - raise InvalidHeader("Sec-WebSocket-Key") - except MultipleValuesError: - raise InvalidHeader( - "Sec-WebSocket-Key", "more than one Sec-WebSocket-Key header found" - ) + s_w_key = headers.get_safe("Sec-WebSocket-Key") 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: - raise InvalidHeader( - "Sec-WebSocket-Version", "more than one Sec-WebSocket-Version header found" - ) + 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: - raise InvalidHeader("Sec-WebSocket-Accept") - except MultipleValuesError: - raise InvalidHeader( - "Sec-WebSocket-Accept", "more than one Sec-WebSocket-Accept header found" - ) + 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: