custom asgi auth middleware

Signed-off-by: Marc 'risson' Schmitt <[email protected]>

authentik/core/urls.py

@@ -1,7 +1,5 @@
 """authentik URL Configuration"""
 
-from channels.auth import AuthMiddleware
-from channels.sessions import CookieMiddleware
 from django.conf import settings
 from django.contrib.auth.decorators import login_required
 from django.urls import path
@@ -25,7 +23,7 @@
     RootRedirectView,
 )
 from authentik.flows.views.interface import FlowInterfaceView
-from authentik.root.asgi_middleware import SessionMiddleware
+from authentik.root.asgi_middleware import AuthMiddlewareStack
 from authentik.root.messages.consumer import MessageConsumer
 from authentik.root.middleware import ChannelsLoggingMiddleware
 
@@ -94,9 +92,7 @@
 websocket_urlpatterns = [
     path(
         "ws/client/",
-        ChannelsLoggingMiddleware(
-            CookieMiddleware(SessionMiddleware(AuthMiddleware(MessageConsumer.as_asgi())))
-        ),
+        ChannelsLoggingMiddleware(AuthMiddlewareStack(MessageConsumer.as_asgi())),
     ),
 ]
 

authentik/providers/rac/urls.py

@@ -1,7 +1,5 @@
 """rac urls"""
 
-from channels.auth import AuthMiddleware
-from channels.sessions import CookieMiddleware
 from django.urls import path
 
 from authentik.outposts.channels import TokenOutpostMiddleware
@@ -12,7 +10,7 @@
 from authentik.providers.rac.consumer_client import RACClientConsumer
 from authentik.providers.rac.consumer_outpost import RACOutpostConsumer
 from authentik.providers.rac.views import RACInterface, RACStartView
-from authentik.root.asgi_middleware import SessionMiddleware
+from authentik.root.asgi_middleware import AuthMiddlewareStack
 from authentik.root.middleware import ChannelsLoggingMiddleware
 
 urlpatterns = [
@@ -31,9 +29,7 @@
 websocket_urlpatterns = [
     path(
         "ws/rac/<str:token>/",
-        ChannelsLoggingMiddleware(
-            CookieMiddleware(SessionMiddleware(AuthMiddleware(RACClientConsumer.as_asgi())))
-        ),
+        ChannelsLoggingMiddleware(AuthMiddlewareStack(RACClientConsumer.as_asgi())),
     ),
     path(
         "ws/outpost_rac/<str:channel>/",

authentik/root/asgi_middleware.py

@@ -1,8 +1,12 @@
 """ASGI middleware"""
 
+from channels.auth import UserLazyObject
 from channels.db import database_sync_to_async
+from channels.middleware import BaseMiddleware
+from channels.sessions import CookieMiddleware
 from channels.sessions import InstanceSessionWrapper as UpstreamInstanceSessionWrapper
 from channels.sessions import SessionMiddleware as UpstreamSessionMiddleware
+from django.contrib.auth.models import AnonymousUser
 
 from authentik.root.middleware import SessionMiddleware as HTTPSessionMiddleware
 
@@ -33,3 +37,48 @@ async def __call__(self, scope, receive, send):
         await wrapper.resolve_session()
 
         return await self.inner(wrapper.scope, receive, wrapper.send)
+
+
+@database_sync_to_async
+def get_user(scope):
+    """
+    Return the user model instance associated with the given scope.
+    If no user is retrieved, return an instance of `AnonymousUser`.
+    """
+    if "session" not in scope:
+        raise ValueError(
+            "Cannot find session in scope. You should wrap your consumer in SessionMiddleware."
+        )
+    user = None
+    if (authenticated_session := scope["session"].get("authenticated_session", None)) is not None:
+        user = authenticated_session.user
+    return user or AnonymousUser()
+
+
+class AuthMiddleware(BaseMiddleware):
+    def populate_scope(self, scope):
+        # Make sure we have a session
+        if "session" not in scope:
+            raise ValueError(
+                "AuthMiddleware cannot find session in scope. SessionMiddleware must be above it."
+            )
+        # Add it to the scope if it's not there already
+        if "user" not in scope:
+            scope["user"] = UserLazyObject()
+
+    async def resolve_scope(self, scope):
+        scope["user"]._wrapped = await get_user(scope)
+
+    async def __call__(self, scope, receive, send):
+        scope = dict(scope)
+        # Scope injection/mutation per this middleware's needs.
+        self.populate_scope(scope)
+        # Grab the finalized/resolved scope
+        await self.resolve_scope(scope)
+
+        return await super().__call__(scope, receive, send)
+
+
+# Handy shortcut for applying all three layers at once
+def AuthMiddlewareStack(inner):
+    return CookieMiddleware(SessionMiddleware(AuthMiddleware(inner)))