Fix more tests

Author: manuel-alvarez-alvarez

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

@@ -1,6 +1,7 @@
 package com.datadog.appsec;
 
 import com.datadog.appsec.api.security.ApiSecurityDownstreamSampler;
+import com.datadog.appsec.api.security.ApiSecurityDownstreamSamplerImpl;
 import com.datadog.appsec.api.security.ApiSecuritySampler;
 import com.datadog.appsec.api.security.ApiSecuritySamplerImpl;
 import com.datadog.appsec.api.security.AppSecSpanPostProcessor;
@@ -218,7 +219,7 @@ private static void maybeInitializeApiSecurity() {
         API_SECURITY_SAMPLER = requestSampler;
 
         final double rate = Config.get().getApiSecurityDownstreamRequestAnalysisSampleRate();
-        API_SECURITY_DOWNSTREAM_SAMPLER = ApiSecurityDownstreamSampler.build(rate);
+        API_SECURITY_DOWNSTREAM_SAMPLER = new ApiSecurityDownstreamSamplerImpl(rate);
       }
     }
   }

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

@@ -20,13 +20,4 @@ public boolean isSampled(AppSecRequestContext ctx, long requestId) {
       return false;
     }
   }
-
-  static ApiSecurityDownstreamSampler build(double rate) {
-    if (rate < 0.0D) {
-      rate = 0.D;
-    } else if (rate > 1.0D) {
-      rate = 1.0D;
-    }
-    return new ApiSecurityDownstreamSamplerImpl(rate);
-  }
 }

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

@@ -12,7 +12,7 @@ public class ApiSecurityDownstreamSamplerImpl implements ApiSecurityDownstreamSa
   private final double threshold;
 
   public ApiSecurityDownstreamSamplerImpl(double rate) {
-    threshold = samplingCutoff(rate);
+    threshold = samplingCutoff(rate < 0.0 ? 0 : (rate > 1.0 ? 1 : rate));
   }
 
   private static double samplingCutoff(double rate) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java

@@ -129,7 +129,7 @@ public interface KnownAddresses {
   Address<Object> IO_NET_REQUEST_BODY = new Address<>("server.io.net.request.body");
 
   /** The status of a network resource being requested (outgoing request) */
-  Address<Integer> IO_NET_RESPONSE_STATUS = new Address<>("server.io.net.response.status");
+  Address<String> IO_NET_RESPONSE_STATUS = new Address<>("server.io.net.response.status");
 
   /** The response headers of a network resource being requested (outgoing request) */
   Address<Map<String, List<String>>> IO_NET_RESPONSE_HEADERS =

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

@@ -149,8 +149,8 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private volatile Long apiSecurityEndpointHash;
   private volatile byte keepType = PrioritySampling.SAMPLER_KEEP;
 
-  private static final AtomicInteger httpClientRequestCount = new AtomicInteger(0);
-  private static final Set<Long> sampledHttpClientRequests = new HashSet<>();
+  private final AtomicInteger httpClientRequestCount = new AtomicInteger(0);
+  private final Set<Long> sampledHttpClientRequests = new HashSet<>();
 
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> WAF_TIMEOUTS_UPDATER =
       AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "wafTimeouts");

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

@@ -386,7 +386,7 @@ private Flow<Void> onHttpClientResponse(RequestContext ctx_, HttpClientResponse
 
     final MapDataBundle.Builder bundleBuilder =
         new MapDataBundle.Builder(CAPACITY_3_4)
-            .add(KnownAddresses.IO_NET_RESPONSE_STATUS, response.getStatus())
+            .add(KnownAddresses.IO_NET_RESPONSE_STATUS, Integer.toString(response.getStatus()))
             .add(KnownAddresses.IO_NET_RESPONSE_HEADERS, response.getHeaders());
     // ignore the response if not sampled
     final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();

dd-java-agent/appsec/src/main/java/com/datadog/appsec/util/BodyParser.java

@@ -33,6 +33,10 @@ static BodyParser forMediaType(final MediaType type) {
     return null;
   }
 
+  static BodyParser forMediaType(final String type) {
+    return forMediaType(MediaType.parse(type));
+  }
+
   class State {
     private int elemsLeft = MAX_ELEMENTS;
     public boolean objectTooDeep = false;
@@ -62,6 +66,7 @@ public BoundedObjectAdapter(final State state) {
         this.state = state;
       }
 
+      @Generated
       @Override
       public void toJson(final JsonWriter writer, @Nullable final Object value) throws IOException {
         throw new UnsupportedOperationException("Parsing-only adapter");
@@ -79,18 +84,16 @@ private Object readValue(final JsonReader r, final int depth) throws IOException
           return null;
         }
 
-        if (state.elemsLeft == 0) {
+        if (state.elemsLeft-- == 0) {
           state.listMapTooLarge = true;
           r.skipValue();
           return null;
         }
-        state.elemsLeft--;
 
         switch (r.peek()) {
           case BEGIN_OBJECT:
             return readObject(r, depth);
           case BEGIN_ARRAY:
-            state.elemsLeft--;
             return readArray(r, depth);
           case STRING:
             String value = r.nextString();
@@ -116,12 +119,9 @@ private Map<String, Object> readObject(final JsonReader r, final int depth)
         r.beginObject();
         while (r.hasNext()) {
           String name = r.nextName();
-          if (state.elemsLeft > 0) {
-            Object val = readValue(r, depth + 1);
+          Object val = readValue(r, depth + 1);
+          if (!state.listMapTooLarge) {
             map.put(name, val);
-          } else {
-            state.listMapTooLarge = true;
-            r.skipValue();
           }
         }
         r.endObject();
@@ -132,11 +132,9 @@ private List<Object> readArray(final JsonReader r, final int depth) throws IOExc
         List<Object> list = new ArrayList<>();
         r.beginArray();
         while (r.hasNext()) {
-          if (state.elemsLeft > 0) {
-            list.add(readValue(r, depth + 1));
-          } else {
-            state.listMapTooLarge = true;
-            r.skipValue();
+          Object value = readValue(r, depth + 1);
+          if (!state.listMapTooLarge) {
+            list.add(value);
           }
         }
         r.endArray();

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

@@ -26,20 +26,21 @@ class ApiSecurityDownstreamSamplerTest extends DDSpecification {
   void 'test sampling algorithm'() {
     given:
     final epsilon = 0.05
+    final expectedRate = rate < 0 ? 0.0 : rate > 1 ? 1.0 : rate
     final ctx = Mock(AppSecRequestContext) {
       sampleHttpClientRequest(_ as long) >> true
       isHttpClientRequestSampled(_ as long) >> true
     }
-    final sampler = ApiSecurityDownstreamSampler.build(expectedRate)
+    final sampler = new ApiSecurityDownstreamSamplerImpl(rate)
 
     when:
     final samples = (1..100).collect { sampler.sampleHttpClientRequest(ctx, it)}
 
     then:
-    final rate = samples.count { it } / samples.size()
-    rate.subtract(expectedRate).abs() <= epsilon
+    final receivedRate = samples.count { it } / samples.size()
+    receivedRate.subtract(expectedRate).abs() <= epsilon
 
     where:
-    expectedRate << [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]
+    rate << [-1.0, 0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0, 2.0]
   }
 }

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

@@ -924,7 +924,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
     { a, b, db, gw -> bundle = db; gatewayContext = gw; NoopFlow.INSTANCE }
     bundle.size() == (sampled ? 3 : 2)
-    bundle.get(KnownAddresses.IO_NET_RESPONSE_STATUS) == status
+    bundle.get(KnownAddresses.IO_NET_RESPONSE_STATUS) == Integer.toString(status)
     bundle.get(KnownAddresses.IO_NET_RESPONSE_HEADERS) == headers
     if (sampled) {
       bundle.get(KnownAddresses.IO_NET_RESPONSE_BODY) == ['Hello': 'World!']

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/util/BodyParserSpecification.groovy

@@ -7,6 +7,19 @@ import java.nio.charset.StandardCharsets
 
 class BodyParserSpecification extends Specification {
 
+  void 'test body parser per media type'() {
+    when:
+    final result = BodyParser.forMediaType(media)
+
+    then:
+    (result != null) == nonNull
+
+    where:
+    media              | nonNull
+    'application/json' | true
+    'application/xml'  | false
+  }
+
   void 'test parse simple JSON object'() {
     given:
     def parser = BodyParser.forJson()