Propagate HTTPClient.Task<Response> failures to subsequent redirect tasks (#814)

Discussed in
https://github.com/swift-server/async-http-client/issues/753

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -956,7 +956,7 @@ internal struct RedirectHandler<ResponseType> {
         status: HTTPResponseStatus,
         to redirectURL: URL,
         promise: EventLoopPromise<ResponseType>
-    ) {
+    ) -> HTTPClient.Task<ResponseType>? {
         do {
             var redirectState = self.redirectState
             try redirectState.redirect(to: redirectURL.absoluteString)
@@ -976,13 +976,19 @@ internal struct RedirectHandler<ResponseType> {
                 headers: headers,
                 body: body
             )
-            self.execute(newRequest, redirectState).futureResult.whenComplete { result in
+
+            let newTask = self.execute(newRequest, redirectState)
+
+            newTask.futureResult.whenComplete { result in
                 promise.futureResult.eventLoop.execute {
                     promise.completeWith(result)
                 }
             }
+
+            return newTask
         } catch {
             promise.fail(error)
+            return nil
         }
     }
 }

Sources/AsyncHTTPClient/RequestBag.swift

@@ -43,6 +43,9 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
     // the consume body part stack depth is synchronized on the task event loop.
     private var consumeBodyPartStackDepth: Int
 
+    // if a redirect occurs, we store the task for it so we can propagate cancellation
+    private var redirectTask: HTTPClient.Task<Delegate.Response>? = nil
+
     // MARK: HTTPClientTask properties
 
     var logger: Logger {
@@ -234,7 +237,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             executor.demandResponseBodyStream(self)
 
         case .redirect(let executor, let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
             executor.cancelRequest(self)
 
         case .forwardResponseHead(let head):
@@ -258,7 +261,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             executor.demandResponseBodyStream(self)
 
         case .redirect(let executor, let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
             executor.cancelRequest(self)
 
         case .forwardResponsePart(let part):
@@ -294,7 +297,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             }
 
         case .redirect(let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
         }
     }
 
@@ -358,6 +361,8 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
         let action = self.state.fail(error)
 
         self.executeFailAction0(action)
+
+        self.redirectTask?.fail(reason: error)
     }
 
     private func executeFailAction0(_ action: RequestBag<Delegate>.StateMachine.FailAction) {

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -4132,6 +4132,70 @@ final class HTTPClientTests: XCTestCaseHTTPClientTestsBaseClass {
         XCTAssertNoThrow(try client.execute(request: request).wait())
     }
 
+    func testCancelingRequestAfterRedirect() throws {
+        let request = try Request(
+            url: self.defaultHTTPBinURLPrefix + "redirect/target",
+            method: .GET,
+            headers: ["X-Target-Redirect-URL": self.defaultHTTPBinURLPrefix + "wait"],
+            body: nil
+        )
+
+        class CancelAfterRedirect: HTTPClientResponseDelegate {
+            init() {}
+            func didFinishRequest(task: AsyncHTTPClient.HTTPClient.Task<Void>) throws {}
+        }
+
+        let task = defaultClient.execute(
+            request: request,
+            delegate: CancelAfterRedirect(),
+            deadline: .now() + .seconds(1)
+        )
+
+        // there is currently no HTTPClientResponseDelegate method to ensure the redirect occurs before we cancel, so we just sleep for 500ms
+        Thread.sleep(forTimeInterval: 0.5)
+
+        task.cancel()
+
+        XCTAssertThrowsError(try task.wait()) { error in
+            guard case let error = error as? HTTPClientError, error == .cancelled else {
+                return XCTFail("Should fail with cancelled")
+            }
+        }
+    }
+
+    func testFailingRequestAfterRedirect() throws {
+        let request = try Request(
+            url: self.defaultHTTPBinURLPrefix + "redirect/target",
+            method: .GET,
+            headers: ["X-Target-Redirect-URL": self.defaultHTTPBinURLPrefix + "wait"],
+            body: nil
+        )
+
+        class FailAfterRedirect: HTTPClientResponseDelegate {
+            init() {}
+            func didFinishRequest(task: AsyncHTTPClient.HTTPClient.Task<Void>) throws {}
+        }
+
+        let task = defaultClient.execute(
+            request: request,
+            delegate: FailAfterRedirect(),
+            deadline: .now() + .seconds(1)
+        )
+
+        // there is currently no HTTPClientResponseDelegate method to ensure the redirect occurs before we fail, so we just sleep for 500ms
+        Thread.sleep(forTimeInterval: 0.5)
+
+        struct TestError: Error {}
+
+        task.fail(reason: TestError())
+
+        XCTAssertThrowsError(try task.wait()) { error in
+            guard error is TestError else {
+                return XCTFail("Should fail with TestError")
+            }
+        }
+    }
+
     func testCancelingHTTP1RequestAfterHeaderSend() throws {
         var request = try HTTPClient.Request(url: self.defaultHTTPBin.baseURL + "/wait", method: .POST)
         // non-empty body is important