Merge pull request onevcat#2048 from onevcat/feature/response-delegate

Feature response delegate

Sources/General/KingfisherError.swift

@@ -88,6 +88,10 @@ public enum KingfisherError: Error {
         /// The task is done but no URL response found. Code 2005.
         /// - task: The failed task.
         case noURLResponse(task: SessionDataTask)
+
+        /// The task is cancelled by `ImageDownloaderDelegate` due to the `.cancel` response disposition is
+        /// specified by the delegate method. Code 2006.
+        case cancelledByDelegate(response: URLResponse)
     }
 
     /// Represents the error reason during Kingfisher caching system.
@@ -345,7 +349,10 @@ extension KingfisherError.ResponseErrorReason {
         case .dataModifyingFailed(let task):
             return "The data modifying delegate returned `nil` for the downloaded data. Task: \(task)."
         case .noURLResponse(let task):
-            return "No URL response received. Task: \(task),"
+            return "No URL response received. Task: \(task)."
+        case .cancelledByDelegate(let response):
+            return "The downloading task is cancelled by the downloader delegate. Response: \(response)."
+
         }
     }
 
@@ -356,6 +363,7 @@ extension KingfisherError.ResponseErrorReason {
         case .URLSessionError: return 2003
         case .dataModifyingFailed: return 2004
         case .noURLResponse: return 2005
+        case .cancelledByDelegate: return 2006
         }
     }
 }

Sources/Networking/ImageDownloader.swift

@@ -184,6 +184,9 @@ open class ImageDownloader {
         sessionDelegate.onValidStatusCode.delegate(on: self) { (self, code) in
             return (self.delegate ?? self).isValidStatusCode(code, for: self)
         }
+        sessionDelegate.onResponseReceived.delegate(on: self) { (self, invoke) in
+            (self.delegate ?? self).imageDownloader(self, didReceive: invoke.0, completionHandler: invoke.1)
+        }
         sessionDelegate.onDownloadingFinished.delegate(on: self) { (self, value) in
             let (url, result) = value
             do {

Sources/Networking/ImageDownloaderDelegate.swift

@@ -116,6 +116,28 @@ public protocol ImageDownloaderDelegate: AnyObject {
     /// - Note: If the default 200 to 400 valid code does not suit your need,
     ///         you can implement this method to change that behavior.
     func isValidStatusCode(_ code: Int, for downloader: ImageDownloader) -> Bool
+
+    /// Called when the task has received a valid HTTP response after it passes other checks such as the status code.
+    /// You can perform additional checks or verification on the response to determine if the download should be allowed.
+    ///
+    /// For example, it is useful if you want to verify some header values in the response before actually starting the
+    /// download.
+    ///
+    /// If implemented, it is your responsibility to call the `completionHandler` with a proper response disposition,
+    /// such as `.allow` to start the actual downloading or `.cancel` to cancel the task. If `.cancel` is used as the
+    /// disposition, the downloader will raise an `KingfisherError` with
+    /// `ResponseErrorReason.cancelledByDelegate` as its reason. If not implemented, any response which passes other
+    /// checked will be allowed and the download starts.
+    ///
+    /// - Parameters:
+    ///   - downloader: The `ImageDownloader` object which is used for the downloading operation.
+    ///   - response: The original response object of the downloading process.
+    ///   - completionHandler: A completion handler that receives the disposition for the download task. You must call
+    ///   this handler with either `.allow` or `.cancel`.
+    func imageDownloader(
+        _ downloader: ImageDownloader,
+        didReceive response: URLResponse,
+        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void)
 }
 
 // Default implementation for `ImageDownloaderDelegate`.
@@ -151,4 +173,12 @@ extension ImageDownloaderDelegate {
     public func imageDownloader(_ downloader: ImageDownloader, didDownload data: Data, for url: URL) -> Data? {
         return data
     }
+
+    public func imageDownloader(
+        _ downloader: ImageDownloader,
+        didReceive response: URLResponse,
+        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
+            completionHandler(.allow)
+        }
+
 }

Sources/Networking/SessionDelegate.swift

@@ -47,6 +47,7 @@ open class SessionDelegate: NSObject {
     private let lock = NSLock()
 
     let onValidStatusCode = Delegate<Int, Bool>()
+    let onResponseReceived = Delegate<(URLResponse, (URLSession.ResponseDisposition) -> Void), Void>()
     let onDownloadingFinished = Delegate<(URL, Result<URLResponse, KingfisherError>), Void>()
     let onDidDownloadData = Delegate<SessionDataTask, Data?>()
 
@@ -169,7 +170,15 @@ extension SessionDelegate: URLSessionDataDelegate {
             completionHandler(.cancel)
             return
         }
-        completionHandler(.allow)
+
+        let inspectedHandler: (URLSession.ResponseDisposition) -> Void = { disposition in
+            if disposition == .cancel {
+                let error = KingfisherError.responseError(reason: .cancelledByDelegate(response: response))
+                self.onCompleted(task: dataTask, result: .failure(error))
+            }
+            completionHandler(disposition)
+        }
+        onResponseReceived.call((response, inspectedHandler))
     }
 
     open func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {

Tests/KingfisherTests/ImageDownloaderTests.swift

@@ -547,9 +547,9 @@ class ImageDownloaderTests: XCTestCase {
 
 
     func testSessionDelegate() {
-        class ExtensionDelegate:SessionDelegate {
+        class ExtensionDelegate: SessionDelegate {
             //'exp' only for test
-            public let exp:XCTestExpectation
+            public let exp: XCTestExpectation
             init(_ expectation:XCTestExpectation) {
                 exp = expectation
             }
@@ -566,6 +566,78 @@ class ImageDownloaderTests: XCTestCase {
         }
         waitForExpectations(timeout: 3, handler: nil)
     }
+
+    func testDownloaderReceiveResponsePass() {
+
+        let exp = expectation(description: #function)
+
+        let url = testURLs[0]
+        stub(url, data: testImageData, headers: ["someKey": "someValue"])
+
+        let handler = TaskResponseCompletion()
+        let obj = NSObject()
+        handler.onReceiveResponse.delegate(on: obj) { (obj, response) in
+            guard let httpResponse = response as? HTTPURLResponse else {
+                XCTFail("Should be an HTTP response.")
+                return .cancel
+            }
+            XCTAssertEqual(httpResponse.statusCode, 200)
+            XCTAssertEqual(httpResponse.url, url)
+            XCTAssertEqual(httpResponse.allHeaderFields["someKey"] as? String, "someValue")
+
+            return .allow
+        }
+
+        downloader.delegate = handler
+        downloader.downloadImage(with: url) { result in
+            XCTAssertNotNil(result.value)
+            XCTAssertNil(result.error)
+
+            self.downloader.delegate = nil
+            // hold delegate
+            _ = handler
+            exp.fulfill()
+        }
+        waitForExpectations(timeout: 3, handler: nil)
+    }
+
+    func testDownloaderReceiveResponseFailure() {
+        let exp = expectation(description: #function)
+
+        let url = testURLs[0]
+        stub(url, data: testImageData, headers: ["someKey": "someValue"])
+
+        let handler = TaskResponseCompletion()
+        let obj = NSObject()
+        handler.onReceiveResponse.delegate(on: obj) { (obj, response) in
+            guard let httpResponse = response as? HTTPURLResponse else {
+                XCTFail("Should be an HTTP response.")
+                return .cancel
+            }
+            XCTAssertEqual(httpResponse.statusCode, 200)
+            XCTAssertEqual(httpResponse.url, url)
+            XCTAssertEqual(httpResponse.allHeaderFields["someKey"] as? String, "someValue")
+
+            return .cancel
+        }
+
+        downloader.delegate = handler
+        downloader.downloadImage(with: url) { result in
+            XCTAssertNil(result.value)
+            XCTAssertNotNil(result.error)
+
+            if case .responseError(reason: .cancelledByDelegate) = result.error! {
+            } else {
+                XCTFail()
+            }
+
+            self.downloader.delegate = nil
+            // hold delegate
+            _ = handler
+            exp.fulfill()
+        }
+        waitForExpectations(timeout: 3, handler: nil)
+    }
 }
 
 class URLNilDataModifier: ImageDownloaderDelegate {
@@ -580,6 +652,19 @@ class TaskNilDataModifier: ImageDownloaderDelegate {
     }
 }
 
+class TaskResponseCompletion: ImageDownloaderDelegate {
+
+    let onReceiveResponse = Delegate<URLResponse, URLSession.ResponseDisposition>()
+
+    func imageDownloader(
+        _ downloader: ImageDownloader,
+        didReceive response: URLResponse,
+        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
+    ) {
+        completionHandler(onReceiveResponse.call(response)!)
+    }
+}
+
 class URLModifier: ImageDownloadRequestModifier {
     var url: URL? = nil
     func modified(for request: URLRequest) -> URLRequest? {

Tests/KingfisherTests/Utils/StubHelpers.swift

@@ -27,16 +27,29 @@
 import Foundation
 
 @discardableResult
-func stub(_ url: URL, data: Data, statusCode: Int = 200, length: Int? = nil) -> LSStubResponseDSL {
-    var stubResult = stubRequest("GET", url.absoluteString as NSString).andReturn(statusCode)?.withBody(data as NSData)
+func stub(_ url: URL,
+          data: Data,
+          statusCode: Int = 200,
+          length: Int? = nil,
+          headers: [String: String] = [:]
+) -> LSStubResponseDSL {
+    var stubResult = stubRequest("GET", url.absoluteString as NSString)
+        .andReturn(statusCode)?
+        .withHeaders(headers)?
+        .withBody(data as NSData)
     if let length = length {
         stubResult = stubResult?.withHeader("Content-Length", "\(length)")
     }
     return stubResult!
 }
 
-func delayedStub(_ url: URL, data: Data, statusCode: Int = 200, length: Int? = nil) -> LSStubResponseDSL {
-    let result = stub(url, data: data, statusCode: statusCode, length: length)
+func delayedStub(_ url: URL,
+                 data: Data,
+                 statusCode: Int = 200,
+                 length: Int? = nil,
+                 headers: [String: String] = [:]
+) -> LSStubResponseDSL {
+    let result = stub(url, data: data, statusCode: statusCode, length: length, headers: headers)
     return result.delay()!
 }