From da2d141e8d68fd698ddb2dd5d54c4b1a7622d4df Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguelgrinberg50@gmail.com>
Date: Sun, 4 Sep 2016 12:31:58 -0700
Subject: [PATCH] put clients in a pre-disconnect state while their disconnect
 handler runs

This avoids potential endless recursion.
See https://github.com/miguelgrinberg/Flask-SocketIO/issues/312
---
 socketio/base_manager.py   | 21 +++++++++++++++++++++
 socketio/server.py         |  1 +
 tests/test_base_manager.py | 14 ++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/socketio/base_manager.py b/socketio/base_manager.py
index cdbaa80..5b98990 100644
--- a/socketio/base_manager.py
+++ b/socketio/base_manager.py
@@ -16,6 +16,7 @@ class BaseManager(object):
         self.server = None
         self.rooms = {}
         self.callbacks = {}
+        self.pending_disconnect = {}
 
     def initialize(self, server):
         self.server = server
@@ -35,11 +36,26 @@ class BaseManager(object):
         self.enter_room(sid, namespace, sid)
 
     def is_connected(self, sid, namespace):
+        if namespace in self.pending_disconnect and \
+                sid in self.pending_disconnect[namespace]:
+            # the client is in the process of being disconnected
+            return False
         try:
             return self.rooms[namespace][None][sid]
         except KeyError:
             pass
 
+    def pre_disconnect(self, sid, namespace):
+        """Put the client in the to-be-disconnected list.
+
+        This allows the client data structures to be present while the
+        disconnect handler is invoked, but still recognize the fact that the
+        client is soon going away.
+        """
+        if namespace not in self.pending_disconnect:
+            self.pending_disconnect[namespace] = []
+        self.pending_disconnect[namespace].append(sid)
+
     def disconnect(self, sid, namespace):
         """Register a client disconnect from a namespace."""
         rooms = []
@@ -52,6 +68,11 @@ class BaseManager(object):
             del self.callbacks[sid][namespace]
             if len(self.callbacks[sid]) == 0:
                 del self.callbacks[sid]
+        if namespace in self.pending_disconnect and \
+                sid in self.pending_disconnect[namespace]:
+            self.pending_disconnect[namespace].remove(sid)
+            if len(self.pending_disconnect[namespace]) == 0:
+                del self.pending_disconnect[namespace]
 
     def enter_room(self, sid, namespace, room):
         """Add a client to a room."""
diff --git a/socketio/server.py b/socketio/server.py
index 21e067e..52e0f37 100644
--- a/socketio/server.py
+++ b/socketio/server.py
@@ -298,6 +298,7 @@ class Server(object):
         namespace = namespace or '/'
         if self.manager.is_connected(sid, namespace=namespace):
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
+            self.manager.pre_disconnect(sid, namespace=namespace)
             self._send_packet(sid, packet.Packet(packet.DISCONNECT,
                                                  namespace=namespace))
             self._trigger_event('disconnect', namespace, sid)
diff --git a/tests/test_base_manager.py b/tests/test_base_manager.py
index b916ddc..2a801be 100644
--- a/tests/test_base_manager.py
+++ b/tests/test_base_manager.py
@@ -24,6 +24,20 @@ class TestBaseManager(unittest.TestCase):
         self.assertEqual(self.bm.rooms['/foo'], {None: {'123': True},
                                                  '123': {'123': True}})
 
+    def test_pre_disconnect(self):
+        self.bm.connect('123', '/foo')
+        self.bm.connect('456', '/foo')
+        self.bm.pre_disconnect('123', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['123']})
+        self.assertFalse(self.bm.is_connected('123', '/foo'))
+        self.bm.pre_disconnect('456', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['123', '456']})
+        self.assertFalse(self.bm.is_connected('456', '/foo'))
+        self.bm.disconnect('123', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['456']})
+        self.bm.disconnect('456', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {})
+
     def test_disconnect(self):
         self.bm.connect('123', '/foo')
         self.bm.connect('456', '/foo')