Improve detection of missing request end events

Author: smola

dd-java-agent/appsec/build.gradle

@@ -23,6 +23,7 @@ dependencies {
   testImplementation project(':utils:test-utils')
   testImplementation group: 'org.hamcrest', name: 'hamcrest', version: '2.2'
   testImplementation group: 'com.flipkart.zjsonpatch', name: 'zjsonpatch', version: '0.4.11'
+  testImplementation libs.logback.classic
 
   testFixturesApi project(':dd-java-agent:testing')
 }

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

@@ -143,6 +143,9 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   // keep a reference to the last published usr.session_id
   private volatile String sessionId;
 
+  // Used to detect missing request-end event at close.
+  private volatile boolean requestEndCalled;
+
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> WAF_TIMEOUTS_UPDATER =
       AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "wafTimeouts");
   private static final AtomicIntegerFieldUpdater<AppSecRequestContext> RASP_TIMEOUTS_UPDATER =
@@ -584,10 +587,15 @@ public void close() {
   /* Should be accessible from the modules */
 
   public void close(boolean requiresPostProcessing) {
-    if (additive != null || derivatives != null) {
+    if (!requestEndCalled) {
+      log.debug(SEND_TELEMETRY, "Request end event was not called before close");
+    }
+    if (additive != null) {
       log.debug(
           SEND_TELEMETRY, "WAF object had not been closed (probably missed request-end event)");
       closeAdditive();
+    }
+    if (derivatives != null) {
       derivatives = null;
     }
 
@@ -699,4 +707,9 @@ public boolean isThrottled(RateLimiter rateLimiter) {
   public boolean isAdditiveClosed() {
     return additiveClosed;
   }
+
+  /** Must be called during request end event processing. */
+  void setRequestEndCalled() {
+    requestEndCalled = true;
+  }
 }

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

@@ -767,6 +767,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
     if (ctx == null) {
       return NoopFlow.INSTANCE;
     }
+    ctx.setRequestEndCalled();
 
     maybeExtractSchemas(ctx);
 

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

@@ -5,6 +5,8 @@ import com.datadog.appsec.config.CurrentAppSecConfig
 import com.datadog.appsec.event.data.KnownAddresses
 import com.datadog.appsec.event.data.MapDataBundle
 import com.datadog.appsec.report.AppSecEvent
+import datadog.trace.api.telemetry.LogCollector
+import datadog.trace.test.logging.TestLogCollector
 import datadog.trace.util.stacktrace.StackTraceEvent
 import com.datadog.appsec.test.StubAppSecConfigService
 import datadog.trace.test.util.DDSpecification
@@ -316,4 +318,21 @@ class AppSecRequestContextSpecification extends DDSpecification {
     ctx.getWafError(AppSecRequestContext.DD_WAF_RUN_INVALID_ARGUMENT_ERROR) == 0
     ctx.getWafError(0) == 0
   }
+
+  void 'close logs if request end was not called'() {
+    given:
+    TestLogCollector.enable()
+    def ctx = new AppSecRequestContext()
+
+    when:
+    ctx.close()
+
+    then:
+    def log = TestLogCollector.drainCapturedLogs().find { it.message.contains('Request end event was not called before close') }
+    log != null
+    log.marker == LogCollector.SEND_TELEMETRY
+
+    cleanup:
+    TestLogCollector.disable()
+  }
 }

dd-java-agent/appsec/src/test/resources/logback-test.xml

@@ -0,0 +1,12 @@
+<configuration>
+  <appender name="TEST" class="datadog.trace.test.logging.TestLogbackAppender"/>
+  <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+    <encoder>
+      <pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern>
+    </encoder>
+  </appender>
+  <root level="DEBUG">
+    <appender-ref ref="TEST"/>
+    <appender-ref ref="STDOUT"/>
+  </root>
+</configuration>