URLSession instrumentation: Fix custom request manipulation

Separate the possibility to modify URLRequest to another `injectCustomHeaders` callback, `shouldInjectTracingHeaders` can be called more than once, and if the request was being modified in the method it would happen twice.

Author: Ignacio Bonafonte

Sources/Instrumentation/URLSession/README.md

@@ -13,7 +13,9 @@ This behaviour can be modified or augmented by using the optional callbacks defi
 
 `shouldRecordPayload: ((URLSession) -> (Bool)?)?`: Implement if you want the session to record payload data, false by default.
 
-`shouldInjectTracingHeaders: ((inout URLRequest) -> (Bool)?)?`: Allows filtering which requests you want to inject headers to follow the trace, true by default. You can also modify the request or add other headers in this method.
+`shouldInjectTracingHeaders: ((URLRequest) -> (Bool)?)?`: Allows filtering which requests you want to inject headers to follow the trace, true by default. You must also return true if you want to inject custom headers.
+
+`injectCustomHeaders: ((inout URLRequest, Span?) -> Void)?`: Implement this callback to inject custom headers or modify the request in any other way
 
 `nameSpan: ((URLRequest) -> (String)?)?` - Modifies the name for the given request instead of stantard Opentelemetry name
 

Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift

@@ -107,8 +107,8 @@ public class URLSessionInstrumentation {
 
             let block: @convention(block) (URLSession, AnyObject) -> URLSessionTask = { session, argument in
                 if let url = argument as? URL {
-                    var request = URLRequest(url: url)
-                    if self.configuration.shouldInjectTracingHeaders?(&request) ?? true {
+                    let request = URLRequest(url: url)
+                    if self.configuration.shouldInjectTracingHeaders?(request) ?? true {
                         if selector == #selector(URLSession.dataTask(with:) as (URLSession) -> (URL) -> URLSessionDataTask) {
                             return session.dataTask(with: request)
                         } else {
@@ -184,9 +184,9 @@ public class URLSessionInstrumentation {
             let block: @convention(block) (URLSession, AnyObject, ((Any?, URLResponse?, Error?) -> Void)?) -> URLSessionTask = { session, argument, completion in
 
                 if let url = argument as? URL {
-                    var request = URLRequest(url: url)
+                    let request = URLRequest(url: url)
 
-                    if self.configuration.shouldInjectTracingHeaders?(&request) ?? true {
+                    if self.configuration.shouldInjectTracingHeaders?(request) ?? true {
                         if selector == #selector(URLSession.dataTask(with:completionHandler:) as (URLSession) -> (URL, @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTask) {
                             if let completion = completion {
                                 return session.dataTask(with: request, completionHandler: completion)

Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift

@@ -15,14 +15,16 @@ public struct URLSessionInstrumentationConfiguration {
     public init(shouldRecordPayload: ((URLSession) -> (Bool)?)? = nil,
                 shouldInstrument: ((URLRequest) -> (Bool)?)? = nil,
                 nameSpan: ((URLRequest) -> (String)?)? = nil,
-                shouldInjectTracingHeaders: ((inout URLRequest) -> (Bool)?)? = nil,
+                shouldInjectTracingHeaders: ((URLRequest) -> (Bool)?)? = nil,
+                injectCustomHeaders: ((inout URLRequest, Span?) -> Void)? = nil,
                 createdRequest: ((URLRequest, Span) -> Void)? = nil,
                 receivedResponse: ((URLResponse, DataOrFile?, Span) -> Void)? = nil,
                 receivedError: ((Error, DataOrFile?, HTTPStatus, Span) -> Void)? = nil)
     {
         self.shouldRecordPayload = shouldRecordPayload
         self.shouldInstrument = shouldInstrument
         self.shouldInjectTracingHeaders = shouldInjectTracingHeaders
+        self.injectCustomHeaders = injectCustomHeaders
         self.nameSpan = nameSpan
         self.createdRequest = createdRequest
         self.receivedResponse = receivedResponse
@@ -39,9 +41,12 @@ public struct URLSessionInstrumentationConfiguration {
     public var shouldRecordPayload: ((URLSession) -> (Bool)?)?
 
     /// Implement this callback to filter which requests you want to inject headers to follow the trace,
-    /// you can also modify the request or add other headers in this method.
+    /// also must implement it if you want to inject custom headers
     /// Instruments all requests by default
-    public var shouldInjectTracingHeaders: ((inout URLRequest) -> (Bool)?)?
+    public var shouldInjectTracingHeaders: ((URLRequest) -> (Bool)?)?
+
+    /// Implement this callback to inject custom headers or modify the request in any other way
+    public var injectCustomHeaders: ((inout URLRequest, Span?) -> Void)?
 
     /// Implement this callback to override the default span name for a given request, return nil to use default.
     /// default name: `HTTP {method}` e.g. `HTTP PUT`

Sources/Instrumentation/URLSession/URLSessionLogger.swift

@@ -138,10 +138,11 @@ class URLSessionLogger {
 
     private static func instrumentedRequest(for request: URLRequest, span: Span?, instrumentation: URLSessionInstrumentation) -> URLRequest? {
         var request = request
-        guard instrumentation.configuration.shouldInjectTracingHeaders?(&request) ?? true
+        guard instrumentation.configuration.shouldInjectTracingHeaders?(request) ?? true
         else {
             return nil
         }
+        instrumentation.configuration.injectCustomHeaders?(&request, span)
         var instrumentedRequest = request
         objc_setAssociatedObject(instrumentedRequest, &URLSessionInstrumentation.instrumentedKey, true, .OBJC_ASSOCIATION_COPY_NONATOMIC)
         var traceHeaders = tracePropagationHTTPHeaders(span: span, textMapPropagator: instrumentation.tracer.textFormat)