Identifying clients with their User-Agent header

LavalinkServer/src/main/java/lavalink/server/io/SocketContext.kt

@@ -39,6 +39,7 @@ import moe.kyokobot.koe.KoeClient
 import moe.kyokobot.koe.KoeEventAdapter
 import moe.kyokobot.koe.MediaConnection
 import org.slf4j.LoggerFactory
+import org.slf4j.MDC
 import org.springframework.web.socket.CloseStatus
 import org.springframework.web.socket.WebSocketSession
 import org.springframework.web.socket.adapter.standard.StandardWebSocketSession
@@ -55,6 +56,7 @@ class SocketContext(
     statsCollector: StatsCollector,
     override val userId: Long,
     override val clientName: String?,
+    override val userAgent: String?,
     val koe: KoeClient,
     eventHandlers: Collection<PluginEventHandler>,
     private val pluginInfoModifiers: List<AudioPluginInfoModifier>,
@@ -159,6 +161,10 @@ class SocketContext(
             return
         }
 
+        if (userAgent != null) {
+            MDC.put("userAgent", userAgent)
+        }
+
         if (!session.isOpen) return
 
         val undertowSession = (session as StandardWebSocketSession).nativeSession as UndertowSession
@@ -194,14 +200,20 @@ class SocketContext(
     }
 
     internal fun shutdown() {
-        log.info("Shutting down ${playingPlayers.size} playing players.")
+        if (userAgent != null) {
+            MDC.put("userAgent", userAgent)
+        }
+
+        log.info("Shutting down ${playingPlayers.size} playing players for session $sessionId")
+
         executor.shutdown()
         playerUpdateService.shutdown()
         players.values.forEach {
             this.destroyPlayer(it.guildId)
         }
         koe.close()
         eventEmitter.onSocketContextDestroyed()
+        MDC.remove("userAgent")
     }
 
     override fun closeWebSocket(closeCode: Int, reason: String?) {

LavalinkServer/src/main/java/lavalink/server/io/SocketServer.kt

@@ -33,6 +33,7 @@ import lavalink.server.player.LavalinkPlayer
 import moe.kyokobot.koe.Koe
 import moe.kyokobot.koe.KoeOptions
 import org.slf4j.LoggerFactory
+import org.slf4j.MDC
 import org.springframework.stereotype.Service
 import org.springframework.web.socket.CloseStatus
 import org.springframework.web.socket.TextMessage
@@ -68,8 +69,8 @@ final class SocketServer(
 
             val connection = socketContext.getMediaConnection(player).gatewayConnection
             socketContext.sendMessage(
-                    Message.Serializer,
-                    Message.PlayerUpdateEvent(
+                Message.Serializer,
+                Message.PlayerUpdateEvent(
                     PlayerState(
                         System.currentTimeMillis(),
                         player.audioPlayer.playingTrack?.position ?: 0,
@@ -100,6 +101,11 @@ final class SocketServer(
         val clientName = session.handshakeHeaders.getFirst("Client-Name")
         val userAgent = session.handshakeHeaders.getFirst("User-Agent")
 
+        if (userAgent != null) {
+            session.attributes["userAgent"] = userAgent
+            MDC.put("userAgent", userAgent)
+        }
+
         var resumable: SocketContext? = null
         if (sessionId != null) resumable = resumableSessions.remove(sessionId)
 
@@ -124,6 +130,7 @@ final class SocketServer(
             statsCollector,
             userId,
             clientName,
+            userAgent,
             koe.newClient(userId),
             eventHandlers,
             pluginInfoModifiers
@@ -146,6 +153,12 @@ final class SocketServer(
 
     override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
         val context = sessions.remove(session.attributes["sessionId"]) ?: return
+        val userAgent = session.attributes["userAgent"] as? String
+
+        if (userAgent != null) {
+            MDC.put("userAgent", userAgent)
+        }
+
         if (context.resumable) {
             resumableSessions.remove(context.sessionId)?.let { removed ->
                 log.warn(

LavalinkServer/src/main/java/lavalink/server/io/UserAgentContextFilter.kt

@@ -0,0 +1,27 @@
+package lavalink.server.io
+
+import jakarta.servlet.http.HttpServletRequest
+import jakarta.servlet.http.HttpServletResponse
+import org.slf4j.MDC
+import org.springframework.stereotype.Component
+import org.springframework.web.filter.OncePerRequestFilter
+
+@Component
+class UserAgentContextFilter : OncePerRequestFilter() {
+
+    override fun doFilterInternal(
+        request: HttpServletRequest,
+        response: HttpServletResponse,
+        filterChain: jakarta.servlet.FilterChain
+    ) {
+        try {
+            val userAgent = request.getHeader("User-Agent")
+            if (userAgent != null) {
+                MDC.put("userAgent", userAgent)
+            }
+            filterChain.doFilter(request, response)
+        } finally {
+            MDC.remove("userAgent")
+        }
+    }
+}

LavalinkServer/src/main/java/lavalink/server/io/UserAgentWebSocketInterceptor.kt

@@ -0,0 +1,24 @@
+package lavalink.server.io
+
+import org.springframework.stereotype.Component
+import org.springframework.web.socket.WebSocketHandler
+import org.springframework.web.socket.WebSocketSession
+import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory
+
+@Component
+class UserAgentWebSocketInterceptor : WebSocketHandlerDecoratorFactory {
+    override fun decorate(handler: WebSocketHandler): WebSocketHandler {
+        return UserAgentWebSocketHandlerDecorator(handler)
+    }
+
+    private class UserAgentWebSocketHandlerDecorator(private val delegate: WebSocketHandler) :
+        WebSocketHandler by delegate {
+        override fun afterConnectionEstablished(session: WebSocketSession) {
+            val userAgent = session.handshakeHeaders.getFirst("User-Agent")
+            if (userAgent != null) {
+                session.attributes["userAgent"] = userAgent
+            }
+            delegate.afterConnectionEstablished(session)
+        }
+    }
+}

LavalinkServer/src/main/resources/application.yml

@@ -9,3 +9,6 @@ server:
     include-binding-errors: ALWAYS
     include-stacktrace: ON_PARAM
     include-exception: false
+logging:
+  pattern:
+    console: "%clr(%d{${LOG_DATEFORMAT_PATTERN:-yyyy-MM-dd'T'HH:mm:ss.SSSXXX}}){faint} %clr(${LOG_LEVEL_PATTERN:-%5p}) %clr(${PID:- }){magenta} %clr(---){faint} %clr([%15.15t]){faint} %clr(%-40.40logger{39}){cyan} %clr(%replace(%X{userAgent}){'^(.+)$',' ($1):'}){magenta} %m%n${LOG_EXCEPTION_CONVERSION_WORD:-%wEx}"

docs/api/rest.md

@@ -6,9 +6,11 @@ description: Lavalink REST API documentation.
 
 Lavalink exposes a REST API to allow for easy control of the players.
 Most routes require the `Authorization` header with the configured password.
+An aditional `User-Agent` header can be set to easily identify the client on the server logs.
 
 ```
 Authorization: youshallnotpass
+User-Agent: TestBot/1.0.0
 ```
 
 Routes are prefixed with `/v3` as of `v3.7.0` and `/v4` as of `v4.0.0`. Routes without an API prefix were removed in v4 (except `/version`).

docs/api/websocket.md

@@ -15,6 +15,7 @@ When opening a websocket connection, you must supply 3 required headers:
 | `Authorization` | The password you set in your Lavalink config    |
 | `User-Id`       | The user id of the bot                          |
 | `Client-Name`   | The name of the client in `NAME/VERSION` format |
+| `User-Agent`? * | Information about the connecting client         |
 | `Session-Id`? * | The id of the previous session to resume        |
 
 **\*For more information on resuming see [Resuming](index.md#resuming)**
@@ -26,6 +27,7 @@ When opening a websocket connection, you must supply 3 required headers:
 Authorization: youshallnotpass
 User-Id: 170939974227541168
 Client-Name: lavalink-client/2.0.0
+User-Agent: TestBot/1.0.0
 ```
 
 </details>

plugin-api/src/main/java/dev/arbjerg/lavalink/api/ISocketContext.kt

@@ -22,6 +22,11 @@ interface ISocketContext {
      */
     val clientName: String?
 
+    /**
+     * The User Agent of the Client if specified.
+     */
+    val userAgent: String?
+
     /**
      * A read-only map of all players associated by their guild.
      */