From 23466e648b66111f5466b87fe3fa4cd36602d95a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Mon, 16 Sep 2024 19:06:18 +0200 Subject: [PATCH] Fix failing tests --- .../appsec/event/data/KnownAddresses.java | 4 +-- .../data/KnownAddressesSpecification.groovy | 5 ++- .../src/test/groovy/GrizzlyAsyncTest.groovy | 18 ++++++++++ .../api/gateway/InstrumentationGateway.java | 36 +++++++++++++++++++ .../gateway/InstrumentationGatewayTest.java | 16 +++++++++ 5 files changed, 76 insertions(+), 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 index c29536252bd..65649e149ec 100644 --- 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 DB_SQL_QUERY = new Address<>("server.db.statement"); - Address> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor"); - /** Login failure business event */ Address LOGIN_FAILURE = new Address<>("server.business_logic.users.login.failure"); /** Login success business event */ Address LOGIN_SUCCESS = new Address<>("server.business_logic.users.login.success"); + Address> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor"); + static Address forName(String name) { switch (name) { case "server.request.body": diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy index 8c11222ef94..1d285fc7913 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy +++ b/dd-java-agent/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 index 1eb132a13ff..e3e05c61f60 100644 --- 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,4 +1,5 @@ import datadog.appsec.api.blocking.Blocking +import org.glassfish.grizzly.http.server.Request import javax.ws.rs.GET import javax.ws.rs.HeaderParam @@ -6,11 +7,13 @@ 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 index 72166ca1718..45d734f0d11 100644 --- 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>() { + @Override + public Flow apply(RequestContext ctx, UserIdCollectionMode mode, String arg) { + try { + return ((TriFunction>) + 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>() { + @Override + public Flow apply(RequestContext ctx, String arg) { + try { + return ((BiFunction>) 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 index 153a9b8a66a..468a8971980 100644 --- 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); }