From 23a8164e90701d92c3a513620f28accfa8b33979 Mon Sep 17 00:00:00 2001 From: ValentinZakharov Date: Fri, 5 Jul 2024 13:25:09 +0200 Subject: [PATCH] Exploit prevention for SQL injection (blocking support) (#7231) * Introduced SQL-injection blocking * Fixed re-throwing BlockingException to change execution flow * Fixed test * Added debug log message when suppress exception in StatementInstrumentation * Logger in separated class * No blocking in database onConnection flow * add smoke test for rasp stack trace * [wip] smoke test rasp blocking * Missing return * Fix blocking test * Fix test with groovy + jdk 11 * SQLi RASP in one shot * Forbidden method invocation: java.lang.Class#forName * Exclude SQL-injection test code * fix appsec.blocked * remove debug level in smoke tests (increases flakiness under load) * add assert * fix tests * Fixed suppress exception logic in StatementInstrumentation * Update dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java Co-authored-by: Santiago M. Mola * Fixed typo in field name * Added RASP info in StatusLogger --------- Co-authored-by: Santiago Mola --- .../decorator/DatabaseClientDecorator.java | 21 +++- .../appsec/gateway/AppSecRequestContext.java | 11 +- .../datadog/appsec/gateway/GatewayBridge.java | 37 +++--- .../appsec/powerwaf/PowerWAFResultData.java | 2 +- .../gateway/GatewayBridgeSpecification.groovy | 2 +- .../jdbc/InstrumentationLogger.java | 11 ++ .../jdbc/StatementInstrumentation.java | 15 ++- dd-smoke-tests/appsec/springboot/build.gradle | 1 + .../springboot/controller/WebController.java | 18 +++ .../appsec/SpringBootSmokeTest.groovy | 109 ++++++++++++++++++ .../AbstractAppSecServerSmokeTest.groovy | 2 + .../java/datadog/trace/core/StatusLogger.java | 2 + gradle/spotbugFilters/exclude.xml | 1 + .../datadog/trace/api/gateway/Events.java | 4 +- .../api/gateway/InstrumentationGateway.java | 17 ++- .../gateway/InstrumentationGatewayTest.java | 4 +- 16 files changed, 220 insertions(+), 37 deletions(-) create mode 100644 dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java 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 151acbb1063..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 @@ -3,9 +3,12 @@ import static datadog.trace.api.gateway.Events.EVENTS; import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE; +import datadog.appsec.api.blocking.BlockingException; import datadog.trace.api.Config; import datadog.trace.api.cache.DDCache; import datadog.trace.api.cache.DDCaches; +import datadog.trace.api.gateway.BlockResponseFunction; +import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.naming.NamingSchema; @@ -15,6 +18,7 @@ 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 { protected static class NamingEntry { @@ -116,14 +120,27 @@ public AgentSpan onStatement(final AgentSpan span, final CharSequence statement) */ public void onRawStatement(AgentSpan span, String sql) { if (Config.get().isAppSecRaspEnabled() && sql != null && !sql.isEmpty()) { - BiConsumer sqlQueryCallback = + BiFunction> sqlQueryCallback = AgentTracer.get() .getCallbackProvider(RequestContextSlot.APPSEC) .getCallback(EVENTS.databaseSqlQuery()); if (sqlQueryCallback != null) { RequestContext ctx = span.getRequestContext(); if (ctx != null) { - sqlQueryCallback.accept(ctx, sql); + Flow flow = sqlQueryCallback.apply(ctx, sql); + 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 SQL query)"); + } } } } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index f01cdf7009b..7a82a1eb8d2 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -94,6 +94,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private String inferredClientIp; private volatile StoredBodySupplier storedRequestBodySupplier; + private String dbType; private int responseStatus; @@ -104,7 +105,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private boolean pathParamsPublished; private Map apiSchemas; - private AtomicBoolean rateLimited = new AtomicBoolean(false); + private final AtomicBoolean rateLimited = new AtomicBoolean(false); private volatile boolean throttled; // should be guarded by this @@ -341,6 +342,14 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier) this.storedRequestBodySupplier = storedRequestBodySupplier; } + public String getDbType() { + return dbType; + } + + public void setDbType(String dbType) { + this.dbType = dbType; + } + public int getResponseStatus() { return responseStatus; } 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 8234f1e4425..67c545d7d5f 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 @@ -1,5 +1,6 @@ package com.datadog.appsec.gateway; +import static com.datadog.appsec.event.data.MapDataBundle.Builder.CAPACITY_0_2; import static com.datadog.appsec.event.data.MapDataBundle.Builder.CAPACITY_6_10; import static com.datadog.appsec.gateway.AppSecRequestContext.DEFAULT_REQUEST_HEADERS_ALLOW_LIST; import static com.datadog.appsec.gateway.AppSecRequestContext.REQUEST_HEADERS_ALLOW_LIST; @@ -83,7 +84,6 @@ public class GatewayBridge { private volatile DataSubscriberInfo grpcServerRequestMsgSubInfo; private volatile DataSubscriberInfo graphqlServerRequestMsgSubInfo; private volatile DataSubscriberInfo requestEndSubInfo; - private volatile DataSubscriberInfo dbConnectionSubInfo; private volatile DataSubscriberInfo dbSqlQuerySubInfo; public GatewayBridge( @@ -454,23 +454,7 @@ public void init() { if (ctx == null) { return; } - while (true) { - DataSubscriberInfo subInfo = dbConnectionSubInfo; - if (subInfo == null) { - subInfo = producerService.getDataSubscribers(KnownAddresses.DB_TYPE); - dbConnectionSubInfo = subInfo; - } - if (subInfo == null || subInfo.isEmpty()) { - return; - } - DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_TYPE, dbType); - try { - producerService.publishDataEvent(subInfo, ctx, bundle, false); - return; - } catch (ExpiredSubscriberInfoException e) { - dbConnectionSubInfo = null; - } - } + ctx.setDbType(dbType); }); subscriptionService.registerCallback( @@ -478,21 +462,26 @@ public void init() { (ctx_, sql) -> { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null) { - return; + return NoopFlow.INSTANCE; } while (true) { DataSubscriberInfo subInfo = dbSqlQuerySubInfo; if (subInfo == null) { - subInfo = producerService.getDataSubscribers(KnownAddresses.DB_SQL_QUERY); + subInfo = + producerService.getDataSubscribers( + KnownAddresses.DB_TYPE, KnownAddresses.DB_SQL_QUERY); dbSqlQuerySubInfo = subInfo; } if (subInfo == null || subInfo.isEmpty()) { - return; + return NoopFlow.INSTANCE; } - DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_SQL_QUERY, sql); + DataBundle bundle = + new MapDataBundle.Builder(CAPACITY_0_2) + .add(KnownAddresses.DB_TYPE, ctx.getDbType()) + .add(KnownAddresses.DB_SQL_QUERY, sql) + .build(); try { - producerService.publishDataEvent(subInfo, ctx, bundle, false); - return; + return producerService.publishDataEvent(subInfo, ctx, bundle, false); } catch (ExpiredSubscriberInfoException e) { dbSqlQuerySubInfo = null; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFResultData.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFResultData.java index 4426514e4ad..4de3a3c2ed2 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFResultData.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFResultData.java @@ -21,7 +21,7 @@ public static class Rule { } public static class Parameter extends MatchInfo { - MatchInfo resources; + MatchInfo resource; MatchInfo params; MatchInfo db_type; List highlight; 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 9017a99a676..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 @@ -79,7 +79,7 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction> grpcServerRequestMessageCB BiFunction, Flow> graphqlServerRequestMessageCB BiConsumer databaseConnectionCB - BiConsumer databaseSqlQueryCB + BiFunction> databaseSqlQueryCB void setup() { callInitAndCaptureCBs() diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java new file mode 100644 index 00000000000..e3e99824f48 --- /dev/null +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java @@ -0,0 +1,11 @@ +package datadog.trace.instrumentation.jdbc; + +import org.slf4j.LoggerFactory; + +public class InstrumentationLogger { + public static void debug( + String instrumentation, final Class target, final Throwable throwable) { + LoggerFactory.getLogger(instrumentation) + .debug("Failed to handle exception in instrumentation for " + target, throwable); + } +} diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index fb55315d191..20b513ab7d6 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -14,6 +14,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.appsec.api.blocking.BlockingException; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.CallDepthThreadLocalMap; @@ -55,7 +56,9 @@ public Map contextStore() { @Override public String[] helperClassNames() { return new String[] { - packageName + ".JDBCDecorator", packageName + ".SQLCommenter", + packageName + ".JDBCDecorator", + packageName + ".SQLCommenter", + packageName + ".InstrumentationLogger", }; } @@ -70,7 +73,7 @@ public void methodAdvice(MethodTransformer transformer) { } public static class StatementAdvice { - @Advice.OnMethodEnter(suppress = Throwable.class) + @Advice.OnMethodEnter() public static AgentScope onEnter( @Advice.Argument(value = 0, readOnly = false) String sql, @Advice.This final Statement statement) { @@ -116,6 +119,14 @@ public static AgentScope onEnter( } catch (SQLException e) { // if we can't get the connection for any reason return null; + } catch (BlockingException e) { + // re-throw blocking exceptions + throw e; + } catch (Throwable e) { + // suppress anything else + InstrumentationLogger.debug( + "datadog.trace.instrumentation.jdbc.StatementInstrumentation", statement.getClass(), e); + return null; } } diff --git a/dd-smoke-tests/appsec/springboot/build.gradle b/dd-smoke-tests/appsec/springboot/build.gradle index 10885e2a635..3dd2fc106d5 100644 --- a/dd-smoke-tests/appsec/springboot/build.gradle +++ b/dd-smoke-tests/appsec/springboot/build.gradle @@ -16,6 +16,7 @@ jar { dependencies { implementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.6.0' implementation(group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.6.0') + implementation group: 'com.h2database', name: 'h2', version: '2.1.212' testImplementation project(':dd-smoke-tests:appsec') } diff --git a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java index e8efc385dfe..fb3e57f90f0 100644 --- a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java +++ b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java @@ -1,10 +1,14 @@ package datadog.smoketest.appsec.springboot.controller; +import java.sql.Connection; +import java.sql.DriverManager; import org.springframework.web.bind.annotation.GetMapping; import org.springframework.web.bind.annotation.PathVariable; import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestHeader; import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RequestParam; import org.springframework.web.bind.annotation.RestController; @RestController @@ -27,4 +31,18 @@ static class BodyMappedClass { public String requestBody(@RequestBody BodyMappedClass obj) { return obj.v; } + + @GetMapping("/sqli/query") + public String sqliQuery(@RequestParam("id") String id) throws Exception { + Connection conn = DriverManager.getConnection("jdbc:h2:mem:testdb", "sa", ""); + conn.createStatement().execute("SELECT 1 FROM DUAL WHERE '1' = '" + id + "'"); + return "EXECUTED"; + } + + @GetMapping("/sqli/header") + public String sqliHeader(@RequestHeader("x-custom-header") String id) throws Exception { + Connection conn = DriverManager.getConnection("jdbc:h2:mem:testdb", "sa", ""); + conn.createStatement().execute("SELECT 1 FROM DUAL WHERE '1' = '" + id + "'"); + return "EXECUTED"; + } } diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy index 85cbdd9046e..8add944b625 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy @@ -37,6 +37,48 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { ], transformers: [], on_match : ['block'] + ], + [ + id : '__test_sqli_stacktrace_on_query', + name : 'test rule to generate stacktrace on sqli', + tags : [ + type : 'test', + category : 'test', + confidence: '1', + ], + conditions: [ + [ + parameters: [ + resource: [[address: "server.db.statement"]], + params: [[ address: "server.request.query" ]], + db_type: [[ address: "server.db.system" ]], + ], + operator: "sqli_detector", + ], + ], + transformers: [], + on_match : ['stack_trace'] + ], + [ + id : '__test_sqli_block_on_header', + name : 'test rule to block on sqli', + tags : [ + type : 'test', + category : 'test', + confidence: '1', + ], + conditions: [ + [ + parameters: [ + resource: [[address: "server.db.statement"]], + params: [[ address: "server.request.headers.no_cookies" ]], + db_type: [[ address: "server.db.system" ]], + ], + operator: "sqli_detector", + ], + ], + transformers: [], + on_match : ['block'] ] ]) } @@ -197,4 +239,71 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { assert it.meta.get('appsec.blocked') != null, 'appsec.blocked is not set' } } + + void 'rasp reports stacktrace on sql injection'() { + when: + String url = "http://localhost:${httpPort}/sqli/query?id=' OR 1=1 --" + def request = new Request.Builder() + .url(url) + .get() + .build() + def response = client.newCall(request).execute() + def responseBodyStr = response.body().string() + + then: + response.code() == 200 + responseBodyStr == 'EXECUTED' + + when: + waitForTraceCount(1) + + then: + def rootSpans = this.rootSpans.toList() + rootSpans.size() == 1 + def rootSpan = rootSpans[0] + assert rootSpan.meta.get('appsec.blocked') == null, 'appsec.blocked is set' + assert rootSpan.meta.get('_dd.appsec.json') != null, '_dd.appsec.json is not set' + def trigger + for (t in rootSpan.triggers) { + if (t['rule']['id'] == '__test_sqli_stacktrace_on_query') { + trigger = t + break + } + } + assert trigger != null, 'test trigger not found' + } + + void 'rasp blocks on sql injection'() { + when: + String url = "http://localhost:${httpPort}/sqli/header" + def request = new Request.Builder() + .url(url) + .header("x-custom-header", "' OR 1=1 --") + .get() + .build() + def response = client.newCall(request).execute() + def responseBodyStr = response.body().string() + + then: + response.code() == 403 + responseBodyStr == '{"errors":[{"title":"You\'ve been blocked","detail":"Sorry, you cannot access this page. Please contact the customer service team. Security provided by Datadog."}]}\n' + + when: + waitForTraceCount(1) + + then: + def rootSpans = this.rootSpans.toList() + rootSpans.size() == 1 + def rootSpan = rootSpans[0] + assert rootSpan.meta.get('appsec.blocked') == 'true', 'appsec.blocked is not set' + assert rootSpan.meta.get('_dd.appsec.json') != null, '_dd.appsec.json is not set' + def trigger = null + for (t in rootSpan.triggers) { + if (t['rule']['id'] == '__test_sqli_block_on_header') { + trigger = t + break + } + } + assert trigger != null, 'test trigger not found' + } } diff --git a/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy b/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy index 3e2b1fcdf33..26182d9a049 100644 --- a/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy +++ b/dd-smoke-tests/appsec/src/main/groovy/datadog/smoketest/appsec/AbstractAppSecServerSmokeTest.groovy @@ -43,6 +43,8 @@ abstract class AbstractAppSecServerSmokeTest extends AbstractServerSmokeTest { @Shared protected String[] defaultAppSecProperties = [ "-Ddd.appsec.enabled=${System.getProperty('smoke_test.appsec.enabled') ?: 'true'}", + // TODO: rely on default once its true + "-Ddd.appsec.rasp.enabled=true", "-Ddd.profiling.enabled=false", // disable AppSec rate limit "-Ddd.appsec.trace.rate.limit=-1" diff --git a/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java b/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java index 567b7f26dff..f7e63505d23 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java @@ -118,6 +118,8 @@ public void toJson(JsonWriter writer, Config config) throws IOException { writer.value(config.getAppSecActivation().toString()); writer.name("appsec_rules_file_path"); writer.value(config.getAppSecRulesFile()); + writer.name("rasp_enabled"); + writer.value(config.isAppSecRaspEnabled()); writer.name("telemetry_enabled"); writer.value(config.isTelemetryEnabled()); writer.name("telemetry_dependency_collection_enabled"); diff --git a/gradle/spotbugFilters/exclude.xml b/gradle/spotbugFilters/exclude.xml index 1baf683dbca..051bd4fc1bc 100644 --- a/gradle/spotbugFilters/exclude.xml +++ b/gradle/spotbugFilters/exclude.xml @@ -5,6 +5,7 @@ + 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 99412ab7143..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 @@ -221,8 +221,8 @@ public EventType> databaseConnection() { new ET<>("database.query", DATABASE_SQL_QUERY_ID); /** A database sql query */ @SuppressWarnings("unchecked") - public EventType> databaseSqlQuery() { - return (EventType>) DATABASE_SQL_QUERY; + public EventType>> databaseSqlQuery() { + return (EventType>>) DATABASE_SQL_QUERY; } static final int GRPC_SERVER_METHOD_ID = 18; 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 970583460e7..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 @@ -153,7 +153,7 @@ public void cancel() { } /** Ensure that callbacks don't leak exceptions */ - @SuppressWarnings({"unchecked", "rawtypes", "DuplicateBranchesInSwitch"}) + @SuppressWarnings({"unchecked", "DuplicateBranchesInSwitch"}) public static C wrap(final EventType eventType, final C callback) { switch (eventType.getId()) { case REQUEST_STARTED_ID: @@ -366,7 +366,6 @@ public Flow apply(RequestContext ctx, Integer status) { } }; case DATABASE_CONNECTION_ID: - case DATABASE_SQL_QUERY_ID: return (C) new BiConsumer() { @Override @@ -378,6 +377,20 @@ public void accept(RequestContext ctx, String arg) { } } }; + case DATABASE_SQL_QUERY_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(); + } + } + }; default: log.warn("Unwrapped callback for {}", eventType); return callback; 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 69bc77a4887..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 @@ -203,7 +203,7 @@ public void testNormalCalls() { ss.registerCallback(events.databaseConnection(), callback); cbp.getCallback(events.databaseConnection()).accept(null, null); ss.registerCallback(events.databaseSqlQuery(), callback); - cbp.getCallback(events.databaseSqlQuery()).accept(null, null); + cbp.getCallback(events.databaseSqlQuery()).apply(null, null); assertThat(callback.count).isEqualTo(Events.MAX_EVENTS); } @@ -259,7 +259,7 @@ public void testThrowableBlocking() { ss.registerCallback(events.databaseConnection(), throwback); cbp.getCallback(events.databaseConnection()).accept(null, null); ss.registerCallback(events.databaseSqlQuery(), throwback); - cbp.getCallback(events.databaseSqlQuery()).accept(null, null); + cbp.getCallback(events.databaseSqlQuery()).apply(null, null); assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS); }