Alternative API Security sampling algorithm

Author: smola

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

@@ -5,6 +5,7 @@
 import com.datadog.appsec.blocking.BlockingServiceImpl;
 import com.datadog.appsec.config.AppSecConfigService;
 import com.datadog.appsec.config.AppSecConfigServiceImpl;
+import com.datadog.appsec.config.TraceSegmentPostProcessor;
 import com.datadog.appsec.ddwaf.WAFModule;
 import com.datadog.appsec.event.EventDispatcher;
 import com.datadog.appsec.event.ReplaceableEventProducerService;
@@ -22,7 +23,6 @@
 import datadog.trace.api.telemetry.ProductChange;
 import datadog.trace.api.telemetry.ProductChangeCollector;
 import datadog.trace.bootstrap.ActiveSubsystems;
-import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -68,18 +68,6 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     EventDispatcher eventDispatcher = new EventDispatcher();
     REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher);
 
-    ApiSecuritySampler requestSampler;
-    if (Config.get().isApiSecurityEnabled()) {
-      requestSampler = new ApiSecuritySamplerImpl();
-      // When DD_API_SECURITY_ENABLED=true, ths post-processor is set even when AppSec is inactive.
-      // This should be low overhead since the post-processor exits early if there's no AppSec
-      // context.
-      SpanPostProcessor.Holder.INSTANCE =
-          new ApiSecurityProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
-    } else {
-      requestSampler = new ApiSecuritySampler.NoOp();
-    }
-
     ConfigurationPoller configurationPoller = sco.configurationPoller(config);
     // may throw and abort startup
     APP_SEC_CONFIG_SERVICE =
@@ -89,11 +77,15 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
 
     sco.createRemaining(config);
 
+    TraceSegmentPostProcessor apiSecurityPostProcessor =
+        Config.get().isApiSecurityEnabled()
+            ? new ApiSecurityProcessor(new ApiSecuritySampler(), REPLACEABLE_EVENT_PRODUCER)
+            : null;
     GatewayBridge gatewayBridge =
         new GatewayBridge(
             gw,
             REPLACEABLE_EVENT_PRODUCER,
-            requestSampler,
+            apiSecurityPostProcessor,
             APP_SEC_CONFIG_SERVICE.getTraceSegmentPostProcessors());
 
     loadModules(eventDispatcher, sco.monitoring);

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

@@ -1,23 +1,25 @@
 package com.datadog.appsec.api.security;
 
+import com.datadog.appsec.config.TraceSegmentPostProcessor;
 import com.datadog.appsec.event.EventProducerService;
 import com.datadog.appsec.event.ExpiredSubscriberInfoException;
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.event.data.KnownAddresses;
 import com.datadog.appsec.event.data.SingletonDataBundle;
 import com.datadog.appsec.gateway.AppSecRequestContext;
 import com.datadog.appsec.gateway.GatewayContext;
-import datadog.trace.api.gateway.RequestContext;
-import datadog.trace.api.gateway.RequestContextSlot;
+import com.datadog.appsec.report.AppSecEvent;
+import datadog.trace.api.Config;
+import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.internal.TraceSegment;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import java.util.Collection;
 import java.util.Collections;
 import javax.annotation.Nonnull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class ApiSecurityProcessor {
+public class ApiSecurityProcessor implements TraceSegmentPostProcessor {
 
   private static final Logger log = LoggerFactory.getLogger(ApiSecurityProcessor.class);
   private final ApiSecuritySampler sampler;
@@ -28,39 +30,22 @@ public ApiSecurityProcessor(ApiSecuritySampler sampler, EventProducerService pro
     this.producerService = producerService;
   }
 
-  public void process(@Nonnull AgentSpan span) {
-    final RequestContext ctx_ = span.getRequestContext();
-    if (ctx_ == null) {
+  @Override
+  public void processTraceSegment(
+      TraceSegment segment, AppSecRequestContext ctx, Collection<AppSecEvent> collectedEvents) {
+    if (segment == null || ctx == null) {
       return;
     }
-    final AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
-    if (ctx == null) {
+    if (!sampler.sample(ctx)) {
+      log.debug("Request not sampled, skipping API security post-processing");
       return;
     }
-
-    try {
-      if (!sampler.sample(ctx)) {
-        log.debug("Request not sampled, skipping API security post-processing");
-        return;
-      }
-      log.debug("Request sampled, processing API security post-processing");
-      extractSchemas(ctx, ctx_.getTraceSegment());
-    } finally {
-      ctx.setKeepOpenForApiSecurityPostProcessing(false);
-      try {
-        // XXX: Close the additive first. This is not strictly needed, but it'll prevent getting it
-        // detected as a
-        // missed request-ended event.
-        ctx.closeWafContext();
-        ctx.close();
-      } catch (Exception e) {
-        log.debug("Error closing AppSecRequestContext", e);
-      }
-      sampler.releaseOne();
-    }
+    log.debug("Request sampled, processing API security post-processing");
+    extractSchemas(ctx, segment);
   }
 
-  private void extractSchemas(final AppSecRequestContext ctx, final TraceSegment traceSegment) {
+  private void extractSchemas(
+      final @Nonnull AppSecRequestContext ctx, final @Nonnull TraceSegment traceSegment) {
     final EventProducerService.DataSubscriberInfo sub =
         producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR);
     if (sub == null || sub.isEmpty()) {
@@ -74,7 +59,12 @@ private void extractSchemas(final AppSecRequestContext ctx, final TraceSegment t
     try {
       GatewayContext gwCtx = new GatewayContext(false);
       producerService.publishDataEvent(sub, ctx, bundle, gwCtx);
-      ctx.commitDerivatives(traceSegment);
+      // TODO: Perhaps do this if schemas have actually been extracted (check when committing
+      // derivatives)
+      traceSegment.setTagTop(Tags.ASM_KEEP, true);
+      if (!Config.get().isApmTracingEnabled()) {
+        traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
+      }
     } catch (ExpiredSubscriberInfoException e) {
       log.debug("Subscriber info expired", e);
     }

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

@@ -2,7 +2,6 @@
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
 import datadog.trace.util.AgentTaskScheduler;
-
 import java.util.Random;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -11,10 +10,10 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
- * Internal map for API Security sampling.
- * See "[RFC-1021] API Security Sampling Algorithm for thread-based concurrency".
+ * Internal map for API Security sampling. See "[RFC-1021] API Security Sampling Algorithm for
+ * thread-based concurrency".
  */
-final public class ApiSecuritySampler {
+public class ApiSecuritySampler {
 
   private static final int DEFAULT_MAX_ITEM_COUNT = 4096;
   private static final int DEFAULT_INTERVAL_SECONDS = 30;
@@ -28,10 +27,20 @@ final public class ApiSecuritySampler {
   private final long maxItemCount;
 
   public ApiSecuritySampler() {
-    this(DEFAULT_MAX_ITEM_COUNT, DEFAULT_INTERVAL_SECONDS, new Random().nextLong(), new DefaultMonotonicClock(), AgentTaskScheduler.INSTANCE);
+    this(
+        DEFAULT_MAX_ITEM_COUNT,
+        DEFAULT_INTERVAL_SECONDS,
+        new Random().nextLong(),
+        new DefaultMonotonicClock(),
+        AgentTaskScheduler.INSTANCE);
   }
 
-  public ApiSecuritySampler(final int maxItemCount, final int intervalSeconds, final long zero, final MonotonicClock clock, Executor executor) {
+  public ApiSecuritySampler(
+      final int maxItemCount,
+      final int intervalSeconds,
+      final long zero,
+      final MonotonicClock clock,
+      Executor executor) {
     table = new AtomicReference<>(new Table(maxItemCount));
     this.maxItemCount = maxItemCount;
     this.intervalSeconds = intervalSeconds;
@@ -154,7 +163,7 @@ public FindSlotResult findSlot(final long key) {
           index = 0;
         }
       } while (index != startIndex);
-      return new FindSlotResult(table[(int)(maxItemCount * 2)], false);
+      return new FindSlotResult(table[(int) (maxItemCount * 2)], false);
     }
 
     static class FindSlotResult {

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

@@ -142,9 +142,6 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   // Used to detect missing request-end event at close.
   private volatile boolean requestEndCalled;
 
-  private volatile boolean keepOpenForApiSecurityPostProcessing;
-  private volatile Long apiSecurityEndpointHash;
-
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> WAF_TIMEOUTS_UPDATER =
       AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "wafTimeouts");
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> RASP_TIMEOUTS_UPDATER =
@@ -343,22 +340,6 @@ public void setRoute(String route) {
     this.route = route;
   }
 
-  public void setKeepOpenForApiSecurityPostProcessing(final boolean flag) {
-    this.keepOpenForApiSecurityPostProcessing = flag;
-  }
-
-  public boolean isKeepOpenForApiSecurityPostProcessing() {
-    return this.keepOpenForApiSecurityPostProcessing;
-  }
-
-  public void setApiSecurityEndpointHash(long hash) {
-    this.apiSecurityEndpointHash = hash;
-  }
-
-  public Long getApiSecurityEndpointHash() {
-    return this.apiSecurityEndpointHash;
-  }
-
   void addRequestHeader(String name, String value) {
     if (finishedRequestHeaders) {
       throw new IllegalStateException("Request headers were said to be finished before");
@@ -554,23 +535,18 @@ public void close() {
     if (!requestEndCalled) {
       log.debug(SEND_TELEMETRY, "Request end event was not called before close");
     }
-    // For API Security, we sometimes keep contexts open for late processing. In that case, this
-    // flag needs to be
-    // later reset by the API Security post-processor and close must be called again.
-    if (!keepOpenForApiSecurityPostProcessing) {
-      if (wafContext != null) {
-        log.debug(
-            SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)");
-        closeWafContext();
-      }
-      collectedCookies = null;
-      requestHeaders.clear();
-      responseHeaders.clear();
-      persistentData.clear();
-      if (derivatives != null) {
-        derivatives.clear();
-        derivatives = null;
-      }
+    if (wafContext != null) {
+      log.debug(
+          SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)");
+      closeWafContext();
+    }
+    collectedCookies = null;
+    requestHeaders.clear();
+    responseHeaders.clear();
+    persistentData.clear();
+    if (derivatives != null) {
+      derivatives.clear();
+      derivatives = null;
     }
   }
 

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

@@ -88,7 +88,7 @@ public class GatewayBridge {
 
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
-  private final ApiSecuritySampler requestSampler;
+  private final TraceSegmentPostProcessor apiSecurityPostProcessor;
   private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors;
 
   // subscriber cache
@@ -114,11 +114,11 @@ public class GatewayBridge {
   public GatewayBridge(
       SubscriptionService subscriptionService,
       EventProducerService producerService,
-      ApiSecuritySampler requestSampler,
+      TraceSegmentPostProcessor apiSecurityPostProcessor,
       List<TraceSegmentPostProcessor> traceSegmentPostProcessors) {
     this.subscriptionService = subscriptionService;
     this.producerService = producerService;
-    this.requestSampler = requestSampler;
+    this.apiSecurityPostProcessor = apiSecurityPostProcessor;
     this.traceSegmentPostProcessors = traceSegmentPostProcessors;
   }
 
@@ -679,22 +679,20 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
     TraceSegment traceSeg = ctx_.getTraceSegment();
     Map<String, Object> tags = spanInfo.getTags();
 
-    if (maybeSampleForApiSecurity(ctx, spanInfo, tags)) {
-      if (!Config.get().isApmTracingEnabled()) {
-        traceSeg.setTagTop(Tags.ASM_KEEP, true);
-        traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
-      }
-    } else {
-      ctx.closeWafContext();
-    }
-
     // AppSec report metric and events for web span only
     if (traceSeg != null) {
       traceSeg.setTagTop("_dd.appsec.enabled", 1);
       traceSeg.setTagTop("_dd.runtime_family", "jvm");
 
       Collection<AppSecEvent> collectedEvents = ctx.transferCollectedEvents();
 
+      final Object route = tags.get(Tags.HTTP_ROUTE);
+      if (route != null) {
+        ctx.setRoute(route.toString());
+      }
+      // TODO: Move this to traceSegmentPostProcessors
+      apiSecurityPostProcessor.processTraceSegment(traceSeg, ctx, null);
+
       for (TraceSegmentPostProcessor pp : this.traceSegmentPostProcessors) {
         pp.processTraceSegment(traceSeg, ctx, collectedEvents);
       }
@@ -748,6 +746,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
         writeRequestHeaders(
             traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders(), false);
       }
+
       // If extracted any derivatives - commit them
       if (!ctx.commitDerivatives(traceSeg)) {
         log.debug("Unable to commit, derivatives will be skipped {}", ctx.getDerivativeKeys());
@@ -765,21 +764,11 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
               );
     }
 
+    ctx.closeWafContext();
     ctx.close();
     return NoopFlow.INSTANCE;
   }
 
-  private boolean maybeSampleForApiSecurity(
-      AppSecRequestContext ctx, IGSpanInfo spanInfo, Map<String, Object> tags) {
-    log.debug("Checking API Security for end of request handler on span: {}", spanInfo.getSpanId());
-    // API Security sampling requires http.route tag.
-    final Object route = tags.get(Tags.HTTP_ROUTE);
-    if (route != null) {
-      ctx.setRoute(route.toString());
-    }
-    return requestSampler.preSampleRequest(ctx);
-  }
-
   private Flow<Void> onRequestHeadersDone(RequestContext ctx_) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null || ctx.isReqDataPublished()) {