Skip to content

Commit ca9e74d

Browse files
keankean
committed
Simplify how URLSessionSwizzler gets enabled
1 parent 380f222 commit ca9e74d

File tree

1 file changed

+101
-96
lines changed

1 file changed

+101
-96
lines changed

‎Sources/PulseProxy/URLSessionSwizzler.swift

+101-96
Original file line numberDiff line numberDiff line change
@@ -15,126 +15,131 @@ extension NetworkLogger {
1515
///
1616
/// - parameter logger: The network logger to be used for recording the requests. By default, uses shared logger.
1717
public static func enableProxy(logger: NetworkLogger? = nil) {
18-
guard Thread.isMainThread else {
19-
return DispatchQueue.main.async { NetworkLogger.URLSessionSwizzler.enable(logger: logger) }
20-
}
21-
MainActor.assumeIsolated {
22-
NetworkLogger.URLSessionSwizzler.enable(logger: logger)
23-
}
18+
URLSessionSwizzler.enable(logger: logger)
2419
}
2520
}
2621

27-
extension NetworkLogger {
28-
@MainActor
29-
final class URLSessionSwizzler {
30-
static var shared: URLSessionSwizzler?
22+
final class URLSessionSwizzler {
23+
static var shared: URLSessionSwizzler?
24+
25+
private var logger: NetworkLogger { _logger ?? .shared }
26+
private let _logger: NetworkLogger?
27+
28+
init(logger: NetworkLogger?) {
29+
self._logger = logger
30+
}
3131

32-
private var logger: NetworkLogger { _logger ?? .shared }
33-
private let _logger: NetworkLogger?
32+
static let lock = NSLock()
33+
static var isEnabled = false
3434

35-
init(logger: NetworkLogger?) {
36-
self._logger = logger
35+
static func enable(logger: NetworkLogger?) {
36+
lock.lock()
37+
if isEnabled {
38+
lock.unlock()
39+
NSLog("Error: Pulse proxy is already enabled")
40+
return
3741
}
42+
isEnabled = true
43+
lock.unlock()
3844

39-
@MainActor
40-
static func enable(logger: NetworkLogger?) {
41-
guard NetworkLogger.URLSessionSwizzler.shared == nil else {
42-
NSLog("Error: Pulse.URLSessionProxy already enabled")
43-
return
44-
}
45+
let proxy = URLSessionSwizzler(logger: logger)
46+
proxy.enable()
47+
URLSessionSwizzler.shared = proxy
48+
}
4549

46-
let proxy = URLSessionSwizzler(logger: logger)
47-
proxy.enable()
48-
URLSessionSwizzler.shared = proxy
50+
func enable() {
51+
swizzleURLSessionTaskResume()
52+
// "__NSCFURLLocalSessionConnection"
53+
if let sessionClass = NSClassFromString(["__", "NS", "CFURL", "Local", "Session", "Connection"].joined()) {
54+
swizzleDataTaskDidReceiveData(baseClass: sessionClass)
55+
swizzleDataDataDidCompleteWithError(baseClass: sessionClass)
56+
} else {
57+
NSLog("Pulse.URLSessionSwizzler failed to initialize. Please report at https://github.com/kean/Pulse/issues.")
4958
}
59+
}
5060

51-
func enable() {
52-
swizzleURLSessionTaskResume()
53-
// "__NSCFURLLocalSessionConnection"
54-
if let sessionClass = NSClassFromString(["__", "NS", "CFURL", "Local", "Session", "Connection"].joined()) {
55-
swizzleDataTaskDidReceiveData(baseClass: sessionClass)
56-
swizzleDataDataDidCompleteWithError(baseClass: sessionClass)
57-
} else {
58-
NSLog("Pulse.URLSessionSwizzler failed to initialize. Please report at https://github.com/kean/Pulse/issues.")
59-
}
61+
// - `resume` (optional)
62+
private func swizzleURLSessionTaskResume() {
63+
var methods = [Method]()
64+
if let method = class_getInstanceMethod(URLSessionTask.self, #selector(URLSessionTask.resume)) {
65+
methods.append(method)
66+
}
67+
// "__NSCFURLSessionTask"
68+
if let sessionTaskClass = NSClassFromString(["__", "NS", "CFURL", "Session", "Task"].joined()),
69+
let method = class_getInstanceMethod(sessionTaskClass, NSSelectorFromString("resume")) {
70+
methods.append(method)
6071
}
72+
methods.forEach {
73+
let method = $0
74+
var originalImplementation: IMP?
75+
let block: @convention(block) (URLSessionTask) -> Void = { [weak self] task in
76+
self?.logger.logTaskCreated(task)
6177

62-
// - `resume` (optional)
63-
private func swizzleURLSessionTaskResume() {
64-
var methods = [Method]()
65-
if let method = class_getInstanceMethod(URLSessionTask.self, #selector(URLSessionTask.resume)) {
66-
methods.append(method)
67-
}
68-
// "__NSCFURLSessionTask"
69-
if let sessionTaskClass = NSClassFromString(["__", "NS", "CFURL", "Session", "Task"].joined()),
70-
let method = class_getInstanceMethod(sessionTaskClass, NSSelectorFromString("resume")) {
71-
methods.append(method)
72-
}
73-
methods.forEach {
74-
let method = $0
75-
var originalImplementation: IMP?
76-
let block: @convention(block) (URLSessionTask) -> Void = { [weak self] task in
77-
self?.logger.logTaskCreated(task)
78-
79-
guard task.currentRequest != nil else { return }
80-
let key = String(method.hashValue)
81-
objc_setAssociatedObject(task, key, true, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
82-
let castedIMP = unsafeBitCast(originalImplementation, to: (@convention(c) (Any) -> Void).self)
83-
castedIMP(task)
84-
objc_setAssociatedObject(task, key, nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
85-
}
86-
let swizzledIMP = imp_implementationWithBlock(unsafeBitCast(block, to: AnyObject.self))
87-
originalImplementation = method_setImplementation(method, swizzledIMP)
78+
guard task.currentRequest != nil else { return }
79+
let key = String(method.hashValue)
80+
objc_setAssociatedObject(task, key, true, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
81+
let castedIMP = unsafeBitCast(originalImplementation, to: (@convention(c) (Any) -> Void).self)
82+
castedIMP(task)
83+
objc_setAssociatedObject(task, key, nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
8884
}
85+
let swizzledIMP = imp_implementationWithBlock(unsafeBitCast(block, to: AnyObject.self))
86+
originalImplementation = method_setImplementation(method, swizzledIMP)
87+
}
88+
}
89+
90+
// - `urlSession(_:task:didCompleteWithError:)`
91+
func swizzleDataDataDidCompleteWithError(baseClass: AnyClass) {
92+
// "_didFinishWithError:"
93+
let selector = NSSelectorFromString(["_", "didFinish", "With", "Error", ":"].joined())
94+
guard let method = class_getInstanceMethod(baseClass, selector),
95+
baseClass.instancesRespond(to: selector) else {
96+
return
8997
}
98+
typealias MethodSignature = @convention(c) (AnyObject, Selector, AnyObject?) -> Void
99+
let originalImp: IMP = method_getImplementation(method)
100+
let closure: @convention(block) (AnyObject, AnyObject?) -> Void = { [weak self] object, error in
101+
let original: MethodSignature = unsafeBitCast(originalImp, to: MethodSignature.self)
102+
original(object, selector, error)
90103

91-
// - `urlSession(_:task:didCompleteWithError:)`
92-
func swizzleDataDataDidCompleteWithError(baseClass: AnyClass) {
93-
// "_didFinishWithError:"
94-
let selector = NSSelectorFromString(["_", "didFinish", "With", "Error", ":"].joined())
95-
guard let method = class_getInstanceMethod(baseClass, selector),
96-
baseClass.instancesRespond(to: selector) else {
97-
return
98-
}
99-
typealias MethodSignature = @convention(c) (AnyObject, Selector, AnyObject?) -> Void
100-
let originalImp: IMP = method_getImplementation(method)
101-
let closure: @convention(block) (AnyObject, AnyObject?) -> Void = { [weak self] object, error in
102-
let original: MethodSignature = unsafeBitCast(originalImp, to: MethodSignature.self)
103-
original(object, selector, error)
104-
105-
if let task = object.value(forKey: "task") as? URLSessionTask {
106-
// "_incompleteTaskMetrics"
107-
if let metrics = task.value(forKey: ["_", "incomplete", "Task", "Metrics"].joined()) as? URLSessionTaskMetrics {
108-
self?.logger.logTask(task, didFinishCollecting: metrics)
104+
if let task = object.value(forKey: "task") as? URLSessionTask {
105+
// "_incompleteTaskMetrics"
106+
if let metrics = task.value(forKey: ["_", "incomplete", "Task", "Metrics"].joined()) as? URLSessionTaskMetrics {
107+
self?.logger.logTask(task, didFinishCollecting: metrics)
108+
}
109+
if var error = error as? NSError {
110+
if error.domain == "kCFErrorDomainCFNetwork" {
111+
// Satifsy LogggerStore (needs refactoring)
112+
error = NSError(domain: URLError.errorDomain, code: error.code, userInfo: error.userInfo)
109113
}
110-
let error = error as? Error
111114
self?.logger.logTask(task, didCompleteWithError: error)
115+
} else {
116+
self?.logger.logTask(task, didCompleteWithError: error as? Error)
112117
}
113118
}
114-
method_setImplementation(method, imp_implementationWithBlock(closure))
115119
}
120+
method_setImplementation(method, imp_implementationWithBlock(closure))
121+
}
116122

117-
// - `urlSession(_:dataTask:didReceive:)`
118-
func swizzleDataTaskDidReceiveData(baseClass: AnyClass) {
119-
// "_didReceiveData"
120-
let selector = NSSelectorFromString(["_", "did", "Receive", "Data", ":"].joined())
121-
guard let method = class_getInstanceMethod(baseClass, selector),
122-
baseClass.instancesRespond(to: selector) else {
123-
return
124-
}
123+
// - `urlSession(_:dataTask:didReceive:)`
124+
func swizzleDataTaskDidReceiveData(baseClass: AnyClass) {
125+
// "_didReceiveData"
126+
let selector = NSSelectorFromString(["_", "did", "Receive", "Data", ":"].joined())
127+
guard let method = class_getInstanceMethod(baseClass, selector),
128+
baseClass.instancesRespond(to: selector) else {
129+
return
130+
}
125131

126-
typealias MethodSignature = @convention(c) (AnyObject, Selector, AnyObject) -> Void
127-
let originalImp: IMP = method_getImplementation(method)
128-
let closure: @convention(block) (AnyObject, AnyObject) -> Void = { [weak self] (object, data) in
129-
let original: MethodSignature = unsafeBitCast(originalImp, to: MethodSignature.self)
130-
original(object, selector, data)
132+
typealias MethodSignature = @convention(c) (AnyObject, Selector, AnyObject) -> Void
133+
let originalImp: IMP = method_getImplementation(method)
134+
let closure: @convention(block) (AnyObject, AnyObject) -> Void = { [weak self] (object, data) in
135+
let original: MethodSignature = unsafeBitCast(originalImp, to: MethodSignature.self)
136+
original(object, selector, data)
131137

132-
if let task = object.value(forKey: "task") as? URLSessionDataTask {
133-
let data = (data as? Data) ?? Data()
134-
self?.logger.logDataTask(task, didReceive: data)
135-
}
138+
if let task = object.value(forKey: "task") as? URLSessionDataTask {
139+
let data = (data as? Data) ?? Data()
140+
self?.logger.logDataTask(task, didReceive: data)
136141
}
137-
method_setImplementation(method, imp_implementationWithBlock(closure))
138142
}
143+
method_setImplementation(method, imp_implementationWithBlock(closure))
139144
}
140145
}

0 commit comments

Comments
 (0)