Block SQL queries that fail tokenization

Our SQL injection detection tokenizes queries to check them. If the
tokenizer can't parse a query, we skip it and the query goes through.
Some databases still execute partially valid queries though: ClickHouse
ignores junk after ; and SQLite runs everything before an unclosed /*.

Now when user input shows up in a query we can't tokenize, we treat
it as an attack. On by default, opt out with AIKIDO_BLOCK_INVALID_SQL=false.

Author: hansott

agent_api/src/main/java/dev/aikido/agent_api/helpers/env/BlockInvalidSqlEnv.java

@@ -0,0 +1,10 @@
+package dev.aikido.agent_api.helpers.env;
+
+public class BlockInvalidSqlEnv extends BooleanEnv {
+    private static final String environmentName = "AIKIDO_BLOCK_INVALID_SQL";
+    private static final boolean defaultValue = true;
+
+    public BlockInvalidSqlEnv() {
+        super(environmentName, defaultValue);
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/RustSQLInterface.java

@@ -24,20 +24,19 @@ public interface SqlLib {
         int detect_sql_injection(String query, long queryLen, String userinput, long userinputLen, int dialect);
     }
 
-    public static boolean detectSqlInjection(String query, String userInput, Dialect dialect) {
+    public static int detectSqlInjection(String query, String userInput, Dialect dialect) {
         int dialectInteger = dialect.getDialectInteger();
         try {
             SqlLib lib = loadLibrary();
             if (lib != null) {
                 long queryLen = query != null ? query.getBytes(StandardCharsets.UTF_8).length : 0;
                 long userInputLen = userInput != null ? userInput.getBytes(StandardCharsets.UTF_8).length : 0;
-                int result = lib.detect_sql_injection(query, queryLen, userInput, userInputLen, dialectInteger);
-                return result == 1;
+                return lib.detect_sql_injection(query, queryLen, userInput, userInputLen, dialectInteger);
             }
         } catch (Throwable e) {
             logger.trace(e);
         }
-        return false;
+        return 0;
     }
     public static SqlLib loadLibrary() {
         String path = getPathForBinary();

agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java

@@ -1,9 +1,11 @@
 package dev.aikido.agent_api.vulnerabilities.sql_injection;
 
+import dev.aikido.agent_api.helpers.env.BlockInvalidSqlEnv;
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
 import dev.aikido.agent_api.vulnerabilities.Detector;
 
+import java.util.HashMap;
 import java.util.Map;
 import java.util.regex.Pattern;
 
@@ -23,21 +25,32 @@ public DetectorResult run(String userInput, String[] arguments) {
         }
         String query = arguments[0];
         Dialect dialect = new Dialect(arguments[1]);
-        boolean detectedAttack = detectSqlInjection(query, userInput, dialect);
-        if (detectedAttack) {
+        int result = detectSqlInjection(query, userInput, dialect);
+
+        if (result == 1) {
             Map<String, String> metadata = Map.of(
                 "sql", query,
                 "dialect", dialect.getHumanName()
             );
-            return new DetectorResult(/* detectedAttack*/ true, metadata, SQLInjectionException.get(dialect));
+            return new DetectorResult(true, metadata, SQLInjectionException.get(dialect));
+        }
+
+        if (result == 3 && new BlockInvalidSqlEnv().getValue()) {
+            Map<String, String> metadata = new HashMap<>();
+            metadata.put("sql", query);
+            metadata.put("dialect", dialect.getHumanName());
+            metadata.put("failedToTokenize", "true");
+            return new DetectorResult(true, metadata, SQLInjectionException.get(dialect));
         }
+
         return new DetectorResult();
     }
-    public static boolean detectSqlInjection(String query, String userInput, Dialect dialect) {
+
+    public static int detectSqlInjection(String query, String userInput, Dialect dialect) {
         String queryLower = query.toLowerCase();
         String userInputLower = userInput.toLowerCase();
         if (shouldReturnEarly(queryLower, userInputLower)) {
-            return false;
+            return 0;
         }
         return RustSQLInterface.detectSqlInjection(queryLower, userInputLower, dialect);
     }

agent_api/src/test/java/vulnerabilities/RustSQLInterfaceTest.java

@@ -4,15 +4,14 @@
 import dev.aikido.agent_api.vulnerabilities.sql_injection.RustSQLInterface;
 import org.junit.jupiter.api.Test;
 
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class RustSQLInterfaceTest {
     @Test
     public void testItWorks() {
-        boolean injectionResult = RustSQLInterface.detectSqlInjection("SELECT * FROM table;", "table;", new Dialect("postgresql"));
-        assertTrue(injectionResult);
+        int injectionResult = RustSQLInterface.detectSqlInjection("SELECT * FROM table;", "table;", new Dialect("postgresql"));
+        assertEquals(1, injectionResult);
         injectionResult = RustSQLInterface.detectSqlInjection("SELECT * FROM table;", "table", new Dialect("postgresql"));
-        assertFalse(injectionResult);
+        assertEquals(0, injectionResult);
     }
 }

agent_api/src/test/java/vulnerabilities/SqlInjectionTest.java

@@ -5,30 +5,30 @@
 import org.junit.jupiter.api.Test;
 
 import static dev.aikido.agent_api.vulnerabilities.sql_injection.SqlDetector.detectSqlInjection;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 
 public class SqlInjectionTest {
     private void isNotSqlInjection(String sql, String input, String dialect) {
-        boolean result;
+        int result;
         if ("mysql".equals(dialect) || "all".equals(dialect)) {
             result = detectSqlInjection(sql, input, new Dialect("mysql"));
-            assertFalse(result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
+            assertNotEquals(1, result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
         }
         if ("postgresql".equals(dialect) || "all".equals(dialect)) {
             result = detectSqlInjection(sql, input, new Dialect("postgresql"));
-            assertFalse(result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
+            assertNotEquals(1, result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
         }
     }
     private void isSqlInjection(String sql, String input, String dialect) {
-        boolean result;
+        int result;
         if ("mysql".equals(dialect) || "all".equals(dialect)) {
             result = detectSqlInjection(sql, input, new Dialect("mysql"));
-            assertTrue(result, String.format("Expected SQL injection for SQL: %s and input: %s", sql, input));
+            assertEquals(1, result, String.format("Expected SQL injection for SQL: %s and input: %s", sql, input));
         }
         if ("postgresql".equals(dialect) || "all".equals(dialect)) {
             result = detectSqlInjection(sql, input, new Dialect("postgresql"));
-            assertTrue(result, String.format("Expected SQL injection for SQL: %s and input: %s", sql, input));
+            assertEquals(1, result, String.format("Expected SQL injection for SQL: %s and input: %s", sql, input));
         }
     }
 

docs/invalid-sql-queries.md

@@ -0,0 +1,8 @@
+# Blocking invalid SQL queries
+
+Zen blocks SQL queries that it can't tokenize when they contain user input. This prevents attackers from bypassing SQL injection detection with malformed queries. For example, ClickHouse ignores invalid SQL after `;`, and SQLite runs queries before an unclosed `/*` comment.
+
+This is on by default. In blocking mode, these queries are blocked. In detection-only mode, they are reported but still executed.
+
+If you see false positives (legitimate queries being blocked), set the
+`AIKIDO_BLOCK_INVALID_SQL` environment variable to `false`.