[3.8] Update close status codes ws (#5198). (#5279)

(cherry picked from commit f82f084) Co-authored-by: Dmitry Erlikh <[email protected]> Co-authored-by: Dmitry Erlikh <[email protected]>
aio-libs · Nov 23, 2020 · d88aa16 · d88aa16
1 parent 5ebc985
commit d88aa16
Show file tree

Hide file tree

Showing 8 changed files with 55 additions and 38 deletions.
diff --git a/CHANGES/5192.bugfix b/CHANGES/5192.bugfix
@@ -0,0 +1 @@
+Add ABNORMAL_CLOSURE and BAD_GATEWAY to WSCloseCode
diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py
@@ -12,6 +12,7 @@
     WS_CLOSED_MESSAGE,
     WS_CLOSING_MESSAGE,
     WebSocketError,
+    WSCloseCode,
     WSMessage,
     WSMsgType,
 )
@@ -101,7 +102,7 @@ def _send_heartbeat(self) -> None:
     def _pong_not_received(self) -> None:
         if not self._closed:
             self._closed = True
-            self._close_code = 1006
+            self._close_code = WSCloseCode.ABNORMAL_CLOSURE
             self._exception = asyncio.TimeoutError()
             self._response.close()
 
@@ -163,7 +164,7 @@ async def send_json(
     ) -> None:
         await self.send_str(dumps(data), compress=compress)
 
-    async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
+    async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool:
         # we need to break `receive()` cycle first,
         # `close()` may be called from different task
         if self._waiting is not None and not self._closed:
@@ -176,11 +177,11 @@ async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
             try:
                 await self._writer.close(code, message)
             except asyncio.CancelledError:
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 self._response.close()
                 raise
             except Exception as exc:
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 self._exception = exc
                 self._response.close()
                 return True
@@ -194,11 +195,11 @@ async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
                     with async_timeout.timeout(self._timeout, loop=self._loop):
                         msg = await self._reader.read()
                 except asyncio.CancelledError:
-                    self._close_code = 1006
+                    self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                     self._response.close()
                     raise
                 except Exception as exc:
-                    self._close_code = 1006
+                    self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                     self._exception = exc
                     self._response.close()
                     return True
@@ -234,15 +235,15 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
                     self._waiting = None
                     set_result(waiter, True)
             except (asyncio.CancelledError, asyncio.TimeoutError):
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 raise
             except EofStream:
-                self._close_code = 1000
+                self._close_code = WSCloseCode.OK
                 await self.close()
                 return WSMessage(WSMsgType.CLOSED, None, None)
             except ClientError:
                 self._closed = True
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 return WS_CLOSED_MESSAGE
             except WebSocketError as exc:
                 self._close_code = exc.code
@@ -251,7 +252,7 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
             except Exception as exc:
                 self._exception = exc
                 self._closing = True
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 await self.close()
                 return WSMessage(WSMsgType.ERROR, exc, None)
 

diff --git a/aiohttp/http_websocket.py b/aiohttp/http_websocket.py
@@ -33,13 +33,15 @@ class WSCloseCode(IntEnum):
     GOING_AWAY = 1001
     PROTOCOL_ERROR = 1002
     UNSUPPORTED_DATA = 1003
+    ABNORMAL_CLOSURE = 1006
     INVALID_TEXT = 1007
     POLICY_VIOLATION = 1008
     MESSAGE_TOO_BIG = 1009
     MANDATORY_EXTENSION = 1010
     INTERNAL_ERROR = 1011
     SERVICE_RESTART = 1012
     TRY_AGAIN_LATER = 1013
+    BAD_GATEWAY = 1014
 
 
 ALLOWED_CLOSE_CODES = {int(i) for i in WSCloseCode}

diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py
@@ -19,6 +19,7 @@
     WebSocketError,
     WebSocketReader,
     WebSocketWriter,
+    WSCloseCode,
     WSMessage,
     WSMsgType as WSMsgType,
     ws_ext_gen,
@@ -122,7 +123,7 @@ def _send_heartbeat(self) -> None:
     def _pong_not_received(self) -> None:
         if self._req is not None and self._req.transport is not None:
             self._closed = True
-            self._close_code = 1006
+            self._close_code = WSCloseCode.ABNORMAL_CLOSURE
             self._exception = asyncio.TimeoutError()
             self._req.transport.close()
 
@@ -324,7 +325,7 @@ async def write_eof(self) -> None:  # type: ignore
         await self.close()
         self._eof_sent = True
 
-    async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
+    async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool:
         if self._writer is None:
             raise RuntimeError("Call .prepare() first")
 
@@ -346,10 +347,10 @@ async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
                 assert writer is not None
                 await writer.drain()
             except (asyncio.CancelledError, asyncio.TimeoutError):
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 raise
             except Exception as exc:
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 self._exception = exc
                 return True
 
@@ -362,18 +363,18 @@ async def close(self, *, code: int = 1000, message: bytes = b"") -> bool:
                 with async_timeout.timeout(self._timeout, loop=self._loop):
                     msg = await reader.read()
             except asyncio.CancelledError:
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 raise
             except Exception as exc:
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 self._exception = exc
                 return True
 
             if msg.type == WSMsgType.CLOSE:
                 self._close_code = msg.data
                 return True
 
-            self._close_code = 1006
+            self._close_code = WSCloseCode.ABNORMAL_CLOSURE
             self._exception = asyncio.TimeoutError()
             return True
         else:
@@ -410,10 +411,10 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
                     set_result(waiter, True)
                     self._waiting = None
             except (asyncio.CancelledError, asyncio.TimeoutError):
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 raise
             except EofStream:
-                self._close_code = 1000
+                self._close_code = WSCloseCode.OK
                 await self.close()
                 return WSMessage(WSMsgType.CLOSED, None, None)
             except WebSocketError as exc:
@@ -423,7 +424,7 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
             except Exception as exc:
                 self._exception = exc
                 self._closing = True
-                self._close_code = 1006
+                self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 await self.close()
                 return WSMessage(WSMsgType.ERROR, exc, None)
 

diff --git a/docs/client_reference.rst b/docs/client_reference.rst
@@ -1578,14 +1578,14 @@ manually.
          The method is converted into :term:`coroutine`,
          *compress* parameter added.
 
-   .. comethod:: close(*, code=1000, message=b'')
+   .. comethod:: close(*, code=WSCloseCode.OK, message=b'')
 
       A :ref:`coroutine<coroutine>` that initiates closing handshake by sending
       :const:`~aiohttp.WSMsgType.CLOSE` message. It waits for
       close response from server. To add a timeout to `close()` call
       just wrap the call with `asyncio.wait()` or `asyncio.wait_for()`.
 
-      :param int code: closing code
+      :param int code: closing code. See also :class:`~aiohttp.WSCloseCode`.
 
       :param message: optional payload of *close* message,
          :class:`str` (converted to *UTF-8* encoded bytes) or :class:`bytes`.

diff --git a/docs/web_reference.rst b/docs/web_reference.rst
@@ -1098,14 +1098,14 @@ WebSocketResponse
          The method is converted into :term:`coroutine`,
          *compress* parameter added.
 
-   .. comethod:: close(*, code=1000, message=b'')
+   .. comethod:: close(*, code=WSCloseCode.OK, message=b'')
 
       A :ref:`coroutine<coroutine>` that initiates closing
       handshake by sending :const:`~aiohttp.WSMsgType.CLOSE` message.
 
       It is safe to call `close()` from different task.
 
-      :param int code: closing code
+      :param int code: closing code. See also :class:`~aiohttp.WSCloseCode`.
 
       :param message: optional payload of *close* message,
                       :class:`str` (converted to *UTF-8* encoded bytes)

diff --git a/docs/websocket_utilities.rst b/docs/websocket_utilities.rst
@@ -79,6 +79,18 @@ WebSocket utilities
        connect to a different IP (when there are multiple for the
        target) or reconnect to the same IP upon user action.
 
+    .. attribute:: ABNORMAL_CLOSURE
+
+       Used to indicate that a connection was closed abnormally
+       (that is, with no close frame being sent) when a status code
+       is expected.
+
+    .. attribute:: BAD_GATEWAY
+
+       The server was acting as a gateway or proxy and received
+       an invalid response from the upstream server.
+       This is similar to 502 HTTP Status Code.
+
 
 .. class:: WSMsgType
 

diff --git a/tests/test_web_websocket_functional.py b/tests/test_web_websocket_functional.py
@@ -6,7 +6,7 @@
 
 import aiohttp
 from aiohttp import web
-from aiohttp.http import WSMsgType
+from aiohttp.http import WSCloseCode, WSMsgType
 
 
 async def test_websocket_can_prepare(loop, aiohttp_client) -> None:
@@ -153,11 +153,11 @@ async def handler(request):
 
     msg = await ws.receive()
     assert msg.type == aiohttp.WSMsgType.CLOSE
-    assert msg.data == 1000
+    assert msg.data == WSCloseCode.OK
     assert msg.extra == ""
 
     assert ws.closed
-    assert ws.close_code == 1000
+    assert ws.close_code == WSCloseCode.OK
 
     await closed
 
@@ -188,11 +188,11 @@ async def handler(request):
 
     msg = await ws.receive()
     assert msg.type == aiohttp.WSMsgType.CLOSE
-    assert msg.data == 1000
+    assert msg.data == WSCloseCode.OK
     assert msg.extra == ""
 
     assert ws.closed
-    assert ws.close_code == 1000
+    assert ws.close_code == WSCloseCode.OK
 
     await closed
 
@@ -223,7 +223,7 @@ async def handler(request):
 
     msg = await ws.receive()
     assert msg.type == aiohttp.WSMsgType.CLOSE
-    assert msg.data == 1000
+    assert msg.data == WSCloseCode.OK
     assert msg.extra == ""
 
     await ws.close()
@@ -244,7 +244,7 @@ async def handler(request):
         begin = ws._loop.time()
         assert await ws.close()
         elapsed = ws._loop.time() - begin
-        assert ws.close_code == 1006
+        assert ws.close_code == WSCloseCode.ABNORMAL_CLOSURE
         assert isinstance(ws.exception(), asyncio.TimeoutError)
         aborted.set_result(1)
         return ws
@@ -300,7 +300,7 @@ async def handler(request):
 
     ws = await client.ws_connect("/", autoclose=False, protocols=("eggs", "bar"))
 
-    await srv_ws.close(code=1007)
+    await srv_ws.close(code=WSCloseCode.INVALID_TEXT)
 
     msg = await ws.receive()
     assert msg.type == WSMsgType.CLOSE
@@ -321,7 +321,7 @@ async def handler(request):
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSE
-        assert msg.data == 1000
+        assert msg.data == WSCloseCode.OK
         assert msg.extra == "exit message"
         closed.set_result(None)
         return ws
@@ -336,7 +336,7 @@ async def handler(request):
 
     msg = await ws.receive()
     assert msg.type == WSMsgType.PONG
-    await ws.close(code=1000, message="exit message")
+    await ws.close(code=WSCloseCode.OK, message="exit message")
     await closed
 
 
@@ -407,7 +407,7 @@ async def handler(request):
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSE
-        assert msg.data == 1000
+        assert msg.data == WSCloseCode.OK
         assert msg.extra == "exit message"
         closed.set_result(None)
         return ws
@@ -423,7 +423,7 @@ async def handler(request):
     assert msg.type == WSMsgType.PONG
     assert msg.data == b"data"
 
-    await ws.close(code=1000, message="exit message")
+    await ws.close(code=WSCloseCode.OK, message="exit message")
 
     await closed
 
@@ -511,7 +511,7 @@ async def handler(request):
         assert not ws.closed
         await ws.close()
         assert ws.closed
-        assert ws.close_code == 1007
+        assert ws.close_code == WSCloseCode.INVALID_TEXT
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSED
@@ -525,7 +525,7 @@ async def handler(request):
 
     ws = await client.ws_connect("/", autoclose=False, protocols=("eggs", "bar"))
 
-    await ws.close(code=1007)
+    await ws.close(code=WSCloseCode.INVALID_TEXT)
     msg = await ws.receive()
     assert msg.type == WSMsgType.CLOSED
     await closed
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Add ABNORMAL_CLOSURE and BAD_GATEWAY to WSCloseCode