Enable and fix test coverage for SCA types in internal API

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<String, String> 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;
     }

internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java

@@ -5,15 +5,20 @@
 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;
 
@@ -576,6 +581,40 @@ void testFinishWithAllHeaders() {
     inferredProxySpan.finish();
   }
 
+  @Test
+  @DisplayName("optional API Gateway metadata should exercise ARN tag computation")
+  void testStartWithOptionalApiGatewayMetadata() {
+    Map<String, String> 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() {
+    InferredProxySpan inferredProxySpan = fromHeaders(null);
+
+    assertEquals(
+        "arn:aws:apigateway:us-east-1::/restapis/api-id",
+        inferredProxySpan.computeArn("aws-apigateway", "us-east-1", "api-id"));
+    assertEquals(
+        "arn:aws:apigateway:us-east-1::/apis/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
   @DisplayName("Multiple InferredProxySpan instances should finish independently")
   void testMultipleProxySpansFinishIndependently() {
@@ -616,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);
@@ -632,4 +672,52 @@ 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);
+    inferredProxySpan.span = 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<String, String> validHeaders() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    return headers;
+  }
 }

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<Void> 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<String, String> extraHeaders;
+    private String securityResponseId;
+
+    @Override
+    public boolean tryCommitBlockingResponse(
+        TraceSegment segment,
+        int statusCode,
+        BlockingContentType templateType,
+        Map<String, String> extraHeaders,
+        String securityResponseId) {
+      this.segment = segment;
+      this.statusCode = statusCode;
+      this.templateType = templateType;
+      this.extraHeaders = extraHeaders;
+      this.securityResponseId = securityResponseId;
+      return true;
+    }
+  }
+
   private static class Throwback<D, T>
       implements Supplier<Flow<D>>,
           BiConsumer<RequestContext, T>,

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<ScaReachabilityDependencyRegistry.DependencySnapshot> 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<ScaReachabilityDependencyRegistry.DependencySnapshot> 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<ScaReachabilityDependencyRegistry.DependencySnapshot> 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(() -> {});