Introduced GatewayContext (#7304)

DataDog · Jul 11, 2024 · 8b00e3e · 8b00e3e
1 parent e25c56d
commit 8b00e3e
Show file tree

Hide file tree

Showing 14 changed files with 193 additions and 116 deletions.
diff --git a/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/EventDispatcherBenchmark.java b/dd-java-agent/appsec/src/jmh/java/datadog/appsec/benchmark/EventDispatcherBenchmark.java
@@ -11,6 +11,7 @@
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.event.data.KnownAddresses;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 import java.util.Collections;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -87,8 +88,7 @@ public void onDataAvailable(
                 ChangeableFlow flow,
                 AppSecRequestContext context,
                 DataBundle dataBundle,
-                boolean isTransient,
-                boolean isRasp) {}
+                GatewayContext gatewayContext) {}
 
             @Override
             public Priority getPriority() {

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/blocking/BlockingServiceImpl.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/blocking/BlockingServiceImpl.java
@@ -5,6 +5,7 @@
 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.appsec.api.blocking.BlockingContentType;
 import datadog.appsec.api.blocking.BlockingDetails;
 import datadog.appsec.api.blocking.BlockingService;
@@ -45,7 +46,8 @@ public BlockingDetails shouldBlockUser(@Nonnull String userId) {
       }
       SingletonDataBundle<String> db = new SingletonDataBundle<>(KnownAddresses.USER_ID, userId);
       try {
-        flow = eventProducer.publishDataEvent(subInfo, reqCtx, db, true, false);
+        GatewayContext gwCtx = new GatewayContext(true, false);
+        flow = eventProducer.publishDataEvent(subInfo, reqCtx, db, gwCtx);
         break;
       } catch (ExpiredSubscriberInfoException e) {
         subInfo = null;
@@ -65,7 +67,9 @@ public BlockingDetails shouldBlockUser(@Nonnull String userId) {
 
   @Override
   public boolean tryCommitBlockingResponse(
-      int statusCode, BlockingContentType templateType, Map<String, String> extraHeaders) {
+      int statusCode,
+      @Nonnull BlockingContentType templateType,
+      @Nonnull Map<String, String> extraHeaders) {
     log.info(
         "Will try to commit blocking response statusCode={} templateType={} extraHeaders={}",
         statusCode,

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/DataListener.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/DataListener.java
@@ -2,12 +2,12 @@
 
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 
 public interface DataListener extends OrderedCallback {
   void onDataAvailable(
       ChangeableFlow flow,
       AppSecRequestContext context,
       DataBundle dataBundle,
-      boolean isTransient,
-      boolean isRasp);
+      GatewayContext gatewayContext);
 }
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/EventDispatcher.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/EventDispatcher.java
@@ -4,6 +4,7 @@
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.event.data.KnownAddresses;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 import datadog.trace.api.gateway.Flow;
 import java.util.ArrayList;
 import java.util.BitSet;
@@ -132,20 +133,19 @@ public Flow publishDataEvent(
       DataSubscriberInfo subscribers,
       AppSecRequestContext ctx,
       DataBundle newData,
-      boolean isTransient,
-      boolean isRasp)
+      GatewayContext gwCtx)
       throws ExpiredSubscriberInfoException {
     if (!((DataSubscriberInfoImpl) subscribers).isEventDispatcher(this)) {
       throw new ExpiredSubscriberInfoException();
     }
 
-    if (!isTransient) {
+    if (!gwCtx.isTransient) {
       ctx.addAll(newData);
     }
     ChangeableFlow flow = new ChangeableFlow();
     for (int idx : ((DataSubscriberInfoImpl) subscribers).listenerIndices) {
       try {
-        dataListenersIdx.get(idx).onDataAvailable(flow, ctx, newData, isTransient, isRasp);
+        dataListenersIdx.get(idx).onDataAvailable(flow, ctx, newData, gwCtx);
       } catch (RuntimeException rte) {
         log.warn("AppSec callback exception", rte);
       }

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/EventProducerService.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/EventProducerService.java
@@ -3,6 +3,7 @@
 import com.datadog.appsec.event.data.Address;
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 import datadog.trace.api.gateway.Flow;
 import java.util.Collection;
 
@@ -14,7 +15,7 @@ public interface EventProducerService {
    * set of addresses is identical.
    *
    * <p>The return value is to be passed to {@link #publishDataEvent(DataSubscriberInfo,
-   * AppSecRequestContext, DataBundle, boolean, boolean)}.
+   * AppSecRequestContext, DataBundle, GatewayContext)}.
    *
    * @param newAddresses the addresses contained in the {@link DataBundle} that is to be passed to
    *     <code>publishDataEvent()</code>.
@@ -33,10 +34,9 @@ public interface EventProducerService {
    */
   Flow<Void> publishDataEvent(
       DataSubscriberInfo subscribers,
-      AppSecRequestContext ctx,
+      AppSecRequestContext appSecRequestContext,
       DataBundle newData,
-      boolean isTransient,
-      boolean isRasp)
+      GatewayContext gatewayContext)
       throws ExpiredSubscriberInfoException;
 
   interface DataSubscriberInfo {

diff --git a/...-agent/appsec/src/main/java/com/datadog/appsec/event/ReplaceableEventProducerService.java b/...-agent/appsec/src/main/java/com/datadog/appsec/event/ReplaceableEventProducerService.java
@@ -3,6 +3,7 @@
 import com.datadog.appsec.event.data.Address;
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 import datadog.trace.api.gateway.Flow;
 import java.util.Collection;
 
@@ -21,12 +22,11 @@ public DataSubscriberInfo getDataSubscribers(Address<?>... newAddresses) {
   @Override
   public Flow<Void> publishDataEvent(
       DataSubscriberInfo subscribers,
-      AppSecRequestContext ctx,
+      AppSecRequestContext reqCtx,
       DataBundle newData,
-      boolean isTransient,
-      boolean isRasp)
+      GatewayContext gwCtx)
       throws ExpiredSubscriberInfoException {
-    return cur.publishDataEvent(subscribers, ctx, newData, isTransient, isRasp);
+    return cur.publishDataEvent(subscribers, reqCtx, newData, gwCtx);
   }
 
   @Override

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
@@ -235,7 +235,8 @@ public void init() {
               DataBundle bundle =
                   new SingletonDataBundle<>(KnownAddresses.REQUEST_PATH_PARAMS, data);
               try {
-                return producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+                GatewayContext gwCtx = new GatewayContext(false, false);
+                return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
               } catch (ExpiredSubscriberInfoException e) {
                 pathParamsSubInfo = null;
               }
@@ -269,7 +270,8 @@ public void init() {
             DataBundle bundle =
                 new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_RAW, bodyContent);
             try {
-              return producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+              GatewayContext gwCtx = new GatewayContext(false, false);
+              return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
             } catch (ExpiredSubscriberInfoException e) {
               rawRequestBodySubInfo = null;
             }
@@ -306,7 +308,8 @@ public void init() {
                   new SingletonDataBundle<>(
                       KnownAddresses.REQUEST_BODY_OBJECT, ObjectIntrospection.convert(obj));
               try {
-                return producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+                GatewayContext gwCtx = new GatewayContext(false, false);
+                return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
               } catch (ExpiredSubscriberInfoException e) {
                 requestBodySubInfo = null;
               }
@@ -385,7 +388,8 @@ public void init() {
             DataBundle bundle =
                 new SingletonDataBundle<>(KnownAddresses.GRPC_SERVER_METHOD, method);
             try {
-              return producerService.publishDataEvent(subInfo, ctx, bundle, true, false);
+              GatewayContext gwCtx = new GatewayContext(true, false);
+              return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
             } catch (ExpiredSubscriberInfoException e) {
               grpcServerMethodSubInfo = null;
             }
@@ -413,7 +417,8 @@ public void init() {
             DataBundle bundle =
                 new SingletonDataBundle<>(KnownAddresses.GRPC_SERVER_REQUEST_MESSAGE, convObj);
             try {
-              return producerService.publishDataEvent(subInfo, ctx, bundle, true, false);
+              GatewayContext gwCtx = new GatewayContext(true, false);
+              return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
             } catch (ExpiredSubscriberInfoException e) {
               grpcServerRequestMsgSubInfo = null;
             }
@@ -440,7 +445,8 @@ public void init() {
             DataBundle bundle =
                 new SingletonDataBundle<>(KnownAddresses.GRAPHQL_SERVER_ALL_RESOLVERS, data);
             try {
-              return producerService.publishDataEvent(subInfo, ctx, bundle, true, false);
+              GatewayContext gwCtx = new GatewayContext(true, false);
+              return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
             } catch (ExpiredSubscriberInfoException e) {
               graphqlServerRequestMsgSubInfo = null;
             }
@@ -481,7 +487,8 @@ public void init() {
                     .add(KnownAddresses.DB_SQL_QUERY, sql)
                     .build();
             try {
-              return producerService.publishDataEvent(subInfo, ctx, bundle, false, true);
+              GatewayContext gwCtx = new GatewayContext(false, true);
+              return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
             } catch (ExpiredSubscriberInfoException e) {
               dbSqlQuerySubInfo = null;
             }
@@ -678,7 +685,8 @@ private Flow<Void> maybePublishRequestData(AppSecRequestContext ctx) {
       }
 
       try {
-        return producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+        GatewayContext gwCtx = new GatewayContext(false, false);
+        return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
       } catch (ExpiredSubscriberInfoException e) {
         this.initialReqDataSubInfo = null;
       }
@@ -710,7 +718,8 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
       }
 
       try {
-        return producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+        GatewayContext gwCtx = new GatewayContext(false, false);
+        return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
       } catch (ExpiredSubscriberInfoException e) {
         respDataSubInfo = null;
       }
@@ -742,7 +751,8 @@ private void maybeExtractSchemas(AppSecRequestContext ctx) {
               KnownAddresses.WAF_CONTEXT_PROCESSOR,
               Collections.singletonMap("extract-schema", true));
       try {
-        producerService.publishDataEvent(subInfo, ctx, bundle, false, false);
+        GatewayContext gwCtx = new GatewayContext(false, false);
+        producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
         return;
       } catch (ExpiredSubscriberInfoException e) {
         requestEndSubInfo = null;

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayContext.java
@@ -0,0 +1,11 @@
+package com.datadog.appsec.gateway;
+
+public class GatewayContext {
+  public final boolean isTransient;
+  public final boolean isRasp;
+
+  public GatewayContext(final boolean isTransient, final boolean isRasp) {
+    this.isTransient = isTransient;
+    this.isRasp = isRasp;
+  }
+}
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java
@@ -13,6 +13,7 @@
 import com.datadog.appsec.event.data.DataBundle;
 import com.datadog.appsec.event.data.KnownAddresses;
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import com.datadog.appsec.gateway.GatewayContext;
 import com.datadog.appsec.gateway.RateLimiter;
 import com.datadog.appsec.report.AppSecEvent;
 import com.datadog.appsec.stack_trace.StackTraceEvent;
@@ -417,8 +418,7 @@ public void onDataAvailable(
         ChangeableFlow flow,
         AppSecRequestContext reqCtx,
         DataBundle newData,
-        boolean isTransient,
-        boolean isRasp) {
+        GatewayContext gwCtx) {
       Powerwaf.ResultWithData resultWithData;
       CtxAndAddresses ctxAndAddr = ctxAndAddresses.get();
       if (ctxAndAddr == null) {
@@ -433,7 +433,7 @@ public void onDataAvailable(
       }
 
       try {
-        resultWithData = doRunPowerwaf(reqCtx, newData, ctxAndAddr, isTransient, isRasp);
+        resultWithData = doRunPowerwaf(reqCtx, newData, ctxAndAddr, gwCtx);
       } catch (TimeoutPowerwafException tpe) {
         reqCtx.increaseTimeouts();
         log.debug(LogCollector.EXCLUDE_TELEMETRY, "Timeout calling the WAF", tpe);
@@ -587,20 +587,20 @@ private Powerwaf.ResultWithData doRunPowerwaf(
         AppSecRequestContext reqCtx,
         DataBundle newData,
         CtxAndAddresses ctxAndAddr,
-        boolean isTransient,
-        boolean isRasp)
+        GatewayContext gwCtx)
         throws AbstractPowerwafException {
 
-      Additive additive = reqCtx.getOrCreateAdditive(ctxAndAddr.ctx, wafMetricsEnabled, isRasp);
+      Additive additive =
+          reqCtx.getOrCreateAdditive(ctxAndAddr.ctx, wafMetricsEnabled, gwCtx.isRasp);
       PowerwafMetrics metrics;
-      if (isRasp) {
+      if (gwCtx.isRasp) {
         metrics = reqCtx.getRaspMetrics();
         reqCtx.getRaspMetricsCounter().incrementAndGet();
       } else {
         metrics = reqCtx.getWafMetrics();
       }
 
-      if (isTransient) {
+      if (gwCtx.isTransient) {
         return runPowerwafTransient(additive, metrics, newData, ctxAndAddr);
       } else {
         return runPowerwafAdditive(additive, metrics, newData, ctxAndAddr);

diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecModuleSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecModuleSpecification.groovy
@@ -5,6 +5,7 @@ import com.datadog.appsec.event.OrderedCallback
 import com.datadog.appsec.event.data.Address
 import com.datadog.appsec.event.data.DataBundle
 import com.datadog.appsec.gateway.AppSecRequestContext
+import com.datadog.appsec.gateway.GatewayContext
 import datadog.trace.test.util.DDSpecification
 
 import static com.datadog.appsec.event.OrderedCallback.Priority.DEFAULT
@@ -33,7 +34,7 @@ class AppSecModuleSpecification extends DDSpecification {
     }
 
     @Override
-    void onDataAvailable(ChangeableFlow flow, AppSecRequestContext context, DataBundle dataBundle, boolean isTransient, boolean isRasp) {
+    void onDataAvailable(ChangeableFlow flow, AppSecRequestContext reqCtx, DataBundle dataBundle, GatewayContext gwCtx) {
     }
 
     @Override

diff --git a/...ppsec/src/test/groovy/com/datadog/appsec/blocking/BlockingServiceImplSpecification.groovy b/...ppsec/src/test/groovy/com/datadog/appsec/blocking/BlockingServiceImplSpecification.groovy
@@ -8,6 +8,7 @@ import com.datadog.appsec.event.OrderedCallback
 import com.datadog.appsec.event.data.DataBundle
 import com.datadog.appsec.event.data.KnownAddresses
 import com.datadog.appsec.gateway.AppSecRequestContext
+import com.datadog.appsec.gateway.GatewayContext
 import datadog.appsec.api.blocking.BlockingContentType
 import datadog.appsec.api.blocking.BlockingDetails
 import datadog.trace.api.gateway.BlockResponseFunction
@@ -41,7 +42,7 @@ class BlockingServiceImplSpecification extends DDSpecification {
       final OrderedCallback.Priority priority = OrderedCallback.Priority.DEFAULT
 
       @Override
-      void onDataAvailable(ChangeableFlow flow, AppSecRequestContext context, DataBundle dataBundle, boolean isTransient, boolean isRasp) {
+      void onDataAvailable(ChangeableFlow flow, AppSecRequestContext reqCtx, DataBundle dataBundle, GatewayContext gwCtx) {
         if (dataBundle.get(KnownAddresses.USER_ID) == 'blocked.user') {
           flow.action = new Flow.Action.RequestBlockingAction(405, BlockingContentType.HTML)
         }