Fix appsec.waf.requests telemetry metric

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);
+        incrementErrorCodeMetric(reqCtx, 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();
     }
   }
 

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() {

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();

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:
@@ -1787,7 +1786,6 @@ class WAFModuleSpecification extends DDSpecification {
     1 * wafMetricCollector.wafInit(Waf.LIB_VERSION, _, true)
     1 * ctx.getRaspMetrics()
     1 * ctx.getRaspMetricsCounter()
-    (0..1) * WafMetricCollector.get().wafRequestError() // TODO: remove this line when WAFModule is removed
     1 * wafMetricCollector.raspErrorCode(RuleType.SQL_INJECTION, wafErrorCode.code)
     0 * _
 
@@ -1816,8 +1814,8 @@ class WAFModuleSpecification extends DDSpecification {
     1 * wafContext.run(_, _, _) >> { throw createWafException(wafErrorCode) }
     1 * wafMetricCollector.wafInit(Waf.LIB_VERSION, _, true)
     2 * ctx.getWafMetrics()
-    (0..1) * WafMetricCollector.get().wafRequestError() // TODO: remove this line when WAFModule is removed
     1 * wafMetricCollector.wafErrorCode(wafErrorCode.code)
+    1 * ctx.setWafErrors()
     0 * _
 
     where:

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -20,10 +20,12 @@ import datadog.trace.api.gateway.SubscriptionService
 import datadog.trace.api.http.StoredBodySupplier
 import datadog.trace.api.internal.TraceSegment
 import datadog.trace.api.telemetry.LoginEvent
+import datadog.trace.api.telemetry.WafMetricCollector
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase
 import datadog.trace.test.util.DDSpecification
+import spock.lang.Shared
 
 import java.util.function.BiConsumer
 import java.util.function.BiFunction
@@ -37,6 +39,9 @@ import static datadog.trace.api.telemetry.LoginEvent.SIGN_UP
 
 class GatewayBridgeSpecification extends DDSpecification {
 
+  @Shared
+  protected static final ORIGINAL_METRIC_COLLECTOR = WafMetricCollector.get()
+
   private static final String USER_ID = 'user'
 
   SubscriptionService ig = Mock()
@@ -106,12 +111,16 @@ class GatewayBridgeSpecification extends DDSpecification {
   BiFunction<RequestContext, String, Flow<Void>> userCB
   TriFunction<RequestContext, LoginEvent, String, Flow<Void>> loginEventCB
 
+  WafMetricCollector wafMetricCollector = Mock(WafMetricCollector)
+
   void setup() {
     callInitAndCaptureCBs()
     AppSecSystem.active = true
+    WafMetricCollector.INSTANCE = wafMetricCollector
   }
 
   void cleanup() {
+    WafMetricCollector.INSTANCE  = ORIGINAL_METRIC_COLLECTOR
     bridge.stop()
   }
 
@@ -169,6 +178,13 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * traceSegment.setTagTop('http.request.headers.accept', 'header_value')
     1 * traceSegment.setTagTop('http.response.headers.content-type', 'text/html; charset=UTF-8')
     1 * traceSegment.setTagTop('network.client.ip', '2001::1')
+    1 * mockAppSecCtx.isWafBlocked()
+    1 * mockAppSecCtx.hasWafErrors()
+    1 * mockAppSecCtx.getWafTimeouts()
+    1 * mockAppSecCtx.isWafRequestBlockFailure()
+    1 * mockAppSecCtx.isWafRateLimited()
+    1 * mockAppSecCtx.isWafTruncated()
+    1 * wafMetricCollector.wafRequest(_, _, _, _, _, _, _) // call waf request metric
     flow.result == null
     flow.action == Flow.Action.Noop.INSTANCE
   }

internal-api/build.gradle

@@ -203,6 +203,8 @@ excludedClassesCoverage += [
   'datadog.trace.api.telemetry.LogCollector.RawLogMessage',
   "datadog.trace.api.telemetry.MetricCollector.DistributionSeriesPoint",
   "datadog.trace.api.telemetry.MetricCollector",
+  //Enum
+  "datadog.trace.api.telemetry.WafTruncatedType",
   // stubs
   'datadog.trace.api.profiling.Timing.NoOp',
   'datadog.trace.api.profiling.Timer.NoOp',