From 0585000cf3bd61c747ea961c1c1257ff90d2b820 Mon Sep 17 00:00:00 2001
From: Mitchell Syer <Syer10@users.noreply.github.com>
Date: Fri, 17 Oct 2025 12:07:02 -0400
Subject: [PATCH] Fix header/cookie based websocket auth (#1722)

* Fix header/cookie based websocket auth

* Lint
---
 .../ApolloSubscriptionProtocolHandler.kt             | 12 +++++++++---
 .../kotlin/suwayomi/tachidesk/server/JavalinSetup.kt |  7 +++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/server/src/main/kotlin/suwayomi/tachidesk/graphql/server/subscriptions/ApolloSubscriptionProtocolHandler.kt b/server/src/main/kotlin/suwayomi/tachidesk/graphql/server/subscriptions/ApolloSubscriptionProtocolHandler.kt
index 5cb20dd9..38a0ef76 100644
--- a/server/src/main/kotlin/suwayomi/tachidesk/graphql/server/subscriptions/ApolloSubscriptionProtocolHandler.kt
+++ b/server/src/main/kotlin/suwayomi/tachidesk/graphql/server/subscriptions/ApolloSubscriptionProtocolHandler.kt
@@ -37,6 +37,8 @@ import suwayomi.tachidesk.graphql.server.subscriptions.SubscriptionOperationMess
 import suwayomi.tachidesk.graphql.server.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_ERROR
 import suwayomi.tachidesk.graphql.server.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_NEXT
 import suwayomi.tachidesk.graphql.server.toGraphQLContext
+import suwayomi.tachidesk.server.JavalinSetup.Attribute
+import suwayomi.tachidesk.server.JavalinSetup.getAttributeOrSet
 import suwayomi.tachidesk.server.user.UserType
 import suwayomi.tachidesk.server.user.getUserFromToken
 
@@ -152,10 +154,14 @@ class ApolloSubscriptionProtocolHandler(
         context: WsContext,
     ): Flow<SubscriptionOperationMessage> {
         @Suppress("UNCHECKED_CAST")
-        val payload = operationMessage.payload as? Map<String, Any?>
-        val token = payload?.let { it[Header.AUTHORIZATION] as? String }
+        val user =
+            context.getAttributeOrSet(Attribute.TachideskUser) {
+                val payload = operationMessage.payload as? Map<String, Any?>
+                val token = payload?.let { it[Header.AUTHORIZATION] as? String }
+                getUserFromToken(token)
+            }
 
-        saveContext(getUserFromToken(token), context)
+        saveContext(user, context)
         return flowOf(acknowledgeMessage)
     }
 
diff --git a/server/src/main/kotlin/suwayomi/tachidesk/server/JavalinSetup.kt b/server/src/main/kotlin/suwayomi/tachidesk/server/JavalinSetup.kt
index 73f2cf95..a2c2d27a 100644
--- a/server/src/main/kotlin/suwayomi/tachidesk/server/JavalinSetup.kt
+++ b/server/src/main/kotlin/suwayomi/tachidesk/server/JavalinSetup.kt
@@ -310,4 +310,11 @@ object JavalinSetup {
     fun <T : Any> Context.getAttribute(attribute: Attribute<T>): T = attribute(attribute.name)!!
 
     fun <T : Any> WsContext.getAttribute(attribute: Attribute<T>): T = attribute(attribute.name)!!
+
+    fun <T : Any> WsContext.getAttributeOrSet(
+        attribute: Attribute<T>,
+        set: () -> T,
+    ): T =
+        attribute(attribute.name)
+            ?: set().also { setAttribute(attribute, it) }
 }