Skip to content

Instantly share code, notes, and snippets.

@buranmert
Last active October 26, 2022 06:36
Show Gist options
  • Save buranmert/61eac46e8bad3d0e7c71b7d1ba4fa524 to your computer and use it in GitHub Desktop.
Save buranmert/61eac46e8bad3d0e7c71b7d1ba4fa524 to your computer and use it in GitHub Desktop.
Adding hooks to URLSession
/*
* Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
* This product includes software developed at Datadog (https://www.datadoghq.com/).
* Copyright 2019-2020 Datadog, Inc.
*/
import Foundation
public extension URLSession {
internal typealias RequestInterceptor = HookedSession.RequestInterceptor
internal typealias TaskObserver = HookedSession.TaskObserver
static func tracedSession(
configuration: URLSessionConfiguration = .default,
delegate: URLSessionDelegate? = nil,
delegateQueue: OperationQueue? = nil
) -> URLSession {
let session = URLSession(
configuration: configuration,
delegate: delegate,
delegateQueue: delegateQueue
)
let requestInterceptor: RequestInterceptor = {
// TODO
return $0
}
let taskObserver: TaskObserver = { _, _ in
// TODO
}
let hookedSession = HookedSession(
session: session,
requestInterceptor: requestInterceptor,
taskObserver: taskObserver
)
return hookedSession.asURLSession()
}
}
/*
HookedSession is a NSObject subclass.
It keeps an URLSession instance inside.
It relays all the method calls to URLSession,
except those that are implemented by HookedSession.
It intercepts and observes requests and tasks respectively.
*/
internal final class HookedSession: NSObject {
typealias RequestInterceptor = (URLRequest) -> URLRequest
typealias TaskObserver = (URLSessionTask, URLSessionTask.State?) -> Void
private let session: URLSession
let requestInterceptor: RequestInterceptor
let taskObserver: TaskObserver
private var observations = [Int: NSKeyValueObservation](minimumCapacity: 100)
init(session: URLSession,
requestInterceptor: @escaping RequestInterceptor,
taskObserver: @escaping TaskObserver) {
self.session = session
self.requestInterceptor = requestInterceptor
self.taskObserver = taskObserver
super.init()
}
func asURLSession() -> URLSession {
let castedSession: URLSession = unsafeBitCast(self, to: URLSession.self)
return castedSession
}
// MARK: - Transparent messaging
/*
As a NSObject subclass yet exposed as URLSession (ref: URLSession.tracedSession)
all the method calls are `unrecognized_selector` for HookedSession, except
those which are implemented in HookedSession.
forwardingTarget passes those unimplemented method calls to session
*/
override func forwardingTarget(for aSelector: Selector!) -> Any? { // swiftlint:disable:this implicitly_unwrapped_optional
return session
}
// MARK: - URLSessionDataTask
/*
IMPORTANT NOTE:
@objc enables dynamic method dispatch and
make sure @objc names match NSURLSession Obj-C interface.
Otherwise, these methods are not called.
*/
@objc(dataTaskWithURL:)
func dataTask(with url: URL) -> URLSessionDataTask {
return dataTask(with: URLRequest(url: url))
}
@objc(dataTaskWithURL:completionHandler:)
func dataTask(
with url: URL,
completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void
) -> URLSessionDataTask {
return dataTask(with: URLRequest(url: url), completionHandler: completionHandler)
}
@objc(dataTaskWithRequest:)
func dataTask(with request: URLRequest) -> URLSessionDataTask {
return observed(session.dataTask(with: requestInterceptor(request)))
}
@objc(dataTaskWithRequest:completionHandler:)
func dataTask(
with request: URLRequest,
completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void
) -> URLSessionDataTask {
let task = session.dataTask(with: requestInterceptor(request),
completionHandler: completionHandler)
return observed(task)
}
// MARK: - Helpers
/*
IMPORTANT NOTE:
If you create an URLSessionTask instance from an URLSession instance,
the task stays alive EVEN IF you nullify the session instance.
This happens because URLSessionTask has private `__taskGroup` property which keeps
session instance alive as long as the task is alive.
This is not the case for HookedSession instances.
let task = hookedSession.dataTask(...)
hookedSession = nil
task.resume()
In the case above, task and
hookedSession.session (private property) will stay alive
yet hookedSession will be deallocated.
That means observations will be deallocated too.
In order to keep observing until the task completes,
observationBlock below captures `self` on purpose.
Therefore, observationBlock keeps `self` alive
until the block is removed from self.observations dict.
*/
private func observed<T: URLSessionTask>(_ task: T) -> T {
let observer = taskObserver
var previousState: URLSessionTask.State? = nil
let observation = task.observe(\.state, options: [.initial]) { observed, _ in
observer(observed, previousState)
previousState = observed.state
switch observed.state {
case .canceling, .completed:
self.observations[observed.taskIdentifier] = nil
default:
break
}
}
observations[task.taskIdentifier] = observation
return task
}
}
extension HookedSession: CustomReflectable {
/*
HookedSession imitates URLSession from outside in case that anyone does type-check
such as isKindOf:/isMemberOf: or checks superclass at runtime
*/
var customMirror: Mirror {
return Mirror(reflecting: session)
}
override func isKind(of aClass: AnyClass) -> Bool {
return session.isKind(of: aClass)
}
override class func isKind(of aClass: AnyClass) -> Bool {
return URLSession.isKind(of: aClass)
}
override func isMember(of aClass: AnyClass) -> Bool {
return session.isMember(of: aClass)
}
override class func isMember(of aClass: AnyClass) -> Bool {
return URLSession.isMember(of: aClass)
}
override var superclass: AnyClass? { return session.superclass }
override class func superclass() -> AnyClass? { return URLSession.superclass() }
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment