New Api Security Sampling mechanism

Author: ValentinZakharov

dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java

@@ -72,10 +72,7 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     // may throw and abort startup
     APP_SEC_CONFIG_SERVICE =
         new AppSecConfigServiceImpl(
-            config,
-            configurationPoller,
-            requestSampler,
-            () -> reloadSubscriptions(REPLACEABLE_EVENT_PRODUCER));
+            config, configurationPoller, () -> reloadSubscriptions(REPLACEABLE_EVENT_PRODUCER));
     APP_SEC_CONFIG_SERVICE.init();
 
     sco.createRemaining(config);

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

@@ -0,0 +1,77 @@
+package com.datadog.appsec.api.security;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * The ApiAccessTracker class provides a mechanism to track API access events, managing them within
+ * a specified capacity limit. Each event is associated with a unique combination of route, method,
+ * and status code, which is used to generate a unique key for tracking access timestamps.
+ *
+ * <p>Usage: - When an API access event occurs, the `updateApiAccessIfExpired` method is called with
+ * the route, method, and status code of the API request. - If the access event for the given
+ * parameters is new or has expired (based on the expirationTimeInMs threshold), the event's
+ * timestamp is updated, effectively moving the event to the end of the tracking list. - If the
+ * tracker's capacity is reached, the oldest event is automatically removed to make room for new
+ * events. - This mechanism ensures that the tracker always contains the most recent access events
+ * within the specified capacity limit, with older, less relevant events being discarded.
+ */
+public class ApiAccessTracker {
+  private static final int INTERVAL_SECONDS = 30;
+  private static final int MAX_SIZE = 4096;
+  private final Map<Long, Long> apiAccessLog; // Map<hash, timestamp>
+  private final long expirationTimeInMs;
+
+  public ApiAccessTracker() {
+    this(MAX_SIZE, INTERVAL_SECONDS * 1000);
+  }
+
+  public ApiAccessTracker(int capacity, long expirationTimeInMs) {
+    this.expirationTimeInMs = expirationTimeInMs;
+    this.apiAccessLog =  Collections.synchronizedMap(new LinkedHashMap<Long, Long>() {
+      @Override
+      protected boolean removeEldestEntry(Map.Entry<Long, Long> eldest) {
+        return size() > capacity;
+      }
+    });
+  }
+
+  /**
+   * Updates the API access log with the given route, method, and status code. If the record exists
+   * and is outdated, it is updated by moving to the end of the list. If the record does not exist,
+   * a new record is added. If the capacity limit is reached, the oldest record is removed. Returns
+   * true if the record was updated or added, false otherwise.
+   *
+   * @param route         The route of the API endpoint request
+   * @param method        The method of the API request
+   * @param statusCode    The HTTP response status code of the API request
+   * @return return true if the record was updated or added, false otherwise
+   */
+  public boolean updateApiAccessIfExpired(String route, String method, int statusCode) {
+    long currentTime = System.currentTimeMillis();
+    long hash = computeApiHash(route, method, statusCode);
+
+    synchronized (apiAccessLog) {
+      if (apiAccessLog.containsKey(hash)) {
+        long lastAccessTime = apiAccessLog.get(hash);
+        if (currentTime - lastAccessTime > expirationTimeInMs) {
+          apiAccessLog.put(hash, currentTime);
+          return true;
+        }
+        return false;
+      } else {
+        apiAccessLog.put(hash, currentTime);
+        return true;
+      }
+    }
+  }
+
+  private long computeApiHash(String route, String method, int statusCode) {
+    long result = 17;
+    result = 31 * result + route.hashCode();
+    result = 31 * result + method.hashCode();
+    result = 31 * result + statusCode;
+    return result;
+  }
+}

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

@@ -1,56 +1,45 @@
 package com.datadog.appsec.api.security;
 
 import datadog.trace.api.Config;
-import java.util.concurrent.atomic.AtomicLong;
+import datadog.trace.api.gateway.IGSpanInfo;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import java.util.Map;
 
 public class ApiSecurityRequestSampler {
 
-  private volatile int sampling;
-  private final AtomicLong cumulativeCounter = new AtomicLong();
+  private final ApiAccessTracker apiAccessTracker;
+  private final Config config;
 
   public ApiSecurityRequestSampler(final Config config) {
-    sampling = computeSamplingParameter(config.getApiSecurityRequestSampleRate());
+    this.apiAccessTracker = new ApiAccessTracker();
+    this.config = config;
   }
 
-  /**
-   * Sets the new sampling parameter
-   *
-   * @return {@code true} if the value changed
-   */
-  public boolean setSampling(final float newSamplingFloat) {
-    int newSampling = computeSamplingParameter(newSamplingFloat);
-    if (newSampling != sampling) {
-      sampling = newSampling;
-      cumulativeCounter.set(0); // Reset current sampling counter
-      return true;
+  public boolean sampleRequest(IGSpanInfo span) {
+    if (!config.isApiSecurityEnabled() || span == null) {
+      return false;
     }
-    return false;
-  }
 
-  public int getSampling() {
-    return sampling;
-  }
+    Map<String, Object> tags = span.getTags();
 
-  public boolean sampleRequest() {
-    long prevValue = cumulativeCounter.getAndAdd(sampling);
-    long newValue = prevValue + sampling;
-    if (newValue / 100 == prevValue / 100 + 1) {
-      // Sample request
-      return true;
+    Object routeObj = tags.get(Tags.HTTP_ROUTE);
+    String route = routeObj instanceof String ? (String) routeObj : null;
+    if (route == null) {
+      return false;
     }
-    // Skipped by sampling
-    return false;
-  }
 
-  static int computeSamplingParameter(final float pct) {
-    if (pct >= 1) {
-      return 100;
+    Object methodObj = tags.get(Tags.HTTP_METHOD);
+    String method = methodObj instanceof String ? (String) methodObj : null;
+    if (method == null) {
+      return false;
     }
-    if (pct < 0) {
-      // Api security can only be disabled by setting the sampling to zero, so we set it to 100%.
-      // TODO: We probably want a warning here.
-      return 100;
+
+    Object statusCodeObj = tags.get(Tags.HTTP_STATUS);
+    int statusCode = statusCodeObj instanceof Integer ? (Integer) statusCodeObj : 0;
+    if (statusCode == 0) {
+      return false;
     }
-    return (int) (pct * 100);
+
+    return apiAccessTracker.updateApiAccessIfExpired(route, method, statusCode);
   }
 }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java

@@ -24,7 +24,6 @@
 import static datadog.remoteconfig.Capabilities.CAPABILITY_ENDPOINT_FINGERPRINT;
 
 import com.datadog.appsec.AppSecSystem;
-import com.datadog.appsec.api.security.ApiSecurityRequestSampler;
 import com.datadog.appsec.config.AppSecModuleConfigurer.SubconfigListener;
 import com.datadog.appsec.config.CurrentAppSecConfig.DirtyStatus;
 import com.datadog.appsec.util.AbortStartupException;
@@ -72,8 +71,6 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
   private final Config tracerConfig;
   private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors = new ArrayList<>();
   private final AppSecModuleConfigurer.Reconfiguration reconfiguration;
-  private final ApiSecurityRequestSampler apiSecurityRequestSampler;
-
   private final ConfigurationEndListener applyRemoteConfigListener =
       this::applyRemoteConfigListener;
 
@@ -82,12 +79,10 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
   public AppSecConfigServiceImpl(
       Config tracerConfig,
       ConfigurationPoller configurationPoller,
-      ApiSecurityRequestSampler apiSecurityRequestSampler,
       AppSecModuleConfigurer.Reconfiguration reconfig) {
     this.tracerConfig = tracerConfig;
     this.configurationPoller = configurationPoller;
     this.reconfiguration = reconfig;
-    this.apiSecurityRequestSampler = apiSecurityRequestSampler;
   }
 
   private void subscribeConfigurationPoller() {
@@ -385,7 +380,6 @@ private void applyRemoteConfigListener() {
     // apply ASM_FEATURES configuration first as they might enable AppSec
     final AppSecFeatures features = mergedAsmFeatures.getMergedData();
     setAppSecActivation(features.asm);
-    setApiSecuritySampling(features.apiSecurity);
     setUserIdCollectionMode(features.autoUserInstrum);
 
     if (!AppSecSystem.isActive() || !currentAppSecConfig.dirtyStatus.isAnyDirty()) {
@@ -415,25 +409,6 @@ private void setAppSecActivation(final AppSecFeatures.Asm asm) {
     }
   }
 
-  private void setApiSecuritySampling(final AppSecFeatures.ApiSecurity apiSecurity) {
-    final float newSampling;
-    if (apiSecurity == null) {
-      newSampling = tracerConfig.getApiSecurityRequestSampleRate();
-    } else {
-      newSampling = apiSecurity.requestSampleRate;
-    }
-    if (apiSecurityRequestSampler.setSampling(newSampling)) {
-      int pct = apiSecurityRequestSampler.getSampling();
-      if (pct == 0) {
-        log.info("Api Security is disabled via remote-config");
-      } else {
-        log.info(
-            "Api Security changed via remote-config. New sampling rate is {}% of all requests.",
-            pct);
-      }
-    }
-  }
-
   private void setUserIdCollectionMode(final AppSecFeatures.AutoUserInstrum autoUserInstrum) {
     UserIdCollectionMode current = UserIdCollectionMode.get();
     UserIdCollectionMode newMode;

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecFeatures.java

@@ -3,9 +3,6 @@
 public class AppSecFeatures {
   public Asm asm;
 
-  @com.squareup.moshi.Json(name = "api_security")
-  public ApiSecurity apiSecurity;
-
   @com.squareup.moshi.Json(name = "auto_user_instrum")
   public AutoUserInstrum autoUserInstrum;
 
@@ -18,16 +15,6 @@ public String toString() {
     }
   }
 
-  public static class ApiSecurity {
-    @com.squareup.moshi.Json(name = "request_sample_rate")
-    public Float requestSampleRate;
-
-    @Override
-    public String toString() {
-      return "ApiSecurity{" + "requestSampleRate=" + requestSampleRate + '}';
-    }
-  }
-
   public static class AutoUserInstrum {
     public String mode;
 
@@ -39,13 +26,6 @@ public String toString() {
 
   @Override
   public String toString() {
-    return "AppSecFeatures{"
-        + "asm="
-        + asm
-        + ", apiSecurity="
-        + apiSecurity
-        + ", autoUserInstrum="
-        + autoUserInstrum
-        + '}';
+    return "AppSecFeatures{" + "asm=" + asm + ", autoUserInstrum=" + autoUserInstrum + '}';
   }
 }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/MergedAsmFeatures.java

@@ -31,7 +31,6 @@ public AppSecFeatures getMergedData() {
 
   private AppSecFeatures merge(final AppSecFeatures target, final AppSecFeatures newFeatures) {
     mergeAsm(target, newFeatures.asm);
-    mergeApiSecurity(target, newFeatures.apiSecurity);
     mergeAutoUserInstrum(target, newFeatures.autoUserInstrum);
     return target;
   }
@@ -43,14 +42,6 @@ private void mergeAsm(final AppSecFeatures target, final AppSecFeatures.Asm newV
     target.asm = newValue;
   }
 
-  private void mergeApiSecurity(
-      final AppSecFeatures target, final AppSecFeatures.ApiSecurity newValue) {
-    if (newValue == null || newValue.requestSampleRate == null) {
-      return;
-    }
-    target.apiSecurity = newValue;
-  }
-
   private void mergeAutoUserInstrum(
       final AppSecFeatures target, final AppSecFeatures.AutoUserInstrum newValue) {
     if (newValue == null || newValue.mode == null) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -740,7 +740,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       return NoopFlow.INSTANCE;
     }
 
-    maybeExtractSchemas(ctx);
+    maybeExtractSchemas(ctx, spanInfo);
 
     // WAF call
     ctx.closeAdditive();
@@ -1031,10 +1031,10 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
     }
   }
 
-  private void maybeExtractSchemas(AppSecRequestContext ctx) {
+  private void maybeExtractSchemas(AppSecRequestContext ctx, IGSpanInfo spanInfo) {
     boolean extractSchema = false;
-    if (Config.get().isApiSecurityEnabled() && requestSampler != null) {
-      extractSchema = requestSampler.sampleRequest();
+    if (Config.get().isApiSecurityEnabled() && requestSampler != null && spanInfo != null) {
+      extractSchema = requestSampler.sampleRequest(spanInfo);
     }
 
     if (!extractSchema) {

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

@@ -0,0 +1,55 @@
+package com.datadog.appsec.api.security
+
+import datadog.trace.test.util.DDSpecification
+
+class ApiAccessTrackerTest extends DDSpecification {
+  def "should add new api access and update if expired"() {
+    given: "An ApiAccessTracker with capacity 2 and expiration time 1 second"
+    def tracker = new ApiAccessTracker(2, 1000)
+
+    when: "Adding new api access"
+    tracker.updateApiAccessIfExpired("route1", "GET", 200)
+    def firstAccessTime = tracker.apiAccessLog.values().iterator().next()
+
+    then: "The access is added"
+    tracker.apiAccessLog.size() == 1
+
+    when: "Waiting more than expiration time and adding another access with the same key"
+    Thread.sleep(1100) // Waiting more than 1 second to ensure expiration
+    tracker.updateApiAccessIfExpired("route1", "GET", 200)
+    def secondAccessTime = tracker.apiAccessLog.values().iterator().next()
+
+    then: "The access is updated and moved to the end"
+    tracker.apiAccessLog.size() == 1
+    secondAccessTime > firstAccessTime
+  }
+
+  def "should remove the oldest access when capacity is exceeded"() {
+    given: "An ApiAccessTracker with capacity 1"
+    def tracker = new ApiAccessTracker(1, 1000)
+
+    when: "Adding two api accesses"
+    tracker.updateApiAccessIfExpired("route1", "GET", 200)
+    Thread.sleep(100) // Delay to ensure different timestamps
+    tracker.updateApiAccessIfExpired("route2", "POST", 404)
+
+    then: "The oldest access is removed"
+    tracker.apiAccessLog.size() == 1
+    !tracker.apiAccessLog.containsKey(tracker.computeApiHash("route1", "GET", 200))
+    tracker.apiAccessLog.containsKey(tracker.computeApiHash("route2", "POST", 404))
+  }
+
+  def "should not update access if not expired"() {
+    given: "An ApiAccessTracker with a short expiration time"
+    def tracker = new ApiAccessTracker(2, 2000) // 2 seconds expiration
+
+    when: "Adding an api access and updating it before it expires"
+    tracker.updateApiAccessIfExpired("route1", "GET", 200)
+    def updateTime = System.currentTimeMillis()
+    boolean updatedBeforeExpiration = tracker.updateApiAccessIfExpired("route1", "GET", 200)
+
+    then: "The access is not updated"
+    !updatedBeforeExpiration
+    tracker.apiAccessLog.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime
+  }
+}