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/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..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 @@ -20,6 +20,7 @@ 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.naming.SpanNaming; @@ -147,20 +148,32 @@ 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); } + protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) { + InferredProxySpan span; + if (!Config.get().isInferredProxyPropagationEnabled() + || (span = InferredProxySpan.fromContext(context)) == null) { + return extracted; + } + return span.start(extracted); + } + public AgentSpan onRequest( final AgentSpan span, final CONNECTION connection, @@ -381,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()); @@ -518,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-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..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 @@ -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; @@ -91,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; @@ -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/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 aae53ec7e07..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; @@ -827,6 +828,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 +1732,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 +3122,10 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public boolean isInferredProxyPropagationEnabled() { + return traceInferredProxyEnabled; + } + public boolean isBaggageExtract() { return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); } @@ -5402,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/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); 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"); + } +}