DeflateDecompressor: fix inflater memory leak (#1137)

http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/DeflateSpec.scala

@@ -16,12 +16,16 @@ package org.apache.pekko.http.scaladsl.coding
 import org.apache.pekko
 import pekko.util.ByteString
 import java.io.{ InputStream, OutputStream }
+import java.util.concurrent.{ CountDownLatch, TimeUnit }
+import java.util.concurrent.atomic.AtomicInteger
 import java.util.zip._
 
 import pekko.http.scaladsl.model.HttpMethods.POST
 import pekko.http.scaladsl.model.{ HttpEntity, HttpRequest }
 import pekko.http.impl.util._
 import pekko.http.scaladsl.model.headers.{ `Content-Encoding`, HttpEncodings }
+import pekko.stream.SystemMaterializer
+import pekko.stream.scaladsl.{ Sink, Source }
 import pekko.testkit._
 import scala.annotation.nowarn
 
@@ -61,6 +65,62 @@ class DeflateSpec extends CoderSpec {
         3.seconds.dilated)
         .awaitResult(3.seconds.dilated) should equal(request)
     }
+    "release the inflater when decoding completes" in {
+      val inflater = new TrackingInflater
+      decodeWith(inflater, streamEncode(smallTextBytes)) should readAs(smallText)
+      inflater.endCalls.get() shouldEqual 1
+    }
+    "release the inflater when decoding is cancelled early" in {
+      val inflater = new TrackingInflater
+      val compressed = streamEncode(largeTextBytes)
+
+      Source.single(compressed)
+        .via(decoderWith(inflater).withMaxBytesPerChunk(1).decoderFlow)
+        .take(1)
+        .runWith(Sink.ignore)
+        .awaitResult(3.seconds.dilated)
+
+      // postStop() (which calls end()) is dispatched to the stage actor after the
+      // Sink.ignore future completes, so we must wait for end() itself rather than
+      // for the stream future to avoid a race.
+      inflater.awaitEnd(3.seconds.dilated)
+      inflater.endCalls.get() shouldEqual 1
+    }
+    "release the inflater when decoding is truncated" in {
+      val inflater = new TrackingInflater
+      // Truncated deflate streams complete without exception (completeStage on truncation)
+      decodeWith(inflater, streamEncode(smallTextBytes).dropRight(5))
+      inflater.endCalls.get() shouldEqual 1
+    }
+  }
+
+  private def decodeWith(inflater: TrackingInflater, bytes: ByteString): ByteString =
+    decoderWith(inflater).decode(bytes)(SystemMaterializer(system).materializer).awaitResult(3.seconds.dilated)
+
+  @nowarn("msg=deprecated")
+  private def decoderWith(inflater: TrackingInflater): StreamDecoder =
+    new StreamDecoder {
+      override val encoding = HttpEncodings.deflate
+
+      override def newDecompressorStage(maxBytesPerChunk: Int) =
+        () =>
+          new DeflateDecompressor(maxBytesPerChunk) {
+            override protected[coding] def createInflater(noWrap: Boolean) = inflater
+          }
+    }
+
+  private class TrackingInflater extends java.util.zip.Inflater(false) {
+    val endCalls = new AtomicInteger
+    private val endLatch = new CountDownLatch(1)
+
+    def awaitEnd(atMost: FiniteDuration): Unit =
+      endLatch.await(atMost.toMillis, TimeUnit.MILLISECONDS)
+
+    override def end(): Unit = {
+      endCalls.incrementAndGet()
+      endLatch.countDown()
+      super.end()
+    }
   }
 
   private def encodeMessage(request: HttpRequest, compressionLevel: Int, noWrap: Boolean): HttpRequest = {

http/src/main/scala/org/apache/pekko/http/scaladsl/coding/DeflateCompressor.scala

@@ -104,14 +104,28 @@ private[coding] object DeflateCompressor {
 private[coding] class DeflateDecompressor(
     maxBytesPerChunk: Int = Decoder.MaxBytesPerChunkDefault) extends DeflateDecompressorBase(maxBytesPerChunk) {
 
+  protected[coding] def createInflater(noWrap: Boolean): Inflater = new Inflater(noWrap)
+
   override def createLogic(attr: Attributes) = new ParsingLogic {
+    private var currentInflater: Inflater = null
+    private var inflaterEnded = false
+
+    private def cleanupInflater(): Unit =
+      if (!inflaterEnded && currentInflater != null) {
+        inflaterEnded = true
+        currentInflater.end()
+      }
+
+    override def postStop(): Unit = cleanupInflater()
 
     /** Step that probes if the deflate stream contains a zlib wrapper */
     case object ProbeWrapping extends ParseStep[ByteString] {
       override def onTruncation(): Unit = completeStage()
 
       override def parse(reader: ByteStringParser.ByteReader): ParseResult[ByteString] = {
         val inflater = examineAndBuildInflater(reader.remainingData)
+        cleanupInflater()
+        currentInflater = inflater
         ParseResult(None, new Inflate(inflater, noPostProcessing = true, ProbeWrapping))
       }
     }
@@ -132,7 +146,7 @@ private[coding] class DeflateDecompressor(
      */
     private def examineAndBuildInflater(bytes: ByteString): Inflater = {
       val wrapped = (bytes.head & 0x0F) == 0x08
-      new Inflater(!wrapped)
+      createInflater(!wrapped)
     }
 
     startWith(ProbeWrapping)