diff --git a/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift b/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift index 7737f3c36f..41d4b61be8 100644 --- a/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift +++ b/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift @@ -44,7 +44,7 @@ public class URLSessionInterceptor: URLSessionInterceptorType { internal let injectTracingHeadersToFirstPartyRequests: Bool /// Additional header injected to intercepted 1st party requests. /// Set to `x-datadog-origin: rum` if both RUM and Tracing instrumentations are enabled and `nil` in all other cases. - internal let additionalHeadersForFirstPartyRequests: [String: String]? + internal let datadogOriginHeader: (field: String, value: String)? /// Tracing sampler used to sample traces generated by the SDK or RUM BE. internal let tracingSampler: Sampler @@ -88,15 +88,13 @@ public class URLSessionInterceptor: URLSessionInterceptorType { if configuration.instrumentRUM { // If RUM instrumentation is enabled, additional `x-datadog-origin: rum` header is injected to the user request, // so that user's backend instrumentation can further process it and count on RUM quota (w/o affecting APM Index Spans counts). - self.additionalHeadersForFirstPartyRequests = [ - TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value - ] + self.datadogOriginHeader = TracingHTTPHeaders.ddOrigin } else { - self.additionalHeadersForFirstPartyRequests = nil + self.datadogOriginHeader = nil } } else { self.injectTracingHeadersToFirstPartyRequests = false - self.additionalHeadersForFirstPartyRequests = nil + self.datadogOriginHeader = nil } } @@ -250,6 +248,9 @@ public class URLSessionInterceptor: URLSessionInterceptorType { switch $0 { case .datadog: writer = HTTPHeadersWriter(sampler: tracingSampler) + if let datadogOriginHeader = datadogOriginHeader { + newRequest.setValue(datadogOriginHeader.value, forHTTPHeaderField: datadogOriginHeader.field) + } case .b3: writer = OTelHTTPHeadersWriter( sampler: tracingSampler, @@ -268,10 +269,6 @@ public class URLSessionInterceptor: URLSessionInterceptorType { writer.tracePropagationHTTPHeaders.forEach { field, value in newRequest.setValue(value, forHTTPHeaderField: field) } - - additionalHeadersForFirstPartyRequests?.forEach { field, value in - newRequest.setValue(value, forHTTPHeaderField: field) - } } return newRequest } diff --git a/Tests/DatadogTests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift b/Tests/DatadogTests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift index 7a382ee9eb..85cf2baa2d 100644 --- a/Tests/DatadogTests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift +++ b/Tests/DatadogTests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift @@ -55,7 +55,7 @@ class URLSessionInterceptorTests: XCTestCase { "Tracing headers should be injected when only Tracing instrumentation is enabled." ) XCTAssertNil( - interceptor.additionalHeadersForFirstPartyRequests, + interceptor.datadogOriginHeader, "Just the tracing headers should be injected when only Tracing instrumentation is enabled." ) } @@ -79,7 +79,7 @@ class URLSessionInterceptorTests: XCTestCase { "Tracing headers should not be injected when only RUM instrumentation is enabled." ) XCTAssertNil( - interceptor.additionalHeadersForFirstPartyRequests, + interceptor.datadogOriginHeader, "No additional headers should be injected when only RUM instrumentation is enabled." ) } @@ -103,9 +103,14 @@ class URLSessionInterceptorTests: XCTestCase { "Tracing headers should be injected when both Tracing and RUM instrumentations are enabled." ) XCTAssertEqual( - interceptor.additionalHeadersForFirstPartyRequests, - [TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value], - "Additional `x-datadog-origin: rum` header should be injected when both Tracing and RUM instrumentations are enabled." + interceptor.datadogOriginHeader?.field, + TracingHTTPHeaders.ddOrigin.field, + "Additional `x-datadog-origin` header field should be injected when both Tracing and RUM instrumentations are enabled." + ) + XCTAssertEqual( + interceptor.datadogOriginHeader?.value, + TracingHTTPHeaders.ddOrigin.value, + "Additional `rum` header value should be injected when both Tracing and RUM instrumentations are enabled." ) }