Fix more tests

Author: manuel-alvarez-alvarez

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

@@ -6,21 +6,21 @@
 public class ApiSecurityDownstreamSamplerImpl implements ApiSecurityDownstreamSampler {
 
   private static final long KNUTH_FACTOR = 1111111111111111111L;
-  private static final double SAMPLING_MAX = Math.pow(2, 64) - 1;
-
-  private final AtomicLong globalRequestCount = new AtomicLong(0);
+  private final AtomicLong globalRequestCount;
   private final double threshold;
 
   public ApiSecurityDownstreamSamplerImpl(double rate) {
     threshold = samplingCutoff(rate < 0.0 ? 0 : (rate > 1.0 ? 1 : rate));
+    globalRequestCount = new AtomicLong(0);
   }
 
   private static double samplingCutoff(double rate) {
+    final double max = Math.pow(2, 64) - 1;
     if (rate < 0.5) {
-      return (long) (rate * SAMPLING_MAX) + Long.MIN_VALUE;
+      return (long) (rate * max) + Long.MIN_VALUE;
     }
     if (rate < 1.0) {
-      return (long) ((rate * SAMPLING_MAX) + Long.MIN_VALUE);
+      return (long) ((rate * max) + Long.MIN_VALUE);
     }
     return Long.MAX_VALUE;
   }

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

@@ -238,11 +238,8 @@ public void increaseRaspTimeouts() {
     RASP_TIMEOUTS_UPDATER.incrementAndGet(this);
   }
 
-  public void increaseHttpClientRequestCount() {
-    httpClientRequestCount.incrementAndGet();
-  }
-
   public boolean sampleHttpClientRequest(final long id) {
+    httpClientRequestCount.incrementAndGet();
     synchronized (sampledHttpClientRequests) {
       if (sampledHttpClientRequests.size()
           < Config.get().getApiSecurityMaxDownstreamRequestBodyAnalysis()) {

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

@@ -329,9 +329,8 @@ private Flow<Void> onRequestSession(final RequestContext ctx_, final String sess
   private Flow<Boolean> onHttpClientSampling(RequestContext ctx_, final long requestId) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null) {
-      return new Flow.ResultFlow<>(null);
+      return new Flow.ResultFlow<>(false);
     }
-    ctx.increaseHttpClientRequestCount();
     final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
     return new Flow.ResultFlow<>(sampler.sampleHttpClientRequest(ctx, requestId));
   }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy

@@ -10,6 +10,7 @@ import com.datadog.ddwaf.WafContext
 import com.squareup.moshi.JsonAdapter
 import com.squareup.moshi.Moshi
 import com.squareup.moshi.Types
+import datadog.trace.api.Config
 import datadog.trace.api.telemetry.LogCollector
 import datadog.trace.test.logging.TestLogCollector
 import datadog.trace.test.util.DDSpecification
@@ -433,4 +434,21 @@ class AppSecRequestContextSpecification extends DDSpecification {
     then:
     keys.size() == 0 // No attributes should be added for invalid key paths
   }
+
+  void 'test sampling of requests'() {
+    given:
+    final maxRequests = Config.get().apiSecurityMaxDownstreamRequestBodyAnalysis
+    final context = new AppSecRequestContext()
+    final random = new Random()
+    final requestIds = (0..maxRequests).collect { random.nextLong() }
+
+    when:
+    final map = requestIds.collectEntries{ [(it) : context.sampleHttpClientRequest(it)] }
+
+    then:
+    map.values().count { it } == maxRequests
+    map.each { requestId, sampled ->
+      assert context.isHttpClientRequestSampled(requestId) == sampled
+    }
+  }
 }