From 239092b872e8e8183e419d6f6dfb4ae50361c58e Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Mon, 10 Oct 2022 13:18:08 +0300
Subject: [PATCH] core: fix messages not being shown when no client is
 connected

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/core/templates/base/header_js.html  | 20 +++++++++++++++++++
 authentik/core/templates/if/admin.html        |  7 +------
 authentik/core/templates/if/flow.html         |  5 +----
 authentik/core/templates/if/user.html         |  7 +------
 authentik/core/templates/login/base_full.html |  1 +
 authentik/flows/api/stages.py                 |  2 +-
 authentik/root/messages/storage.py            | 15 +++++++++-----
 web/src/elements/messages/MessageContainer.ts |  5 +++--
 8 files changed, 38 insertions(+), 24 deletions(-)
 create mode 100644 authentik/core/templates/base/header_js.html

diff --git a/authentik/core/templates/base/header_js.html b/authentik/core/templates/base/header_js.html
new file mode 100644
index 000000000..ce92fc8b8
--- /dev/null
+++ b/authentik/core/templates/base/header_js.html
@@ -0,0 +1,20 @@
+<script>
+    window.authentik = {};
+    window.authentik.locale = "{{ tenant.default_locale }}";
+    window.authentik.config = JSON.parse('{{ config_json|escapejs }}');
+    window.authentik.tenant = JSON.parse('{{ tenant_json|escapejs }}');
+    window.addEventListener("DOMContentLoaded", () => {
+        {% for message in messages %}
+        window.dispatchEvent(
+            new CustomEvent("ak-message", {
+                bubbles: true,
+                composed: true,
+                detail: {
+                    level: "{{ message.tags|escapejs }}",
+                    message: "{{ message.message|escapejs }}",
+                },
+            }),
+        );
+        {% endfor %}
+    });
+</script>
diff --git a/authentik/core/templates/if/admin.html b/authentik/core/templates/if/admin.html
index b0fe05922..880ea3bb6 100644
--- a/authentik/core/templates/if/admin.html
+++ b/authentik/core/templates/if/admin.html
@@ -9,12 +9,7 @@
 <meta name="theme-color" content="#ffffff" media="(prefers-color-scheme: light)">
 <link rel="icon" href="{{ tenant.branding_favicon }}">
 <link rel="shortcut icon" href="{{ tenant.branding_favicon }}">
-<script>
-window.authentik = {};
-window.authentik.locale = "{{ tenant.default_locale }}";
-window.authentik.config = JSON.parse('{{ config_json|escapejs }}');
-window.authentik.tenant = JSON.parse('{{ tenant_json|escapejs }}');
-</script>
+{% include "base/header_js.html" %}
 {% endblock %}
 
 {% block body %}
diff --git a/authentik/core/templates/if/flow.html b/authentik/core/templates/if/flow.html
index 657c21579..05a03692a 100644
--- a/authentik/core/templates/if/flow.html
+++ b/authentik/core/templates/if/flow.html
@@ -11,11 +11,8 @@
 {% if flow.compatibility_mode and not inspector %}
 <script>ShadyDOM = { force: !navigator.webdriver };</script>
 {% endif %}
+{% include "base/header_js.html" %}
 <script>
-window.authentik = {};
-window.authentik.locale = "{{ tenant.default_locale }}";
-window.authentik.config = JSON.parse('{{ config_json|escapejs }}');
-window.authentik.tenant = JSON.parse('{{ tenant_json|escapejs }}');
 window.authentik.flow = {
     "layout": "{{ flow.layout }}",
 };
diff --git a/authentik/core/templates/if/user.html b/authentik/core/templates/if/user.html
index ba4a7a7e5..7c0907820 100644
--- a/authentik/core/templates/if/user.html
+++ b/authentik/core/templates/if/user.html
@@ -9,12 +9,7 @@
 <meta name="theme-color" content="#151515" media="(prefers-color-scheme: dark)">
 <link rel="icon" href="{{ tenant.branding_favicon }}">
 <link rel="shortcut icon" href="{{ tenant.branding_favicon }}">
-<script>
-window.authentik = {};
-window.authentik.locale = "{{ tenant.default_locale }}";
-window.authentik.config = JSON.parse('{{ config_json|escapejs }}');
-window.authentik.tenant = JSON.parse('{{ tenant_json|escapejs }}');
-</script>
+{% include "base/header_js.html" %}
 {% endblock %}
 
 {% block body %}
diff --git a/authentik/core/templates/login/base_full.html b/authentik/core/templates/login/base_full.html
index e5562df7d..a1c58e7dd 100644
--- a/authentik/core/templates/login/base_full.html
+++ b/authentik/core/templates/login/base_full.html
@@ -6,6 +6,7 @@
 {% block head_before %}
 <link rel="prefetch" href="/static/dist/assets/images/flow_background.jpg" />
 <link rel="stylesheet" type="text/css" href="{% static 'dist/patternfly.min.css' %}">
+{% include "base/header_js.html" %}
 {% endblock %}
 
 {% block head %}
diff --git a/authentik/flows/api/stages.py b/authentik/flows/api/stages.py
index d7a394859..674e2cf36 100644
--- a/authentik/flows/api/stages.py
+++ b/authentik/flows/api/stages.py
@@ -55,7 +55,7 @@ class StageViewSet(
 ):
     """Stage Viewset"""
 
-    queryset = Stage.objects.all().select_related("flow_set")
+    queryset = Stage.objects.none()
     serializer_class = StageSerializer
     search_fields = ["name"]
     filterset_fields = ["name"]
diff --git a/authentik/root/messages/storage.py b/authentik/root/messages/storage.py
index cdbc92627..fdc15ef5c 100644
--- a/authentik/root/messages/storage.py
+++ b/authentik/root/messages/storage.py
@@ -1,12 +1,15 @@
 """Channels Messages storage"""
 from asgiref.sync import async_to_sync
 from channels.layers import get_channel_layer
-from django.contrib.messages.storage.base import BaseStorage, Message
+from django.contrib.messages.storage.base import Message
+from django.contrib.messages.storage.session import SessionStorage
 from django.core.cache import cache
 from django.http.request import HttpRequest
 
+SESSION_KEY = "_messages"
 
-class ChannelsStorage(BaseStorage):
+
+class ChannelsStorage(SessionStorage):
     """Send contrib.messages over websocket"""
 
     def __init__(self, request: HttpRequest) -> None:
@@ -14,12 +17,13 @@ class ChannelsStorage(BaseStorage):
         super().__init__(request)
         self.channel = get_channel_layer()
 
-    def _get(self):
-        return [], True
-
     def _store(self, messages: list[Message], response, *args, **kwargs):
         prefix = f"user_{self.request.session.session_key}_messages_"
         keys = cache.keys(f"{prefix}*")
+        # if no active connections are open, fallback to storing messages in the
+        # session, so they can always be retrieved
+        if len(keys) < 1:
+            return super()._store(messages, response, *args, **kwargs)
         for key in keys:
             uid = key.replace(prefix, "")
             for message in messages:
@@ -33,3 +37,4 @@ class ChannelsStorage(BaseStorage):
                         "message": message.message,
                     },
                 )
+        return []
diff --git a/web/src/elements/messages/MessageContainer.ts b/web/src/elements/messages/MessageContainer.ts
index 459c71edd..ef666e040 100644
--- a/web/src/elements/messages/MessageContainer.ts
+++ b/web/src/elements/messages/MessageContainer.ts
@@ -44,11 +44,11 @@ export class MessageContainer extends AKElement {
 
     constructor() {
         super();
-        this.addEventListener(EVENT_WS_MESSAGE, ((e: CustomEvent<WSMessage>) => {
+        window.addEventListener(EVENT_WS_MESSAGE, ((e: CustomEvent<WSMessage>) => {
             if (e.detail.message_type !== WS_MSG_TYPE_MESSAGE) return;
             this.addMessage(e.detail as unknown as APIMessage);
         }) as EventListener);
-        this.addEventListener(EVENT_MESSAGE, ((e: CustomEvent<APIMessage>) => {
+        window.addEventListener(EVENT_MESSAGE, ((e: CustomEvent<APIMessage>) => {
             this.addMessage(e.detail);
         }) as EventListener);
     }
@@ -61,6 +61,7 @@ export class MessageContainer extends AKElement {
             }
         }
         this.messages.push(message);
+        this.requestUpdate();
     }
 
     render(): TemplateResult {