From b115f8fa04e93562b69c2dfe968922cadc430445 Mon Sep 17 00:00:00 2001
From: Sam Bolgert <sbolgert@gmail.com>
Date: Thu, 29 Sep 2016 11:08:44 -0700
Subject: [PATCH] Update channel_session decorator to rehydrate http_session
 (#348)

* Update channel_session decorator to rehydrate http_session

Update the http_session decorator to write the http session key
to the channel_session when available. This allows the channel_session
decorator to rehydrate the http_session after the initial websocket
connection.

closes #318

* Add persist=True option to http_session

* Add explicit option to store the session key in the channel session
* Update docs
* Add test case

* Add channel_and_http_session decorator

This decorator enables both sessions and maintains
the http_session for the lifetime of the websocket connection.
---
 channels/sessions.py            | 24 +++++++++++++++++++++
 channels/tests/test_sessions.py | 38 ++++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/channels/sessions.py b/channels/sessions.py
index ea2e696..8e3c254 100644
--- a/channels/sessions.py
+++ b/channels/sessions.py
@@ -183,3 +183,27 @@ def http_session(func):
             session.save()
         return result
     return inner
+
+
+def channel_and_http_session(func):
+    """
+    Enables both the channel_session and http_session.
+
+    Stores the http session key in the channel_session on websocket.connect messages.
+    It will then hydrate the http_session from that same key on subsequent messages.
+    """
+    @http_session
+    @channel_session
+    @functools.wraps(func)
+    def inner(message, *args, **kwargs):
+        # Store the session key in channel_session
+        if message.http_session is not None and settings.SESSION_COOKIE_NAME not in message.channel_session:
+            message.channel_session[settings.SESSION_COOKIE_NAME] = message.http_session.session_key
+        # Hydrate the http_session from session_key
+        elif message.http_session is None and settings.SESSION_COOKIE_NAME in message.channel_session:
+            session_engine = import_module(settings.SESSION_ENGINE)
+            session = session_engine.SessionStore(session_key=message.channel_session[settings.SESSION_COOKIE_NAME])
+            message.http_session = session
+        # Run the consumer
+        return func(message, *args, **kwargs)
+    return inner
diff --git a/channels/tests/test_sessions.py b/channels/tests/test_sessions.py
index da65f7a..9102aea 100644
--- a/channels/tests/test_sessions.py
+++ b/channels/tests/test_sessions.py
@@ -3,7 +3,8 @@ from __future__ import unicode_literals
 from django.conf import settings
 from django.test import override_settings
 from channels.message import Message
-from channels.sessions import channel_session, http_session, enforce_ordering, session_for_reply_channel
+from channels.sessions import channel_session, channel_and_http_session, http_session, enforce_ordering, \
+    session_for_reply_channel
 from channels.tests import ChannelTestCase
 from channels import DEFAULT_CHANNEL_LAYER, channel_layers
 
@@ -105,6 +106,41 @@ class SessionTests(ChannelTestCase):
         session2 = session_for_reply_channel("test-reply")
         self.assertEqual(session2["species"], "horse")
 
+    def test_channel_and_http_session(self):
+        """
+        Tests that channel_and_http_session decorator stores the http session key and hydrates it when expected
+        """
+        # Make a session to try against
+        session = session_for_reply_channel("test-reply-session")
+        # Construct message to send
+        message = Message({
+            "reply_channel": "test-reply-session",
+            "http_version": "1.1",
+            "method": "GET",
+            "path": "/test2/",
+            "headers": {
+                "host": b"example.com",
+                "cookie": ("%s=%s" % (settings.SESSION_COOKIE_NAME, session.session_key)).encode("ascii"),
+            },
+        }, None, None)
+
+        @channel_and_http_session
+        def inner(message):
+            pass
+
+        inner(message)
+
+        # It should store the session key
+        self.assertEqual(message.channel_session[settings.SESSION_COOKIE_NAME], session.session_key)
+
+        # Construct a new message
+        message2 = Message({"reply_channel": "test-reply-session", "path": "/"}, None, None)
+
+        inner(message2)
+
+        # It should hydrate the http_session
+        self.assertEqual(message2.http_session.session_key, session.session_key)
+
     def test_enforce_ordering_slight(self):
         """
         Tests that slight mode of enforce_ordering works