Skip to content

Instantly share code, notes, and snippets.

@kean
Created December 31, 2022 01:29
Show Gist options
  • Save kean/3154a5bde8e0c5e9dc3322f21ba86757 to your computer and use it in GitHub Desktop.
Save kean/3154a5bde8e0c5e9dc3322f21ba86757 to your computer and use it in GitHub Desktop.
// The MIT License (MIT)
//
// Copyright (c) 2020–2022 Alexander Grebenyuk (github.com/kean).
import Pulse
import Foundation
import ObjectiveC.runtime
#if DEBUG
extension Experimental {
static var networkLogger = NetworkLogger()
/// Swizzles `URLSession` APIs to track _all_ tasks, even the ones created
/// using the Async/Await and completion-based APIs.
///
/// - warning: This is sample code that is _not_ designed to be used in
/// production because it contains some private APIs. Do _not_ in your
/// production builds. It's only for testing purposes.
static func swizzleURLSession() {
swizzleInit()
swizzleReceiveData()
swizzleFinishWithError()
}
private static func swizzleReceiveData() {
let selector = NSSelectorFromString("_didReceiveData:")
let klass: AnyClass = NSClassFromString("__NSCFURLLocalSessionConnection")!
let originalImp = class_getMethodImplementation(klass, selector)
typealias IMPType = @convention(c) (AnyObject, Selector, AnyObject) -> Void
let originalImpCallable = unsafeBitCast(originalImp, to: IMPType.self)
let block: @convention(block) (AnyObject, AnyObject) -> Void = {
if let this = $0 as? NSObject, let task = this.value(forKey: "task") as? URLSessionDataTask, let data = $1 as? Data {
Experimental.networkLogger.logDataTask(task, didReceive: data)
}
originalImpCallable($0, selector, $1)
}
setNewIMPWithBlock(block, forSelector: selector, toClass: klass)
}
private static func swizzleInit() {
let selector = NSSelectorFromString("initWithTask:delegate:delegateQueue:")
let klass: AnyClass = NSClassFromString("__NSCFURLLocalSessionConnection")!
let originalImp = class_getMethodImplementation(klass, selector)
typealias IMPType = @convention(c) (AnyObject, Selector, AnyObject, AnyObject, AnyObject) -> AnyObject
let originalImpCallable = unsafeBitCast(originalImp, to: IMPType.self)
let block: @convention(block) (AnyObject, AnyObject, AnyObject, AnyObject) -> AnyObject = {
if let task = $1 as? URLSessionTask {
Experimental.networkLogger.logTaskCreated(task)
}
return originalImpCallable($0, selector, $1, $2, $3)
}
setNewIMPWithBlock(block, forSelector: selector, toClass: klass)
}
private static func swizzleFinishWithError() {
let selector = NSSelectorFromString("_didFinishWithError:")
let klass: AnyClass = NSClassFromString("__NSCFURLLocalSessionConnection")!
let originalImp = class_getMethodImplementation(klass, selector)
typealias IMPType = @convention(c) (AnyObject, Selector, AnyObject?) -> Void
let originalImpCallable = unsafeBitCast(originalImp, to: IMPType.self)
let block: @convention(block) (AnyObject, AnyObject?) -> Void = {
if let this = $0 as? NSObject,
let task = this.value(forKey: "task") as? URLSessionTask,
let metrics = task.value(forKey: "_incompleteTaskMetrics") as? URLSessionTaskMetrics {
Experimental.networkLogger.logTask(task, didFinishCollecting: metrics)
}
if let this = $0 as? NSObject, let task = this.value(forKey: "task") as? URLSessionTask, let error = $1 as? (Error?) {
Experimental.networkLogger.logTask(task, didCompleteWithError: error)
}
originalImpCallable($0, selector, $1)
}
setNewIMPWithBlock(block, forSelector: selector, toClass: klass)
}
private static func setNewIMPWithBlock<T>(_ block: T, forSelector selector: Selector, toClass klass: AnyClass) {
let method = class_getInstanceMethod(klass, selector)
let imp = imp_implementationWithBlock(unsafeBitCast(block, to: AnyObject.self))
if !class_addMethod(klass, selector, imp, method_getTypeEncoding(method!)) {
method_setImplementation(method!, imp)
}
}
}
#endif
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment