RUMM-1234 DelegateProviding protocol added

DataDog · Apr 1, 2021 · b05dd10 · b05dd10
1 parent c9a837f
commit b05dd10
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 10 deletions.
diff --git a/...ios/URLSession/NSURLSessionAutoInstrumentation/ObjcSendFirstPartyRequestsViewController.m b/...ios/URLSession/NSURLSessionAutoInstrumentation/ObjcSendFirstPartyRequestsViewController.m
@@ -20,7 +20,7 @@ - (void)viewDidLoad {
 
     self.testScenario = SwiftGlobals.currentTestScenario;
 
-    self.session = [self.testScenario buildURLSession];
+    self.session = [self.testScenario buildNSURLSession];
     assert(self.testScenario != nil);
 }
 

diff --git a/...ios/URLSession/NSURLSessionAutoInstrumentation/ObjcSendThirdPartyRequestsViewController.m b/...ios/URLSession/NSURLSessionAutoInstrumentation/ObjcSendThirdPartyRequestsViewController.m
@@ -18,7 +18,7 @@ @implementation ObjcSendThirdPartyRequestsViewController
 - (void)viewDidLoad {
     [super viewDidLoad];
     self.testScenario = SwiftGlobals.currentTestScenario;
-    self.session = [self.testScenario buildURLSession];
+    self.session = [self.testScenario buildNSURLSession];
     assert(self.testScenario != nil);
 }
 

diff --git a/Datadog/Example/Scenarios/URLSession/URLSessionScenarios.swift b/Datadog/Example/Scenarios/URLSession/URLSessionScenarios.swift
@@ -6,6 +6,7 @@
 
 import Foundation
 import Datadog
+import DatadogObjc
 
 /// Base scenario for `URLSession` and `NSURLSession` instrumentation.  It makes
 /// both Swift and Objective-C tests share the same endpoints and SDK configuration.
@@ -83,7 +84,6 @@ class URLSessionBaseScenario: NSObject {
         }
     }
 
-    @objc
     func buildURLSession() -> URLSession {
         let delegate: DDURLSessionDelegate
         if feedAdditionalFirstyPartyHosts {
@@ -103,4 +103,25 @@ class URLSessionBaseScenario: NSObject {
             delegateQueue: nil
         )
     }
+
+    @objc
+    func buildNSURLSession() -> URLSession {
+        let delegate: DDNSURLSessionDelegate
+        if feedAdditionalFirstyPartyHosts {
+            delegate = DDNSURLSessionDelegate(
+                additionalFirstPartyHosts: [
+                    customGETResourceURL.host!,
+                    customPOSTRequest.url!.host!,
+                    badResourceURL.host!
+                ]
+            )
+        } else {
+            delegate = DDNSURLSessionDelegate()
+        }
+        return URLSession(
+            configuration: .default,
+            delegate: delegate,
+            delegateQueue: nil
+        )
+    }
 }
diff --git a/Sources/Datadog/URLSessionAutoInstrumentation/DDURLSessionDelegate.swift b/Sources/Datadog/URLSessionAutoInstrumentation/DDURLSessionDelegate.swift
@@ -6,12 +6,20 @@
 
 import Foundation
 
+public protocol DDURLSessionDelegateProviding {
+    var delegate: DDURLSessionDelegate { get }
+}
+
 /// The `URLSession` delegate object which enables network requests instrumentation. **It must be
 /// used together with** `Datadog.Configuration.trackURLSession(firstPartyHosts:)`.
 ///
 /// All requests made with the `URLSession` instrumented with this delegate will be intercepted by the SDK.
 @objc
-open class DDURLSessionDelegate: NSObject, URLSessionTaskDelegate {
+open class DDURLSessionDelegate: NSObject, URLSessionTaskDelegate, DDURLSessionDelegateProviding {
+    public var delegate: DDURLSessionDelegate {
+        return self
+    }
+
     var interceptor: URLSessionInterceptorType?
     let firstPartyURLsFilter: FirstPartyURLsFilter?
 

diff --git a/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift b/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift
@@ -178,7 +178,8 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
     // MARK: - Private
 
     private func isFirstParty(request: URLRequest, for session: URLSession?) -> Bool {
-        let delegateURLFilter = (session?.delegate as? DDURLSessionDelegate)?.firstPartyURLsFilter
+        let delegateProvider = session?.delegate as? DDURLSessionDelegateProviding
+        let delegateURLFilter = delegateProvider?.delegate.firstPartyURLsFilter
         let isFirstPartyForDelegate = (delegateURLFilter?.isFirstParty(url: request.url)) ?? false
         let isFirstPartyForInterceptor = self.defaultFirstPartyURLsFilter.isFirstParty(url: request.url)
         return isFirstPartyForDelegate || isFirstPartyForInterceptor

diff --git a/Sources/Datadog/URLSessionAutoInstrumentation/URLSessionSwizzler.swift b/Sources/Datadog/URLSessionAutoInstrumentation/URLSessionSwizzler.swift
@@ -70,7 +70,8 @@ internal class URLSessionSwizzler {
             typealias Signature = @convention(block) (URLSession, URLRequest, CompletionHandler?) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
                 return { session, urlRequest, completionHandler -> URLSessionDataTask in
-                    guard let interceptor = (session.delegate as? DDURLSessionDelegate)?.interceptor else {
+                    guard let delegateProvider = session.delegate as? DDURLSessionDelegateProviding,
+                          let interceptor = delegateProvider.delegate.interceptor else {
                         return previousImplementation(session, Self.selector, urlRequest, completionHandler)
                     }
                     let task: URLSessionDataTask
@@ -129,7 +130,8 @@ internal class URLSessionSwizzler {
             typealias Signature = @convention(block) (URLSession, URL, CompletionHandler?) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
                 return { session, url, completionHandler -> URLSessionDataTask in
-                    guard let interceptor = (session.delegate as? DDURLSessionDelegate)?.interceptor else {
+                    guard let delegateProvider = session.delegate as? DDURLSessionDelegateProviding,
+                          let interceptor = delegateProvider.delegate.interceptor else {
                         return previousImplementation(session, Self.selector, url, completionHandler)
                     }
                     let task: URLSessionDataTask
@@ -183,7 +185,8 @@ internal class URLSessionSwizzler {
             typealias Signature = @convention(block) (URLSession, URLRequest) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
                 return { session, urlRequest -> URLSessionDataTask in
-                    guard let interceptor = (session.delegate as? DDURLSessionDelegate)?.interceptor else {
+                    guard let delegateProvider = session.delegate as? DDURLSessionDelegateProviding,
+                          let interceptor = delegateProvider.delegate.interceptor else {
                         return previousImplementation(session, Self.selector, urlRequest)
                     }
                     let newRequest = interceptor.modify(request: urlRequest, session: session)
@@ -227,7 +230,8 @@ internal class URLSessionSwizzler {
             typealias Signature = @convention(block) (URLSession, URL) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
                 return { session, url -> URLSessionDataTask in
-                    guard let interceptor = (session.delegate as? DDURLSessionDelegate)?.interceptor else {
+                    guard let delegateProvider = session.delegate as? DDURLSessionDelegateProviding,
+                          let interceptor = delegateProvider.delegate.interceptor else {
                         return previousImplementation(session, Self.selector, url)
                     }
                     let task = previousImplementation(session, Self.selector, url)

diff --git a/Sources/DatadogObjc/DDURLSessionDelegate+objc.swift b/Sources/DatadogObjc/DDURLSessionDelegate+objc.swift
@@ -8,4 +8,26 @@ import Foundation
 import Datadog
 
 @objc
-public class DDNSURLSessionDelegate: DDURLSessionDelegate {}
+open class DDNSURLSessionDelegate: NSObject, URLSessionTaskDelegate, DDURLSessionDelegateProviding {
+    public let delegate: DDURLSessionDelegate
+
+    @objc
+    override public init() {
+        self.delegate = DDURLSessionDelegate()
+    }
+
+    @objc
+    public init(additionalFirstPartyHosts: Set<String>) {
+        self.delegate = DDURLSessionDelegate(additionalFirstPartyHosts: additionalFirstPartyHosts)
+    }
+
+    @objc
+    open func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) {
+        delegate.urlSession(session, task: task, didFinishCollecting: metrics)
+    }
+
+    @objc
+    open func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
+        delegate.urlSession(session, task: task, didCompleteWithError: error)
+    }
+}