Better handling of close packets from the client

Fixes #41
miguelgrinberg · Mar 8, 2017 · db988f1 · db988f1
1 parent 6fcf926
commit db988f1
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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')