Implemented time-based sampling per endpoint

Author: ValentinZakharov

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

@@ -0,0 +1,102 @@
+package com.datadog.appsec.api.security;
+
+import java.util.Deque;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
+
+/**
+ * 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> apiAccessMap; // Map<hash, timestamp>
+  private final Deque<Long> apiAccessQueue; // hashes ordered by access time
+  private final long expirationTimeInMs;
+  private final int capacity;
+
+  public ApiAccessTracker() {
+    this(MAX_SIZE, INTERVAL_SECONDS * 1000);
+  }
+
+  public ApiAccessTracker(int capacity, long expirationTimeInMs) {
+    this.capacity = capacity;
+    this.expirationTimeInMs = expirationTimeInMs;
+    this.apiAccessMap = new ConcurrentHashMap<>();
+    this.apiAccessQueue = new ConcurrentLinkedDeque<>();
+  }
+
+  /**
+   * 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);
+
+    // New or updated record
+    boolean isNewOrUpdated = false;
+    if (!apiAccessMap.containsKey(hash)
+        || currentTime - apiAccessMap.get(hash) > expirationTimeInMs) {
+
+      cleanupExpiredEntries(currentTime);
+
+      apiAccessMap.put(hash, currentTime); // Update timestamp
+      // move hash to the end of the queue
+      apiAccessQueue.remove(hash);
+      apiAccessQueue.addLast(hash);
+      isNewOrUpdated = true;
+
+      // Remove the oldest hash if capacity is reached
+      while (apiAccessQueue.size() > this.capacity) {
+        Long oldestHash = apiAccessQueue.pollFirst();
+        if (oldestHash != null) {
+          apiAccessMap.remove(oldestHash);
+        }
+      }
+    }
+
+    return isNewOrUpdated;
+  }
+
+  private void cleanupExpiredEntries(long currentTime) {
+    while (!apiAccessQueue.isEmpty()) {
+      Long oldestHash = apiAccessQueue.peekFirst();
+      if (oldestHash == null) break;
+
+      Long lastAccessTime = apiAccessMap.get(oldestHash);
+      if (lastAccessTime == null || currentTime - lastAccessTime > expirationTimeInMs) {
+        apiAccessQueue.pollFirst(); // remove from head
+        apiAccessMap.remove(oldestHash);
+      } else {
+        break; // is up-to-date
+      }
+    }
+  }
+
+  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,38 @@
 package com.datadog.appsec.api.security;
 
+import com.datadog.appsec.gateway.AppSecRequestContext;
 import datadog.trace.api.Config;
-import java.util.concurrent.atomic.AtomicLong;
 
 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(AppSecRequestContext ctx) {
+    if (!config.isApiSecurityEnabled() || ctx == null) {
+      return false;
     }
-    return false;
-  }
-
-  public int getSampling() {
-    return sampling;
-  }
 
-  public boolean sampleRequest() {
-    long prevValue = cumulativeCounter.getAndAdd(sampling);
-    long newValue = prevValue + sampling;
-    if (newValue / 100 == prevValue / 100 + 1) {
-      // Sample request
-      return true;
+    String route = ctx.getRoute();
+    if (route == null) {
+      return false;
     }
-    // Skipped by sampling
-    return false;
-  }
 
-  static int computeSamplingParameter(final float pct) {
-    if (pct >= 1) {
-      return 100;
+    String method = ctx.getMethod();
+    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;
+
+    int statusCode = ctx.getResponseStatus();
+    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/gateway/AppSecRequestContext.java

@@ -88,6 +88,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private String scheme;
   private String method;
   private String savedRawURI;
+  private String route;
   private final Map<String, List<String>> requestHeaders = new LinkedHashMap<>();
   private final Map<String, List<String>> responseHeaders = new LinkedHashMap<>();
   private volatile Map<String, List<String>> collectedCookies;
@@ -270,15 +271,15 @@ void setScheme(String scheme) {
     this.scheme = scheme;
   }
 
-  String getMethod() {
+  public String getMethod() {
     return method;
   }
 
   void setMethod(String method) {
     this.method = method;
   }
 
-  String getSavedRawURI() {
+  public String getSavedRawURI() {
     return savedRawURI;
   }
 
@@ -290,6 +291,18 @@ void setRawURI(String savedRawURI) {
     this.savedRawURI = savedRawURI;
   }
 
+  public String getRoute() {
+    return route;
+  }
+
+  void setRoute(String route) {
+    if (this.route != null && this.route.compareToIgnoreCase(route) != 0) {
+      throw new IllegalStateException(
+          "Forbidden attempt to set different route for given request context");
+    }
+    this.route = route;
+  }
+
   void addRequestHeader(String name, String value) {
     if (finishedRequestHeaders) {
       throw new IllegalStateException("Request headers were said to be finished before");

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

@@ -828,8 +828,10 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       }
     }
 
-    if (!spanInfo.isRequiresPostProcessing()) {
-      ctx.close();
+    // Route used in post-processing
+    Object route = spanInfo.getTags().get(Tags.HTTP_ROUTE);
+    if (route instanceof String) {
+      ctx.setRoute((String) route);
     }
     return NoopFlow.INSTANCE;
   }
@@ -889,6 +891,7 @@ private void onRequestHeader(RequestContext ctx_, String name, String value) {
     }
   }
 
+  // This handler is executed in a separate thread due the computation possible overhead
   private void onPostProcessing(RequestContext ctx_) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null) {
@@ -1068,7 +1071,7 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
   private void maybeExtractSchemas(AppSecRequestContext ctx) {
     boolean extractSchema = false;
     if (Config.get().isApiSecurityEnabled() && requestSampler != null) {
-      extractSchema = requestSampler.sampleRequest();
+      extractSchema = requestSampler.sampleRequest(ctx);
     }
 
     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.apiAccessMap.values().iterator().next()
+
+    then: "The access is added"
+    tracker.apiAccessMap.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.apiAccessMap.values().iterator().next()
+
+    then: "The access is updated and moved to the end"
+    tracker.apiAccessMap.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.apiAccessMap.size() == 1
+    !tracker.apiAccessMap.containsKey(tracker.computeApiHash("route1", "GET", 200))
+    tracker.apiAccessMap.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.apiAccessMap.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime
+  }
+}

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

@@ -1,55 +1,35 @@
 package com.datadog.appsec.api.security
 
+import com.datadog.appsec.gateway.AppSecRequestContext
 import datadog.trace.api.Config
 import datadog.trace.test.util.DDSpecification
-import spock.lang.Shared
 
 class ApiSecurityRequestSamplerTest extends DDSpecification {
 
-  @Shared
-  static final float DEFAULT_SAMPLE_RATE = Config.get().getApiSecurityRequestSampleRate()
+  def config = Mock(Config) {
+    isApiSecurityEnabled() >> true
+  }
 
-  void 'Api Security Request Sample Rate'() {
-    given:
-    def config = Spy(Config.get())
-    config.getApiSecurityRequestSampleRate() >> sampleRate
-    def sampler = new ApiSecurityRequestSampler(config)
+  def sampler = new ApiSecurityRequestSampler(config)
 
+  void 'Api Security Sample Request'() {
     when:
-    def numOfRequest = expectedSampledRequests.size()
-    def results = new int[numOfRequest]
-    for (int i = 0; i < numOfRequest; i++) {
-      results[i] = sampler.sampleRequest() ? 1 : 0
+    def span = Mock(AppSecRequestContext) {
+      getSavedRawURI() >> route
+      getMethod() >> method
+      getResponseStatus() >> statusCode
     }
+    def sample = sampler.sampleRequest(span)
 
     then:
-    results == expectedSampledRequests as int[]
+    sample == sampleResult
 
     where:
-    sampleRate               | expectedSampledRequests
-    DEFAULT_SAMPLE_RATE      | [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]  // Default sample rate - 10%
-    0.0                      | [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
-    0.1                      | [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
-    0.25                     | [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]
-    0.33                     | [0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1]
-    0.5                      | [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
-    0.75                     | [0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1]
-    0.9                      | [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]
-    0.99                     | [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
-    1.0                      | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
-    1.25                     | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]    // Wrong sample rate - use 100%
-    -0.5                     | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]    // Wrong sample rate - use 100%
-  }
-
-  void 'update sample rate'() {
-    given:
-    def config = Spy(Config.get())
-    def sampler = new ApiSecurityRequestSampler(config)
-
-    when:
-    sampler.setSampling(0.2)
-
-    then:
-    sampler.sampling == 20
+    method | route    | statusCode | sampleResult
+    'GET'  | 'route1' | 200  | true
+    'GET'  | 'route2' | null | false
+    'GET'  | null     | 404  | false
+    'TOP'  | 999      | 404  | true
+    null   | '999'    | 404  | false
   }
-}
+}