From 922626636c2f85eceb64cea24694b0599419bff9 Mon Sep 17 00:00:00 2001
From: Cory Benfield <lukasa@apple.com>
Date: Fri, 8 Apr 2022 15:57:37 +0100
Subject: [PATCH] Tolerate the request stream being started after .finished

Motivation

The RequestBag intermediates between two different threads. This means
it can get requests that were reasonable when they were made but have
been superseded with newer information since then. These generally have
to be tolerated.

Unfortunately if we received a request to resume the request body stream
_after_ the need for that stream has been invalidated, we could hit a
crash. That's unnecessary, and we should tolerate it better.

Modifications

Tolerated receiving requests to resume body streaming in the finished
state.

Result

Fewer crashes
---
 .../RequestBag+StateMachine.swift             |  6 ++-
 .../RequestBagTests+XCTest.swift              |  1 +
 .../RequestBagTests.swift                     | 39 +++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/Sources/AsyncHTTPClient/RequestBag+StateMachine.swift b/Sources/AsyncHTTPClient/RequestBag+StateMachine.swift
index 456317f11..59beed926 100644
--- a/Sources/AsyncHTTPClient/RequestBag+StateMachine.swift
+++ b/Sources/AsyncHTTPClient/RequestBag+StateMachine.swift
@@ -127,7 +127,11 @@ extension RequestBag.StateMachine {
             return .none
 
         case .finished:
-            preconditionFailure("Invalid state: \(self.state)")
+            // If this task has been cancelled we may be in an error state. As a matter of
+            // defensive programming, we also tolerate receiving this notification if we've ended cleanly:
+            // while it shouldn't happen, nothing will go wrong if we just ignore it.
+            // All paths through this state machine should cancel our request body stream to get here anyway.
+            return .none
 
         case .modifying:
             preconditionFailure("Invalid state: \(self.state)")
diff --git a/Tests/AsyncHTTPClientTests/RequestBagTests+XCTest.swift b/Tests/AsyncHTTPClientTests/RequestBagTests+XCTest.swift
index b2919081c..0a8f850ad 100644
--- a/Tests/AsyncHTTPClientTests/RequestBagTests+XCTest.swift
+++ b/Tests/AsyncHTTPClientTests/RequestBagTests+XCTest.swift
@@ -31,6 +31,7 @@ extension RequestBagTests {
             ("testCancelFailsTaskAfterRequestIsSent", testCancelFailsTaskAfterRequestIsSent),
             ("testCancelFailsTaskWhenTaskIsQueued", testCancelFailsTaskWhenTaskIsQueued),
             ("testFailsTaskWhenTaskIsWaitingForMoreFromServer", testFailsTaskWhenTaskIsWaitingForMoreFromServer),
+            ("testChannelBecomingWritableDoesntCrashCancelledTask", testChannelBecomingWritableDoesntCrashCancelledTask),
             ("testHTTPUploadIsCancelledEvenThoughRequestSucceeds", testHTTPUploadIsCancelledEvenThoughRequestSucceeds),
             ("testRaceBetweenConnectionCloseAndDemandMoreData", testRaceBetweenConnectionCloseAndDemandMoreData),
         ]
diff --git a/Tests/AsyncHTTPClientTests/RequestBagTests.swift b/Tests/AsyncHTTPClientTests/RequestBagTests.swift
index 2f2eb8cf5..ed50ae02d 100644
--- a/Tests/AsyncHTTPClientTests/RequestBagTests.swift
+++ b/Tests/AsyncHTTPClientTests/RequestBagTests.swift
@@ -336,6 +336,45 @@ final class RequestBagTests: XCTestCase {
         }
     }
 
+    func testChannelBecomingWritableDoesntCrashCancelledTask() {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        var maybeRequest: HTTPClient.Request?
+        XCTAssertNoThrow(maybeRequest = try HTTPClient.Request(
+            url: "https://swift.org",
+            body: .bytes([1, 2, 3, 4, 5])
+        ))
+        guard let request = maybeRequest else { return XCTFail("Expected to have a request") }
+
+        let delegate = UploadCountingDelegate(eventLoop: embeddedEventLoop)
+        var maybeRequestBag: RequestBag<UploadCountingDelegate>?
+        XCTAssertNoThrow(maybeRequestBag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        ))
+        guard let bag = maybeRequestBag else { return XCTFail("Expected to be able to create a request bag.") }
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+        executor.runRequest(bag)
+
+        // This simulates a race between the user cancelling the task (which invokes `RequestBag.cancel`) and the
+        // call to `resumeRequestBodyStream` (which comes from the `Channel` event loop and so may have to hop.
+        bag.cancel()
+        bag.resumeRequestBodyStream()
+
+        XCTAssertEqual(executor.isCancelled, true)
+        XCTAssertThrowsError(try bag.task.futureResult.wait()) {
+            XCTAssertEqual($0 as? HTTPClientError, .cancelled)
+        }
+    }
+
     func testHTTPUploadIsCancelledEvenThoughRequestSucceeds() {
         let embeddedEventLoop = EmbeddedEventLoop()
         defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }