test incorrect response body stream count, fixes #2392

Sources/Vapor/HTTP/Server/HTTPServerHandler.swift

@@ -36,9 +36,14 @@ final class HTTPServerHandler: ChannelInboundHandler, RemovableChannelHandler {
         default:
             response.headers.add(name: .connection, value: request.isKeepAlive ? "keep-alive" : "close")
             let done = context.write(self.wrapOutboundOut(response))
-            if !request.isKeepAlive {
-                done.whenComplete { _ in
-                    context.close(mode: .output, promise: nil)
+            done.whenComplete { result in
+                switch result {
+                case .success:
+                    if !request.isKeepAlive {
+                        context.close(mode: .output, promise: nil)
+                    }
+                case .failure(let error):
+                    self.errorCaught(context: context, error: error)
                 }
             }
         }

Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift

@@ -66,7 +66,8 @@ final class HTTPServerResponseEncoder: ChannelOutboundHandler, RemovableChannelH
                 let channelStream = ChannelResponseBodyStream(
                     context: context,
                     handler: self,
-                    promise: promise
+                    promise: promise,
+                    count: stream.count
                 )
                 stream.callback(channelStream)
             }
@@ -87,29 +88,52 @@ private final class ChannelResponseBodyStream: BodyStreamWriter {
     let context: ChannelHandlerContext
     let handler: HTTPServerResponseEncoder
     let promise: EventLoopPromise<Void>?
+    let count: Int
+    var currentCount: Int
     var isComplete: Bool
 
     var eventLoop: EventLoop {
         return self.context.eventLoop
     }
 
+    enum Error: Swift.Error {
+        case tooManyBytes
+        case notEnoughBytes
+    }
+
     init(
         context: ChannelHandlerContext,
         handler: HTTPServerResponseEncoder,
-        promise: EventLoopPromise<Void>?
+        promise: EventLoopPromise<Void>?,
+        count: Int
     ) {
         self.context = context
         self.handler = handler
         self.promise = promise
+        self.count = count
+        self.currentCount = 0
         self.isComplete = false
     }
 
     func write(_ result: BodyStreamResult, promise: EventLoopPromise<Void>?) {
         switch result {
         case .buffer(let buffer):
             self.context.writeAndFlush(self.handler.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: promise)
+            self.currentCount += buffer.readableBytes
+            guard self.currentCount <= self.count else {
+                self.promise?.fail(Error.tooManyBytes)
+                promise?.fail(Error.notEnoughBytes)
+                return
+                // self.context.close(promise: nil)
+            }
         case .end:
             self.isComplete = true
+            guard self.currentCount == self.count else {
+                self.promise?.fail(Error.notEnoughBytes)
+                promise?.fail(Error.notEnoughBytes)
+                return
+                // self.context.close(promise: nil)
+            }
             self.context.writeAndFlush(self.handler.wrapOutboundOut(.end(nil)), promise: promise)
             self.context.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
             self.promise?.succeed(())

Tests/VaporTests/PipelineTests.swift

@@ -85,6 +85,30 @@ final class PipelineTests: XCTestCase {
         try XCTAssertContains(channel.readOutbound(as: ByteBuffer.self)?.string, "HTTP/1.1 200 OK")
     }
 
+    func testBadStreamLength() throws {
+        let app = Application(.testing)
+        defer { app.shutdown() }
+
+        app.on(.POST, "echo", body: .stream) { request -> Response in
+            Response(body: .init(stream: { writer in
+                writer.write(.buffer(.init(string: "a")), promise: nil)
+                writer.write(.end, promise: nil)
+            }, count: 2))
+        }
+
+        let channel = EmbeddedChannel()
+        try channel.pipeline.addVaporHTTP1Handlers(
+            application: app,
+            responder: app.responder,
+            configuration: app.http.server.configuration
+        ).wait()
+
+        try channel.writeInbound(ByteBuffer(string: "POST /echo HTTP/1.1\r\n\r\n"))
+        try XCTAssertContains(channel.readOutbound(as: ByteBuffer.self)?.string, "HTTP/1.1 200 OK")
+        try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "a")
+        try XCTAssertNil(channel.readOutbound(as: ByteBuffer.self)?.string)
+    }
+
     override class func setUp() {
         XCTAssert(isLoggingConfigured)
     }