Assert specific return values in SQL injection tests

Instead of assertNotEquals(1, result), assert the exact expected
return value: 0 for safe queries, 1 for injections, 3 for tokenize
errors. Unterminated string cases now use isSqlTokenizeError.

Author: hansott

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

@@ -7,32 +7,39 @@
 import static dev.aikido.agent_api.vulnerabilities.sql_injection.SqlDetector.detectSqlInjection;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class SqlInjectionTest {
     private void isNotSqlInjection(String sql, String input, String dialect) {
-        int result;
         if ("mysql".equals(dialect) || "all".equals(dialect)) {
-            result = detectSqlInjection(sql, input, new Dialect("mysql"));
-            assertNotEquals(1, result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
+            int result = detectSqlInjection(sql, input, new Dialect("mysql"));
+            assertEquals(0, 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"));
-            assertNotEquals(1, result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
+            int result = detectSqlInjection(sql, input, new Dialect("postgresql"));
+            assertEquals(0, result, String.format("Expected no SQL injection for SQL: %s and input: %s", sql, input));
         }
     }
     private void isSqlInjection(String sql, String input, String dialect) {
-        int result;
         if ("mysql".equals(dialect) || "all".equals(dialect)) {
-            result = detectSqlInjection(sql, input, new Dialect("mysql"));
+            int result = detectSqlInjection(sql, input, new Dialect("mysql"));
             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"));
+            int result = detectSqlInjection(sql, input, new Dialect("postgresql"));
             assertEquals(1, result, String.format("Expected SQL injection for SQL: %s and input: %s", sql, input));
         }
     }
+    private void isSqlTokenizeError(String sql, String input, String dialect) {
+        if ("mysql".equals(dialect) || "all".equals(dialect)) {
+            int result = detectSqlInjection(sql, input, new Dialect("mysql"));
+            assertEquals(3, result, String.format("Expected SQL tokenize error for SQL: %s and input: %s", sql, input));
+        }
+        if ("postgresql".equals(dialect) || "all".equals(dialect)) {
+            int result = detectSqlInjection(sql, input, new Dialect("postgresql"));
+            assertEquals(3, result, String.format("Expected SQL tokenize error for SQL: %s and input: %s", sql, input));
+        }
+    }
 
 
     /**
@@ -166,17 +173,8 @@ public void testShouldReturnEarly() {
         assertFalse(SqlDetector.shouldReturnEarly("SELECT * FROM users; DROP TABLE", "users; DROP TABLE"));
     }
 
-    /**
-     * Moved :
-     * is_sql_injection("SELECT * FROM users WHERE id = 'users\\'", "users\\")
-     * is_sql_injection("SELECT * FROM users WHERE id = 'users\\\\'", "users\\\\")
-     * to is_not_sql_injection. Reason : Invalid SQL.
-     */
     @Test
     public void testAllowEscapeSequences() {
-        // Invalid queries:
-        isNotSqlInjection("SELECT * FROM users WHERE id = 'users\\'", "users\\", "all");
-        isNotSqlInjection("SELECT * FROM users WHERE id = 'users\\\\'", "users\\\\", "all");
         isNotSqlInjection("SELECT * FROM users WHERE id = '\nusers'", "\nusers", "all");
         isNotSqlInjection("SELECT * FROM users WHERE id = '\rusers'", "\rusers", "all");
         isNotSqlInjection("SELECT * FROM users WHERE id = '\tusers'", "\tusers", "all");
@@ -220,8 +218,7 @@ public void testCheckStringSafelyEscaped() {
                 "SELECT * FROM comments WHERE comment = \"I\\`m writing you\"", "I`m writing you", "all"
         );
 
-        // Invalid query (strings don't terminate)
-        isNotSqlInjection(
+        isSqlTokenizeError(
                 "SELECT * FROM comments WHERE comment = 'I'm writing you'", "I'm writing you", "all"
         );
 
@@ -375,6 +372,12 @@ public void testLowercasedInputSqlInjection() {
         isSqlInjection(sql, expectedSqlInjection, "all");
     }
 
+    @Test
+    public void testUnterminatedStrings() {
+        isSqlTokenizeError("SELECT * FROM users WHERE id = 'users\\'", "users\\", "all");
+        isSqlTokenizeError("SELECT * FROM users WHERE id = 'users\\\\'", "users\\\\", "all");
+    }
+
     /**
      * Marked the following as SQL injection since this would result in 2 or more tokens becoming one :
      * is_not_sql_injection("foobar)", "foobar)")