From 4127f2414f25631a14a8bd9d362303be27c605ea Mon Sep 17 00:00:00 2001
From: topi314 <git@topi.wtf>
Date: Mon, 25 Dec 2023 22:31:09 +0100
Subject: [PATCH] add new plugin ISocketServer interface to expose lavalink
 sessions to plugins

---
 .../java/lavalink/server/io/SocketServer.kt     | 17 +++++++++--------
 .../src/main/java/lavalink/server/util/util.kt  |  2 +-
 .../dev/arbjerg/lavalink/api/ISocketServer.kt   | 17 +++++++++++++++++
 3 files changed, 27 insertions(+), 9 deletions(-)
 create mode 100644 plugin-api/src/main/java/dev/arbjerg/lavalink/api/ISocketServer.kt

diff --git a/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt b/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
index 5cf087783..6ffbdfc61 100644
--- a/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
+++ b/LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt
@@ -24,6 +24,7 @@ package lavalink.server.io
 
 import com.sedmelluq.discord.lavaplayer.player.AudioPlayerManager
 import dev.arbjerg.lavalink.api.AudioPluginInfoModifier
+import dev.arbjerg.lavalink.api.ISocketServer
 import dev.arbjerg.lavalink.api.PluginEventHandler
 import dev.arbjerg.lavalink.protocol.v4.Message
 import dev.arbjerg.lavalink.protocol.v4.PlayerState
@@ -46,11 +47,11 @@ final class SocketServer(
     koeOptions: KoeOptions,
     private val eventHandlers: List<PluginEventHandler>,
     private val pluginInfoModifiers: List<AudioPluginInfoModifier>
-) : TextWebSocketHandler() {
+) : TextWebSocketHandler(), ISocketServer {
 
     // sessionID <-> Session
-    val contextMap = ConcurrentHashMap<String, SocketContext>()
-    private val resumableSessions = mutableMapOf<String, SocketContext>()
+    override val sessions = ConcurrentHashMap<String, SocketContext>()
+    override val resumableSessions = mutableMapOf<String, SocketContext>()
     private val koe = Koe.koe(koeOptions)
     private val statsCollector = StatsCollector(this)
     private val charPool = ('a'..'z') + ('0'..'9')
@@ -81,12 +82,12 @@ final class SocketServer(
         var sessionId: String
         do {
             sessionId = List(16) { charPool.random() }.joinToString("")
-        } while (contextMap[sessionId] != null)
+        } while (sessions[sessionId] != null)
         return sessionId
     }
 
     val contexts: Collection<SocketContext>
-        get() = contextMap.values
+        get() = sessions.values
 
     @Suppress("UastIncorrectHttpHeaderInspection")
     override fun afterConnectionEstablished(session: WebSocketSession) {
@@ -100,7 +101,7 @@ final class SocketServer(
 
         if (resumable != null) {
             session.attributes["sessionId"] = resumable.sessionId
-            contextMap[resumable.sessionId] = resumable
+            sessions[resumable.sessionId] = resumable
             resumable.resume(session)
             log.info("Resumed session with id $sessionId")
             resumable.eventEmitter.onWebSocketOpen(true)
@@ -123,7 +124,7 @@ final class SocketServer(
             eventHandlers,
             pluginInfoModifiers
         )
-        contextMap[sessionId] = socketContext
+        sessions[sessionId] = socketContext
         socketContext.sendMessage(Message.Serializer, Message.ReadyEvent(false, sessionId))
         socketContext.eventEmitter.onWebSocketOpen(false)
         if (clientName != null) {
@@ -140,7 +141,7 @@ final class SocketServer(
     }
 
     override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
-        val context = contextMap.remove(session.attributes["sessionId"]) ?: return
+        val context = sessions.remove(session.attributes["sessionId"]) ?: return
         if (context.resumable) {
             resumableSessions.remove(context.sessionId)?.let { removed ->
                 log.warn(
diff --git a/LavalinkServer/src/main/java/lavalink/server/util/util.kt b/LavalinkServer/src/main/java/lavalink/server/util/util.kt
index a3b63d39e..040b532a8 100644
--- a/LavalinkServer/src/main/java/lavalink/server/util/util.kt
+++ b/LavalinkServer/src/main/java/lavalink/server/util/util.kt
@@ -122,7 +122,7 @@ fun getRootCause(throwable: Throwable?): Throwable {
 }
 
 fun socketContext(socketServer: SocketServer, sessionId: String) =
-    socketServer.contextMap[sessionId] ?: throw ResponseStatusException(HttpStatus.NOT_FOUND, "Session not found")
+    socketServer.sessions[sessionId] ?: throw ResponseStatusException(HttpStatus.NOT_FOUND, "Session not found")
 
 fun existingPlayer(socketContext: SocketContext, guildId: Long) =
     socketContext.players[guildId] ?: throw ResponseStatusException(HttpStatus.NOT_FOUND, "Player not found")
diff --git a/plugin-api/src/main/java/dev/arbjerg/lavalink/api/ISocketServer.kt b/plugin-api/src/main/java/dev/arbjerg/lavalink/api/ISocketServer.kt
new file mode 100644
index 000000000..d54627e20
--- /dev/null
+++ b/plugin-api/src/main/java/dev/arbjerg/lavalink/api/ISocketServer.kt
@@ -0,0 +1,17 @@
+package dev.arbjerg.lavalink.api
+
+/**
+ * Represents a Lavalink server which handles WebSocket connections.
+ */
+interface ISocketServer {
+    /**
+     * A map of all active sessions by their session id.
+     */
+    val sessions: Map<String, ISocketContext>
+
+    /**
+     * A map of all resumable sessions by their session id.
+     * A session is resumable if the client configured resuming and has disconnected.
+     */
+    val resumableSessions: Map<String, ISocketContext>
+}
\ No newline at end of file