From b6436826f614109725ecc30a0e17a402ccfbefd3 Mon Sep 17 00:00:00 2001
From: Chris Meyers <chris.meyers.fsu@gmail.com>
Date: Wed, 14 Feb 2024 14:14:20 -0500
Subject: [PATCH] Support websocket per-endpoint auth

* Channels doesn't really give you an interface to support per-endpoint
  auth ... so this adds one.
* The web browser and node <--> node communication have different auth
  needs.
---
 awx/main/routing.py | 38 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/awx/main/routing.py b/awx/main/routing.py
index 239db32fad..05edd2a802 100644
--- a/awx/main/routing.py
+++ b/awx/main/routing.py
@@ -27,14 +27,50 @@ class AWXProtocolTypeRouter(ProtocolTypeRouter):
         super().__init__(*args, **kwargs)
 
 
+class MultipleURLRouterAdapter:
+    """
+    Django channels doesn't nicely support Auth_1(urls_1), Auth_2(urls_2), ..., Auth_n(urls_n)
+    This class allows assocating a websocket url with an auth
+    Ordering matters. The first matching url will be used.
+    """
+
+    def __init__(self, *auths):
+        self._auths = [a for a in auths]
+
+    async def __call__(self, scope, receive, send):
+        """
+        Loop through the list of passed in URLRouter's (they may or may not be wrapped by auth).
+        We know we have exhausted the list of URLRouter patterns when we get a
+        ValueError('No route found for path %s'). When that happens, move onto the next
+        URLRouter.
+        If the final URLRouter raises an error, re-raise it in the end.
+
+        We know that we found a match when no error is raised, end the loop.
+        """
+        last_index = len(self._auths) - 1
+        for i, auth in enumerate(self._auths):
+            try:
+                return await auth.__call__(scope, receive, send)
+            except ValueError as e:
+                if str(e).startswith('No route found for path'):
+                    # Only surface the error if on the last URLRouter
+                    if i == last_index:
+                        raise
+
+
 websocket_urlpatterns = [
     re_path(r'api/websocket/$', consumers.EventConsumer.as_asgi()),
     re_path(r'websocket/$', consumers.EventConsumer.as_asgi()),
+]
+websocket_relay_urlpatterns = [
     re_path(r'websocket/relay/$', consumers.RelayConsumer.as_asgi()),
 ]
 
 application = AWXProtocolTypeRouter(
     {
-        'websocket': DrfAuthMiddlewareStack(URLRouter(websocket_urlpatterns)),
+        'websocket': MultipleURLRouterAdapter(
+            URLRouter(websocket_relay_urlpatterns),
+            DrfAuthMiddlewareStack(URLRouter(websocket_urlpatterns)),
+        )
     }
 )