diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java index f77bdfddfd0..bef0194b1b4 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java @@ -17,6 +17,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import java.util.function.BiConsumer; import java.util.function.BiFunction; public abstract class DatabaseClientDecorator extends ClientDecorator { @@ -151,27 +152,14 @@ protected void processDatabaseType(AgentSpan span, String dbType) { postProcessServiceAndOperationName(span, namingEntry); if (Config.get().isAppSecRaspEnabled() && dbType != null) { - BiFunction> connectDbCallback = + BiConsumer connectDbCallback = AgentTracer.get() .getCallbackProvider(RequestContextSlot.APPSEC) .getCallback(EVENTS.databaseConnection()); if (connectDbCallback != null) { RequestContext ctx = span.getRequestContext(); if (ctx != null) { - Flow flow = connectDbCallback.apply(ctx, dbType); - Flow.Action action = flow.getAction(); - if (action instanceof Flow.Action.RequestBlockingAction) { - BlockResponseFunction brf = ctx.getBlockResponseFunction(); - if (brf != null) { - Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; - brf.tryCommitBlockingResponse( - ctx.getTraceSegment(), - rba.getStatusCode(), - rba.getBlockingContentType(), - rba.getExtraHeaders()); - } - throw new BlockingException("Blocked request (for DB connection)"); - } + connectDbCallback.accept(ctx, dbType); } } } 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 index 6d4c630ddd1..847dc398408 100644 --- 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 @@ -459,7 +459,7 @@ public void init() { (ctx_, dbType) -> { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null) { - return NoopFlow.INSTANCE; + return; } while (true) { DataSubscriberInfo subInfo = dbConnectionSubInfo; @@ -468,11 +468,11 @@ public void init() { dbConnectionSubInfo = subInfo; } if (subInfo == null || subInfo.isEmpty()) { - return NoopFlow.INSTANCE; + return; } DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_TYPE, dbType); try { - return producerService.publishDataEvent(subInfo, ctx, bundle, false); + producerService.publishDataEvent(subInfo, ctx, bundle, false); } catch (ExpiredSubscriberInfoException e) { dbConnectionSubInfo = null; } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 3768afc29de..cab14238ed1 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -24,6 +24,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase import datadog.trace.test.util.DDSpecification +import java.util.function.BiConsumer import java.util.function.BiFunction import java.util.function.Function import java.util.function.Supplier @@ -79,7 +80,7 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction> grpcServerMethodCB BiFunction> grpcServerRequestMessageCB BiFunction, Flow> graphqlServerRequestMessageCB - BiFunction> databaseConnectionCB + BiConsumer databaseConnectionCB BiFunction> databaseSqlQueryCB void setup() { diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java index d4e0f333082..a31c6c09121 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java @@ -6,6 +6,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -209,8 +210,8 @@ public EventType>> grpcServerReque new ET<>("database.connection", DATABASE_CONNECTION_ID); /** A database connection */ @SuppressWarnings("unchecked") - public EventType>> databaseConnection() { - return (EventType>>) DATABASE_CONNECTION; + public EventType> databaseConnection() { + return (EventType>) DATABASE_CONNECTION; } static final int DATABASE_SQL_QUERY_ID = 17;