Repository: tipsy/javalin
Pull request #1463 (2022-01-12 15:52:55)

WsContext: Added closeSession() methods.


Patch:
diff --git a/javalin/src/main/java/io/javalin/http/JavalinServlet.kt b/javalin/src/main/java/io/javalin/http/JavalinServlet.kt
index 03f2d3ab..61600c0a 100644
--- a/javalin/src/main/java/io/javalin/http/JavalinServlet.kt
+++ b/javalin/src/main/java/io/javalin/http/JavalinServlet.kt
@@ -97,7 +97,9 @@ class JavalinServlet(val config: JavalinConfig) : HttpServlet() {
             resultFuture.exceptionally { throwable ->
                 exceptionMapper.handleFutureException(ctx, throwable)
             }.thenAccept { futureValue ->
-                ctx.futureConsumer?.accept(futureValue) // this consumer can set result, status, etc
+                tryWithExceptionMapper {
+                    ctx.futureConsumer?.accept(futureValue) // this consumer can set result, status, etc
+                }
                 finishUpResponse(asyncContext.response).also { asyncContext.complete() }
             }.exceptionally { throwable -> // exception might occur when writing response
                 exceptionMapper.handleUnexpectedThrowable(res, throwable).also { asyncContext.complete() }
diff --git a/javalin/src/main/java/io/javalin/websocket/WsContext.kt b/javalin/src/main/java/io/javalin/websocket/WsContext.kt
index 61197cb1..99bd8974 100644
--- a/javalin/src/main/java/io/javalin/websocket/WsContext.kt
+++ b/javalin/src/main/java/io/javalin/websocket/WsContext.kt
@@ -10,6 +10,7 @@ import io.javalin.http.Context
 import io.javalin.jetty.upgradeContextKey
 import io.javalin.jetty.upgradeSessionAttrsKey
 import io.javalin.plugin.json.jsonMapper
+import org.eclipse.jetty.websocket.api.CloseStatus
 import org.eclipse.jetty.websocket.api.RemoteEndpoint
 import org.eclipse.jetty.websocket.api.Session
 import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest
@@ -59,6 +60,10 @@ abstract class WsContext(val sessionId: String, @JvmField val session: Session)
     fun <T> sessionAttribute(key: String): T? = sessionAttributeMap()[key] as T
     fun sessionAttributeMap(): Map<String, Any?> = sessionAttributes ?: mapOf()
 
+    fun closeSession() = session.close()
+    fun closeSession(closeStatus: CloseStatus) = session.close(closeStatus)
+    fun closeSession(code: Int, reason: String?) = session.close(code, reason)
+
     override fun equals(other: Any?) = session == (other as WsContext).session
     override fun hashCode() = session.hashCode()
 }


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestFuture.kt b/javalin/src/test/java/io/javalin/TestFuture.kt
index f512fce2..a418a6f4 100644
--- a/javalin/src/test/java/io/javalin/TestFuture.kt
+++ b/javalin/src/test/java/io/javalin/TestFuture.kt
@@ -2,6 +2,7 @@ package io.javalin
 
 import io.javalin.core.util.Header
 import io.javalin.http.ContentType
+import io.javalin.http.NotFoundResponse
 import io.javalin.testing.TestUtil
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
@@ -111,6 +112,14 @@ class TestFuture {
         assertThat(noContentResponse.headers.getFirst(Header.CONTENT_TYPE)).isEqualTo(ContentType.PLAIN)
     }
 
+    @Test
+    fun `exceptions in future callback are mapped`() = TestUtil.test { app, http ->
+        app.get("/") { ctx ->
+            ctx.future(getFuture("result")) { throw NotFoundResponse() }
+        }
+        assertThat(http.get("/").status).isEqualTo(404)
+    }
+
     private val impatientServer: Javalin by lazy { Javalin.create { it.asyncRequestTimeout = 5 } }
 
     @Test
@@ -169,4 +178,3 @@ class TestFuture {
     }
 
 }
-
diff --git a/javalin/src/test/java/io/javalin/TestWebSocket.kt b/javalin/src/test/java/io/javalin/TestWebSocket.kt
index 573bc635..4b4c7753 100644
--- a/javalin/src/test/java/io/javalin/TestWebSocket.kt
+++ b/javalin/src/test/java/io/javalin/TestWebSocket.kt
@@ -16,6 +16,7 @@ import io.javalin.testing.TestUtil
 import io.javalin.testing.TypedException
 import io.javalin.websocket.WsContext
 import org.assertj.core.api.Assertions.assertThat
+import org.eclipse.jetty.websocket.api.CloseStatus
 import org.eclipse.jetty.websocket.api.MessageTooLargeException
 import org.eclipse.jetty.websocket.api.StatusCode
 import org.eclipse.jetty.websocket.api.WebSocketConstants
@@ -529,6 +530,54 @@ class TestWebSocket {
         assertThat(response.headers.getFirst(WebSocketConstants.SEC_WEBSOCKET_PROTOCOL)).isEqualTo("mqtt")
     }
 
+    @Test
+    fun `websocket closeSession() methods`() = TestUtil.test { app, http ->
+        app.ws("/websocket") {
+            it.onConnect {
+                app.logger().log.add("Connected")
+            }
+            it.onMessage {
+                val message = it.message()
+
+                if (message.equals("closeSession()")) {
+                    it.closeSession();
+                } else if (message.equals("closeSession(CloseStatus)")) {
+                    it.closeSession(CloseStatus(3456, "Test"))
+                } else if (message.equals("closeSession(code: Int, reason: String?)")) {
+                    it.closeSession(3456, "Test")
+                }
+
+                it.closeSession(4567, "Unexpected Message")
+            }
+            it.onClose {
+                app.logger().log.add("${it.status()}:${it.reason()}")
+            }
+        }
+
+        val testClient = TestClient(app, "/websocket")
+
+        //region closeSession()
+        doAndSleepWhile({ testClient.connect() }, { "Connected" !in app.logger().log })
+        doAndSleepWhile({ testClient.send("closeSession()") }, { "Closing" !in app.logger().log })
+        doAndSleepWhile({  }, { "3456:Test" !in app.logger().log })
+        assertThat(app.logger().log).contains("3456:Test")
+        //endregion
+
+        //region closeSession(CloseStatus)
+        doAndSleepWhile({ testClient.connect() }, { "Connected" !in app.logger().log })
+        doAndSleepWhile({ testClient.send("closeSession(CloseStatus)") }, { "Closing" !in app.logger().log })
+        doAndSleepWhile({  }, { "3456:Test" !in app.logger().log })
+        assertThat(app.logger().log).contains("3456:Test")
+        //endregion
+
+        //region closeSession(code: Int, reason: String?)
+        doAndSleepWhile({ testClient.connect() }, { "Connected" !in app.logger().log })
+        doAndSleepWhile({ testClient.send("closeSession(code: Int, reason: String?)") }, { "Closing" !in app.logger().log })
+        doAndSleepWhile({  }, { "3456:Test" !in app.logger().log })
+        assertThat(app.logger().log).contains("3456:Test")
+        //endregion
+    }
+
     // ********************************************************************************************
     // Helpers
     // ********************************************************************************************