From db988f1c24db10db0a5969b0292834c4ac7b8882 Mon Sep 17 00:00:00 2001 From: Miguel Grinberg Date: Tue, 7 Mar 2017 22:15:25 -0800 Subject: [PATCH] Better handling of close packets from the client Fixes #41 --- engineio/asyncio_socket.py | 2 ++ engineio/socket.py | 2 ++ tests/test_asyncio_socket.py | 8 ++++++++ tests/test_socket.py | 8 ++++++++ 4 files changed, 20 insertions(+) diff --git a/engineio/asyncio_socket.py b/engineio/asyncio_socket.py index 2b4256fc..d431cd88 100644 --- a/engineio/asyncio_socket.py +++ b/engineio/asyncio_socket.py @@ -41,6 +41,8 @@ async def receive(self, pkt): await self.server._trigger_event('message', self.sid, pkt.data) elif pkt.packet_type == packet.UPGRADE: await self.send(packet.Packet(packet.NOOP)) + elif pkt.packet_type == packet.CLOSE: + await self.close(wait=False, abort=True) else: raise ValueError diff --git a/engineio/socket.py b/engineio/socket.py index 0a7e6a8e..a7cdb10f 100755 --- a/engineio/socket.py +++ b/engineio/socket.py @@ -55,6 +55,8 @@ def receive(self, pkt): async=self.server.async_handlers) elif pkt.packet_type == packet.UPGRADE: self.send(packet.Packet(packet.NOOP)) + elif pkt.packet_type == packet.CLOSE: + self.close(wait=False, abort=True) else: raise ValueError diff --git a/tests/test_asyncio_socket.py b/tests/test_asyncio_socket.py index d5d7d9ad..376bee89 100644 --- a/tests/test_asyncio_socket.py +++ b/tests/test_asyncio_socket.py @@ -250,6 +250,14 @@ def test_upgrade_not_supported(self): _run(s._upgrade_websocket(environ)) mock_server._bad_request.assert_called_once_with() + def test_close_packet(self): + mock_server = self._get_mock_server() + s = asyncio_socket.AsyncSocket(mock_server, 'sid') + s.connected = True + s.close = AsyncMock() + _run(s.receive(packet.Packet(packet.CLOSE))) + s.close.mock.assert_called_once_with(wait=False, abort=True) + def test_websocket_read_write(self): mock_server = self._get_mock_server() s = asyncio_socket.AsyncSocket(mock_server, 'sid') diff --git a/tests/test_socket.py b/tests/test_socket.py index 0537f34f..589b855c 100644 --- a/tests/test_socket.py +++ b/tests/test_socket.py @@ -227,6 +227,14 @@ def test_upgrade_no_upgrade_packet(self): self.assertEqual(s.queue.get().packet_type, packet.NOOP) self.assertFalse(s.upgraded) + def test_close_packet(self): + mock_server = self._get_mock_server() + s = socket.Socket(mock_server, 'sid') + s.connected = True + s.close = mock.MagicMock() + s.receive(packet.Packet(packet.CLOSE)) + s.close.assert_called_once_with(wait=False, abort=True) + def test_invalid_packet_type(self): mock_server = self._get_mock_server() s = socket.Socket(mock_server, 'sid')