Skip to content

Commit

Permalink
restore swizzling
Browse files Browse the repository at this point in the history
  • Loading branch information
@maxep
maxep committed Dec 15, 2023
1 parent 8174c17 commit b8155aa
Show file tree
Hide file tree
Showing 5 changed files with 206 additions and 21 deletions.
45 changes: 31 additions & 14 deletions DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down
15 changes: 15 additions & 0 deletions DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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 {
Expand Down
121 changes: 121 additions & 0 deletions DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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()
}
Expand All @@ -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)
Expand Down Expand Up @@ -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)
}
}
}
}
}
3 changes: 0 additions & 3 deletions DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTaskInterception.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
43 changes: 39 additions & 4 deletions DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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()

Expand All @@ -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 {
Expand All @@ -106,6 +133,7 @@ class URLSessionSwizzlerTests: XCTestCase {

let delegate = MockDelegate()
let didFinishCollecting = XCTestExpectation(description: "didFinishCollecting")
let didCompleteWithError = XCTestExpectation(description: "didCompleteWithError")

let swizzler = URLSessionSwizzler()

Expand All @@ -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)
}
}

0 comments on commit b8155aa

Please sign in to comment.