From bf09cacd9bd03c2d6f191a6dceb8c57cff2a4c3a Mon Sep 17 00:00:00 2001
From: Todd Wolfson <todd@twolfson.com>
Date: Fri, 30 Oct 2015 14:42:27 -0500
Subject: [PATCH] Added session destruction and regeneration support

---
 flask_session/sessions.py | 27 ++++++++++++++++----
 test_session.py           | 52 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/flask_session/sessions.py b/flask_session/sessions.py
index 1a850788..e7b40133 100644
--- a/flask_session/sessions.py
+++ b/flask_session/sessions.py
@@ -123,14 +123,31 @@ def open_session(self, app, request):
                 return self.session_class(sid=sid)
         return self.session_class(sid=sid)
 
+    def destroy(self, session):
+        # Delete from our session store
+        self.redis.delete(self.key_prefix + session.sid)
+
+        # Remove `sid` so `save_session` knows to delete the cookie
+        session.sid = None
+
+    def regenerate(self, session):
+        # Delete old session info
+        self.redis.delete(self.key_prefix + session.sid)
+
+        # Generate a new sid and mark the session as modified
+        session.sid = self._generate_sid()
+        session.modified = True
+
+        # Session data will be preserved on the `session` dict
+        # `save_session` will take care of updating the cookie
+
     def save_session(self, app, session, response):
         domain = self.get_cookie_domain(app)
         path = self.get_cookie_path(app)
-        if not session:
-            if session.modified:
-                self.redis.delete(self.key_prefix + session.sid)
-                response.delete_cookie(app.session_cookie_name,
-                                       domain=domain, path=path)
+
+        if not session.sid:
+            response.delete_cookie(app.session_cookie_name,
+                                   domain=domain, path=path)
             return
 
         # Modification case.  There are upsides and downsides to
diff --git a/test_session.py b/test_session.py
index 1c0b20c6..a3dd584c 100644
--- a/test_session.py
+++ b/test_session.py
@@ -6,7 +6,12 @@
 
 
 class FlaskSessionTestCase(unittest.TestCase):
-    
+    def _get_cookie_dict(self, test_client):
+        return {
+            cookie.name: cookie
+            for cookie in test_client.cookie_jar
+        }
+
     def test_null_session(self):
         app = flask.Flask(__name__)
         Session(app)
@@ -32,18 +37,57 @@ def set():
             return 'value set'
         @app.route('/get')
         def get():
-            return flask.session['value']
+            return flask.session.get('value', '')
         @app.route('/delete', methods=['POST'])
         def delete():
             del flask.session['value']
             return 'value deleted'
+        @app.route('/destroy', methods=['POST'])
+        def destroy():
+            app.session_interface.destroy(flask.session)
+            return 'session destroyed'
+        @app.route('/regenerate', methods=['POST'])
+        def regenerate():
+            app.session_interface.regenerate(flask.session)
+            return 'session regenerated'
+        @app.errorhandler(500)
+        def errorhandler_500(exc):
+            raise exc
 
+        # Create, retrieve, delete tests
         c = app.test_client()
         self.assertEqual(c.post('/set', data={'value': '42'}).data, b'value set')
         self.assertEqual(c.get('/get').data, b'42')
         c.post('/delete')
-    
-    
+        self.assertEqual(c.get('/get').data, b'')
+
+        # Destruction test
+        # Verify destruction works
+        c = app.test_client()
+        self.assertEqual(c.post('/set', data={'value': '42'}).data, b'value set')
+        session_cookie = self._get_cookie_dict(c)['session']
+        c.post('/destroy')
+        self.assertNotIn('session', self._get_cookie_dict(c))
+
+        # Verify our session was erased from the underlying store
+        # `session=abcdef-original-session-id`
+        cookie_header = 'session={value}'.format(value=session_cookie.value)
+        self.assertEqual(app.test_client().get('/get', headers={'Cookie': cookie_header}).data, b'')
+
+        # Regeneration test
+        # Verify regeneration preserves data but gives us a new session id
+        c = app.test_client()
+        self.assertEqual(c.post('/set', data={'value': '42'}).data, b'value set')
+        original_session_cookie = self._get_cookie_dict(c)['session']
+        c.post('/regenerate')
+        self.assertEqual(c.get('/get').data, b'42')
+        self.assertNotEqual(self._get_cookie_dict(c)['session'].value, original_session_cookie.value)
+
+        # Verify our original session was erased from the underlying store
+        # `session=abcdef-original-session-id`
+        cookie_header = 'session={value}'.format(value=original_session_cookie.value)
+        self.assertEqual(app.test_client().get('/get', headers={'Cookie': cookie_header}).data, b'')
+
     def test_memcached_session(self):
         app = flask.Flask(__name__)
         app.config['SESSION_TYPE'] = 'memcached'