From bb6eed0db194651f366571218bc604312784e4fd Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Fri, 19 Nov 2021 23:19:07 +0100
Subject: [PATCH] root: properly catch 404 errors for websocket connections

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/root/asgi.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/authentik/root/asgi.py b/authentik/root/asgi.py
index f0ea34fee..6412c114e 100644
--- a/authentik/root/asgi.py
+++ b/authentik/root/asgi.py
@@ -42,11 +42,28 @@ class LifespanApp:
                     return
 
 
+class RouteNotFoundMiddleware:
+    """Middleware to ignore 404s for websocket requests
+    taken from https://github.com/django/daphne/issues/165#issuecomment-808284950"""
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        try:
+            return await self.app(scope, receive, send)
+        except ValueError as exc:
+            if "No route found for path" in str(exc) and scope["type"] == "websocket":
+                await send({"type": "websocket.close"})
+            else:
+                raise exc
+
+
 application = SentryAsgiMiddleware(
     ProtocolTypeRouter(
         {
             "http": get_asgi_application(),
-            "websocket": URLRouter(websocket.websocket_urlpatterns),
+            "websocket": RouteNotFoundMiddleware(URLRouter(websocket.websocket_urlpatterns)),
             "lifespan": LifespanApp,
         }
     )