Prevent TaskHandler state change after .endOrError

Motivation:
Right now if task handler encounters an error, it changes state to
`.endOrError`. We gate on that state to make sure that we do not
process errors in the pipeline twice. Unfortunately, that state
can be reset when we upload body or receive response parts.

Modifications:
Adds state validation before state is updated to a new value
Adds a test

Result:
Fixes #297

Author: artemredkin

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -840,17 +840,27 @@ extension TaskHandler: ChannelDuplexHandler {
             self.writeBody(request: request, context: context)
         }.flatMap {
             context.eventLoop.assertInEventLoop()
-            self.state = .bodySent
-            if let expectedBodyLength = self.expectedBodyLength, expectedBodyLength != self.actualBodyLength {
-                let error = HTTPClientError.bodyLengthMismatch
-                self.errorCaught(context: context, error: error)
-                return context.eventLoop.makeFailedFuture(error)
+            switch self.state {
+            case .endOrError:
+                return context.eventLoop.makeSucceededFuture(())
+            default:
+                self.state = .bodySent
+                if let expectedBodyLength = self.expectedBodyLength, expectedBodyLength != self.actualBodyLength {
+                    let error = HTTPClientError.bodyLengthMismatch
+                    self.errorCaught(context: context, error: error)
+                    return context.eventLoop.makeFailedFuture(error)
+                }
+                return context.writeAndFlush(self.wrapOutboundOut(.end(nil)))
             }
-            return context.writeAndFlush(self.wrapOutboundOut(.end(nil)))
         }.map {
             context.eventLoop.assertInEventLoop()
-            self.state = .sent
-            self.callOutToDelegateFireAndForget(self.delegate.didSendRequest)
+            switch self.state {
+            case .endOrError:
+                break
+            default:
+                self.state = .sent
+                self.callOutToDelegateFireAndForget(self.delegate.didSendRequest)
+            }
         }.flatMapErrorThrowing { error in
             context.eventLoop.assertInEventLoop()
             self.errorCaught(context: context, error: error)
@@ -926,7 +936,7 @@ extension TaskHandler: ChannelDuplexHandler {
         case .head(let head):
             switch self.state {
             case .endOrError:
-                preconditionFailure("unexpected state on .head")
+                break
             default:
                 if !head.isKeepAlive {
                     self.closing = true
@@ -945,10 +955,8 @@ extension TaskHandler: ChannelDuplexHandler {
             }
         case .body(let body):
             switch self.state {
-            case .redirected:
+            case .redirected, .endOrError:
                 break
-            case .endOrError:
-                preconditionFailure("unexpected state on .body")
             default:
                 self.state = .body
                 self.mayRead = false
@@ -960,7 +968,7 @@ extension TaskHandler: ChannelDuplexHandler {
         case .end:
             switch self.state {
             case .endOrError:
-                preconditionFailure("unexpected state on .end")
+                break
             case .redirected(let head, let redirectURL):
                 self.state = .endOrError
                 self.task.releaseAssociatedConnection(delegateType: Delegate.self, closing: self.closing).whenSuccess {

Tests/AsyncHTTPClientTests/HTTPClientInternalTests+XCTest.swift

@@ -47,6 +47,7 @@ extension HTTPClientInternalTests {
             ("testInternalRequestURI", testInternalRequestURI),
             ("testBodyPartStreamStateChangedBeforeNotification", testBodyPartStreamStateChangedBeforeNotification),
             ("testHandlerDoubleError", testHandlerDoubleError),
+            ("testTaskHandlerStateChangeAfterError", testTaskHandlerStateChangeAfterError),
         ]
     }
 }

Tests/AsyncHTTPClientTests/HTTPClientInternalTests.swift

@@ -1119,4 +1119,48 @@ class HTTPClientInternalTests: XCTestCase {
 
         XCTAssertEqual(delegate.count, 1)
     }
+
+    func testTaskHandlerStateChangeAfterError() throws {
+        let channel = EmbeddedChannel()
+        let task = Task<Void>(eventLoop: channel.eventLoop, logger: HTTPClient.loggingDisabled)
+
+        let handler = TaskHandler(task: task,
+                                  kind: .host,
+                                  delegate: TestHTTPDelegate(),
+                                  redirectHandler: nil,
+                                  ignoreUncleanSSLShutdown: false,
+                                  logger: HTTPClient.loggingDisabled)
+
+        try channel.pipeline.addHandler(handler).wait()
+
+        var request = try Request(url: "http://localhost:8080/get")
+        request.headers.add(name: "X-Test-Header", value: "X-Test-Value")
+        request.body = .stream(length: 4) { writer in
+            writer.write(.byteBuffer(channel.allocator.buffer(string: "1234"))).map {
+                handler.state = .endOrError
+            }
+        }
+
+        XCTAssertNoThrow(try channel.writeOutbound(request))
+
+        try channel.writeInbound(HTTPClientResponsePart.head(.init(version: .init(major: 1, minor: 1), status: .ok)))
+        XCTAssertTrue(handler.state.isEndOrError)
+
+        try channel.writeInbound(HTTPClientResponsePart.body(channel.allocator.buffer(string: "1234")))
+        XCTAssertTrue(handler.state.isEndOrError)
+
+        try channel.writeInbound(HTTPClientResponsePart.end(nil))
+        XCTAssertTrue(handler.state.isEndOrError)
+    }
+}
+
+extension TaskHandler.State {
+    var isEndOrError: Bool {
+        switch self {
+        case .endOrError:
+            return true
+        default:
+            return false
+        }
+    }
 }