Improved API Security schema computation performance (#6765)

* Compute API Security schema extraction at the end of request * Better formatting
DataDog · Mar 8, 2024 · b174715 · b174715
1 parent 8887043
commit b174715
Showing 1 changed file with 37 additions and 9 deletions.
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java
@@ -73,6 +73,7 @@ public class GatewayBridge {
   private volatile DataSubscriberInfo pathParamsSubInfo;
   private volatile DataSubscriberInfo respDataSubInfo;
   private volatile DataSubscriberInfo grpcServerRequestMsgSubInfo;
+  private volatile DataSubscriberInfo requestEndSubInfo;
 
   public GatewayBridge(
       SubscriptionService subscriptionService,
@@ -110,6 +111,8 @@ public void init() {
             return NoopFlow.INSTANCE;
           }
 
+          maybeExtractSchemas(ctx);
+
           // WAF call
           ctx.closeAdditive();
 
@@ -558,18 +561,10 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
 
     ctx.setRespDataPublished(true);
 
-    boolean extractSchema = false;
-    if (Config.get().isApiSecurityEnabled() && requestSampler != null) {
-      extractSchema = requestSampler.sampleRequest();
-    }
-
     MapDataBundle bundle =
         MapDataBundle.of(
             KnownAddresses.RESPONSE_STATUS, String.valueOf(ctx.getResponseStatus()),
-            KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, ctx.getResponseHeaders(),
-            // Extract api schema on response stage
-            KnownAddresses.WAF_CONTEXT_PROCESSOR,
-                Collections.singletonMap("extract-schema", extractSchema));
+            KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, ctx.getResponseHeaders());
 
     while (true) {
       DataSubscriberInfo subInfo = respDataSubInfo;
@@ -588,6 +583,39 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
     }
   }
 
+  private void maybeExtractSchemas(AppSecRequestContext ctx) {
+    boolean extractSchema = false;
+    if (Config.get().isApiSecurityEnabled() && requestSampler != null) {
+      extractSchema = requestSampler.sampleRequest();
+    }
+
+    if (!extractSchema) {
+      return;
+    }
+
+    while (true) {
+      DataSubscriberInfo subInfo = requestEndSubInfo;
+      if (subInfo == null) {
+        subInfo = producerService.getDataSubscribers(KnownAddresses.WAF_CONTEXT_PROCESSOR);
+        requestEndSubInfo = subInfo;
+      }
+      if (subInfo == null || subInfo.isEmpty()) {
+        return;
+      }
+
+      DataBundle bundle =
+          new SingletonDataBundle<>(
+              KnownAddresses.WAF_CONTEXT_PROCESSOR,
+              Collections.singletonMap("extract-schema", true));
+      try {
+        producerService.publishDataEvent(subInfo, ctx, bundle, false);
+        return;
+      } catch (ExpiredSubscriberInfoException e) {
+        requestEndSubInfo = null;
+      }
+    }
+  }
+
   private static Map<String, List<String>> parseQueryStringParams(
       String queryString, Charset uriEncoding) {
     if (queryString == null) {