Repository: tipsy/javalin
Pull request #1353 (2021-08-18 20:13:12)

[context] Add custom cookie class (samesite!)


Patch:
diff --git a/javalin/src/main/java/io/javalin/core/util/Headers.kt b/javalin/src/main/java/io/javalin/core/util/Headers.kt
index f0146650..00965ba1 100644
--- a/javalin/src/main/java/io/javalin/core/util/Headers.kt
+++ b/javalin/src/main/java/io/javalin/core/util/Headers.kt
@@ -100,7 +100,7 @@ class Headers {
         headers[Header.CROSS_ORIGIN_OPENER_POLICY] = policy.name.toHttpHeaderValue()
     }
 
-    // Cross-Origin-Resource-Policy: same-site	| same-origin | cross-origin
+    // Cross-Origin-Resource-Policy: same-site | same-origin | cross-origin
     enum class CrossOriginResourcePolicy {
         SAME_SITE, SAME_ORIGIN, CROSS_ORIGIN;
     }
diff --git a/javalin/src/main/java/io/javalin/http/Context.kt b/javalin/src/main/java/io/javalin/http/Context.kt
index c9278cb7..d98195eb 100644
--- a/javalin/src/main/java/io/javalin/http/Context.kt
+++ b/javalin/src/main/java/io/javalin/http/Context.kt
@@ -22,7 +22,6 @@ import java.nio.charset.Charset
 import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.function.Consumer
-import javax.servlet.http.Cookie
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
 
@@ -400,6 +399,10 @@ open class Context(@JvmField val req: HttpServletRequest, @JvmField val res: Htt
         }
     }
 
+    /** Sets the response status. */
+    fun status(httpCode: HttpCode): Context =
+        status(httpCode.status)
+
     /** Sets the response status. */
     fun status(statusCode: Int): Context {
         res.status = statusCode
@@ -411,24 +414,18 @@ open class Context(@JvmField val req: HttpServletRequest, @JvmField val res: Htt
 
     /** Sets a cookie with name, value, and (overloaded) max-age. */
     @JvmOverloads
-    fun cookie(name: String, value: String, maxAge: Int = -1): Context = cookie(Cookie(name, value).apply { setMaxAge(maxAge) })
+    fun cookie(name: String, value: String, maxAge: Int = -1) = cookie(Cookie(name = name, value = value, maxAge= maxAge))
 
     /** Sets a Cookie. */
     fun cookie(cookie: Cookie): Context {
-        cookie.path = cookie.path ?: "/"
-        res.addCookie(cookie) // we rely on this method for formatting the header
-        (res.getHeaders(Header.SET_COOKIE) ?: listOf()).toMutableList().let { cookies -> // mutable list of all cookies
-            cookies.removeIf { it.startsWith("${cookie.name}=") && !it.contains(cookie.value) } // remove old cookie if duplicate name
-            cookies.removeFirst()?.let { res.setHeader(Header.SET_COOKIE, it) } // remove first cookie and use it to clear the header
-            cookies.forEach { res.addHeader(Header.SET_COOKIE, it) } // add all remaining cookies
-        }
+        res.setJavalinCookie(cookie)
         return this
     }
 
     /** Removes cookie specified by name and path (optional). */
     @JvmOverloads
     fun removeCookie(name: String, path: String? = "/"): Context {
-        res.addCookie(Cookie(name, "").apply {
+        res.addCookie(javax.servlet.http.Cookie(name, "").apply {
             this.path = path
             this.maxAge = 0
         })
diff --git a/javalin/src/main/java/io/javalin/http/Cookie.kt b/javalin/src/main/java/io/javalin/http/Cookie.kt
new file mode 100644
index 00000000..d02f559e
--- /dev/null
+++ b/javalin/src/main/java/io/javalin/http/Cookie.kt
@@ -0,0 +1,56 @@
+package io.javalin.http
+
+import io.javalin.core.util.Header
+import javax.servlet.http.HttpServletResponse
+import javax.servlet.http.Cookie as ServletCookie
+
+const val SAME_SITE = "SameSite"
+
+enum class SameSite(val value: String) {
+    NONE("$SAME_SITE=None"),
+    STRICT("$SAME_SITE=Strict"),
+    LAX("$SAME_SITE=Lax");
+
+    override fun toString() = this.value
+}
+
+data class Cookie(
+    var name: String,
+    var value: String,
+    var path: String = "/",
+    var maxAge: Int = -1,
+    var secure: Boolean = false,
+    var version: Int = 0,
+    var isHttpOnly: Boolean = false,
+    var comment: String? = null,
+    var domain: String? = null,
+    var sameSite: SameSite? = null
+)
+
+fun HttpServletResponse.setJavalinCookie(javalinCookie: Cookie) {
+    val cookie = ServletCookie(javalinCookie.name, javalinCookie.value).apply {
+        this.path = javalinCookie.path
+        this.maxAge = javalinCookie.maxAge
+        this.secure = javalinCookie.secure
+        this.version = javalinCookie.version
+        this.isHttpOnly = javalinCookie.isHttpOnly
+        this.comment = javalinCookie.comment
+        if (javalinCookie.domain != null) {
+            // Null check required as Servlet API does domain.toLowerCase()
+            // with the amusing comment "IE allegedly needs this"
+            this.domain = javalinCookie.domain
+        }
+    }
+    this.addCookie(cookie) // we rely on this method for formatting the cookie header
+    (this.getHeaders(Header.SET_COOKIE) ?: listOf()).toMutableList().let { cookies -> // mutable list of all cookies
+        cookies.removeIf { it.startsWith("${cookie.name}=") && !it.contains(cookie.value) } // remove old cookie if duplicate name
+        cookies.removeFirst()?.let { first -> this.setHeader(Header.SET_COOKIE, first.addSameSite(javalinCookie)) } // remove first cookie and use it to clear the header
+        cookies.forEach { remaining -> this.addHeader(Header.SET_COOKIE, remaining.addSameSite(javalinCookie)) } // add all remaining cookies
+    }
+
+}
+
+fun String.addSameSite(cookie: Cookie): String {
+    if (cookie.sameSite == null || this.contains(SAME_SITE)) return this
+    return "$this; ${cookie.sameSite}"
+}
diff --git a/javalin/src/main/java/io/javalin/http/sse/Emitter.kt b/javalin/src/main/java/io/javalin/http/sse/Emitter.kt
index 1df31e7e..e797401f 100644
--- a/javalin/src/main/java/io/javalin/http/sse/Emitter.kt
+++ b/javalin/src/main/java/io/javalin/http/sse/Emitter.kt
@@ -1,6 +1,7 @@
 package io.javalin.http.sse
 
 import java.io.IOException
+import java.io.InputStream
 import javax.servlet.AsyncContext
 import javax.servlet.ServletOutputStream
 
@@ -18,19 +19,15 @@ class Emitter(private var asyncContext: AsyncContext) {
         }
     }
 
-    @JvmOverloads
-    fun emit(event: String, data: String, id: String? = null) = synchronized(this) {
+    fun emit(event: String, data: InputStream, id: String?) = synchronized(this) {
         try {
-            val sb = StringBuilder()
             if (id != null) {
-                sb.append("id: $id$newline")
+                output.print("id: $id$newline")
             }
-            sb.append("event: $event$newline")
-            data.lines().forEach { line ->
-                sb.append("data: $line$newline")
-            }
-            sb.append("$newline$newline")
-            output.print(sb.toString())
+            output.print("event: $event$newline")
+            output.print("data: " )
+            data.copyTo(output)
+            output.print(newline)
             asyncContext.response.flushBuffer()
         } catch (e: IOException) {
             closed = true
diff --git a/javalin/src/main/java/io/javalin/http/sse/SseClient.kt b/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
index dc1f34e2..5f4468d5 100644
--- a/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
+++ b/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
@@ -1,6 +1,8 @@
 package io.javalin.http.sse
 
 import io.javalin.http.Context
+import io.javalin.plugin.json.jsonMapper
+import java.io.InputStream
 
 class SseClient(@JvmField val ctx: Context) {
 
@@ -11,11 +13,15 @@ class SseClient(@JvmField val ctx: Context) {
         this.closeCallback = closeCallback
     }
 
-    fun sendEvent(data: String) = sendEvent("message", data)
+    fun sendEvent(data: Any) = sendEvent("message", data)
 
     @JvmOverloads
-    fun sendEvent(event: String, data: String, id: String? = null) {
-        emitter.emit(event, data, id)
+    fun sendEvent(event: String, data: Any, id: String? = null) {
+        when (data) {
+            is InputStream -> emitter.emit(event, data, id)
+            is String -> emitter.emit(event, data.byteInputStream(), id)
+            else -> emitter.emit(event, ctx.jsonMapper().toJsonString(data).byteInputStream(), id)
+        }
         if (emitter.isClosed()) { // can't detect if closed before we try emitting?
             closeCallback.run()
         }
diff --git a/javalin/src/main/java/io/javalin/http/util/CookieStore.kt b/javalin/src/main/java/io/javalin/http/util/CookieStore.kt
index 64016b9d..44428c3c 100644
--- a/javalin/src/main/java/io/javalin/http/util/CookieStore.kt
+++ b/javalin/src/main/java/io/javalin/http/util/CookieStore.kt
@@ -6,9 +6,9 @@
 
 package io.javalin.http.util
 
+import io.javalin.http.Cookie
 import io.javalin.plugin.json.JsonMapper
 import java.util.*
-import javax.servlet.http.Cookie
 
 @Suppress("UNCHECKED_CAST")
 class CookieStore(val jsonMapper: JsonMapper, cookie: String?) {
@@ -19,7 +19,7 @@ class CookieStore(val jsonMapper: JsonMapper, cookie: String?) {
 
     private val cookieMap = deserialize(cookie)
 
-    fun serializeToCookie() = Cookie(COOKIE_NAME, serialize(cookieMap)).apply { path = "/" }
+    fun serializeToCookie() = Cookie(COOKIE_NAME, serialize(cookieMap))
 
     operator fun <T> get(key: String) = cookieMap[key] as T
 


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestCookie.kt b/javalin/src/test/java/io/javalin/TestCookie.kt
index 47804c08..ea186590 100644
--- a/javalin/src/test/java/io/javalin/TestCookie.kt
+++ b/javalin/src/test/java/io/javalin/TestCookie.kt
@@ -2,16 +2,17 @@ package io.javalin
 
 import com.mashape.unirest.http.Unirest
 import io.javalin.core.util.Header
+import io.javalin.http.Cookie
+import io.javalin.http.SameSite
 import io.javalin.testing.TestUtil
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Test
-import javax.servlet.http.Cookie
 
 class TestCookie {
 
     @Test
     fun `cookie set on unspecified path is set to root URL`() = TestUtil.test { app, http ->
-        app.get("/cookie") { it.cookie(Cookie("key", "value").apply { path = "/" }) }
+        app.get("/cookie") { it.cookie(Cookie("key", "value")) }
 
         val setCookieResponse = http.get("/cookie")
         val cookiePath = setCookieResponse.headers.getFirst(Header.SET_COOKIE).split(";")[1].replaceFirst(" ", "")
@@ -21,9 +22,7 @@ class TestCookie {
 
     @Test
     fun `removing cookie without specifying path will remove cookie set to root URL`() = TestUtil.test { app, http ->
-        val cookie = Cookie("key", "value").apply {
-            path = "/"
-        }
+        val cookie = Cookie("key", "value")
         app.get("/cookie") { it.cookie(cookie) }
         app.get("/cookie-remove") { it.removeCookie(cookie.name) }
         val setCookieResponse = http.get("/cookie")
@@ -36,9 +35,7 @@ class TestCookie {
 
     @Test
     fun `removing cookie without specifying path will not remove cookie set to non-root URL`() = TestUtil.test { app, http ->
-        val cookie = Cookie("key", "value").apply {
-            path = "/some-path"
-        }
+        val cookie = Cookie("key", "value", path = "/some-path")
         app.get("/cookie") { it.cookie(cookie) }
         app.get("/cookie-remove") { it.removeCookie(cookie.name) }
         val setCookieResponse = http.get("/cookie")
@@ -91,7 +88,7 @@ class TestCookie {
 
     @Test
     fun `setting a Cookie object works`() = TestUtil.test { app, http ->
-        app.get("/create-cookie") { ctx -> ctx.cookie(Cookie("Hest", "Hast").apply { maxAge = 7 }) }
+        app.get("/create-cookie") { ctx -> ctx.cookie(Cookie("Hest", "Hast", maxAge = 7)) }
         assertThat(http.get("/create-cookie").headers.getFirst(Header.SET_COOKIE)).contains("Hest=Hast")
         assertThat(http.get("/create-cookie").headers.getFirst(Header.SET_COOKIE)).contains("Max-Age=7")
     }
@@ -122,6 +119,20 @@ class TestCookie {
         assertThat(response.headers[Header.SET_COOKIE]!!.size).isEqualTo(1)
     }
 
+    @Test
+    fun `can set samesite easily`() = TestUtil.test { app, http ->
+        app.get("/create-cookie") { it.cookie(Cookie("Test", "Tast", sameSite = SameSite.STRICT)) }
+        val cookie = http.get("/create-cookie").headers.getFirst(Header.SET_COOKIE)
+        assertThat(cookie).isEqualTo("Test=Tast; Path=/; SameSite=Strict")
+    }
+
+    @Test
+    fun `can set samesite and other properties`() = TestUtil.test { app, http ->
+        app.get("/create-cookie") { it.cookie(Cookie("Test", "Tast", sameSite = SameSite.NONE, isHttpOnly = true, domain = "localhost")) }
+        val cookie = http.get("/create-cookie").headers.getFirst(Header.SET_COOKIE)
+        assertThat(cookie).isEqualTo("Test=Tast; Path=/; Domain=localhost; HttpOnly; SameSite=None")
+    }
+
     private fun cookieIsEffectivelyRemoved(cookie: String, path: String): Boolean {
         val pathMatches = cookie.split(";")[1].split("=")[1] == path
         val expiresEpoch = cookie.split(";")[2] == " Expires=Thu, 01-Jan-1970 00:00:00 GMT"
diff --git a/javalin/src/test/java/io/javalin/TestSse.kt b/javalin/src/test/java/io/javalin/TestSse.kt
index d06297e7..569ac176 100644
--- a/javalin/src/test/java/io/javalin/TestSse.kt
+++ b/javalin/src/test/java/io/javalin/TestSse.kt
@@ -1,10 +1,10 @@
 package io.javalin
 
 import io.javalin.http.sse.SseClient
+import io.javalin.testing.SerializableObject
 import io.javalin.testing.TestUtil
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Test
-import java.util.*
 
 class TestSse {
 
@@ -20,6 +20,18 @@ class TestSse {
         assertThat(body).contains("data: $data")
     }
 
+    @Test
+    fun `sending input stream works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
+        app.sse("/sse") { it.sendEvent(event, "MY DATA".byteInputStream()) }
+        assertThat(http.sse("/sse").get().body).contains("data: MY DATA")
+    }
+
+    @Test
+    fun `sending json works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
+        app.sse("/sse") { it.sendEvent(event, SerializableObject()) }
+        assertThat(http.sse("/sse").get().body).contains("""data: {"value1":"FirstValue","value2":"SecondValue"}""")
+    }
+
     @Test
     fun `sending events with ID works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
         app.sse("/sse") { it.sendEvent(event, data, id = "SOME_ID") }
@@ -68,5 +80,5 @@ class TestSse {
         assertThat(body).contains("event: $event")
         assertThat(body).contains("data: " + "my-qp")
     }
-    
+
 }