wip

DataDog · smola · Mar 27, 2025 · Feb 9, 2025 · Feb 10, 2025 · Feb 11, 2025
commit 64539edea675fcfece7211efd9fefdcab291fdbb
diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java
@@ -2,7 +2,7 @@
 
 import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED;
 
-import datadog.trace.util.NonBlockingSemaphore;
+import com.datadog.iast.util.NonBlockingSemaphore;
 
 public class OverheadContext {
 

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java
@@ -4,6 +4,7 @@
 
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.IastSystem;
+import com.datadog.iast.util.NonBlockingSemaphore;
 import datadog.trace.api.Config;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
@@ -12,7 +13,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.util.AgentTaskScheduler;
-import datadog.trace.util.NonBlockingSemaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.Nullable;

diff --git a/...adog/trace/util/NonBlockingSemaphore.java → ...tadog/iast/util/NonBlockingSemaphore.java b/...adog/trace/util/NonBlockingSemaphore.java → ...tadog/iast/util/NonBlockingSemaphore.java
@@ -1,4 +1,4 @@
-package datadog.trace.util;
+package com.datadog.iast.util;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;

diff --git a/...race/util/NonBlockingSemaphoreTest.groovy → ...iast/util/NonBlockingSemaphoreTest.groovy b/...race/util/NonBlockingSemaphoreTest.groovy → ...iast/util/NonBlockingSemaphoreTest.groovy
@@ -1,4 +1,4 @@
-package datadog.trace.util
+package com.datadog.iast.util
 
 
 import groovy.transform.CompileDynamic

@@ -82,6 +82,7 @@ ext {
     'com.datadog.appsec.config.AppSecFeatures.ApiSecurity',
     'com.datadog.appsec.config.AppSecFeatures.AutoUserInstrum',
     'com.datadog.appsec.event.ReplaceableEventProducerService',
+    'com.datadog.appsec.api.security.ApiSecurityRequestSampler.NoOp',
   ]
   excludedClassesBranchCoverage = [
     'com.datadog.appsec.gateway.GatewayBridge',

@@ -1,6 +1,7 @@
 package com.datadog.appsec;
 
 import com.datadog.appsec.api.security.ApiSecurityRequestSampler;
+import com.datadog.appsec.api.security.AppSecSpanPostProcessor;
 import com.datadog.appsec.blocking.BlockingServiceImpl;
 import com.datadog.appsec.config.AppSecConfigService;
 import com.datadog.appsec.config.AppSecConfigServiceImpl;
@@ -21,6 +22,7 @@
 import datadog.trace.api.telemetry.ProductChange;
 import datadog.trace.api.telemetry.ProductChangeCollector;
 import datadog.trace.bootstrap.ActiveSubsystems;
+import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -66,7 +68,14 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     EventDispatcher eventDispatcher = new EventDispatcher();
     REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher);
 
-    ApiSecurityRequestSampler requestSampler = new ApiSecurityRequestSampler(config);
+    ApiSecurityRequestSampler requestSampler;
+    if (Config.get().isApiSecurityEnabled()) {
+      requestSampler = new ApiSecurityRequestSampler();
+      SpanPostProcessor.Holder.INSTANCE =
+          new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
+    } else {
+      requestSampler = new ApiSecurityRequestSampler.NoOp();
+    }
 
     ConfigurationPoller configurationPoller = sco.configurationPoller(config);
     // may throw and abort startup

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiAccessTracker.java
@@ -2,40 +2,191 @@
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
 import datadog.trace.api.Config;
+import datadog.trace.api.time.SystemTimeSource;
+import datadog.trace.api.time.TimeSource;
+import java.util.Deque;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.Semaphore;
+import javax.annotation.Nonnull;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class ApiSecurityRequestSampler {
 
-  private final ApiAccessTracker apiAccessTracker;
-  private final Config config;
+  private static final Logger log = LoggerFactory.getLogger(ApiSecurityRequestSampler.class);
 
-  public ApiSecurityRequestSampler(final Config config) {
-    this.apiAccessTracker = new ApiAccessTracker();
-    this.config = config;
+  /**
+   * A maximum number of request contexts we'll keep open past the end of request at any given time.
+   * This will avoid excessive memory usage in case of a high number of concurrent requests, and
+   * should also prevent memory leaks.
+   */
+  private static final int MAX_POST_PROCESSING_TASKS = 4;
+  /** Maximum number of entries in the access map. */
+  private static final int MAX_SIZE = 4096;
+  /** Mapping from endpoint hash to last access timestamp in millis. */
+  private final ConcurrentHashMap<Long, Long> accessMap;
+  /** Deque of endpoint hashes ordered by access time. Oldest is always first. */
+  private final Deque<Long> accessDeque;
+
+  private final long expirationTimeInMs;
+  private final int capacity;
+  private final TimeSource timeSource;
+  private final Semaphore counter = new Semaphore(MAX_POST_PROCESSING_TASKS);
+
+  public ApiSecurityRequestSampler() {
+    this(
+        MAX_SIZE,
+        (long) (Config.get().getApiSecuritySampleDelay() * 1_000),
+        SystemTimeSource.INSTANCE);
+  }
+
+  public ApiSecurityRequestSampler(
+      int capacity, long expirationTimeInMs, @Nonnull TimeSource timeSource) {
+    this.capacity = capacity;
+    this.expirationTimeInMs = expirationTimeInMs;
+    this.accessMap = new ConcurrentHashMap<>();
+    this.accessDeque = new ConcurrentLinkedDeque<>();
+    this.timeSource = timeSource;
+  }
+
+  /**
+   * Prepare a request context for later sampling decision. This method should be called at request
+   * end, and is thread-safe. If a request can potentially be sampled, this method will return true.
+   * If this method returns true, the caller MUST call {@link #releaseOne()} once the context is not
+   * needed anymore.
+   */
+  public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) {
+    final String route = ctx.getRoute();
+    if (route == null) {
+      log.debug("Route is null, skipping API security sampling");
+      return false;
+    }
+    final String method = ctx.getMethod();
+    if (method == null) {
+      log.debug("Method is null, skipping API security sampling");
+      return false;
+    }
+    final int statusCode = ctx.getResponseStatus();
+    if (statusCode == 0) {
+      log.debug("Status code is 0, skipping API security sampling");
+      return false;
+    }
+    long hash = computeApiHash(route, method, statusCode);
+    ctx.setApiSecurityEndpointHash(hash);
+    if (!isApiAccessExpired(hash)) {
+      log.debug("API security sampling is not required for this request");
+      return false;
+    }
+    if (counter.tryAcquire()) {
+      log.debug("API security sampling is required for this request (presampled)");
+      ctx.setKeepOpenForApiSecurityPostProcessing(true);
+      return true;
+    }
+    return false;
   }
 
+  /** Get the final sampling decision. This method is NOT thread-safe. */
   public boolean sampleRequest(AppSecRequestContext ctx) {
-    if (!isValid(ctx)) {
+    if (ctx == null) {
       return false;
     }
+    final Long hash = ctx.getApiSecurityEndpointHash();
+    if (hash == null) {
+      // This should never happen, it should have been short-circuited before.
+      return false;
+    }
+    return updateApiAccessIfExpired(hash);
+  }
 
-    return apiAccessTracker.updateApiAccessIfExpired(
-        ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus());
+  /** Release one permit for the sampler. This must be called after processing a span. */
+  public void releaseOne() {
+    counter.release();
   }
 
-  public boolean preSampleRequest(AppSecRequestContext ctx) {
-    if (!isValid(ctx)) {
+  private boolean updateApiAccessIfExpired(final long hash) {
+    final long currentTime = timeSource.getCurrentTimeMillis();
+
+    Long lastAccess = accessMap.get(hash);
+    if (lastAccess != null && currentTime - lastAccess < expirationTimeInMs) {
       return false;
     }
 
-    return apiAccessTracker.isApiAccessExpired(
-        ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus());
+    if (accessMap.put(hash, currentTime) == null) {
+      accessDeque.addLast(hash);
+      // If we added a new entry, we perform purging.
+      cleanupExpiredEntries(currentTime);
+    } else {
+      // This is now the most recently accessed entry.
+      accessDeque.remove(hash);
+      accessDeque.addLast(hash);
+    }
+
+    return true;
+  }
+
+  private boolean isApiAccessExpired(final long hash) {
+    final long currentTime = timeSource.getCurrentTimeMillis();
+    final Long lastAccess = accessMap.get(hash);
+    return lastAccess == null || currentTime - lastAccess >= expirationTimeInMs;
+  }
+
+  private void cleanupExpiredEntries(final long currentTime) {
+    // Purge all expired entries.
+    while (!accessDeque.isEmpty()) {
+      final Long oldestHash = accessDeque.peekFirst();
+      if (oldestHash == null) {
+        // Should never happen
+        continue;
+      }
+
+      final Long lastAccessTime = accessMap.get(oldestHash);
+      if (lastAccessTime == null) {
+        // Should never  happen
+        continue;
+      }
+
+      if (currentTime - lastAccessTime < expirationTimeInMs) {
+        // The oldest hash is up-to-date, so stop here.
+        break;
+      }
+
+      accessDeque.pollFirst();
+      accessMap.remove(oldestHash);
+    }
+
+    // If we went over capacity, remove the oldest entries until we are within the limit.
+    // This should never be more than 1.
+    final int toRemove = accessMap.size() - this.capacity;
+    for (int i = 0; i < toRemove; i++) {
+      Long oldestHash = accessDeque.pollFirst();
+      if (oldestHash != null) {
+        accessMap.remove(oldestHash);
+      }
+    }
+  }
+
+  private long computeApiHash(final String route, final String method, final int statusCode) {
+    long result = 17;
+    result = 31 * result + route.hashCode();
+    result = 31 * result + method.hashCode();
+    result = 31 * result + statusCode;
+    return result;
   }
 
-  private boolean isValid(AppSecRequestContext ctx) {
-    return config.isApiSecurityEnabled()
-        && ctx != null
-        && ctx.getRoute() != null
-        && ctx.getMethod() != null
-        && ctx.getResponseStatus() != 0;
+  public static final class NoOp extends ApiSecurityRequestSampler {
+    public NoOp() {
+      super(0, 0, SystemTimeSource.INSTANCE);
+    }
+
+    @Override
+    public boolean preSampleRequest(@Nonnull AppSecRequestContext ctx) {
+      return false;
+    }
+
+    @Override
+    public boolean sampleRequest(AppSecRequestContext ctx) {
+      return false;
+    }
   }
 }