Repository: tipsy/javalin
Pull request #1507 (2022-04-06 21:03:48)

[servlet] Experimental refactor of JavalinResponseWrapper


Patch:
diff --git a/javalin/src/main/java/io/javalin/core/compression/CompressedStream.kt b/javalin/src/main/java/io/javalin/core/compression/CompressedStream.kt
new file mode 100644
index 00000000..a0504ae7
--- /dev/null
+++ b/javalin/src/main/java/io/javalin/core/compression/CompressedStream.kt
@@ -0,0 +1,33 @@
+package io.javalin.core.compression
+
+import com.nixxcode.jvmbrotli.enc.BrotliOutputStream
+import com.nixxcode.jvmbrotli.enc.Encoder
+import io.javalin.core.util.Header
+import io.javalin.http.Context
+import io.javalin.jetty.JettyPrecompressingResourceHandler.CompressType
+import io.javalin.jetty.JettyPrecompressingResourceHandler.CompressType.BR
+import io.javalin.jetty.JettyPrecompressingResourceHandler.CompressType.GZIP
+import java.io.OutputStream
+import java.util.zip.GZIPOutputStream
+
+data class CompressedStream(val type: CompressType, val outputStream: OutputStream) {
+    companion object {
+        fun tryBrotli(compression: CompressionStrategy, ctx: Context): CompressedStream? =
+            if (compression.brotli != null && ctx.header(Header.ACCEPT_ENCODING)?.contains(BR.typeName, ignoreCase = true) == true)
+                CompressedStream(BR, LeveledBrotliStream(ctx.res.outputStream, compression.brotli.level))
+            else null
+
+        fun tryGzip(compression: CompressionStrategy, ctx: Context): CompressedStream? =
+            if (compression.gzip != null && ctx.header(Header.ACCEPT_ENCODING)?.contains(GZIP.typeName, ignoreCase = true) == true)
+                CompressedStream(GZIP, LeveledGzipStream(ctx.res.outputStream, compression.gzip.level))
+            else null
+    }
+}
+
+class LeveledGzipStream(out: OutputStream, level: Int) : GZIPOutputStream(out) {
+    init {
+        this.def.setLevel(level)
+    }
+}
+
+class LeveledBrotliStream(out: OutputStream, level: Int) : BrotliOutputStream(out, Encoder.Parameters().setQuality(level))
diff --git a/javalin/src/main/java/io/javalin/core/compression/CompressionStrategy.kt b/javalin/src/main/java/io/javalin/core/compression/CompressionStrategy.kt
index 4281ee9f..d43b59ce 100644
--- a/javalin/src/main/java/io/javalin/core/compression/CompressionStrategy.kt
+++ b/javalin/src/main/java/io/javalin/core/compression/CompressionStrategy.kt
@@ -4,6 +4,7 @@ import com.nixxcode.jvmbrotli.common.BrotliLoader
 import io.javalin.core.util.JavalinLogger
 import io.javalin.core.util.OptionalDependency
 import io.javalin.core.util.Util
+import java.util.Arrays
 
 /**
  * This class is a settings container for Javalin's content compression.
@@ -17,9 +18,6 @@ import io.javalin.core.util.Util
  */
 class CompressionStrategy(brotli: Brotli? = null, gzip: Gzip? = null) {
 
-    val brotli: Brotli?
-    val gzip: Gzip?
-
     companion object {
         @JvmField
         val NONE = CompressionStrategy()
@@ -28,12 +26,31 @@ class CompressionStrategy(brotli: Brotli? = null, gzip: Gzip? = null) {
         val GZIP = CompressionStrategy(null, Gzip())
     }
 
+    val brotli: Brotli?
+    val gzip: Gzip?
+
     init {
         //Enabling brotli requires special handling since jvm-brotli is platform dependent
         this.brotli = if (brotli != null) tryLoadBrotli(brotli) else null
         this.gzip = gzip
     }
 
+    /** 1500 is the size of a packet, compressing responses smaller than this serves no purpose */
+    var minSizeForCompression = 1500
+    /** Those mime types will be processed using NONE compression strategy */
+    var excludedMimeTypesFromCompression = listOf(
+        "image/",
+        "audio/",
+        "video/",
+        "application/compress",
+        "application/zip",
+        "application/gzip",
+        "application/bzip2",
+        "application/brotli",
+        "application/x-xz",
+        "application/x-rar-compressed"
+    )
+
     /**
      * When enabling Brotli, we try loading the jvm-brotli native library first.
      * If this fails, we keep Brotli disabled and warn the user.
diff --git a/javalin/src/main/java/io/javalin/http/JavalinResponseWrapper.kt b/javalin/src/main/java/io/javalin/http/JavalinResponseWrapper.kt
index c712c30c..122caca1 100644
--- a/javalin/src/main/java/io/javalin/http/JavalinResponseWrapper.kt
+++ b/javalin/src/main/java/io/javalin/http/JavalinResponseWrapper.kt
@@ -1,17 +1,14 @@
 package io.javalin.http
 
-import com.nixxcode.jvmbrotli.enc.BrotliOutputStream
-import com.nixxcode.jvmbrotli.enc.Encoder
 import io.javalin.core.JavalinConfig
-import io.javalin.core.util.Header.ACCEPT_ENCODING
+import io.javalin.core.compression.CompressedStream
 import io.javalin.core.util.Header.CONTENT_ENCODING
 import io.javalin.core.util.Header.ETAG
 import io.javalin.core.util.Header.IF_NONE_MATCH
 import io.javalin.core.util.Util
 import io.javalin.http.HandlerType.GET
+import io.javalin.http.HttpCode.NOT_MODIFIED
 import java.io.InputStream
-import java.io.OutputStream
-import java.util.zip.GZIPOutputStream
 import javax.servlet.ServletOutputStream
 import javax.servlet.WriteListener
 import javax.servlet.http.HttpServletResponse
@@ -19,10 +16,10 @@ import javax.servlet.http.HttpServletResponseWrapper
 
 class JavalinResponseWrapper(private val ctx: Context, private val config: JavalinConfig, private val requestType: HandlerType) : HttpServletResponseWrapper(ctx.res) {
 
-    private val outputStreamWrapper by lazy { OutputStreamWrapper(ctx, config) }
+    private val outputStreamWrapper by lazy { OutputStreamWrapper(config, ctx) }
     override fun getOutputStream() = outputStreamWrapper
 
-    private val serverEtag by lazy { ctx.res.getHeader(ETAG) }
+    private val serverEtag by lazy { getHeader(ETAG) }
     private val clientEtag by lazy { ctx.req.getHeader(IF_NONE_MATCH) }
 
     fun write(resultStream: InputStream?) = when {
@@ -33,10 +30,9 @@ class JavalinResponseWrapper(private val ctx: Context, private val config: Javal
     }
 
     private fun generateEtagWriteAndClose(resultStream: InputStream) {
-        // TODO: https://github.com/tipsy/javalin/issues/1505
-        val inputStream = resultStream.readBytes().inputStream().also { resultStream.close() }
+        val inputStream = resultStream.use { it.readBytes().inputStream() } // TODO: https://github.com/tipsy/javalin/issues/1505
         val generatedEtag = Util.getChecksumAndReset(inputStream)
-        ctx.res.setHeader(ETAG, generatedEtag)
+        setHeader(ETAG, generatedEtag)
         when (generatedEtag) {
             clientEtag -> closeWith304(inputStream)
             else -> writeToWrapperAndClose(inputStream)
@@ -44,67 +40,38 @@ class JavalinResponseWrapper(private val ctx: Context, private val config: Javal
     }
 
     private fun writeToWrapperAndClose(inputStream: InputStream) {
-        inputStream.copyTo(outputStreamWrapper).also { inputStream.close() }
-        outputStreamWrapper.finalizeCompression()
+        inputStream.use { input ->
+            outputStreamWrapper.use { output ->
+                input.copyTo(output)
+            }
+        }
     }
 
-    private fun closeWith304(inputStream: InputStream) = inputStream.close().also { ctx.res.status = 304 }
+    private fun closeWith304(inputStream: InputStream) {
+        inputStream.use { ctx.status(NOT_MODIFIED) }
+    }
 
 }
 
-private const val BR = "br"
-private const val GZIP = "gzip"
-
-class OutputStreamWrapper(private val ctx: Context, config: JavalinConfig) : ServletOutputStream() {
-
-    private val acceptsBrotli by lazy { ctx.req.getHeader(ACCEPT_ENCODING)?.contains(BR, ignoreCase = true) == true }
-    private val acceptsGzip by lazy { ctx.req.getHeader(ACCEPT_ENCODING)?.contains(GZIP, ignoreCase = true) == true }
+class OutputStreamWrapper(val config: JavalinConfig, val ctx: Context, val response: HttpServletResponse = ctx.res) : ServletOutputStream() {
     private val compression = config.inner.compressionStrategy
-    private val response: HttpServletResponse by lazy { ctx.res }
-    private val isAllowedType by lazy { response.contentType == null || excludedMimeTypes.none { response.contentType.contains(it, ignoreCase = true) } }
-
-    private var compressedStream: OutputStream? = null
+    private var compressedStream: CompressedStream? = null
 
     override fun write(bytes: ByteArray, offset: Int, length: Int) {
-        if (isAllowedType && compressedStream == null && length >= minSizeForCompression) {
-            if (acceptsBrotli && compression.brotli != null) {
-                response.setHeader(CONTENT_ENCODING, BR)
-                compressedStream = LeveledBrotliStream(response.outputStream, compression.brotli.level)
-            } else if (acceptsGzip && compression.gzip != null) {
-                response.setHeader(CONTENT_ENCODING, GZIP)
-                compressedStream = LeveledGzipStream(response.outputStream, compression.gzip.level)
-            }
+        if (compressedStream == null && length >= compression.minSizeForCompression && response.contentType.allowsForCompression()) {
+            compressedStream = CompressedStream.tryBrotli(compression, ctx) ?: CompressedStream.tryGzip(compression, ctx)
+            compressedStream?.let { response.setHeader(CONTENT_ENCODING, it.type.typeName) }
         }
-        (compressedStream ?: response.outputStream).write(bytes, offset, length) // fall back to default stream if no compression
+        (compressedStream?.outputStream ?: response.outputStream).write(bytes, offset, length) // fall back to default stream if no compression
     }
 
-    fun finalizeCompression() = compressedStream?.close()
+    private fun String?.allowsForCompression(): Boolean =
+        this == null || compression.excludedMimeTypesFromCompression.none { excluded -> this.contains(excluded, ignoreCase = true) }
 
-    override fun isReady(): Boolean = response.outputStream.isReady
-    override fun setWriteListener(writeListener: WriteListener?) = response.outputStream.setWriteListener(writeListener)
     override fun write(byte: Int) = response.outputStream.write(byte)
-
-    companion object {
-        var minSizeForCompression = 1500 // 1500 is the size of a packet, compressing responses smaller than this serves no purpose
-        val excludedMimeTypes = setOf(
-            "image/",
-            "audio/",
-            "video/",
-            "application/compress",
-            "application/zip",
-            "application/gzip",
-            "application/bzip2",
-            "application/brotli",
-            "application/x-xz",
-            "application/x-rar-compressed"
-        )
-    }
-}
-
-class LeveledGzipStream(out: OutputStream, level: Int) : GZIPOutputStream(out) {
-    init {
-        this.def.setLevel(level)
+    override fun setWriteListener(writeListener: WriteListener?) = response.outputStream.setWriteListener(writeListener)
+    override fun isReady(): Boolean = response.outputStream.isReady
+    override fun close() {
+        compressedStream?.outputStream?.close()
     }
 }
-
-class LeveledBrotliStream(out: OutputStream, level: Int) : BrotliOutputStream(out, Encoder.Parameters().setQuality(level))
diff --git a/javalin/src/main/java/io/javalin/jetty/JettyPrecompressingResourceHandler.kt b/javalin/src/main/java/io/javalin/jetty/JettyPrecompressingResourceHandler.kt
index e6a7edc1..bb9c2947 100644
--- a/javalin/src/main/java/io/javalin/jetty/JettyPrecompressingResourceHandler.kt
+++ b/javalin/src/main/java/io/javalin/jetty/JettyPrecompressingResourceHandler.kt
@@ -1,11 +1,11 @@
 package io.javalin.jetty
 
+import io.javalin.core.compression.LeveledBrotliStream
+import io.javalin.core.compression.LeveledGzipStream
 import io.javalin.core.util.Header
 import io.javalin.core.util.JavalinLogger
 import io.javalin.core.util.OptionalDependency
 import io.javalin.core.util.Util
-import io.javalin.http.LeveledBrotliStream
-import io.javalin.http.LeveledGzipStream
 import org.eclipse.jetty.http.MimeTypes
 import org.eclipse.jetty.util.resource.Resource
 import java.io.ByteArrayOutputStream
diff --git a/javalin/src/main/java/io/javalin/jetty/JettyResourceHandler.kt b/javalin/src/main/java/io/javalin/jetty/JettyResourceHandler.kt
index b044cbf6..682684aa 100644
--- a/javalin/src/main/java/io/javalin/jetty/JettyResourceHandler.kt
+++ b/javalin/src/main/java/io/javalin/jetty/JettyResourceHandler.kt
@@ -8,7 +8,6 @@ package io.javalin.jetty
 
 import io.javalin.core.util.JavalinException
 import io.javalin.core.util.JavalinLogger
-import io.javalin.http.JavalinResponseWrapper
 import io.javalin.http.staticfiles.Location
 import io.javalin.http.staticfiles.StaticFileConfig
 import org.eclipse.jetty.server.Request
@@ -46,9 +45,8 @@ class JettyResourceHandler : JavalinResourceHandler {
                         return true
                     }
                     httpResponse.contentType = null // Jetty will only set the content-type if it's null
-                    handler.handle(target, baseRequest, httpRequest, httpResponse)
+                    httpResponse.outputStream.use { handler.handle(target, baseRequest, httpRequest, httpResponse) }
                     httpRequest.setAttribute("handled-as-static-file", true)
-                    (httpResponse as JavalinResponseWrapper).outputStream.finalizeCompression()
                     return true
                 }
             } catch (e: Exception) { // it's fine, we'll just 404


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestCompression.kt b/javalin/src/test/java/io/javalin/TestCompression.kt
index 58532b96..ee9fb349 100644
--- a/javalin/src/test/java/io/javalin/TestCompression.kt
+++ b/javalin/src/test/java/io/javalin/TestCompression.kt
@@ -9,11 +9,11 @@ package io.javalin
 import com.nixxcode.jvmbrotli.common.BrotliLoader
 import com.nixxcode.jvmbrotli.dec.BrotliInputStream
 import io.javalin.core.compression.Brotli
+import io.javalin.core.compression.CompressionStrategy
 import io.javalin.core.compression.Gzip
 import io.javalin.core.util.FileUtil
 import io.javalin.core.util.Header
 import io.javalin.core.util.OptionalDependency
-import io.javalin.http.OutputStreamWrapper
 import io.javalin.http.staticfiles.Location
 import io.javalin.testing.TestUtil
 import kong.unirest.Unirest
@@ -23,7 +23,6 @@ import okhttp3.Request
 import okhttp3.Response
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Assumptions.assumeTrue
-import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import java.util.zip.GZIPInputStream
 
@@ -32,30 +31,34 @@ class TestCompression {
     private data class SillyObject(var fieldOne: String, var fieldTwo: String, var fieldThree: String)
 
     private fun getSomeObjects(numberOfObjects: Int) = (1..numberOfObjects).map { i -> SillyObject("f$i", "f$i", "f$i") }.toList()
+    private val tinyLength = getSomeObjects(10).toString().length
+    private val hugeLength = getSomeObjects(1000).toString().length
 
     private val testDocument = FileUtil.readResource("/public/html.html")
 
-    @BeforeEach
-    fun reset() {
-        OutputStreamWrapper.minSizeForCompression = testDocument.length
-    }
+    private fun customCompresionApp(limit: Int): Javalin =
+        Javalin.create {
+            it.inner.compressionStrategy.minSizeForCompression = limit
+            it.addStaticFiles("/public", Location.CLASSPATH)
+        }.addTestEndpoints()
 
-    val fullCompressionApp by lazy {
+    private val fullCompressionApp by lazy {
         Javalin.create {
-            it.compressionStrategy(Brotli(), Gzip())
+            it.compressionStrategy(CompressionStrategy(Brotli(), Gzip()).apply { minSizeForCompression = testDocument.length })
             it.addStaticFiles("/public", Location.CLASSPATH)
         }.addTestEndpoints()
     }
 
-    val brotliDisabledApp by lazy {
+    private val brotliDisabledApp by lazy {
         Javalin.create {
-            it.compressionStrategy(null, Gzip())
+            it.compressionStrategy(CompressionStrategy(null, Gzip()).apply { minSizeForCompression = testDocument.length })
             it.addStaticFiles("/public", Location.CLASSPATH)
         }.addTestEndpoints()
     }
 
-    val etagApp by lazy {
+    private val etagApp by lazy {
         Javalin.create {
+            it.inner.compressionStrategy.minSizeForCompression = testDocument.length
             it.addStaticFiles("/public", Location.CLASSPATH)
             it.autogenerateEtags = true
         }.addTestEndpoints()
@@ -66,9 +69,6 @@ class TestCompression {
         get("/tiny") { ctx -> ctx.result(getSomeObjects(10).toString()) }
     }
 
-    private val tinyLength = getSomeObjects(10).toString().length
-    private val hugeLength = getSomeObjects(1000).toString().length
-
     @Test
     fun `doesn't compress when Accepts is not set`() = TestUtil.test(fullCompressionApp) { _, http ->
         assertThat(Unirest.get(http.origin + "/huge").header(Header.ACCEPT_ENCODING, "null").asString().body.length).isEqualTo(hugeLength)
@@ -79,12 +79,13 @@ class TestCompression {
     }
 
     @Test
-    fun `doesn't compress when response is too small`() = TestUtil.test(fullCompressionApp) { _, http ->
-        OutputStreamWrapper.minSizeForCompression = tinyLength + 1 // Ensure tiny response length is too short for compression
+    fun `doesn't compress when response is too small`() = TestUtil.test(customCompresionApp(tinyLength + 1)) { _, http -> // Ensure tiny response length is too short for compression
         assertThat(Unirest.get(http.origin + "/tiny").asString().body.length).isEqualTo(tinyLength)
         assertThat(getResponse(http.origin, "/tiny", "br, gzip").headers.get(Header.CONTENT_ENCODING)).isNull()
+    }
 
-        OutputStreamWrapper.minSizeForCompression = testDocument.length + 1 // Ensure static file length is too short for compression
+    @Test
+    fun `doesn't compress small static files`() = TestUtil.test(customCompresionApp(testDocument.length + 1)) { _, http ->  // Ensure static file length is too short for compression
         assertThat(Unirest.get(http.origin + "/html.html").header(Header.ACCEPT_ENCODING, "br, gzip").asString().body.length).isEqualTo(testDocument.length)
         assertThat(getResponse(http.origin, "/html.html", "br, gzip").headers.get(Header.CONTENT_ENCODING)).isNull()
     }
@@ -153,6 +154,7 @@ class TestCompression {
     @Test
     fun `does gzip when CompressionStrategy not set`() {
         val defaultApp = Javalin.create {
+            it.inner.compressionStrategy.minSizeForCompression = testDocument.length
             it.addStaticFiles("/public", Location.CLASSPATH)
         }.addTestEndpoints()
         TestUtil.test(defaultApp) { _, http ->