Tests for inflator failure.

okhttp/src/main/kotlin/okhttp3/internal/ws/MessageInflater.kt

@@ -55,6 +55,11 @@ class MessageInflater(
     do {
       inflaterSource.readOrInflate(buffer, Long.MAX_VALUE)
     } while (inflater.bytesRead < totalBytesToRead && !inflater.finished())
+
+    if (noContextTakeover && inflater.bytesRead < totalBytesToRead) {
+      val left = totalBytesToRead - inflater.bytesRead
+      throw IOException("additional bytes ($left) after inflation")
+    }
   }
 
   @Throws(IOException::class)

okhttp/src/test/java/okhttp3/internal/ws/MessageDeflaterInflaterTest.kt

@@ -16,7 +16,9 @@
 package okhttp3.internal.ws
 
 import assertk.assertThat
+import assertk.assertions.hasMessage
 import assertk.assertions.isEqualTo
+import assertk.assertions.isGreaterThan
 import assertk.assertions.isLessThan
 import java.io.EOFException
 import kotlin.test.assertFailsWith
@@ -25,10 +27,14 @@ import okio.Buffer
 import okio.ByteString
 import okio.ByteString.Companion.decodeHex
 import okio.ByteString.Companion.encodeUtf8
+import okio.FileSystem
+import okio.IOException
+import okio.Path.Companion.toPath
 import org.junit.jupiter.api.Test
 
 internal class MessageDeflaterInflaterTest {
-  @Test fun `inflate golden value`() {
+  @Test
+  fun `inflate golden value`() {
     val inflater = MessageInflater(false)
     val message = "f248cdc9c957c8cc4bcb492cc9cccf530400".decodeHex()
     assertThat(inflater.inflate(message)).isEqualTo("Hello inflation!".encodeUtf8())
@@ -38,19 +44,22 @@ internal class MessageDeflaterInflaterTest {
    * We had a bug where self-finishing inflater streams would infinite loop!
    * https://github.com/square/okhttp/issues/8078
    */
-  @Test fun `inflate returns finished before bytesRead reaches input length`() {
+  @Test
+  fun `inflate returns finished before bytesRead reaches input length`() {
     val inflater = MessageInflater(false)
     val message = "53621260020000".decodeHex()
     assertThat(inflater.inflate(message)).isEqualTo("22021002".decodeHex())
   }
 
-  @Test fun `deflate golden value`() {
+  @Test
+  fun `deflate golden value`() {
     val deflater = MessageDeflater(false)
     val deflated = deflater.deflate("Hello deflate!".encodeUtf8())
     assertThat(deflated.hex()).isEqualTo("f248cdc9c95748494dcb492c49550400")
   }
 
-  @Test fun `inflate deflate`() {
+  @Test
+  fun `inflate deflate`() {
     val deflater = MessageDeflater(false)
     val inflater = MessageInflater(false)
 
@@ -63,7 +72,8 @@ internal class MessageDeflaterInflaterTest {
     assertThat(inflated).isEqualTo(goldenValue)
   }
 
-  @Test fun `inflate deflate empty message`() {
+  @Test
+  fun `inflate deflate empty message`() {
     val deflater = MessageDeflater(false)
     val inflater = MessageInflater(false)
 
@@ -76,7 +86,8 @@ internal class MessageDeflaterInflaterTest {
     assertThat(inflated).isEqualTo(goldenValue)
   }
 
-  @Test fun `inflate deflate with context takeover`() {
+  @Test
+  fun `inflate deflate with context takeover`() {
     val deflater = MessageDeflater(false)
     val inflater = MessageInflater(false)
 
@@ -91,7 +102,8 @@ internal class MessageDeflaterInflaterTest {
     assertThat(deflatedValue2.size).isLessThan(deflatedValue1.size)
   }
 
-  @Test fun `inflate deflate with no context takeover`() {
+  @Test
+  fun `inflate deflate with no context takeover`() {
     val deflater = MessageDeflater(true)
     val inflater = MessageInflater(true)
 
@@ -106,7 +118,8 @@ internal class MessageDeflaterInflaterTest {
     assertThat(deflatedValue2).isEqualTo(deflatedValue1)
   }
 
-  @Test fun `deflate after close`() {
+  @Test
+  fun `deflate after close`() {
     val deflater = MessageDeflater(true)
     deflater.close()
 
@@ -115,7 +128,8 @@ internal class MessageDeflaterInflaterTest {
     }
   }
 
-  @Test fun `inflate after close`() {
+  @Test
+  fun `inflate after close`() {
     val inflater = MessageInflater(false)
 
     inflater.close()
@@ -125,11 +139,44 @@ internal class MessageDeflaterInflaterTest {
     }
   }
 
+  @Test
+  fun `per message compression left with deflated bytes`() {
+    // per message compression
+    val inflater = MessageInflater(noContextTakeover = true)
+    val message =
+      FileSystem.RESOURCES.read("/okhttp3/internal/ws/inflater_notworking.txt".toPath()) {
+        readUtf8().trim().decodeHex()
+      }
+    val buffer = fragmentBuffer(Buffer().write(message))
+    val thrown = assertFailsWith<IOException> {
+      inflater.inflate(buffer)
+    }
+    assertThat(thrown).hasMessage("additional bytes (5) after inflation")
+  }
+
+  @Test
+  fun `per message compression left with deflated bytes working message`() {
+    // per message compression
+    val inflater = MessageInflater(noContextTakeover = true)
+    val message =
+      FileSystem.RESOURCES.read("/okhttp3/internal/ws/inflater_working.txt".toPath()) {
+        readUtf8().trim().decodeHex()
+      }
+    val buffer = fragmentBuffer(Buffer().write(message))
+    inflater.inflate(buffer)
+    assertThat(buffer.size).isGreaterThan(0)
+
+    val bufferAgain = fragmentBuffer(Buffer().write(message))
+    inflater.inflate(bufferAgain)
+    assertThat(bufferAgain.size).isGreaterThan(0)
+  }
+
   /**
    * Test for an [EOFException] caused by mishandling of fragmented buffers in web socket
    * compression. https://github.com/square/okhttp/issues/5965
    */
-  @Test fun `inflate golden value in buffer that has been fragmented`() {
+  @Test
+  fun `inflate golden value in buffer that has been fragmented`() {
     val inflater = MessageInflater(false)
     val buffer = fragmentBuffer(Buffer().write("f248cdc9c957c8cc4bcb492cc9cccf530400".decodeHex()))
     inflater.inflate(buffer)