Skip to content

Commit

Permalink
Introduced SQL-injection blocking
Browse files Browse the repository at this point in the history
  • Loading branch information
@ValentinZakharov
ValentinZakharov committed Jun 21, 2024
1 parent 901334b commit c2b0115
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 23 deletions.
39 changes: 34 additions & 5 deletions .../main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -14,7 +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<CONNECTION> extends ClientDecorator {
protected static class NamingEntry {
Expand Down Expand Up @@ -116,14 +119,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<RequestContext, String> sqlQueryCallback =
BiFunction<RequestContext, String, Flow<Void>> sqlQueryCallback =
AgentTracer.get()
.getCallbackProvider(RequestContextSlot.APPSEC)
.getCallback(EVENTS.databaseSqlQuery());
if (sqlQueryCallback != null) {
RequestContext ctx = span.getRequestContext();
if (ctx != null) {
sqlQueryCallback.accept(ctx, sql);
Flow<Void> 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)");
}
}
}
}
Expand All @@ -135,14 +151,27 @@ protected void processDatabaseType(AgentSpan span, String dbType) {
postProcessServiceAndOperationName(span, namingEntry);

if (Config.get().isAppSecRaspEnabled() && dbType != null) {
BiConsumer<RequestContext, String> connectDbCallback =
BiFunction<RequestContext, String, Flow<Void>> connectDbCallback =
AgentTracer.get()
.getCallbackProvider(RequestContextSlot.APPSEC)
.getCallback(EVENTS.databaseConnection());
if (connectDbCallback != null) {
RequestContext ctx = span.getRequestContext();
if (ctx != null) {
connectDbCallback.accept(ctx, dbType);
Flow<Void> 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)");
}
}
}
}
Expand Down
14 changes: 6 additions & 8 deletions dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -459,7 +459,7 @@ public void init() {
(ctx_, dbType) -> {
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
if (ctx == null) {
return;
return NoopFlow.INSTANCE;
}
while (true) {
DataSubscriberInfo subInfo = dbConnectionSubInfo;
Expand All @@ -468,12 +468,11 @@ public void init() {
dbConnectionSubInfo = subInfo;
}
if (subInfo == null || subInfo.isEmpty()) {
return;
return NoopFlow.INSTANCE;
}
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_TYPE, dbType);
try {
producerService.publishDataEvent(subInfo, ctx, bundle, false);
return;
return producerService.publishDataEvent(subInfo, ctx, bundle, false);
} catch (ExpiredSubscriberInfoException e) {
dbConnectionSubInfo = null;
}
Expand All @@ -485,7 +484,7 @@ public void init() {
(ctx_, sql) -> {
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
if (ctx == null) {
return;
return NoopFlow.INSTANCE;
}
while (true) {
DataSubscriberInfo subInfo = dbSqlQuerySubInfo;
Expand All @@ -494,12 +493,11 @@ public void init() {
dbSqlQuerySubInfo = subInfo;
}
if (subInfo == null || subInfo.isEmpty()) {
return;
return NoopFlow.INSTANCE;
}
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_SQL_QUERY, sql);
try {
producerService.publishDataEvent(subInfo, ctx, bundle, false);
return;
return producerService.publishDataEvent(subInfo, ctx, bundle, false);
} catch (ExpiredSubscriberInfoException e) {
dbSqlQuerySubInfo = null;
}
Expand Down
9 changes: 4 additions & 5 deletions internal-api/src/main/java/datadog/trace/api/gateway/Events.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
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;
Expand Down Expand Up @@ -210,8 +209,8 @@ public EventType<BiFunction<RequestContext, Object, Flow<Void>>> grpcServerReque
new ET<>("database.connection", DATABASE_CONNECTION_ID);
/** A database connection */
@SuppressWarnings("unchecked")
public EventType<BiConsumer<RequestContext, String>> databaseConnection() {
return (EventType<BiConsumer<RequestContext, String>>) DATABASE_CONNECTION;
public EventType<BiFunction<RequestContext, String, Flow<Void>>> databaseConnection() {
return (EventType<BiFunction<RequestContext, String, Flow<Void>>>) DATABASE_CONNECTION;
}

static final int DATABASE_SQL_QUERY_ID = 17;
Expand All @@ -221,8 +220,8 @@ public EventType<BiConsumer<RequestContext, String>> databaseConnection() {
new ET<>("database.query", DATABASE_SQL_QUERY_ID);
/** A database sql query */
@SuppressWarnings("unchecked")
public EventType<BiConsumer<RequestContext, String>> databaseSqlQuery() {
return (EventType<BiConsumer<RequestContext, String>>) DATABASE_SQL_QUERY;
public EventType<BiFunction<RequestContext, String, Flow<Void>>> databaseSqlQuery() {
return (EventType<BiFunction<RequestContext, String, Flow<Void>>>) DATABASE_SQL_QUERY;
}

static final int GRPC_SERVER_METHOD_ID = 18;
Expand Down
10 changes: 5 additions & 5 deletions internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
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;
Expand Down Expand Up @@ -153,7 +152,7 @@ public void cancel() {
}

/** Ensure that callbacks don't leak exceptions */
@SuppressWarnings({"unchecked", "rawtypes", "DuplicateBranchesInSwitch"})
@SuppressWarnings({"unchecked", "DuplicateBranchesInSwitch"})
public static <C> C wrap(final EventType<C> eventType, final C callback) {
switch (eventType.getId()) {
case REQUEST_STARTED_ID:
Expand Down Expand Up @@ -368,13 +367,14 @@ public Flow<Void> apply(RequestContext ctx, Integer status) {
case DATABASE_CONNECTION_ID:
case DATABASE_SQL_QUERY_ID:
return (C)
new BiConsumer<RequestContext, String>() {
new BiFunction<RequestContext, String, Flow<Void>>() {
@Override
public void accept(RequestContext ctx, String arg) {
public Flow<Void> apply(RequestContext ctx, String arg) {
try {
((BiConsumer<RequestContext, String>) callback).accept(ctx, arg);
return ((BiFunction<RequestContext, String, Flow<Void>>) callback).apply(ctx, arg);
} catch (Throwable t) {
log.warn("Callback for {} threw.", eventType, t);
return Flow.ResultFlow.empty();
}
}
};
Expand Down

0 comments on commit c2b0115

Please sign in to comment.