From d642f14a6a1df6855cb583d7ab0b60cf0863b470 Mon Sep 17 00:00:00 2001
From: Fabian Fett <fabianfett@apple.com>
Date: Mon, 24 Jan 2022 12:27:03 +0100
Subject: [PATCH] Fix race between connection close and scheduling new request

---
 .../HTTP1/HTTP1ConnectionStateMachine.swift   | 34 ++++++++++++++-----
 ...P1ConnectionStateMachineTests+XCTest.swift |  1 +
 .../HTTP1ConnectionStateMachineTests.swift    | 14 ++++++++
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift
index 743626607..19825aec7 100644
--- a/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift
+++ b/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift
@@ -158,17 +158,35 @@ struct HTTP1ConnectionStateMachine {
         head: HTTPRequestHead,
         metadata: RequestFramingMetadata
     ) -> Action {
-        guard case .idle = self.state else {
+        switch self.state {
+        case .initialized, .closing, .inRequest:
+            // These states are unreachable as the connection pool state machine has put the
+            // connection into these states. In other words the connection pool state machine must
+            // be aware about these states before the connection itself. For this reason the
+            // connection pool state machine must not send a new request to the connection, if the
+            // connection is `.initialized`, `.closing` or `.inRequest`
             preconditionFailure("Invalid state: \(self.state)")
-        }
 
-        var requestStateMachine = HTTPRequestStateMachine(isChannelWritable: self.isChannelWritable)
-        let action = requestStateMachine.startRequest(head: head, metadata: metadata)
+        case .closed:
+            // The remote may have closed the connection and the connection pool state machine
+            // was not updated yet because of a race condition. New request vs. marking connection
+            // as closed.
+            //
+            // TODO: AHC should support a fast rescheduling mechanism here.
+            return .failRequest(HTTPClientError.remoteConnectionClosed, .none)
+
+        case .idle:
+            var requestStateMachine = HTTPRequestStateMachine(isChannelWritable: self.isChannelWritable)
+            let action = requestStateMachine.startRequest(head: head, metadata: metadata)
 
-        // by default we assume a persistent connection. however in `requestVerified`, we read the
-        // "connection" header.
-        self.state = .inRequest(requestStateMachine, close: metadata.connectionClose)
-        return self.state.modify(with: action)
+            // by default we assume a persistent connection. however in `requestVerified`, we read the
+            // "connection" header.
+            self.state = .inRequest(requestStateMachine, close: metadata.connectionClose)
+            return self.state.modify(with: action)
+
+        case .modifying:
+            preconditionFailure("Invalid state: \(self.state)")
+        }
     }
 
     mutating func requestStreamPartReceived(_ part: IOData) -> Action {
diff --git a/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests+XCTest.swift b/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests+XCTest.swift
index 75bc6c017..76a37936b 100644
--- a/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests+XCTest.swift
+++ b/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests+XCTest.swift
@@ -42,6 +42,7 @@ extension HTTP1ConnectionStateMachineTests {
             ("testConnectionIsClosedIfErrorHappensWhileInRequest", testConnectionIsClosedIfErrorHappensWhileInRequest),
             ("testConnectionIsClosedAfterSwitchingProtocols", testConnectionIsClosedAfterSwitchingProtocols),
             ("testWeDontCrashAfterEarlyHintsAndConnectionClose", testWeDontCrashAfterEarlyHintsAndConnectionClose),
+            ("testWeDontCrashInRaceBetweenSchedulingNewRequestAndConnectionClose", testWeDontCrashInRaceBetweenSchedulingNewRequestAndConnectionClose),
         ]
     }
 }
diff --git a/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift b/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift
index 1ceb9f2bf..c8ad3d510 100644
--- a/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift
+++ b/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift
@@ -267,6 +267,20 @@ class HTTP1ConnectionStateMachineTests: XCTestCase {
         XCTAssertEqual(state.channelRead(.head(responseHead)), .wait)
         XCTAssertEqual(state.channelInactive(), .failRequest(HTTPClientError.remoteConnectionClosed, .none))
     }
+
+    func testWeDontCrashInRaceBetweenSchedulingNewRequestAndConnectionClose() {
+        var state = HTTP1ConnectionStateMachine()
+        XCTAssertEqual(state.channelActive(isWritable: true), .fireChannelActive)
+        XCTAssertEqual(state.channelInactive(), .fireChannelInactive)
+
+        let requestHead = HTTPRequestHead(version: .http1_1, method: .GET, uri: "/")
+        let metadata = RequestFramingMetadata(connectionClose: false, body: .fixedSize(0))
+        let newRequestAction = state.runNewRequest(head: requestHead, metadata: metadata)
+        guard case .failRequest(let error, .none) = newRequestAction else {
+            return XCTFail("Unexpected test case")
+        }
+        XCTAssertEqual(error as? HTTPClientError, .remoteConnectionClosed)
+    }
 }
 
 extension HTTP1ConnectionStateMachine.Action: Equatable {