Add onUpgrade method to WebSocket RequestLogger

javalin/src/main/java/io/javalin/http/servlet/JavalinWsServletContext.kt

@@ -10,9 +10,9 @@ class JavalinWsServletContext(
     cfg: JavalinServletContextConfig,
     req: HttpServletRequest,
     res: HttpServletResponse,
-    routeRoles: Set<RouteRole>,
-    matchedPath: String,
-    pathParamMap: Map<String, String>,
+    routeRoles: Set<RouteRole> = emptySet(),
+    matchedPath: String = "",
+    pathParamMap: Map<String, String> = emptyMap(),
 ) : JavalinServletContext(
     cfg = cfg,
     req = req,

javalin/src/main/java/io/javalin/jetty/JavalinJettyServlet.kt

@@ -51,16 +51,25 @@ class JavalinJettyServlet(val cfg: JavalinConfig) : JettyWebSocketServlet() {
     }
 
     private fun serviceWebSocketRequest(req: HttpServletRequest, res: HttpServletResponse) {
+        val requestStartTime = System.nanoTime()
         val requestUri = req.requestURI.removePrefix(req.contextPath)
         val wsRouterHandlerEntry = cfg.pvt.wsRouter.wsPathMatcher.findEndpointHandlerEntry(requestUri)
-            ?: return res.sendError(404, "WebSocket handler not found")
+        if (wsRouterHandlerEntry == null) {
+            res.sendError(404, "WebSocket handler not found")
+            // Still need to call upgrade logger for 404 cases
+            val upgradeContext = JavalinWsServletContext(servletContextConfig, req, res)
+            val executionTimeMs = calculateExecutionTimeMs(requestStartTime)
+            cfg.pvt.wsLogger?.wsUpgradeLogger?.handle(upgradeContext, executionTimeMs)
+            return
+        }
+
         val upgradeContext = JavalinWsServletContext(
             cfg = servletContextConfig,
             req = req,
             res = res,
+            routeRoles = wsRouterHandlerEntry.roles,
             matchedPath = wsRouterHandlerEntry.path,
             pathParamMap = wsRouterHandlerEntry.extractPathParams(requestUri),
-            routeRoles = wsRouterHandlerEntry.roles,
         )
         req.setAttribute(upgradeContextKey, upgradeContext)
         res.setWsProtocolHeader(req)
@@ -72,14 +81,21 @@ class JavalinJettyServlet(val cfg: JavalinConfig) : JettyWebSocketServlet() {
         // add after handlers
         cfg.pvt.internalRouter.findHttpHandlerEntries(HandlerType.WEBSOCKET_AFTER_UPGRADE, requestUri)
             .forEach { handler -> upgradeContext.tasks.offer(Task { handler.handle(upgradeContext, requestUri) }) }
-        while (upgradeContext.tasks.isNotEmpty()) { // execute all tasks in order
-            try {
-                val task = upgradeContext.tasks.poll()
-                task.handler.handle()
-            } catch (e: Exception) {
-                cfg.pvt.internalRouter.handleHttpException(upgradeContext, e)
-                break
+
+        try {
+            while (upgradeContext.tasks.isNotEmpty()) { // execute all tasks in order
+                try {
+                    val task = upgradeContext.tasks.poll()
+                    task.handler.handle()
+                } catch (e: Exception) {
+                    cfg.pvt.internalRouter.handleHttpException(upgradeContext, e)
+                    break
+                }
             }
+        } finally {
+            // Call the WebSocket upgrade logger whether successful or not
+            val executionTimeMs = calculateExecutionTimeMs(requestStartTime)
+            cfg.pvt.wsLogger?.wsUpgradeLogger?.handle(upgradeContext, executionTimeMs)
         }
     }
 
@@ -89,4 +105,7 @@ class JavalinJettyServlet(val cfg: JavalinConfig) : JettyWebSocketServlet() {
         this.setHeader(WebSocketConstants.SEC_WEBSOCKET_PROTOCOL, firstProtocol)
     }
 
+    private fun calculateExecutionTimeMs(startTimeNanos: Long): Float = 
+        (System.nanoTime() - startTimeNanos) / 1_000_000.0f
+
 }

javalin/src/main/java/io/javalin/websocket/WsConfig.java

@@ -18,6 +18,7 @@ public class WsConfig {
     WsBinaryMessageHandler wsBinaryMessageHandler = null;
     WsCloseHandler wsCloseHandler = null;
     WsErrorHandler wsErrorHandler = null;
+    public WsUpgradeLogger wsUpgradeLogger = null;
 
     /**
      * Add a WsConnectHandler to the WsHandler.
@@ -64,4 +65,15 @@ public void onError(@NotNull WsErrorHandler wsErrorHandler) {
         this.wsErrorHandler = wsErrorHandler;
     }
 
+    /**
+     * Add a WsUpgradeLogger to the WsHandler.
+     * The handler is called when an HTTP request attempts to upgrade to WebSocket.
+     * This is invoked after the upgrade process completes, whether the upgrade succeeds or fails.
+     */
+    public void onUpgrade(@NotNull WsUpgradeLogger wsUpgradeLogger) {
+        this.wsUpgradeLogger = wsUpgradeLogger;
+    }
+
+
+
 }

javalin/src/main/java/io/javalin/websocket/WsUpgradeLogger.java

@@ -0,0 +1,28 @@
+/*
+ * Javalin - https://javalin.io
+ * Copyright 2017 David Åse
+ * Licensed under Apache 2.0: https://github.com/tipsy/javalin/blob/master/LICENSE
+ */
+
+package io.javalin.websocket;
+
+import io.javalin.http.Context;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Interface for logging WebSocket upgrade requests.
+ *
+ * @see Context
+ * @see <a href="https://javalin.io/documentation#request-loggers">RequestLogger in documentation</a>
+ */
+@FunctionalInterface
+public interface WsUpgradeLogger {
+    /**
+     * Handles a WebSocket upgrade request
+     *
+     * @param ctx             the current request context during upgrade
+     * @param executionTimeMs the requests' execution time in milliseconds
+     * @throws Exception any exception while logging information about the request
+     */
+    void handle(@NotNull Context ctx, @NotNull Float executionTimeMs) throws Exception;
+}

javalin/src/test/java/io/javalin/TestWebSocket.kt

@@ -337,6 +337,57 @@ class TestWebSocket {
         )
     }
 
+    @Test
+    fun `web socket upgrade logging works`() = TestUtil.test(Javalin.create().apply {
+        this.unsafeConfig().requestLogger.ws { ws ->
+            ws.onUpgrade { ctx, executionTimeMs ->
+                this.logger().log.add("${ctx.path()} upgrade attempted (${ctx.status()})")
+            }
+        }
+    }) { app, _ ->
+        app.ws("/upgrade-test") {}
+        // Test successful upgrade
+        TestClient(app, "/upgrade-test").connectAndDisconnect()
+        // Test failed upgrade (404)
+        val response = Unirest.get("http://localhost:${app.port()}/non-existent-ws")
+            .header(Header.SEC_WEBSOCKET_KEY, "test-key")
+            .header(Header.UPGRADE, "websocket")
+            .header(Header.CONNECTION, "upgrade")
+            .asString()
+        assertThat(response.status).isEqualTo(404)
+
+        // Verify that both successful and failed upgrades are logged
+        assertThat(app.logger().log).containsExactlyInAnyOrder(
+            "/upgrade-test upgrade attempted (101 Switching Protocols)",
+            "/non-existent-ws upgrade attempted (404 Not Found)"
+        )
+    }
+
+    @Test
+    fun `web socket upgrade logging works for wsBeforeUpgrade errors`() = TestUtil.test(Javalin.create().apply {
+        this.unsafeConfig().requestLogger.ws { ws ->
+            ws.onUpgrade { ctx, executionTimeMs ->
+                this.logger().log.add("${ctx.path()} upgrade attempted (${ctx.status()})")
+            }
+        }
+    }) { app, _ ->
+        app.wsBeforeUpgrade("/auth-ws") { ctx ->
+            throw UnauthorizedResponse()
+        }
+        app.ws("/auth-ws") {}
+
+        // Test failed upgrade due to authentication
+        val response = Unirest.get("http://localhost:${app.port()}/auth-ws")
+            .header(Header.SEC_WEBSOCKET_KEY, "test-key")
+            .header(Header.UPGRADE, "websocket")
+            .header(Header.CONNECTION, "upgrade")
+            .asString()
+        assertThat(response.status).isEqualTo(401)
+
+        // Verify that the failed upgrade due to authentication error is logged
+        assertThat(app.logger().log).containsExactly("/auth-ws upgrade attempted (401 Unauthorized)")
+    }
+
     @Test
     fun `dev logging works for web sockets`() = TestUtil.test(Javalin.create { it.registerPlugin(DevLoggingPlugin()) }) { app, _ ->
         app.ws("/path/{param}") {}