关于django-channels的token认证实现

django-channels只支持基于session的认证，在实际使用中可能需要基于token或者jwt的认证。

参照 https://www.machinalis.com/blog/jwt-django-channels/ 实现基于django token认证。

# ws_authentication.py
from functools import wraps
from json import loads, dumps

from django.utils.translation import ugettext_lazy as _
from django.core.exceptions import ObjectDoesNotExist

from rest_framework import exceptions
from channels.handler import AsgiRequest

from multi_token.models import Token


def _close_reply_channel(message):
    message.reply_channel.send({"close": True})


def authenticate(token):
    """
    Tries to authenticate user based on the supplied token. It also checks
    the token structure and validity.
    """
    try:
        auth_token = Token.objects.get(key=token)
    except Token.DoesNotExist:
        msg = _('Invalid token')
        raise exceptions.AuthenticationFailed(msg)
    """Other authenticate operation"""

    return auth_token.user, auth_token


def ws_auth_request_token(func):
    """
    Checks the presence of a "token" request parameter and tries to
    authenticate the user based on its content.
    The request url must include token.
    eg: /v1/channel/1/?token=abcdefghijklmn
    """
    @wraps(func)
    def inner(message, *args, **kwargs):
        try:
            if "method" not in message.content:
                message.content['method'] = "FAKE"
            request = AsgiRequest(message)
        except Exception as e:
            raise ValueError("Cannot parse HTTP message - are you sure this is a HTTP consumer? %s" % e)

        token = request.GET.get("token", None)

        print request.path,request.GET

        if token is None:
            _close_reply_channel(message)
            raise ValueError("Missing token request parameter. Closing channel.")

        user, token = authenticate(token)

        message.token = token
        message.user = user

        return func(message, *args, **kwargs)
    return inner


def ws_auth_message_token(func):
    """
    Checks the presence of a "token" field on the message's text field and
    tries to authenticate the user based on its content.
    The text must include "token" field.
    eg:{'text':{'token': 'abcdefg','otherdata':''}}.
    """
    @wraps(func)
    def inner(message, *args, **kwargs):
        message_text = message.get('text', None)
        if message_text is None:
            _close_reply_channel(message)
            raise ValueError("Missing text field. Closing channel.")

        try:
            message_text_json = loads(message_text)
        except ValueError:
            _close_reply_channel(message)
            raise

        token = message_text_json.pop('token', None)
        if token is None:
            _close_reply_channel(message)
            raise ValueError("Missing token field. Closing channel.")

        user, token = authenticate(token)

        message.token = token
        message.user = user
        message.text = dumps(message_text_json)

        return func(message, *args, **kwargs)
    return inner

使用：

from channels import Group
from .ws_authentication import ws_auth_request_token, ws_auth_message_token


@ws_auth_request_token
def ws_connect(message, room_id):
    Group("room-%s" % room_id).add(message.reply_channel)
    message.reply_channel.send({"accept": True})


@ws_auth_message_token
def ws_receive(message, room_id):
    message.reply_channel.send({'text': message['text']})


def ws_disconnect(message,room_id):
    Group("room-%s" % room_id).discard(message.reply_channel)

ws测试：

from channels.tests import ChannelTestCase, HttpClient

from rest_framework import status
from rest_framework.test import APITestCase


class WsNotifyTests(ChannelTestCase, APITestCase):

    @override_settings(DEBUG=True)
    def get_token(self):
        """Get token"""
        return token

    def test_notify(self):

        token = self.get_token()

        token_str = "token=%s" % token
        client = HttpClient()
        client.send_and_consume(u'websocket.connect',
                                {'path': '/v1/channel/,
                                 'method': 'GET',
                                 "query_string": token_str})

        self.assertIsNone(client.receive())
        Group('restaurant-%s' % restaurant_id).send({'text': 'ok'}, immediately=True)
        self.assertEqual(client.receive(json=False), 'ok')