Skip to content

Prevent multiple attack events for SQL injection (detection-only mode) #352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hansott wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Prevent multiple attack events for SQL injection (detection-only mode) #352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
75c25b9
Prevent multiple attack events for SQL injection (detection-only mode)
hansott Aug 31, 2026
54d8cb6
Add health checks to databases
hansott Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/gradle-tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions agent/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
}
Expand Down
37 changes: 34 additions & 3 deletions agent/src/main/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdvice.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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<Deque<String>> jdbcCallStack = ThreadLocal.withInitial(ArrayDeque::new);
private JDBCConnectionAdvice() {}
public static ElementMatcher<? super MethodDescription> getMatcher(String module) {
ElementMatcher.Junction<? super MethodDescription> statementMatcher =
Expand All @@ -41,12 +54,18 @@ public static ElementMatcher<? super MethodDescription> 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<String> 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;
Expand All @@ -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);
}
Comment thread
hansott marked this conversation as resolved.
}
return sql;
}

@Advice.OnMethodExit(onThrowable = Throwable.class)
public static void after(@Advice.Enter String enteredSql) {
if (enteredSql == null) {
return;
}
Deque<String> sqlCalls = jdbcCallStack.get();
if (!sqlCalls.isEmpty() && sqlCalls.peek() == enteredSql) {
sqlCalls.pop();
}
}
}
59 changes: 59 additions & 0 deletions agent/src/test/java/dev/aikido/agent/wrappers/jdbc/JDBCConnectionAdviceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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());
}
}
26 changes: 25 additions & 1 deletion agent_api/src/test/java/wrappers/MSSQLWrapperTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -44,6 +45,7 @@ public void tearDown() throws SQLException {
}
Context.set(null);
StatisticsStore.clear();
ServiceConfigStore.updateBlocking(true);
}

@Test
Expand All @@ -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"));
Expand Down
32 changes: 32 additions & 0 deletions agent_api/src/test/java/wrappers/MariadbWrapperTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -11,6 +13,7 @@

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;

Expand All @@ -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
Expand All @@ -37,6 +42,8 @@ public void tearDown() throws SQLException {
connection.close();
}
Context.set(null);
StatisticsStore.clear();
ServiceConfigStore.updateBlocking(true);
}

@Test
Expand All @@ -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"));
Expand Down
28 changes: 28 additions & 0 deletions agent_api/src/test/java/wrappers/MysqlCJWrapperTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -26,14 +28,18 @@ 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
public void tearDown() throws SQLException {
if (connection != null) {
connection.close();
}
StatisticsStore.clear();
Context.set(null);
ServiceConfigStore.updateBlocking(true);
}

@Test
Expand Down Expand Up @@ -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();
Expand Down
Loading
Loading