Merge pull request #221 from nachoBonafonte/fix-custom-request-manipulation

URLSession instrumentation: Fix request modification
Separate the possibility to modify URLRequest to a new callback: injectCustomHeaders
The previous used callback:shouldInjectTracingHeaders cannot modify it anymore, because for the same request it can be called more than once, and if the request was modified there it could end up processing it twice.
This is a breaking change for current users of shouldInjectTracingHeaders, but not for the rest.

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)