Repository: tipsy/javalin
Pull request #558 (2019-04-18 10:35:06)

Websocket before after handler

tipsy: Looks great @peterbencze, thank you!

Patch:
diff --git a/src/main/java/io/javalin/Javalin.java b/src/main/java/io/javalin/Javalin.java
index b12a6f60..6d80e586 100644
--- a/src/main/java/io/javalin/Javalin.java
+++ b/src/main/java/io/javalin/Javalin.java
@@ -18,6 +18,8 @@ import io.javalin.core.JavalinEvent;
 import io.javalin.core.JavalinServer;
 import io.javalin.core.JavalinServlet;
 import io.javalin.core.JavalinWsServlet;
+import io.javalin.core.WsHandlerMetaInfo;
+import io.javalin.core.WsHandlerType;
 import io.javalin.core.util.CorsBeforeHandler;
 import io.javalin.core.util.CorsOptionsHandler;
 import io.javalin.core.util.RouteOverviewRenderer;
@@ -26,8 +28,10 @@ import io.javalin.security.AccessManager;
 import io.javalin.security.CoreRoles;
 import io.javalin.security.Role;
 import io.javalin.security.SecurityUtil;
+import static io.javalin.security.SecurityUtil.roles;
 import io.javalin.serversentevent.SseClient;
 import io.javalin.serversentevent.SseHandler;
+import io.javalin.websocket.WsExceptionHandler;
 import io.javalin.websocket.WsHandler;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -37,7 +41,6 @@ import java.util.function.Consumer;
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import static io.javalin.security.SecurityUtil.roles;
 
 public class Javalin {
 
@@ -196,6 +199,14 @@ public class Javalin {
         return this;
     }
 
+    /**
+     * Adds a WebSocket exception mapper to the instance.
+     */
+    public <T extends Exception> Javalin wsException(@NotNull Class<T> exceptionClass, @NotNull WsExceptionHandler<? super T> exceptionHandler) {
+        wsServlet.getWsExceptionMapper().addHandler(exceptionClass, (WsExceptionHandler<Exception>) exceptionHandler);
+        return this;
+    }
+
     /**
      * Adds an error mapper to the instance.
      * Useful for turning error-codes (404, 500) into standardized messages/pages
@@ -254,6 +265,21 @@ public class Javalin {
         return addHandler(httpMethod, path, handler, new HashSet<>()); // no roles set for this route (open to everyone with default access manager)
     }
 
+    /**
+     * Adds a specific WebSocket handler for the given path to the instance.
+     * Requires an access manager to be set on the instance.
+     */
+    public Javalin addHandler(@NotNull WsHandlerType handlerType, @NotNull String path, @NotNull Consumer<WsHandler> wsHandler, @NotNull Set<Role> roles) {
+        wsServlet.addHandler(handlerType, path, wsHandler, roles);
+        eventManager.fireWsHandlerAddedEvent(new WsHandlerMetaInfo(handlerType, Util.prefixContextPath(servlet.getConfig().contextPath, path), wsHandler, roles));
+        return this;
+    }
+
+    /** Adds a specific WebSocket handler for the given path to the instance. */
+    public Javalin addHandler(@NotNull WsHandlerType handlerType, @NotNull String path, @NotNull Consumer<WsHandler> wsHandler) {
+        return addHandler(handlerType, path, wsHandler, new HashSet<>());
+    }
+
     // ********************************************************************************************
     // HTTP verbs
     // ********************************************************************************************
@@ -522,9 +548,27 @@ 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, @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;
+        return addHandler(WsHandlerType.WEBSOCKET, path, ws, permittedRoles);
+    }
+
+    /** Adds a WebSocket before handler for the specified path to the instance. */
+    public Javalin wsBefore(@NotNull String path, @NotNull Consumer<WsHandler> wsHandler) {
+        return addHandler(WsHandlerType.WEBSOKET_BEFORE, path, wsHandler);
+    }
+
+    /** Adds a WebSocket before handler for all routes in the instance. */
+    public Javalin wsBefore(@NotNull Consumer<WsHandler> wsHandler) {
+        return wsBefore("*", wsHandler);
+    }
+
+    /** Adds a WebSocket after handler for the specified path to the instance. */
+    public Javalin wsAfter(@NotNull String path, @NotNull Consumer<WsHandler> wsHandler) {
+        return addHandler(WsHandlerType.WEBSOCKET_AFTER, path, wsHandler);
+    }
+
+    /** Adds a WebSocket after handler for all routes in the instance. */
+    public Javalin wsAfter(@NotNull Consumer<WsHandler> wsHandler) {
+        return wsAfter("*", wsHandler);
     }
 
 }
diff --git a/src/main/java/io/javalin/apibuilder/ApiBuilder.java b/src/main/java/io/javalin/apibuilder/ApiBuilder.java
index 8c5b0695..681f3c65 100644
--- a/src/main/java/io/javalin/apibuilder/ApiBuilder.java
+++ b/src/main/java/io/javalin/apibuilder/ApiBuilder.java
@@ -8,6 +8,7 @@ package io.javalin.apibuilder;
 
 import io.javalin.Handler;
 import io.javalin.Javalin;
+import io.javalin.core.WsHandlerType;
 import io.javalin.security.AccessManager;
 import io.javalin.security.Role;
 import io.javalin.serversentevent.SseClient;
@@ -405,6 +406,38 @@ public class ApiBuilder {
         staticInstance().ws(prefixPath(""), ws, permittedRoles);
     }
 
+    /**
+     * Adds a WebSocket before handler for the specified path to the {@link Javalin} instance.
+     * The method can only be called inside a {@link Javalin#routes(EndpointGroup)}.
+     */
+    public Javalin wsBefore(@NotNull String path, @NotNull Consumer<WsHandler> wsHandler) {
+        return staticInstance().wsBefore(prefixPath(path), wsHandler);
+    }
+
+    /**
+     * Adds a WebSocket before handler for the current path to the {@link Javalin} instance.
+     * The method can only be called inside a {@link Javalin#routes(EndpointGroup)}.
+     */
+    public Javalin wsBefore(@NotNull Consumer<WsHandler> wsHandler) {
+        return staticInstance().wsBefore(prefixPath("/*"), wsHandler);
+    }
+
+    /**
+     * Adds a WebSocket after handler for the specified path to the {@link Javalin} instance.
+     * The method can only be called inside a {@link Javalin#routes(EndpointGroup)}.
+     */
+    public Javalin wsAfter(@NotNull String path, @NotNull Consumer<WsHandler> wsHandler) {
+        return staticInstance().wsAfter(prefixPath(path), wsHandler);
+    }
+
+    /**
+     * Adds a WebSocket after handler for the current path to the {@link Javalin} instance.
+     * The method can only be called inside a {@link Javalin#routes(EndpointGroup)}.
+     */
+    public Javalin wsAfter(@NotNull Consumer<WsHandler> wsHandler) {
+        return staticInstance().wsAfter(prefixPath("/*"), wsHandler);
+    }
+
     // ********************************************************************************************
     // Server-sent events
     // ********************************************************************************************
diff --git a/src/main/java/io/javalin/core/EventAttacher.java b/src/main/java/io/javalin/core/EventAttacher.java
index b8317e63..c7046c22 100644
--- a/src/main/java/io/javalin/core/EventAttacher.java
+++ b/src/main/java/io/javalin/core/EventAttacher.java
@@ -31,6 +31,12 @@ public class EventAttacher {
         return eventManager.getParentJavalin();
     }
 
+    public Javalin wsHandlerAdded(@NotNull Consumer<WsHandlerMetaInfo> callback) {
+        eventManager.getWsHandlerAddedCallbacks().add(callback);
+
+        return eventManager.getParentJavalin();
+    }
+
     private Javalin addLifecycleEvent(@NotNull JavalinEvent event, @NotNull Runnable callback) {
         eventManager.getCallbackMap().get(event).add(callback);
         return eventManager.getParentJavalin();
diff --git a/src/main/java/io/javalin/core/EventManager.kt b/src/main/java/io/javalin/core/EventManager.kt
index 6f41e7fd..75ffedf4 100644
--- a/src/main/java/io/javalin/core/EventManager.kt
+++ b/src/main/java/io/javalin/core/EventManager.kt
@@ -17,7 +17,13 @@ class EventManager(val parentJavalin: Javalin) {
     fun fireEvent(javalinEvent: JavalinEvent) = callbackMap[javalinEvent]!!.forEach { callback -> callback.run() }
     var handlerAddedCallbacks = mutableSetOf<Consumer<HandlerMetaInfo>>()
     fun fireHandlerAddedEvent(handlerMetaInfo: HandlerMetaInfo) = handlerAddedCallbacks.apply { this.forEach { it.accept(handlerMetaInfo) } }
+
+    val wsHandlerAddedCallbacks = mutableSetOf<Consumer<WsHandlerMetaInfo>>()
+    fun fireWsHandlerAddedEvent(wsHandlerMetaInfo: WsHandlerMetaInfo) =
+            wsHandlerAddedCallbacks.apply { this.forEach { it.accept(wsHandlerMetaInfo) } }
 }
 
 enum class JavalinEvent { SERVER_STARTING, SERVER_STARTED, SERVER_START_FAILED, SERVER_STOPPING, SERVER_STOPPED }
 data class HandlerMetaInfo(val httpMethod: HandlerType, val path: String, val handler: Any, val roles: Set<Role>)
+
+data class WsHandlerMetaInfo(val handlerType: WsHandlerType, val path: String, val handler: Any, val roles: Set<Role>)
diff --git a/src/main/java/io/javalin/core/HandlerType.kt b/src/main/java/io/javalin/core/HandlerType.kt
index 7410dbf3..718abc53 100644
--- a/src/main/java/io/javalin/core/HandlerType.kt
+++ b/src/main/java/io/javalin/core/HandlerType.kt
@@ -10,7 +10,7 @@ import io.javalin.core.util.Header
 import javax.servlet.http.HttpServletRequest
 
 enum class HandlerType {
-    GET, POST, PUT, PATCH, DELETE, HEAD, TRACE, CONNECT, OPTIONS, BEFORE, AFTER, INVALID, WEBSOCKET;
+    GET, POST, PUT, PATCH, DELETE, HEAD, TRACE, CONNECT, OPTIONS, BEFORE, AFTER, INVALID;
 
     fun isHttpMethod() = when (this) {
         GET, HEAD, POST, PUT, DELETE, CONNECT, OPTIONS, TRACE, PATCH -> true
diff --git a/src/main/java/io/javalin/core/JavalinWsServlet.kt b/src/main/java/io/javalin/core/JavalinWsServlet.kt
index 1febe374..54099247 100644
--- a/src/main/java/io/javalin/core/JavalinWsServlet.kt
+++ b/src/main/java/io/javalin/core/JavalinWsServlet.kt
@@ -13,9 +13,7 @@ import io.javalin.core.util.ContextUtil
 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
+import io.javalin.websocket.*
 import org.eclipse.jetty.websocket.servlet.WebSocketCreator
 import org.eclipse.jetty.websocket.servlet.WebSocketServlet
 import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory
@@ -25,21 +23,24 @@ import javax.servlet.http.HttpServletResponse
 
 class JavalinWsServlet(val config: JavalinConfig) : WebSocketServlet() {
 
-    var wsPathMatcher = WsPathMatcher(config)
+    val wsExceptionMapper = WsExceptionMapper()
+
+    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))
-            wsPathMatcher // this is a long-lived object handling multiple connections
+
+            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.findEntry(requestUri) ?: return res.sendError(404, "WebSocket handler not found")
+            val entry = wsPathMatcher.findEndpointHandlerEntry(requestUri) ?: return res.sendError(404, "WebSocket handler not found")
             try {
                 val upgradeContext = Context(req, res).apply {
                     pathParamMap = entry.extractPathParams(requestUri)
@@ -58,8 +59,8 @@ class JavalinWsServlet(val config: JavalinConfig) : WebSocketServlet() {
         }
     }
 
-    fun addHandler(path: String, ws: Consumer<WsHandler>, permittedRoles: Set<Role>) {
-        wsPathMatcher.wsEntries.add(WsEntry(path, WsHandler().apply { ws.accept(this) }, permittedRoles))
+    fun addHandler(handlerType: WsHandlerType, path: String, ws: Consumer<WsHandler>, permittedRoles: Set<Role>) {
+        wsPathMatcher.add(WsEntry(handlerType, path, WsHandler().apply { ws.accept(this) }, permittedRoles))
     }
 }
 
diff --git a/src/main/java/io/javalin/core/WsHandlerType.kt b/src/main/java/io/javalin/core/WsHandlerType.kt
new file mode 100644
index 00000000..b3fb9b7d
--- /dev/null
+++ b/src/main/java/io/javalin/core/WsHandlerType.kt
@@ -0,0 +1,14 @@
+/*
+ * Javalin - https://javalin.io
+ * Copyright 2017 David Åse
+ * Licensed under Apache 2.0: https://github.com/tipsy/javalin/blob/master/LICENSE
+ */
+
+package io.javalin.core
+
+/**
+ * The available WebSocket handler types.
+ */
+enum class WsHandlerType {
+    WEBSOKET_BEFORE, WEBSOCKET, WEBSOCKET_AFTER
+}
diff --git a/src/main/java/io/javalin/websocket/WsExceptionHandler.java b/src/main/java/io/javalin/websocket/WsExceptionHandler.java
new file mode 100644
index 00000000..28553134
--- /dev/null
+++ b/src/main/java/io/javalin/websocket/WsExceptionHandler.java
@@ -0,0 +1,20 @@
+/*
+ * 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 org.jetbrains.annotations.NotNull;
+
+/**
+ * A handler for use with {@link io.javalin.Javalin#wsException(Class, WsExceptionHandler)}. Is
+ * triggered when an exception is thrown by a {@link WsHandler}.
+ *
+ * @see WsContext
+ */
+@FunctionalInterface
+public interface WsExceptionHandler<T extends Exception> {
+    void handle(@NotNull T exception, @NotNull WsContext ctx);
+}
diff --git a/src/main/java/io/javalin/websocket/WsExceptionMapper.kt b/src/main/java/io/javalin/websocket/WsExceptionMapper.kt
new file mode 100644
index 00000000..14c40a72
--- /dev/null
+++ b/src/main/java/io/javalin/websocket/WsExceptionMapper.kt
@@ -0,0 +1,46 @@
+/*
+ * 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.Javalin
+import org.eclipse.jetty.websocket.api.StatusCode
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.reflect.full.superclasses
+
+/**
+ * Maps exception types to exception handlers.
+ */
+class WsExceptionMapper {
+
+    private val handlers = ConcurrentHashMap<Class<out Exception>, WsExceptionHandler<Exception>?>()
+
+    /** Associates an [exceptionHandler] with the specific [exceptionClass]. */
+    fun addHandler(exceptionClass: Class<out Exception>, exceptionHandler: WsExceptionHandler<Exception>?) =
+            handlers.put(exceptionClass, exceptionHandler)
+
+    /**
+     * Handles the specific [exception]. If no handler is associated with the exception, then the
+     * socket is closed with a status code that indicates internal error.
+     */
+    fun handle(exception: Exception, ctx: WsContext) {
+        val handler: WsExceptionHandler<Exception>? = findHandler(exception.javaClass)
+        if (handler != null) {
+            handler.handle(exception, ctx)
+        } else {
+            Javalin.log.warn("Uncaught exception in WebSocket handler", exception)
+            ctx.session.close(StatusCode.SERVER_ERROR, exception.message)
+        }
+    }
+
+    private fun findHandler(exceptionClass: Class<out Exception>): WsExceptionHandler<Exception>? {
+        return handlers.getOrElse(exceptionClass) {
+            exceptionClass.kotlin.superclasses
+                    .firstOrNull { superclass -> handlers.containsKey(superclass.java) }
+                    ?.let { superclass -> handlers[superclass.java] }
+        }
+    }
+}
diff --git a/src/main/java/io/javalin/websocket/WsHandlerController.kt b/src/main/java/io/javalin/websocket/WsHandlerController.kt
new file mode 100644
index 00000000..26abb5a7
--- /dev/null
+++ b/src/main/java/io/javalin/websocket/WsHandlerController.kt
@@ -0,0 +1,159 @@
+/*
+ * 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 org.eclipse.jetty.websocket.api.Session
+import org.eclipse.jetty.websocket.api.annotations.*
+import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest
+import java.util.*
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * A special WebSocket handler which delegates the handling of the specific WebSocket event to the
+ * matching custom handlers.
+ */
+@WebSocket
+class WsHandlerController(private val matcher: WsPathMatcher,
+                          private val exceptionMapper: WsExceptionMapper,
+                          private val wsLogger: WsHandler?) {
+
+    private val sessionIds = ConcurrentHashMap<Session, String>()
+
+    @OnWebSocketConnect
+    fun onConnect(session: Session) {
+        // Associate a unique ID with this new session
+        val sessionId = UUID.randomUUID().toString()
+        sessionIds[session] = sessionId
+
+        val requestUri = session.uriNoContextPath()
+        val ctx = WsConnectContext(sessionId, session)
+
+        try {
+            // Invoke before handlers (if any)
+            matcher.findBeforeHandlerEntries(requestUri).forEach { it.handler.connectHandler?.handleConnect(ctx) }
+
+            // Invoke endpoint handler (can't be null, otherwise the ws servlet would not have invoked the controller)
+            matcher.findEndpointHandlerEntry(requestUri)!!.let { it.handler.connectHandler?.handleConnect(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        // Invoke after handlers (if any)
+        try {
+            matcher.findAfterHandlerEntries(requestUri).forEach { it.handler.connectHandler?.handleConnect(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        wsLogger?.connectHandler?.handleConnect(ctx)
+    }
+
+    @OnWebSocketMessage
+    fun onMessage(session: Session, message: String) {
+        val requestUri = session.uriNoContextPath()
+        val ctx = WsMessageContext(sessionIds[session]!!, session, message)
+
+        try {
+            // Invoke before handlers (if any)
+            matcher.findBeforeHandlerEntries(requestUri).forEach { it.handler.messageHandler?.handleMessage(ctx) }
+
+            // Invoke endpoint handler (can't be null, otherwise the ws servlet would not have invoked the controller)
+            matcher.findEndpointHandlerEntry(requestUri)!!.let { it.handler.messageHandler?.handleMessage(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        // Invoke after handlers (if any)
+        try {
+            matcher.findAfterHandlerEntries(requestUri).forEach { it.handler.messageHandler?.handleMessage(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        wsLogger?.messageHandler?.handleMessage(ctx)
+    }
+
+    @OnWebSocketMessage
+    fun onMessage(session: Session, buffer: ByteArray, offset: Int, length: Int) {
+        val requestUri = session.uriNoContextPath()
+        val ctx = WsBinaryMessageContext(sessionIds[session]!!, session, buffer.toTypedArray(), offset, length)
+
+        try {
+            // Invoke before handlers (if any)
+            matcher.findBeforeHandlerEntries(requestUri).forEach { it.handler.binaryMessageHandler?.handleBinaryMessage(ctx) }
+
+            // Invoke endpoint handler (can't be null, otherwise the ws servlet would not have invoked the controller)
+            matcher.findEndpointHandlerEntry(requestUri)!!.let { it.handler.binaryMessageHandler?.handleBinaryMessage(ctx) }
+        } catch (e: java.lang.Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        // Invoke after handlers (if any)
+        try {
+            matcher.findAfterHandlerEntries(requestUri).forEach { it.handler.binaryMessageHandler?.handleBinaryMessage(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        wsLogger?.binaryMessageHandler?.handleBinaryMessage(ctx)
+    }
+
+    @OnWebSocketClose
+    fun onClose(session: Session, statusCode: Int, reason: String?) {
+        val requestUri = session.uriNoContextPath()
+        val ctx = WsCloseContext(sessionIds[session]!!, session, statusCode, reason)
+
+        try {
+            // Invoke before handlers (if any)
+            matcher.findBeforeHandlerEntries(requestUri).forEach { it.handler.closeHandler?.handleClose(ctx) }
+
+            // Invoke endpoint handler (can't be null, otherwise the ws servlet would not have invoked the controller)
+            matcher.findEndpointHandlerEntry(requestUri)!!.let { it.handler.closeHandler?.handleClose(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        // Invoke after handlers (if any)
+        try {
+            matcher.findAfterHandlerEntries(requestUri).forEach { it.handler.closeHandler?.handleClose(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        wsLogger?.closeHandler?.handleClose(ctx)
+
+        // The socket has been closed, we no longer need to keep track of the session ID
+        sessionIds.remove(session)
+    }
+
+    @OnWebSocketError
+    fun onError(session: Session, throwable: Throwable?) {
+        val requestUri = session.uriNoContextPath()
+        val ctx = WsErrorContext(sessionIds[session]!!, session, throwable)
+
+        try {
+            // Invoke before handlers (if any)
+            matcher.findBeforeHandlerEntries(requestUri).forEach { it.handler.errorHandler?.handleError(ctx) }
+
+            // Invoke endpoint handler (can't be null, otherwise the ws servlet would not have invoked the controller)
+            matcher.findEndpointHandlerEntry(requestUri)!!.let { it.handler.errorHandler?.handleError(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        // Invoke after handlers (if any)
+        try {
+            matcher.findAfterHandlerEntries(requestUri).forEach { it.handler.errorHandler?.handleError(ctx) }
+        } catch (e: Exception) {
+            exceptionMapper.handle(e, ctx)
+        }
+
+        wsLogger?.errorHandler?.handleError(ctx)
+    }
+}
+
+private fun Session.uriNoContextPath() = this.upgradeRequest.requestURI.path.removePrefix((this.upgradeRequest as ServletUpgradeRequest).httpServletRequest.contextPath)
diff --git a/src/main/java/io/javalin/websocket/WsPathMatcher.kt b/src/main/java/io/javalin/websocket/WsPathMatcher.kt
index b34ebfeb..c79a6b9e 100644
--- a/src/main/java/io/javalin/websocket/WsPathMatcher.kt
+++ b/src/main/java/io/javalin/websocket/WsPathMatcher.kt
@@ -6,79 +6,45 @@
 
 package io.javalin.websocket
 
-import io.javalin.core.JavalinConfig
 import io.javalin.core.PathParser
+import io.javalin.core.WsHandlerType
 import io.javalin.security.Role
-import org.eclipse.jetty.websocket.api.Session
-import org.eclipse.jetty.websocket.api.annotations.*
-import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest
 import java.util.*
-import java.util.concurrent.ConcurrentHashMap
 
-data class WsEntry(val path: String, val handler: WsHandler, val permittedRoles: Set<Role>) {
+data class WsEntry(val type: WsHandlerType, 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)
+    fun matches(path: String) = pathParser.matches(path)
+    fun extractPathParams(path: String) = pathParser.extractPathParams(path)
 }
 
 /**
- * Every WebSocket request passes through a single instance of this class.
- * Session IDs are generated and tracked here, and path-parameters are cached for performance.
+ * Performs match operations on WebSocket paths.
  */
-@WebSocket
-class WsPathMatcher(val config: JavalinConfig) {
+class WsPathMatcher {
 
-    val wsEntries = mutableListOf<WsEntry>()
-    private val sessionIds = ConcurrentHashMap<Session, String>()
+    private val wsHandlerEntries = WsHandlerType.values()
+            .associateTo(EnumMap<WsHandlerType, MutableList<WsEntry>>(WsHandlerType::class.java)) {
+                it to mutableListOf()
+            }
 
-    @OnWebSocketConnect
-    fun webSocketConnect(session: Session) {
-        val entry = cacheSessionAndFindEntry(session)
-        val ctx = WsConnectContext(sessionIds[session]!!, session)
-        entry.handler.connectHandler?.handleConnect(ctx)
-        config.inner.wsLogger?.connectHandler?.handleConnect(ctx)
-    }
-
-    @OnWebSocketMessage
-    fun webSocketMessage(session: Session, message: String) {
-        val entry = cacheSessionAndFindEntry(session)
-        val ctx = WsMessageContext(sessionIds[session]!!, session, message)
-        entry.handler.messageHandler?.handleMessage(ctx)
-        config.inner.wsLogger?.messageHandler?.handleMessage(ctx)
-    }
-
-    @OnWebSocketMessage
-    fun webSocketBinaryMessage(session: Session, buffer: ByteArray, offset: Int, length: Int) {
-        val entry = cacheSessionAndFindEntry(session)
-        val ctx = WsBinaryMessageContext(sessionIds[session]!!, session, buffer.toTypedArray(), offset, length)
-        entry.handler.binaryMessageHandler?.handleBinaryMessage(ctx)
-        config.inner.wsLogger?.binaryMessageHandler?.handleBinaryMessage(ctx)
-    }
+    fun add(entry: WsEntry) {
+        if (wsHandlerEntries[entry.type]!!.find { it.type == entry.type && it.path == entry.path } != null) {
+            throw IllegalArgumentException("Handler with type='${entry.type}' and path='${entry.path}' already exists.")
+        }
 
-    @OnWebSocketError
-    fun webSocketError(session: Session, throwable: Throwable?) {
-        val entry = cacheSessionAndFindEntry(session)
-        val ctx = WsErrorContext(sessionIds[session]!!, session, throwable)
-        entry.handler.errorHandler?.handleError(ctx)
-        config.inner.wsLogger?.errorHandler?.handleError(ctx)
+        wsHandlerEntries[entry.type]!!.add(entry)
     }
 
-    @OnWebSocketClose
-    fun webSocketClose(session: Session, statusCode: Int, reason: String?) {
-        val entry = cacheSessionAndFindEntry(session)
-        val ctx = WsCloseContext(sessionIds[session]!!, session, statusCode, reason)
-        entry.handler.closeHandler?.handleClose(ctx)
-        config.inner.wsLogger?.closeHandler?.handleClose(ctx)
-        sessionIds.remove(session)
-    }
+    /** Returns all the before handlers that match the given [path]. */
+    fun findBeforeHandlerEntries(path: String) = findEntries(WsHandlerType.WEBSOKET_BEFORE, path)
 
-    fun findEntry(uri: String) = wsEntries.find { it.matches(uri) }
+    /** Returns the first endpoint handler that match the given [path], or `null`. */
+    fun findEndpointHandlerEntry(path: String) = findEntries(WsHandlerType.WEBSOCKET, path).firstOrNull()
 
-    private fun cacheSessionAndFindEntry(session: Session): WsEntry {
-        sessionIds.putIfAbsent(session, UUID.randomUUID().toString())
-        return findEntry(session.uriNoContextPath())!! // can't be null, has 404 handler in front
-    }
+    /** Returns all the after handlers that match the given [path]. */
+    fun findAfterHandlerEntries(path: String) = findEntries(WsHandlerType.WEBSOCKET_AFTER, path)
 
+    /** Returns all the handlers of type [handlerType] that match the given [path]. */
+    private fun findEntries(handlerType: WsHandlerType, path: String) =
+            wsHandlerEntries[handlerType]!!.filter { entry -> entry.path == "*" || entry.matches(path) }
 }
-
-private fun Session.uriNoContextPath() = this.upgradeRequest.requestURI.path.removePrefix((this.upgradeRequest as ServletUpgradeRequest).httpServletRequest.contextPath)


Test patch:
diff --git a/src/test/java/io/javalin/TestWebSocket.kt b/src/test/java/io/javalin/TestWebSocket.kt
index 741dee30..ff959835 100644
--- a/src/test/java/io/javalin/TestWebSocket.kt
+++ b/src/test/java/io/javalin/TestWebSocket.kt
@@ -10,15 +10,19 @@ import com.mashape.unirest.http.Unirest
 import io.javalin.apibuilder.ApiBuilder.ws
 import io.javalin.json.JavalinJson
 import io.javalin.misc.SerializeableObject
+import io.javalin.misc.TypedException
 import io.javalin.websocket.WsContext
 import org.assertj.core.api.Assertions.assertThat
 import org.eclipse.jetty.websocket.api.MessageTooLargeException
+import org.eclipse.jetty.websocket.api.StatusCode
 import org.java_websocket.client.WebSocketClient
 import org.java_websocket.drafts.Draft_6455
 import org.java_websocket.handshake.ServerHandshake
 import org.junit.Test
 import java.net.URI
+import java.time.Duration
 import java.util.*
+import java.util.concurrent.TimeoutException
 import java.util.concurrent.atomic.AtomicInteger
 
 /**
@@ -280,8 +284,13 @@ class TestWebSocket {
     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")) } }
         app.ws("/other-path/:param") { ws -> ws.onConnect { ctx -> app.logger().log.add(ctx.pathParam("param")) } }
-        TestClient(app, "/PaTh/my-param").connectAndDisconnect()
+
+        val client = TestClient(app, "/PaTh/my-param")
+
+        doAndSleepWhile({ client.connect() }, { !client.isClosed })
+
         TestClient(app, "/other-path/My-PaRaM").connectAndDisconnect()
+
         assertThat(app.logger().log).doesNotContain("my-param")
         assertThat(app.logger().log).contains("My-PaRaM")
     }
@@ -346,7 +355,10 @@ class TestWebSocket {
 
     @Test
     fun `accessmanager rejects invalid request`() = TestUtil.test(accessManagedJavalin()) { app, _ ->
-        TestClient(app, "/").connectAndDisconnect()
+        val client = TestClient(app, "/")
+
+        doAndSleepWhile({ client.connect() }, { !client.isClosed })
+
         assertThat(app.logger().log.size).isEqualTo(2)
         assertThat(app.logger().log).containsExactlyInAnyOrder("handling upgrade request ...", "upgrade request invalid!")
     }
@@ -360,7 +372,10 @@ class TestWebSocket {
 
     @Test
     fun `accessmanager doesn't crash on exception`() = TestUtil.test(accessManagedJavalin()) { app, _ ->
-        TestClient(app, "/?exception=true").connectAndDisconnect()
+        val client = TestClient(app, "/?exception=true")
+
+        doAndSleepWhile({ client.connect() }, { !client.isClosed })
+
         assertThat(app.logger().log.size).isEqualTo(1)
     }
 
@@ -376,11 +391,110 @@ class TestWebSocket {
         assertThat(app.logger().log).containsExactly("value", "cookieMapSize:3")
     }
 
+    @Test
+    fun `before handlers work`() = TestUtil.test { app, _ ->
+        app.wsBefore { ws ->
+            ws.onConnect { app.logger().log.add("before handler: onConnect") }
+            ws.onMessage { app.logger().log.add("before handler: onMessage") }
+            ws.onClose { app.logger().log.add("before handler: onClose") }
+        }
+
+        app.ws("/ws") { ws ->
+            ws.onConnect { app.logger().log.add("endpoint handler: onConnect") }
+            ws.onMessage { app.logger().log.add("endpoint handler: onMessage") }
+            ws.onClose { app.logger().log.add("endpoint handler: onClose") }
+        }
+
+        val client = TestClient(app, "/ws")
+
+        doAndSleepWhile({ client.connect() }, { !client.isOpen })
+        doAndSleep { client.send("test") }
+        doAndSleepWhile({ client.close() }, { client.isClosing })
+
+        assertThat(app.logger().log).containsExactly(
+                "before handler: onConnect", "endpoint handler: onConnect",
+                "before handler: onMessage", "endpoint handler: onMessage",
+                "before handler: onClose", "endpoint handler: onClose"
+        )
+    }
+
+    @Test
+    fun `after handlers work`() = TestUtil.test { app, _ ->
+        app.ws("/ws") { ws ->
+            ws.onConnect { app.logger().log.add("endpoint handler: onConnect") }
+            ws.onMessage { app.logger().log.add("endpoint handler: onMessage") }
+            ws.onClose { app.logger().log.add("endpoint handler: onClose") }
+        }
+
+        app.wsAfter { ws ->
+            ws.onConnect { app.logger().log.add("after handler: onConnect") }
+            ws.onMessage { app.logger().log.add("after handler: onMessage") }
+            ws.onClose { app.logger().log.add("after handler: onClose") }
+        }
+
+        val client = TestClient(app, "/ws")
+
+        doAndSleepWhile({ client.connect() }, { !client.isOpen })
+        doAndSleep { client.send("test") }
+        doAndSleepWhile({ client.close() }, { client.isClosing })
+
+        assertThat(app.logger().log).containsExactly(
+                "endpoint handler: onConnect", "after handler: onConnect",
+                "endpoint handler: onMessage", "after handler: onMessage",
+                "endpoint handler: onClose", "after handler: onClose"
+        )
+    }
+
+    @Test
+    fun `unmapped exceptions are caught by default handler`() = TestUtil.test { app, _ ->
+        val exception = Exception("Error message")
+
+        app.ws("/ws") { it.onConnect { throw exception } }
+
+        val client = object : TestClient(app, "/ws") {
+            override fun onClose(i: Int, s: String, b: Boolean) {
+                this.app.logger().log.add("Status code: $i")
+                this.app.logger().log.add("Reason: $s")
+            }
+        }
+
+        doAndSleepWhile({ client.connect() }, { !client.isClosed })
+
+        assertThat(client.app.logger().log).containsExactly(
+                "Status code: ${StatusCode.SERVER_ERROR}",
+                "Reason: ${exception.message}"
+        )
+    }
+
+    @Test
+    fun `mapped exceptions are handled`() = TestUtil.test { app, _ ->
+        app.ws("/ws") { it.onConnect { throw Exception() } }
+                .wsException(Exception::class.java) { _, _ ->
+                    app.logger().log.add("Exception handler called") }
+
+        TestClient(app, "/ws").connectAndDisconnect()
+
+        assertThat(app.logger().log).containsExactly("Exception handler called")
+    }
+
+    @Test
+    fun `most specific exception handler handles exception`() = TestUtil.test { app, _ ->
+        app.ws("/ws") { it.onConnect { throw TypedException() } }
+                .wsException(Exception::class.java) { _, _ ->
+                    app.logger().log.add("Exception handler called") }
+                .wsException(TypedException::class.java) { _, _ ->
+                    app.logger().log.add("TypedException handler called") }
+
+        TestClient(app, "/ws").connectAndDisconnect()
+
+        assertThat(app.logger().log).containsExactly("TypedException handler called")
+    }
+
     // ********************************************************************************************
     // Helpers
     // ********************************************************************************************
 
-    internal inner class TestClient(var app: Javalin, path: String, headers: Map<String, String> = emptyMap()) : WebSocketClient(URI.create("ws://localhost:" + app.port() + path), Draft_6455(), headers, 0) {
+    internal open inner class TestClient(var app: Javalin, path: String, headers: Map<String, String> = emptyMap()) : WebSocketClient(URI.create("ws://localhost:" + app.port() + path), Draft_6455(), headers, 0) {
 
         override fun onOpen(serverHandshake: ServerHandshake) {}
         override fun onClose(i: Int, s: String, b: Boolean) {}
@@ -390,15 +504,22 @@ class TestWebSocket {
         }
 
         fun connectAndDisconnect() {
-            doAndSleepWhile({ this.connect() }, { !this.isOpen })
-            doAndSleepWhile({ this.close() }, { this.isClosing })
+            doAndSleepWhile({ connect() }, { !isOpen })
+            doAndSleepWhile({ close() }, { !isClosed })
         }
     }
 
-    private fun doAndSleepWhile(slowFunction: () -> Unit, conditionFunction: () -> Boolean) {
+    private fun doAndSleepWhile(slowFunction: () -> Unit, conditionFunction: () -> Boolean, timeout: Duration = Duration.ofSeconds(5)) {
         val startTime = System.currentTimeMillis()
+        val limitTime = startTime + timeout.toMillis();
+
         slowFunction.invoke()
-        while (conditionFunction.invoke() && System.currentTimeMillis() < startTime + 250) {
+
+        while (conditionFunction.invoke()) {
+            if (System.currentTimeMillis() > limitTime) {
+                throw TimeoutException("Wait for condition has timed out")
+            }
+
             Thread.sleep(2)
         }
     }