Adding traceparent HTTP requests header (#198)

.swiftlint.yml

@@ -2,6 +2,7 @@ disabled_rules: # rule identifiers turned on by default to exclude from running
   - identifier_name
   - function_parameter_count
   - notification_center_detachment
+  - compiler_protocol_init
 
 opt_in_rules:
   - conditional_returns_on_newline
@@ -17,16 +18,16 @@ excluded: # case-sensitive paths to ignore during linting. Takes precedence over
 reporter: "xcode" # reporter type (xcode, json, csv, checkstyle, codeclimate, junit, html, emoji, sonarqube, markdown, github-actions-logging, summary)
 
 file_length:
-  warning: 600
-  error: 800
-  
+  warning: 800
+  error: 1000
+
 line_length:
   ignores_comments: true
   ignores_interpolated_strings: true
 
 type_body_length:
-  warning: 400
-  error: 600
+  warning: 600
+  error: 800
 
 type_name:
   max_length:

Sources/EmbraceConfig/EmbraceConfig.swift

@@ -52,6 +52,10 @@ public class EmbraceConfig {
         return isEnabled(threshold: payload.backgroundSessionThreshold)
     }
 
+    public var isNetworkSpansForwardingEnabled: Bool {
+        return isEnabled(threshold: payload.networkSpansForwardingThreshold)
+    }
+
     // MARK: - Update
     @discardableResult
     public func updateIfNeeded() -> Bool {

Sources/EmbraceConfig/RemoteConfigPayload.swift

@@ -9,25 +9,33 @@ import Foundation
 struct RemoteConfigPayload: Decodable {
     var sdkEnabledThreshold: Float
     var backgroundSessionThreshold: Float
+    var networkSpansForwardingThreshold: Float
 
     enum CodingKeys: String, CodingKey {
         case sdkEnabledThreshold = "threshold"
         case background
+        case networkSpansForwarding = "network_span_forwarding"
 
         enum BackgroundCodingKeys: String, CodingKey {
             case threshold
         }
+
+        enum NetworkSpansForwardingCodigKeys: String, CodingKey {
+            case threshold = "pct_enabled"
+        }
     }
 
     public init(from decoder: Decoder) throws {
         let defaultPayload = RemoteConfigPayload()
 
+        // sdk enabled
         let rootContainer = try decoder.container(keyedBy: CodingKeys.self)
         sdkEnabledThreshold = try rootContainer.decodeIfPresent(
             Float.self,
             forKey: .sdkEnabledThreshold
         ) ?? defaultPayload.sdkEnabledThreshold
 
+        // background session
         if rootContainer.contains(.background) {
             let backgroundContainer = try rootContainer.nestedContainer(
                 keyedBy: CodingKeys.BackgroundCodingKeys.self,
@@ -40,12 +48,27 @@ struct RemoteConfigPayload: Decodable {
         } else {
             backgroundSessionThreshold = defaultPayload.backgroundSessionThreshold
         }
+
+        // network span forwarding
+        if rootContainer.contains(.networkSpansForwarding) {
+            let networkSpansForwardingContainer = try rootContainer.nestedContainer(
+                keyedBy: CodingKeys.NetworkSpansForwardingCodigKeys.self,
+                forKey: .networkSpansForwarding
+            )
+            networkSpansForwardingThreshold = try networkSpansForwardingContainer.decodeIfPresent(
+                Float.self,
+                forKey: CodingKeys.NetworkSpansForwardingCodigKeys.threshold
+            ) ?? defaultPayload.networkSpansForwardingThreshold
+        } else {
+            networkSpansForwardingThreshold = defaultPayload.networkSpansForwardingThreshold
+        }
     }
 
     // defaults
     public init() {
         sdkEnabledThreshold = 100.0
         backgroundSessionThreshold = 0.0
+        networkSpansForwardingThreshold = 0.0
     }
 }
 

Sources/EmbraceCore/Capture/Network/URLSessionCaptureService+Options.swift

@@ -0,0 +1,26 @@
+//
+//  Copyright © 2023 Embrace Mobile, Inc. All rights reserved.
+//
+
+import Foundation
+
+extension URLSessionCaptureService {
+    /// Class used to setup a URLSessionCaptureService.
+    @objc public final class Options: NSObject {
+        /// Defines wether or not the Embrace SDK should inject the `traceparent` header into all network requests
+        @objc public let injectTracingHeader: Bool
+
+        /// `URLSessionRequestsDataSource` instance that will manipuate all network requests
+        /// before the Embrace SDK captures their data.
+        @objc public let requestsDataSource: URLSessionRequestsDataSource?
+
+        @objc public init(injectTracingHeader: Bool, requestsDataSource: URLSessionRequestsDataSource?) {
+            self.injectTracingHeader = injectTracingHeader
+            self.requestsDataSource = requestsDataSource
+        }
+
+        @objc convenience override init() {
+            self.init(injectTracingHeader: true, requestsDataSource: nil)
+        }
+    }
+}

Sources/EmbraceCore/Capture/Network/URLSessionCaptureService.swift

@@ -21,16 +21,26 @@ protocol URLSessionSwizzler: Swizzlable {
 
 @objc public final class URLSessionCaptureService: CaptureService, URLSessionTaskHandlerDataSource {
 
+    public let options: URLSessionCaptureService.Options
     private let lock: NSLocking
     private let swizzlerProvider: URLSessionSwizzlerProvider
     private var swizzlers: [any URLSessionSwizzler] = []
     private var handler: URLSessionTaskHandler?
 
+    @objc public convenience init(options: URLSessionCaptureService.Options) {
+        self.init(options: options, lock: NSLock(), swizzlerProvider: DefaultURLSessionSwizzlerProvider())
+    }
+
     public convenience override init() {
         self.init(lock: NSLock(), swizzlerProvider: DefaultURLSessionSwizzlerProvider())
     }
 
-    init(lock: NSLocking, swizzlerProvider: URLSessionSwizzlerProvider) {
+    init(
+        options: URLSessionCaptureService.Options = URLSessionCaptureService.Options(),
+        lock: NSLocking,
+        swizzlerProvider: URLSessionSwizzlerProvider
+    ) {
+        self.options = options
         self.lock = lock
         self.swizzlerProvider = swizzlerProvider
     }
@@ -60,6 +70,20 @@ protocol URLSessionSwizzler: Swizzlable {
             }
         }
     }
+
+    var injectTracingHeader: Bool {
+        // check remote config
+        guard Embrace.client?.config.isNetworkSpansForwardingEnabled == true else {
+            return false
+        }
+
+        // check local config
+        return options.injectTracingHeader
+    }
+
+    var requestsDataSource: URLSessionRequestsDataSource? {
+        return options.requestsDataSource
+    }
 }
 
 // swiftlint:disable line_length
@@ -108,23 +132,19 @@ struct SessionTaskResumeSwizzler: URLSessionSwizzler {
         if #available(iOS 15.0, tvOS 15.0, macOS 12, watchOS 8, *) {
             try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
                 return { [weak handler = self.handler] task in
-                    if let handler = handler {
-                        handler.create(task: task) { captured in
-                            // if the task wasn't captured by other swizzlers it probably means
-                            // it was an async/await task
-                            // we set a proxy delegate to get a callback when the task finishes
-                            if captured {
-                                let originalDelegate = task.delegate
-                                task.delegate = URLSessionDelegateProxy(originalDelegate: originalDelegate, handler: handler)
-                            }
-
-                            // call original
-                            originalImplementation(task, Self.selector)
-                        }
-                    } else {
-                        // call original
-                        originalImplementation(task, Self.selector)
+                    let captured = handler?.create(task: task) ?? true
+
+                    // if the task wasn't captured by other swizzlers
+                    // by the time resume was called it probably means
+                    // it was an async/await task
+                    // we set a proxy delegate to get a callback when the task finishes
+                    if !captured, let handler = handler {
+                        let originalDelegate = task.delegate
+                        task.delegate = URLSessionDelegateProxy(originalDelegate: originalDelegate, handler: handler)
                     }
+
+                    // call original
+                    originalImplementation(task, Self.selector)
                 }
             }
         } else {
@@ -148,18 +168,10 @@ struct DataTaskWithURLSwizzler: URLSessionSwizzler {
     }
 
     func install() throws {
-        try swizzleInstanceMethod { originalImplementation in
-            return { [weak handler = self.handler] urlSession, url -> URLSessionDataTask in
-                // TODO: For this cases, we'll need to have the `URLSessionDelegate` swizzled/proxied.
-                // Note: two things to take into account:
-                // 1. We cannot ensure that this request will be tracked with this implementation.
-                //  For example, when using `URLSession.shared.dataTask(with:)`.
-                // 2. As this is has `URL` as parameter, and not `URLRequest`, we cannot add custom headers.
-                //  The old sdk would call the original method using `URLRequest`, but that leads to other kind
-                //  of edge cases.
-                let dataTask = originalImplementation(urlSession, Self.selector, url)
-                handler?.create(task: dataTask)
-                return dataTask
+        try swizzleInstanceMethod { _ in
+            return { urlSession, url -> URLSessionDataTask in
+                // create the task using a URLRequest and let the other swizzler handle it
+                return urlSession.dataTask(with: URLRequest(url: url))
             }
         }
     }
@@ -183,7 +195,6 @@ struct DataTaskWithURLRequestSwizzler: URLSessionSwizzler {
         try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
             return { [weak handler = self.handler] urlSession, urlRequest -> URLSessionDataTask in
                 let request = urlRequest.addEmbraceHeaders()
-                // TODO: For this cases, we'll need to have the `URLSessionDelegate` swizzled/proxied.
                 let dataTask = originalImplementation(urlSession, Self.selector, request)
                 handler?.create(task: dataTask)
                 return dataTask
@@ -208,27 +219,12 @@ struct DataTaskWithURLAndCompletionSwizzler: URLSessionSwizzler {
     }
 
     func install() throws {
-        try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
-            return { [weak handler = self.handler] urlSession, url, completion -> URLSessionDataTask in
-                // TODO: For this cases, we'll need to have the `URLSessionDelegate` swizzled/proxied.
-                guard let completion = completion else {
-                    let task = originalImplementation(urlSession, Self.selector, url, completion)
-                    handler?.create(task: task)
-                    return task
+        try swizzleInstanceMethod { _ in
+            return { urlSession, url, completion -> URLSessionDataTask in
+                // create the task using a URLRequest and let the other swizzler handle it
+                return urlSession.dataTask(with: URLRequest(url: url)) { data, response, error in
+                    completion?(data, response, error)
                 }
-
-                var originalTask: URLSessionDataTask?
-
-                let dataTask = originalImplementation(urlSession, Self.selector, url) { data, response, error in
-                    if let task = originalTask {
-                        handler?.finish(task: task, data: data, error: error)
-                    }
-                    completion(data, response, error)
-                }
-
-                originalTask = dataTask
-                handler?.create(task: dataTask)
-                return dataTask
             }
         }
     }
@@ -254,21 +250,24 @@ struct DataTaskWithURLRequestAndCompletionSwizzler: URLSessionSwizzler {
     func install() throws {
         try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
             return { [weak handler = self.handler] urlSession, urlRequest, completion -> URLSessionDataTask in
-                // TODO: For this cases, we'll need to have the `URLSessionDelegate` swizzled/proxied.
+
+                let request = urlRequest.addEmbraceHeaders()
+
                 guard let completion = completion else {
-                    let task = originalImplementation(urlSession, Self.selector, urlRequest, completion)
+                    let task = originalImplementation(urlSession, Self.selector, request, completion)
                     handler?.create(task: task)
                     return task
                 }
 
                 var originalTask: URLSessionDataTask?
-                let request = urlRequest.addEmbraceHeaders()
+
                 let dataTask = originalImplementation(urlSession, Self.selector, request) { data, response, error in
                     if let task = originalTask {
                         handler?.finish(task: task, data: data, error: error)
                     }
                     completion(data, response, error)
                 }
+
                 originalTask = dataTask
                 handler?.create(task: dataTask)
                 return dataTask
@@ -392,13 +391,15 @@ struct UploadTaskWithRequestFromFileWithCompletionSwizzler: URLSessionSwizzler {
     func install() throws {
         try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
             return { [weak handler = self.handler] urlSession, urlRequest, url, completion -> URLSessionUploadTask in
+
+                let request = urlRequest.addEmbraceHeaders()
+
                 guard let completion = completion else {
-                    let task = originalImplementation(urlSession, Self.selector, urlRequest, url, completion)
+                    let task = originalImplementation(urlSession, Self.selector, request, url, completion)
                     handler?.create(task: task)
                     return task
                 }
 
-                let request = urlRequest.addEmbraceHeaders()
                 var originalTask: URLSessionUploadTask?
                 let uploadTask = originalImplementation(urlSession, Self.selector, request, url) { data, response, error in
                     if let task = originalTask {
@@ -461,14 +462,15 @@ struct DownloadTaskWithURLRequestWithCompletionSwizzler: URLSessionSwizzler {
     func install() throws {
         try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
             return { [weak handler = self.handler] urlSession, urlRequest, completion -> URLSessionDownloadTask in
+
+                let request = urlRequest.addEmbraceHeaders()
+
                 guard let completion = completion else {
-                    let task = originalImplementation(urlSession, Self.selector, urlRequest, completion)
+                    let task = originalImplementation(urlSession, Self.selector, request, completion)
                     handler?.create(task: task)
                     return task
                 }
 
-                let request = urlRequest.addEmbraceHeaders()
-
                 var originalTask: URLSessionDownloadTask?
                 let downloadTask = originalImplementation(urlSession, Self.selector, request) { url, response, error in
                     if let task = originalTask {

Sources/EmbraceCore/Capture/Network/URLSessionRequestsDataSource.swift

@@ -0,0 +1,15 @@
+//
+//  Copyright © 2023 Embrace Mobile, Inc. All rights reserved.
+//
+
+import Foundation
+
+/// This protocol can be used to modify requests before the Embrace SDK
+/// captures their data into OTel spans.
+///
+/// Example:
+/// This could be useful if you need to obfuscate certains parts of a request path
+/// if it contains sensitive data.
+@objc public protocol URLSessionRequestsDataSource: NSObjectProtocol {
+    @objc func modifiedRequest(for request: URLRequest) -> URLRequest
+}