diff --git a/DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift b/DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift index 210f4935fb..13fe608c55 100644 --- a/DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift +++ b/DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift @@ -94,12 +94,37 @@ internal final class NetworkInstrumentationFeature: DatadogFeature { } ) - try swizzler.swizzle( - delegateClass: configuration.delegateClass, - interceptDidFinishCollecting: { [weak self] session, task, metrics in - self?.task(task, didFinishCollecting: metrics) - } - ) + if #available(iOS 15, tvOS 15, *) { + try swizzler.swizzle( + delegateClass: configuration.delegateClass, + interceptDidFinishCollecting: { [weak self] session, task, metrics in + self?.task(task, didFinishCollecting: metrics) + // iOS 15 and above, didCompleteWithError is not called hence we use task state to detect task completion + // while prior to iOS 15, task state doesn't change to completed hence we use didCompleteWithError to detect task completion + self?.task(task, didCompleteWithError: task.error) + } + ) + } else { + try swizzler.swizzle( + delegateClass: configuration.delegateClass, + interceptDidFinishCollecting: { [weak self] session, task, metrics in + self?.task(task, didFinishCollecting: metrics) + } + ) + + try swizzler.swizzle( + delegateClass: configuration.delegateClass, + interceptDidCompleteWithError: { [weak self] session, task, error in + self?.task(task, didCompleteWithError: error) + } + ) + + try swizzler.swizzle( + interceptCompletionHandler: { [weak self] task, _, error in + self?.task(task, didCompleteWithError: error) + } + ) + } } /// Unswizzles `URLSessionTaskDelegate`, `URLSessionDataDelegate`, `URLSessionTask` and `URLSession` methods @@ -149,14 +174,6 @@ extension NetworkInstrumentationFeature { isFirstParty: firstPartyHosts.isFirstParty(url: request.url) ) - // observe the state for completion - // note: all task properties support Key-Value Observing - interception.stateValueObserver = task.observe(\.state, options: [.initial, .new]) { [weak self] task, _ in - if task.state == .completed { - self?.task(task, didCompleteWithError: task.error) - } - } - interception.register(request: request) if let trace = self.extractTrace(firstPartyHosts: firstPartyHosts, request: request) { diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift index 933bc7316e..9c84a82232 100644 --- a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift +++ b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift @@ -98,6 +98,11 @@ open class DatadogURLSessionDelegate: NSObject, URLSessionDataDelegate { open func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) { interceptor?.task(task, didFinishCollecting: metrics) + if #available(iOS 15, tvOS 15, *) { + // iOS 15 and above, didCompleteWithError is not called hence we use task state to detect task completion + // while prior to iOS 15, task state doesn't change to completed hence we use didCompleteWithError to detect task completion + interceptor?.task(task, didCompleteWithError: task.error) + } } open func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) { @@ -129,6 +134,16 @@ open class DatadogURLSessionDelegate: NSObject, URLSessionDataDelegate { interceptor.intercept(task: task, additionalFirstPartyHosts: firstPartyHosts) } ) + + if #unavailable(iOS 15, tvOS 15) { + // prior to iOS 15, task state doesn't change to completed + // hence we use didCompleteWithError to detect task completion + try swizzler.swizzle( + interceptCompletionHandler: { [weak self] task, _, error in + self?.interceptor?.task(task, didCompleteWithError: error) + } + ) + } } deinit { diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift index 54061e006f..d69bf82f59 100644 --- a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift +++ b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift @@ -10,9 +10,21 @@ import Foundation internal final class URLSessionSwizzler { private let lock = NSRecursiveLock() + private var dataTaskCompletionHandler: DataTaskCompletionHandler? private var taskResume: TaskResume? private var didFinishCollecting: DidFinishCollecting? private var didReceive: DidReceive? + private var didCompleteWithError: DidCompleteWithError? + + /// Swizzles `URLSession.dataTask(with:completionHandler:)` method. + func swizzle( + interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void + ) throws { + lock.lock() + dataTaskCompletionHandler = try DataTaskCompletionHandler.build() + dataTaskCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler) + lock.unlock() + } /// Swizzles `URLSessionTask.resume()` method. func swizzle( @@ -46,13 +58,26 @@ internal final class URLSessionSwizzler { lock.unlock() } + /// Swizzles `URLSessionTaskDelegate.urlSession(_:task:didCompleteWithError:)` method. + func swizzle( + delegateClass: AnyClass, + interceptDidCompleteWithError: @escaping (URLSession, URLSessionTask, Error?) -> Void + ) throws { + lock.lock() + didCompleteWithError = try DidCompleteWithError.build(klass: delegateClass) + didCompleteWithError?.swizzle(intercept: interceptDidCompleteWithError) + lock.unlock() + } + /// Unswizzles all. /// /// This method is called during deinit. func unswizzle() { lock.lock() + dataTaskCompletionHandler?.unswizzle() taskResume?.unswizzle() didFinishCollecting?.unswizzle() + didCompleteWithError?.unswizzle() didReceive?.unswizzle() lock.unlock() } @@ -61,6 +86,58 @@ internal final class URLSessionSwizzler { unswizzle() } + typealias CompletionHandler = (Data?, URLResponse?, Error?) -> Void + + /// Swizzles `URLSession.dataTask(with:completionHandler:)` method. + class DataTaskCompletionHandler: MethodSwizzler<@convention(c) (URLSession, Selector, URLRequest, CompletionHandler?) -> URLSessionDataTask, @convention(block) (URLSession, URLRequest, CompletionHandler?) -> URLSessionDataTask> { + private static let selector = #selector( + URLSession.dataTask(with:completionHandler:) as (URLSession) -> (URLRequest, @escaping CompletionHandler) -> URLSessionDataTask + ) + + private let method: Method + + static func build() throws -> DataTaskCompletionHandler { + return try DataTaskCompletionHandler( + selector: self.selector, + klass: URLSession.self + ) + } + + private init(selector: Selector, klass: AnyClass) throws { + self.method = try dd_class_getInstanceMethod(klass, selector) + super.init() + } + + func swizzle( + interceptCompletion: @escaping (URLSessionTask, Data?, Error?) -> Void + ) { + typealias Signature = @convention(block) (URLSession, URLRequest, CompletionHandler?) -> URLSessionDataTask + swizzle(method) { previousImplementation -> Signature in + return { session, request, completionHandler -> URLSessionDataTask in + guard let completionHandler = completionHandler else { + // The `completionHandler` can be `nil` in two cases: + // - on iOS 11 or 12, where `dataTask(with:)` (for `URL` and `URLRequest`) calls + // the `dataTask(with:completionHandler:)` (for `URLRequest`) internally by nullifying the completion block. + // - when `[session dataTaskWithURL:completionHandler:]` is called in Objective-C with explicitly passing + // `nil` as the `completionHandler` (it produces a warning, but compiles). + return previousImplementation(session, Self.selector, request, completionHandler) + } + + var _task: URLSessionDataTask? + let task = previousImplementation(session, Self.selector, request) { data, response, error in + completionHandler(data, response, error) + + if let task = _task { // sanity check, should always succeed + interceptCompletion(task, data, error) + } + } + _task = task + return task + } + } + } + } + /// Swizzles `URLSessionTask.resume()` method. class TaskResume: MethodSwizzler<@convention(c) (URLSessionTask, Selector) -> Void, @convention(block) (URLSessionTask) -> Void> { private static let selector = #selector(URLSessionTask.resume) @@ -178,4 +255,48 @@ internal final class URLSessionSwizzler { } } } + + class DidCompleteWithError: MethodSwizzler<@convention(c) (URLSessionTaskDelegate, Selector, URLSession, URLSessionTask, Error?) -> Void, @convention(block) (URLSessionTaskDelegate, URLSession, URLSessionTask, Error?) -> Void> { + private static let selector = #selector(URLSessionTaskDelegate.urlSession(_:task:didCompleteWithError:)) + + private let method: Method + + static func build(klass: AnyClass) throws -> DidCompleteWithError { + return try DidCompleteWithError(selector: self.selector, klass: klass) + } + + private init(selector: Selector, klass: AnyClass) throws { + do { + method = try dd_class_getInstanceMethod(klass, selector) + } catch { + // URLSessionTaskDelegate doesn't implement the selector, so we inject it and swizzle it + let block: @convention(block) (URLSessionTaskDelegate, URLSession, URLSessionTask, Error?) -> Void = { delegate, session, task, error in + } + let imp = imp_implementationWithBlock(block) + /* + v@:@@@ means: + v - return type is void + @ - self + : - selector + @ - first argument is an object + @ - second argument is an object + @ - third argument is an object + */ + class_addMethod(klass, selector, imp, "v@:@@@") + method = try dd_class_getInstanceMethod(klass, selector) + } + + super.init() + } + + func swizzle(intercept: @escaping (URLSession, URLSessionTask, Error?) -> Void) { + typealias Signature = @convention(block) (URLSessionTaskDelegate, URLSession, URLSessionTask, Error?) -> Void + swizzle(method) { previousImplementation -> Signature in + return { delegate, session, task, error in + intercept(session, task, error) + return previousImplementation(delegate, Self.selector, session, task, error) + } + } + } + } } diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTaskInterception.swift b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTaskInterception.swift index aadec49a7c..375da553d5 100644 --- a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTaskInterception.swift +++ b/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTaskInterception.swift @@ -31,9 +31,6 @@ public class URLSessionTaskInterception { /// /// Setting the value to 'rum' will indicate that the span is reported as a RUM Resource. public private(set) var origin: String? - /// A KVO token observing the ``URLSessionTask``'s state. - /// The token is invalidated at deinit. - internal var stateValueObserver: NSKeyValueObservation? init(request: URLRequest, isFirstParty: Bool) { self.identifier = UUID() diff --git a/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift b/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift index 97dac13d77..2bdb7c8d96 100644 --- a/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift +++ b/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift @@ -9,6 +9,25 @@ import XCTest @testable import DatadogInternal class URLSessionSwizzlerTests: XCTestCase { + func testSwizzling_dataTaskWithCompletion() throws { + let didInterceptCompletion = XCTestExpectation(description: "interceptCompletion") + + let swizzler = URLSessionSwizzler() + + try swizzler.swizzle( + interceptCompletionHandler: { _, _, _ in + didInterceptCompletion.fulfill() + } + ) + + let session = URLSession(configuration: .default) + let request = URLRequest(url: URL(string: "https://www.datadoghq.com/")!) + let task = session.dataTask(with: request) { _, _, _ in } + task.resume() + + wait(for: [didInterceptCompletion], timeout: 5) + } + func testSwizzling_whenDidReceiveDataIsImplemented() throws { class MockDelegate: NSObject, URLSessionDataDelegate { func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) { @@ -77,12 +96,13 @@ class URLSessionSwizzlerTests: XCTestCase { func testSwizzling_taskDelegate_whenMethodsAreImplemented() throws { class MockDelegate: NSObject, URLSessionTaskDelegate { - func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) { - } + func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) { } + func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) { } } let delegate = MockDelegate() let didFinishCollecting = XCTestExpectation(description: "didFinishCollecting") + let didCompleteWithError = XCTestExpectation(description: "didCompleteWithError") let swizzler = URLSessionSwizzler() @@ -93,11 +113,18 @@ class URLSessionSwizzlerTests: XCTestCase { } ) + try swizzler.swizzle( + delegateClass: MockDelegate.self, + interceptDidCompleteWithError: { _, _, _ in + didCompleteWithError.fulfill() + } + ) + let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: nil) let task = session.dataTask(with: URL(string: "https://www.datadoghq.com/")!) task.resume() - wait(for: [didFinishCollecting], timeout: 5) + wait(for: [didFinishCollecting, didCompleteWithError], timeout: 5) } func testSwizzling_taskDelegate_whenMethodsAreNotImplemented() throws { @@ -106,6 +133,7 @@ class URLSessionSwizzlerTests: XCTestCase { let delegate = MockDelegate() let didFinishCollecting = XCTestExpectation(description: "didFinishCollecting") + let didCompleteWithError = XCTestExpectation(description: "didCompleteWithError") let swizzler = URLSessionSwizzler() @@ -116,10 +144,17 @@ class URLSessionSwizzlerTests: XCTestCase { } ) + try swizzler.swizzle( + delegateClass: MockDelegate.self, + interceptDidCompleteWithError: { _, _, _ in + didCompleteWithError.fulfill() + } + ) + let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: nil) let task = session.dataTask(with: URL(string: "https://www.datadoghq.com/")!) task.resume() - wait(for: [didFinishCollecting], timeout: 5) + wait(for: [didFinishCollecting, didCompleteWithError], timeout: 5) } }