Fix failing tests

DataDog · Sep 18, 2024 · 23466e6 · 23466e6
1 parent faa2009
commit 23466e6
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 3 deletions.
diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java
@@ -125,14 +125,14 @@ public interface KnownAddresses {
   /** The SQL query being executed */
   Address<String> DB_SQL_QUERY = new Address<>("server.db.statement");
 
-  Address<Map<String, Object>> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor");
-
   /** Login failure business event */
   Address<String> LOGIN_FAILURE = new Address<>("server.business_logic.users.login.failure");
 
   /** Login success business event */
   Address<String> LOGIN_SUCCESS = new Address<>("server.business_logic.users.login.success");
 
+  Address<Map<String, Object>> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor");
+
   static Address<?> forName(String name) {
     switch (name) {
       case "server.request.body":

diff --git a/...t/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy b/...t/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy
@@ -36,13 +36,16 @@ class KnownAddressesSpecification extends Specification {
       'server.db.system',
       'server.db.statement',
       'usr.id',
+      'usr.session_id',
+      'server.business_logic.users.login.failure',
+      'server.business_logic.users.login.success',
       'waf.context.processor',
     ]
   }
 
   void 'number of known addresses is expected number'() {
     expect:
-    Address.instanceCount() == 32
+    Address.instanceCount() == 35
     KnownAddresses.WAF_CONTEXT_PROCESSOR.serial == Address.instanceCount() - 1
   }
 }
diff --git a/dd-java-agent/instrumentation/grizzly-2/src/test/groovy/GrizzlyAsyncTest.groovy b/dd-java-agent/instrumentation/grizzly-2/src/test/groovy/GrizzlyAsyncTest.groovy
@@ -1,16 +1,19 @@
 import datadog.appsec.api.blocking.Blocking
+import org.glassfish.grizzly.http.server.Request
 
 import javax.ws.rs.GET
 import javax.ws.rs.HeaderParam
 import javax.ws.rs.Path
 import javax.ws.rs.QueryParam
 import javax.ws.rs.container.AsyncResponse
 import javax.ws.rs.container.Suspended
+import javax.ws.rs.core.Context
 import javax.ws.rs.core.Response
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
 
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SESSION_ID
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED
@@ -121,5 +124,20 @@ class GrizzlyAsyncTest extends GrizzlyTest {
         }
       }
     }
+
+    @GET
+    @Path("session")
+    Response session(@Context Request request, @Suspended AsyncResponse ar) {
+      executor.execute {
+        try {
+          controller(SESSION_ID) {
+            final session = request.getSession(true)
+            ar.resume(Response.status(SESSION_ID.status).entity(session.idInternal).build())
+          }
+        } catch (Exception e) {
+          ar.resume(e)
+        }
+      }
+    }
   }
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java
@@ -6,6 +6,8 @@
 import static datadog.trace.api.gateway.Events.GRAPHQL_SERVER_REQUEST_MESSAGE_ID;
 import static datadog.trace.api.gateway.Events.GRPC_SERVER_METHOD_ID;
 import static datadog.trace.api.gateway.Events.GRPC_SERVER_REQUEST_MESSAGE_ID;
+import static datadog.trace.api.gateway.Events.LOGIN_FAILURE_ID;
+import static datadog.trace.api.gateway.Events.LOGIN_SUCCESS_ID;
 import static datadog.trace.api.gateway.Events.MAX_EVENTS;
 import static datadog.trace.api.gateway.Events.NETWORK_CONNECTION_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_BODY_CONVERTED_ID;
@@ -18,11 +20,14 @@
 import static datadog.trace.api.gateway.Events.REQUEST_INFERRED_CLIENT_ADDRESS_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_METHOD_URI_RAW_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_PATH_PARAMS_ID;
+import static datadog.trace.api.gateway.Events.REQUEST_SESSION_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_STARTED_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_STARTED_ID;
+import static datadog.trace.api.gateway.Events.USER_ID;
 
+import datadog.trace.api.UserIdCollectionMode;
 import datadog.trace.api.function.TriConsumer;
 import datadog.trace.api.function.TriFunction;
 import datadog.trace.api.http.StoredBodySupplier;
@@ -379,6 +384,37 @@ public void accept(RequestContext ctx, String arg) {
                 }
               }
             };
+      case USER_ID:
+      case LOGIN_SUCCESS_ID:
+      case LOGIN_FAILURE_ID:
+        return (C)
+            new TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>>() {
+              @Override
+              public Flow<Void> apply(RequestContext ctx, UserIdCollectionMode mode, String arg) {
+                try {
+                  return ((TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>>)
+                          callback)
+                      .apply(ctx, mode, arg);
+                } catch (Throwable t) {
+                  log.warn("Callback for {} threw.", eventType, t);
+                  return Flow.ResultFlow.empty();
+                }
+              }
+            };
+      case REQUEST_SESSION_ID:
+        return (C)
+            new BiFunction<RequestContext, String, Flow<Void>>() {
+              @Override
+              public Flow<Void> apply(RequestContext ctx, String arg) {
+                try {
+                  return ((BiFunction<RequestContext, String, Flow<Void>>) callback)
+                      .apply(ctx, arg);
+                } catch (Throwable t) {
+                  log.warn("Callback for {} threw.", eventType, t);
+                  return Flow.ResultFlow.empty();
+                }
+              }
+            };
       case DATABASE_SQL_QUERY_ID:
       case NETWORK_CONNECTION_ID:
       case FILE_LOADED_ID:

diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java
@@ -208,6 +208,14 @@ public void testNormalCalls() {
     cbp.getCallback(events.networkConnection()).apply(null, null);
     ss.registerCallback(events.fileLoaded(), callback);
     cbp.getCallback(events.fileLoaded()).apply(null, null);
+    ss.registerCallback(events.userId(), callback);
+    cbp.getCallback(events.userId()).apply(null, null, null);
+    ss.registerCallback(events.requestSession(), callback);
+    cbp.getCallback(events.requestSession()).apply(null, null);
+    ss.registerCallback(events.loginSuccess(), callback);
+    cbp.getCallback(events.loginSuccess()).apply(null, null, null);
+    ss.registerCallback(events.loginFailure(), callback);
+    cbp.getCallback(events.loginFailure()).apply(null, null, null);
     assertThat(callback.count).isEqualTo(Events.MAX_EVENTS);
   }
 
@@ -268,6 +276,14 @@ public void testThrowableBlocking() {
     cbp.getCallback(events.networkConnection()).apply(null, null);
     ss.registerCallback(events.fileLoaded(), throwback);
     cbp.getCallback(events.fileLoaded()).apply(null, null);
+    ss.registerCallback(events.userId(), throwback);
+    cbp.getCallback(events.userId()).apply(null, null, null);
+    ss.registerCallback(events.requestSession(), throwback);
+    cbp.getCallback(events.requestSession()).apply(null, null);
+    ss.registerCallback(events.loginSuccess(), throwback);
+    cbp.getCallback(events.loginSuccess()).apply(null, null, null);
+    ss.registerCallback(events.loginFailure(), throwback);
+    cbp.getCallback(events.loginFailure()).apply(null, null, null);
     assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS);
   }