test EOF framing, fixes #2391

vapor · Jun 23, 2020 · 205663c · 205663c
1 parent fe55229
commit 205663c
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 20 deletions.
diff --git a/Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift b/Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift
@@ -83,27 +83,45 @@ final class HTTPServerResponseEncoder: ChannelOutboundHandler, RemovableChannelH
     }
 }
 
-private struct ChannelResponseBodyStream: BodyStreamWriter {
+private final class ChannelResponseBodyStream: BodyStreamWriter {
     let context: ChannelHandlerContext
     let handler: HTTPServerResponseEncoder
     let promise: EventLoopPromise<Void>?
+    var isComplete: Bool
 
     var eventLoop: EventLoop {
         return self.context.eventLoop
     }
+
+    init(
+        context: ChannelHandlerContext,
+        handler: HTTPServerResponseEncoder,
+        promise: EventLoopPromise<Void>?
+    ) {
+        self.context = context
+        self.handler = handler
+        self.promise = promise
+        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)
         case .end:
+            self.isComplete = true
             self.context.writeAndFlush(self.handler.wrapOutboundOut(.end(nil)), promise: promise)
             self.context.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
             self.promise?.succeed(())
         case .error(let error):
+            self.isComplete = true
             self.context.writeAndFlush(self.handler.wrapOutboundOut(.end(nil)), promise: promise)
             self.context.fireUserInboundEventTriggered(HTTPServerResponseEncoder.ResponseEndSentEvent())
             self.promise?.fail(error)
         }
     }
+
+    deinit {
+        assert(self.isComplete, "Response body stream writer deinitialized before .end or .error was sent.")
+    }
 }
diff --git a/Sources/Vapor/Request/Request+Body.swift b/Sources/Vapor/Request/Request+Body.swift
@@ -30,7 +30,8 @@ extension Request {
             case .collected(let buffer):
                 _ = handler(.buffer(buffer))
                 _ = handler(.end)
-            case .none: break
+            case .none:
+                _ = handler(.end)
             }
         }
 

diff --git a/Tests/VaporTests/PipelineTests.swift b/Tests/VaporTests/PipelineTests.swift
@@ -28,22 +28,14 @@ final class PipelineTests: XCTestCase {
             configuration: app.http.server.configuration
         ).wait()
 
-        //
-        // First chunk.
-        //
         try channel.writeInbound(ByteBuffer(string: "POST /echo HTTP/1.1\r\ntransfer-encoding: chunked\r\n\r\n1\r\na\r\n"))
         try XCTAssertNil(channel.readOutbound(as: ByteBuffer.self)?.string)
 
-        //
-        // Second chunk.
-        //
         try channel.writeInbound(ByteBuffer(string: "1\r\nb\r\n"))
-        do {
-            let chunk = try channel.readOutbound(as: ByteBuffer.self)?.string
-            XCTAssertContains(chunk, "HTTP/1.1 200 OK")
-            XCTAssertContains(chunk, "connection: keep-alive")
-            XCTAssertContains(chunk, "transfer-encoding: chunked")
-        }
+        let chunk = try channel.readOutbound(as: ByteBuffer.self)?.string
+        XCTAssertContains(chunk, "HTTP/1.1 200 OK")
+        XCTAssertContains(chunk, "connection: keep-alive")
+        XCTAssertContains(chunk, "transfer-encoding: chunked")
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "1\r\n")
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "a")
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "\r\n")
@@ -52,20 +44,48 @@ final class PipelineTests: XCTestCase {
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "\r\n")
         try XCTAssertNil(channel.readOutbound(as: ByteBuffer.self)?.string)
 
-        //
-        // Third chunk.
-        //
         try channel.writeInbound(ByteBuffer(string: "1\r\nc\r\n"))
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "1\r\n")
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "c")
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "\r\n")
         try XCTAssertNil(channel.readOutbound(as: ByteBuffer.self)?.string)
 
-        //
-        // Final chunk.
-        //
         try channel.writeInbound(ByteBuffer(string: "0\r\n\r\n"))
         try XCTAssertEqual(channel.readOutbound(as: ByteBuffer.self)?.string, "0\r\n\r\n")
         try XCTAssertNil(channel.readOutbound(as: ByteBuffer.self)?.string)
     }
+
+    func testEOFFraming() throws {
+        let app = Application(.testing)
+        defer { app.shutdown() }
+
+        app.on(.POST, "echo", body: .stream) { request -> Response in
+            Response(body: .init(stream: { writer in
+                request.body.drain { body in
+                    switch body {
+                    case .buffer(let buffer):
+                        return writer.write(.buffer(buffer))
+                    case .error(let error):
+                        return writer.write(.error(error))
+                    case .end:
+                        return writer.write(.end)
+                    }
+                }
+            }))
+        }
+
+        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")
+    }
+
+    override class func setUp() {
+        XCTAssert(isLoggingConfigured)
+    }
 }