Custom URLProtocol with URLSession

URLProtocol + URLSession.swift

//
//  MyURLProtocol.swift
//  NSURLProtocolExample
//
//  Created by Chung Tran on 26/08/2017.
//  Copyright © 2017 Zedenem. All rights reserved.
//
import UIKit
var requestCount = 0

class MyURLProtocol: URLProtocol, URLSessionDataDelegate, URLSessionTaskDelegate {
    var dataTask: URLSessionDataTask!
    var receivedData: NSMutableData!
    var urlResponse: URLResponse!
    class var HandledKey: String {
        return "MyURLProtocolHandledKey"
    }
    
    override class func canInit(with request: URLRequest) -> Bool {
        if URLProtocol.property(forKey: HandledKey, in: request) != nil {
            return false
        }
        requestCount += 1
        print("Request #\(requestCount). URL: \(request.url!.absoluteString)")
        
        return true
    }
    
    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        return request
    }
    
    override class func requestIsCacheEquivalent(_ a: URLRequest, to b: URLRequest) -> Bool {
        return super.requestIsCacheEquivalent(a, to: b)
    }
    
    override func startLoading() {
        let newRequest = NSMutableURLRequest(url: self.request.url!, cachePolicy: .useProtocolCachePolicy, timeoutInterval: 240.0)
        URLProtocol.setProperty(true, forKey: MyURLProtocol.HandledKey, in: newRequest)
        let session = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
        dataTask = session.dataTask(with: self.request)
        dataTask.resume()
    }
    
    override func stopLoading() {
        dataTask.cancel()
        dataTask = nil
    }
    
    // MARK: - URLSessionDataDelegate
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        client?.urlProtocol(self, didLoad: data)
        receivedData.append(data)
    }
    
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse, completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
        urlResponse = response
        receivedData = NSMutableData()
        completionHandler(.allow)
    }
    
    // MARK: - URLSessionTaskDelegate
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if let error = error {
            client?.urlProtocol(self, didFailWithError: error)
            return
        }
        // Do what ever you want like saveCachedResponse()
        client?.urlProtocolDidFinishLoading(self)
    }
}