Skip to content

Refactor header exception logic into get_safe method #779

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
cool-RR wants to merge 2 commits into from
Closed

Refactor header exception logic into get_safe method #779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
222a9d0
Fix exception causes in handshake.py
cool-RR Jun 13, 2020
ba9e901
Refactor header exception logic into `get_safe` method
cool-RR Jun 13, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 7 additions & 28 deletions src/websockets/handshake.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions src/websockets/http.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,16 @@ def copy(self) -> "Headers":
copy._list = self._list.copy()
return copy

def get_safe(self, header: str) -> str:

@aaugustin aaugustin Jun 14, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The only difference between get_safe and __getitem__ is that get_safe translates KeyError or MultipleValuesError into InvalidHeader. Having two methods that do the same thing in the same class is weird.

websockets.http.Headers is intended as a low-level utility that only knows about HTTP. This is why it uses websockets.http.MultipleValuesError instead of one of the websockets-specific exceptions in websockets.exceptions.

To factor out this behavior, I'd write a utility function in websockets.handshake. The job of this function would be to raise an exception from websockets.exceptions and to provide a good error message when the WebSocket protocol specifies that a given header must be provided exactly once, but it's there zero or multiple times.

You can call it get_unique_header for example.

from .exceptions import InvalidHeader

@aaugustin aaugustin Jun 14, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess you've had to do this to avoid an import cycle. This is a very common issue since I implemented static type checking with mypy.

The websockets.exceptions module is imported at the bottom at websockets.http, so you remove this import use websockets.exceptions.InvalidHeader.

It's ugly but I'd rather have a module-level import at the bottom of a module than scattered imports inside functions.

But if you're moving the code to websockets.handshake, the problem goes away.


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:
Expand Down