wip

Author: smola

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

@@ -70,9 +70,12 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     EventDispatcher eventDispatcher = new EventDispatcher();
     REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher);
 
-    ApiSecurityRequestSampler requestSampler = new ApiSecurityRequestSampler(config);
+    ApiSecurityRequestSampler requestSampler;
     if (Config.get().isApiSecurityEnabled()) {
-      SpanPostProcessor.Holder.INSTANCE = new AppSecSpanPostProcessor();
+      requestSampler = new ApiSecurityRequestSampler();
+      SpanPostProcessor.Holder.INSTANCE = new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
+    } else {
+      requestSampler = new ApiSecurityRequestSampler.NoOp();
     }
 
     ConfigurationPoller configurationPoller = sco.configurationPoller(config);

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

@@ -1,7 +1,6 @@
 package com.datadog.appsec.api.security;
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
-import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.util.NonBlockingSemaphore;
 
 import java.util.Deque;
@@ -11,15 +10,20 @@
 
 public class ApiSecurityRequestSampler {
 
-  private static final int MAX_POST_PROCESSING_TASKS = 8;
+  /**
+   * A maximum number of request contexts we'll keep open past the end of request at any given time. This will avoid
+   * excessive memory usage in case of a high number of concurrent requests, and should also prevent memory leaks in
+   * case of a bug.
+   */
+  private static final int MAX_POST_PROCESSING_TASKS = 4;
   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;
 
-  private final NonBlockingSemaphore counter = NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS);
+  final NonBlockingSemaphore counter = NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS);
 
   public ApiSecurityRequestSampler() {
     this(MAX_SIZE, INTERVAL_SECONDS * 1000);
@@ -32,12 +36,7 @@ public ApiSecurityRequestSampler(int capacity, long expirationTimeInMs) {
     this.apiAccessQueue = new ConcurrentLinkedDeque<>();
   }
 
-  public void preSampleRequest(final AppSecRequestContext ctx, final Map<String, Object> tags) {
-    final Object route = tags.get(Tags.HTTP_ROUTE);
-    if (route instanceof String) {
-      ctx.setRoute((String) route);
-    }
-
+  public void preSampleRequest(final AppSecRequestContext ctx) {
     if (!isValid(ctx)) {
       return;
     }
@@ -138,4 +137,19 @@ private long computeApiHash(String route, String method, int statusCode) {
     return result;
   }
 
+  public static final class NoOp extends ApiSecurityRequestSampler {
+    public NoOp() {
+      super(0, 0);
+    }
+
+    @Override
+    public void preSampleRequest(AppSecRequestContext ctx) {
+    }
+
+    @Override
+    public boolean sampleRequest(AppSecRequestContext ctx) {
+      return false;
+    }
+  }
+
 }

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

@@ -7,69 +7,77 @@
 import com.datadog.appsec.event.data.SingletonDataBundle;
 import com.datadog.appsec.gateway.AppSecRequestContext;
 import com.datadog.appsec.gateway.GatewayContext;
-import datadog.trace.api.Config;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Collections;
 import java.util.function.BooleanSupplier;
 
 public class AppSecSpanPostProcessor implements SpanPostProcessor {
 
+  private static final Logger log = LoggerFactory.getLogger(AppSecSpanPostProcessor.class);
+  private final ApiSecurityRequestSampler sampler;
+  private final EventProducerService producerService;
+
+  public AppSecSpanPostProcessor(ApiSecurityRequestSampler sampler, EventProducerService producerService) {
+    this.sampler = sampler;
+    this.producerService = producerService;
+  }
+
   @Override
   public void process(AgentSpan span, BooleanSupplier timeoutCheck) {
-    if (timeoutCheck.getAsBoolean()) {
+    final RequestContext ctx_ = span.getRequestContext();
+    if (ctx_ == null) {
       return;
     }
-    final RequestContext ctx = span.getRequestContext();
+    final AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null) {
       return;
     }
-    final AppSecRequestContext appsecCtx = ctx.getData(RequestContextSlot.APPSEC);
-    if (appsecCtx == null) {
+
+    if (!ctx.isKeepOpenForApiSecurityPostProcessing()) {
       return;
     }
 
-    maybeExtractSchemas(appsecCtx);
-    ctx.close();
-    // Decrease the counter to allow the next request to be post-processed
-    postProcessingCounter.release();
+    try {
+      if (timeoutCheck.getAsBoolean()) {
+        return;
+      }
+      if (!sampler.sampleRequest(ctx)) {
+        return;
+      }
+      maybeExtractSchemas(ctx);
+    } finally {
+      ctx.setKeepOpenForApiSecurityPostProcessing(false);
+      try {
+        ctx.close();
+      } catch (Exception e) {
+        log.debug("Error closing AppSecRequestContext", e);
+      }
+      sampler.counter.release();
+    }
   }
 
   private void maybeExtractSchemas(AppSecRequestContext ctx) {
-    boolean extractSchema = false;
-    if (Config.get().isApiSecurityEnabled() && requestSampler != null) {
-      extractSchema = requestSampler.sampleRequest(ctx);
-    }
-
-    if (!extractSchema) {
-      return;
-    }
-
-    while (true) {
-      EventProducerService.DataSubscriberInfo subInfo = requestEndSubInfo;
-      if (subInfo == null) {
-        subInfo = producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR);
-        requestEndSubInfo = subInfo;
-      }
-      if (subInfo == null || subInfo.isEmpty()) {
+      final EventProducerService.DataSubscriberInfo sub = producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR);
+      if (sub == null|| sub.isEmpty()) {
         return;
       }
 
-      DataBundle bundle =
+      final DataBundle bundle =
           new SingletonDataBundle<>(
               KnownAddresses.WAF_CONTEXT_PROCESSOR,
               Collections.singletonMap("extract-schema", true));
       try {
         GatewayContext gwCtx = new GatewayContext(false);
-        producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
-        return;
+        producerService.publishDataEvent(sub, ctx, bundle, gwCtx);
       } catch (ExpiredSubscriberInfoException e) {
-        requestEndSubInfo = null;
+        log.debug("Subscriber info expired", e);
       }
-    }
   }
 
 }

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

@@ -399,6 +399,10 @@ public void setKeepOpenForApiSecurityPostProcessing(final boolean flag) {
     this.keepOpenForApiSecurityPostProcessing = flag;
   }
 
+  public boolean isKeepOpenForApiSecurityPostProcessing() {
+    return this.keepOpenForApiSecurityPostProcessing;
+  }
+
   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/AppSecRequestContext.java

@@ -831,7 +831,12 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
       }
     }
 
-    requestSampler.preSampleRequest(ctx, tags);
+    // API Security sampling requires http.route tag.
+    final Object route = tags.get(Tags.HTTP_ROUTE);
+    if (route instanceof String) {
+      ctx.setRoute((String) route);
+      requestSampler.preSampleRequest(ctx);
+    }
 
     return NoopFlow.INSTANCE;
   }

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

@@ -1,35 +1,73 @@
 package com.datadog.appsec.api.security
 
 import com.datadog.appsec.gateway.AppSecRequestContext
-import datadog.trace.api.Config
 import datadog.trace.test.util.DDSpecification
 
 class ApiSecurityRequestSamplerTest extends DDSpecification {
 
-  def config = Mock(Config) {
-    isApiSecurityEnabled() >> true
+  void 'happy path with single request'() {
+    given:
+    def ctx = Mock(AppSecRequestContext)
+    def sampler = new ApiSecurityRequestSampler()
+
+    when:
+    sampler.preSampleRequest(ctx)
+
+    then:
+    _ * ctx.getRoute() >> 'route1'
+    _ * ctx.getMethod() >> 'GET'
+    _ * ctx.getResponseStatus() >> 200
+    1 * ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    0 * _
+
+    when:
+    def sampleDecision = sampler.sampleRequest(ctx)
+
+    then:
+    sampleDecision
+    _ * ctx.getRoute() >> 'route1'
+    _ * ctx.getMethod() >> 'GET'
+    _ * ctx.getResponseStatus() >> 200
+    _ * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
+    0 * _
   }
 
-  def sampler = new ApiSecurityRequestSampler(config)
+  void 'second request is not sampled for the same endpoint'() {
+    given:
+    AppSecRequestContext ctx1 = Mock(AppSecRequestContext)
+    AppSecRequestContext ctx2 = Mock(AppSecRequestContext)
+    def sampler = new ApiSecurityRequestSampler()
 
-  void 'Api Security Sample Request'() {
     when:
-    def span = Mock(AppSecRequestContext) {
-      getRoute() >> route
-      getMethod() >> method
-      getResponseStatus() >> statusCode
-    }
-    def sample = sampler.sampleRequest(span)
+    sampler.preSampleRequest(ctx1)
+    def sampleDecision = sampler.sampleRequest(ctx1)
 
     then:
-    sample == sampleResult
-
-    where:
-    method | route    | statusCode | sampleResult
-    'GET'  | 'route1' | 200  | true
-    'GET'  | 'route2' | null | false
-    'GET'  | null     | 404  | false
-    'TOP'  | 999      | 404  | true
-    null   | '999'    | 404  | false
+    sampleDecision
+    _ * ctx1.getRoute() >> 'route1'
+    _ * ctx1.getMethod() >> 'GET'
+    _ * ctx1.getResponseStatus() >> 200
+    _ * _
+
+    when:
+    sampler.preSampleRequest(ctx2)
+
+    then:
+    _ * ctx2.getRoute() >> 'route1'
+    _ * ctx2.getMethod() >> 'GET'
+    _ * ctx2.getResponseStatus() >> 200
+    0 * ctx2.setKeepOpenForApiSecurityPostProcessing(_)
+    0 * _
+
+    when:
+    sampleDecision = sampler.sampleRequest(ctx2)
+
+    then:
+    !sampleDecision
+    _ * ctx2.getRoute() >> 'route1'
+    _ * ctx2.getMethod() >> 'GET'
+    _ * ctx2.getResponseStatus() >> 200
+    0 * _
   }
-}
+
+}