Add support for dataTask with URL and completion handler

DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift

@@ -94,37 +94,31 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {
             }
         )
 
-        if #available(iOS 15, tvOS 15, *) {
-            try swizzler.swizzle(
-                delegateClass: configuration.delegateClass,
-                interceptDidFinishCollecting: { [weak self] session, task, metrics in
-                    self?.task(task, didFinishCollecting: metrics)
+        try swizzler.swizzle(
+            delegateClass: configuration.delegateClass,
+            interceptDidFinishCollecting: { [weak self] session, task, metrics in
+                self?.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
                     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(
+            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)
-                }
-            )
-        }
+        try swizzler.swizzle(
+            interceptCompletionHandler: { [weak self] task, _, error in
+                self?.task(task, didCompleteWithError: error)
+            }
+        )
     }
 
     /// Unswizzles `URLSessionTaskDelegate`, `URLSessionDataDelegate`, `URLSessionTask` and `URLSession` methods

DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift

@@ -135,15 +135,11 @@ open class DatadogURLSessionDelegate: NSObject, URLSessionDataDelegate {
             }
         )
 
-        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)
-                }
-            )
-        }
+        try swizzler.swizzle(
+            interceptCompletionHandler: { [weak self] task, _, error in
+                self?.interceptor?.task(task, didCompleteWithError: error)
+            }
+        )
     }
 
     deinit {

DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift

@@ -10,19 +10,22 @@ import Foundation
 internal final class URLSessionSwizzler {
     private let lock = NSRecursiveLock()
 
-    private var dataTaskCompletionHandler: DataTaskCompletionHandler?
+    private var dataTaskURLRequestCompletionHandler: DataTaskURLRequestCompletionHandler?
+    private var dataTaskURLCompletionHandler: DataTaskURLCompletionHandler?
     private var taskResume: TaskResume?
     private var didFinishCollecting: DidFinishCollecting?
     private var didReceive: DidReceive?
     private var didCompleteWithError: DidCompleteWithError?
 
-    /// Swizzles `URLSession.dataTask(with:completionHandler:)` method.
+    /// Swizzles `URLSession.dataTask(with:completionHandler:)` methods (with `URL` and `URLRequest`).
     func swizzle(
         interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void
     ) throws {
         lock.lock()
-        dataTaskCompletionHandler = try DataTaskCompletionHandler.build()
-        dataTaskCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler)
+        dataTaskURLRequestCompletionHandler = try DataTaskURLRequestCompletionHandler.build()
+        dataTaskURLRequestCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler)
+        dataTaskURLCompletionHandler = try DataTaskURLCompletionHandler.build()
+        dataTaskURLCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler)
         lock.unlock()
     }
 
@@ -74,7 +77,8 @@ internal final class URLSessionSwizzler {
     /// This method is called during deinit.
     func unswizzle() {
         lock.lock()
-        dataTaskCompletionHandler?.unswizzle()
+        dataTaskURLRequestCompletionHandler?.unswizzle()
+        dataTaskURLCompletionHandler?.unswizzle()
         taskResume?.unswizzle()
         didFinishCollecting?.unswizzle()
         didCompleteWithError?.unswizzle()
@@ -88,16 +92,16 @@ internal final class URLSessionSwizzler {
 
     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> {
+    /// Swizzles `URLSession.dataTask(with:completionHandler:)` (with `URLRequest`) method.
+    class DataTaskURLRequestCompletionHandler: 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(
+        static func build() throws -> DataTaskURLRequestCompletionHandler {
+            return try DataTaskURLRequestCompletionHandler(
                 selector: self.selector,
                 klass: URLSession.self
             )
@@ -138,6 +142,56 @@ internal final class URLSessionSwizzler {
         }
     }
 
+    /// Swizzles `URLSession.dataTask(with:completionHandler:)` (with `URL`) method.
+    class DataTaskURLCompletionHandler: MethodSwizzler<@convention(c) (URLSession, Selector, URL, CompletionHandler?) -> URLSessionDataTask, @convention(block) (URLSession, URL, CompletionHandler?) -> URLSessionDataTask> {
+        private static let selector = #selector(
+            URLSession.dataTask(with:completionHandler:) as (URLSession) -> (URL, @escaping CompletionHandler) -> URLSessionDataTask
+        )
+
+        private let method: Method
+
+        static func build() throws -> DataTaskURLCompletionHandler {
+            return try DataTaskURLCompletionHandler(
+                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, URL, CompletionHandler?) -> URLSessionDataTask
+            swizzle(method) { previousImplementation -> Signature in
+                return { session, url, 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, url, completionHandler)
+                    }
+
+                    var _task: URLSessionDataTask?
+                    let task = previousImplementation(session, Self.selector, url) { 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)

DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift

@@ -11,6 +11,7 @@ import XCTest
 class URLSessionSwizzlerTests: XCTestCase {
     func testSwizzling_dataTaskWithCompletion() throws {
         let didInterceptCompletion = XCTestExpectation(description: "interceptCompletion")
+        didInterceptCompletion.expectedFulfillmentCount = 2
 
         let swizzler = URLSessionSwizzler()
 
@@ -21,9 +22,9 @@ class URLSessionSwizzlerTests: XCTestCase {
         )
 
         let session = URLSession(configuration: .default)
-        let request = URLRequest(url: URL(string: "https://www.datadoghq.com/")!)
-        let task = session.dataTask(with: request) { _, _, _ in }
-        task.resume()
+        let url = URL(string: "https://www.datadoghq.com/")!
+        session.dataTask(with: url) { _, _, _ in }.resume()
+        session.dataTask(with: URLRequest(url: url)) { _, _, _ in }.resume()
 
         wait(for: [didInterceptCompletion], timeout: 5)
     }