From 57cc8f68168dbad34b4d2965484229ee22473c75 Mon Sep 17 00:00:00 2001 From: Brice Dutheil Date: Mon, 29 Jun 2026 15:44:02 +0200 Subject: [PATCH 1/2] test: Coverage tests for proxy span metadata, ARN computation, SCV dep handling --- .../api/gateway/InferredProxySpanTests.java | 93 +++++++++++++++++++ .../gateway/InstrumentationGatewayTest.java | 44 +++++++++ ...ScaReachabilityDependencyRegistryTest.java | 92 ++++++++++++++++++ 3 files changed, 229 insertions(+) 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 index c9adcb83551..93d2e2bc92a 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java @@ -5,21 +5,27 @@ 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 datadog.trace.bootstrap.instrumentation.api.ErrorPriorities.HTTP_SERVER_DECORATOR; import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_REQUEST_HEADERS_X_DATADOG_ENDPOINT_SCAN; import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_REQUEST_HEADERS_X_DATADOG_SECURITY_TEST; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_USER_AGENT; import static java.util.Collections.emptyMap; +import static org.junit.jupiter.api.Assertions.assertEquals; 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 static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.lang.reflect.Field; +import java.lang.reflect.Method; import java.util.HashMap; import java.util.Map; import java.util.stream.Stream; @@ -576,6 +582,43 @@ void testFinishWithAllHeaders() { inferredProxySpan.finish(); } + @Test + @DisplayName("optional API Gateway metadata should exercise ARN tag computation") + void testStartWithOptionalApiGatewayMetadata() { + Map headers = new HashMap<>(); + headers.put(PROXY_START_TIME_MS, "12345"); + headers.put(PROXY_SYSTEM, "aws-apigateway"); + headers.put(InferredProxySpan.PROXY_HTTP_METHOD, "GET"); + headers.put(InferredProxySpan.PROXY_PATH, "/api/users/123"); + headers.put(InferredProxySpan.PROXY_API_ID, "api-id"); + headers.put(InferredProxySpan.PROXY_REGION, "us-east-1"); + headers.put(InferredProxySpan.PROXY_ACCOUNT_ID, "123456789012"); + + InferredProxySpan inferredProxySpan = fromHeaders(headers); + assertNotNull(inferredProxySpan.start(null)); + + inferredProxySpan.finish(); + } + + @Test + @DisplayName("computeArn should support known proxy systems and reject unknown input") + void testComputeArn() throws Exception { + InferredProxySpan inferredProxySpan = fromHeaders(null); + Method computeArn = + InferredProxySpan.class.getDeclaredMethod( + "computeArn", String.class, String.class, String.class); + computeArn.setAccessible(true); + + assertEquals( + "arn:aws:apigateway:us-east-1::/restapis/api-id", + computeArn.invoke(inferredProxySpan, "aws-apigateway", "us-east-1", "api-id")); + assertEquals( + "arn:aws:apigateway:us-east-1::/apis/api-id", + computeArn.invoke(inferredProxySpan, "aws-httpapi", "us-east-1", "api-id")); + assertNull(computeArn.invoke(inferredProxySpan, "unknown", "us-east-1", "api-id")); + assertNull(computeArn.invoke(inferredProxySpan, "aws-apigateway", null, "api-id")); + } + @Test @DisplayName("Multiple InferredProxySpan instances should finish independently") void testMultipleProxySpansFinishIndependently() { @@ -632,4 +675,54 @@ void testFinishForwardsSecurityTestingHeaders() throws Exception { verify(mockInferredSpan).setTag(HTTP_REQUEST_HEADERS_X_DATADOG_ENDPOINT_SCAN, "scan-uuid"); verify(mockInferredSpan).setTag(HTTP_REQUEST_HEADERS_X_DATADOG_SECURITY_TEST, "test-uuid"); } + + @Test + @DisplayName("finish should phase child calls and publish after service-entry completion") + void testFinishPhasesChildCallAndPublishesOnServiceEntry() throws Exception { + InferredProxySpan inferredProxySpan = fromHeaders(validHeaders()); + + AgentSpan inferredSpan = mock(AgentSpan.class); + Field spanField = InferredProxySpan.class.getDeclaredField("span"); + spanField.setAccessible(true); + spanField.set(inferredProxySpan, inferredSpan); + + AgentSpan childSpan = mock(AgentSpan.class); + when(childSpan.getTag("_dd.appsec.enabled")).thenReturn(Boolean.TRUE); + when(childSpan.getTag("_dd.appsec.json")).thenReturn("{\"triggers\":[]}"); + + AgentSpan serviceEntrySpan = mock(AgentSpan.class); + when(serviceEntrySpan.getHttpStatusCode()).thenReturn((short) 503); + when(serviceEntrySpan.getTag(HTTP_USER_AGENT)).thenReturn("curl/8.0"); + when(serviceEntrySpan.getTag(HTTP_REQUEST_HEADERS_X_DATADOG_ENDPOINT_SCAN)) + .thenReturn("scan-uuid"); + when(serviceEntrySpan.getTag(HTTP_REQUEST_HEADERS_X_DATADOG_SECURITY_TEST)) + .thenReturn("test-uuid"); + + inferredProxySpan.registerServiceEntrySpan(serviceEntrySpan); + + inferredProxySpan.finish(childSpan); + inferredProxySpan.finish(childSpan); + verify(inferredSpan).setMetric("_dd.appsec.enabled", 1); + verify(inferredSpan).setTag("_dd.appsec.json", "{\"triggers\":[]}"); + verify(inferredSpan, times(1)).phasedFinish(); + verify(inferredSpan, never()).finish(); + verify(inferredSpan, never()).publish(); + + inferredProxySpan.finish(serviceEntrySpan); + + verify(inferredSpan).setHttpStatusCode(503); + verify(inferredSpan).setError(true, HTTP_SERVER_DECORATOR); + verify(inferredSpan).setTag(HTTP_USER_AGENT, "curl/8.0"); + verify(inferredSpan).setTag(HTTP_REQUEST_HEADERS_X_DATADOG_ENDPOINT_SCAN, "scan-uuid"); + verify(inferredSpan).setTag(HTTP_REQUEST_HEADERS_X_DATADOG_SECURITY_TEST, "test-uuid"); + verify(inferredSpan).publish(); + verify(inferredSpan, never()).finish(); + } + + private static Map validHeaders() { + Map headers = new HashMap<>(); + headers.put(PROXY_START_TIME_MS, "12345"); + headers.put(PROXY_SYSTEM, "aws-apigateway"); + return headers; + } } diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java index 4767a0051b9..145b3b117f6 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java @@ -17,6 +17,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import java.util.Collections; +import java.util.Map; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; @@ -169,6 +170,26 @@ public void testRequestBlockingAction() { assertEquals("https://www.google.com/", rba.getExtraHeaders().get("Location")); } + @Test + public void blockResponseFunctionDefaultMethodDelegatesRequestBlockingAction() { + TraceSegment segment = TraceSegment.NoOp.INSTANCE; + Flow.Action.RequestBlockingAction action = + new Flow.Action.RequestBlockingAction( + 451, + BlockingContentType.JSON, + Collections.singletonMap("x-blocked", "true"), + "security-response-id"); + + CapturingBlockResponseFunction blockResponseFunction = new CapturingBlockResponseFunction(); + + assertTrue(blockResponseFunction.tryCommitBlockingResponse(segment, action)); + assertSame(segment, blockResponseFunction.segment); + assertEquals(451, blockResponseFunction.statusCode); + assertEquals(BlockingContentType.JSON, blockResponseFunction.templateType); + assertEquals("true", blockResponseFunction.extraHeaders.get("x-blocked")); + assertEquals("security-response-id", blockResponseFunction.securityResponseId); + } + @Test public void testNormalCalls() { // check that we pass through normal calls @@ -564,6 +585,29 @@ public Flow apply(RequestContext requestContext, T t, T t2) { } } + private static final class CapturingBlockResponseFunction implements BlockResponseFunction { + private TraceSegment segment; + private int statusCode; + private BlockingContentType templateType; + private Map extraHeaders; + private String securityResponseId; + + @Override + public boolean tryCommitBlockingResponse( + TraceSegment segment, + int statusCode, + BlockingContentType templateType, + Map extraHeaders, + String securityResponseId) { + this.segment = segment; + this.statusCode = statusCode; + this.templateType = templateType; + this.extraHeaders = extraHeaders; + this.securityResponseId = securityResponseId; + return true; + } + } + private static class Throwback implements Supplier>, BiConsumer, diff --git a/internal-api/src/test/java/datadog/trace/api/telemetry/ScaReachabilityDependencyRegistryTest.java b/internal-api/src/test/java/datadog/trace/api/telemetry/ScaReachabilityDependencyRegistryTest.java index ddf61e0ad38..3d022cc9144 100644 --- a/internal-api/src/test/java/datadog/trace/api/telemetry/ScaReachabilityDependencyRegistryTest.java +++ b/internal-api/src/test/java/datadog/trace/api/telemetry/ScaReachabilityDependencyRegistryTest.java @@ -115,6 +115,54 @@ void registerCve_addsEntryAndMarksPending() { assertNull(dep.cves.get(0).hit, "class-load registration has no callsite yet"); } + @Test + void recordHit_snapshotContainsFullHitMetadata() { + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "com.example:lib", "2.0.0", "GHSA-0001", "com.myapp.Ctrl", "handle", 42); + + List snapshots = + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies(); + + assertEquals(1, snapshots.size()); + ScaReachabilityDependencyRegistry.DependencySnapshot dep = snapshots.get(0); + assertEquals("com.example:lib", dep.artifact); + assertEquals("2.0.0", dep.version); + assertEquals(1, dep.cves.size()); + + ScaReachabilityDependencyRegistry.CveSnapshot cve = dep.cves.get(0); + assertEquals("GHSA-0001", cve.vulnId); + assertNotNull(cve.hit); + assertEquals("GHSA-0001", cve.hit.vulnId()); + assertEquals("com.example:lib", cve.hit.artifact()); + assertEquals("2.0.0", cve.hit.version()); + assertEquals("com.myapp.Ctrl", cve.hit.className()); + assertEquals("handle", cve.hit.symbolName()); + assertEquals(42, cve.hit.line()); + } + + @Test + void peekSnapshot_returnsCurrentStateWithoutClearingPendingFlag() { + assertNull(ScaReachabilityDependencyRegistry.INSTANCE.peekSnapshot("missing", "1.0.0")); + + ScaReachabilityDependencyRegistry.INSTANCE.registerCve("com.example:lib", "2.0.0", "GHSA-0001"); + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "com.example:lib", "2.0.0", "GHSA-0001", "com.myapp.Ctrl", "handle", 42); + + ScaReachabilityDependencyRegistry.DependencySnapshot peeked = + ScaReachabilityDependencyRegistry.INSTANCE.peekSnapshot("com.example:lib", "2.0.0"); + + assertNotNull(peeked); + assertEquals("com.example:lib", peeked.artifact); + assertEquals("2.0.0", peeked.version); + assertEquals(1, peeked.cves.size()); + assertEquals("GHSA-0001", peeked.cves.get(0).vulnId); + assertNotNull(peeked.cves.get(0).hit); + + List snapshots = + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies(); + assertEquals(1, snapshots.size(), "peek must not clear pending state"); + } + @Test void drainPendingDependencies_secondDrainEmpty_untilNewHit() { ScaReachabilityDependencyRegistry.INSTANCE.registerCve("com.example:lib", "2.0.0", "GHSA-0001"); @@ -138,6 +186,26 @@ void drainPendingDependencies_secondDrainEmpty_untilNewHit() { assertNotNull(third.get(0).cves.get(0).hit, "hit callsite must be recorded"); } + @Test + void recordHit_firstHitWinsAndDuplicateDoesNotMarkPendingAgain() { + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "com.example:lib", "2.0.0", "GHSA-0001", "com.myapp.First", "first", 1); + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies(); + + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "com.example:lib", "2.0.0", "GHSA-0001", "com.myapp.Second", "second", 2); + + assertTrue( + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies().isEmpty(), + "duplicate hit for the same CVE must not mark the dependency pending"); + + ScaReachabilityDependencyRegistry.DependencySnapshot snapshot = + ScaReachabilityDependencyRegistry.INSTANCE.peekSnapshot("com.example:lib", "2.0.0"); + assertNotNull(snapshot); + assertEquals("com.myapp.First", snapshot.cves.get(0).hit.className()); + assertEquals("first", snapshot.cves.get(0).hit.symbolName()); + } + @Test void registerCve_atCap_newKeysRejected() { int cap = Config.get().getAppSecScaMaxTrackedDependencies(); @@ -177,6 +245,30 @@ void registerCve_atCap_existingKeyStillUpdated() { assertEquals(2, snapshots.get(0).cves.size(), "both CVEs must be present"); } + @Test + void recordHit_atCap_newKeysRejectedButExistingKeyStillUpdated() { + int cap = Config.get().getAppSecScaMaxTrackedDependencies(); + + for (int i = 0; i < cap; i++) { + ScaReachabilityDependencyRegistry.INSTANCE.registerCve("art" + i, "1.0", "GHSA-" + i); + } + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies(); + + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "art-over-cap", "1.0", "GHSA-over", "com.myapp.Ctrl", "handle", 42); + assertTrue( + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies().isEmpty(), + "over-cap hit for a new dependency must be rejected"); + + ScaReachabilityDependencyRegistry.INSTANCE.recordHit( + "art0", "1.0", "GHSA-0", "com.myapp.Ctrl", "handle", 42); + List snapshots = + ScaReachabilityDependencyRegistry.INSTANCE.drainPendingDependencies(); + assertEquals(1, snapshots.size(), "existing dependency can still be updated at cap"); + assertEquals("art0", snapshots.get(0).artifact); + assertNotNull(snapshots.get(0).cves.get(0).hit); + } + @Test void resetForTesting_clearsPeriodicWorkCallback() { ScaReachabilityDependencyRegistry.INSTANCE.setPeriodicWorkCallback(() -> {}); From ca5c9f9bf3a6a13577a05bac27ef233d8d28071f Mon Sep 17 00:00:00 2001 From: Brice Dutheil Date: Tue, 30 Jun 2026 11:09:48 +0200 Subject: [PATCH 2/2] test: address sca review --- .../trace/api/gateway/InferredProxySpan.java | 6 ++++-- .../api/gateway/InferredProxySpanTests.java | 21 +++++++------------ 2 files changed, 12 insertions(+), 15 deletions(-) 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 index aef0fb81a70..5a893e7d760 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java @@ -18,6 +18,7 @@ import datadog.context.ContextKey; import datadog.context.ImplicitContextKeyed; import datadog.trace.api.Config; +import datadog.trace.api.internal.VisibleForTesting; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; @@ -49,7 +50,7 @@ public class InferredProxySpan implements ImplicitContextKeyed { } private final Map headers; - private AgentSpan span; + @VisibleForTesting AgentSpan span; // Service-entry span registered at startSpan() time; used to guard against premature finishing // by child spans (e.g., Spring MVC handler spans) before the response status is known. private AgentSpan registeredServiceEntrySpan; @@ -178,7 +179,8 @@ private String header(String name) { * arn:aws:apigateway:{region}::/restapis/{api-id} Format for v2 HTTP: * arn:aws:apigateway:{region}::/apis/{api-id} */ - private String computeArn(String proxySystem, String region, String apiId) { + @VisibleForTesting + String computeArn(String proxySystem, String region, String apiId) { if (proxySystem == null || region == null || apiId == null) { return null; } 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 index 93d2e2bc92a..8463136c63b 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java @@ -25,7 +25,6 @@ import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.lang.reflect.Field; -import java.lang.reflect.Method; import java.util.HashMap; import java.util.Map; import java.util.stream.Stream; @@ -602,21 +601,18 @@ void testStartWithOptionalApiGatewayMetadata() { @Test @DisplayName("computeArn should support known proxy systems and reject unknown input") - void testComputeArn() throws Exception { + void testComputeArn() { InferredProxySpan inferredProxySpan = fromHeaders(null); - Method computeArn = - InferredProxySpan.class.getDeclaredMethod( - "computeArn", String.class, String.class, String.class); - computeArn.setAccessible(true); assertEquals( "arn:aws:apigateway:us-east-1::/restapis/api-id", - computeArn.invoke(inferredProxySpan, "aws-apigateway", "us-east-1", "api-id")); + inferredProxySpan.computeArn("aws-apigateway", "us-east-1", "api-id")); assertEquals( "arn:aws:apigateway:us-east-1::/apis/api-id", - computeArn.invoke(inferredProxySpan, "aws-httpapi", "us-east-1", "api-id")); - assertNull(computeArn.invoke(inferredProxySpan, "unknown", "us-east-1", "api-id")); - assertNull(computeArn.invoke(inferredProxySpan, "aws-apigateway", null, "api-id")); + inferredProxySpan.computeArn("aws-httpapi", "us-east-1", "api-id")); + assertNull(inferredProxySpan.computeArn("unknown", "us-east-1", "api-id")); + assertNull(inferredProxySpan.computeArn("aws-apigateway", null, "api-id")); + assertNull(inferredProxySpan.computeArn("aws-apigateway", "us-east-1", null)); } @Test @@ -659,6 +655,7 @@ void testFinishForwardsSecurityTestingHeaders() throws Exception { // Replace the real (noop) inferred span with a mock we can verify against. Drive through // the public finish() API so the test stays valid if the internal copy-helper is renamed. AgentSpan mockInferredSpan = mock(AgentSpan.class); + // Keep this reflected field name in sync with InferredProxySpan.span. Field spanField = InferredProxySpan.class.getDeclaredField("span"); spanField.setAccessible(true); spanField.set(inferredProxySpan, mockInferredSpan); @@ -682,9 +679,7 @@ void testFinishPhasesChildCallAndPublishesOnServiceEntry() throws Exception { InferredProxySpan inferredProxySpan = fromHeaders(validHeaders()); AgentSpan inferredSpan = mock(AgentSpan.class); - Field spanField = InferredProxySpan.class.getDeclaredField("span"); - spanField.setAccessible(true); - spanField.set(inferredProxySpan, inferredSpan); + inferredProxySpan.span = inferredSpan; AgentSpan childSpan = mock(AgentSpan.class); when(childSpan.getTag("_dd.appsec.enabled")).thenReturn(Boolean.TRUE);