From 6f270eef955e28289285c03138187511bc148547 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Sun, 7 Aug 2022 19:46:21 +0100
Subject: [PATCH] Support new Flask 2.2 session structure (Fixes #1856)

---
 example/requirements.txt       | 21 +++++++++++++--------
 src/flask_socketio/__init__.py | 11 +++++++++--
 test_socketio.py               | 24 ++++++++++++++++++++++--
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/example/requirements.txt b/example/requirements.txt
index 421b91b8..740b70fb 100644
--- a/example/requirements.txt
+++ b/example/requirements.txt
@@ -1,11 +1,16 @@
-Flask==1.0.2
-Flask-Login==0.4.1
-Flask-Session==0.3.1
-Flask_SocketIO
-itsdangerous==1.1.0
-Jinja2==2.10
-MarkupSafe==1.1.0
+bidict==0.22.0
+cachelib==0.9.0
+click==8.1.3
+Flask==2.2.1
+Flask-Login==0.6.2
+Flask-Session==0.4.0
+Flask-SocketIO
+importlib-metadata==4.12.0
+itsdangerous==2.1.2
+Jinja2==3.1.2
+MarkupSafe==2.1.1
 python-engineio
 python-socketio
 six==1.11.0
-Werkzeug==0.14.1
+Werkzeug==2.2.1
+zipp==3.8.1
diff --git a/src/flask_socketio/__init__.py b/src/flask_socketio/__init__.py
index 570add47..83a93ffa 100644
--- a/src/flask_socketio/__init__.py
+++ b/src/flask_socketio/__init__.py
@@ -16,7 +16,7 @@
     sys.exit(1)
 
 import flask
-from flask import _request_ctx_stack, has_request_context, json as flask_json
+from flask import has_request_context, json as flask_json
 from flask.sessions import SessionMixin
 import socketio
 from socketio.exceptions import ConnectionRefusedError  # noqa: F401
@@ -797,6 +797,14 @@ def _handle_event(self, handler, message, namespace, sid, *args):
                 if 'saved_session' not in environ:
                     environ['saved_session'] = _ManagedSession(flask.session)
                 session_obj = environ['saved_session']
+                if hasattr(flask, 'globals') and \
+                        hasattr(flask.globals, 'request_ctx'):
+                    # update session for Flask >= 2.2
+                    ctx = flask.globals.request_ctx._get_current_object()
+                else:  # pragma: no cover
+                    # update session for Flask < 2.2
+                    ctx = flask._request_ctx_stack.top
+                ctx.session = session_obj
             else:
                 # let Flask handle the user session
                 # for cookie based sessions, this effectively freezes the
@@ -804,7 +812,6 @@ def _handle_event(self, handler, message, namespace, sid, *args):
                 # for server-side sessions, this allows HTTP and Socket.IO to
                 # share the session, with both having read/write access to it
                 session_obj = flask.session._get_current_object()
-            _request_ctx_stack.top.session = session_obj
             flask.request.sid = sid
             flask.request.namespace = namespace
             flask.request.event = {'message': message, 'args': args}
diff --git a/test_socketio.py b/test_socketio.py
index 14e590ff..7c1d952c 100644
--- a/test_socketio.py
+++ b/test_socketio.py
@@ -48,7 +48,12 @@ def on_disconnect_test():
 def message(message):
     send(message)
     if message == 'test session':
-        session['a'] = 'b'
+        if not socketio.manage_session and 'a' in session:
+            raise RuntimeError('session is being stored')
+        if 'a' not in session:
+            session['a'] = 'b'
+        else:
+            session['a'] = 'c'
     if message not in "test noackargs":
         return message
 
@@ -474,7 +479,7 @@ def test_broadcast_namespace(self):
         self.assertEqual(received[0]['args'][0]['a'], 'b')
         self.assertEqual(len(client3.get_received()), 0)
 
-    def test_session(self):
+    def test_managed_session(self):
         flask_client = app.test_client()
         flask_client.get('/session')
         client = socketio.test_client(app, flask_test_client=flask_client,
@@ -488,6 +493,21 @@ def test_session(self):
         self.assertEqual(
             socketio.server.environ[client.eio_sid]['saved_session'],
             {'a': 'b', 'foo': 'bar'})
+        client.send('test session')
+        self.assertEqual(
+            socketio.server.environ[client.eio_sid]['saved_session'],
+            {'a': 'c', 'foo': 'bar'})
+
+    def test_unmanaged_session(self):
+        socketio.manage_session = False
+        flask_client = app.test_client()
+        flask_client.get('/session')
+        client = socketio.test_client(app, flask_test_client=flask_client,
+                                      auth={'foo': 'bar'})
+        client.get_received()
+        client.send('test session')
+        client.send('test session')
+        socketio.manage_session = True
 
     def test_room(self):
         client1 = socketio.test_client(app, auth={'foo': 'bar'})