From 54cacd784c3379a19cf7520fb77fdf67ee8430b1 Mon Sep 17 00:00:00 2001
From: Jens L <jens@goauthentik.io>
Date: Mon, 20 Mar 2023 23:39:25 +0100
Subject: [PATCH] *: load websocket paths similarly to URLs (#5018)

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
---
 authentik/core/apps.py      |  1 +
 authentik/core/urls.py      | 10 ++++++++++
 authentik/outposts/apps.py  |  1 +
 authentik/outposts/urls.py  |  8 ++++++++
 authentik/root/websocket.py | 30 ++++++++++++++++++------------
 5 files changed, 38 insertions(+), 12 deletions(-)
 create mode 100644 authentik/outposts/urls.py

diff --git a/authentik/core/apps.py b/authentik/core/apps.py
index 719b2abb1..9042069e6 100644
--- a/authentik/core/apps.py
+++ b/authentik/core/apps.py
@@ -11,6 +11,7 @@ class AuthentikCoreConfig(ManagedAppConfig):
     label = "authentik_core"
     verbose_name = "authentik Core"
     mountpoint = ""
+    ws_mountpoint = "authentik.outposts.urls"
     default = True
 
     def reconcile_load_core_signals(self):
diff --git a/authentik/core/urls.py b/authentik/core/urls.py
index 2252674c6..8dc412465 100644
--- a/authentik/core/urls.py
+++ b/authentik/core/urls.py
@@ -1,4 +1,6 @@
 """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
@@ -9,6 +11,8 @@ from authentik.core.views import apps, impersonate
 from authentik.core.views.debug import AccessDeniedView
 from authentik.core.views.interface import FlowInterfaceView, InterfaceView
 from authentik.core.views.session import EndSessionView
+from authentik.root.asgi_middleware import SessionMiddleware
+from authentik.root.messages.consumer import MessageConsumer
 
 urlpatterns = [
     path(
@@ -64,6 +68,12 @@ urlpatterns = [
     ),
 ]
 
+websocket_urlpatterns = [
+    path(
+        "ws/client/", CookieMiddleware(SessionMiddleware(AuthMiddleware(MessageConsumer.as_asgi())))
+    ),
+]
+
 if settings.DEBUG:
     urlpatterns += [
         path("debug/policy/deny/", AccessDeniedView.as_view(), name="debug-policy-deny"),
diff --git a/authentik/outposts/apps.py b/authentik/outposts/apps.py
index 6898a170a..ea29ebdca 100644
--- a/authentik/outposts/apps.py
+++ b/authentik/outposts/apps.py
@@ -24,6 +24,7 @@ class AuthentikOutpostConfig(ManagedAppConfig):
     label = "authentik_outposts"
     verbose_name = "authentik Outpost"
     default = True
+    ws_mountpoint = "authentik.outposts.urls"
 
     def reconcile_load_outposts_signals(self):
         """Load outposts signals"""
diff --git a/authentik/outposts/urls.py b/authentik/outposts/urls.py
new file mode 100644
index 000000000..696fd7ff6
--- /dev/null
+++ b/authentik/outposts/urls.py
@@ -0,0 +1,8 @@
+"""Outpost Websocket URLS"""
+from django.urls import path
+
+from authentik.outposts.channels import OutpostConsumer
+
+websocket_urlpatterns = [
+    path("ws/outpost/<uuid:pk>/", OutpostConsumer.as_asgi()),
+]
diff --git a/authentik/root/websocket.py b/authentik/root/websocket.py
index d7591ff96..8a044056f 100644
--- a/authentik/root/websocket.py
+++ b/authentik/root/websocket.py
@@ -1,15 +1,21 @@
 """root Websocket URLS"""
-from channels.auth import AuthMiddleware
-from channels.sessions import CookieMiddleware
-from django.urls import path
+from importlib import import_module
 
-from authentik.outposts.channels import OutpostConsumer
-from authentik.root.asgi_middleware import SessionMiddleware
-from authentik.root.messages.consumer import MessageConsumer
+from structlog.stdlib import get_logger
 
-websocket_urlpatterns = [
-    path("ws/outpost/<uuid:pk>/", OutpostConsumer.as_asgi()),
-    path(
-        "ws/client/", CookieMiddleware(SessionMiddleware(AuthMiddleware(MessageConsumer.as_asgi())))
-    ),
-]
+from authentik.lib.utils.reflection import get_apps
+
+LOGGER = get_logger()
+
+websocket_urlpatterns = []
+for _authentik_app in get_apps():
+    mountpoint = getattr(_authentik_app, "ws_mountpoint", None)
+    if not mountpoint:
+        continue
+    ws_paths = import_module(mountpoint)
+    websocket_urlpatterns.extend(getattr(ws_paths, "websocket_urlpatterns"))
+    LOGGER.debug(
+        "Mounted URLs",
+        app_name=_authentik_app.name,
+        app_mountpoint=mountpoint,
+    )