Rework metric appsec.waf.requests

jandro996 · jandro996 · commit c233763637e8 · 2025-04-02T13:23:33.000+02:00
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/ddwaf/WAFModule.java
@@ -40,7 +40,6 @@
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.api.telemetry.LogCollector;
 import datadog.trace.api.telemetry.WafMetricCollector;
-import datadog.trace.api.telemetry.WafTruncatedType;
 import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -440,20 +439,17 @@ public void onDataAvailable(
           WafMetricCollector.get().raspTimeout(gwCtx.raspRuleType);
         } else {
           reqCtx.increaseWafTimeouts();
-          WafMetricCollector.get().wafRequestTimeout();
           log.debug(LogCollector.EXCLUDE_TELEMETRY, "Timeout calling the WAF", tpe);
         }
         return;
       } catch (UnclassifiedWafException e) {
         if (!reqCtx.isWafContextClosed()) {
           log.error("Error calling WAF", e);
         }
-        // TODO this is wrong and will be fixed in another PR
-        WafMetricCollector.get().wafRequestError();
         incrementErrorCodeMetric(gwCtx, e.code);
         return;
       } catch (AbstractWafException e) {
-        incrementErrorCodeMetric(gwCtx, e.code);
+        incrementErrorCodeMetric(reqCtx, gwCtx, e.code);
         return;
       } finally {
         if (log.isDebugEnabled()) {
@@ -467,17 +463,8 @@ public void onDataAvailable(
             final long listMapTooLarge = wafMetrics.getTruncatedListMapTooLargeCount();
             final long objectTooDeep = wafMetrics.getTruncatedObjectTooDeepCount();
 
-            if (stringTooLong > 0) {
-              WafMetricCollector.get()
-                  .wafInputTruncated(WafTruncatedType.STRING_TOO_LONG, stringTooLong);
-            }
-            if (listMapTooLarge > 0) {
-              WafMetricCollector.get()
-                  .wafInputTruncated(WafTruncatedType.LIST_MAP_TOO_LARGE, listMapTooLarge);
-            }
-            if (objectTooDeep > 0) {
-              WafMetricCollector.get()
-                  .wafInputTruncated(WafTruncatedType.OBJECT_TOO_DEEP, objectTooDeep);
+            if (stringTooLong > 0 || listMapTooLarge > 0 || objectTooDeep > 0) {
+              reqCtx.setWafTruncated();
             }
           }
         }
@@ -501,10 +488,12 @@ public void onDataAvailable(
           ActionInfo actionInfo = new ActionInfo(actionType, actionParams);
 
           if ("block_request".equals(actionInfo.type)) {
-            Flow.Action.RequestBlockingAction rba = createBlockRequestAction(actionInfo);
+            Flow.Action.RequestBlockingAction rba =
+                createBlockRequestAction(actionInfo, reqCtx, gwCtx.isRasp);
             flow.setAction(rba);
           } else if ("redirect_request".equals(actionInfo.type)) {
-            Flow.Action.RequestBlockingAction rba = createRedirectRequestAction(actionInfo);
+            Flow.Action.RequestBlockingAction rba =
+                createRedirectRequestAction(actionInfo, reqCtx, gwCtx.isRasp);
             flow.setAction(rba);
           } else if ("generate_stack".equals(actionInfo.type)) {
             if (Config.get().isAppSecStackTraceEnabled()) {
@@ -516,7 +505,9 @@ public void onDataAvailable(
             }
           } else {
             log.info("Ignoring action with type {}", actionInfo.type);
-            WafMetricCollector.get().wafRequestBlockFailure();
+            if (!gwCtx.isRasp) {
+              reqCtx.setWafRequestBlockFailure();
+            }
           }
         }
         Collection<AppSecEvent> events = buildEvents(resultWithData);
@@ -543,12 +534,16 @@ public void onDataAvailable(
             reqCtx.reportEvents(events);
           } else {
             log.debug("Rate limited WAF events");
-            WafMetricCollector.get().wafRequestRateLimited();
+            if (!gwCtx.isRasp) {
+              reqCtx.setWafRateLimited();
+            }
           }
         }
 
         if (flow.isBlocking()) {
-          reqCtx.setBlocked();
+          if (!gwCtx.isRasp) {
+            reqCtx.setWafBlocked();
+          }
         }
       }
 
@@ -557,7 +552,8 @@ public void onDataAvailable(
       }
     }
 
-    private Flow.Action.RequestBlockingAction createBlockRequestAction(ActionInfo actionInfo) {
+    private Flow.Action.RequestBlockingAction createBlockRequestAction(
+        final ActionInfo actionInfo, final AppSecRequestContext reqCtx, final boolean isRasp) {
       try {
         int statusCode;
         Object statusCodeObj = actionInfo.parameters.get("status_code");
@@ -578,12 +574,15 @@ private Flow.Action.RequestBlockingAction createBlockRequestAction(ActionInfo ac
         return new Flow.Action.RequestBlockingAction(statusCode, blockingContentType);
       } catch (RuntimeException cce) {
         log.warn("Invalid blocking action data", cce);
-        WafMetricCollector.get().wafRequestBlockFailure();
+        if (!isRasp) {
+          reqCtx.setWafRequestBlockFailure();
+        }
         return null;
       }
     }
 
-    private Flow.Action.RequestBlockingAction createRedirectRequestAction(ActionInfo actionInfo) {
+    private Flow.Action.RequestBlockingAction createRedirectRequestAction(
+        final ActionInfo actionInfo, final AppSecRequestContext reqCtx, final boolean isRasp) {
       try {
         int statusCode;
         Object statusCodeObj = actionInfo.parameters.get("status_code");
@@ -604,7 +603,9 @@ private Flow.Action.RequestBlockingAction createRedirectRequestAction(ActionInfo
         return Flow.Action.RequestBlockingAction.forRedirect(statusCode, location);
       } catch (RuntimeException cce) {
         log.warn("Invalid blocking action data", cce);
-        WafMetricCollector.get().wafRequestBlockFailure();
+        if (!isRasp) {
+          reqCtx.setWafRequestBlockFailure();
+        }
         return null;
       }
     }
@@ -649,11 +650,13 @@ private Waf.ResultWithData runWafContext(
     }
   }
 
-  private static void incrementErrorCodeMetric(GatewayContext gwCtx, int code) {
+  private static void incrementErrorCodeMetric(
+      AppSecRequestContext reqCtx, GatewayContext gwCtx, int code) {
     if (gwCtx.isRasp) {
       WafMetricCollector.get().raspErrorCode(gwCtx.raspRuleType, code);
     } else {
       WafMetricCollector.get().wafErrorCode(code);
+      reqCtx.setWafErrors();
     }
   }
 
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java
@@ -119,7 +119,13 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private volatile WafMetrics wafMetrics;
   private volatile WafMetrics raspMetrics;
   private final AtomicInteger raspMetricsCounter = new AtomicInteger(0);
-  private volatile boolean blocked;
+
+  private volatile boolean wafBlocked;
+  private volatile boolean wafErrors;
+  private volatile boolean wafTruncated;
+  private volatile boolean wafRequestBlockFailure;
+  private volatile boolean wafRateLimited;
+
   private volatile int wafTimeouts;
   private volatile int raspTimeouts;
 
@@ -174,12 +180,44 @@ public AtomicInteger getRaspMetricsCounter() {
     return raspMetricsCounter;
   }
 
-  public void setBlocked() {
-    this.blocked = true;
+  public void setWafBlocked() {
+    this.wafBlocked = true;
+  }
+
+  public boolean isWafBlocked() {
+    return wafBlocked;
+  }
+
+  public void setWafErrors() {
+    this.wafErrors = true;
+  }
+
+  public boolean hasWafErrors() {
+    return wafErrors;
+  }
+
+  public void setWafTruncated() {
+    this.wafTruncated = true;
+  }
+
+  public boolean isWafTruncated() {
+    return wafTruncated;
+  }
+
+  public void setWafRequestBlockFailure() {
+    this.wafRequestBlockFailure = true;
+  }
+
+  public boolean isWafRequestBlockFailure() {
+    return wafRequestBlockFailure;
+  }
+
+  public void setWafRateLimited() {
+    this.wafRateLimited = true;
   }
 
-  public boolean isBlocked() {
-    return blocked;
+  public boolean isWafRateLimited() {
+    return wafRateLimited;
   }
 
   public void increaseWafTimeouts() {
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
@@ -724,13 +724,16 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
         log.debug("Unable to commit, derivatives will be skipped {}", ctx.getDerivativeKeys());
       }
 
-      if (ctx.isBlocked()) {
-        WafMetricCollector.get().wafRequestBlocked();
-      } else if (!collectedEvents.isEmpty()) {
-        WafMetricCollector.get().wafRequestTriggered();
-      } else {
-        WafMetricCollector.get().wafRequest();
-      }
+      WafMetricCollector.get()
+          .wafRequest(
+              !collectedEvents.isEmpty(), // ruleTriggered
+              ctx.isWafBlocked(), // requestBlocked
+              ctx.hasWafErrors(), // wafError
+              ctx.getWafTimeouts() > 0, // wafTimeout,
+              ctx.isWafRequestBlockFailure(), // blockFailure,
+              ctx.isWafRateLimited(), // rateLimited,
+              ctx.isWafTruncated() // inputTruncated
+              );
     }
 
     ctx.close();
diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/ddwaf/WAFModuleSpecification.groovy
@@ -475,7 +475,7 @@ class WAFModuleSpecification extends DDSpecification {
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext() >> { wafContext.close() }
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     0 * _
 
@@ -560,7 +560,7 @@ class WAFModuleSpecification extends DDSpecification {
     2 * ctx.getWafMetrics()
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     0 * _
   }
@@ -665,7 +665,7 @@ class WAFModuleSpecification extends DDSpecification {
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
     1 * ctx.reportEvents(_)
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     0 * ctx._(*_)
     flow.blocking == true
@@ -731,7 +731,7 @@ class WAFModuleSpecification extends DDSpecification {
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
     1 * ctx.reportEvents(_)
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     0 * ctx._(*_)
     flow.blocking == true
@@ -758,7 +758,7 @@ class WAFModuleSpecification extends DDSpecification {
     1 * ctx.isWafContextClosed() >> false
     1 * ctx.closeWafContext()
     1 * ctx.reportEvents(_)
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     0 * ctx._(*_)
     metrics == null
@@ -817,7 +817,7 @@ class WAFModuleSpecification extends DDSpecification {
     }
     2 * ctx.getWafMetrics() >> metrics
     1 * ctx.reportEvents(*_)
-    1 * ctx.setBlocked()
+    1 * ctx.setWafBlocked()
     1 * ctx.isThrottled(null)
     1 * ctx.isWafContextClosed() >> false
     0 * ctx._(*_)
@@ -1016,7 +1016,6 @@ class WAFModuleSpecification extends DDSpecification {
       wafContext = it[0].openContext() }
     2 * ctx.getWafMetrics()
     1 * ctx.increaseWafTimeouts()
-    1 * wafMetricCollector.get().wafRequestTimeout()
     0 * _
 
     when:
diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java
diff --git a/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy
