django websocket使用JWT登陆验证

使用channels实现websocket，利用restframework-jwt功能进行token验证。
流程如下：

Http登陆-->获取token-->websocket请求连接携带token-->自定义channels的Authentication-->验证token-->允许接入

自定义channels的Authentication官方链接

下面的是rest_framework_jwt（不是channels）的JSONWebTokenAuthentication默认的token验证方法, 此方法是同步的，还执行了user = self.authenticate_credentials(payload)数据库的查询，这些操作都要修改为异步的（channels的异步才能尽其能）

    def authenticate(self, request):
        """
        Returns a two-tuple of `User` and token if a valid signature has been
        supplied using JWT-based authentication.  Otherwise returns `None`.
        """
        jwt_value = self.get_jwt_value(request)
        if jwt_value is None:
            return None

        try:
            payload = jwt_decode_handler(jwt_value)
        except jwt.ExpiredSignature:
            msg = _('Signature has expired.')
            raise exceptions.AuthenticationFailed(msg)
        except jwt.DecodeError:
            msg = _('Error decoding signature.')
            raise exceptions.AuthenticationFailed(msg)
        except jwt.InvalidTokenError:
            raise exceptions.AuthenticationFailed()

        user = self.authenticate_credentials(payload)

        return (user, jwt_value)

上面的代码是从request中获取token的，但websocket提供的是scope。
下面是channels的官方样例的权限验证的中间件。

class QueryAuthMiddleware:
    """
    Custom middleware (insecure) that takes user IDs from the query string.
    """

    def __init__(self, inner):
        # Store the ASGI application we were passed
        self.inner = inner

    def __call__(self, scope):

        # Close old database connections to prevent usage of timed out connections
        #close_old_connections()

        # Look up user from query string (you should also do things like
        # checking if it is a valid user ID, or if scope["user"] is already
        # populated).
        
        user = User.objects.get(id=int(scope["query_string"]))

        # Return the inner application directly and let it run everything else
        return self.inner(dict(scope, user=user))

从上面的JWT和自定义授权类中得知，支持websocket的JWT验证要做的两件事：

• 从socpe中获取token
• 调用JSONWebTokenAuthentication验证Token，并获取用户

从channels的官方文档说明中，执行数据库的查询要使用channels.db.database_sync_to_async ，即是JSONWebTokenAuthentication类中的authenticate方法要变为异步，且其内部调用的用户查询也得是database_sync_to_async。

修改代码，支持异步的websocket支持JWT如下：

新增WebsocketTokenAuthentication类（继承JSONWebTokenAuthentication），修改部分有：

• 从scope中获取url中的token。
• 修改其用户数据库查询代码，支持异步。

class WebsocketTokenAuthentication(JSONWebTokenAuthentication):
    """
    支持channels中间件的token处理和验证。
    """

    def get_jwt_value(self, scope):
        """
        websocket的url：/ws/chat/<str:room_name>?token
        token格式: 'JWT eyJ0eXAiOiJKV1QiLCJhbGciOiJ'
        """
        token = str(scope['query_string'].decode('utf-8'))
        token = urllib.parse.unquote(token)
        token = token.split(" ")[1].encode('utf-8')
        return token
    
    async def authenticate_credentials(self, payload):
        """
        覆盖此方法的user查询，使得其支持异步。
        Returns an active user that matches the payload's user id and email.
        """
        User = get_user_model()
        username = jwt_get_username_from_payload(payload)

        if not username:
            msg = _('Invalid payload.')
            raise exceptions.AuthenticationFailed(msg)

        try:
            # user = User.objects.get_by_natural_key(username)
            user = await self._get_user_by_natural_key(User, username)
        except User.DoesNotExist:
            msg = _('Invalid signature.')
            raise exceptions.AuthenticationFailed(msg)

        if not user.is_active:
            msg = _('User account is disabled.')
            raise exceptions.AuthenticationFailed(msg)

        return user
    
    @database_sync_to_async
    def _get_user_by_natural_key(self, User, username):
        return User.objects.get_by_natural_key(username)

下面就是创建一个QueryAuthMiddleware中间件，从WebsocketTokenAuthentication验证token 并获取用户。

class QueryAuthMiddleware:
    """
    Custom middleware (insecure) that takes user IDs from the query string.
    """

    def __init__(self, inner):
        # Store the ASGI application we were passed
        self.inner = inner

    def __call__(self, scope):

        # Close old database connections to prevent usage of timed out connections
        # close_old_connections()

        # Look up user from query string (you should also do things like
        # checking if it is a valid user ID, or if scope["user"] is already
        # populated).
        
        auth = WebsocketTokenAuthentication()
        user, token = auth.authenticate(scope)

        # Return the inner application directly and let it run everything else
        return self.inner(dict(scope, user=user))

在routing.py中注册中间件

application = ProtocolTypeRouter({
    # 'websocket': AuthMiddlewareStack(
    #     URLRouter(
    #         chat.routing.websocket_urlpatterns
    #     )
    # ),
    'websocket': QueryAuthMiddleware(
        URLRouter(
            chat.routing.websocket_urlpatterns
        )
    ),
})

至此，修改完毕。
使得可以使用AsyncWebsocketConsumer的Consumer和支持JWT验证。
后面只需要在consumer中验证是否存在用户，验证则accept。

class ChatConsumer(AsyncWebsocketConsumer):

    async def connect(self):

        user = await self.scope['user']
        if not user:
            await self.close()
        else:
            self.room_name = self.scope['url_route']['kwargs']['room_name']
            self.room_group_name = 'chat_%s' % self.room_name

            # Join room group
            await self.channel_layer.group_add(
                self.room_group_name,
                self.channel_name
            )
            await self.accept()