Skip to content

Commit 7a6b9b3

Browse files
prevent unnecessary resource allocation
1 parent b29beef commit 7a6b9b3

4 files changed

Lines changed: 71 additions & 0 deletions

File tree

src/socketio/async_server.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,8 @@ async def disconnect(self, sid, namespace=None, ignore_queue=False):
427427
if delete_it:
428428
self.logger.info('Disconnecting %s [%s]', sid, namespace)
429429
eio_sid = self.manager.pre_disconnect(sid, namespace=namespace)
430+
if eio_sid in self._binary_packet:
431+
del self._binary_packet[eio_sid]
430432
await self._send_packet(eio_sid, self.packet_class(
431433
packet.DISCONNECT, namespace=namespace))
432434
await self._trigger_event('disconnect', namespace, sid,
@@ -702,6 +704,9 @@ async def _handle_eio_message(self, eio_sid, data):
702704
pkt.data)
703705
elif pkt.packet_type == packet.BINARY_EVENT or \
704706
pkt.packet_type == packet.BINARY_ACK:
707+
if not self.manager.sid_from_eio_sid(eio_sid,
708+
pkt.namespace or '/'):
709+
raise ValueError('Unexpected binary packet')
705710
self._binary_packet[eio_sid] = pkt
706711
elif pkt.packet_type == packet.CONNECT_ERROR:
707712
raise ValueError('Unexpected CONNECT_ERROR packet.')

src/socketio/server.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,8 @@ def disconnect(self, sid, namespace=None, ignore_queue=False):
402402
if delete_it:
403403
self.logger.info('Disconnecting %s [%s]', sid, namespace)
404404
eio_sid = self.manager.pre_disconnect(sid, namespace=namespace)
405+
if eio_sid in self._binary_packet:
406+
del self._binary_packet[eio_sid]
405407
self._send_packet(eio_sid, self.packet_class(
406408
packet.DISCONNECT, namespace=namespace))
407409
self._trigger_event('disconnect', namespace, sid,
@@ -663,6 +665,9 @@ def _handle_eio_message(self, eio_sid, data):
663665
self._handle_ack(eio_sid, pkt.namespace, pkt.id, pkt.data)
664666
elif pkt.packet_type == packet.BINARY_EVENT or \
665667
pkt.packet_type == packet.BINARY_ACK:
668+
if not self.manager.sid_from_eio_sid(eio_sid,
669+
pkt.namespace or '/'):
670+
raise ValueError('Unexpected binary packet')
666671
self._binary_packet[eio_sid] = pkt
667672
elif pkt.packet_type == packet.CONNECT_ERROR:
668673
raise ValueError('Unexpected CONNECT_ERROR packet.')

tests/async/test_server.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -719,6 +719,20 @@ async def test_handle_event_binary_ack(self, eio):
719719
sid, 321, ['my message', 'a', b'foo']
720720
)
721721

722+
async def test_handle_event_binary_from_unknown(self, eio):
723+
eio.return_value.send = mock.AsyncMock()
724+
s = async_server.AsyncServer(async_handlers=False)
725+
sid = await s.manager.connect('123', '/')
726+
handler = mock.MagicMock()
727+
s.on('my message', handler)
728+
with pytest.raises(ValueError):
729+
await s._handle_eio_message(
730+
'999',
731+
'52-["my message","a",'
732+
'{"_placeholder":true,"num":1},'
733+
'{"_placeholder":true,"num":0}]',
734+
)
735+
722736
async def test_handle_event_with_ack(self, eio):
723737
eio.return_value.send = mock.AsyncMock()
724738
s = async_server.AsyncServer(async_handlers=False)
@@ -923,6 +937,24 @@ async def test_disconnect_twice_namespace(self, eio):
923937
await s.disconnect('1', namespace='/foo')
924938
assert calls == s.eio.send.await_count
925939

940+
async def test_disconnect_with_partial_binary_packet(self, eio):
941+
eio.return_value.send = mock.AsyncMock()
942+
eio.return_value.disconnect = mock.AsyncMock()
943+
s = async_server.AsyncServer()
944+
await s._handle_eio_connect('123', 'environ')
945+
await s._handle_eio_message('123', '0')
946+
await s._handle_eio_message(
947+
'123',
948+
'52-["my message","a",'
949+
'{"_placeholder":true,"num":1},'
950+
'{"_placeholder":true,"num":0}]',
951+
)
952+
await s._handle_eio_message('123', b'foo')
953+
assert s._binary_packet['123'] is not None
954+
await s.disconnect('1')
955+
s.eio.send.assert_any_await('123', '1')
956+
assert '123' not in s._binary_packet
957+
926958
async def test_namespace_handler(self, eio):
927959
eio.return_value.send = mock.AsyncMock()
928960
result = {}

tests/common/test_server.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -676,6 +676,19 @@ def test_handle_event_binary_ack(self, eio):
676676
sid, 321, ['my message', 'a', b'foo']
677677
)
678678

679+
def test_handle_event_binary_from_unknown(self, eio):
680+
s = server.Server(async_handlers=False)
681+
s.manager.connect('123', '/')
682+
handler = mock.MagicMock()
683+
s.on('my message', handler)
684+
with pytest.raises(ValueError):
685+
s._handle_eio_message(
686+
'999',
687+
'52-["my message","a",'
688+
'{"_placeholder":true,"num":1},'
689+
'{"_placeholder":true,"num":0}]',
690+
)
691+
679692
def test_handle_event_with_ack(self, eio):
680693
s = server.Server(async_handlers=False)
681694
sid = s.manager.connect('123', '/')
@@ -849,6 +862,22 @@ def test_disconnect_twice_namespace(self, eio):
849862
s.disconnect('123', namespace='/foo')
850863
assert calls == s.eio.send.call_count
851864

865+
def test_disconnect_with_partial_binary_packet(self, eio):
866+
s = server.Server()
867+
s._handle_eio_connect('123', 'environ')
868+
s._handle_eio_message('123', '0')
869+
s._handle_eio_message(
870+
'123',
871+
'52-["my message","a",'
872+
'{"_placeholder":true,"num":1},'
873+
'{"_placeholder":true,"num":0}]',
874+
)
875+
s._handle_eio_message('123', b'foo')
876+
assert s._binary_packet['123'] is not None
877+
s.disconnect('1')
878+
s.eio.send.assert_any_call('123', '1')
879+
assert '123' not in s._binary_packet
880+
852881
def test_namespace_handler(self, eio):
853882
result = {}
854883

0 commit comments

Comments
 (0)