From f50bf983ea3729010dd041d3b16fcdf99cf37ca0 Mon Sep 17 00:00:00 2001
From: Cory Benfield <lukasa@apple.com>
Date: Mon, 11 Apr 2022 10:13:56 +0100
Subject: [PATCH] Tolerate the request stream being started after .finished
 (#577)

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
Fixes #576
---
 .../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 456317f..59beed9 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 b291908..0a8f850 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 2f2eb8c..ed50ae0 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()) }