From 796b57e503894a1c247ed253148049c1fdfdcc3e Mon Sep 17 00:00:00 2001 From: Jordan Wong <61242306+jordan-wong@users.noreply.github.com> Date: Wed, 7 May 2025 13:37:58 -0400 Subject: [PATCH 1/2] Implement APIGW Inferred Proxy Spans (#8336) * Creates inferred proxy spans for API Gateway calls via presence of http headers --------- Co-authored-by: Zarir Hamza --- .../datadog/context/InferredProxyContext.java | 50 ++ .../propagation/InferredProxyPropagator.java | 74 +++ .../context/InferredProxyHandlingTest.java | 465 ++++++++++++++++++ .../decorator/HttpServerDecorator.java | 451 +++++++++++++++++ .../trace/api/config/TracerConfig.java | 3 + .../java/datadog/trace/core/CoreTracer.java | 5 + .../main/java/datadog/trace/api/Config.java | 7 + .../instrumentation/api/AgentPropagation.java | 2 +- 8 files changed, 1056 insertions(+), 1 deletion(-) create mode 100644 components/context/src/main/java/datadog/context/InferredProxyContext.java create mode 100644 components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java create mode 100644 components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java diff --git a/components/context/src/main/java/datadog/context/InferredProxyContext.java b/components/context/src/main/java/datadog/context/InferredProxyContext.java new file mode 100644 index 00000000000..51eecc4cc02 --- /dev/null +++ b/components/context/src/main/java/datadog/context/InferredProxyContext.java @@ -0,0 +1,50 @@ +package datadog.context; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class InferredProxyContext implements ImplicitContextKeyed { + public static final ContextKey CONTEXT_KEY = + ContextKey.named("inferred-proxy-key"); + private final Map inferredProxy; + + public static InferredProxyContext fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + public InferredProxyContext(Map contextInfo) { + this.inferredProxy = + (contextInfo == null || contextInfo.isEmpty()) + ? new HashMap<>() + : new HashMap<>(contextInfo); + } + + public InferredProxyContext() { + this.inferredProxy = new HashMap<>(); + } + + public Map getInferredProxyContext() { + return Collections.unmodifiableMap(inferredProxy); + } + + public void putInferredProxyInfo(String key, String value) { + inferredProxy.put(key, value); + } + + public void removeInferredProxyInfo(String key) { + inferredProxy.remove(key); + } + + /** + * Creates a new context with this value under its chosen key. + * + * @param context the context to copy the original values from. + * @return the new context with the implicitly keyed value. + * @see Context#with(ImplicitContextKeyed) + */ + @Override + public Context storeInto(Context context) { + return context.with(CONTEXT_KEY, this); + } +} diff --git a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java new file mode 100644 index 00000000000..69e5a0e896e --- /dev/null +++ b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java @@ -0,0 +1,74 @@ +package datadog.context.propagation; + +import datadog.context.Context; +import datadog.context.InferredProxyContext; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; + +public class InferredProxyPropagator implements Propagator { + public static final String INFERRED_PROXY_KEY = "x-dd-proxy"; + /** + * METHOD STUB: InferredProxy is currently not meant to be injected to downstream services Injects + * a context into a downstream service using the given carrier. + * + * @param context the context containing the values to be injected. + * @param carrier the instance that will receive the key/value pairs to propagate. + * @param setter the callback to set key/value pairs into the carrier. + */ + @Override + public void inject(Context context, C carrier, CarrierSetter setter) {} + + /** + * Extracts a context from un upstream service. + * + * @param context the base context to store the extracted values on top, use {@link + * Context#root()} for a default base context. + * @param carrier the instance to fetch the propagated key/value pairs from. + * @param visitor the callback to walk over the carrier and extract its key/value pais. + * @return A context with the extracted values on top of the given base context. + */ + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + if (context == null || carrier == null || visitor == null) { + return context; + } + InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); + visitor.forEachKeyValue(carrier, extractor); + + InferredProxyContext extractedContext = extractor.extractedContext; + if (extractedContext == null) { + return context; + } + return extractedContext.storeInto(context); + } + + public static class InferredProxyContextExtractor implements BiConsumer { + private InferredProxyContext extractedContext; + + InferredProxyContextExtractor() {} + + private Map parseInferredProxyHeaders(String input) { + Map parsedHeaders = new HashMap<>(); + return parsedHeaders; + } + + /** + * Performs this operation on the given arguments. + * + * @param key the first input argument from an http header + * @param value the second input argument from an http header + */ + @Override + public void accept(String key, String value) { + if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY)) { + return; + } + Map inferredProxyMap = parseInferredProxyHeaders(value); + if (extractedContext == null) { + extractedContext = new InferredProxyContext(); + } + extractedContext.putInferredProxyInfo(key, value); + } + } +} diff --git a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java new file mode 100644 index 00000000000..53ddf5cb12a --- /dev/null +++ b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java @@ -0,0 +1,465 @@ +package datadog.context; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.InferredProxyPropagator; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; // For @Test on nested class methods +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class InferredProxyHandlingTest { + + // Define header key constants locally for the test + static final String PROXY_SYSTEM_KEY = "x-dd-proxy-system"; + static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; + static final String PROXY_PATH_KEY = "x-dd-proxy-path"; + static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; + static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; + + private InferredProxyPropagator propagator; + + @BeforeEach + void setUp() { + propagator = new InferredProxyPropagator(); + } + + // Moved @MethodSource providers to the outer class and made them static + static Stream validHeadersProviderForPropagator() { + Map allStandard = new HashMap<>(); + allStandard.put(PROXY_SYSTEM_KEY, "aws-apigw"); // The only currently supported system + allStandard.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + allStandard.put(PROXY_PATH_KEY, "/foo"); + allStandard.put(PROXY_HTTP_METHOD_KEY, "GET"); + allStandard.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); + + return Stream.of( + Arguments.of( + "all standard headers (aws-apigw)", + allStandard, + "aws-apigw", + "12345", + "/foo", + "GET", + "api.example.com", + null, + null)); + } + + static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + missingSystem.put(PROXY_PATH_KEY, "/foo"); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); + missingTime.put(PROXY_PATH_KEY, "/foo"); + + return Stream.of( + Arguments.of("PROXY_SYSTEM_KEY missing", missingSystem), + Arguments.of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); + } + + // Simple Map visitor for tests (can remain static or non-static in outer class) + static class MapVisitor implements CarrierVisitor> { + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + if (carrier == null) { + return; + } + carrier.forEach(visitor); + } + } + + // Custom visitor to test null key path in the extractor - MOVED HERE and made static + static class NullKeyTestVisitor implements CarrierVisitor> { + private final BiConsumer actualExtractorAccept; + + NullKeyTestVisitor(BiConsumer actualExtractorAccept) { + this.actualExtractorAccept = actualExtractorAccept; + } + + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + if (actualExtractorAccept != null) { + actualExtractorAccept.accept(null, "valueForNullKey"); + } + } + } + + @Nested + @DisplayName("InferredProxyPropagator Tests") + class PropagatorTests { // Kept non-static + + @ParameterizedTest(name = "{0}") + @MethodSource( + "datadog.context.InferredProxyHandlingTest#validHeadersProviderForPropagator") // Fully + // qualified + // name + @DisplayName("Should extract InferredProxyContext when valid headers are present") + void testSuccessfulExtraction( + String description, + Map headers, + String expectedSystem, + String expectedTimeMs, + String expectedPath, + String expectedMethod, + String expectedDomain, + String expectedExtraKey, + String expectedExtraValue) { + + Context rootContext = Context.root(); + // Now accesses the outer class's propagator instance field + Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull( + inferredProxyContext, "InferredProxyContext should not be null for: " + description); + Map actualProxyData = inferredProxyContext.getInferredProxyContext(); + assertEquals(expectedSystem, actualProxyData.get(PROXY_SYSTEM_KEY)); + assertEquals(expectedTimeMs, actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + assertEquals(expectedPath, actualProxyData.get(PROXY_PATH_KEY)); + assertEquals(expectedMethod, actualProxyData.get(PROXY_HTTP_METHOD_KEY)); + assertEquals(expectedDomain, actualProxyData.get(PROXY_DOMAIN_NAME_KEY)); + if (expectedExtraKey != null) { + assertEquals(expectedExtraValue, actualProxyData.get(expectedExtraKey)); + } + } + + @ParameterizedTest(name = "{0}") + @MethodSource( + "datadog.context.InferredProxyHandlingTest#invalidOrMissingHeadersProviderForPropagator") // Fully qualified name + @DisplayName("Should create InferredProxyContext even if some critical headers are missing") + void testExtractionWithMissingCriticalHeaders(String description, Map headers) { + Context rootContext = Context.root(); + Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull( + inferredProxyContext, + "InferredProxyContext should still be created if any x-dd-proxy-* header is present for: " + + description); + Map actualProxyData = inferredProxyContext.getInferredProxyContext(); + + if (headers.containsKey(PROXY_SYSTEM_KEY)) { + assertEquals(headers.get(PROXY_SYSTEM_KEY), actualProxyData.get(PROXY_SYSTEM_KEY)); + } else { + assertNull(actualProxyData.get(PROXY_SYSTEM_KEY)); + } + if (headers.containsKey(PROXY_REQUEST_TIME_MS_KEY)) { + assertEquals( + headers.get(PROXY_REQUEST_TIME_MS_KEY), actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + } else { + assertNull(actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + } + } + + @Test + @DisplayName("Should not extract InferredProxyContext if no relevant headers are present") + void testNoRelevantHeaders() { + Map carrier = new HashMap<>(); + carrier.put("x-unrelated-header", "value"); + carrier.put("another-header", "othervalue"); + Context rootContext = Context.root(); + + Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNull( + inferredProxyContext, + "InferredProxyContext should be null if no x-dd-proxy-* headers are found"); + } + + @Test + @DisplayName("Should return original context if carrier is null") + void testNullCarrier() { + InferredProxyContext initialData = + new InferredProxyContext(Collections.singletonMap("test", "value")); + Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); + + Context extractedOuterContext = propagator.extract(rootContext, null, new MapVisitor()); + + assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); + assertEquals( + "value", + InferredProxyContext.fromContext(extractedOuterContext) + .getInferredProxyContext() + .get("test")); + } + + @Test + @DisplayName("Should return original context if visitor is null") + void testNullVisitor() { + Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); + InferredProxyContext initialData = + new InferredProxyContext(Collections.singletonMap("test", "value")); + Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); + + Context extractedOuterContext = propagator.extract(rootContext, carrier, null); + + assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); + assertEquals( + "value", + InferredProxyContext.fromContext(extractedOuterContext) + .getInferredProxyContext() + .get("test")); + } + + @Test + @DisplayName("Should return original context if context is null") + void testNullContext() { + Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); + Context extractedOuterContext = propagator.extract(null, carrier, new MapVisitor()); + assertNull(extractedOuterContext, "Context should remain null if passed as null"); + } + + @Test + @DisplayName("Extractor should handle multiple proxy headers") + void testMultipleProxyHeaders() { + Map carrier = new HashMap<>(); + carrier.put(PROXY_SYSTEM_KEY, "aws-apigw"); + carrier.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + carrier.put("x-dd-proxy-custom", "value1"); // First proxy header + carrier.put("x-dd-proxy-another", "value2"); // Second proxy header + + Context rootContext = Context.root(); + Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull(inferredProxyContext); + // Check if both headers were stored (covers extractedContext == null being false) + assertEquals( + "value1", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-custom")); + assertEquals( + "value2", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-another")); + assertEquals( + "aws-apigw", inferredProxyContext.getInferredProxyContext().get(PROXY_SYSTEM_KEY)); + } + + @Test + @DisplayName("Extractor accept method should handle null/empty keys") + void testExtractorAcceptNullEmptyKeys() { + Context rootContext = Context.root(); + + // Test null key - HashMap doesn't allow null keys. Standard HTTP visitors + // also typically don't yield null keys. Testing this branch is difficult + // without a custom visitor or modifying the source. Relying on coverage report + // or assuming standard carriers won't provide null keys. + + // Test empty key + Map carrierWithEmptyKey = new HashMap<>(); + carrierWithEmptyKey.put("", "emptyKeyValue"); // Add empty key + carrierWithEmptyKey.put(PROXY_SYSTEM_KEY, "aws-apigw"); // Add a valid key too + + Context contextAfterEmpty = + propagator.extract(rootContext, carrierWithEmptyKey, new MapVisitor()); + InferredProxyContext ipcEmpty = InferredProxyContext.fromContext(contextAfterEmpty); + + // The propagator should ignore the empty key entry entirely. + assertNotNull(ipcEmpty, "Context should be created due to valid key"); + assertNull(ipcEmpty.getInferredProxyContext().get(""), "Empty key should not be stored"); + assertEquals( + "aws-apigw", + ipcEmpty.getInferredProxyContext().get(PROXY_SYSTEM_KEY), + "Valid key should still be stored"); + assertEquals(1, ipcEmpty.getInferredProxyContext().size(), "Only valid key should be stored"); + } + + @Test + @DisplayName( + "Extractor accept method should handle explicitly passed null key via custom visitor") + void testExtractorAcceptExplicitNullKey() { + Context rootContext = Context.root(); + Map carrier = new HashMap<>(); // Carrier can be empty for this test + + // We need to get a handle to the internal BiConsumer (the InferredProxyContextExtractor + // instance). + // The extract method will create one. We can pass a visitor that captures it. + + final BiConsumer[] extractorHolder = new BiConsumer[1]; + + CarrierVisitor> capturingVisitor = + (cr, bic) -> { + extractorHolder[0] = bic; // Capture the BiConsumer + // Optionally, call the original MapVisitor if we still want normal processing after + // capture + // new MapVisitor().forEachKeyValue(cr, bic); + }; + + // This first call is primarily to get a reference to the internal extractor + propagator.extract(rootContext, carrier, capturingVisitor); + + assertNotNull(extractorHolder[0], "Failed to capture the internal extractor instance"); + + // Now use a new custom visitor to specifically test the null key path + // on the captured extractor instance (though this isn't how extract is typically used). + // A more direct way to test the BiConsumer if it were accessible or if the design allowed it. + // For now, we directly call accept on the captured one. + extractorHolder[0].accept(null, "valueForNullKey"); + + // The goal is JaCoCo coverage. Asserting internal state of the extractor is hard without + // reflection. + // We can verify that the context remains unchanged or as expected if no valid headers + // processed. + InferredProxyContext ipc = + InferredProxyContext.fromContext( + rootContext); // or context returned by a second extract call + assertNull(ipc, "Context should not have InferredProxyContext from only a null key call"); + } + } + + @Nested + @DisplayName("InferredProxyContext Tests") + class ContextUnitTests { + + @Test + @DisplayName("Default constructor should create an empty context map") + void testDefaultConstructor() { + InferredProxyContext ipc = new InferredProxyContext(); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + + @Test + @DisplayName("Constructor with map should initialize context map") + void testMapConstructor() { + Map initialData = new HashMap<>(); + initialData.put("key1", "value1"); + initialData.put("key2", "value2"); + + InferredProxyContext ipc = new InferredProxyContext(initialData); + assertNotNull(ipc.getInferredProxyContext()); + assertEquals(2, ipc.getInferredProxyContext().size()); + assertEquals("value1", ipc.getInferredProxyContext().get("key1")); + assertEquals("value2", ipc.getInferredProxyContext().get("key2")); + + initialData.put("key3", "value3"); // Modify original map + assertNull(ipc.getInferredProxyContext().get("key3"), "Internal map should be a copy"); + } + + @Test + @DisplayName("putInferredProxyInfo should add to the context map") + void testPutInfo() { + InferredProxyContext ipc = new InferredProxyContext(); + ipc.putInferredProxyInfo("system", "aws-apigw"); + ipc.putInferredProxyInfo("time", "12345"); + + Map contextMap = ipc.getInferredProxyContext(); + assertEquals(2, contextMap.size()); + assertEquals("aws-apigw", contextMap.get("system")); + assertEquals("12345", contextMap.get("time")); + + ipc.putInferredProxyInfo("system", "azure-func"); // Overwrite + assertEquals("azure-func", contextMap.get("system")); + assertEquals(2, contextMap.size()); + } + + @Test + @DisplayName("removeInferredProxyInfo should remove from the context map") + void testRemoveInfo() { + Map initialData = new HashMap<>(); + initialData.put("key1", "value1"); + initialData.put("key2", "value2"); + InferredProxyContext ipc = new InferredProxyContext(initialData); + + ipc.removeInferredProxyInfo("key1"); + Map contextMap = ipc.getInferredProxyContext(); + assertEquals(1, contextMap.size()); + assertNull(contextMap.get("key1")); + assertEquals("value2", contextMap.get("key2")); + + ipc.removeInferredProxyInfo("nonexistent"); // Remove non-existent + assertEquals(1, contextMap.size()); + } + + @Test + @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") + void testStoreAndFromContext() { + InferredProxyContext ipcToStore = new InferredProxyContext(); + ipcToStore.putInferredProxyInfo("customKey", "customValue"); + + Context rootContext = Context.root(); + Context contextWithValue = ipcToStore.storeInto(rootContext); + assertNotNull(contextWithValue); + + InferredProxyContext retrievedIpc = InferredProxyContext.fromContext(contextWithValue); + assertNotNull(retrievedIpc); + assertEquals("customValue", retrievedIpc.getInferredProxyContext().get("customKey")); + + assertNull( + InferredProxyContext.fromContext(rootContext), + "Original root context should not be affected"); + + Context cleanContext = Context.root(); + assertNull( + InferredProxyContext.fromContext(cleanContext), + "fromContext on clean context should be null"); + } + + @Test + @DisplayName("getInferredProxyContext should return an unmodifiable map or a copy") + void testGetInferredProxyContextImmutability() { + InferredProxyContext ipc = new InferredProxyContext(); + ipc.putInferredProxyInfo("key1", "value1"); + + Map retrievedMap = ipc.getInferredProxyContext(); + assertNotNull(retrievedMap); + assertEquals("value1", retrievedMap.get("key1")); + + boolean threwUnsupported = false; + try { + retrievedMap.put("newKey", "newValue"); + } catch (UnsupportedOperationException e) { + threwUnsupported = true; + } + // Depending on whether InferredProxyContext.getInferredProxyContext() returns a direct + // reference or a copy, + // this assertion might change. If it returns a direct mutable reference, threwUnsupported + // will be false. + // If it returns an unmodifiable view or a copy, attempts to modify might throw or simply not + // affect the original. + // For now, we check that the original context was not changed. + assertEquals( + 1, ipc.getInferredProxyContext().size(), "Internal map size should remain unchanged"); + assertEquals( + "value1", + ipc.getInferredProxyContext().get("key1"), + "Internal map content should remain unchanged"); + // If it MUST be unmodifiable, add: assertTrue(threwUnsupported, "Retrieved map should be + // unmodifiable"); + } + + @Test + @DisplayName("Constructor with null map should create an empty context map") + void testNullMapConstructor() { + InferredProxyContext ipc = new InferredProxyContext(null); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + + @Test + @DisplayName("Constructor with empty map should create an empty context map") + void testEmptyMapConstructor() { + Map emptyMap = Collections.emptyMap(); + InferredProxyContext ipc = new InferredProxyContext(emptyMap); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 25684aeb42e..5d8a7ee5974 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -11,8 +11,12 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.context.Context; import datadog.context.propagation.Propagators; +import datadog.context.InferredProxyContext; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.DDTags; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.TraceConfig; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; import datadog.trace.api.gateway.BlockResponseFunction; @@ -22,10 +26,12 @@ import datadog.trace.api.gateway.IGSpanInfo; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.interceptor.MutableSpan; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -38,6 +44,8 @@ import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver; import java.net.InetAddress; import java.util.BitSet; +import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.Locale; import java.util.Map; import java.util.concurrent.ExecutionException; @@ -52,9 +60,401 @@ public abstract class HttpServerDecorator extends ServerDecorator { + class InferredProxySpanGroup implements AgentSpan { + private final AgentSpan inferredProxySpan; + private final AgentSpan serverSpan; + + InferredProxySpanGroup(AgentSpan inferredProxySpan, AgentSpan serverSpan) { + this.inferredProxySpan = inferredProxySpan; + this.serverSpan = serverSpan; + } + + @Override + public DDTraceId getTraceId() { + return serverSpan.getTraceId(); + } + + @Override + public long getSpanId() { + return serverSpan.getSpanId(); + } + + @Override + public AgentSpan setTag(String key, boolean value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, int value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, long value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, double value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, String value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, CharSequence value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, Object value) { + return serverSpan.setTag(key, value); + } + + /** + * @param map + * @return + */ + @Override + public AgentSpan setAllTags(Map map) { + return null; + } + + @Override + public AgentSpan setTag(String key, Number value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, int value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, long value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, double value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setSpanType(CharSequence type) { + return serverSpan.setSpanType(type); + } + + @Override + public Object getTag(String key) { + return serverSpan.getTag(key); + } + + @Override + public AgentSpan setError(boolean error) { + serverSpan.setError(error); + if (inferredProxySpan != null) { + inferredProxySpan.setError(error); + } + return this; + } + + @Override + public AgentSpan setError(boolean error, byte priority) { + serverSpan.setError(error, priority); + if (inferredProxySpan != null) { + inferredProxySpan.setError(error, priority); + } + return this; + } + + @Override + public AgentSpan setMeasured(boolean measured) { + return serverSpan.setMeasured(measured); + } + + @Override + public AgentSpan setErrorMessage(String errorMessage) { + return serverSpan.setErrorMessage(errorMessage); + } + + @Override + public AgentSpan addThrowable(Throwable throwable) { + serverSpan.addThrowable(throwable); + if (inferredProxySpan != null) { + inferredProxySpan.addThrowable(throwable); + } + return this; + } + + @Override + public AgentSpan addThrowable(Throwable throwable, byte errorPriority) { + serverSpan.addThrowable(throwable, errorPriority); + if (inferredProxySpan != null) { + inferredProxySpan.addThrowable(throwable, errorPriority); + } + return this; + } + + @Override + public AgentSpan getLocalRootSpan() { + return serverSpan.getLocalRootSpan(); + } + + @Override + public boolean isSameTrace(AgentSpan otherSpan) { + return serverSpan.isSameTrace(otherSpan); + } + + @Override + public AgentSpanContext context() { + return serverSpan.context(); + } + + @Override + public String getBaggageItem(String key) { + return serverSpan.getBaggageItem(key); + } + + @Override + public AgentSpan setBaggageItem(String key, String value) { + return serverSpan.setBaggageItem(key, value); + } + + @Override + public AgentSpan setHttpStatusCode(int statusCode) { + serverSpan.setHttpStatusCode(statusCode); + if (inferredProxySpan != null) { + inferredProxySpan.setHttpStatusCode(statusCode); + } + return this; + } + + @Override + public short getHttpStatusCode() { + return serverSpan.getHttpStatusCode(); + } + + @Override + public void finish() { + serverSpan.finish(); + if (inferredProxySpan != null) { + inferredProxySpan.finish(); + } + } + + @Override + public void finish(long finishMicros) { + serverSpan.finish(finishMicros); + if (inferredProxySpan != null) { + inferredProxySpan.finish(finishMicros); + } + } + + @Override + public void finishWithDuration(long durationNanos) { + serverSpan.finishWithDuration(durationNanos); + if (inferredProxySpan != null) { + inferredProxySpan.finishWithDuration(durationNanos); + } + } + + @Override + public void beginEndToEnd() { + serverSpan.beginEndToEnd(); + } + + @Override + public void finishWithEndToEnd() { + serverSpan.finishWithEndToEnd(); + if (inferredProxySpan != null) { + inferredProxySpan.finishWithEndToEnd(); + } + } + + @Override + public boolean phasedFinish() { + final boolean ret = serverSpan.phasedFinish(); + if (inferredProxySpan != null) { + inferredProxySpan.phasedFinish(); + } + return ret; + } + + @Override + public void publish() { + serverSpan.publish(); + } + + @Override + public CharSequence getSpanName() { + return serverSpan.getSpanName(); + } + + @Override + public void setSpanName(CharSequence spanName) { + serverSpan.setSpanName(spanName); + } + + @Deprecated + @Override + public boolean hasResourceName() { + return serverSpan.hasResourceName(); + } + + @Override + public byte getResourceNamePriority() { + return serverSpan.getResourceNamePriority(); + } + + @Override + public AgentSpan setResourceName(CharSequence resourceName) { + return serverSpan.setResourceName(resourceName); + } + + @Override + public AgentSpan setResourceName(CharSequence resourceName, byte priority) { + return serverSpan.setResourceName(resourceName, priority); + } + + @Override + public RequestContext getRequestContext() { + return serverSpan.getRequestContext(); + } + + @Override + public Integer forceSamplingDecision() { + return serverSpan.forceSamplingDecision(); + } + + @Override + public AgentSpan setSamplingPriority(int newPriority, int samplingMechanism) { + return serverSpan.setSamplingPriority(newPriority, samplingMechanism); + } + + @Override + public TraceConfig traceConfig() { + return serverSpan.traceConfig(); + } + + @Override + public void addLink(AgentSpanLink link) { + serverSpan.addLink(link); + } + + @Override + public AgentSpan setMetaStruct(String field, Object value) { + return serverSpan.setMetaStruct(field, value); + } + + @Override + public boolean isOutbound() { + return serverSpan.isOutbound(); + } + + @Override + public AgentSpan asAgentSpan() { + return serverSpan.asAgentSpan(); + } + + @Override + public long getStartTime() { + return serverSpan.getStartTime(); + } + + @Override + public long getDurationNano() { + return serverSpan.getDurationNano(); + } + + @Override + public CharSequence getOperationName() { + return serverSpan.getOperationName(); + } + + @Override + public MutableSpan setOperationName(CharSequence serviceName) { + return serverSpan.setOperationName(serviceName); + } + + @Override + public String getServiceName() { + return serverSpan.getServiceName(); + } + + @Override + public MutableSpan setServiceName(String serviceName) { + return serverSpan.setServiceName(serviceName); + } + + @Override + public CharSequence getResourceName() { + return serverSpan.getResourceName(); + } + + @Override + public Integer getSamplingPriority() { + return serverSpan.getSamplingPriority(); + } + + @Deprecated + @Override + public MutableSpan setSamplingPriority(int newPriority) { + return serverSpan.setSamplingPriority(newPriority); + } + + @Override + public String getSpanType() { + return serverSpan.getSpanType(); + } + + @Override + public Map getTags() { + return serverSpan.getTags(); + } + + @Override + public boolean isError() { + return serverSpan.isError(); + } + + @Deprecated + @Override + public MutableSpan getRootSpan() { + return serverSpan.getRootSpan(); + } + + @Override + public void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba) { + serverSpan.setRequestBlockingAction(rba); + } + + @Override + public Flow.Action.RequestBlockingAction getRequestBlockingAction() { + return serverSpan.getRequestBlockingAction(); + } + } + private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class); private static final int UNSET_PORT = 0; + public static final String PROXY_SYSTEM = "x-dd-proxy"; + public static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; + public static final String PROXY_PATH = "x-dd-proxy-path"; + public static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; + public static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; + public static final String STAGE = "x-dd-proxy-stage"; + + public static final Map SUPPORTED_PROXIES; + + static { + SUPPORTED_PROXIES = new HashMap<>(); + SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); + } + public static final String DD_SPAN_ATTRIBUTE = "datadog.span"; public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch"; public static final String DD_RUM_INJECTED = "datadog.rum.injected"; @@ -161,6 +561,57 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { return context.with(span); } + private AgentSpan startSpanWithInferredProxy( + String instrumentationName, + datadog.context.Context fullContextForInferredProxy, + AgentSpanContext.Extracted standardExtractedContext) { + + InferredProxyContext inferredProxy = + InferredProxyContext.fromContext(fullContextForInferredProxy); + + if (inferredProxy == null) { + return null; + } + + Map headers = inferredProxy.getInferredProxyContext(); + + // Check if timestamp and proxy system are present + String startTimeStr = headers.get(PROXY_START_TIME_MS); + String proxySystem = headers.get(PROXY_SYSTEM); + + if (startTimeStr == null + || proxySystem == null + || !SUPPORTED_PROXIES.containsKey(proxySystem)) { + return null; + } + + long startTime; + try { + startTime = Long.parseLong(startTimeStr) * 1000; // Convert to microseconds + } catch (NumberFormatException e) { + return null; // Invalid timestamp + } + + AgentSpan apiGtwSpan = + tracer() + .startSpan( + "inferred_proxy", + SUPPORTED_PROXIES.get(proxySystem), + callIGCallbackStart(standardExtractedContext), + startTime); + + apiGtwSpan.setTag(Tags.COMPONENT, proxySystem); + apiGtwSpan.setTag( + DDTags.RESOURCE_NAME, headers.get(PROXY_HTTP_METHOD) + " " + headers.get(PROXY_PATH)); + apiGtwSpan.setTag(DDTags.SERVICE_NAME, headers.get(PROXY_DOMAIN_NAME)); + apiGtwSpan.setTag(DDTags.SPAN_TYPE, "web"); + apiGtwSpan.setTag(Tags.HTTP_METHOD, headers.get(PROXY_HTTP_METHOD)); + apiGtwSpan.setTag(Tags.HTTP_URL, headers.get(PROXY_DOMAIN_NAME) + headers.get(PROXY_PATH)); + apiGtwSpan.setTag("stage", headers.get(STAGE)); + apiGtwSpan.setTag("_dd.inferred_span", 1); + return apiGtwSpan; + } + public AgentSpan onRequest( final AgentSpan span, final CONNECTION connection, diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java index 0300cef2070..ae5d515cb62 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java @@ -100,6 +100,9 @@ public final class TracerConfig { public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes"; public static final String TRACE_BAGGAGE_TAG_KEYS = "trace.baggage.tag.keys"; + public static final String TRACE_INFERRED_PROXY_SERVICES_ENABLED = + "trace.inferred.proxy.services.enabled"; + public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled"; public static final String CLIENT_IP_ENABLED = "trace.client-ip.enabled"; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 4c81a579f83..56aeedc6e71 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -7,6 +7,7 @@ import static datadog.trace.api.TracePropagationBehaviorExtract.IGNORE; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.BAGGAGE_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; +import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.INFERRED_PROXY_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN; import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator; @@ -22,6 +23,7 @@ import datadog.communication.ddagent.SharedCommunicationObjects; import datadog.communication.monitor.Monitoring; import datadog.communication.monitor.Recording; +import datadog.context.propagation.InferredProxyPropagator; import datadog.context.propagation.Propagators; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; @@ -820,6 +822,9 @@ private CoreTracer( && config.getTracePropagationBehaviorExtract() != IGNORE) { Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator(config)); } + if (config.isInferredProxyPropagationEnabled()) { + Propagators.register(INFERRED_PROXY_CONCERN, new InferredProxyPropagator()); + } if (config.isCiVisibilityEnabled()) { if (config.isCiVisibilityTraceSanitationEnabled()) { diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index aae53ec7e07..a2e58598ea5 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -827,6 +827,7 @@ public static String getHostName() { private final int traceBaggageMaxItems; private final int traceBaggageMaxBytes; private final List traceBaggageTagKeys; + private final boolean traceInferredProxyEnabled; private final int clockSyncPeriod; private final boolean logsInjectionEnabled; @@ -1730,6 +1731,8 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins tracePropagationExtractFirst = configProvider.getBoolean( TRACE_PROPAGATION_EXTRACT_FIRST, DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST); + traceInferredProxyEnabled = + configProvider.getBoolean(TRACE_INFERRED_PROXY_SERVICES_ENABLED, false); clockSyncPeriod = configProvider.getInteger(CLOCK_SYNC_PERIOD, DEFAULT_CLOCK_SYNC_PERIOD); @@ -3118,6 +3121,10 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public boolean isInferredProxyPropagationEnabled() { + return traceInferredProxyEnabled; + } + public boolean isBaggageExtract() { return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java index fba659bea74..0fbb4484353 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java @@ -18,7 +18,7 @@ public final class AgentPropagation { // TODO: remove this priority once we have a story for replacing TagContext with the Context API public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105); public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray"); - + public static final Concern INFERRED_PROXY_CONCERN = named("inferred-proxy"); // TODO DSM propagator should run after the other propagators as it stores the pathway context // TODO into the span context for now. Remove priority after the migration is complete. public static final Concern DSM_CONCERN = withPriority("data-stream-monitoring", 110); From e8c2bafaff3aa416da5b3aa456101ccb0f18092b Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Mon, 12 May 2025 11:26:35 +0200 Subject: [PATCH 2/2] feat(serverless): Fix gateway inferred span design Avoid duplicate expensive context extraction Avoid subclassing tracing span for serverless but used serverless context element instead to store / track inferred span while keep tracing feature untouched Improved propagator to not create / capture inferred span context element on invalid data Rework context element to hold the inferred spans and its captured data Release captured data as soon as they span start (never read after this point so reclaiming memory) Refactor context element and propagator into the right package, not context component (product / feature agnostic) Refactor unit tests --- .github/CODEOWNERS | 6 +- .../datadog/context/InferredProxyContext.java | 50 -- .../propagation/InferredProxyPropagator.java | 74 --- .../context/InferredProxyHandlingTest.java | 465 ----------------- .../decorator/HttpServerDecorator.java | 481 +----------------- .../java/datadog/trace/core/CoreTracer.java | 2 +- .../propagation/InferredProxyPropagator.java | 60 +++ .../InferredProxyPropagatorTests.java | 96 ++++ .../main/java/datadog/trace/api/Config.java | 3 + .../trace/api/gateway/InferredProxySpan.java | 108 ++++ .../api/gateway/InferredProxySpanTests.java | 95 ++++ 11 files changed, 393 insertions(+), 1047 deletions(-) delete mode 100644 components/context/src/main/java/datadog/context/InferredProxyContext.java delete mode 100644 components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java delete mode 100644 components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java create mode 100644 dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java create mode 100644 internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java create mode 100644 internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ef437eebae2..96ab0c6452e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -28,8 +28,10 @@ /internal-api/src/test/groovy/datadog/trace/api/sampling @DataDog/apm-sdk-api-java # @DataDog/apm-serverless -/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless -/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless +**/InferredProxy*.java @DataDog/apm-serverless +**/InferredProxy*.groovy @DataDog/apm-serverless # @DataDog/apm-lang-platform-java /.circleci/ @DataDog/apm-lang-platform-java diff --git a/components/context/src/main/java/datadog/context/InferredProxyContext.java b/components/context/src/main/java/datadog/context/InferredProxyContext.java deleted file mode 100644 index 51eecc4cc02..00000000000 --- a/components/context/src/main/java/datadog/context/InferredProxyContext.java +++ /dev/null @@ -1,50 +0,0 @@ -package datadog.context; - -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; - -public class InferredProxyContext implements ImplicitContextKeyed { - public static final ContextKey CONTEXT_KEY = - ContextKey.named("inferred-proxy-key"); - private final Map inferredProxy; - - public static InferredProxyContext fromContext(Context context) { - return context.get(CONTEXT_KEY); - } - - public InferredProxyContext(Map contextInfo) { - this.inferredProxy = - (contextInfo == null || contextInfo.isEmpty()) - ? new HashMap<>() - : new HashMap<>(contextInfo); - } - - public InferredProxyContext() { - this.inferredProxy = new HashMap<>(); - } - - public Map getInferredProxyContext() { - return Collections.unmodifiableMap(inferredProxy); - } - - public void putInferredProxyInfo(String key, String value) { - inferredProxy.put(key, value); - } - - public void removeInferredProxyInfo(String key) { - inferredProxy.remove(key); - } - - /** - * Creates a new context with this value under its chosen key. - * - * @param context the context to copy the original values from. - * @return the new context with the implicitly keyed value. - * @see Context#with(ImplicitContextKeyed) - */ - @Override - public Context storeInto(Context context) { - return context.with(CONTEXT_KEY, this); - } -} diff --git a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java deleted file mode 100644 index 69e5a0e896e..00000000000 --- a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java +++ /dev/null @@ -1,74 +0,0 @@ -package datadog.context.propagation; - -import datadog.context.Context; -import datadog.context.InferredProxyContext; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; - -public class InferredProxyPropagator implements Propagator { - public static final String INFERRED_PROXY_KEY = "x-dd-proxy"; - /** - * METHOD STUB: InferredProxy is currently not meant to be injected to downstream services Injects - * a context into a downstream service using the given carrier. - * - * @param context the context containing the values to be injected. - * @param carrier the instance that will receive the key/value pairs to propagate. - * @param setter the callback to set key/value pairs into the carrier. - */ - @Override - public void inject(Context context, C carrier, CarrierSetter setter) {} - - /** - * Extracts a context from un upstream service. - * - * @param context the base context to store the extracted values on top, use {@link - * Context#root()} for a default base context. - * @param carrier the instance to fetch the propagated key/value pairs from. - * @param visitor the callback to walk over the carrier and extract its key/value pais. - * @return A context with the extracted values on top of the given base context. - */ - @Override - public Context extract(Context context, C carrier, CarrierVisitor visitor) { - if (context == null || carrier == null || visitor == null) { - return context; - } - InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); - visitor.forEachKeyValue(carrier, extractor); - - InferredProxyContext extractedContext = extractor.extractedContext; - if (extractedContext == null) { - return context; - } - return extractedContext.storeInto(context); - } - - public static class InferredProxyContextExtractor implements BiConsumer { - private InferredProxyContext extractedContext; - - InferredProxyContextExtractor() {} - - private Map parseInferredProxyHeaders(String input) { - Map parsedHeaders = new HashMap<>(); - return parsedHeaders; - } - - /** - * Performs this operation on the given arguments. - * - * @param key the first input argument from an http header - * @param value the second input argument from an http header - */ - @Override - public void accept(String key, String value) { - if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY)) { - return; - } - Map inferredProxyMap = parseInferredProxyHeaders(value); - if (extractedContext == null) { - extractedContext = new InferredProxyContext(); - } - extractedContext.putInferredProxyInfo(key, value); - } - } -} diff --git a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java deleted file mode 100644 index 53ddf5cb12a..00000000000 --- a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java +++ /dev/null @@ -1,465 +0,0 @@ -package datadog.context; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import datadog.context.propagation.CarrierVisitor; -import datadog.context.propagation.InferredProxyPropagator; -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; -import java.util.stream.Stream; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Nested; -import org.junit.jupiter.api.Test; // For @Test on nested class methods -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; - -class InferredProxyHandlingTest { - - // Define header key constants locally for the test - static final String PROXY_SYSTEM_KEY = "x-dd-proxy-system"; - static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; - static final String PROXY_PATH_KEY = "x-dd-proxy-path"; - static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; - static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; - - private InferredProxyPropagator propagator; - - @BeforeEach - void setUp() { - propagator = new InferredProxyPropagator(); - } - - // Moved @MethodSource providers to the outer class and made them static - static Stream validHeadersProviderForPropagator() { - Map allStandard = new HashMap<>(); - allStandard.put(PROXY_SYSTEM_KEY, "aws-apigw"); // The only currently supported system - allStandard.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - allStandard.put(PROXY_PATH_KEY, "/foo"); - allStandard.put(PROXY_HTTP_METHOD_KEY, "GET"); - allStandard.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); - - return Stream.of( - Arguments.of( - "all standard headers (aws-apigw)", - allStandard, - "aws-apigw", - "12345", - "/foo", - "GET", - "api.example.com", - null, - null)); - } - - static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed - Map missingSystem = new HashMap<>(); - missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - missingSystem.put(PROXY_PATH_KEY, "/foo"); - - Map missingTime = new HashMap<>(); - missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); - missingTime.put(PROXY_PATH_KEY, "/foo"); - - return Stream.of( - Arguments.of("PROXY_SYSTEM_KEY missing", missingSystem), - Arguments.of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); - } - - // Simple Map visitor for tests (can remain static or non-static in outer class) - static class MapVisitor implements CarrierVisitor> { - @Override - public void forEachKeyValue(Map carrier, BiConsumer visitor) { - if (carrier == null) { - return; - } - carrier.forEach(visitor); - } - } - - // Custom visitor to test null key path in the extractor - MOVED HERE and made static - static class NullKeyTestVisitor implements CarrierVisitor> { - private final BiConsumer actualExtractorAccept; - - NullKeyTestVisitor(BiConsumer actualExtractorAccept) { - this.actualExtractorAccept = actualExtractorAccept; - } - - @Override - public void forEachKeyValue(Map carrier, BiConsumer visitor) { - if (actualExtractorAccept != null) { - actualExtractorAccept.accept(null, "valueForNullKey"); - } - } - } - - @Nested - @DisplayName("InferredProxyPropagator Tests") - class PropagatorTests { // Kept non-static - - @ParameterizedTest(name = "{0}") - @MethodSource( - "datadog.context.InferredProxyHandlingTest#validHeadersProviderForPropagator") // Fully - // qualified - // name - @DisplayName("Should extract InferredProxyContext when valid headers are present") - void testSuccessfulExtraction( - String description, - Map headers, - String expectedSystem, - String expectedTimeMs, - String expectedPath, - String expectedMethod, - String expectedDomain, - String expectedExtraKey, - String expectedExtraValue) { - - Context rootContext = Context.root(); - // Now accesses the outer class's propagator instance field - Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull( - inferredProxyContext, "InferredProxyContext should not be null for: " + description); - Map actualProxyData = inferredProxyContext.getInferredProxyContext(); - assertEquals(expectedSystem, actualProxyData.get(PROXY_SYSTEM_KEY)); - assertEquals(expectedTimeMs, actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - assertEquals(expectedPath, actualProxyData.get(PROXY_PATH_KEY)); - assertEquals(expectedMethod, actualProxyData.get(PROXY_HTTP_METHOD_KEY)); - assertEquals(expectedDomain, actualProxyData.get(PROXY_DOMAIN_NAME_KEY)); - if (expectedExtraKey != null) { - assertEquals(expectedExtraValue, actualProxyData.get(expectedExtraKey)); - } - } - - @ParameterizedTest(name = "{0}") - @MethodSource( - "datadog.context.InferredProxyHandlingTest#invalidOrMissingHeadersProviderForPropagator") // Fully qualified name - @DisplayName("Should create InferredProxyContext even if some critical headers are missing") - void testExtractionWithMissingCriticalHeaders(String description, Map headers) { - Context rootContext = Context.root(); - Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull( - inferredProxyContext, - "InferredProxyContext should still be created if any x-dd-proxy-* header is present for: " - + description); - Map actualProxyData = inferredProxyContext.getInferredProxyContext(); - - if (headers.containsKey(PROXY_SYSTEM_KEY)) { - assertEquals(headers.get(PROXY_SYSTEM_KEY), actualProxyData.get(PROXY_SYSTEM_KEY)); - } else { - assertNull(actualProxyData.get(PROXY_SYSTEM_KEY)); - } - if (headers.containsKey(PROXY_REQUEST_TIME_MS_KEY)) { - assertEquals( - headers.get(PROXY_REQUEST_TIME_MS_KEY), actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - } else { - assertNull(actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - } - } - - @Test - @DisplayName("Should not extract InferredProxyContext if no relevant headers are present") - void testNoRelevantHeaders() { - Map carrier = new HashMap<>(); - carrier.put("x-unrelated-header", "value"); - carrier.put("another-header", "othervalue"); - Context rootContext = Context.root(); - - Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNull( - inferredProxyContext, - "InferredProxyContext should be null if no x-dd-proxy-* headers are found"); - } - - @Test - @DisplayName("Should return original context if carrier is null") - void testNullCarrier() { - InferredProxyContext initialData = - new InferredProxyContext(Collections.singletonMap("test", "value")); - Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); - - Context extractedOuterContext = propagator.extract(rootContext, null, new MapVisitor()); - - assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); - assertEquals( - "value", - InferredProxyContext.fromContext(extractedOuterContext) - .getInferredProxyContext() - .get("test")); - } - - @Test - @DisplayName("Should return original context if visitor is null") - void testNullVisitor() { - Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); - InferredProxyContext initialData = - new InferredProxyContext(Collections.singletonMap("test", "value")); - Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); - - Context extractedOuterContext = propagator.extract(rootContext, carrier, null); - - assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); - assertEquals( - "value", - InferredProxyContext.fromContext(extractedOuterContext) - .getInferredProxyContext() - .get("test")); - } - - @Test - @DisplayName("Should return original context if context is null") - void testNullContext() { - Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); - Context extractedOuterContext = propagator.extract(null, carrier, new MapVisitor()); - assertNull(extractedOuterContext, "Context should remain null if passed as null"); - } - - @Test - @DisplayName("Extractor should handle multiple proxy headers") - void testMultipleProxyHeaders() { - Map carrier = new HashMap<>(); - carrier.put(PROXY_SYSTEM_KEY, "aws-apigw"); - carrier.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - carrier.put("x-dd-proxy-custom", "value1"); // First proxy header - carrier.put("x-dd-proxy-another", "value2"); // Second proxy header - - Context rootContext = Context.root(); - Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull(inferredProxyContext); - // Check if both headers were stored (covers extractedContext == null being false) - assertEquals( - "value1", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-custom")); - assertEquals( - "value2", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-another")); - assertEquals( - "aws-apigw", inferredProxyContext.getInferredProxyContext().get(PROXY_SYSTEM_KEY)); - } - - @Test - @DisplayName("Extractor accept method should handle null/empty keys") - void testExtractorAcceptNullEmptyKeys() { - Context rootContext = Context.root(); - - // Test null key - HashMap doesn't allow null keys. Standard HTTP visitors - // also typically don't yield null keys. Testing this branch is difficult - // without a custom visitor or modifying the source. Relying on coverage report - // or assuming standard carriers won't provide null keys. - - // Test empty key - Map carrierWithEmptyKey = new HashMap<>(); - carrierWithEmptyKey.put("", "emptyKeyValue"); // Add empty key - carrierWithEmptyKey.put(PROXY_SYSTEM_KEY, "aws-apigw"); // Add a valid key too - - Context contextAfterEmpty = - propagator.extract(rootContext, carrierWithEmptyKey, new MapVisitor()); - InferredProxyContext ipcEmpty = InferredProxyContext.fromContext(contextAfterEmpty); - - // The propagator should ignore the empty key entry entirely. - assertNotNull(ipcEmpty, "Context should be created due to valid key"); - assertNull(ipcEmpty.getInferredProxyContext().get(""), "Empty key should not be stored"); - assertEquals( - "aws-apigw", - ipcEmpty.getInferredProxyContext().get(PROXY_SYSTEM_KEY), - "Valid key should still be stored"); - assertEquals(1, ipcEmpty.getInferredProxyContext().size(), "Only valid key should be stored"); - } - - @Test - @DisplayName( - "Extractor accept method should handle explicitly passed null key via custom visitor") - void testExtractorAcceptExplicitNullKey() { - Context rootContext = Context.root(); - Map carrier = new HashMap<>(); // Carrier can be empty for this test - - // We need to get a handle to the internal BiConsumer (the InferredProxyContextExtractor - // instance). - // The extract method will create one. We can pass a visitor that captures it. - - final BiConsumer[] extractorHolder = new BiConsumer[1]; - - CarrierVisitor> capturingVisitor = - (cr, bic) -> { - extractorHolder[0] = bic; // Capture the BiConsumer - // Optionally, call the original MapVisitor if we still want normal processing after - // capture - // new MapVisitor().forEachKeyValue(cr, bic); - }; - - // This first call is primarily to get a reference to the internal extractor - propagator.extract(rootContext, carrier, capturingVisitor); - - assertNotNull(extractorHolder[0], "Failed to capture the internal extractor instance"); - - // Now use a new custom visitor to specifically test the null key path - // on the captured extractor instance (though this isn't how extract is typically used). - // A more direct way to test the BiConsumer if it were accessible or if the design allowed it. - // For now, we directly call accept on the captured one. - extractorHolder[0].accept(null, "valueForNullKey"); - - // The goal is JaCoCo coverage. Asserting internal state of the extractor is hard without - // reflection. - // We can verify that the context remains unchanged or as expected if no valid headers - // processed. - InferredProxyContext ipc = - InferredProxyContext.fromContext( - rootContext); // or context returned by a second extract call - assertNull(ipc, "Context should not have InferredProxyContext from only a null key call"); - } - } - - @Nested - @DisplayName("InferredProxyContext Tests") - class ContextUnitTests { - - @Test - @DisplayName("Default constructor should create an empty context map") - void testDefaultConstructor() { - InferredProxyContext ipc = new InferredProxyContext(); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - - @Test - @DisplayName("Constructor with map should initialize context map") - void testMapConstructor() { - Map initialData = new HashMap<>(); - initialData.put("key1", "value1"); - initialData.put("key2", "value2"); - - InferredProxyContext ipc = new InferredProxyContext(initialData); - assertNotNull(ipc.getInferredProxyContext()); - assertEquals(2, ipc.getInferredProxyContext().size()); - assertEquals("value1", ipc.getInferredProxyContext().get("key1")); - assertEquals("value2", ipc.getInferredProxyContext().get("key2")); - - initialData.put("key3", "value3"); // Modify original map - assertNull(ipc.getInferredProxyContext().get("key3"), "Internal map should be a copy"); - } - - @Test - @DisplayName("putInferredProxyInfo should add to the context map") - void testPutInfo() { - InferredProxyContext ipc = new InferredProxyContext(); - ipc.putInferredProxyInfo("system", "aws-apigw"); - ipc.putInferredProxyInfo("time", "12345"); - - Map contextMap = ipc.getInferredProxyContext(); - assertEquals(2, contextMap.size()); - assertEquals("aws-apigw", contextMap.get("system")); - assertEquals("12345", contextMap.get("time")); - - ipc.putInferredProxyInfo("system", "azure-func"); // Overwrite - assertEquals("azure-func", contextMap.get("system")); - assertEquals(2, contextMap.size()); - } - - @Test - @DisplayName("removeInferredProxyInfo should remove from the context map") - void testRemoveInfo() { - Map initialData = new HashMap<>(); - initialData.put("key1", "value1"); - initialData.put("key2", "value2"); - InferredProxyContext ipc = new InferredProxyContext(initialData); - - ipc.removeInferredProxyInfo("key1"); - Map contextMap = ipc.getInferredProxyContext(); - assertEquals(1, contextMap.size()); - assertNull(contextMap.get("key1")); - assertEquals("value2", contextMap.get("key2")); - - ipc.removeInferredProxyInfo("nonexistent"); // Remove non-existent - assertEquals(1, contextMap.size()); - } - - @Test - @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") - void testStoreAndFromContext() { - InferredProxyContext ipcToStore = new InferredProxyContext(); - ipcToStore.putInferredProxyInfo("customKey", "customValue"); - - Context rootContext = Context.root(); - Context contextWithValue = ipcToStore.storeInto(rootContext); - assertNotNull(contextWithValue); - - InferredProxyContext retrievedIpc = InferredProxyContext.fromContext(contextWithValue); - assertNotNull(retrievedIpc); - assertEquals("customValue", retrievedIpc.getInferredProxyContext().get("customKey")); - - assertNull( - InferredProxyContext.fromContext(rootContext), - "Original root context should not be affected"); - - Context cleanContext = Context.root(); - assertNull( - InferredProxyContext.fromContext(cleanContext), - "fromContext on clean context should be null"); - } - - @Test - @DisplayName("getInferredProxyContext should return an unmodifiable map or a copy") - void testGetInferredProxyContextImmutability() { - InferredProxyContext ipc = new InferredProxyContext(); - ipc.putInferredProxyInfo("key1", "value1"); - - Map retrievedMap = ipc.getInferredProxyContext(); - assertNotNull(retrievedMap); - assertEquals("value1", retrievedMap.get("key1")); - - boolean threwUnsupported = false; - try { - retrievedMap.put("newKey", "newValue"); - } catch (UnsupportedOperationException e) { - threwUnsupported = true; - } - // Depending on whether InferredProxyContext.getInferredProxyContext() returns a direct - // reference or a copy, - // this assertion might change. If it returns a direct mutable reference, threwUnsupported - // will be false. - // If it returns an unmodifiable view or a copy, attempts to modify might throw or simply not - // affect the original. - // For now, we check that the original context was not changed. - assertEquals( - 1, ipc.getInferredProxyContext().size(), "Internal map size should remain unchanged"); - assertEquals( - "value1", - ipc.getInferredProxyContext().get("key1"), - "Internal map content should remain unchanged"); - // If it MUST be unmodifiable, add: assertTrue(threwUnsupported, "Retrieved map should be - // unmodifiable"); - } - - @Test - @DisplayName("Constructor with null map should create an empty context map") - void testNullMapConstructor() { - InferredProxyContext ipc = new InferredProxyContext(null); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - - @Test - @DisplayName("Constructor with empty map should create an empty context map") - void testEmptyMapConstructor() { - Map emptyMap = Collections.emptyMap(); - InferredProxyContext ipc = new InferredProxyContext(emptyMap); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - } -} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 5d8a7ee5974..91251be57f9 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -11,12 +11,8 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.context.Context; import datadog.context.propagation.Propagators; -import datadog.context.InferredProxyContext; -import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.DDTags; -import datadog.trace.api.DDTraceId; -import datadog.trace.api.TraceConfig; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; import datadog.trace.api.gateway.BlockResponseFunction; @@ -24,14 +20,13 @@ import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction; import datadog.trace.api.gateway.IGSpanInfo; +import datadog.trace.api.gateway.InferredProxySpan; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; -import datadog.trace.api.interceptor.MutableSpan; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; -import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -44,8 +39,6 @@ import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver; import java.net.InetAddress; import java.util.BitSet; -import java.util.HashMap; -import java.util.LinkedHashMap; import java.util.Locale; import java.util.Map; import java.util.concurrent.ExecutionException; @@ -60,401 +53,9 @@ public abstract class HttpServerDecorator extends ServerDecorator { - class InferredProxySpanGroup implements AgentSpan { - private final AgentSpan inferredProxySpan; - private final AgentSpan serverSpan; - - InferredProxySpanGroup(AgentSpan inferredProxySpan, AgentSpan serverSpan) { - this.inferredProxySpan = inferredProxySpan; - this.serverSpan = serverSpan; - } - - @Override - public DDTraceId getTraceId() { - return serverSpan.getTraceId(); - } - - @Override - public long getSpanId() { - return serverSpan.getSpanId(); - } - - @Override - public AgentSpan setTag(String key, boolean value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, int value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, long value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, double value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, String value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, CharSequence value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, Object value) { - return serverSpan.setTag(key, value); - } - - /** - * @param map - * @return - */ - @Override - public AgentSpan setAllTags(Map map) { - return null; - } - - @Override - public AgentSpan setTag(String key, Number value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, int value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, long value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, double value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setSpanType(CharSequence type) { - return serverSpan.setSpanType(type); - } - - @Override - public Object getTag(String key) { - return serverSpan.getTag(key); - } - - @Override - public AgentSpan setError(boolean error) { - serverSpan.setError(error); - if (inferredProxySpan != null) { - inferredProxySpan.setError(error); - } - return this; - } - - @Override - public AgentSpan setError(boolean error, byte priority) { - serverSpan.setError(error, priority); - if (inferredProxySpan != null) { - inferredProxySpan.setError(error, priority); - } - return this; - } - - @Override - public AgentSpan setMeasured(boolean measured) { - return serverSpan.setMeasured(measured); - } - - @Override - public AgentSpan setErrorMessage(String errorMessage) { - return serverSpan.setErrorMessage(errorMessage); - } - - @Override - public AgentSpan addThrowable(Throwable throwable) { - serverSpan.addThrowable(throwable); - if (inferredProxySpan != null) { - inferredProxySpan.addThrowable(throwable); - } - return this; - } - - @Override - public AgentSpan addThrowable(Throwable throwable, byte errorPriority) { - serverSpan.addThrowable(throwable, errorPriority); - if (inferredProxySpan != null) { - inferredProxySpan.addThrowable(throwable, errorPriority); - } - return this; - } - - @Override - public AgentSpan getLocalRootSpan() { - return serverSpan.getLocalRootSpan(); - } - - @Override - public boolean isSameTrace(AgentSpan otherSpan) { - return serverSpan.isSameTrace(otherSpan); - } - - @Override - public AgentSpanContext context() { - return serverSpan.context(); - } - - @Override - public String getBaggageItem(String key) { - return serverSpan.getBaggageItem(key); - } - - @Override - public AgentSpan setBaggageItem(String key, String value) { - return serverSpan.setBaggageItem(key, value); - } - - @Override - public AgentSpan setHttpStatusCode(int statusCode) { - serverSpan.setHttpStatusCode(statusCode); - if (inferredProxySpan != null) { - inferredProxySpan.setHttpStatusCode(statusCode); - } - return this; - } - - @Override - public short getHttpStatusCode() { - return serverSpan.getHttpStatusCode(); - } - - @Override - public void finish() { - serverSpan.finish(); - if (inferredProxySpan != null) { - inferredProxySpan.finish(); - } - } - - @Override - public void finish(long finishMicros) { - serverSpan.finish(finishMicros); - if (inferredProxySpan != null) { - inferredProxySpan.finish(finishMicros); - } - } - - @Override - public void finishWithDuration(long durationNanos) { - serverSpan.finishWithDuration(durationNanos); - if (inferredProxySpan != null) { - inferredProxySpan.finishWithDuration(durationNanos); - } - } - - @Override - public void beginEndToEnd() { - serverSpan.beginEndToEnd(); - } - - @Override - public void finishWithEndToEnd() { - serverSpan.finishWithEndToEnd(); - if (inferredProxySpan != null) { - inferredProxySpan.finishWithEndToEnd(); - } - } - - @Override - public boolean phasedFinish() { - final boolean ret = serverSpan.phasedFinish(); - if (inferredProxySpan != null) { - inferredProxySpan.phasedFinish(); - } - return ret; - } - - @Override - public void publish() { - serverSpan.publish(); - } - - @Override - public CharSequence getSpanName() { - return serverSpan.getSpanName(); - } - - @Override - public void setSpanName(CharSequence spanName) { - serverSpan.setSpanName(spanName); - } - - @Deprecated - @Override - public boolean hasResourceName() { - return serverSpan.hasResourceName(); - } - - @Override - public byte getResourceNamePriority() { - return serverSpan.getResourceNamePriority(); - } - - @Override - public AgentSpan setResourceName(CharSequence resourceName) { - return serverSpan.setResourceName(resourceName); - } - - @Override - public AgentSpan setResourceName(CharSequence resourceName, byte priority) { - return serverSpan.setResourceName(resourceName, priority); - } - - @Override - public RequestContext getRequestContext() { - return serverSpan.getRequestContext(); - } - - @Override - public Integer forceSamplingDecision() { - return serverSpan.forceSamplingDecision(); - } - - @Override - public AgentSpan setSamplingPriority(int newPriority, int samplingMechanism) { - return serverSpan.setSamplingPriority(newPriority, samplingMechanism); - } - - @Override - public TraceConfig traceConfig() { - return serverSpan.traceConfig(); - } - - @Override - public void addLink(AgentSpanLink link) { - serverSpan.addLink(link); - } - - @Override - public AgentSpan setMetaStruct(String field, Object value) { - return serverSpan.setMetaStruct(field, value); - } - - @Override - public boolean isOutbound() { - return serverSpan.isOutbound(); - } - - @Override - public AgentSpan asAgentSpan() { - return serverSpan.asAgentSpan(); - } - - @Override - public long getStartTime() { - return serverSpan.getStartTime(); - } - - @Override - public long getDurationNano() { - return serverSpan.getDurationNano(); - } - - @Override - public CharSequence getOperationName() { - return serverSpan.getOperationName(); - } - - @Override - public MutableSpan setOperationName(CharSequence serviceName) { - return serverSpan.setOperationName(serviceName); - } - - @Override - public String getServiceName() { - return serverSpan.getServiceName(); - } - - @Override - public MutableSpan setServiceName(String serviceName) { - return serverSpan.setServiceName(serviceName); - } - - @Override - public CharSequence getResourceName() { - return serverSpan.getResourceName(); - } - - @Override - public Integer getSamplingPriority() { - return serverSpan.getSamplingPriority(); - } - - @Deprecated - @Override - public MutableSpan setSamplingPriority(int newPriority) { - return serverSpan.setSamplingPriority(newPriority); - } - - @Override - public String getSpanType() { - return serverSpan.getSpanType(); - } - - @Override - public Map getTags() { - return serverSpan.getTags(); - } - - @Override - public boolean isError() { - return serverSpan.isError(); - } - - @Deprecated - @Override - public MutableSpan getRootSpan() { - return serverSpan.getRootSpan(); - } - - @Override - public void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba) { - serverSpan.setRequestBlockingAction(rba); - } - - @Override - public Flow.Action.RequestBlockingAction getRequestBlockingAction() { - return serverSpan.getRequestBlockingAction(); - } - } - private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class); private static final int UNSET_PORT = 0; - public static final String PROXY_SYSTEM = "x-dd-proxy"; - public static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; - public static final String PROXY_PATH = "x-dd-proxy-path"; - public static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; - public static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; - public static final String STAGE = "x-dd-proxy-stage"; - - public static final Map SUPPORTED_PROXIES; - - static { - SUPPORTED_PROXIES = new HashMap<>(); - SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); - } - public static final String DD_SPAN_ATTRIBUTE = "datadog.span"; public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch"; public static final String DD_RUM_INJECTED = "datadog.rum.injected"; @@ -547,69 +148,30 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { instrumentationNames != null && instrumentationNames.length > 0 ? instrumentationNames[0] : DEFAULT_INSTRUMENTATION_NAME; - AgentSpanContext.Extracted extracted = callIGCallbackStart(getExtractedSpanContext(context)); + AgentSpanContext extracted = getExtractedSpanContext(context); + // Call IG callbacks + extracted = callIGCallbackStart(extracted); + // Create gateway inferred span if needed + extracted = startInferredProxySpan(context, extracted); AgentSpan span = tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true); + // Apply RequestBlockingAction if any Flow flow = callIGCallbackRequestHeaders(span, carrier); if (flow.getAction() instanceof RequestBlockingAction) { span.setRequestBlockingAction((RequestBlockingAction) flow.getAction()); } - AgentPropagation.ContextVisitor getter = getter(); - if (null != carrier && null != getter) { - tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); - } + // DSM Checkpoint + tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); return context.with(span); } - private AgentSpan startSpanWithInferredProxy( - String instrumentationName, - datadog.context.Context fullContextForInferredProxy, - AgentSpanContext.Extracted standardExtractedContext) { - - InferredProxyContext inferredProxy = - InferredProxyContext.fromContext(fullContextForInferredProxy); - - if (inferredProxy == null) { - return null; - } - - Map headers = inferredProxy.getInferredProxyContext(); - - // Check if timestamp and proxy system are present - String startTimeStr = headers.get(PROXY_START_TIME_MS); - String proxySystem = headers.get(PROXY_SYSTEM); - - if (startTimeStr == null - || proxySystem == null - || !SUPPORTED_PROXIES.containsKey(proxySystem)) { - return null; - } - - long startTime; - try { - startTime = Long.parseLong(startTimeStr) * 1000; // Convert to microseconds - } catch (NumberFormatException e) { - return null; // Invalid timestamp + protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) { + InferredProxySpan span; + if (!Config.get().isInferredProxyPropagationEnabled() + || (span = InferredProxySpan.fromContext(context)) == null) { + return extracted; } - - AgentSpan apiGtwSpan = - tracer() - .startSpan( - "inferred_proxy", - SUPPORTED_PROXIES.get(proxySystem), - callIGCallbackStart(standardExtractedContext), - startTime); - - apiGtwSpan.setTag(Tags.COMPONENT, proxySystem); - apiGtwSpan.setTag( - DDTags.RESOURCE_NAME, headers.get(PROXY_HTTP_METHOD) + " " + headers.get(PROXY_PATH)); - apiGtwSpan.setTag(DDTags.SERVICE_NAME, headers.get(PROXY_DOMAIN_NAME)); - apiGtwSpan.setTag(DDTags.SPAN_TYPE, "web"); - apiGtwSpan.setTag(Tags.HTTP_METHOD, headers.get(PROXY_HTTP_METHOD)); - apiGtwSpan.setTag(Tags.HTTP_URL, headers.get(PROXY_DOMAIN_NAME) + headers.get(PROXY_PATH)); - apiGtwSpan.setTag("stage", headers.get(STAGE)); - apiGtwSpan.setTag("_dd.inferred_span", 1); - return apiGtwSpan; + return span.start(extracted); } public AgentSpan onRequest( @@ -832,8 +394,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) { return span; } - private AgentSpanContext.Extracted callIGCallbackStart( - @Nullable final AgentSpanContext.Extracted extracted) { + private AgentSpanContext callIGCallbackStart(@Nullable final AgentSpanContext extracted) { AgentTracer.TracerAPI tracer = tracer(); Supplier> startedCbAppSec = tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted()); @@ -969,10 +530,20 @@ private Flow callIGCallbackURI( @Override public AgentSpan beforeFinish(AgentSpan span) { + // TODO Migrate beforeFinish to Context API onRequestEndForInstrumentationGateway(span); + // Close Serverless Gateway Inferred Span if any + // finishInferredProxySpan(context); return super.beforeFinish(span); } + protected void finishInferredProxySpan(Context context) { + InferredProxySpan span; + if ((span = InferredProxySpan.fromContext(context)) != null) { + span.finish(); + } + } + private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) { if (span.getLocalRootSpan() != span) { return; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 56aeedc6e71..71cf4fe76c4 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -23,7 +23,6 @@ import datadog.communication.ddagent.SharedCommunicationObjects; import datadog.communication.monitor.Monitoring; import datadog.communication.monitor.Recording; -import datadog.context.propagation.InferredProxyPropagator; import datadog.context.propagation.Propagators; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; @@ -93,6 +92,7 @@ import datadog.trace.core.monitor.TracerHealthMetrics; import datadog.trace.core.propagation.ExtractedContext; import datadog.trace.core.propagation.HttpCodec; +import datadog.trace.core.propagation.InferredProxyPropagator; import datadog.trace.core.propagation.PropagationTags; import datadog.trace.core.propagation.TracingPropagator; import datadog.trace.core.propagation.XRayPropagator; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java new file mode 100644 index 00000000000..da51bf4643b --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java @@ -0,0 +1,60 @@ +package datadog.trace.core.propagation; + +import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; + +import datadog.context.Context; +import datadog.context.propagation.CarrierSetter; +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.Propagator; +import datadog.trace.api.gateway.InferredProxySpan; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import javax.annotation.ParametersAreNonnullByDefault; + +/** Inferred proxy propagator. Only extract, not meant for injection. */ +@ParametersAreNonnullByDefault +public class InferredProxyPropagator implements Propagator { + private static final String INFERRED_PROXY_KEY_PREFIX = "x-dd-proxy"; + + @Override + public void inject(Context context, C carrier, CarrierSetter setter) {} + + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + if (context == null || carrier == null || visitor == null) { + return context; + } + InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); + visitor.forEachKeyValue(carrier, extractor); + InferredProxySpan inferredProxySpan = extractor.inferredProxySpan(); + if (inferredProxySpan != null) { + context = context.with(inferredProxySpan); + } + return context; + } + + /** Extract inferred proxy related headers into a map. */ + private static class InferredProxyContextExtractor implements BiConsumer { + private Map values; + + @Override + public void accept(String key, String value) { + if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY_PREFIX)) { + return; + } + if (values == null) { + this.values = new HashMap<>(); + } + this.values.put(key, value); + } + + public InferredProxySpan inferredProxySpan() { + if (this.values == null) { + return null; + } + InferredProxySpan inferredProxySpan = fromHeaders(this.values); + return inferredProxySpan.isValid() ? inferredProxySpan : null; + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java new file mode 100644 index 00000000000..5830a160cc7 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java @@ -0,0 +1,96 @@ +package datadog.trace.core.propagation; + +import static datadog.context.Context.root; +import static datadog.trace.api.gateway.InferredProxySpan.fromContext; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.params.provider.Arguments.of; + +import datadog.context.Context; +import datadog.context.propagation.CarrierVisitor; +import datadog.trace.api.gateway.InferredProxySpan; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.stream.Stream; +import javax.annotation.ParametersAreNonnullByDefault; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +@DisplayName("InferredProxyPropagator Tests") +class InferredProxyPropagatorTests { + private static final String PROXY_SYSTEM_KEY = "x-dd-proxy"; + private static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; + private static final String PROXY_PATH_KEY = "x-dd-proxy-path"; + private static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; + private static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; + private static final MapVisitor MAP_VISITOR = new MapVisitor(); + + private InferredProxyPropagator propagator; + + @BeforeEach + void setUp() { + this.propagator = new InferredProxyPropagator(); + } + + @Test + @DisplayName("Should extract InferredProxySpan when valid headers are present") + void testSuccessfulExtraction() { + Map headers = new HashMap<>(); + headers.put(PROXY_SYSTEM_KEY, "aws-apigateway"); + headers.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + headers.put(PROXY_PATH_KEY, "/foo"); + headers.put(PROXY_HTTP_METHOD_KEY, "GET"); + headers.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); + + Context context = this.propagator.extract(root(), headers, MAP_VISITOR); + InferredProxySpan inferredProxySpan = fromContext(context); + assertNotNull(inferredProxySpan); + assertTrue(inferredProxySpan.isValid()); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("invalidOrMissingHeadersProviderForPropagator") + @DisplayName("Should not create InferredProxySpan if some critical headers are missing") + void testExtractionWithMissingCriticalHeaders(String description, Map headers) { + Context rootContext = root(); + Context extractedOuterContext = this.propagator.extract(rootContext, headers, MAP_VISITOR); + InferredProxySpan inferredProxySpan = fromContext(extractedOuterContext); + assertNull(inferredProxySpan, "Invalid inferred proxy span should not be extracted"); + } + + static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + missingSystem.put(PROXY_PATH_KEY, "/foo"); + + Map emptyValue = new HashMap<>(); + emptyValue.put(PROXY_SYSTEM_KEY, ""); + + Map nullValue = new HashMap<>(); + nullValue.put(PROXY_SYSTEM_KEY, null); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); + missingTime.put(PROXY_PATH_KEY, "/foo"); + + return Stream.of( + of("PROXY_SYSTEM_KEY missing", missingSystem), + of("PROXY_SYSTEM_KEY empty", emptyValue), + of("PROXY_SYSTEM_KEY null", nullValue), + of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); + } + + @ParametersAreNonnullByDefault + private static class MapVisitor implements CarrierVisitor> { + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + carrier.forEach(visitor); + } + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index a2e58598ea5..de78bd9dbd8 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -598,6 +598,7 @@ import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_ERROR_STATUSES; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_PATH_RESOURCE_NAME_MAPPING; +import static datadog.trace.api.config.TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_KEEP_LATENCY_THRESHOLD_MS; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_FLUSH_INTERVAL; @@ -5409,6 +5410,8 @@ public String toString() { + tracePropagationBehaviorExtract + ", tracePropagationExtractFirst=" + tracePropagationExtractFirst + + ", traceInferredProxyEnabled=" + + traceInferredProxyEnabled + ", clockSyncPeriod=" + clockSyncPeriod + ", jmxFetchEnabled=" diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java new file mode 100644 index 00000000000..a8256c1b38a --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java @@ -0,0 +1,108 @@ +package datadog.trace.api.gateway; + +import static datadog.context.ContextKey.named; +import static datadog.trace.api.DDTags.RESOURCE_NAME; +import static datadog.trace.api.DDTags.SERVICE_NAME; +import static datadog.trace.api.DDTags.SPAN_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_METHOD; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ImplicitContextKeyed; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class InferredProxySpan implements ImplicitContextKeyed { + private static final ContextKey CONTEXT_KEY = named("inferred-proxy-key"); + static final String PROXY_SYSTEM = "x-dd-proxy"; + static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; + static final String PROXY_PATH = "x-dd-proxy-path"; + static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; + static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; + static final String STAGE = "x-dd-proxy-stage"; + static final Map SUPPORTED_PROXIES; + static final String INSTRUMENTATION_NAME = "inferred_proxy"; + + static { + SUPPORTED_PROXIES = new HashMap<>(); + SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); + } + + private final Map headers; + private AgentSpan span; + + public static InferredProxySpan fromHeaders(Map values) { + return new InferredProxySpan(values); + } + + public static InferredProxySpan fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + private InferredProxySpan(Map headers) { + this.headers = headers == null ? Collections.emptyMap() : headers; + } + + public boolean isValid() { + String startTimeStr = header(PROXY_START_TIME_MS); + String proxySystem = header(PROXY_SYSTEM); + return startTimeStr != null + && proxySystem != null + && SUPPORTED_PROXIES.containsKey(proxySystem); + } + + public AgentSpanContext start(AgentSpanContext extracted) { + if (this.span != null || !isValid()) { + return extracted; + } + + long startTime; + try { + startTime = Long.parseLong(header(PROXY_START_TIME_MS)) * 1000; // Convert to microseconds + } catch (NumberFormatException e) { + return extracted; // Invalid timestamp + } + + String proxySystem = header(PROXY_SYSTEM); + String proxy = SUPPORTED_PROXIES.get(proxySystem); + AgentSpan span = AgentTracer.get().startSpan(INSTRUMENTATION_NAME, proxy, extracted, startTime); + + span.setTag(COMPONENT, proxySystem); + span.setTag(RESOURCE_NAME, header(PROXY_HTTP_METHOD) + " " + header(PROXY_PATH)); + span.setTag(SERVICE_NAME, header(PROXY_DOMAIN_NAME)); + span.setTag(SPAN_TYPE, "web"); + span.setTag(HTTP_METHOD, header(PROXY_HTTP_METHOD)); + span.setTag(HTTP_URL, header(PROXY_DOMAIN_NAME) + header(PROXY_PATH)); + span.setTag("stage", header(STAGE)); + span.setTag("_dd.inferred_span", 1); + + // Free collected headers + this.headers.clear(); + // Store inferred span + this.span = span; + // Return inferred span as new parent context + return this.span.context(); + } + + private String header(String name) { + return this.headers.get(name); + } + + public void finish() { + if (this.span != null) { + this.span.finish(); + this.span = null; + } + } + + @Override + public Context storeInto(Context context) { + return context.with(CONTEXT_KEY, this); + } +} diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java new file mode 100644 index 00000000000..dc61d4a23aa --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java @@ -0,0 +1,95 @@ +package datadog.trace.api.gateway; + +import static datadog.context.Context.root; +import static datadog.trace.api.gateway.InferredProxySpan.PROXY_START_TIME_MS; +import static datadog.trace.api.gateway.InferredProxySpan.PROXY_SYSTEM; +import static datadog.trace.api.gateway.InferredProxySpan.fromContext; +import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; +import static java.util.Collections.emptyMap; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.params.provider.Arguments.of; + +import datadog.context.Context; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Stream; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +@DisplayName("InferredProxyContext Tests") +class InferredProxySpanTests { + @Test + @DisplayName("Valid headers should make valid span") + void testMapConstructor() { + Map headers = new HashMap<>(); + headers.put(PROXY_START_TIME_MS, "12345"); + headers.put(PROXY_SYSTEM, "aws-apigateway"); + + InferredProxySpan inferredProxySpan = InferredProxySpan.fromHeaders(headers); + assertTrue(inferredProxySpan.isValid()); + assertNotNull( + inferredProxySpan.start(null), "inferred proxy span start and return new parent context"); + assertNull(inferredProxySpan.start(null), "inferred proxy span not should start twice"); + inferredProxySpan.finish(); + } + + @ParameterizedTest(name = "{0}") + @DisplayName("Invalid headers should make invalid span") + @MethodSource("invalidHeaders") + void testInvalidHeaders(String useCase, Map headers) { + InferredProxySpan inferredProxySpan = fromHeaders(headers); + assertFalse(inferredProxySpan.isValid(), useCase + " should not be valid"); + assertNull(inferredProxySpan.start(null), "Invalid inferred proxy span should not start"); + } + + static Stream invalidHeaders() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_START_TIME_MS, "12345"); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM, "aws-apigateway"); + Map invalidSystem = new HashMap<>(); + invalidSystem.put(PROXY_START_TIME_MS, "12345"); + invalidSystem.put(PROXY_SYSTEM, "invalidSystem"); + + return Stream.of( + of("Missing system headers", missingSystem), + of("Missing start time headers", missingTime), + of("Invalid system headers", invalidSystem)); + } + + @Test + @DisplayName("Constructor with null should not crash") + void testNullMapConstructor() { + InferredProxySpan inferredProxySpan = fromHeaders(null); + assertNotNull(inferredProxySpan); + assertFalse(inferredProxySpan.isValid()); + } + + @Test + @DisplayName("Constructor with empty map should be invalid") + void testEmptyMapConstructor() { + InferredProxySpan inferredProxySpan = fromHeaders(emptyMap()); + assertNotNull(inferredProxySpan); + assertFalse(inferredProxySpan.isValid()); + } + + @Test + @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") + void testStoreAndFromContext() { + InferredProxySpan inferredProxySpan = fromHeaders(null); + Context context = inferredProxySpan.storeInto(root()); + assertNotNull(context); + + InferredProxySpan retrieved = fromContext(context); + assertNotNull(retrieved); + + assertNull(fromContext(root()), "fromContext on empty context should be null"); + } +}