websockets behind auth, cleaned up errors and bad state handling in websocket.jsx

dispatcharr/asgi.py

@@ -1,14 +1,17 @@
+import django
 import os
 from django.core.asgi import get_asgi_application
 from channels.routing import ProtocolTypeRouter, URLRouter
-from channels.auth import AuthMiddlewareStack
 import dispatcharr.routing
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "dispatcharr.settings")
+django.setup()
+
+from .jwt_ws_auth import JWTAuthMiddleware
 
 application = ProtocolTypeRouter({
     "http": get_asgi_application(),
-    "websocket": AuthMiddlewareStack(
+    "websocket": JWTAuthMiddleware(
         URLRouter(dispatcharr.routing.websocket_urlpatterns)
     ),
 })

dispatcharr/consumers.py

@@ -6,9 +6,15 @@ logger = logging.getLogger(__name__)
 
 class MyWebSocketConsumer(AsyncWebsocketConsumer):
     async def connect(self):
+        self.room_name = "updates"
+
+        user = self.scope["user"]
+        if not user.is_authenticated:
+            await self.close()
+            return
+
         try:
             await self.accept()
-            self.room_name = "updates"
             await self.channel_layer.group_add(self.room_name, self.channel_name)
             # Send a connection confirmation to the client with consistent format
             await self.send(text_data=json.dumps({

dispatcharr/jwt_ws_auth.py

@@ -0,0 +1,36 @@
+from urllib.parse import parse_qs
+from channels.middleware import BaseMiddleware
+from channels.db import database_sync_to_async
+from rest_framework_simplejwt.tokens import UntypedToken
+from django.contrib.auth.models import AnonymousUser
+from django.contrib.auth import get_user_model
+from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
+from rest_framework_simplejwt.authentication import JWTAuthentication
+
+User = get_user_model()
+
+@database_sync_to_async
+def get_user(validated_token):
+    try:
+        jwt_auth = JWTAuthentication()
+        user = jwt_auth.get_user(validated_token)
+        return user
+    except:
+        return AnonymousUser()
+
+class JWTAuthMiddleware(BaseMiddleware):
+    async def __call__(self, scope, receive, send):
+        try:
+            # Extract the token from the query string
+            query_string = parse_qs(scope["query_string"].decode())
+            token = query_string.get("token", [None])[0]
+
+            if token is not None:
+                validated_token = JWTAuthentication().get_validated_token(token)
+                scope["user"] = await get_user(validated_token)
+            else:
+                scope["user"] = AnonymousUser()
+        except (InvalidToken, TokenError):
+            scope["user"] = AnonymousUser()
+
+        return await super().__call__(scope, receive, send)