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 <santiago.mola@datadoghq.com>

* Fixed typo in field name

* Added RASP info in StatusLogger

---------

Co-authored-by: Santiago Mola <santiago.mola@datadoghq.com>

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<CONNECTION> 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<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)");
+          }
         }
       }
     }

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<String, String> 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;
   }

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,45 +454,34 @@ 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(
         EVENTS.databaseSqlQuery(),
         (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;
             }

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<String> highlight;

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -79,7 +79,7 @@ class GatewayBridgeSpecification extends DDSpecification {
   BiFunction<RequestContext, Object, Flow<Void>> grpcServerRequestMessageCB
   BiFunction<RequestContext, Map<String, Object>, Flow<Void>> graphqlServerRequestMessageCB
   BiConsumer<RequestContext, String> databaseConnectionCB
-  BiConsumer<RequestContext, String> databaseSqlQueryCB
+  BiFunction<RequestContext, String, Flow<Void>> databaseSqlQueryCB
 
   void setup() {
     callInitAndCaptureCBs()

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);
+  }
+}

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<String, String> 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;
       }
     }
 

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')
 }

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";
+  }
 }

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'
+  }
 }