prevent unnecessary resource allocation (#1574)

Author: miguelgrinberg

src/socketio/async_server.py

@@ -427,6 +427,8 @@ async def disconnect(self, sid, namespace=None, ignore_queue=False):
         if delete_it:
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
             eio_sid = self.manager.pre_disconnect(sid, namespace=namespace)
+            if eio_sid in self._binary_packet:
+                del self._binary_packet[eio_sid]
             await self._send_packet(eio_sid, self.packet_class(
                 packet.DISCONNECT, namespace=namespace))
             await self._trigger_event('disconnect', namespace, sid,
@@ -702,6 +704,9 @@ async def _handle_eio_message(self, eio_sid, data):
                                        pkt.data)
             elif pkt.packet_type == packet.BINARY_EVENT or \
                     pkt.packet_type == packet.BINARY_ACK:
+                if not self.manager.sid_from_eio_sid(eio_sid,
+                                                     pkt.namespace or '/'):
+                    raise ValueError('Unexpected binary packet')
                 self._binary_packet[eio_sid] = pkt
             elif pkt.packet_type == packet.CONNECT_ERROR:
                 raise ValueError('Unexpected CONNECT_ERROR packet.')

src/socketio/server.py

@@ -402,6 +402,8 @@ def disconnect(self, sid, namespace=None, ignore_queue=False):
         if delete_it:
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
             eio_sid = self.manager.pre_disconnect(sid, namespace=namespace)
+            if eio_sid in self._binary_packet:
+                del self._binary_packet[eio_sid]
             self._send_packet(eio_sid, self.packet_class(
                 packet.DISCONNECT, namespace=namespace))
             self._trigger_event('disconnect', namespace, sid,
@@ -663,6 +665,9 @@ def _handle_eio_message(self, eio_sid, data):
                 self._handle_ack(eio_sid, pkt.namespace, pkt.id, pkt.data)
             elif pkt.packet_type == packet.BINARY_EVENT or \
                     pkt.packet_type == packet.BINARY_ACK:
+                if not self.manager.sid_from_eio_sid(eio_sid,
+                                                     pkt.namespace or '/'):
+                    raise ValueError('Unexpected binary packet')
                 self._binary_packet[eio_sid] = pkt
             elif pkt.packet_type == packet.CONNECT_ERROR:
                 raise ValueError('Unexpected CONNECT_ERROR packet.')

tests/async/test_server.py

@@ -719,6 +719,20 @@ async def test_handle_event_binary_ack(self, eio):
             sid, 321, ['my message', 'a', b'foo']
         )
 
+    async def test_handle_event_binary_from_unknown(self, eio):
+        eio.return_value.send = mock.AsyncMock()
+        s = async_server.AsyncServer(async_handlers=False)
+        await s.manager.connect('123', '/')
+        handler = mock.MagicMock()
+        s.on('my message', handler)
+        with pytest.raises(ValueError):
+            await s._handle_eio_message(
+                '999',
+                '52-["my message","a",'
+                '{"_placeholder":true,"num":1},'
+                '{"_placeholder":true,"num":0}]',
+            )
+
     async def test_handle_event_with_ack(self, eio):
         eio.return_value.send = mock.AsyncMock()
         s = async_server.AsyncServer(async_handlers=False)
@@ -923,6 +937,24 @@ async def test_disconnect_twice_namespace(self, eio):
         await s.disconnect('1', namespace='/foo')
         assert calls == s.eio.send.await_count
 
+    async def test_disconnect_with_partial_binary_packet(self, eio):
+        eio.return_value.send = mock.AsyncMock()
+        eio.return_value.disconnect = mock.AsyncMock()
+        s = async_server.AsyncServer()
+        await s._handle_eio_connect('123', 'environ')
+        await s._handle_eio_message('123', '0')
+        await s._handle_eio_message(
+            '123',
+            '52-["my message","a",'
+            '{"_placeholder":true,"num":1},'
+            '{"_placeholder":true,"num":0}]',
+        )
+        await s._handle_eio_message('123', b'foo')
+        assert s._binary_packet['123'] is not None
+        await s.disconnect('1')
+        s.eio.send.assert_any_await('123', '1')
+        assert '123' not in s._binary_packet
+
     async def test_namespace_handler(self, eio):
         eio.return_value.send = mock.AsyncMock()
         result = {}

tests/common/test_server.py

@@ -676,6 +676,19 @@ def test_handle_event_binary_ack(self, eio):
             sid, 321, ['my message', 'a', b'foo']
         )
 
+    def test_handle_event_binary_from_unknown(self, eio):
+        s = server.Server(async_handlers=False)
+        s.manager.connect('123', '/')
+        handler = mock.MagicMock()
+        s.on('my message', handler)
+        with pytest.raises(ValueError):
+            s._handle_eio_message(
+                '999',
+                '52-["my message","a",'
+                '{"_placeholder":true,"num":1},'
+                '{"_placeholder":true,"num":0}]',
+            )
+
     def test_handle_event_with_ack(self, eio):
         s = server.Server(async_handlers=False)
         sid = s.manager.connect('123', '/')
@@ -849,6 +862,22 @@ def test_disconnect_twice_namespace(self, eio):
         s.disconnect('123', namespace='/foo')
         assert calls == s.eio.send.call_count
 
+    def test_disconnect_with_partial_binary_packet(self, eio):
+        s = server.Server()
+        s._handle_eio_connect('123', 'environ')
+        s._handle_eio_message('123', '0')
+        s._handle_eio_message(
+            '123',
+            '52-["my message","a",'
+            '{"_placeholder":true,"num":1},'
+            '{"_placeholder":true,"num":0}]',
+        )
+        s._handle_eio_message('123', b'foo')
+        assert s._binary_packet['123'] is not None
+        s.disconnect('1')
+        s.eio.send.assert_any_call('123', '1')
+        assert '123' not in s._binary_packet
+
     def test_namespace_handler(self, eio):
         result = {}