wip

Author: smola

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

@@ -1,6 +1,8 @@
 package com.datadog.appsec.api.security;
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import datadog.trace.api.time.SystemTimeSource;
+import datadog.trace.api.time.TimeSource;
 import datadog.trace.util.NonBlockingSemaphore;
 
 import javax.annotation.Nonnull;
@@ -23,18 +25,20 @@ public class ApiSecurityRequestSampler {
   private final Deque<Long> apiAccessQueue; // hashes ordered by access time
   private final long expirationTimeInMs;
   private final int capacity;
+  private final TimeSource timeSource;
 
   final NonBlockingSemaphore counter = NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS);
 
   public ApiSecurityRequestSampler() {
-    this(MAX_SIZE, INTERVAL_SECONDS * 1000);
+    this(MAX_SIZE, INTERVAL_SECONDS * 1000, SystemTimeSource.INSTANCE);
   }
 
-  public ApiSecurityRequestSampler(int capacity, long expirationTimeInMs) {
+  public ApiSecurityRequestSampler(int capacity, long expirationTimeInMs, @Nonnull TimeSource timeSource) {
     this.capacity = capacity;
     this.expirationTimeInMs = expirationTimeInMs;
     this.apiAccessMap = new ConcurrentHashMap<>(MAX_SIZE);
     this.apiAccessQueue = new ConcurrentLinkedDeque<>();
+    this.timeSource = timeSource;
   }
 
   public void preSampleRequest(final @Nonnull AppSecRequestContext ctx) {
@@ -79,12 +83,12 @@ public boolean sampleRequest(AppSecRequestContext ctx) {
    * synchronization for updating data structures is not required.
    */
   public boolean updateApiAccessIfExpired(final long hash) {
-    final long currentTime = System.currentTimeMillis();
+    final long currentTime = timeSource.getCurrentTimeMillis();
 
     // New or updated record
     boolean isNewOrUpdated = false;
     if (!apiAccessMap.containsKey(hash)
-        || currentTime - apiAccessMap.get(hash) > expirationTimeInMs) {
+        || currentTime - apiAccessMap.get(hash) >= expirationTimeInMs) {
 
       cleanupExpiredEntries(currentTime);
 
@@ -107,9 +111,9 @@ public boolean updateApiAccessIfExpired(final long hash) {
   }
 
   public boolean isApiAccessExpired(final long hash) {
-    long currentTime = System.currentTimeMillis();
+    long currentTime = timeSource.getCurrentTimeMillis();
     return !apiAccessMap.containsKey(hash)
-        || currentTime - apiAccessMap.get(hash) > expirationTimeInMs;
+        || currentTime - apiAccessMap.get(hash) >= expirationTimeInMs;
   }
 
   private void cleanupExpiredEntries(final long currentTime) {
@@ -118,7 +122,7 @@ private void cleanupExpiredEntries(final long currentTime) {
       if (oldestHash == null) break;
 
       Long lastAccessTime = apiAccessMap.get(oldestHash);
-      if (lastAccessTime == null || currentTime - lastAccessTime > expirationTimeInMs) {
+      if (lastAccessTime == null || currentTime - lastAccessTime >= expirationTimeInMs) {
         apiAccessQueue.pollFirst(); // remove from head
         apiAccessMap.remove(oldestHash);
       } else {
@@ -137,7 +141,7 @@ private long computeApiHash(final String route, final String method, final int s
 
   public static final class NoOp extends ApiSecurityRequestSampler {
     public NoOp() {
-      super(0, 0);
+      super(0, 0, SystemTimeSource.INSTANCE);
     }
 
     @Override

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

@@ -1,8 +1,11 @@
 package com.datadog.appsec.api.security
 
 import com.datadog.appsec.gateway.AppSecRequestContext
+import datadog.trace.api.time.ControllableTimeSource
 import datadog.trace.test.util.DDSpecification
 
+import java.time.Duration
+
 class ApiSecurityRequestSamplerTest extends DDSpecification {
 
   void 'happy path with single request'() {
@@ -149,6 +152,55 @@ class ApiSecurityRequestSamplerTest extends DDSpecification {
     0 * _
   }
 
+  void 'sampleRequest with null context'() {
+    given:
+    def sampler = new ApiSecurityRequestSampler()
+
+    when:
+    def sampleDecision = sampler.sampleRequest(null)
+
+    then:
+    !sampleDecision
+  }
+
+  void 'sampleRequest honors expiration'() {
+    given:
+    def ctx = createContext('route1', 'GET', 200)
+    ctx.setApiSecurityEndpointHash(42L)
+    ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    ctx = Spy(ctx)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10L
+    final long expirationTimeInNs = expirationTimeInMs * 1_000_000
+    def sampler = new ApiSecurityRequestSampler(10, expirationTimeInMs, timeSource)
+
+    when:
+    def sampleDecision = sampler.sampleRequest(ctx)
+
+    then:
+    sampleDecision
+    1 * ctx.getApiSecurityEndpointHash()
+    0 * _
+
+    when:
+    sampleDecision = sampler.sampleRequest(ctx)
+
+    then: 'second request is not sampled'
+    !sampleDecision
+    1 * ctx.getApiSecurityEndpointHash()
+    0 * _
+
+    when: 'expiration time has passed'
+    timeSource.advance(expirationTimeInNs)
+    sampleDecision = sampler.sampleRequest(ctx)
+
+    then: 'request is sampled again'
+    sampleDecision
+    1 * ctx.getApiSecurityEndpointHash()
+    0 * _
+  }
+
   private AppSecRequestContext createContext(final String route, final String method, int statusCode) {
     final AppSecRequestContext ctx = new AppSecRequestContext()
     ctx.setRoute(route)