基于django-channel构建websocket的web应用、支持token校验

0 前期准备

• 本文参考官方文档：https://channels.readthedocs.io/做了一定的修改，并加上个人的经验及理解
• 搭建一个django项目，假定项目名称为mysite，应用名称为chat
• 安装channels：pip install channels

1 配置asgi文件

asgi文件用于指向接收不同协议的请求要执行的代码。

# mysite/asgi.py
import os

from channels.auth import AuthMiddlewareStack
from channels.routing import ProtocolTypeRouter, URLRouter
from django.core.asgi import get_asgi_application
import chat.routing

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "mysite.settings")

application = ProtocolTypeRouter({
  "http": get_asgi_application(),
  "websocket": AuthMiddlewareStack(
        URLRouter(
            chat.routing.websocket_urlpatterns
        )
    ),
})

AuthMiddlewareStack是用于参数校验的中间件，后面我们会做自定义中间件。
URLRouter定义处理websocket请求的路由，因此下面我们要新建routing文件。

2 在应用中新建routing路由文件

# chat/routing.py
from django.urls import re_path
from base import consumers

websocket_urlpatterns = [
    re_path(r'^ws/chat/(?P<report_id>[^/]+)/$', consumers.ChatConsumer.as_asgi()),
]

指定路由/ws/chat/(?P<report_id>[^/]+)/$通过自定义视图consumers.ChatConsumer（消费者）进行处理，report_id会作为参数传入消费者类。
因此我们要新建消费者ChatConsumer视图进行逻辑处理。

3 编写消费者

在chat下新建consumers.py，可以使用同步或者异步代码编写，这里使用同步，相关解释见注释。

# chat/consumers.py
import json
from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer

class ChatConsumer(WebsocketConsumer):
    # websocket连接时会执行
    def connect(self):
        # 这里会取到url中传入的report_id
        self.report_name = self.scope['url_route']['kwargs']['report_id']
        self.report_group_name = 'chat_%s' % self.report_name

        # group_add 加入一个小组
        # channel_layer下面的方法属于通道层方法是异步的，因此使用async_to_sync进行包装
        async_to_sync(self.channel_layer.group_add)(
            self.report_group_name,
            self.channel_name
        )
        # accept表示建立连接，否则可以使用close
        self.accept()
	
	# 断开websocket连接时会执行
    def disconnect(self, close_code):
        # 断开与小组的连接 group_discard
        async_to_sync(self.channel_layer.group_discard)(
            self.report_group_name,
            self.channel_name
        )

    # 接收到消息会执行
    def receive(self, text_data):
        # text_data 接收到的消息
        text_data_json = json.loads(text_data)
        message = text_data_json['message']

        # group_send 向小组发送消息
        # self.report_group_name 指小组的名称
        # 'type': 'chat_message'  会调用chat_message发送消息
        async_to_sync(self.channel_layer.group_send)(
            self.report_group_name,
            {
                'type': 'chat_message',
                'message': message
            }
        )

    # group_send被调用时，会执行指定选择的方法
    def chat_message(self, event):
        message = event['message']

        # Send message to WebSocket
        self.send(text_data=json.dumps({
            'message': message
        }))

4 配置通道层后端

ASGI_APPLICATION = 'mysite.asgi.application'
CHANNEL_LAYERS = {
    "default": {
        "BACKEND": "channels.layers.InMemoryChannelLayer"
    }
}

ASGI_APPLICATION配置应用接收请求的处理通道。
CHANNEL_LAYERS 配置消息储存的后端，InMemoryChannelLayer是把消息存在内存。
如果要存在redis，使用下面的方法进行配置：

CHANNEL_LAYERS = {
    'default': {
        'BACKEND': 'channels_redis.core.RedisChannelLayer',
        'CONFIG': {
            "hosts": [('127.0.0.1', 6379)],
        },
    },
}

5 通过中间件进行校验

新建chat/middleware.py文件，这里编写一个通过authorization进行校验的中间件，如果校验不通过，则会为websocket请求返回403状态码。

from django.contrib.auth import get_user_model
from django.contrib.auth.models import AnonymousUser
from rest_framework.authtoken.models import Token

User = get_user_model()

# 通过token获取用户对象
def get_user(token):
    try:
        token_obj = Token.objects.get(key=token)
        return User.objects.get(id=token_obj.user_id)
    except User.DoesNotExist:
        return AnonymousUser()
    except Token.DoesNotExist:
        return AnonymousUser()


# 校验中间件
# 请求在进入websocket的connect方法之前，就调用__call__方法
class QueryAuthMiddleware:
    def __init__(self, app):
        self.app = app

    def __call__(self, scope, receive, send):
        # 获取请求头
        headers = scope["headers"]
        headers_dict = {k.decode("utf-8").lower(): v.decode("utf-8") for k, v in headers}
        # 根据token获取user
        # 如果对应token没有user就返回匿名用户状态
        # 判断完后把user对象放到全局变量scope中，在connect中可以再次进行判断和调用
        if "authorization" in headers_dict.keys():
            token = headers_dict["authorization"].split(" ")[1]
            scope['user'] = get_user(token)
        else:
            scope['user'] = AnonymousUser()

        return self.app(scope, receive, send)

消费者connect方法会接收到被中间件处理过的scope对象，我们根据scope中的user的内容进行再次判断，即重写connect方法。

def connect(self):
    # 从user对象中拿到is_active和is_anonymous的值
    user_is_active = self.scope["user"].is_active
    user_is_anonymous = self.scope["user"].is_anonymous
    self.report_name = self.scope['url_route']['kwargs']['report_id']
    self.report_group_name = 'chat_%s' % self.report_name
    # 判断用户账号的活动状态，如果不是活动状态则调用close关闭连接
    if not user_is_active or user_is_anonymous:
        self.close()
    else:
        async_to_sync(self.channel_layer.group_add)(
            self.report_group_name,
            self.channel_name
        )

        await self.accept()