NSURLConnection Swizzling in Swift

NSURLConnectionInterceptor.swift

extension NSURLConnection {
  
    @objc
    static func setupURLConnectionConfiguration() {
      
      guard self == NSURLConnection.self else {
        return
      }
      
      let originalRequestSelector = #selector(self.init(request:delegate:startImmediately:))
      let swizzledRequestSelector = #selector(self.my_init(request:delegate:startImmediately:))
      
      let originalRequestMethod = class_getInstanceMethod(self, originalRequestSelector)
      let swizzledRequestMethod = class_getInstanceMethod(self, swizzledRequestSelector)
      
      if originalRequestMethod == nil || swizzledRequestMethod == nil {
        return
      }
      method_exchangeImplementations(originalRequestMethod!, swizzledRequestMethod!)
      
      let originalRequestImmediatelySelector = #selector(self.init(request:delegate:))
      let sizzledRequestImmediatelySelector = #selector(self.my_initWithOutstartImmediately(request:delegate:))
      
      let originalRequestImmediatelyMethod = class_getInstanceMethod(self, originalRequestImmediatelySelector)
      let swizzledRequestImmediatelyMethod = class_getInstanceMethod(self, sizzledRequestImmediatelySelector)
      
      if originalRequestImmediatelyMethod == nil || swizzledRequestImmediatelyMethod == nil {
        return
      }
      method_exchangeImplementations(originalRequestImmediatelyMethod!, swizzledRequestImmediatelyMethod!)
    }
  
    @objc func my_init(request: NSURLRequest, delegate: NSURLConnectionDelegate?, startImmediately: Bool) -> NSURLConnection? {
      let inspectedDelegate = TempURLConnectionProtocol(actualDelegate: delegate)
      return my_init(request: request, delegate: inspectedDelegate, startImmediately: startImmediately)
    }
  
    @objc func my_initWithOutstartImmediately(request: NSURLRequest, delegate: NSURLConnectionDelegate?) -> NSURLConnection? {
      let inspectedDelegate = TempURLConnectionProtocol(actualDelegate: delegate)
      return my_initWithOutstartImmediately(request: request, delegate: inspectedDelegate)
    }
}


let didReceiveResponseSelector: Selector = #selector((NSURLConnectionDataDelegate.connection(_:didReceive:)) as ((NSURLConnectionDataDelegate) -> (NSURLConnection, URLResponse) -> Void)?)
let didReceiveDataSelector: Selector = #selector((NSURLConnectionDataDelegate.connection(_:didReceive:)) as ((NSURLConnectionDataDelegate) -> (NSURLConnection, Data) -> Void)?)

class TempURLConnectionProtocol: NSObject, NSURLConnectionDelegate, NSURLConnectionDataDelegate {
  
  var received: Data?
  var actualDelegate: NSURLConnectionDelegate?
  var currentRequest: TempDataAccessLayer?
  var response: URLResponse?
  
  init(actualDelegate actual: NSURLConnectionDelegate?) {
    super.init()
    
    received = Data()
    received?.count = 0
    actualDelegate = actual
    currentRequest = nil
    response = nil
  }
  
  public func connection(_ connection: NSURLConnection, didFailWithError error: Error) {
    if actualDelegate!.responds(to: #selector(self.connection(_:didFailWithError:))) {
      self.actualDelegate?.connection!(connection, didFailWithError: error)
    }
    self.cleanUp(connection: connection)
  }
  
  public func connection(_ connection: NSURLConnection, didReceive aResponse: URLResponse) {
    self.response = aResponse
    currentRequest?.initResponse(response: response!)
    if actualDelegate!.responds(to: didReceiveResponseSelector) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        actual.connection?(connection, didReceive: response!)
      }
    }
  }
  
  public func connection(_ connection: NSURLConnection, didReceive data: Data) {
    received?.append(data)
    if actualDelegate!.responds(to: didReceiveDataSelector) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        actual.connection?(connection, didReceive: data)
      }
    }
  }
  
  func connection(_ connection: NSURLConnection, didSendBodyData bytesWritten: Int, totalBytesWritten: Int, totalBytesExpectedToWrite: Int) {
    
    if actualDelegate!.responds(to: #selector(self.connection(_:didSendBodyData:totalBytesWritten:totalBytesExpectedToWrite:))) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        actual.connection?(connection, didSendBodyData: bytesWritten, totalBytesWritten: totalBytesWritten, totalBytesExpectedToWrite: totalBytesExpectedToWrite)
      }
    }
  }
  
  public func connectionDidFinishLoading(_ connection: NSURLConnection) {
    if self.actualDelegate!.responds(to: #selector(self.connectionDidFinishLoading(_:))) {
      if let actual = self.actualDelegate as? NSURLConnectionDataDelegate {
        actual.connectionDidFinishLoading?(connection)
      }
    }
        
    currentRequest?.httpBody = body(from: connection.currentRequest)
    if let startDate = currentRequest?.date {
      currentRequest?.duration = Double(Date().timeIntervalSince(startDate)) * 1000
    }
    //Called event capture method
    DispatchQueue.main.async {    
	self.cleanUp(connection: connection)
    }
  }
  
  public func connection(_ connection: NSURLConnection, willSend request: URLRequest, redirectResponse response: URLResponse?) -> URLRequest? {
    
    self.response = response
    let newRequest = ((request as NSURLRequest).mutableCopy() as? NSMutableURLRequest)!
    currentRequest = TempDataAccessLayer(request: newRequest)
    if actualDelegate!.responds(to: #selector(connection(_:willSend:redirectResponse:))) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        return actual.connection!(connection, willSend: request, redirectResponse: response)
      }
    }
    return request
  }
  
  func connection(_ connection: NSURLConnection, needNewBodyStream request: URLRequest) -> InputStream? {
    if actualDelegate!.responds(to: #selector(self.connection(_:needNewBodyStream:))) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        return actual.connection?(connection, needNewBodyStream: request)
      }
    }
    return nil
  }
  
  func connection(_ connection: NSURLConnection, willCacheResponse cachedResponse: CachedURLResponse) -> CachedURLResponse? {
    if actualDelegate!.responds(to: #selector(self.connection(_:willCacheResponse:))) {
      if let actual = actualDelegate as? NSURLConnectionDataDelegate {
        return actual.connection?(connection, willCacheResponse: cachedResponse)
      }
    }
    return cachedResponse
  }
    
  private func body(from request: URLRequest) -> Data? {
    return request.httpBody ?? request.httpBodyStream.flatMap { stream in
      let data = NSMutableData()
      stream.open()
      while stream.hasBytesAvailable {
        var buffer = [UInt8](repeating: 0, count: 10 * 1024 * 1024)
        let length = stream.read(&buffer, maxLength: buffer.count)
        data.append(buffer, length: length)
      }
      stream.close()
      return data as Data
    }
  }
  
  func cleanUp(connection: NSURLConnection) {
    self.received = nil
    self.actualDelegate = nil
    self.currentRequest = nil
    self.response = nil
  }
}

class TempDataAccessLayer: Codable {

    let url: String
    let date: Date
    let method: String
    var headers: [String: String]?
    var httpBody: Data?
    var stausCode: Int
    var responseHeaders: [String: String]?
    var dataResponse: Data?
    var errorDescription: String?
    var duration: Double?
    
    init(request: NSURLRequest) {
        url = request.url?.absoluteString ?? ""
        date = Date()
        method = request.httpMethod ?? UtilityDefaults.unknown.rawValue
        headers = request.allHTTPHeaderFields
        headers?.removeValue(forKey: NetworkAuthorization.removeAuthorizationHeader.rawValue)
        httpBody = request.httpBody
        stausCode = 0
    }
    
    func initResponse(response: URLResponse) {
        guard let responseHttp = response as? HTTPURLResponse else {return}
        stausCode = responseHttp.statusCode
        responseHeaders = responseHttp.allHeaderFields as? [String: String]
    }
}