Django Channels 中间件使用 ORM

Posted 2023-03-30

【中文标题】Django Channels 中间件使用 ORM

【英文标题】：Django Channels middleware use ORM

【发布时间】：2020-09-08 22:38:18

【问题描述】：

我在检查中间件内部的用户令牌时遇到了麻烦。我从 cookie 中获取令牌，然后我需要查询数据库以检查此令牌是否存在并属于发出请求的用户。

routing.py

from channels.routing import ProtocolTypeRouter, URLRouter
import game.routing
from authentication.utils import TokenAuthMiddlewareStack

application = ProtocolTypeRouter(
    # (http->django views is added by default)
    'websocket': TokenAuthMiddlewareStack(
        URLRouter(
            game.routing.websocket_urlpatterns
        )
    ),
)

middleware.py

from rest_framework.authentication import TokenAuthentication
from rest_framework.exceptions import AuthenticationFailed
from rest_auth.models import TokenModel

from channels.auth import AuthMiddlewareStack
from django.contrib.auth.models import AnonymousUser
from django.db import close_old_connections
...
class TokenAuthMiddleware:
    """
    Token authorization middleware for Django Channels 2
    """

    def __init__(self, inner):
        self.inner = inner


    def __call__(self, scope):
        close_old_connections()

        headers = dict(scope['headers'])
        if b'Authorization' in headers[b'cookie']:
            try:
                cookie_str = headers[b'cookie'].decode('utf-8')

                try:  # no cookie Authorization=Token in the request
                    token_str = [x for x in cookie_str.split(';') if re.search(' Authorization=Token', x)][0].strip()
                except IndexError:
                    scope['user'] = AnonymousUser()
                    return self.inner(scope)

                token_name, token_key = token_str.replace('Authorization=', '').split()
                if token_name == 'Token':
                    token = TokenModel.objects.get(key=token_key)
                    scope['user'] = token.user

            except TokenModel.DoesNotExist:
                scope['user'] = AnonymousUser()
        return self.inner(scope)


TokenAuthMiddlewareStack = lambda inner: TokenAuthMiddleware(AuthMiddlewareStack(inner))

这给了我

django.core.exceptions.SynchronousOnlyOperation: You cannot call this from an async context - use a thread or sync_to_async.

我也尝试了以下方法

async def __call__(self, scope):
    ...
    if token_name == 'Token':
        token = await self.get_token(token_key)
        scope['user'] = token.user
    ...

# approach 1
@sync_to_async
def get_token(self, token_key):
    return TokenModel.objects.get(key=token_key)

# approach 2
@database_sync_to_async
def get_token(self, token_key):
    return TokenModel.objects.get(key=token_key)

这些方法会出现以下错误

[Failure instance: Traceback: <class 'TypeError'>: 'coroutine' object is not callable
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/autobahn/websocket/protocol.py:2847:processHandshake
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/txaio/tx.py:366:as_future
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/twisted/internet/defer.py:151:maybeDeferred
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/daphne/ws_protocol.py:72:onConnect
--- <exception caught here> ---
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/twisted/internet/defer.py:151:maybeDeferred
/Users/nikitatonkoshkur/Documents/work/svoya_igra/venv/lib/python3.8/site-packages/daphne/server.py:206:create_application
]```

【问题讨论】：

将get_token函数保留在类外，使用await get_token(token_key)代替await self.get_token(token_key)

【参考方案1】：

我不确定它是否有效，但你可以尝试

先在类外写get_token函数。

   # approach 1
   @sync_to_async
   def get_token(self, token_key):
     return TokenModel.objects.get(key=token_key)

然后在你的异步函数中写 get_token() 而不是 self.get_token()

  async def __call__(self, scope):
     ...
     if token_name == 'Token':
        token = await get_token(token_key)
        scope['user'] = token.user
     ...

【讨论】：

不，完全一样的结果

【参考方案2】：

查看了django-channels 的源代码，尤其是SessionMiddleware 如何工作和使用 django ORM，我最终以类似的方式编写了我的 TokenMiddleware。

class TokenAuthMiddlewareInstance:
    """
    Token authorization middleware for Django Channels 2
    """

    def __init__(self, scope, middleware):
        self.middleware = middleware
        self.scope = dict(scope)
        self.inner = self.middleware.inner

    async def __call__(self, receive, send):
        close_old_connections()

        headers = dict(self.scope.get('headers', ))
        if b'cookie' in headers:
            cookie_dict = parse_cookie(headers[b'cookie'].decode("ascii"))
            token = cookie_dict.get('Authorization', '')
            token_name, token_key = token.replace('Authorization=', '').split()
            if token_name == 'Token':
                self.scope['user'] = await self.get_user_from_token(token_key)
                inner = self.inner(self.scope)
                return await inner(receive, send)

        inner = self.inner(self.scope)
        self.scope['user'] = AnonymousUser()
        return await inner(receive, send)

    @staticmethod
    @database_sync_to_async
    def get_user_from_token(token_key):
        try:
            return TokenModel.objects.select_related('user').get(key=token_key).user
        except TokenModel.DoesNotExist:
            return AnonymousUser()


class TokenAuthMiddleware:
    def __init__(self, inner):
        self.inner = inner

    def __call__(self, scope):
        return TokenAuthMiddlewareInstance(scope, self)


TokenAuthMiddlewareStack = lambda inner: TokenAuthMiddleware(AuthMiddlewareStack(inner))