URLProtocolでURLSessionの通信をハンドリングする

概要

例えば、通信などが起きたときに、そのログを残しておいたり、通信量を調べたい時があると思います。
そんな通信のハンドリングを行うために、URLProtocolを利用します。

通信をURLProtocol継承クラスに処理させる

通信をハンドリングするには、ハンドリングするためのURLProtocol継承クラスを用意します。ここではMyURLProtocolという名前にしました。
URLProtocol継承クラスの実装に関しては後ほど説明します。

URLProtocol継承クラスを用意したら、各種通信がURLProtocol継承クラスを通して行われるようにします。

URLSession.sharedの場合

もしURLSession.sharedしか使われていないのであれば、URLProtocol.registerClass を使うことで全てのURLSession.sharedの通信がURLProtocol継承クラスを通してハンドリングされます。

URLProtocol.registerClass(MyURLProtocol.self)

個別のURLSessionの場合

個別のURLSessionを作成して通信を行なっている場合、使用するURLSessionConfigurationの protocolClasses にURLProtocol継承クラスを登録するようにします。

let sessionConfiguration = URLSessionConfiguration.default
sessionConfiguration.protocolClasses = [MyURLProtocol.self]
let session = URLSession(configuration: sessionConfiguration)

URLProtocolをオーバーライドしたクラスを用意する

各種通信をハンドリングするURLProtocol継承クラスを用意します。
通信が走るたびに、設定したURLProtocol継承クラスがURLSessionの通信をハンドリングします。

今回の目的はあくまで通信機能はそのまま動作させるようにするので、内部で実際の通信を行うURLSessionを持ち、各種URLProtocolのメソッドをオーバーライドしてstartLoadingの時には通信開始、stopLoadingの時はキャンセルするようにします。

class MyURLProtocol: URLProtocol {
    private var sessionTask: URLSessionDataTask?
    
    private lazy var session: URLSession = { [unowned self] in
        return URLSession(configuration: .default, delegate: self, delegateQueue: nil)
    }()
    
    override class func canInit(with request: URLRequest) -> Bool {
        true
    }
    
    override class func canInit(with task: URLSessionTask) -> Bool {
        task.currentRequest != nil
    }
    
    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        request
    }
    
    override func startLoading() {
        sessionTask = session.dataTask(with: request)
        sessionTask?.resume()
    }
    
    override func stopLoading() {
        sessionTask?.cancel()
        session.finishTasksAndInvalidate()
    }
}

URLSessionDataDelegateの実装

URLSessionDataDelegateの実装では、各種イベントが起きた時にclientにイベントを通知します。

extension MyURLProtocol: URLSessionDataDelegate {
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        client?.urlProtocol(self, didLoad: data)
    }
    
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse, completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
        let policy = URLCache.StoragePolicy(rawValue: request.cachePolicy.rawValue) ?? .notAllowed
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: policy)
        completionHandler(.allow)
    }
    
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if let error = error {
            client?.urlProtocol(self, didFailWithError: error)
        } else {
            client?.urlProtocolDidFinishLoading(self)
        }
    }
    
    func urlSession(_ session: URLSession, task: URLSessionTask, willPerformHTTPRedirection response: HTTPURLResponse, newRequest request: URLRequest, completionHandler: @escaping (URLRequest?) -> Void) {
        client?.urlProtocol(self, wasRedirectedTo: request, redirectResponse: response)
        completionHandler(request)
    }
    
    func urlSession(_ session: URLSession, didBecomeInvalidWithError error: Error?) {
        guard let error = error else { return }
        client?.urlProtocol(self, didFailWithError: error)
    }
    
    func urlSessionDidFinishEvents(forBackgroundURLSession session: URLSession) {
        client?.urlProtocolDidFinishLoading(self)
    }
}

まとめ

URLProtocolを使ってのURLSessionのハンドリングについて説明しました。
各種URLSessionDataDelegateのfuncに各自必要な処理を書くことで、通信量の測定やログの書き込みなどができるかと思います。