Fix for URLProtocolClient doesn't deliver response and data to URLSession delegates

Author: nethraravindran

Foundation/NSURLSession/NSURLSessionTask.swift

@@ -526,7 +526,20 @@ public let URLSessionDownloadTaskResumeData: String = "NSURLSessionDownloadTaskR
 extension _ProtocolClient : URLProtocolClient {
 
     func urlProtocol(_ protocol: URLProtocol, didReceive response: URLResponse, cacheStoragePolicy policy: URLCache.StoragePolicy) {
-        `protocol`.task?.response = response
+        guard let task = `protocol`.task else { fatalError("Received response, but there's no task.") }
+        task.response = response
+        let session = task.session as! URLSession
+        guard let dataTask = task as? URLSessionDataTask else { return }
+        switch session.behaviour(for: task) {
+        case .taskDelegate(let delegate as URLSessionDataDelegate):
+            session.delegateQueue.addOperation {
+                delegate.urlSession(session, dataTask: dataTask, didReceive: response, completionHandler: { _ in
+                    URLSession.printDebug("warning: Ignoring disposition from completion handler.")
+                })
+            }
+        case .noDelegate, .taskDelegate, .dataCompletionHandler, .downloadCompletionHandler:
+            break
+        }
     }
 
     func urlProtocolDidFinishLoading(_ protocol: URLProtocol) {
@@ -567,7 +580,17 @@ extension _ProtocolClient : URLProtocolClient {
 
     func urlProtocol(_ protocol: URLProtocol, didLoad data: Data) {
         `protocol`.properties[.responseData] = data
-        //TODO: this method needs to be extended to call the urlSession(_:dataTask:didReceive:)
+        guard let task = `protocol`.task else { fatalError() }
+        guard let session = task.session as? URLSession else { fatalError() }
+        switch session.behaviour(for: task) {
+        case .taskDelegate(let delegate):
+            let dataDelegate = delegate as? URLSessionDataDelegate
+            let dataTask = task as? URLSessionDataTask
+            session.delegateQueue.addOperation {
+                dataDelegate?.urlSession(session, dataTask: dataTask!, didReceive: data)
+            }
+        default: return
+        }
     }
 
     func urlProtocol(_ protocol: URLProtocol, didFailWithError error: Error) {

TestFoundation/TestNSURLSession.swift

@@ -38,6 +38,7 @@ class TestURLSession : XCTestCase {
             ("test_verifyHttpAdditionalHeaders", test_verifyHttpAdditionalHeaders),
             ("test_timeoutInterval", test_timeoutInterval),
             ("test_customProtocol", test_customProtocol),
+	    ("test_customProtocolResponseWithDelegate", test_customProtocolResponseWithDelegate),
             ("test_httpRedirection", test_httpRedirection),
             ("test_httpRedirectionTimeout", test_httpRedirectionTimeout),
         ]
@@ -344,6 +345,14 @@ class TestURLSession : XCTestCase {
         task.resume()
         waitForExpectations(timeout: 12)
     }
+
+    func test_customProtocolResponseWithDelegate() {
+        let url = URL(string: "http://127.0.0.1:\(TestURLSession.serverPort)/Peru")!
+        let d = DataTask(with: expectation(description: "Custom protocol with delegate"), protocolClasses: [CustomProtocol.self])
+        d.responseReceivedExpectation = expectation(description: "A response wasn't received")
+        d.run(with: url)
+        waitForExpectations(timeout: 12)
+    }
 
     func test_httpRedirection() {
         let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/UnitedStates"
@@ -387,16 +396,22 @@ class DataTask : NSObject {
     var session: URLSession! = nil
     var task: URLSessionDataTask! = nil
     var cancelExpectation: XCTestExpectation?
+    var responseReceivedExpectation: XCTestExpectation?
+    var protocols: [AnyClass]?
 
     public var error = false
 
-    init(with expectation: XCTestExpectation) {
+    init(with expectation: XCTestExpectation, protocolClasses: [AnyClass]? = nil) {
         dataTaskExpectation = expectation
+	protocols = protocolClasses
     }
 
     func run(with request: URLRequest) {
         let config = URLSessionConfiguration.default
         config.timeoutIntervalForRequest = 8
+        if let customProtocols = protocols {
+            config.protocolClasses = customProtocols
+        }
         session = URLSession(configuration: config, delegate: self, delegateQueue: nil)
         task = session.dataTask(with: request)
         task.resume()
@@ -405,6 +420,9 @@ class DataTask : NSObject {
     func run(with url: URL) {
         let config = URLSessionConfiguration.default
         config.timeoutIntervalForRequest = 8
+        if let customProtocols = protocols {
+            config.protocolClasses = customProtocols
+        }
         session = URLSession(configuration: config, delegate: self, delegateQueue: nil)
         task = session.dataTask(with: url)
         task.resume()
@@ -419,6 +437,14 @@ extension DataTask : URLSessionDataDelegate {
     public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
         capital = String(data: data, encoding: String.Encoding.utf8)!
     }
+
+    public func urlSession(_ session: URLSession,
+                    dataTask: URLSessionDataTask,
+                    didReceive response: URLResponse,
+                    completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
+        guard responseReceivedExpectation != nil else { return }
+        responseReceivedExpectation!.fulfill()
+    }
 }
 
 extension DataTask : URLSessionTaskDelegate {