Skip to content

Commit fb72954

Browse files
bdracoara-25
bdraco
and
ara-25
authored
[PR #9511/75ae623 backport][3.11] Fix exceptions in websocket receive_* methods (#9676)
Co-authored-by: J. Nick Koston <[email protected]> Co-authored-by: ara-25 <[email protected]>
1 parent 262fb12 commit fb72954

12 files changed

+166
-21
lines changed

‎CHANGES/6800.bugfix.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Modified websocket :meth:`aiohttp.ClientWebSocketResponse.receive_str`, :py:meth:`aiohttp.ClientWebSocketResponse.receive_bytes`, :py:meth:`aiohttp.web.WebSocketResponse.receive_str` & :py:meth:`aiohttp.web.WebSocketResponse.receive_bytes` methods to raise new :py:exc:`aiohttp.WSMessageTypeError` exception, instead of generic :py:exc:`TypeError`, when websocket messages of incorrect types are received -- by :user:`ara-25`.

‎CONTRIBUTORS.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
- Contributors -
22
----------------
33
A. Jesse Jiryu Davis
4+
Abdur Rehman Ali
45
Adam Bannister
56
Adam Cooper
67
Adam Horacek

‎aiohttp/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
TCPConnector,
4444
TooManyRedirects,
4545
UnixConnector,
46+
WSMessageTypeError,
4647
WSServerHandshakeError,
4748
request,
4849
)
@@ -238,6 +239,7 @@
238239
# workers (imported lazily with __getattr__)
239240
"GunicornUVLoopWebWorker",
240241
"GunicornWebWorker",
242+
"WSMessageTypeError",
241243
)
242244

243245

‎aiohttp/client.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
ServerTimeoutError,
6767
SocketTimeoutError,
6868
TooManyRedirects,
69+
WSMessageTypeError,
6970
WSServerHandshakeError,
7071
)
7172
from .client_reqrep import (
@@ -151,6 +152,7 @@
151152
"ClientTimeout",
152153
"ClientWSTimeout",
153154
"request",
155+
"WSMessageTypeError",
154156
)
155157

156158

‎aiohttp/client_exceptions.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
"NonHttpUrlClientError",
5151
"InvalidUrlRedirectClientError",
5252
"NonHttpUrlRedirectClientError",
53+
"WSMessageTypeError",
5354
)
5455

5556

@@ -410,3 +411,7 @@ def __str__(self) -> str:
410411
"[{0.certificate_error.__class__.__name__}: "
411412
"{0.certificate_error.args}]".format(self)
412413
)
414+
415+
416+
class WSMessageTypeError(TypeError):
417+
"""WebSocket message type is not valid."""

‎aiohttp/client_ws.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
import attr
99

10-
from .client_exceptions import ClientError, ServerTimeoutError
10+
from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError
1111
from .client_reqrep import ClientResponse
1212
from .helpers import calculate_timeout_when, set_result
1313
from .http import (
@@ -377,13 +377,17 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
377377
async def receive_str(self, *, timeout: Optional[float] = None) -> str:
378378
msg = await self.receive(timeout)
379379
if msg.type is not WSMsgType.TEXT:
380-
raise TypeError(f"Received message {msg.type}:{msg.data!r} is not str")
380+
raise WSMessageTypeError(
381+
f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT"
382+
)
381383
return cast(str, msg.data)
382384

383385
async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes:
384386
msg = await self.receive(timeout)
385387
if msg.type is not WSMsgType.BINARY:
386-
raise TypeError(f"Received message {msg.type}:{msg.data!r} is not bytes")
388+
raise WSMessageTypeError(
389+
f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY"
390+
)
387391
return cast(bytes, msg.data)
388392

389393
async def receive_json(

‎aiohttp/web_ws.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from . import hdrs
1313
from ._websocket.writer import DEFAULT_LIMIT
1414
from .abc import AbstractStreamWriter
15+
from .client_exceptions import WSMessageTypeError
1516
from .helpers import calculate_timeout_when, set_exception, set_result
1617
from .http import (
1718
WS_CLOSED_MESSAGE,
@@ -578,17 +579,17 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
578579
async def receive_str(self, *, timeout: Optional[float] = None) -> str:
579580
msg = await self.receive(timeout)
580581
if msg.type is not WSMsgType.TEXT:
581-
raise TypeError(
582-
"Received message {}:{!r} is not WSMsgType.TEXT".format(
583-
msg.type, msg.data
584-
)
582+
raise WSMessageTypeError(
583+
f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT"
585584
)
586585
return cast(str, msg.data)
587586

588587
async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes:
589588
msg = await self.receive(timeout)
590589
if msg.type is not WSMsgType.BINARY:
591-
raise TypeError(f"Received message {msg.type}:{msg.data!r} is not bytes")
590+
raise WSMessageTypeError(
591+
f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY"
592+
)
592593
return cast(bytes, msg.data)
593594

594595
async def receive_json(

‎docs/client_reference.rst

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1689,7 +1689,7 @@ manually.
16891689

16901690
:return str: peer's message content.
16911691

1692-
:raise TypeError: if message is :const:`~aiohttp.WSMsgType.BINARY`.
1692+
:raise aiohttp.WSMessageTypeError: if message is not :const:`~aiohttp.WSMsgType.TEXT`.
16931693

16941694
.. method:: receive_bytes()
16951695
:async:
@@ -1700,7 +1700,7 @@ manually.
17001700

17011701
:return bytes: peer's message content.
17021702

1703-
:raise TypeError: if message is :const:`~aiohttp.WSMsgType.TEXT`.
1703+
:raise aiohttp.WSMessageTypeError: if message is not :const:`~aiohttp.WSMsgType.BINARY`.
17041704

17051705
.. method:: receive_json(*, loads=json.loads)
17061706
:async:
@@ -2257,6 +2257,12 @@ Response errors
22572257

22582258
Derived from :exc:`ClientResponseError`
22592259

2260+
.. exception:: WSMessageTypeError
2261+
2262+
Received WebSocket message of unexpected type
2263+
2264+
Derived from :exc:`TypeError`
2265+
22602266
Connection errors
22612267
^^^^^^^^^^^^^^^^^
22622268

‎docs/web_reference.rst

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1279,7 +1279,7 @@ and :ref:`aiohttp-web-signals` handlers::
12791279

12801280
:return str: peer's message content.
12811281

1282-
:raise TypeError: if message is :const:`~aiohttp.WSMsgType.BINARY`.
1282+
:raise aiohttp.WSMessageTypeError: if message is not :const:`~aiohttp.WSMsgType.TEXT`.
12831283

12841284
.. method:: receive_bytes(*, timeout=None)
12851285
:async:
@@ -1298,7 +1298,7 @@ and :ref:`aiohttp-web-signals` handlers::
12981298

12991299
:return bytes: peer's message content.
13001300

1301-
:raise TypeError: if message is :const:`~aiohttp.WSMsgType.TEXT`.
1301+
:raise aiohttp.WSMessageTypeError: if message is not :const:`~aiohttp.WSMsgType.BINARY`.
13021302

13031303
.. method:: receive_json(*, loads=json.loads, timeout=None)
13041304
:async:

‎tests/test_client_ws_functional.py

Lines changed: 54 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,14 @@
66
import pytest
77

88
import aiohttp
9-
from aiohttp import ClientConnectionResetError, ServerTimeoutError, WSMsgType, hdrs, web
9+
from aiohttp import (
10+
ClientConnectionResetError,
11+
ServerTimeoutError,
12+
WSMessageTypeError,
13+
WSMsgType,
14+
hdrs,
15+
web,
16+
)
1017
from aiohttp.client_ws import ClientWSTimeout
1118
from aiohttp.http import WSCloseCode
1219
from aiohttp.pytest_plugin import AiohttpClient
@@ -58,7 +65,28 @@ async def handler(request):
5865
resp = await client.ws_connect("/")
5966
await resp.send_str("ask")
6067

61-
with pytest.raises(TypeError):
68+
with pytest.raises(WSMessageTypeError):
69+
await resp.receive_bytes()
70+
await resp.close()
71+
72+
73+
async def test_recv_bytes_after_close(aiohttp_client: AiohttpClient) -> None:
74+
async def handler(request: web.Request) -> NoReturn:
75+
ws = web.WebSocketResponse()
76+
await ws.prepare(request)
77+
78+
await ws.close()
79+
assert False
80+
81+
app = web.Application()
82+
app.router.add_route("GET", "/", handler)
83+
client = await aiohttp_client(app)
84+
resp = await client.ws_connect("/")
85+
86+
with pytest.raises(
87+
WSMessageTypeError,
88+
match=f"Received message {WSMsgType.CLOSE}:.+ is not WSMsgType.BINARY",
89+
):
6290
await resp.receive_bytes()
6391
await resp.close()
6492

@@ -103,14 +131,35 @@ async def handler(request):
103131

104132
await resp.send_bytes(b"ask")
105133

106-
with pytest.raises(TypeError):
134+
with pytest.raises(WSMessageTypeError):
107135
await resp.receive_str()
108136

109137
await resp.close()
110138

111139

112-
async def test_send_recv_json(aiohttp_client) -> None:
113-
async def handler(request):
140+
async def test_recv_text_after_close(aiohttp_client: AiohttpClient) -> None:
141+
async def handler(request: web.Request) -> NoReturn:
142+
ws = web.WebSocketResponse()
143+
await ws.prepare(request)
144+
145+
await ws.close()
146+
assert False
147+
148+
app = web.Application()
149+
app.router.add_route("GET", "/", handler)
150+
client = await aiohttp_client(app)
151+
resp = await client.ws_connect("/")
152+
153+
with pytest.raises(
154+
WSMessageTypeError,
155+
match=f"Received message {WSMsgType.CLOSE}:.+ is not WSMsgType.TEXT",
156+
):
157+
await resp.receive_str()
158+
await resp.close()
159+
160+
161+
async def test_send_recv_json(aiohttp_client: AiohttpClient) -> None:
162+
async def handler(request: web.Request) -> web.WebSocketResponse:
114163
ws = web.WebSocketResponse()
115164
await ws.prepare(request)
116165

0 commit comments

Comments
 (0)