KTOR-8528 Fix race condition when Netty removes headers for some responses

Author: Stexxe

ktor-server/ktor-server-netty/jvm/src/io/ktor/server/netty/NettyApplicationResponse.kt

@@ -11,6 +11,8 @@ import io.ktor.server.engine.*
 import io.ktor.utils.io.*
 import io.netty.channel.*
 import io.netty.handler.codec.http.*
+import kotlinx.coroutines.CompletableDeferred
+import java.util.concurrent.CancellationException
 import kotlin.coroutines.*
 
 public abstract class NettyApplicationResponse(
@@ -32,6 +34,8 @@ public abstract class NettyApplicationResponse(
 
     internal var responseChannel: ByteReadChannel = ByteReadChannel.Empty
 
+    internal val sendCompleted: CompletableDeferred<Unit> by lazy { CompletableDeferred() }
+
     override suspend fun respondOutgoingContent(content: OutgoingContent) {
         try {
             super.respondOutgoingContent(content)
@@ -60,6 +64,17 @@ public abstract class NettyApplicationResponse(
         responseMessage = message
         responseReady.setSuccess()
         responseMessageSent = true
+
+        if (isInfoOrNoContentStatus()) {
+            sendCompleted.await()
+        }
+    }
+
+    internal fun isInfoOrNoContentStatus(): Boolean {
+        val status = status()
+        if (status == null) return false
+
+        return status == HttpStatusCode.NoContent || (status.value >= 100 && status.value < 200)
     }
 
     override suspend fun responseChannel(): ByteWriteChannel {
@@ -122,7 +137,7 @@ public abstract class NettyApplicationResponse(
     public fun cancel() {
         if (!responseMessageSent) {
             responseChannel = ByteReadChannel.Empty
-            responseReady.setFailure(java.util.concurrent.CancellationException("Response was cancelled"))
+            responseReady.setFailure(CancellationException("Response was cancelled"))
             responseMessageSent = true
         }
     }

ktor-server/ktor-server-netty/jvm/src/io/ktor/server/netty/cio/NettyHttpResponsePipeline.kt

@@ -71,6 +71,9 @@ internal class NettyHttpResponsePipeline(
             respondWithFailure(call, actualException)
         } finally {
             call.responseWriteJob.cancel()
+            if (call.response.isInfoOrNoContentStatus()) {
+                call.response.sendCompleted.complete(Unit)
+            }
         }
     }