Repository: tipsy/javalin
Pull request #905 (2020-03-23 10:45:25)

[websocket] Add support for Jetty sessions in WebSockets (alt 2)


Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index ac2ea501..8217c8be 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -50,7 +50,7 @@ io.javalin
 │   └── rendering/              // Interface for file rendering, and several template engine implementations
 ├── websocket/                  // Everything related to WebSockets
 │   ├── WsContext.kt            // Wrapper class for WebSockets
-│   ├── JavalinWsServlet.kt     // Responsible for WebSocket upgrade
+│   ├── JavalinWsServlet.kt     // Responsible for WebSocket upgrade, as well as switching between WebSocket and HTTP
 │   └── WsHandlerController.kt  // Responsible for Websocket request lifecycle (before, endpoint, after, logging)
 └── Javalin.java                // Main public API, responsible for setting up and configuring the server
 ```
diff --git a/src/main/java/io/javalin/Javalin.java b/src/main/java/io/javalin/Javalin.java
index 3fa1cdd3..88bb467a 100644
--- a/src/main/java/io/javalin/Javalin.java
+++ b/src/main/java/io/javalin/Javalin.java
@@ -59,7 +59,7 @@ public class Javalin {
 
     protected Javalin() {
         this.server = new JavalinServer(config);
-        this.wsServlet = new JavalinWsServlet(config);
+        this.wsServlet = new JavalinWsServlet(config, servlet);
     }
 
     public Javalin(JavalinServer server, JavalinWsServlet wsServlet) {
@@ -169,7 +169,7 @@ public class Javalin {
         eventManager.fireEvent(JavalinEvent.SERVER_STARTING);
         try {
             Javalin.log.info("Starting Javalin ...");
-            server.start(servlet, wsServlet);
+            server.start(wsServlet);
             Javalin.log.info("Javalin started in " + (System.currentTimeMillis() - startupTimer) + "ms \\o/");
             eventManager.fireEvent(JavalinEvent.SERVER_STARTED);
         } catch (Exception e) {
diff --git a/src/main/java/io/javalin/core/JavalinConfig.java b/src/main/java/io/javalin/core/JavalinConfig.java
index 37f04088..987eda71 100644
--- a/src/main/java/io/javalin/core/JavalinConfig.java
+++ b/src/main/java/io/javalin/core/JavalinConfig.java
@@ -54,7 +54,6 @@ public class JavalinConfig {
     @NotNull public String contextPath = "/";
     @NotNull public Long requestCacheSize = 4096L;
     @NotNull public Long asyncRequestTimeout = 0L;
-    @NotNull public String wsContextPath = "/";
     @NotNull public Inner inner = new Inner();
 
     // it's not bad to access this, the main reason it's hidden
diff --git a/src/main/java/io/javalin/core/JavalinServer.kt b/src/main/java/io/javalin/core/JavalinServer.kt
index 70755a68..cb3639ae 100644
--- a/src/main/java/io/javalin/core/JavalinServer.kt
+++ b/src/main/java/io/javalin/core/JavalinServer.kt
@@ -7,21 +7,17 @@
 package io.javalin.core
 
 import io.javalin.Javalin
-import io.javalin.http.JavalinServlet
 import io.javalin.websocket.JavalinWsServlet
-import io.javalin.websocket.isWebSocket
 import org.eclipse.jetty.server.Handler
 import org.eclipse.jetty.server.LowResourceMonitor
 import org.eclipse.jetty.server.Request
 import org.eclipse.jetty.server.Server
 import org.eclipse.jetty.server.ServerConnector
 import org.eclipse.jetty.server.handler.HandlerCollection
-import org.eclipse.jetty.server.handler.HandlerList
 import org.eclipse.jetty.server.handler.HandlerWrapper
 import org.eclipse.jetty.server.handler.StatisticsHandler
 import org.eclipse.jetty.server.session.SessionHandler
 import org.eclipse.jetty.servlet.ServletContextHandler
-import org.eclipse.jetty.servlet.ServletContextHandler.SESSIONS
 import org.eclipse.jetty.servlet.ServletHolder
 import org.eclipse.jetty.util.thread.QueuedThreadPool
 import java.net.BindException
@@ -41,14 +37,13 @@ class JavalinServer(val config: JavalinConfig) {
     var started = false
 
     @Throws(BindException::class)
-    fun start(javalinServlet: JavalinServlet, javalinWsServlet: JavalinWsServlet) {
+    fun start(wsAndHttpServlet: JavalinWsServlet) {
 
         config.inner.sessionHandler = config.inner.sessionHandler ?: defaultSessionHandler()
         val nullParent = null // javalin handlers are orphans
 
-        val httpHandler = object : ServletContextHandler(nullParent, config.contextPath, SESSIONS) {
+        val wsAndHttpHandler = object : ServletContextHandler(nullParent, config.contextPath, SESSIONS) {
             override fun doHandle(target: String, jettyRequest: Request, request: HttpServletRequest, response: HttpServletResponse) {
-                if (request.isWebSocket()) return // don't touch websocket requests
                 request.setAttribute("jetty-target", target) // used in JettyResourceHandler
                 request.setAttribute("jetty-request", jettyRequest)
                 nextHandle(target, jettyRequest, request, response)
@@ -56,15 +51,11 @@ class JavalinServer(val config: JavalinConfig) {
         }.apply {
             this.sessionHandler = config.inner.sessionHandler
             config.inner.servletContextHandlerConsumer?.accept(this)
-            addServlet(ServletHolder(javalinServlet), "/*")
-        }
-
-        val webSocketHandler = ServletContextHandler(nullParent, javalinWsServlet.config.wsContextPath, SESSIONS).apply {
-            addServlet(ServletHolder(javalinWsServlet), "/*")
+            addServlet(ServletHolder(wsAndHttpServlet), "/*")
         }
 
         server().apply {
-            handler = attachJavalinHandlers(server.handler, HandlerList(httpHandler, webSocketHandler))
+            handler = if (server.handler == null) wsAndHttpHandler else server.handler.attachHandler(wsAndHttpHandler)
             connectors = connectors.takeIf { it.isNotEmpty() } ?: arrayOf(ServerConnector(server).apply {
                 this.port = serverPort
                 this.host = serverHost
@@ -87,21 +78,20 @@ class JavalinServer(val config: JavalinConfig) {
 
     private val ServerConnector.protocol get() = if (protocols.contains("ssl")) "https" else "http"
 
-    private fun attachJavalinHandlers(userHandler: Handler?, javalinHandlers: HandlerList) = when (userHandler) {
-        null -> HandlerWrapper().apply { handler = javalinHandlers } // no custom Handler set, wrap Javalin handlers in a HandlerWrapper
-        is HandlerCollection -> userHandler.apply { addHandler(javalinHandlers) } // user is using a HandlerCollection, add Javalin handlers to it
-        is HandlerWrapper -> userHandler.apply {
-            (unwrap(this) as? HandlerCollection)?.addHandler(javalinHandlers) // if HandlerWrapper unwraps as HandlerCollection, add Javalin handlers
-            (unwrap(this) as? HandlerWrapper)?.handler = javalinHandlers // if HandlerWrapper unwraps as HandlerWrapper, add Javalin last
+    private fun Handler.attachHandler(servletContextHandler: ServletContextHandler) = when (this) {
+        is HandlerCollection -> this.apply { addHandler(servletContextHandler) } // user is using a HandlerCollection, add Javalin handler to it
+        is HandlerWrapper -> this.apply {
+            (this.unwrap() as? HandlerCollection)?.addHandler(servletContextHandler) // if HandlerWrapper unwraps as HandlerCollection, add Javalin handler
+            (this.unwrap() as? HandlerWrapper)?.handler = servletContextHandler // if HandlerWrapper unwraps as HandlerWrapper, add Javalin last
         }
-        else -> throw IllegalStateException("Server has unidentified handler attached to it")
+        else -> throw IllegalStateException("Server has unsupported Handler attached to it (must be HandlerCollection or HandlerWrapper)")
     }
 
-    private fun unwrap(userHandler: HandlerWrapper): Handler = when (userHandler.handler) {
-        null -> userHandler // current HandlerWrapper is last element, return the HandlerWrapper
-        is HandlerCollection -> userHandler.handler // HandlerWrapper wraps HandlerCollection, return HandlerCollection
-        is HandlerWrapper -> unwrap(userHandler.handler as HandlerWrapper) // HandlerWrapper wraps another HandlerWrapper, recursive call required
-        else -> throw IllegalStateException("Cannot insert Javalin handlers into a Handler that is not a HandlerCollection or HandlerWrapper")
+    private fun HandlerWrapper.unwrap(): Handler = when (this.handler) {
+        null -> this // current HandlerWrapper is last element, return the HandlerWrapper itself
+        is HandlerCollection -> this.handler // HandlerWrapper wraps HandlerCollection, return HandlerCollection
+        is HandlerWrapper -> (this.handler as HandlerWrapper).unwrap() // HandlerWrapper wraps another HandlerWrapper, recursive call required
+        else -> throw IllegalStateException("HandlerWrapper has unsupported Handler type (must be HandlerCollection or HandlerWrapper")
     }
 
 }
diff --git a/src/main/java/io/javalin/websocket/JavalinWsServlet.kt b/src/main/java/io/javalin/websocket/JavalinWsServlet.kt
index ec596b5b..5ea5cd0b 100644
--- a/src/main/java/io/javalin/websocket/JavalinWsServlet.kt
+++ b/src/main/java/io/javalin/websocket/JavalinWsServlet.kt
@@ -6,12 +6,11 @@
 
 package io.javalin.websocket
 
-import io.javalin.Javalin
 import io.javalin.core.JavalinConfig
 import io.javalin.core.security.Role
 import io.javalin.core.util.Header
-import io.javalin.core.util.Util
 import io.javalin.http.Context
+import io.javalin.http.JavalinServlet
 import io.javalin.http.UnauthorizedResponse
 import io.javalin.http.util.ContextUtil
 import org.eclipse.jetty.websocket.servlet.WebSocketCreator
@@ -21,40 +20,48 @@ import java.util.function.Consumer
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
 
-class JavalinWsServlet(val config: JavalinConfig) : WebSocketServlet() {
+internal const val upgradeAllowedKey = "javalin-ws-upgrade-allowed"
+internal const val upgradeContextKey = "javalin-ws-upgrade-context"
+internal const val upgradeSessionAttrsKey = "javalin-ws-upgrade-http-session"
+
+/**
+ * The JavalinWsServlet is responsible for both WebSocket and HTTP requests.
+ * It extends Jetty's WebSocketServlet, and has a HTTP Servlet as a constructor arg.
+ * It switches between WebSocket and HTTP in the `service` method.
+ */
+class JavalinWsServlet(val config: JavalinConfig, private val httpServlet: JavalinServlet) : WebSocketServlet() {
 
     val wsExceptionMapper = WsExceptionMapper()
 
-    val wsPathMatcher = WsPathMatcher()
+    private val wsPathMatcher = WsPathMatcher()
 
     override fun configure(factory: WebSocketServletFactory) {
         config.inner.wsFactoryConfig?.accept(factory)
         factory.creator = WebSocketCreator { req, res ->
-            val preUpgradeContext = req.httpServletRequest.getAttribute("javalin-ws-upgrade-context") as Context
-            req.httpServletRequest.setAttribute("javalin-ws-upgrade-context", ContextUtil.changeBaseRequest(preUpgradeContext, req.httpServletRequest))
-            WsHandlerController(wsPathMatcher, wsExceptionMapper, config.inner.wsLogger)
+            val preUpgradeContext = req.httpServletRequest.getAttribute(upgradeContextKey) as Context
+            req.httpServletRequest.setAttribute(upgradeContextKey, ContextUtil.changeBaseRequest(preUpgradeContext, req.httpServletRequest))
+            req.httpServletRequest.setAttribute(upgradeSessionAttrsKey, req.session?.attributeNames?.asSequence()?.associateWith { req.session.getAttribute(it) })
+            return@WebSocketCreator WsHandlerController(wsPathMatcher, wsExceptionMapper, config.inner.wsLogger)
         }
     }
 
-    override fun service(req: HttpServletRequest, res: HttpServletResponse) {
-        if (req.isWebSocket()) {
-            val requestUri = req.requestURI.removePrefix(req.contextPath)
-            val entry = wsPathMatcher.findEndpointHandlerEntry(requestUri) ?: return res.sendError(404, "WebSocket handler not found")
-            try {
-                val upgradeContext = Context(req, res).apply {
-                    pathParamMap = entry.extractPathParams(requestUri)
-                    matchedPath = entry.path
-                }
-                config.inner.accessManager.manage({ ctx -> ctx.req.setAttribute("javalin-ws-upgrade-allowed", "true") }, upgradeContext, entry.permittedRoles)
-                if (req.getAttribute("javalin-ws-upgrade-allowed") != "true") throw UnauthorizedResponse() // if set to true, the access manager ran the handler (== valid)
-                req.setAttribute("javalin-ws-upgrade-context", upgradeContext)
-                super.service(req, res)
-            } catch (e: Exception) {
-                res.sendError(401, "Unauthorized")
+    override fun service(req: HttpServletRequest, res: HttpServletResponse) { // this handles both http and websocket
+        if (!req.isWebSocket()) {
+            return httpServlet.service(req, res)
+        }
+        val requestUri = req.requestURI.removePrefix(req.contextPath)
+        val entry = wsPathMatcher.findEndpointHandlerEntry(requestUri) ?: return res.sendError(404, "WebSocket handler not found")
+        try {
+            val upgradeContext = Context(req, res).apply {
+                pathParamMap = entry.extractPathParams(requestUri)
+                matchedPath = entry.path
             }
-        } else { // if not websocket (and not handled by http-handler), this request is below the context path
-            Util.writeResponse(res, "Not found. Request is below context-path", 404)
-            Javalin.log?.warn("Received a request below context-path. Returned 404.")
+            config.inner.accessManager.manage({ ctx -> ctx.req.setAttribute(upgradeAllowedKey, true) }, upgradeContext, entry.permittedRoles)
+            if (req.getAttribute(upgradeAllowedKey) != true) throw UnauthorizedResponse() // if set to true, the access manager ran the handler (== valid)
+            req.setAttribute(upgradeContextKey, upgradeContext)
+            super.service(req, res) // everything is okay, perform websocket upgrade
+        } catch (e: Exception) {
+            res.sendError(401, "Unauthorized")
         }
     }
 
diff --git a/src/main/java/io/javalin/websocket/WsContext.kt b/src/main/java/io/javalin/websocket/WsContext.kt
index 37b53a8e..7f701cf4 100644
--- a/src/main/java/io/javalin/websocket/WsContext.kt
+++ b/src/main/java/io/javalin/websocket/WsContext.kt
@@ -20,8 +20,9 @@ import java.nio.ByteBuffer
  */
 abstract class WsContext(val sessionId: String, @JvmField val session: Session) {
 
-    private val upgradeReq = session.upgradeRequest as ServletUpgradeRequest
-    private val upgradeCtx = upgradeReq.httpServletRequest.getAttribute("javalin-ws-upgrade-context") as Context
+    private val upgradeReq by lazy { session.upgradeRequest as ServletUpgradeRequest }
+    private val upgradeCtx by lazy { upgradeReq.httpServletRequest.getAttribute(upgradeContextKey) as Context }
+    private val sessionAttributes by lazy { upgradeReq.httpServletRequest.getAttribute(upgradeSessionAttrsKey) as Map<String, Any>? }
 
     fun matchedPath() = upgradeCtx.matchedPath
 
@@ -54,6 +55,9 @@ abstract class WsContext(val sessionId: String, @JvmField val session: Session)
     fun <T> attribute(key: String): T? = upgradeCtx.attribute(key)
     fun <T> attributeMap(): Map<String, T?> = upgradeCtx.attributeMap()
 
+    fun <T> sessionAttribute(key: String): T? = sessionAttributeMap()[key] as T
+    fun sessionAttributeMap(): Map<String, Any?> = sessionAttributes ?: mapOf()
+
     override fun equals(other: Any?) = session == (other as WsContext).session
     override fun hashCode() = session.hashCode()
 }


Test patch:
diff --git a/src/test/java/io/javalin/TestConfiguration.kt b/src/test/java/io/javalin/TestConfiguration.kt
index c97629fa..915099d8 100644
--- a/src/test/java/io/javalin/TestConfiguration.kt
+++ b/src/test/java/io/javalin/TestConfiguration.kt
@@ -45,7 +45,6 @@ class TestConfiguration {
             it.requestLogger { ctx, executionTimeMs -> }
             it.sessionHandler { SessionHandler() }
             // WsServlet
-            it.wsContextPath = "/"
             it.wsFactoryConfig { }
             it.wsLogger { }
             // Server
diff --git a/src/test/java/io/javalin/TestContextPath.kt b/src/test/java/io/javalin/TestContextPath.kt
index babb7f11..b3813fd7 100644
--- a/src/test/java/io/javalin/TestContextPath.kt
+++ b/src/test/java/io/javalin/TestContextPath.kt
@@ -42,7 +42,7 @@ class TestContextPath {
         val javalin = Javalin.create { it.contextPath = "/context-path" }
         TestUtil.test(javalin) { app, http ->
             app.get("/hello") { ctx -> ctx.result("Hello World") }
-            assertThat(http.getBody("/hello")).isEqualTo("Not found. Request is below context-path")
+            assertThat(http.get("/hello").status).isEqualTo(404)
             assertThat(http.getBody("/context-path/hello")).isEqualTo("Hello World")
         }
     }
diff --git a/src/test/java/io/javalin/TestCustomJetty.kt b/src/test/java/io/javalin/TestCustomJetty.kt
index 636bbcf2..b21f3727 100644
--- a/src/test/java/io/javalin/TestCustomJetty.kt
+++ b/src/test/java/io/javalin/TestCustomJetty.kt
@@ -15,8 +15,6 @@ import org.eclipse.jetty.server.RequestLog
 import org.eclipse.jetty.server.Server
 import org.eclipse.jetty.server.handler.ContextHandlerCollection
 import org.eclipse.jetty.server.handler.HandlerCollection
-import org.eclipse.jetty.server.handler.HandlerList
-import org.eclipse.jetty.server.handler.HandlerWrapper
 import org.eclipse.jetty.server.handler.RequestLogHandler
 import org.eclipse.jetty.server.handler.StatisticsHandler
 import org.eclipse.jetty.server.session.DefaultSessionCache
@@ -118,7 +116,7 @@ class TestCustomJetty {
             it.sessionHandler { fileSessionHandler }
             it.server { newServer }
         }.start(0)
-        val httpHandler = (((newServer.handlers[0] as HandlerWrapper).handler as HandlerList).handlers.first() as ServletContextHandler)
+        val httpHandler = (newServer.handlers[0] as ServletContextHandler)
         assertThat(httpHandler.sessionHandler).isEqualTo(fileSessionHandler)
         javalin.stop()
 
diff --git a/src/test/java/io/javalin/TestWebSocket.kt b/src/test/java/io/javalin/TestWebSocket.kt
index 761117a7..328b2be6 100644
--- a/src/test/java/io/javalin/TestWebSocket.kt
+++ b/src/test/java/io/javalin/TestWebSocket.kt
@@ -40,7 +40,7 @@ class TestWebSocket {
         return this.attribute(TestLogger::class.java)
     }
 
-    fun contextPathJavalin() = Javalin.create { it.wsContextPath = "/websocket" }
+    fun contextPathJavalin() = Javalin.create { it.contextPath = "/websocket" }
 
     fun javalinWithWsLogger() = Javalin.create().apply {
         this.config.wsLogger { ws ->
@@ -282,6 +282,20 @@ class TestWebSocket {
         assertThat(app.logger().log).containsExactly("Success")
     }
 
+    @Test
+    fun `getting session attributes works`() = TestUtil.test { app, http ->
+        app.get("/") { ctx -> ctx.sessionAttribute("session-key", "session-value") }
+        app.ws("/") { ws ->
+            ws.onConnect {
+                app.logger().log.add(it.sessionAttribute("session-key")!!)
+                app.logger().log.add("sessionAttributeMapSize:${it.sessionAttributeMap().size}")
+            }
+        }
+        val sessionCookie = http.get("/").headers["Set-Cookie"]!!.first().removePrefix("JSESSIONID=")
+        TestClient(app, "/", mapOf("Cookie" to "JSESSIONID=${sessionCookie}")).connectAndDisconnect()
+        assertThat(app.logger().log).containsExactly("session-value", "sessionAttributeMapSize:1")
+    }
+
     @Test
     fun `routing and path-params case sensitive works`() = TestUtil.test { app, _ ->
         app.ws("/pAtH/:param") { ws -> ws.onConnect { ctx -> app.logger().log.add(ctx.pathParam("param")) } }