v5 protocol: handle per-namespace sids in base manager

Author: miguelgrinberg

setup.py

@@ -30,6 +30,7 @@
     platforms='any',
     install_requires=[
         'six>=1.9.0',
+        'bidict>=0.21.0',
         'python-engineio>=3.13.0,<4'
     ],
     extras_require={

socketio/asyncio_client.py

@@ -353,7 +353,7 @@ async def _handle_error(self, namespace, data):
         if namespace in self.namespaces:
             self.namespaces.remove(namespace)
         if namespace == '/':
-            self.namespaces = []
+            self.namespaces = {}
             self.connected = False
 
     async def _trigger_event(self, event, namespace, *args):
@@ -456,7 +456,7 @@ async def _handle_eio_disconnect(self):
         if self.connected:
             for n in self.namespaces:
                 await self._trigger_event('disconnect', namespace=n)
-            self.namespaces = []
+            self.namespaces = {}
             self.connected = False
         self.callbacks = {}
         self._binary_packet = None

socketio/asyncio_manager.py

@@ -20,13 +20,13 @@ async def emit(self, event, data, namespace, room=None, skip_sid=None,
         tasks = []
         if not isinstance(skip_sid, list):
             skip_sid = [skip_sid]
-        for sid in self.get_participants(namespace, room):
+        for sid, eio_sid in self.get_participants(namespace, room):
             if sid not in skip_sid:
                 if callback is not None:
                     id = self._generate_ack_id(sid, namespace, callback)
                 else:
                     id = None
-                tasks.append(self.server._emit_internal(sid, event, data,
+                tasks.append(self.server._emit_internal(eio_sid, event, data,
                                                         namespace, id))
         if tasks == []:  # pragma: no cover
             return

socketio/asyncio_server.py

@@ -336,13 +336,11 @@ async def disconnect(self, sid, namespace=None, ignore_queue=False):
             delete_it = await self.manager.can_disconnect(sid, namespace)
         if delete_it:
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
-            self.manager.pre_disconnect(sid, namespace=namespace)
-            await self._send_packet(sid, packet.Packet(packet.DISCONNECT,
-                                                       namespace=namespace))
+            eio_sid = self.manager.pre_disconnect(sid, namespace=namespace)
+            await self._send_packet(eio_sid, packet.Packet(
+                packet.DISCONNECT, namespace=namespace))
             await self._trigger_event('disconnect', namespace, sid)
             self.manager.disconnect(sid, namespace=namespace)
-            if namespace == '/':
-                await self.eio.disconnect(sid)
 
     async def handle_request(self, *args, **kwargs):
         """Handle an HTTP request from the client.
@@ -396,67 +394,61 @@ async def _emit_internal(self, sid, event, data, namespace=None, id=None):
         await self._send_packet(sid, packet.Packet(
             packet.EVENT, namespace=namespace, data=[event] + data, id=id))
 
-    async def _send_packet(self, sid, pkt):
+    async def _send_packet(self, eio_sid, pkt):
         """Send a Socket.IO packet to a client."""
         encoded_packet = pkt.encode()
         if isinstance(encoded_packet, list):
             for ep in encoded_packet:
-                await self.eio.send(sid, ep)
+                await self.eio.send(eio_sid, ep)
         else:
-            await self.eio.send(sid, encoded_packet)
+            await self.eio.send(eio_sid, encoded_packet)
 
-    async def _handle_connect(self, sid, namespace):
+    async def _handle_connect(self, eio_sid, namespace):
         """Handle a client connection request."""
         namespace = namespace or '/'
-        self.manager.connect(sid, namespace)
+        sid = self.manager.connect(eio_sid, namespace)
         if self.always_connect:
-            await self._send_packet(sid, packet.Packet(packet.CONNECT,
-                                                       namespace=namespace))
+            await self._send_packet(eio_sid, packet.Packet(
+                packet.CONNECT, {'sid': sid}, namespace=namespace))
         fail_reason = None
         try:
             success = await self._trigger_event('connect', namespace, sid,
-                                                self.environ[sid])
+                                                self.environ[eio_sid])
         except exceptions.ConnectionRefusedError as exc:
             fail_reason = exc.error_args
             success = False
 
         if success is False:
             if self.always_connect:
                 self.manager.pre_disconnect(sid, namespace)
-                await self._send_packet(sid, packet.Packet(
+                await self._send_packet(eio_sid, packet.Packet(
                     packet.DISCONNECT, data=fail_reason, namespace=namespace))
             elif namespace != '/':
-                await self._send_packet(sid, packet.Packet(
+                await self._send_packet(eio_sid, packet.Packet(
                     packet.CONNECT_ERROR, data=fail_reason,
                     namespace=namespace))
             self.manager.disconnect(sid, namespace)
-            if namespace == '/' and sid in self.environ:  # pragma: no cover
-                del self.environ[sid]
+            if namespace == '/' and \
+                    eio_sid in self.environ:  # pragma: no cover
+                del self.environ[eio_sid]
             return fail_reason or False
         elif not self.always_connect:
-            await self._send_packet(sid, packet.Packet(packet.CONNECT,
-                                                       namespace=namespace))
+            await self._send_packet(eio_sid, packet.Packet(
+                packet.CONNECT, {'sid': sid}, namespace=namespace))
 
-    async def _handle_disconnect(self, sid, namespace):
+    async def _handle_disconnect(self, eio_sid, namespace):
         """Handle a client disconnect."""
         namespace = namespace or '/'
-        if namespace == '/':
-            namespace_list = list(self.manager.get_namespaces())
-        else:
-            namespace_list = [namespace]
-        for n in namespace_list:
-            if n != '/' and self.manager.is_connected(sid, n):
-                self.manager.pre_disconnect(sid, namespace=n)
-                await self._trigger_event('disconnect', n, sid)
-                self.manager.disconnect(sid, n)
-        if namespace == '/' and self.manager.is_connected(sid, namespace):
-            self.manager.pre_disconnect(sid, namespace='/')
-            await self._trigger_event('disconnect', '/', sid)
-            self.manager.disconnect(sid, '/')
-
-    async def _handle_event(self, sid, namespace, id, data):
+        sid = self.manager.sid_from_eio_sid(eio_sid, namespace)
+        if self.manager.is_connected(sid, namespace):
+            self.manager.pre_disconnect(sid, namespace=namespace)
+            await self._trigger_event('disconnect', namespace, sid)
+            self.manager.disconnect(sid, namespace)
+
+    async def _handle_event(self, eio_sid, namespace, id, data):
         """Handle an incoming client event."""
         namespace = namespace or '/'
+        sid = self.manager.sid_from_eio_sid(eio_sid, namespace)
         self.logger.info('received event "%s" from %s [%s]', data[0], sid,
                          namespace)
         if not self.manager.is_connected(sid, namespace):
@@ -465,11 +457,13 @@ async def _handle_event(self, sid, namespace, id, data):
             return
         if self.async_handlers:
             self.start_background_task(self._handle_event_internal, self, sid,
-                                       data, namespace, id)
+                                       eio_sid, data, namespace, id)
         else:
-            await self._handle_event_internal(self, sid, data, namespace, id)
+            await self._handle_event_internal(self, sid, eio_sid, data,
+                                              namespace, id)
 
-    async def _handle_event_internal(self, server, sid, data, namespace, id):
+    async def _handle_event_internal(self, server, sid, eio_sid, data,
+                                     namespace, id):
         r = await server._trigger_event(data[0], namespace, sid, *data[1:])
         if id is not None:
             # send ACK packet with the response returned by the handler
@@ -480,13 +474,13 @@ async def _handle_event_internal(self, server, sid, data, namespace, id):
                 data = list(r)
             else:
                 data = [r]
-            await server._send_packet(sid, packet.Packet(packet.ACK,
-                                                         namespace=namespace,
-                                                         id=id, data=data))
+            await server._send_packet(eio_sid, packet.Packet(
+                packet.ACK, namespace=namespace, id=id, data=data))
 
-    async def _handle_ack(self, sid, namespace, id, data):
+    async def _handle_ack(self, eio_sid, namespace, id, data):
         """Handle ACK packets from the client."""
         namespace = namespace or '/'
+        sid = self.manager.sid_from_eio_sid(eio_sid, namespace)
         self.logger.info('received ack from %s [%s]', sid, namespace)
         await self.manager.trigger_callback(sid, namespace, id, data)
 
@@ -509,48 +503,50 @@ async def _trigger_event(self, event, namespace, *args):
             return await self.namespace_handlers[namespace].trigger_event(
                 event, *args)
 
-    async def _handle_eio_connect(self, sid, environ):
+    async def _handle_eio_connect(self, eio_sid, environ):
         """Handle the Engine.IO connection event."""
         if not self.manager_initialized:
             self.manager_initialized = True
             self.manager.initialize()
-        self.environ[sid] = environ
+        self.environ[eio_sid] = environ
 
-    async def _handle_eio_message(self, sid, data):
+    async def _handle_eio_message(self, eio_sid, data):
         """Dispatch Engine.IO messages."""
-        if sid in self._binary_packet:
-            pkt = self._binary_packet[sid]
+        if eio_sid in self._binary_packet:
+            pkt = self._binary_packet[eio_sid]
             if pkt.add_attachment(data):
-                del self._binary_packet[sid]
+                del self._binary_packet[eio_sid]
                 if pkt.packet_type == packet.BINARY_EVENT:
-                    await self._handle_event(sid, pkt.namespace, pkt.id,
+                    await self._handle_event(eio_sid, pkt.namespace, pkt.id,
                                              pkt.data)
                 else:
-                    await self._handle_ack(sid, pkt.namespace, pkt.id,
+                    await self._handle_ack(eio_sid, pkt.namespace, pkt.id,
                                            pkt.data)
         else:
             pkt = packet.Packet(encoded_packet=data)
             if pkt.packet_type == packet.CONNECT:
-                await self._handle_connect(sid, pkt.namespace)
+                await self._handle_connect(eio_sid, pkt.namespace)
             elif pkt.packet_type == packet.DISCONNECT:
-                await self._handle_disconnect(sid, pkt.namespace)
+                await self._handle_disconnect(eio_sid, pkt.namespace)
             elif pkt.packet_type == packet.EVENT:
-                await self._handle_event(sid, pkt.namespace, pkt.id, pkt.data)
+                await self._handle_event(eio_sid, pkt.namespace, pkt.id,
+                                         pkt.data)
             elif pkt.packet_type == packet.ACK:
-                await self._handle_ack(sid, pkt.namespace, pkt.id, pkt.data)
+                await 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:
-                self._binary_packet[sid] = pkt
+                self._binary_packet[eio_sid] = pkt
             elif pkt.packet_type == packet.CONNECT_ERROR:
                 raise ValueError('Unexpected CONNECT_ERROR packet.')
             else:
                 raise ValueError('Unknown packet type.')
 
-    async def _handle_eio_disconnect(self, sid):
+    async def _handle_eio_disconnect(self, eio_sid):
         """Handle Engine.IO disconnect event."""
-        await self._handle_disconnect(sid, '/')
-        if sid in self.environ:
-            del self.environ[sid]
+        await self._handle_disconnect(eio_sid, '/')
+        if eio_sid in self.environ:
+            del self.environ[eio_sid]
 
     def _engineio_server_class(self):
         return engineio.AsyncServer

socketio/base_manager.py

@@ -1,6 +1,7 @@
 import itertools
 import logging
 
+from bidict import bidict
 import six
 
 default_logger = logging.getLogger('socketio')
@@ -18,7 +19,8 @@ class BaseManager(object):
     def __init__(self):
         self.logger = None
         self.server = None
-        self.rooms = {}
+        self.rooms = {}  # self.rooms[namespace][room][sio_sid] = eio_sid
+        self.eio_to_sid = {}
         self.callbacks = {}
         self.pending_disconnect = {}
 
@@ -37,13 +39,15 @@ def get_namespaces(self):
 
     def get_participants(self, namespace, room):
         """Return an iterable with the active participants in a room."""
-        for sid, active in six.iteritems(self.rooms[namespace][room].copy()):
-            yield sid
+        for sid, eio_sid in self.rooms[namespace][room].copy().items():
+            yield sid, eio_sid
 
-    def connect(self, sid, namespace):
+    def connect(self, eio_sid, namespace):
         """Register a client connection to a namespace."""
-        self.enter_room(sid, namespace, None)
-        self.enter_room(sid, namespace, sid)
+        sid = self.server.eio.generate_id()
+        self.enter_room(sid, namespace, None, eio_sid=eio_sid)
+        self.enter_room(sid, namespace, sid, eio_sid=eio_sid)
+        return sid
 
     def is_connected(self, sid, namespace):
         if namespace in self.pending_disconnect and \
@@ -55,6 +59,9 @@ def is_connected(self, sid, namespace):
         except KeyError:
             pass
 
+    def sid_from_eio_sid(self, eio_sid, namespace):
+        return self.rooms[namespace][None].inverse.get(eio_sid)
+
     def can_disconnect(self, sid, namespace):
         return self.is_connected(sid, namespace)
 
@@ -68,6 +75,7 @@ def pre_disconnect(self, sid, namespace):
         if namespace not in self.pending_disconnect:
             self.pending_disconnect[namespace] = []
         self.pending_disconnect[namespace].append(sid)
+        return self.rooms[namespace][None].get(sid)
 
     def disconnect(self, sid, namespace):
         """Register a client disconnect from a namespace."""
@@ -89,13 +97,15 @@ def disconnect(self, sid, namespace):
             if len(self.pending_disconnect[namespace]) == 0:
                 del self.pending_disconnect[namespace]
 
-    def enter_room(self, sid, namespace, room):
+    def enter_room(self, sid, namespace, room, eio_sid=None):
         """Add a client to a room."""
         if namespace not in self.rooms:
             self.rooms[namespace] = {}
         if room not in self.rooms[namespace]:
-            self.rooms[namespace][room] = {}
-        self.rooms[namespace][room][sid] = True
+            self.rooms[namespace][room] = bidict()
+        if eio_sid is None:
+            eio_sid = self.rooms[namespace][None][sid]
+        self.rooms[namespace][room][sid] = eio_sid
 
     def leave_room(self, sid, namespace, room):
         """Remove a client from a room."""
@@ -111,7 +121,7 @@ def leave_room(self, sid, namespace, room):
     def close_room(self, room, namespace):
         """Remove all participants from a room."""
         try:
-            for sid in self.get_participants(namespace, room):
+            for sid, _ in self.get_participants(namespace, room):
                 self.leave_room(sid, namespace, room)
         except KeyError:
             pass
@@ -121,7 +131,7 @@ def get_rooms(self, sid, namespace):
         r = []
         try:
             for room_name, room in six.iteritems(self.rooms[namespace]):
-                if room_name is not None and sid in room and room[sid]:
+                if room_name is not None and sid in room:
                     r.append(room_name)
         except KeyError:
             pass
@@ -135,13 +145,13 @@ def emit(self, event, data, namespace, room=None, skip_sid=None,
             return
         if not isinstance(skip_sid, list):
             skip_sid = [skip_sid]
-        for sid in self.get_participants(namespace, room):
+        for sid, eio_sid in self.get_participants(namespace, room):
             if sid not in skip_sid:
                 if callback is not None:
                     id = self._generate_ack_id(sid, namespace, callback)
                 else:
                     id = None
-                self.server._emit_internal(sid, event, data, namespace, id)
+                self.server._emit_internal(eio_sid, event, data, namespace, id)
 
     def trigger_callback(self, sid, namespace, id, data):
         """Invoke an application callback."""

socketio/client.py

@@ -533,7 +533,7 @@ def _handle_error(self, namespace, data):
         if namespace in self.namespaces:
             del self.namespaces[namespace]
         if namespace == '/':
-            self.namespaces = []
+            self.namespaces = {}
             self.connected = False
 
     def _trigger_event(self, event, namespace, *args):
@@ -625,7 +625,7 @@ def _handle_eio_disconnect(self):
         if self.connected:
             for n in self.namespaces:
                 self._trigger_event('disconnect', namespace=n)
-            self.namespaces = []
+            self.namespaces = {}
             self.connected = False
         self.callbacks = {}
         self._binary_packet = None