From 75c25b976fd1af930860b363ab90cf762b2cd7ec Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Mon, 31 Aug 2026 14:05:01 +0200 Subject: [PATCH 1/2] Prevent multiple attack events for SQL injection (detection-only mode) --- agent/build.gradle | 1 + .../wrappers/jdbc/JDBCConnectionAdvice.java | 37 +++++++++++- .../jdbc/JDBCConnectionAdviceTest.java | 59 +++++++++++++++++++ .../test/java/wrappers/MSSQLWrapperTest.java | 26 +++++++- .../java/wrappers/MariadbWrapperTest.java | 32 ++++++++++ .../java/wrappers/MysqlCJWrapperTest.java | 28 +++++++++ .../java/wrappers/PostgresWrapperTest.java | 54 +++++++++++++++++ 7 files changed, 233 insertions(+), 4 deletions(-) create mode 100644 agent/src/test/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdviceTest.java diff --git a/agent/build.gradle b/agent/build.gradle index 2e4e217cd..226cd16ba 100644 --- a/agent/build.gradle +++ b/agent/build.gradle @@ -14,6 +14,7 @@ dependencies { compileOnly 'org.springframework:spring-web:5.3.20' testImplementation 'org.junit.jupiter:junit-jupiter:5.9.2' + testImplementation 'org.postgresql:postgresql:42.2.23' testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' testRuntimeOnly 'org.junit.platform:junit-platform-launcher:1.9.2' } diff --git a/agent/src/main/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdvice.java b/agent/src/main/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdvice.java index b6b9d4cd7..261ff49c0 100644 --- a/agent/src/main/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdvice.java +++ b/agent/src/main/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdvice.java @@ -11,6 +11,8 @@ import java.lang.reflect.Executable; import java.sql.Connection; +import java.util.ArrayDeque; +import java.util.Deque; import java.sql.DatabaseMetaData; import java.sql.Statement; @@ -19,6 +21,17 @@ public final class JDBCConnectionAdvice { public static final Logger logger = LogManager.getLogger(JDBCConnectionAdvice.class); + + // Database drivers often delegate one prepareStatement overload to another. + // Skip only nested calls with the same SQL to avoid repeated WASM checks and duplicate detection-only events. + // + // Example for sql1 = "SELECT 1", where both overloads receive the same sql1 object: + // prepareStatement(sql1) -> prepareStatement(sql1, options) + // Enter outer: [] -> [sql1] (check) + // Enter inner: [sql1] -> [sql1, sql1] (skip) + // Exit inner: [sql1, sql1] -> [sql1] + // Exit outer: [sql1] -> [] + public static final ThreadLocal> jdbcCallStack = ThreadLocal.withInitial(ArrayDeque::new); private JDBCConnectionAdvice() {} public static ElementMatcher getMatcher(String module) { ElementMatcher.Junction statementMatcher = @@ -41,12 +54,18 @@ public static ElementMatcher getMatcher(String module * addBatch(sql), execute(sql, [...]), executeLargeUpdate(sql, [...]), executeQuery(sql), executeUpdate(sql, [...]) */ @Advice.OnMethodEnter - public static void before( + public static String before( @Advice.This(typing = DYNAMIC, optional = true) Object obj, @Advice.Origin Executable method, @Advice.Argument(0) String sql ) throws Throwable { - if (sql != null) { + if (sql == null) { + return null; + } + Deque sqlCalls = jdbcCallStack.get(); + boolean isDelegatedCall = !sqlCalls.isEmpty() && sqlCalls.peek() == sql; + sqlCalls.push(sql); + if (!isDelegatedCall) { try { // Get connection whether it's from a Statement or not: Connection databaseConnection = null; @@ -62,12 +81,24 @@ public static void before( String operation = "(" + metaData.getDriverName() + ") java.sql." + methodName; String dialect = metaData.getDatabaseProductName().toLowerCase(); SQLCollector.report(sql, dialect, operation); - } catch (AikidoException e) { + sqlCalls.pop(); throw e; } catch (Throwable e) { logger.debug(e); } } + return sql; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class) + public static void after(@Advice.Enter String enteredSql) { + if (enteredSql == null) { + return; + } + Deque sqlCalls = jdbcCallStack.get(); + if (!sqlCalls.isEmpty() && sqlCalls.peek() == enteredSql) { + sqlCalls.pop(); + } } } \ No newline at end of file diff --git a/agent/src/test/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdviceTest.java b/agent/src/test/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdviceTest.java new file mode 100644 index 000000000..571432f42 --- /dev/null +++ b/agent/src/test/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdviceTest.java @@ -0,0 +1,59 @@ +package dev.aikido.agent.wrappers.jdbc; + +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.storage.ServiceConfigStore; +import dev.aikido.agent_api.storage.statistics.StatisticsStore; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; +import java.sql.Connection; +import java.sql.DriverManager; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class JDBCConnectionAdviceTest { + private Connection connection; + private Method prepareStatement; + + @BeforeEach + void setUp() throws Exception { + connection = DriverManager.getConnection("jdbc:postgresql://localhost:5432/db", "user", "password"); + prepareStatement = connection.getClass().getMethod("prepareStatement", String.class); + Context.set(null); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(false); + } + + @AfterEach + void tearDown() throws Exception { + connection.close(); + Context.set(null); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); + JDBCConnectionAdvice.jdbcCallStack.remove(); + } + + @Test + void reportsNestedCallsWithDifferentSql() throws Throwable { + String outerSql = "SELECT 1"; + String nestedSql = "SELECT 2"; + + String outerCall = JDBCConnectionAdvice.before(connection, prepareStatement, outerSql); + try { + String nestedCall = JDBCConnectionAdvice.before(connection, prepareStatement, nestedSql); + JDBCConnectionAdvice.after(nestedCall); + } finally { + JDBCConnectionAdvice.after(outerCall); + } + + var operation = StatisticsStore.getStatsRecord().operations() + .get("(PostgreSQL JDBC Driver) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(2, operation.total()); + assertTrue(JDBCConnectionAdvice.jdbcCallStack.get().isEmpty()); + } +} diff --git a/agent_api/src/test/java/wrappers/MSSQLWrapperTest.java b/agent_api/src/test/java/wrappers/MSSQLWrapperTest.java index dc3c3466b..abfff5fc3 100644 --- a/agent_api/src/test/java/wrappers/MSSQLWrapperTest.java +++ b/agent_api/src/test/java/wrappers/MSSQLWrapperTest.java @@ -35,6 +35,7 @@ public void setUp() throws SQLException { String password = "Strong!Passw0rd"; // Change to your password connection = DriverManager.getConnection(url, user, password); StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @AfterEach @@ -44,6 +45,7 @@ public void tearDown() throws SQLException { } Context.set(null); StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @Test @@ -61,12 +63,34 @@ public void testSelectSqlWithPrepareStatement() throws SQLException { }); assertEquals("Aikido Zen has blocked SQL Injection, Dialect: Microsoft SQL", exception.getMessage()); var operation = StatisticsStore.getStatsRecord().operations().get("(Microsoft JDBC Driver 10.2 for SQL Server) java.sql.Connection.prepareStatement"); - assertEquals(5, operation.total()); + assertEquals(3, operation.total()); assertEquals(1, operation.getAttacksDetected().get("blocked")); assertEquals(1, operation.getAttacksDetected().get("total")); assertEquals(OperationKind.SQL_OP, operation.getKind()); } + @Test + public void testPrepareStatementReportsOnceInDetectionOnlyMode() throws SQLException { + String payload = "Malicious Pet', 'Gru from the Minions') -- "; + String sql = "INSERT INTO pets (pet_name, owner) VALUES ('" + payload + "', 'Aikido Security')"; + Context.set(new EmptySampleContextObject(payload)); + ServiceConfigStore.updateBlocking(false); + + assertDoesNotThrow(() -> connection.prepareStatement(sql)); + + var stats = StatisticsStore.getStatsRecord(); + assertEquals(1, stats.requests().attacksDetected().total()); + assertEquals(1, stats.operations().values().stream() + .filter(record -> record.getKind() == OperationKind.SQL_OP) + .count()); + var operation = stats.operations() + .get("(Microsoft JDBC Driver 10.2 for SQL Server) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(1, operation.total()); + assertEquals(1, operation.getAttacksDetected().get("total")); + assertEquals(0, operation.getAttacksDetected().get("blocked")); + } + @Test public void testSelectSqlSafeWithPrepareStatement() throws SQLException { Context.set(new EmptySampleContextObject("FROM")); diff --git a/agent_api/src/test/java/wrappers/MariadbWrapperTest.java b/agent_api/src/test/java/wrappers/MariadbWrapperTest.java index 5ece5f2b0..c7ea41ef5 100644 --- a/agent_api/src/test/java/wrappers/MariadbWrapperTest.java +++ b/agent_api/src/test/java/wrappers/MariadbWrapperTest.java @@ -2,6 +2,8 @@ import dev.aikido.agent_api.context.Context; import dev.aikido.agent_api.storage.ServiceConfigStore; +import dev.aikido.agent_api.storage.statistics.OperationKind; +import dev.aikido.agent_api.storage.statistics.StatisticsStore; import dev.aikido.agent_api.vulnerabilities.sql_injection.SQLInjectionException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -11,6 +13,7 @@ import java.sql.Connection; import java.sql.DriverManager; +import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; @@ -29,6 +32,8 @@ public static void clean() { public void setUp() throws SQLException { // Connect to the MySQL database connection = DriverManager.getConnection("jdbc:mariadb://localhost:3306/db?allowPublicKeyRetrieval=true&useSSL=false", "user", "password"); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @AfterEach @@ -37,6 +42,8 @@ public void tearDown() throws SQLException { connection.close(); } Context.set(null); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @Test @@ -55,6 +62,31 @@ public void testSelectSqlWithPrepareStatement() throws SQLException { assertEquals("Aikido Zen has blocked SQL Injection, Dialect: MySQL", exception.getMessage()); } + @Test + public void testPrepareStatementReportsOnceInDetectionOnlyMode() throws SQLException { + String payload = "Malicious Pet', 'Gru from the Minions') -- "; + String sql = "INSERT INTO pets (pet_name, owner) VALUES ('" + payload + "', 'Aikido Security')"; + Context.set(new EmptySampleContextObject(payload)); + ServiceConfigStore.updateBlocking(false); + + assertDoesNotThrow(() -> connection.prepareStatement(sql, + ResultSet.TYPE_FORWARD_ONLY, + ResultSet.CONCUR_READ_ONLY, + ResultSet.CLOSE_CURSORS_AT_COMMIT)); + + var stats = StatisticsStore.getStatsRecord(); + assertEquals(1, stats.requests().attacksDetected().total()); + assertEquals(1, stats.operations().values().stream() + .filter(record -> record.getKind() == OperationKind.SQL_OP) + .count()); + var operation = stats.operations() + .get("(MariaDB Connector/J) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(1, operation.total()); + assertEquals(1, operation.getAttacksDetected().get("total")); + assertEquals(0, operation.getAttacksDetected().get("blocked")); + } + @Test public void testSelectSqlSafeWithPrepareStatement() throws SQLException { Context.set(new EmptySampleContextObject("FROM")); diff --git a/agent_api/src/test/java/wrappers/MysqlCJWrapperTest.java b/agent_api/src/test/java/wrappers/MysqlCJWrapperTest.java index 2f0d7811b..273491903 100644 --- a/agent_api/src/test/java/wrappers/MysqlCJWrapperTest.java +++ b/agent_api/src/test/java/wrappers/MysqlCJWrapperTest.java @@ -2,6 +2,8 @@ import dev.aikido.agent_api.context.Context; import dev.aikido.agent_api.storage.ServiceConfigStore; +import dev.aikido.agent_api.storage.statistics.OperationKind; +import dev.aikido.agent_api.storage.statistics.StatisticsStore; import dev.aikido.agent_api.vulnerabilities.sql_injection.SQLInjectionException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -26,6 +28,8 @@ public static void clean() { public void setUp() throws SQLException { // Connect to the MySQL database connection = DriverManager.getConnection("jdbc:mysql://localhost:3306/db", "user", "password"); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @AfterEach @@ -33,7 +37,9 @@ public void tearDown() throws SQLException { if (connection != null) { connection.close(); } + StatisticsStore.clear(); Context.set(null); + ServiceConfigStore.updateBlocking(true); } @Test @@ -82,6 +88,28 @@ public void testSelectSqlWithPreparedStatementWithoutExecute() throws SQLExcepti assertEquals("Aikido Zen has blocked SQL Injection, Dialect: MySQL", exception.getMessage()); } + @Test + public void testPrepareStatementReportsOnceInDetectionOnlyMode() throws SQLException { + String payload = "Malicious Pet', 'Gru from the Minions') -- "; + String sql = "INSERT INTO pets (pet_name, owner) VALUES ('" + payload + "', 'Aikido Security')"; + Context.set(new EmptySampleContextObject(payload)); + ServiceConfigStore.updateBlocking(false); + + assertDoesNotThrow(() -> connection.prepareStatement(sql)); + + var stats = StatisticsStore.getStatsRecord(); + assertEquals(1, stats.requests().attacksDetected().total()); + assertEquals(1, stats.operations().values().stream() + .filter(record -> record.getKind() == OperationKind.SQL_OP) + .count()); + var operation = stats.operations() + .get("(MySQL Connector/J) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(1, operation.total()); + assertEquals(1, operation.getAttacksDetected().get("total")); + assertEquals(0, operation.getAttacksDetected().get("blocked")); + } + @Test public void testExecute() throws SQLException { Statement stmt = connection.createStatement(); diff --git a/agent_api/src/test/java/wrappers/PostgresWrapperTest.java b/agent_api/src/test/java/wrappers/PostgresWrapperTest.java index 850e0e34b..fb4717e8b 100644 --- a/agent_api/src/test/java/wrappers/PostgresWrapperTest.java +++ b/agent_api/src/test/java/wrappers/PostgresWrapperTest.java @@ -2,6 +2,8 @@ import dev.aikido.agent_api.context.Context; import dev.aikido.agent_api.storage.ServiceConfigStore; +import dev.aikido.agent_api.storage.statistics.OperationKind; +import dev.aikido.agent_api.storage.statistics.StatisticsStore; import dev.aikido.agent_api.vulnerabilities.sql_injection.SQLInjectionException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -25,6 +27,8 @@ public static void clean() { public void setUp() throws SQLException { // Connect to the PostgreSQL database connection = DriverManager.getConnection("jdbc:postgresql://localhost:5432/db", "user", "password"); + StatisticsStore.clear(); + ServiceConfigStore.updateBlocking(true); } @AfterEach @@ -32,7 +36,9 @@ public void tearDown() throws SQLException { if (connection != null) { connection.close(); } + StatisticsStore.clear(); Context.set(null); + ServiceConfigStore.updateBlocking(true); } @Test @@ -81,6 +87,54 @@ public void testSelectSqlWithPreparedStatementWithoutExecute() throws SQLExcepti assertEquals("Aikido Zen has blocked SQL Injection, Dialect: PostgreSQL", exception.getMessage()); } + @Test + public void testPrepareStatementReportsOnceInDetectionOnlyMode() throws SQLException { + String payload = "Malicious Pet', 'Gru from the Minions') -- "; + String sql = "INSERT INTO pets (pet_name, owner) VALUES ('" + payload + "', 'Aikido Security')"; + Context.set(new EmptySampleContextObject(payload)); + ServiceConfigStore.updateBlocking(false); + + assertDoesNotThrow(() -> connection.prepareStatement(sql)); + + var stats = StatisticsStore.getStatsRecord(); + assertEquals(1, stats.requests().attacksDetected().total()); + assertEquals(1, stats.operations().values().stream() + .filter(record -> record.getKind() == OperationKind.SQL_OP) + .count()); + var operation = stats.operations() + .get("(PostgreSQL JDBC Driver) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(1, operation.total()); + assertEquals(1, operation.getAttacksDetected().get("total")); + assertEquals(0, operation.getAttacksDetected().get("blocked")); + + ServiceConfigStore.updateBlocking(true); + assertThrows(SQLInjectionException.class, + () -> connection.prepareStatement(sql)); + } + + @Test + public void testPrepareStatementReportsAfterBlocking() throws SQLException { + String payload = "Malicious Pet', 'Gru from the Minions') -- "; + String sql = "INSERT INTO pets (pet_name, owner) VALUES ('" + payload + "', 'Aikido Security')"; + Context.set(new EmptySampleContextObject(payload)); + + assertThrows(SQLInjectionException.class, + () -> connection.prepareStatement(sql)); + + ServiceConfigStore.updateBlocking(false); + assertDoesNotThrow(() -> connection.prepareStatement(sql)); + + var stats = StatisticsStore.getStatsRecord(); + assertEquals(2, stats.requests().attacksDetected().total()); + var operation = stats.operations() + .get("(PostgreSQL JDBC Driver) java.sql.Connection.prepareStatement"); + assertNotNull(operation); + assertEquals(2, operation.total()); + assertEquals(2, operation.getAttacksDetected().get("total")); + assertEquals(1, operation.getAttacksDetected().get("blocked")); + } + @Test public void testExecute() throws SQLException { Statement stmt = connection.createStatement(); From 54d8cb6f6520fe2285dea5260eb5226846a493a0 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Mon, 31 Aug 2026 16:15:01 +0200 Subject: [PATCH 2/2] Add health checks to databases --- .github/workflows/gradle-tests.yml | 2 +- sample-apps/databases/docker-compose.yml | 15 +++++++++++++++ .../databases/mssql_database/setup_database.sh | 13 ++++++++++--- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/.github/workflows/gradle-tests.yml b/.github/workflows/gradle-tests.yml index f2afccba6..92bcf5455 100644 --- a/.github/workflows/gradle-tests.yml +++ b/.github/workflows/gradle-tests.yml @@ -63,7 +63,7 @@ jobs: - name: Start databases working-directory: ./sample-apps/databases - run: docker compose up --build -d postgres_database mysql_database mssql_database && sleep 10 + run: docker compose up --build -d --wait --wait-timeout 120 postgres_database mysql_database mssql_database - name: Start mock server working-directory: ./end2end/server diff --git a/sample-apps/databases/docker-compose.yml b/sample-apps/databases/docker-compose.yml index e8424e87b..f89da6269 100644 --- a/sample-apps/databases/docker-compose.yml +++ b/sample-apps/databases/docker-compose.yml @@ -14,6 +14,11 @@ services: POSTGRES_PASSWORD: 'password' ports: - "5432:5432" + healthcheck: + test: ["CMD-SHELL", "PGPASSWORD=password psql -h 127.0.0.1 -U user -d db -c 'SELECT 1 FROM pets LIMIT 1' >/dev/null 2>&1"] + interval: 2s + timeout: 5s + retries: 60 mysql_database: image: mysql container_name: sample_mysql_db @@ -29,6 +34,11 @@ services: ports: - "3306:3306" command: --init-file /data/application/init.sql + healthcheck: + test: ["CMD-SHELL", "MYSQL_PWD=password mysql -h 127.0.0.1 -u user db -e 'SELECT 1 FROM pets LIMIT 1' >/dev/null 2>&1"] + interval: 2s + timeout: 5s + retries: 60 mssql_database: build: ./mssql_database container_name: sample_mssql_db @@ -40,6 +50,11 @@ services: - "1433:1433" volumes: - mssql_db_data:/var/opt/mssql + healthcheck: + test: ["CMD-SHELL", "/opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P \"$${SA_PASSWORD}\" -d db -Q 'SELECT TOP 1 1 FROM pets' -N -C -b >/dev/null 2>&1"] + interval: 2s + timeout: 5s + retries: 60 mongo_database: image: mongo:latest container_name: sample_mongo_db_java diff --git a/sample-apps/databases/mssql_database/setup_database.sh b/sample-apps/databases/mssql_database/setup_database.sh index 2a705068a..6fd9cb065 100755 --- a/sample-apps/databases/mssql_database/setup_database.sh +++ b/sample-apps/databases/mssql_database/setup_database.sh @@ -1,4 +1,11 @@ #!/usr/bin/env bash -# Wait for database to startup -sleep 15 -./opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P "Strong!Passw0rd" -i setup.sql -N -C + +for attempt in {1..60}; do + if /opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P "$SA_PASSWORD" -Q "SELECT 1" -N -C >/dev/null 2>&1; then + exec /opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P "$SA_PASSWORD" -i /setup.sql -N -C -b + fi + sleep 2 +done + +echo "SQL Server did not become ready in time" >&2 +exit 1