diff --git a/src/socketio/async_pubsub_manager.py b/src/socketio/async_pubsub_manager.py index 9ccc338..1bfcf9d 100644 --- a/src/socketio/async_pubsub_manager.py +++ b/src/socketio/async_pubsub_manager.py @@ -1,10 +1,12 @@ import asyncio +import base64 from functools import partial import uuid from engineio import json from .async_manager import AsyncManager +from .packet import Packet class AsyncPubSubManager(AsyncManager): @@ -64,8 +66,12 @@ class AsyncPubSubManager(AsyncManager): callback = (room, namespace, id) else: callback = None + binary = Packet.data_is_binary(data) + if binary: + data, attachments = Packet.deconstruct_binary(data) + data = [data, *[base64.b64encode(a).decode() for a in attachments]] message = {'method': 'emit', 'event': event, 'data': data, - 'namespace': namespace, 'room': room, + 'binary': binary, 'namespace': namespace, 'room': room, 'skip_sid': skip_sid, 'callback': callback, 'host_id': self.host_id} await self._handle_emit(message) # handle in this host @@ -145,7 +151,11 @@ class AsyncPubSubManager(AsyncManager): *remote_callback) else: callback = None - await super().emit(message['event'], message['data'], + data = message['data'] + if message.get('binary'): + attachments = [base64.b64decode(a) for a in data[1:]] + data = Packet.reconstruct_binary(data[0], attachments) + await super().emit(message['event'], data, namespace=message.get('namespace'), room=message.get('room'), skip_sid=message.get('skip_sid'), diff --git a/src/socketio/packet.py b/src/socketio/packet.py index f7ad87e..3deba7f 100644 --- a/src/socketio/packet.py +++ b/src/socketio/packet.py @@ -29,7 +29,7 @@ class Packet: self.namespace = namespace self.id = id if self.uses_binary_events and \ - (binary or (binary is None and self._data_is_binary( + (binary or (binary is None and self.data_is_binary( self.data))): if self.packet_type == EVENT: self.packet_type = BINARY_EVENT @@ -51,7 +51,7 @@ class Packet: """ encoded_packet = str(self.packet_type) if self.packet_type == BINARY_EVENT or self.packet_type == BINARY_ACK: - data, attachments = self._deconstruct_binary(self.data) + data, attachments = self.deconstruct_binary(self.data) encoded_packet += str(len(attachments)) + '-' else: data = self.data @@ -119,61 +119,65 @@ class Packet: raise ValueError('Unexpected binary attachment') self.attachments.append(attachment) if self.attachment_count == len(self.attachments): - self.reconstruct_binary(self.attachments) + self.data = self.reconstruct_binary(self.data, self.attachments) return True return False - def reconstruct_binary(self, attachments): + @classmethod + def reconstruct_binary(cls, data, attachments): """Reconstruct a decoded packet using the given list of binary attachments. """ - self.data = self._reconstruct_binary_internal(self.data, - self.attachments) + return cls._reconstruct_binary_internal(data, attachments) - def _reconstruct_binary_internal(self, data, attachments): + @classmethod + def _reconstruct_binary_internal(cls, data, attachments): if isinstance(data, list): - return [self._reconstruct_binary_internal(item, attachments) + return [cls._reconstruct_binary_internal(item, attachments) for item in data] elif isinstance(data, dict): if data.get('_placeholder') and 'num' in data: return attachments[data['num']] else: - return {key: self._reconstruct_binary_internal(value, - attachments) + return {key: cls._reconstruct_binary_internal(value, + attachments) for key, value in data.items()} else: return data - def _deconstruct_binary(self, data): + @classmethod + def deconstruct_binary(cls, data): """Extract binary components in the packet.""" attachments = [] - data = self._deconstruct_binary_internal(data, attachments) + data = cls._deconstruct_binary_internal(data, attachments) return data, attachments - def _deconstruct_binary_internal(self, data, attachments): + @classmethod + def _deconstruct_binary_internal(cls, data, attachments): if isinstance(data, bytes): attachments.append(data) return {'_placeholder': True, 'num': len(attachments) - 1} elif isinstance(data, list): - return [self._deconstruct_binary_internal(item, attachments) + return [cls._deconstruct_binary_internal(item, attachments) for item in data] elif isinstance(data, dict): - return {key: self._deconstruct_binary_internal(value, attachments) + return {key: cls._deconstruct_binary_internal(value, attachments) for key, value in data.items()} else: return data - def _data_is_binary(self, data): + @classmethod + def data_is_binary(cls, data): """Check if the data contains binary components.""" if isinstance(data, bytes): return True elif isinstance(data, list): return functools.reduce( - lambda a, b: a or b, [self._data_is_binary(item) + lambda a, b: a or b, [cls.data_is_binary(item) for item in data], False) elif isinstance(data, dict): return functools.reduce( - lambda a, b: a or b, [self._data_is_binary(item) + lambda a, b: a or b, [cls.data_is_binary(item) for item in data.values()], False) else: diff --git a/src/socketio/pubsub_manager.py b/src/socketio/pubsub_manager.py index 80744f2..3528d22 100644 --- a/src/socketio/pubsub_manager.py +++ b/src/socketio/pubsub_manager.py @@ -1,9 +1,11 @@ +import base64 from functools import partial import uuid from engineio import json from .manager import Manager +from .packet import Packet class PubSubManager(Manager): @@ -61,8 +63,12 @@ class PubSubManager(Manager): callback = (room, namespace, id) else: callback = None + binary = Packet.data_is_binary(data) + if binary: + data, attachments = Packet.deconstruct_binary(data) + data = [data, *[base64.b64encode(a).decode() for a in attachments]] message = {'method': 'emit', 'event': event, 'data': data, - 'namespace': namespace, 'room': room, + 'binary': binary, 'namespace': namespace, 'room': room, 'skip_sid': skip_sid, 'callback': callback, 'host_id': self.host_id} self._handle_emit(message) # handle in this host @@ -141,7 +147,11 @@ class PubSubManager(Manager): *remote_callback) else: callback = None - super().emit(message['event'], message['data'], + data = message['data'] + if message.get('binary'): + attachments = [base64.b64decode(a) for a in data[1:]] + data = Packet.reconstruct_binary(data[0], attachments) + super().emit(message['event'], data, namespace=message.get('namespace'), room=message.get('room'), skip_sid=message.get('skip_sid'), callback=callback) diff --git a/tests/async/test_pubsub_manager.py b/tests/async/test_pubsub_manager.py index 4a7012e..5d34e53 100644 --- a/tests/async/test_pubsub_manager.py +++ b/tests/async/test_pubsub_manager.py @@ -57,6 +57,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': None, @@ -66,6 +67,36 @@ class TestAsyncPubSubManager: } ) + async def test_emit_binary(self): + await self.pm.emit('foo', b'bar') + self.pm._publish.assert_awaited_once_with( + { + 'method': 'emit', + 'event': 'foo', + 'binary': True, + 'data': [{'_placeholder': True, 'num': 0}, 'YmFy'], + 'namespace': '/', + 'room': None, + 'skip_sid': None, + 'callback': None, + 'host_id': '123456', + } + ) + await self.pm.emit('foo', {'foo': b'bar'}) + self.pm._publish.assert_awaited_with( + { + 'method': 'emit', + 'event': 'foo', + 'binary': True, + 'data': [{'foo': {'_placeholder': True, 'num': 0}}, 'YmFy'], + 'namespace': '/', + 'room': None, + 'skip_sid': None, + 'callback': None, + 'host_id': '123456', + } + ) + async def test_emit_with_to(self): sid = 'room-mate' await self.pm.emit('foo', 'bar', to=sid) @@ -73,6 +104,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': sid, @@ -88,6 +120,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/baz', 'room': None, @@ -103,6 +136,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': 'baz', @@ -118,6 +152,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': None, @@ -136,6 +171,7 @@ class TestAsyncPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': 'baz', @@ -241,6 +277,37 @@ class TestAsyncPubSubManager: callback=None, ) + async def test_handle_emit_binary(self): + with mock.patch.object( + async_manager.AsyncManager, 'emit' + ) as super_emit: + await self.pm._handle_emit({ + 'event': 'foo', + 'binary': True, + 'data': [{'_placeholder': True, 'num': 0}, 'YmFy'], + }) + super_emit.assert_awaited_once_with( + 'foo', + b'bar', + namespace=None, + room=None, + skip_sid=None, + callback=None, + ) + await self.pm._handle_emit({ + 'event': 'foo', + 'binary': True, + 'data': [{'foo': {'_placeholder': True, 'num': 0}}, 'YmFy'], + }) + super_emit.assert_awaited_with( + 'foo', + {'foo': b'bar'}, + namespace=None, + room=None, + skip_sid=None, + callback=None, + ) + async def test_handle_emit_with_namespace(self): with mock.patch.object( async_manager.AsyncManager, 'emit' diff --git a/tests/common/test_packet.py b/tests/common/test_packet.py index 5682dab..e15babf 100644 --- a/tests/common/test_packet.py +++ b/tests/common/test_packet.py @@ -266,16 +266,24 @@ class TestPacket: assert pkt.data["a"] == "0123456789-" assert pkt.attachment_count == 0 + def test_deconstruct_binary(self): + datas = [b'foo', [b'foo', b'bar'], ['foo', b'bar'], {'foo': b'bar'}, + {'foo': 'bar', 'baz': b'qux'}, {'foo': [b'bar']}] + for data in datas: + bdata, attachments = packet.Packet.deconstruct_binary(data) + rdata = packet.Packet.reconstruct_binary(bdata, attachments) + assert data == rdata + def test_data_is_binary_list(self): pkt = packet.Packet() - assert not pkt._data_is_binary(['foo']) - assert not pkt._data_is_binary([]) - assert pkt._data_is_binary([b'foo']) - assert pkt._data_is_binary(['foo', b'bar']) + assert not pkt.data_is_binary(['foo']) + assert not pkt.data_is_binary([]) + assert pkt.data_is_binary([b'foo']) + assert pkt.data_is_binary(['foo', b'bar']) def test_data_is_binary_dict(self): pkt = packet.Packet() - assert not pkt._data_is_binary({'a': 'foo'}) - assert not pkt._data_is_binary({}) - assert pkt._data_is_binary({'a': b'foo'}) - assert pkt._data_is_binary({'a': 'foo', 'b': b'bar'}) + assert not pkt.data_is_binary({'a': 'foo'}) + assert not pkt.data_is_binary({}) + assert pkt.data_is_binary({'a': b'foo'}) + assert pkt.data_is_binary({'a': 'foo', 'b': b'bar'}) diff --git a/tests/common/test_pubsub_manager.py b/tests/common/test_pubsub_manager.py index 32c7d41..5ef5473 100644 --- a/tests/common/test_pubsub_manager.py +++ b/tests/common/test_pubsub_manager.py @@ -69,6 +69,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': None, @@ -78,6 +79,36 @@ class TestPubSubManager: } ) + def test_emit_binary(self): + self.pm.emit('foo', b'bar') + self.pm._publish.assert_called_once_with( + { + 'method': 'emit', + 'event': 'foo', + 'binary': True, + 'data': [{'_placeholder': True, 'num': 0}, 'YmFy'], + 'namespace': '/', + 'room': None, + 'skip_sid': None, + 'callback': None, + 'host_id': '123456', + } + ) + self.pm.emit('foo', {'foo': b'bar'}) + self.pm._publish.assert_called_with( + { + 'method': 'emit', + 'event': 'foo', + 'binary': True, + 'data': [{'foo': {'_placeholder': True, 'num': 0}}, 'YmFy'], + 'namespace': '/', + 'room': None, + 'skip_sid': None, + 'callback': None, + 'host_id': '123456', + } + ) + def test_emit_with_to(self): sid = "ferris" self.pm.emit('foo', 'bar', to=sid) @@ -85,6 +116,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': sid, @@ -100,6 +132,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/baz', 'room': None, @@ -115,6 +148,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': 'baz', @@ -130,6 +164,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': None, @@ -148,6 +183,7 @@ class TestPubSubManager: { 'method': 'emit', 'event': 'foo', + 'binary': False, 'data': 'bar', 'namespace': '/', 'room': 'baz', @@ -250,6 +286,35 @@ class TestPubSubManager: callback=None, ) + def test_handle_emit_binary(self): + with mock.patch.object(manager.Manager, 'emit') as super_emit: + self.pm._handle_emit({ + 'event': 'foo', + 'binary': True, + 'data': [{'_placeholder': True, 'num': 0}, 'YmFy'], + }) + super_emit.assert_called_once_with( + 'foo', + b'bar', + namespace=None, + room=None, + skip_sid=None, + callback=None, + ) + self.pm._handle_emit({ + 'event': 'foo', + 'binary': True, + 'data': [{'foo': {'_placeholder': True, 'num': 0}}, 'YmFy'], + }) + super_emit.assert_called_with( + 'foo', + {'foo': b'bar'}, + namespace=None, + room=None, + skip_sid=None, + callback=None, + ) + def test_handle_emit_with_namespace(self): with mock.patch.object(manager.Manager, 'emit') as super_emit: self.pm._handle_emit(