Skip to content

Commit

Permalink
fix: upgrade is not websocket and dependencies are installed, should …
Browse files Browse the repository at this point in the history
…not warning (#2360)

* fix: upgrade is not websocket, handle httptools with the same logic as h11

* test: add case

* refactor: clearer judgment

* Refactor(http): streamline upgrade warning logic

* ruff

* fix: simplify httptools_impl methods

---------

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
  • Loading branch information
vvanglro and Kludex authored Aug 13, 2024
1 parent cee31a6 commit 587a1cc
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 18 deletions.
44 changes: 44 additions & 0 deletions tests/protocols/test_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,18 @@
b"".join([b"x" * 32 * 1024 + b"\r\n", b"\r\n", b"\r\n"]),
]

UPGRADE_REQUEST_ERROR_FIELD = b"\r\n".join(
[
b"GET / HTTP/1.1",
b"Host: example.org",
b"Connection: upgrade",
b"Upgrade: not-websocket",
b"Sec-WebSocket-Version: 11",
b"",
b"",
]
)


class MockTransport:
def __init__(self, sockname=None, peername=None, sslcontext=False):
Expand Down Expand Up @@ -1080,3 +1092,35 @@ async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
assert b"Hi!" in protocol.transport.buffer

assert not expected_states # consumed


async def test_header_upgrade_is_not_websocket_depend_installed(
caplog: pytest.LogCaptureFixture, http_protocol_cls: HTTPProtocol
):
caplog.set_level(logging.WARNING, logger="uvicorn.error")
app = Response("Hello, world", media_type="text/plain")

protocol = get_connected_protocol(app, http_protocol_cls)
protocol.data_received(UPGRADE_REQUEST_ERROR_FIELD)
await protocol.loop.run_one()
assert "Unsupported upgrade request." in caplog.text
msg = "No supported WebSocket library detected. Please use \"pip install 'uvicorn[standard]'\", or install 'websockets' or 'wsproto' manually." # noqa: E501
assert msg not in caplog.text
assert b"HTTP/1.1 200 OK" in protocol.transport.buffer
assert b"Hello, world" in protocol.transport.buffer


async def test_header_upgrade_is_websocket_depend_not_installed(
caplog: pytest.LogCaptureFixture, http_protocol_cls: HTTPProtocol
):
caplog.set_level(logging.WARNING, logger="uvicorn.error")
app = Response("Hello, world", media_type="text/plain")

protocol = get_connected_protocol(app, http_protocol_cls, ws="none")
protocol.data_received(UPGRADE_REQUEST_ERROR_FIELD)
await protocol.loop.run_one()
assert "Unsupported upgrade request." in caplog.text
msg = "No supported WebSocket library detected. Please use \"pip install 'uvicorn[standard]'\", or install 'websockets' or 'wsproto' manually." # noqa: E501
assert msg in caplog.text
assert b"HTTP/1.1 200 OK" in protocol.transport.buffer
assert b"Hello, world" in protocol.transport.buffer
24 changes: 16 additions & 8 deletions uvicorn/protocols/http/h11_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,14 +147,24 @@ def _get_upgrade(self) -> bytes | None:

def _should_upgrade_to_ws(self) -> bool:
if self.ws_protocol_class is None:
if self.config.ws == "auto":
msg = "Unsupported upgrade request."
self.logger.warning(msg)
msg = "No supported WebSocket library detected. Please use \"pip install 'uvicorn[standard]'\", or install 'websockets' or 'wsproto' manually." # noqa: E501
self.logger.warning(msg)
return False
return True

def _unsupported_upgrade_warning(self) -> None:
msg = "Unsupported upgrade request."
self.logger.warning(msg)
if not self._should_upgrade_to_ws():
msg = "No supported WebSocket library detected. Please use \"pip install 'uvicorn[standard]'\", or install 'websockets' or 'wsproto' manually." # noqa: E501
self.logger.warning(msg)

def _should_upgrade(self) -> bool:
upgrade = self._get_upgrade()
if upgrade == b"websocket" and self._should_upgrade_to_ws():
return True
if upgrade is not None:
self._unsupported_upgrade_warning()
return False

def data_received(self, data: bytes) -> None:
self._unset_keepalive_if_required()

Expand Down Expand Up @@ -206,9 +216,7 @@ def handle_events(self) -> None:
"headers": self.headers,
"state": self.app_state.copy(),
}

upgrade = self._get_upgrade()
if upgrade == b"websocket" and self._should_upgrade_to_ws():
if self._should_upgrade():
self.handle_websocket_upgrade(event)
return

Expand Down
22 changes: 12 additions & 10 deletions uvicorn/protocols/http/httptools_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,19 +141,20 @@ def _get_upgrade(self) -> bytes | None:
return upgrade
return None # pragma: full coverage

def _should_upgrade_to_ws(self, upgrade: bytes | None) -> bool:
if upgrade == b"websocket" and self.ws_protocol_class is not None:
return True
if self.config.ws == "auto":
msg = "Unsupported upgrade request."
self.logger.warning(msg)
def _should_upgrade_to_ws(self) -> bool:
if self.ws_protocol_class is None:
return False
return True

def _unsupported_upgrade_warning(self) -> None:
self.logger.warning("Unsupported upgrade request.")
if not self._should_upgrade_to_ws():
msg = "No supported WebSocket library detected. Please use \"pip install 'uvicorn[standard]'\", or install 'websockets' or 'wsproto' manually." # noqa: E501
self.logger.warning(msg)
return False

def _should_upgrade(self) -> bool:
upgrade = self._get_upgrade()
return self._should_upgrade_to_ws(upgrade)
return upgrade == b"websocket" and self._should_upgrade_to_ws()

def data_received(self, data: bytes) -> None:
self._unset_keepalive_if_required()
Expand All @@ -166,9 +167,10 @@ def data_received(self, data: bytes) -> None:
self.send_400_response(msg)
return
except httptools.HttpParserUpgrade:
upgrade = self._get_upgrade()
if self._should_upgrade_to_ws(upgrade):
if self._should_upgrade():
self.handle_websocket_upgrade()
else:
self._unsupported_upgrade_warning()

def handle_websocket_upgrade(self) -> None:
if self.logger.level <= TRACE_LOG_LEVEL:
Expand Down

0 comments on commit 587a1cc

Please sign in to comment.