Network Autoinstrumentation

NetworkAutoinstrumentation.swift

import Foundation
import ObjectiveC
import PlaygroundSupport
PlaygroundPage.current.needsIndefiniteExecution = true

func swizzleInit() {
    let selector = #selector(URLSession.init(configuration:delegate:delegateQueue:))
    let method = class_getClassMethod(URLSession.self, selector)!

    typealias ImpBlock = @convention(block) (AnyClass, URLSessionConfiguration, URLSessionDelegate?, OperationQueue?) -> URLSession
    typealias Imp = @convention(c) (AnyClass, Selector, URLSessionConfiguration, URLSessionDelegate?, OperationQueue?) -> URLSession

    let origImp = unsafeBitCast(method_getImplementation(method), to: Imp.self)
    let newImpBlock: ImpBlock = { sessionClass, config, delegate, opQueue in
        return origImp(sessionClass, selector, config, swizzling(delegate), opQueue)
    }
    let newImp = imp_implementationWithBlock(newImpBlock)
    method_setImplementation(method, newImp)
}

func swizzling(_ delegate: URLSessionDelegate?) -> URLSessionDelegate {
    if delegate?.isKind(of: DefaultInterceptor.self) == true {
        return delegate!
    }
    guard let someObj = delegate,
          let baseClass = object_getClass(someObj) else {
        return DefaultInterceptor()
    }
    let newClass: AnyClass = delegateClass(basedOn: baseClass)
    object_setClass(someObj, newClass)
    return someObj
}

func delegateClass(basedOn superclass: AnyClass) -> AnyClass {
    typealias ImpBlock = @convention(block) (URLSessionDelegate, URLSession, URLSessionTask, Error?) -> Void
    typealias Imp = @convention(c) (URLSessionDelegate, Selector, URLSession, URLSessionTask, Error?) -> Void

    let selector = #selector(URLSessionTaskDelegate.urlSession(_:task:didCompleteWithError:))

    let defaultMethod = class_getInstanceMethod(DefaultInterceptor.self, selector)!
    let defaultImp = method_getImplementation(defaultMethod)
    let callableDefaultImp = unsafeBitCast(defaultImp, to: Imp.self)
    let defaultEncoding = method_getTypeEncoding(defaultMethod)

    let newClass: AnyClass = objc_allocateClassPair(superclass, "DynamicDelegate", 0)!

    if let origMethod = class_getInstanceMethod(newClass, selector) {
        let origImp = method_getImplementation(origMethod)
        let callableOrigImp = unsafeBitCast(origImp, to: Imp.self)
        let newImpBlock: ImpBlock = { delegate, session, task, error in
            callableDefaultImp(delegate, selector, session, task, error)
            callableOrigImp(delegate, selector, session, task, error)
        }
        let newImp = imp_implementationWithBlock(newImpBlock)
        class_addMethod(newClass, selector, newImp, defaultEncoding)
    } else {
        class_addMethod(newClass, selector, defaultImp, defaultEncoding)
    }

    objc_registerClassPair(newClass)
    return newClass
}

class DefaultInterceptor: NSObject, URLSessionTaskDelegate {
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        print("default delegate")
    }
}

class SubDefaultDelegate: DefaultInterceptor {
    override func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        print("sub-default delegate")
        super.urlSession(session, task: task, didCompleteWithError: error)
    }
}

class MyCustomDelegate: NSObject, URLSessionTaskDelegate {
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        print("custom delegate")
    }
}

swizzleInit()

let url = URL(string: "https://news.ycombinator.com")!

let nilDelegateSession = URLSession(configuration: .default, delegate: nil, delegateQueue: nil)
let task1 = nilDelegateSession.dataTask(with: url)
task1.resume()
/* prints:
default delegate
 */

//let customDelegateSession = URLSession(configuration: .default, delegate: MyCustomDelegate(), delegateQueue: nil)
//let task2 = customDelegateSession.dataTask(with: url)
//task2.resume()
/* prints:
default delegate
custom delegate
 */

//let subDefaultDelegateSession = URLSession(configuration: .default, delegate: SubDefaultDelegate(), delegateQueue: nil)
//let task3 = subDefaultDelegateSession.dataTask(with: url)
//task3.resume()
/* prints:
sub-default delegate
default delegate
 */