Channel and http session from http (#564)

* Add a new auth decorator to get the user and rehydrate the http session

* Add http_user_and_session, taking precedence over http_user, applying the channel_and_http_session_user_from_http decorator (a superset of http user functionality)

* Only set session cookies on the first send, since subsequent real requests don't have access to HTTP information

* Add a test for new http_user_and_session WebsocketConsumer attribute

* Fix isort check

channels/auth.py

@@ -2,7 +2,7 @@ import functools
 
 from django.contrib import auth
 
-from .sessions import channel_session, http_session
+from .sessions import channel_and_http_session, channel_session, http_session
 
 
 def transfer_user(from_session, to_session):
@@ -88,3 +88,19 @@ def channel_session_user_from_http(func):
             transfer_user(message.http_session, message.channel_session)
         return func(message, *args, **kwargs)
     return inner
+
+
+def channel_and_http_session_user_from_http(func):
+    """
+    Decorator that automatically transfers the user from HTTP sessions to
+    channel-based sessions, rehydrates the HTTP session, and returns the
+    user as message.user as well.
+    """
+    @http_session_user
+    @channel_and_http_session
+    @functools.wraps(func)
+    def inner(message, *args, **kwargs):
+        if message.http_session is not None:
+            transfer_user(message.http_session, message.channel_session)
+        return func(message, *args, **kwargs)
+    return inner

channels/generic/websockets.py

@@ -1,6 +1,6 @@
 from django.core.serializers.json import DjangoJSONEncoder, json
 
-from ..auth import channel_session_user_from_http
+from ..auth import channel_and_http_session_user_from_http, channel_session_user_from_http
 from ..channel import Group
 from ..exceptions import SendNotAvailableOnDemultiplexer
 from ..sessions import enforce_ordering
@@ -23,6 +23,7 @@ class WebsocketConsumer(BaseConsumer):
     # Turning this on passes the user over from the HTTP session on connect,
     # implies channel_session_user
     http_user = False
+    http_user_and_session = False
 
     # Set to True if you want the class to enforce ordering for you
     strict_ordering = False
@@ -35,13 +36,15 @@ class WebsocketConsumer(BaseConsumer):
         adds the ordering decorator.
         """
         # HTTP user implies channel session user
-        if self.http_user:
+        if self.http_user or self.http_user_and_session:
             self.channel_session_user = True
         # Get super-handler
         self.path = message['path']
         handler = super(WebsocketConsumer, self).get_handler(message, **kwargs)
         # Optionally apply HTTP transfer
-        if self.http_user:
+        if self.http_user_and_session:
+            handler = channel_and_http_session_user_from_http(handler)
+        elif self.http_user:
             handler = channel_session_user_from_http(handler)
         # Ordering decorators
         if self.strict_ordering:

channels/test/http.py

@@ -24,6 +24,7 @@ class HttpClient(Client):
         self._session = None
         self._headers = {}
         self._cookies = {}
+        self._session_cookie = True
 
     def set_cookie(self, key, value):
         """
@@ -42,7 +43,7 @@ class HttpClient(Client):
     def get_cookies(self):
         """Return cookies"""
         cookies = copy.copy(self._cookies)
-        if apps.is_installed('django.contrib.sessions'):
+        if self._session_cookie and apps.is_installed('django.contrib.sessions'):
             cookies[settings.SESSION_COOKIE_NAME] = self.session.session_key
         return cookies
 
@@ -86,6 +87,7 @@ class HttpClient(Client):
             else:
                 content['text'] = text
         self.channel_layer.send(to, content)
+        self._session_cookie = False
 
     def send_and_consume(self, channel, content={}, text=None, path='/', fail_on_none=True, check_accept=True):
         """

tests/test_generic.py

@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 import json
 
 from django.test import override_settings
+from django.contrib.auth import get_user_model
 
 from channels import route_class
 from channels.exceptions import SendNotAvailableOnDemultiplexer
@@ -87,6 +88,29 @@ class GenericTests(ChannelTestCase):
             self.assertEqual(client.consume('websocket.connect').order, 0)
             self.assertEqual(client.consume('websocket.connect').order, 1)
 
+    def test_websockets_http_session_and_channel_session(self):
+
+        class WebsocketConsumer(websockets.WebsocketConsumer):
+            http_user_and_session = True
+
+        user_model = get_user_model()
+        user = user_model.objects.create_user(username='test', email='test@test.com', password='123456')
+
+        client = HttpClient()
+        client.force_login(user)
+        with apply_routes([route_class(WebsocketConsumer, path='/path')]):
+            connect = client.send_and_consume('websocket.connect', {'path': '/path'})
+            receive = client.send_and_consume('websocket.receive', {'path': '/path'}, text={'key': 'value'})
+            disconnect = client.send_and_consume('websocket.disconnect', {'path': '/path'})
+        self.assertEqual(
+            connect.message.http_session.session_key,
+            receive.message.http_session.session_key
+        )
+        self.assertEqual(
+            connect.message.http_session.session_key,
+            disconnect.message.http_session.session_key
+        )
+
     def test_simple_as_route_method(self):
 
         class WebsocketConsumer(websockets.WebsocketConsumer):