Repository: tipsy/javalin
Pull request #540 (2019-04-09 21:04:25)

[websocket] Create access manager for websocket

tipsy: @dherges @rcornell @Playacem @klassm @SE7-KN8 you have all worked on the WebSocket implementation in some shape or form. Any opinions on this approach? 

I intercept the upgrade request and add a `Handler` which sets an attribute on the context if the current `AccessManager` accets the request. If this attribute is not set (`AccessManager` did not accept the request), I throw an exception and return 401.

Patch:
diff --git a/src/main/java/io/javalin/Javalin.java b/src/main/java/io/javalin/Javalin.java
index aea50a86..1b6ede98 100644
--- a/src/main/java/io/javalin/Javalin.java
+++ b/src/main/java/io/javalin/Javalin.java
@@ -537,8 +537,12 @@ public class Javalin {
      * @see <a href="https://javalin.io/documentation#websockets">WebSockets in docs</a>
      */
     public Javalin ws(@NotNull String path, @NotNull Consumer<WsHandler> ws) {
-        wsServlet.addHandler(path, ws);
-        eventManager.fireHandlerAddedEvent(new HandlerMetaInfo(HandlerType.WEBSOCKET, Util.prefixContextPath(wsServlet.getConfig().contextPath, path), ws, new HashSet<>()));
+        return ws(path, ws, new HashSet<>());
+    }
+
+    public Javalin ws(@NotNull String path, @NotNull Consumer<WsHandler> ws, @NotNull Set<Role> permittedRoles) {
+        wsServlet.addHandler(path, ws, permittedRoles);
+        eventManager.fireHandlerAddedEvent(new HandlerMetaInfo(HandlerType.WEBSOCKET, Util.prefixContextPath(wsServlet.getConfig().contextPath, path), ws, permittedRoles));
         return this;
     }
 
diff --git a/src/main/java/io/javalin/core/JavalinConfig.java b/src/main/java/io/javalin/core/JavalinConfig.java
index f929ae45..75b906c5 100644
--- a/src/main/java/io/javalin/core/JavalinConfig.java
+++ b/src/main/java/io/javalin/core/JavalinConfig.java
@@ -41,7 +41,7 @@ public class JavalinConfig {
     public List<String> corsOrigins = new ArrayList<>();
     RequestLogger requestLogger;
     ResourceHandler resourceHandler;
-    AccessManager accessManager = SecurityUtil::noopAccessManager;
+    public AccessManager accessManager = SecurityUtil::noopAccessManager;
     SinglePageHandler singlePageHandler = new SinglePageHandler();
     SessionHandler sessionHandler;
 
diff --git a/src/main/java/io/javalin/core/JavalinWsServlet.kt b/src/main/java/io/javalin/core/JavalinWsServlet.kt
index 78b58ca1..8c5c98bc 100644
--- a/src/main/java/io/javalin/core/JavalinWsServlet.kt
+++ b/src/main/java/io/javalin/core/JavalinWsServlet.kt
@@ -6,9 +6,12 @@
 
 package io.javalin.core
 
+import io.javalin.Context
 import io.javalin.Javalin
+import io.javalin.UnauthorizedResponse
 import io.javalin.core.util.Header
 import io.javalin.core.util.Util
+import io.javalin.security.Role
 import io.javalin.websocket.WsEntry
 import io.javalin.websocket.WsHandler
 import io.javalin.websocket.WsPathMatcher
@@ -16,8 +19,6 @@ import org.eclipse.jetty.websocket.servlet.WebSocketCreator
 import org.eclipse.jetty.websocket.servlet.WebSocketServlet
 import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory
 import java.util.function.Consumer
-import javax.servlet.ServletRequest
-import javax.servlet.ServletResponse
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
 
@@ -27,20 +28,29 @@ class JavalinWsServlet(val config: JavalinConfig) : WebSocketServlet() {
 
     override fun configure(factory: WebSocketServletFactory) {
         config.wsFactoryConfig?.accept(factory)
-        factory.creator = WebSocketCreator { req, res ->
-            wsPathMatcher.findEntry(req) ?: res.sendError(404, "WebSocket handler not found")
-            wsPathMatcher // this is a long-lived object handling multiple connections
-        }
+        factory.creator = WebSocketCreator { req, res -> wsPathMatcher } // this is a long-lived object handling multiple connections
     }
 
-    override fun service(req: ServletRequest?, res: ServletResponse?) = if ((req as HttpServletRequest).isWebSocket()) {
-        super.service(req, res) // if websocket, handle normally
-    } else { // if not websocket (and not handled by http-handler), this request is below the context path
-        Util.writeResponse(res as HttpServletResponse, "Not found. Request is below context-path", 404)
-        Javalin.log.warn("Received a request below context-path. Returned 404.")
+    override fun service(req: HttpServletRequest, res: HttpServletResponse) {
+        if (req.isWebSocket()) {
+            val entry = wsPathMatcher.findEntry(req.requestURI.removePrefix(req.contextPath)) ?: return res.sendError(404, "WebSocket handler not found")
+            try {
+                config.accessManager.manage({ ctx -> ctx.req.setAttribute("javalin-ws-upgrade-allowed", "true") }, Context(req, res, mapOf()), entry.permittedRoles)
+                if (req.getAttribute("javalin-ws-upgrade-allowed") != "true") throw UnauthorizedResponse() // if set to true, the access manager ran the handler (== valid)
+                super.service(req, res)
+            } catch (e: Exception) {
+                res.sendError(401, "Unauthorized")
+            }
+        } 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.")
+        }
     }
 
-    fun addHandler(path: String, ws: Consumer<WsHandler>) = wsPathMatcher.add(WsEntry(path, WsHandler().apply { ws.accept(this) }))
+    fun addHandler(path: String, ws: Consumer<WsHandler>, permittedRoles: Set<Role>) {
+        wsPathMatcher.add(WsEntry(path, WsHandler().apply { ws.accept(this) }, permittedRoles))
+    }
 }
 
 fun HttpServletRequest.isWebSocket(): Boolean = this.getHeader(Header.SEC_WEBSOCKET_KEY) != null
+
diff --git a/src/main/java/io/javalin/core/util/ContextUtil.kt b/src/main/java/io/javalin/core/util/ContextUtil.kt
index c7b9a18f..f2bcf6e6 100644
--- a/src/main/java/io/javalin/core/util/ContextUtil.kt
+++ b/src/main/java/io/javalin/core/util/ContextUtil.kt
@@ -38,12 +38,6 @@ object ContextUtil {
 
     fun urlDecode(s: String): String = URLDecoder.decode(s.replace("+", "%2B"), "UTF-8").replace("%2B", "+")
 
-    fun mapKeysOrReturnNullIfAnyNulls(keys: Array<out String>, f: (s: String) -> String?): List<String>? = try {
-        keys.map { f.invoke(it) }.requireNoNulls().toList()
-    } catch (e: IllegalArgumentException) {
-        null
-    }
-
     fun getBasicAuthCredentials(header: String?): BasicAuthCredentials? = try {
         val (username, password) = String(Base64.getDecoder().decode(header!!.removePrefix("Basic "))).split(":")
         BasicAuthCredentials(username, password)
diff --git a/src/main/java/io/javalin/websocket/WsPathMatcher.kt b/src/main/java/io/javalin/websocket/WsPathMatcher.kt
index d7da80a5..0286511b 100644
--- a/src/main/java/io/javalin/websocket/WsPathMatcher.kt
+++ b/src/main/java/io/javalin/websocket/WsPathMatcher.kt
@@ -8,6 +8,7 @@ package io.javalin.websocket
 
 import io.javalin.core.JavalinConfig
 import io.javalin.core.PathParser
+import io.javalin.security.Role
 import org.eclipse.jetty.websocket.api.Session
 import org.eclipse.jetty.websocket.api.UpgradeRequest
 import org.eclipse.jetty.websocket.api.annotations.*
@@ -15,7 +16,7 @@ import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest
 import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 
-data class WsEntry(val path: String, val handler: WsHandler) {
+data class WsEntry(val path: String, val handler: WsHandler, val permittedRoles: Set<Role>) {
     private val pathParser = PathParser(path)
     fun matches(requestUri: String) = pathParser.matches(requestUri)
     fun extractPathParams(requestUri: String) = pathParser.extractPathParams(requestUri)
@@ -84,8 +85,8 @@ class WsPathMatcher(val config: JavalinConfig) {
     }
 
     private fun findEntry(session: Session) = findEntry(session.upgradeRequest)?.also { cacheSession(session, it) }
-
-    fun findEntry(req: UpgradeRequest) = wsEntries.find { it.matches(req.uriNoContextPath()) }
+    fun findEntry(req: UpgradeRequest) = findEntry(req.uriNoContextPath())
+    fun findEntry(uri: String) = wsEntries.find { it.matches(uri) }
 
     private fun cacheSession(session: Session, wsEntry: WsEntry) {
         sessionIds.putIfAbsent(session, UUID.randomUUID().toString())


Test patch:
diff --git a/src/test/java/io/javalin/TestWebSocket.kt b/src/test/java/io/javalin/TestWebSocket.kt
index 0adf238e..2793f117 100644
--- a/src/test/java/io/javalin/TestWebSocket.kt
+++ b/src/test/java/io/javalin/TestWebSocket.kt
@@ -29,17 +29,38 @@ import java.util.concurrent.atomic.AtomicInteger
  */
 class TestWebSocket {
 
-    private val contextPathJavalin = Javalin.create().configure { it.wsContextPath = "/websocket" }
-    private val javalinWithWsLogger = Javalin.create().configure {
-        it.wsLogger { ws ->
-            ws.onConnect { ctx -> log.add(ctx.pathParam("param") + " connected") }
-            ws.onClose { ctx -> log.add(ctx.pathParam("param") + " disconnected") }
+    private val contextPathJavalin: Javalin by lazy { Javalin.create().configure { it.wsContextPath = "/websocket" } }
+    private val javalinWithWsLogger: Javalin by lazy {
+        Javalin.create().configure {
+            it.wsLogger { ws ->
+                ws.onConnect { ctx -> log.add(ctx.pathParam("param") + " connected") }
+                ws.onClose { ctx -> log.add(ctx.pathParam("param") + " disconnected") }
+            }
+        }
+    }
+    val accessManagedJavalin: Javalin by lazy {
+        Javalin.create().configure {
+            it.accessManager { handler, ctx, permittedRoles ->
+                log.add("handling upgrade request ...")
+                if (ctx.queryParam("allowed") == "true") {
+                    log.add("upgrade request valid!")
+                    handler.handle(ctx)
+                } else if (ctx.queryParam("exception") == "true") {
+                    throw UnauthorizedResponse()
+                } else {
+                    log.add("upgrade request invalid!")
+                }
+            }
+        }.ws("/*") { ws ->
+            ws.onConnect { log.add("connected with upgrade request") }
         }
     }
+
     private var log = mutableListOf<String>()
 
     @Before
     fun resetLog() {
+        Thread.sleep(25)
         log = ArrayList()
     }
 
@@ -311,6 +332,26 @@ class TestWebSocket {
         assertThat(err).isExactlyInstanceOf(MessageTooLargeException::class.java)
     }
 
+    @Test
+    fun `accessmanager rejects invalid request`() = TestUtil.test(accessManagedJavalin) { app, _ ->
+        connectAndDisconnect(TestClient(URI.create("ws://localhost:" + app.port() + "/")))
+        assertThat(log.size).isEqualTo(2)
+        assertThat(log).containsExactlyInAnyOrder("handling upgrade request ...", "upgrade request invalid!")
+    }
+
+    @Test
+    fun `accessmanager accepts valid request`() = TestUtil.test(accessManagedJavalin) { app, _ ->
+        connectAndDisconnect(TestClient(URI.create("ws://localhost:" + app.port() + "/?allowed=true")))
+        assertThat(log.size).isEqualTo(3)
+        assertThat(log).containsExactlyInAnyOrder("handling upgrade request ...", "upgrade request valid!", "connected with upgrade request")
+    }
+
+    @Test
+    fun `accessmanager doesn't crash on exception`() = TestUtil.test(accessManagedJavalin) { app, _ ->
+        connectAndDisconnect(TestClient(URI.create("ws://localhost:" + app.port() + "/?exception=true")))
+        assertThat(log.size).isEqualTo(1)
+    }
+
     // ********************************************************************************************
     // Helpers
     // ********************************************************************************************