Fix API Security sampling in standalone mode

dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java

@@ -4,6 +4,8 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.api.time.TimeSource;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.Deque;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
@@ -73,8 +75,15 @@ public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) {
       return false;
     }
     if (counter.tryAcquire()) {
-      log.debug("API security sampling is required for this request (presampled)");
       ctx.setKeepOpenForApiSecurityPostProcessing(true);
+      if (!Config.get().isApmTracingEnabled()) {
+        boolean sampled = updateApiAccessIfExpired(hash);
+        if (sampled) {
+          logSamplingDecision("preSampleRequest", hash);
+        }
+        return sampled;
+      }
+      logSamplingDecision("preSampleRequest", hash);
       return true;
     }
     return false;
@@ -91,7 +100,11 @@ public boolean sampleRequest(AppSecRequestContext ctx) {
       // This should never happen, it should have been short-circuited before.
       return false;
     }
-    return updateApiAccessIfExpired(hash);
+    boolean sampled = Config.get().isApmTracingEnabled() ? updateApiAccessIfExpired(hash) : true;
+    if (sampled) {
+      logSamplingDecision("sampleRequest", hash);
+    }
+    return sampled;
   }
 
   @Override
@@ -168,4 +181,23 @@ private long computeApiHash(final String route, final String method, final int s
     result = 31 * result + statusCode;
     return result;
   }
+
+  private void logSamplingDecision(final String method, final long hash) {
+    if (!log.isDebugEnabled()) {
+      return;
+    }
+    AgentSpan activeSpan = AgentTracer.get().activeSpan();
+
+    if (activeSpan != null) {
+      log.debug(
+          "API security sampling decision in {}: hash={}, traceId={}, spanId={}",
+          method,
+          hash,
+          activeSpan.getTraceId(),
+          activeSpan.getSpanId());
+    } else {
+      log.debug(
+          "API security sampling decision in {}: hash={}, traceId=null, spanId=null", method, hash);
+    }
+  }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy

@@ -209,6 +209,161 @@ class ApiSecuritySamplerTest extends DDSpecification {
     }
   }
 
+  void 'preSampleRequest with tracing disabled updates access map immediately'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'first request is presampled with tracing disabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'request is sampled and access map is updated immediately'
+    preSampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(ctx.getApiSecurityEndpointHash())
+
+    when: 'second request for same endpoint is attempted'
+    def ctx2 = createContext('route1', 'GET', 200)
+    sampler.releaseOne()
+    def preSampled2 = sampler.preSampleRequest(ctx2)
+
+    then: 'second request is not sampled because endpoint was already updated in first preSampleRequest'
+    !preSampled2
+  }
+
+  void 'sampleRequest with tracing disabled returns true without updating access map'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    ctx.setApiSecurityEndpointHash(42L)
+    ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'sampleRequest is called with tracing disabled'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'request is sampled without updating access map'
+    sampled
+    sampler.accessMap.size() == 0
+  }
+
+  void 'preSampleRequest with tracing enabled does not update access map immediately'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'true')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'request is presampled with tracing enabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'request is sampled but access map is NOT updated yet'
+    preSampled
+    sampler.accessMap.size() == 0
+
+    when: 'sampleRequest is called to finalize sampling'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'access map is updated in sampleRequest'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(ctx.getApiSecurityEndpointHash())
+  }
+
+  void 'sampleRequest with tracing enabled updates access map'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'true')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    ctx.setApiSecurityEndpointHash(42L)
+    ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'sampleRequest is called with tracing enabled'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'request is sampled and access map is updated'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(42L)
+
+    when: 'second request for same endpoint is made'
+    def sampled2 = sampler.sampleRequest(ctx)
+
+    then: 'second request is not sampled'
+    !sampled2
+  }
+
+  void 'concurrent requests with tracing disabled do not see expired state'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx1 = createContext('route1', 'GET', 200)
+    def ctx2 = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'first request is presampled'
+    def preSampled1 = sampler.preSampleRequest(ctx1)
+
+    then: 'first request is sampled and access map is updated immediately'
+    preSampled1
+    ctx1.getApiSecurityEndpointHash() != null
+    sampler.accessMap.size() == 1
+
+    when: 'concurrent second request tries to presample same endpoint'
+    sampler.releaseOne()
+    def preSampled2 = sampler.preSampleRequest(ctx2)
+
+    then: 'second request is not sampled because endpoint is already in access map'
+    !preSampled2
+  }
+
+  void 'full flow with tracing disabled updates map only in preSampleRequest'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'request goes through full sampling flow with tracing disabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'preSampleRequest returns true and updates access map'
+    preSampled
+    sampler.accessMap.size() == 1
+    def hash = ctx.getApiSecurityEndpointHash()
+    sampler.accessMap.containsKey(hash)
+
+    when: 'sampleRequest is called'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'sampleRequest returns true without modifying access map'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.get(hash) == 0L // Still has the value from preSampleRequest
+  }
+
   private static AppSecRequestContext createContext(final String route, final String method, int statusCode) {
     final AppSecRequestContext ctx = new AppSecRequestContext()
     ctx.setRoute(route)