Repository: tipsy/javalin (kotlin)

Issue:
writeJsonStream doesn't correctly compress response stream
**Actual behavior (the bug)**
Sending a request to a handler which writes it's response using `writeJsonStream` fails when requesting compressed output.
Specifically sending a request with `curl --compressed` which sets `Accept-Encoding: deflate, gzip` fails with `curl: (61) Error while processing content unencoding: incorrect header check`.
Javalin responds with `Content-Type: gzip` but I believe sends uncompressed bytes which causes cURL to error out with the above error message.

**Expected behavior**
The server should handle `curl --compressed` for streaming endpoints without additional configuration.

**To Reproduce**
Create a handler that uses `ctx.writeJsonStream` that consumes a relatively large stream.
Send a request to said handler with `curl --compressed`.

**Additional context**
Using a custom compression strategy and setting `minSizeForCompression = 0` works around the issue.

Pull request description:
A request scope attribute `ALWAYS_COMPRESS` will signal that a compression stream should always be created, if allowed. This attribute is now used for the `writeJsonStream()` call, where the total bytes to write cannot be known in advance.

The evaluation of creating a compressed stream or not was reworked to only evaluate only once.

closes #1950

Patch:
diff --git a/javalin/src/main/java/io/javalin/compression/CompressedStream.kt b/javalin/src/main/java/io/javalin/compression/CompressedStream.kt
index 74129c825..80e3b33ae 100644
--- a/javalin/src/main/java/io/javalin/compression/CompressedStream.kt
+++ b/javalin/src/main/java/io/javalin/compression/CompressedStream.kt
@@ -6,22 +6,41 @@ import jakarta.servlet.ServletOutputStream
 import jakarta.servlet.WriteListener
 import java.io.OutputStream
 
-internal class CompressedOutputStream(val compression: CompressionStrategy, val ctx: Context) : ServletOutputStream() {
+internal class CompressedOutputStream(
+    val minSizeForCompression: Int,
+    val compression: CompressionStrategy,
+    val ctx: Context,
+) : ServletOutputStream() {
 
     private val originStream = ctx.res().outputStream
     private var compressedStream: OutputStream? = null
+    private var isCompressionDecisionMade = false
 
-    override fun write(bytes: ByteArray, offset: Int, length: Int) {
-        if (compressedStream == null && length >= compression.minSizeForCompression && compression.allowsForCompression(ctx.res().contentType) && !ctx.res().containsHeader(Header.CONTENT_ENCODING)) {
-            tryMatchCompression(compression, ctx)?.also {
-                this.compressedStream = it.compress(originStream)
-                ctx.header(Header.CONTENT_ENCODING, it.encoding())
+    private fun maybeCreateCompressionStreamOnFirstWrite(length: Int) {
+        if (!isCompressionDecisionMade) {
+            val isCompressionAllowed = !ctx.res().containsHeader(Header.CONTENT_ENCODING) &&
+                compression.allowsForCompression(ctx.res().contentType)
+            val isCompressionDesired = length >= minSizeForCompression
+            if (isCompressionAllowed && isCompressionDesired) {
+                findMatchingCompressor(compression, ctx)?.also {
+                    this.compressedStream = it.compress(originStream)
+                    ctx.header(Header.CONTENT_ENCODING, it.encoding())
+                }
             }
+            isCompressionDecisionMade = true
         }
-        (compressedStream ?: originStream).write(bytes, offset, length) // fall back to default stream if no compression
     }
 
-    override fun write(byte: Int) = originStream.write(byte)
+    override fun write(bytes: ByteArray, offset: Int, length: Int) {
+        maybeCreateCompressionStreamOnFirstWrite(length)
+        (compressedStream ?: originStream).write(bytes, offset, length)
+    }
+
+    override fun write(byte: Int) {
+        maybeCreateCompressionStreamOnFirstWrite(1)
+        (compressedStream ?: originStream).write(byte)
+    }
+
     override fun setWriteListener(writeListener: WriteListener?) = originStream.setWriteListener(writeListener)
     override fun isReady(): Boolean = originStream.isReady
     override fun close() {
@@ -30,7 +49,7 @@ internal class CompressedOutputStream(val compression: CompressionStrategy, val
 
 }
 
-private fun tryMatchCompression(
+private fun findMatchingCompressor(
     compression: CompressionStrategy,
     ctx: Context,
 ): Compressor? =
@@ -40,5 +59,6 @@ private fun tryMatchCompression(
 
 
 private fun CompressionStrategy.allowsForCompression(contentType: String?): Boolean =
-    contentType == null || excludedMimeTypesFromCompression.none { excluded -> contentType.contains(excluded, ignoreCase = true) }
-
+    contentType == null || excludedMimeTypesFromCompression.none { excluded ->
+        contentType.contains(excluded, ignoreCase = true)
+    }
diff --git a/javalin/src/main/java/io/javalin/compression/CompressionStrategy.kt b/javalin/src/main/java/io/javalin/compression/CompressionStrategy.kt
index c85e9d7c9..58279ea9c 100644
--- a/javalin/src/main/java/io/javalin/compression/CompressionStrategy.kt
+++ b/javalin/src/main/java/io/javalin/compression/CompressionStrategy.kt
@@ -55,7 +55,7 @@ class CompressionStrategy(brotli: Brotli? = null, gzip: Gzip? = null) {
     }
 
     /** 1500 is the size of a packet, compressing responses smaller than this serves no purpose */
-    var minSizeForCompression = 1500
+    var defaultMinSizeForCompression = 1500
 
     /** Those mime types will be processed using NONE compression strategy */
     var excludedMimeTypesFromCompression = listOf(
diff --git a/javalin/src/main/java/io/javalin/http/Context.kt b/javalin/src/main/java/io/javalin/http/Context.kt
index dd03202bf..49a184cf8 100644
--- a/javalin/src/main/java/io/javalin/http/Context.kt
+++ b/javalin/src/main/java/io/javalin/http/Context.kt
@@ -293,13 +293,22 @@ interface Context {
     fun responseCharset(): Charset = runCatching { Charset.forName(res().characterEncoding) }.getOrElse { Charset.defaultCharset() }
 
     /**
-     * Gets output-stream you can write to.
-     * This stream by default uses compression specified in Javalin configuration,
-     * if you're looking for raw, uncompressed servlet's output-stream, use `ctx.res().outputStream`.
+     * Gets the output stream you can write to.
+     * This stream by default uses compression specified in the Javalin configuration.
+     * If you're looking for the servlet's raw, uncompressed output stream, use `ctx.res().outputStream`.
      * @see [HttpServletResponse.getOutputStream]
      */
     fun outputStream(): ServletOutputStream
 
+    /**
+     * The output stream returned by outputStream() will use compression (as specified in Javalin configuration), but
+     * compression will happen only if the first write to the output stream is larger than `minSizeForCompression`.
+     * Setting this value to zero will cause compression to always be used. This value must be assigned before calling
+     * outputStream() for the first time. The default value is set to the value of
+     * [io.javalin.config.CompressionStrategy.defaultMinSizeForCompression].
+     */
+    var minSizeForCompression: Int
+
     /**
      * Writes the specified inputStream as a seekable stream.
      * You can change this default in [io.javalin.config.JavalinConfig].
@@ -434,10 +443,10 @@ interface Context {
      * Consumes the specified stream with the configured JsonMapper, which transforms the stream's
      * content to JSON, writing the results directly to the response's `outputStream` as the stream
      * is consumed. This function call is synchronous, and may be wrapped in `ctx.async { }` if needed.
+     * The response will always be compressed regardless of size, given that compression is enabled in
+     * the Javalin configuration.
      */
-    fun writeJsonStream(stream: Stream<*>) = jsonMapper().writeToOutputStream(
-        stream, this.contentType(ContentType.APPLICATION_JSON).outputStream()
-    )
+    fun writeJsonStream(stream: Stream<*>)
 
     /** Sets context result to specified html string and sets content-type to text/html. */
     fun html(html: String): Context = contentType(ContentType.TEXT_HTML).result(html)
diff --git a/javalin/src/main/java/io/javalin/http/servlet/JavalinServletContext.kt b/javalin/src/main/java/io/javalin/http/servlet/JavalinServletContext.kt
index d6dbb49d2..9ddb6cb41 100644
--- a/javalin/src/main/java/io/javalin/http/servlet/JavalinServletContext.kt
+++ b/javalin/src/main/java/io/javalin/http/servlet/JavalinServletContext.kt
@@ -15,6 +15,7 @@ import io.javalin.http.HandlerType.AFTER
 import io.javalin.http.Header
 import io.javalin.http.HttpResponseException
 import io.javalin.http.HttpStatus
+import io.javalin.json.jsonMapper
 import io.javalin.routing.HandlerEntry
 import io.javalin.security.BasicAuthCredentials
 import io.javalin.util.JavalinLogger
@@ -29,6 +30,7 @@ import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.function.Supplier
+import java.util.stream.Stream
 
 data class JavalinServletContextConfig(
     val appAttributes: Map<String, Any>,
@@ -51,6 +53,7 @@ class JavalinServletContext(
     internal var endpointHandlerPath: String = "",
     internal var userFutureSupplier: Supplier<out CompletableFuture<*>>? = null,
     private var resultStream: InputStream? = null,
+    override var minSizeForCompression: Int = cfg.compressionStrategy.defaultMinSizeForCompression,
 ) : Context {
 
     init {
@@ -112,7 +115,9 @@ class JavalinServletContext(
     private val queryParams by lazy { super.queryParamMap() }
     override fun queryParamMap(): Map<String, List<String>> = queryParams
 
-    internal val outputStreamWrapper = lazy { CompressedOutputStream(cfg.compressionStrategy, this) }
+    internal val outputStreamWrapper = lazy {
+        CompressedOutputStream(minSizeForCompression, cfg.compressionStrategy, this)
+    }
     override fun outputStream(): ServletOutputStream = outputStreamWrapper.value
 
     override fun redirect(location: String, status: HttpStatus) {
@@ -136,6 +141,10 @@ class JavalinServletContext(
 
     override fun skipRemainingHandlers(): Context = also { tasks.clear() }
 
+    override fun writeJsonStream(stream: Stream<*>) {
+        minSizeForCompression = 0
+        jsonMapper().writeToOutputStream(stream, this.contentType(ContentType.APPLICATION_JSON).outputStream())
+    }
 }
 
 // this header is semicolon separated, like: "text/html; charset=UTF-8"


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestCompression.kt b/javalin/src/test/java/io/javalin/TestCompression.kt
index e3a299b6a..e82b7c944 100644
--- a/javalin/src/test/java/io/javalin/TestCompression.kt
+++ b/javalin/src/test/java/io/javalin/TestCompression.kt
@@ -10,6 +10,8 @@ package io.javalin
 import io.javalin.compression.Brotli
 import io.javalin.compression.CompressionStrategy
 import io.javalin.compression.Gzip
+import io.javalin.http.ContentType
+import io.javalin.http.Handler
 import io.javalin.http.Header
 import io.javalin.http.staticfiles.Location
 import io.javalin.testing.TestDependency
@@ -23,6 +25,7 @@ import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.condition.EnabledIf
 import java.util.zip.GZIPInputStream
+import kotlin.streams.asStream
 import com.aayushatharva.brotli4j.decoder.BrotliInputStream as Brotli4jInputStream
 
 class TestCompression {
@@ -36,22 +39,22 @@ class TestCompression {
     private val testDocument = FileUtil.readResource("/public/html.html")
 
     private fun customCompressionApp(limit: Int): Javalin = Javalin.create {
-        it.pvt.compressionStrategy.minSizeForCompression = limit
+        it.pvt.compressionStrategy.defaultMinSizeForCompression = limit
         it.staticFiles.add("/public", Location.CLASSPATH)
     }.addTestEndpoints()
 
     private fun superCompressingApp() = Javalin.create {
-        it.compression.custom(CompressionStrategy(Brotli(), Gzip()).apply { minSizeForCompression = 1 })
+        it.compression.custom(CompressionStrategy(Brotli(), Gzip()).apply { defaultMinSizeForCompression = 1 })
         it.staticFiles.add("/public", Location.CLASSPATH)
     }.addTestEndpoints()
 
     private fun brotliDisabledApp() = Javalin.create {
-        it.compression.custom(CompressionStrategy(null, Gzip()).apply { minSizeForCompression = testDocument.length })
+        it.compression.custom(CompressionStrategy(null, Gzip()).apply { defaultMinSizeForCompression = testDocument.length })
         it.staticFiles.add("/public", Location.CLASSPATH)
     }.addTestEndpoints()
 
     private fun etagApp() = Javalin.create {
-        it.pvt.compressionStrategy.minSizeForCompression = testDocument.length
+        it.pvt.compressionStrategy.defaultMinSizeForCompression = testDocument.length
         it.staticFiles.add("/public", Location.CLASSPATH)
         it.http.generateEtags = true
     }.addTestEndpoints()
@@ -283,13 +286,114 @@ class TestCompression {
                 staticFiles.directory = "/public"
                 staticFiles.location = Location.CLASSPATH
             }
-            it.pvt.compressionStrategy.minSizeForCompression = 0 // minSize to enable automatic compress
+            it.pvt.compressionStrategy.defaultMinSizeForCompression = 0 // minSize to enable automatic compress
         }
         TestUtil.test(gzipWebjars) { _, http ->
             assertValidGzipResponse(http.origin, path)
         }
     }
 
+    private fun buildSampleJson(chunksOf20Chars: Int): String {
+        val obj = """{"value":123456789}""" // 19 chars, will be 20 with comma
+        return List(chunksOf20Chars) { obj }.joinToString(",", "[", "]")
+    }
+    private val sampleJson100 = buildSampleJson(5)
+    private val sampleJson10k = buildSampleJson(500)
+    private fun testValidCompressionHandler(handler: Handler) {
+        val gzipTestApp = Javalin.create {
+            it.compression.gzipOnly()
+        }.apply {
+            get("/gzip-test", handler)
+        }
+        TestUtil.test(gzipTestApp) { _, http ->
+            assertValidGzipResponse(http.origin, "/gzip-test")
+        }
+        val brotliTestApp = Javalin.create {
+            it.compression.brotliOnly()
+        }.apply {
+            get("/brotli-test", handler)
+        }
+        TestUtil.test(brotliTestApp) { _, http ->
+            assertValidBrotliResponse(http.origin, "/brotli-test")
+        }
+    }
+    private fun testValidUncompressedHandler(handler: Handler) {
+        val uncompressedTestApp = Javalin.create {
+            it.compression.gzipOnly() // compression is enabled so that we can test minSizeForCompression thresholds
+        }.apply {
+            get("/uncompressed-test", handler)
+        }
+        TestUtil.test(uncompressedTestApp) { _, http ->
+            assertUncompressedResponse(http.origin, "/uncompressed-test")
+        }
+    }
+
+    @Test
+    fun `compresses a large string of JSON`() {
+        testValidCompressionHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON).result(sampleJson10k)
+        }
+        testValidUncompressedHandler { ctx ->
+            ctx.minSizeForCompression = sampleJson10k.length + 1
+            ctx.contentType(ContentType.APPLICATION_JSON).result(sampleJson10k)
+        }
+    }
+
+    @Test
+    fun `compresses a large string of JSON with direct single byte writes to outputStream`() {
+        testValidCompressionHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            ctx.minSizeForCompression = 0 // must force compression to use single byte writes
+            val out = ctx.outputStream()
+            sampleJson10k.forEach { out.write(it.code) }
+        }
+        testValidUncompressedHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            val out = ctx.outputStream()
+            sampleJson10k.forEach { out.write(it.code) } // first write is one byte, so no compression
+        }
+    }
+
+    @Test
+    fun `compresses a large string of JSON with direct byte array writes to outputStream`() {
+        testValidCompressionHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            val out = ctx.outputStream()
+            sampleJson10k.map { it.code.toByte() }.toByteArray().let { bytes -> out.write(bytes) }
+        }
+        testValidUncompressedHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            ctx.minSizeForCompression = sampleJson10k.length + 1
+            val out = ctx.outputStream()
+            sampleJson10k.map { it.code.toByte() }.toByteArray().let { bytes -> out.write(bytes) }
+        }
+    }
+
+    @Test
+    fun `compresses a small string of JSON with direct byte array writes to outputStream`() {
+        testValidCompressionHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            ctx.minSizeForCompression = 0 // must force compression since small string is below default threshold
+            val out = ctx.outputStream()
+            sampleJson100.map { it.code.toByte() }.toByteArray().let { bytes -> out.write(bytes) }
+        }
+        testValidUncompressedHandler { ctx ->
+            ctx.contentType(ContentType.APPLICATION_JSON)
+            val out = ctx.outputStream()
+            sampleJson100.map { it.code.toByte() }.toByteArray().let { bytes -> out.write(bytes) }
+        }
+    }
+
+    @Test
+    fun `compresses a large Stream of JSON`() {
+        data class Foo(val value: Long) // will become 19 chars in JSON, 20 with comma separator
+        fun createLargeJsonStream() = generateSequence { Foo(123456789) }.take(500).asStream() // > 10,000 chars
+        testValidCompressionHandler { ctx ->
+            ctx.writeJsonStream(createLargeJsonStream())
+        }
+        // no test for uncompressed since writing a Stream<T> forces compression
+    }
+
     private fun assertUncompressedResponse(origin: String, url: String) {
         val response = getResponse(origin, url, "br, gzip")
         assertThat(response.code).isLessThan(400)