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 b71fd1b1b8e..7f09c528f50 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 @@ -451,7 +451,7 @@ public void init() { (ctx_, dbType) -> { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null) { - return NoopFlow.INSTANCE; + return; } while (true) { DataSubscriberInfo subInfo = dbConnectionSubInfo; @@ -460,11 +460,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 a6ad943a72b..2a046cdfada 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 @@ -23,6 +23,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 @@ -77,7 +78,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; 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 00ea5a10e7f..932df32abbe 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 @@ -27,6 +27,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import java.util.Map; import java.util.concurrent.atomic.AtomicReferenceArray; +import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -365,6 +366,17 @@ public Flow apply(RequestContext ctx, Integer status) { } }; case DATABASE_CONNECTION_ID: + return (C) + new BiConsumer() { + @Override + public void accept(RequestContext ctx, String arg) { + try { + ((BiConsumer) callback).accept(ctx, arg); + } catch (Throwable t) { + log.warn("Callback for {} threw.", eventType, t); + } + } + }; case DATABASE_SQL_QUERY_ID: return (C) new BiFunction>() { 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 07ca0602b70..ff69eaf18e6 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 @@ -201,7 +201,7 @@ public void testNormalCalls() { assertThat(cbp.getCallback(events.graphqlServerRequestMessage()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE); ss.registerCallback(events.databaseConnection(), callback); - cbp.getCallback(events.databaseConnection()).apply(null, null); + cbp.getCallback(events.databaseConnection()).accept(null, null); ss.registerCallback(events.databaseSqlQuery(), callback); cbp.getCallback(events.databaseSqlQuery()).apply(null, null); assertThat(callback.count).isEqualTo(Events.MAX_EVENTS); @@ -257,7 +257,7 @@ public void testThrowableBlocking() { assertThat(cbp.getCallback(events.graphqlServerRequestMessage()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE); ss.registerCallback(events.databaseConnection(), throwback); - cbp.getCallback(events.databaseConnection()).apply(null, null); + cbp.getCallback(events.databaseConnection()).accept(null, null); ss.registerCallback(events.databaseSqlQuery(), throwback); cbp.getCallback(events.databaseSqlQuery()).apply(null, null); assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS);