Added support for SQLi exploit prevention (#7051)

* Added support for SQLi exploit prevention * Fixed collecting raw sql queries * Fixed tests * Spotless * Fixed test * Capabilities for SQLi exploit prevention * Added config option to disable RASP * Db type and Sql query collection moved to DatabaseClientDecorator * Fixed tests * Fixed database type callback * Refactoring * Refactoring * Added comments * Refactoring
DataDog · Jun 3, 2024 · 148ac0f · 148ac0f
1 parent f800fc6
commit 148ac0f
Show file tree

Hide file tree

Showing 16 changed files with 182 additions and 10 deletions.
diff --git a/.../main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java b/.../main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java
@@ -1,15 +1,20 @@
 package datadog.trace.bootstrap.instrumentation.decorator;
 
+import static datadog.trace.api.gateway.Events.EVENTS;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.naming.NamingSchema;
 import datadog.trace.api.naming.SpanNaming;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+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;
 
 public abstract class DatabaseClientDecorator<CONNECTION> extends ClientDecorator {
   protected static class NamingEntry {
@@ -105,10 +110,42 @@ public AgentSpan onStatement(final AgentSpan span, final CharSequence statement)
     return span;
   }
 
+  /**
+   * The method used to provide raw sql to prevent SQL-injection attacks SQL query should never be
+   * exposed because it may contain sensitive data.
+   */
+  public void onRawStatement(AgentSpan span, String sql) {
+    if (Config.get().getAppSecRaspEnabled() && sql != null && !sql.isEmpty()) {
+      BiConsumer<RequestContext, String> sqlQueryCallback =
+          AgentTracer.get()
+              .getCallbackProvider(RequestContextSlot.APPSEC)
+              .getCallback(EVENTS.databaseSqlQuery());
+      if (sqlQueryCallback != null) {
+        RequestContext ctx = span.getRequestContext();
+        if (ctx != null) {
+          sqlQueryCallback.accept(ctx, sql);
+        }
+      }
+    }
+  }
+
   protected void processDatabaseType(AgentSpan span, String dbType) {
     final NamingEntry namingEntry = CACHE.computeIfAbsent(dbType, NamingEntry::new);
     span.setTag(DB_TYPE, namingEntry.dbType);
     postProcessServiceAndOperationName(span, namingEntry);
+
+    if (Config.get().getAppSecRaspEnabled() && dbType != null) {
+      BiConsumer<RequestContext, String> connectDbCallback =
+          AgentTracer.get()
+              .getCallbackProvider(RequestContextSlot.APPSEC)
+              .getCallback(EVENTS.databaseConnection());
+      if (connectDbCallback != null) {
+        RequestContext ctx = span.getRequestContext();
+        if (ctx != null) {
+          connectDbCallback.accept(ctx, dbType);
+        }
+      }
+    }
   }
 
   protected void postProcessServiceAndOperationName(AgentSpan span, NamingEntry namingEntry) {}

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java
@@ -8,6 +8,7 @@
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_DD_RULES;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_EXCLUSIONS;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_IP_BLOCKING;
+import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_RASP_SQLI;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_REQUEST_BLOCKING;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_TRUSTED_IPS;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_USER_BLOCKING;
@@ -96,7 +97,8 @@ private void subscribeConfigurationPoller() {
             | CAPABILITY_ASM_USER_BLOCKING
             | CAPABILITY_ASM_CUSTOM_RULES
             | CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE
-            | CAPABILITY_ASM_TRUSTED_IPS);
+            | CAPABILITY_ASM_TRUSTED_IPS
+            | CAPABILITY_ASM_RASP_SQLI);
   }
 
   private void subscribeRulesAndData() {
@@ -334,7 +336,8 @@ public void close() {
             | CAPABILITY_ASM_CUSTOM_RULES
             | CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE
             | CAPABILITY_ASM_TRUSTED_IPS
-            | CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE);
+            | CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE
+            | CAPABILITY_ASM_RASP_SQLI);
     this.configurationPoller.removeListeners(Product.ASM_DD);
     this.configurationPoller.removeListeners(Product.ASM_DATA);
     this.configurationPoller.removeListeners(Product.ASM);

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
@@ -109,6 +109,12 @@ public interface KnownAddresses {
 
   Address<String> USER_ID = new Address<>("usr.id");
 
+  /** The database type (ex: mysql, postgresql, sqlite) */
+  Address<String> DB_TYPE = new Address<>("server.db.system");
+
+  /** The SQL query being executed */
+  Address<String> DB_SQL_QUERY = new Address<>("server.db.statement");
+
   Address<Map<String, Object>> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor");
 
   static Address<?> forName(String name) {
@@ -165,6 +171,10 @@ static Address<?> forName(String name) {
         return SERVER_GRAPHQL_ALL_RESOLVERS;
       case "usr.id":
         return USER_ID;
+      case "server.db.system":
+        return DB_TYPE;
+      case "server.db.statement":
+        return DB_SQL_QUERY;
       case "waf.context.processor":
         return WAF_CONTEXT_PROCESSOR;
       default:

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
@@ -82,6 +82,8 @@ 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(
       SubscriptionService subscriptionService,
@@ -413,6 +415,58 @@ public void init() {
             }
           }
         });
+
+    subscriptionService.registerCallback(
+        EVENTS.databaseConnection(),
+        (ctx_, dbType) -> {
+          AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
+          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;
+            }
+          }
+        });
+
+    subscriptionService.registerCallback(
+        EVENTS.databaseSqlQuery(),
+        (ctx_, sql) -> {
+          AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
+          if (ctx == null) {
+            return;
+          }
+          while (true) {
+            DataSubscriberInfo subInfo = dbSqlQuerySubInfo;
+            if (subInfo == null) {
+              subInfo = producerService.getDataSubscribers(KnownAddresses.DB_SQL_QUERY);
+              dbSqlQuerySubInfo = subInfo;
+            }
+            if (subInfo == null || subInfo.isEmpty()) {
+              return;
+            }
+            DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.DB_SQL_QUERY, sql);
+            try {
+              producerService.publishDataEvent(subInfo, ctx, bundle, false);
+              return;
+            } catch (ExpiredSubscriberInfoException e) {
+              dbSqlQuerySubInfo = null;
+            }
+          }
+        });
   }
 
   public void stop() {

diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java
@@ -368,6 +368,8 @@ private static Collection<Address<?>> getUsedAddresses(PowerwafContext ctx) {
     addressList.add(KnownAddresses.RESPONSE_HEADERS_NO_COOKIES);
     addressList.add(KnownAddresses.RESPONSE_BODY_OBJECT);
     addressList.add(KnownAddresses.GRAPHQL_SERVER_ALL_RESOLVERS);
+    addressList.add(KnownAddresses.DB_TYPE);
+    addressList.add(KnownAddresses.DB_SQL_QUERY);
 
     return addressList;
   }

diff --git a/...sec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy b/...sec/src/test/groovy/com/datadog/appsec/config/AppSecConfigServiceImplSpecification.groovy
@@ -229,7 +229,7 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     }
     1 * poller.addConfigurationEndListener(_) >> { listeners.savedConfEndListener = it[0] }
     1 * poller.addCapabilities(2L)
-    1 * poller.addCapabilities(1980L)
+    1 * poller.addCapabilities(2099132L)
     0 * _._
     initialWafConfig.get() != null
 
@@ -366,7 +366,7 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     }
     1 * poller.addConfigurationEndListener(_) >> { listeners.savedConfEndListener = it[0] }
     1 * poller.addCapabilities(2L)
-    1 * poller.addCapabilities(1980L)
+    1 * poller.addCapabilities(2099132L)
     0 * _._
 
     when:
@@ -422,7 +422,7 @@ class AppSecConfigServiceImplSpecification extends DDSpecification {
     poller = null
 
     then:
-    1 * poller.removeCapabilities(4030L)
+    1 * poller.removeCapabilities(2101182L)
     4 * poller.removeListeners(_)
     1 * poller.removeConfigurationEndListener(_)
     1 * poller.stop()

diff --git a/...t/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy b/...t/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy
@@ -32,14 +32,16 @@ class KnownAddressesSpecification extends Specification {
       'grpc.server.request.metadata',
       'graphql.server.all_resolvers',
       'graphql.server.resolver',
+      'server.db.system',
+      'server.db.statement',
       'usr.id',
       'waf.context.processor',
     ]
   }
 
   void 'number of known addresses is expected number'() {
     expect:
-    Address.instanceCount() == 27
+    Address.instanceCount() == 29
     KnownAddresses.WAF_CONTEXT_PROCESSOR.serial == Address.instanceCount() - 1
   }
 }
diff --git a/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/...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
@@ -78,6 +79,8 @@ class GatewayBridgeSpecification extends DDSpecification {
   Function<RequestContext, Flow<Void>> respHeadersDoneCB
   BiFunction<RequestContext, Object, Flow<Void>> grpcServerRequestMessageCB
   BiFunction<RequestContext, Map<String, Object>, Flow<Void>> graphqlServerRequestMessageCB
+  BiConsumer<RequestContext, String> databaseConnectionCB
+  BiConsumer<RequestContext, String> databaseSqlQueryCB
 
   void setup() {
     callInitAndCaptureCBs()
@@ -412,6 +415,8 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.grpcServerRequestMessage(), _) >> { grpcServerRequestMessageCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.graphqlServerRequestMessage(), _) >> { graphqlServerRequestMessageCB = it[1]; null }
+    1 * ig.registerCallback(EVENTS.databaseConnection(), _) >> { databaseConnectionCB = it[1]; null }
+    1 * ig.registerCallback(EVENTS.databaseSqlQuery(), _) >> { databaseSqlQueryCB = it[1]; null }
     0 * ig.registerCallback(_, _)
 
     bridge.init()

diff --git a/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java b/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
@@ -205,7 +205,9 @@ public static DBInfo parseDBInfoFromConnection(final Connection connection) {
     return dbInfo;
   }
 
-  public AgentSpan onStatement(AgentSpan span, DBQueryInfo dbQueryInfo) {
+  public AgentSpan onStatement(AgentSpan span, final String statement) {
+    onRawStatement(span, statement);
+    DBQueryInfo dbQueryInfo = DBQueryInfo.ofStatement(statement);
     return withQueryInfo(span, dbQueryInfo, JDBC_STATEMENT);
   }
 

diff --git a/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
@@ -21,7 +21,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
-import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -112,7 +111,7 @@ public static AgentScope onEnter(
                   injectTraceContext,
                   appendComment);
         }
-        DECORATE.onStatement(span, DBQueryInfo.ofStatement(copy));
+        DECORATE.onStatement(span, copy);
         return activateSpan(span);
       } catch (SQLException e) {
         // if we can't get the connection for any reason

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -101,7 +101,7 @@ public final class ConfigDefaults {
   static final int DEFAULT_APPSEC_WAF_TIMEOUT = 100000; // 0.1 s
   static final boolean DEFAULT_API_SECURITY_ENABLED = false;
   static final float DEFAULT_API_SECURITY_REQUEST_SAMPLE_RATE = 0.1f; // 10 %
-
+  static final boolean DEFAULT_APPSEC_RASP_ENABLED = false;
   static final String DEFAULT_IAST_ENABLED = "false";
   static final boolean DEFAULT_IAST_DEBUG_ENABLED = false;
   public static final int DEFAULT_IAST_MAX_CONCURRENT_REQUESTS = 4;

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/AppSecConfig.java
@@ -27,6 +27,7 @@ public final class AppSecConfig {
   public static final String API_SECURITY_REQUEST_SAMPLE_RATE = "api-security.request.sample.rate";
 
   public static final String APPSEC_SCA_ENABLED = "appsec.sca.enabled";
+  public static final String APPSEC_RASP_ENABLED = "appsec.rasp.enabled";
 
   private AppSecConfig() {}
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -6,6 +6,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_ANALYTICS_SAMPLE_RATE;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_API_SECURITY_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_API_SECURITY_REQUEST_SAMPLE_RATE;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_APPSEC_RASP_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_APPSEC_REPORTING_INBAND;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_APPSEC_TRACE_RATE_LIMIT;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_APPSEC_WAF_METRICS;
@@ -147,6 +148,7 @@
 import static datadog.trace.api.config.AppSecConfig.APPSEC_IP_ADDR_HEADER;
 import static datadog.trace.api.config.AppSecConfig.APPSEC_OBFUSCATION_PARAMETER_KEY_REGEXP;
 import static datadog.trace.api.config.AppSecConfig.APPSEC_OBFUSCATION_PARAMETER_VALUE_REGEXP;
+import static datadog.trace.api.config.AppSecConfig.APPSEC_RASP_ENABLED;
 import static datadog.trace.api.config.AppSecConfig.APPSEC_REPORTING_INBAND;
 import static datadog.trace.api.config.AppSecConfig.APPSEC_REPORT_TIMEOUT_SEC;
 import static datadog.trace.api.config.AppSecConfig.APPSEC_RULES_FILE;
@@ -736,6 +738,7 @@ static class HostNameHolder {
   private final String appSecHttpBlockedTemplateJson;
   private final UserEventTrackingMode appSecUserEventsTracking;
   private final Boolean appSecScaEnabled;
+  private final Boolean appSecRaspEnabled;
   private final boolean apiSecurityEnabled;
   private final float apiSecurityRequestSampleRate;
 
@@ -1633,6 +1636,7 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
             configProvider.getStringNotEmpty(
                 APPSEC_AUTOMATED_USER_EVENTS_TRACKING, SAFE.toString()));
     appSecScaEnabled = configProvider.getBoolean(APPSEC_SCA_ENABLED);
+    appSecRaspEnabled = configProvider.getBoolean(APPSEC_RASP_ENABLED, DEFAULT_APPSEC_RASP_ENABLED);
     apiSecurityEnabled =
         configProvider.getBoolean(
             API_SECURITY_ENABLED, DEFAULT_API_SECURITY_ENABLED, API_SECURITY_ENABLED_EXPERIMENTAL);
@@ -3976,6 +3980,10 @@ public Boolean getAppSecScaEnabled() {
     return appSecScaEnabled;
   }
 
+  public Boolean getAppSecRaspEnabled() {
+    return appSecRaspEnabled;
+  }
+
   private <T> Set<T> getSettingsSetFromEnvironment(
       String name, Function<String, T> mapper, boolean splitOnWS) {
     final String value = configProvider.getString(name, "");
@@ -4612,6 +4620,8 @@ public String toString() {
         + telemetryMetricsEnabled
         + ", appSecScaEnabled="
         + appSecScaEnabled
+        + ", appSecRaspEnabled="
+        + appSecRaspEnabled
         + ", dataJobsEnabled="
         + dataJobsEnabled
         + '}';

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
@@ -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;
@@ -202,6 +203,28 @@ public EventType<BiFunction<RequestContext, Object, Flow<Void>>> grpcServerReque
         GRAPHQL_SERVER_REQUEST_MESSAGE;
   }
 
+  static final int DATABASE_CONNECTION_ID = 16;
+
+  @SuppressWarnings("rawtypes")
+  private static final EventType DATABASE_CONNECTION =
+      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;
+  }
+
+  static final int DATABASE_SQL_QUERY_ID = 17;
+
+  @SuppressWarnings("rawtypes")
+  private static final EventType DATABASE_SQL_QUERY =
+      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;
+  }
+
   static final int MAX_EVENTS = nextId.get();
 
   private static final class ET<T> extends EventType<T> {