mirror of
https://github.com/alibaba/druid.git
synced 2023-03-04 19:17:16 +00:00
improved WallFilter.
This commit is contained in:
@@ -31,7 +31,9 @@ import java.util.concurrent.locks.Lock;
|
||||
import java.util.concurrent.locks.ReentrantReadWriteLock;
|
||||
|
||||
import com.alibaba.druid.sql.SQLUtils;
|
||||
import com.alibaba.druid.sql.ast.SQLName;
|
||||
import com.alibaba.druid.sql.ast.SQLStatement;
|
||||
import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
|
||||
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlHintStatement;
|
||||
import com.alibaba.druid.sql.parser.Lexer;
|
||||
import com.alibaba.druid.sql.parser.NotAllowCommentException;
|
||||
@@ -669,16 +671,39 @@ public abstract class WallProvider {
|
||||
violations.addAll(visitor.getViolations());
|
||||
}
|
||||
|
||||
Map<String, WallSqlTableStat> tableStat = context.getTableStats();
|
||||
|
||||
boolean updateCheckHandlerEnable = false;
|
||||
{
|
||||
WallUpdateCheckHandler updateCheckHandler = config.getUpdateCheckHandler();
|
||||
if (updateCheckHandler != null) {
|
||||
for (SQLStatement stmt : statementList) {
|
||||
if (stmt instanceof SQLUpdateStatement) {
|
||||
SQLUpdateStatement updateStmt = (SQLUpdateStatement) stmt;
|
||||
SQLName table = updateStmt.getTableName();
|
||||
if (table != null) {
|
||||
String tableName = table.getSimpleName();
|
||||
Set<String> updateCheckColumns = config.getUpdateCheckTable(tableName);
|
||||
if (updateCheckColumns != null && updateCheckColumns.size() > 0) {
|
||||
updateCheckHandlerEnable = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
WallSqlStat sqlStat = null;
|
||||
if (violations.size() > 0) {
|
||||
violationCount.incrementAndGet();
|
||||
|
||||
if (sql.length() < MAX_SQL_LENGTH) {
|
||||
sqlStat = addBlackSql(sql, context.getTableStats(), context.getFunctionStats(), violations, syntaxError);
|
||||
if ((!updateCheckHandlerEnable) && sql.length() < MAX_SQL_LENGTH) {
|
||||
sqlStat = addBlackSql(sql, tableStat, context.getFunctionStats(), violations, syntaxError);
|
||||
}
|
||||
} else {
|
||||
if (sql.length() < MAX_SQL_LENGTH) {
|
||||
sqlStat = addWhiteSql(sql, context.getTableStats(), context.getFunctionStats(), syntaxError);
|
||||
if ((!updateCheckHandlerEnable) && sql.length() < MAX_SQL_LENGTH) {
|
||||
sqlStat = addWhiteSql(sql, tableStat, context.getFunctionStats(), syntaxError);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -712,6 +737,10 @@ public abstract class WallProvider {
|
||||
}
|
||||
|
||||
private WallCheckResult checkWhiteAndBlackList(String sql) {
|
||||
if (config.getUpdateCheckHandler() != null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// check black list
|
||||
if (blackListEnable) {
|
||||
WallSqlStat sqlStat = getBlackSql(sql);
|
||||
|
||||
@@ -19,6 +19,13 @@ public class WallUpdateCheckTest extends TestCase {
|
||||
Properties properties = new Properties();
|
||||
properties.put("druid.wall.updateCheckColumns", "t_orders.status");
|
||||
wallProvider.getConfig().configFromProperties(properties);
|
||||
|
||||
}
|
||||
public void test_update_check_handler() throws Exception {
|
||||
{
|
||||
WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status = 4");
|
||||
assertTrue(result.getViolations().size() == 0);
|
||||
}
|
||||
wallProvider.getConfig().setUpdateCheckHandler(new WallUpdateCheckHandler() {
|
||||
|
||||
@Override
|
||||
@@ -26,9 +33,22 @@ public class WallUpdateCheckTest extends TestCase {
|
||||
return false;
|
||||
}
|
||||
});
|
||||
}
|
||||
public void test_update_check_handler() throws Exception {
|
||||
WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status = 4");
|
||||
assertTrue(result.getViolations().size() > 0);
|
||||
{
|
||||
WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status = 4");
|
||||
assertTrue(result.getViolations().size() > 0);
|
||||
}
|
||||
wallProvider.getConfig().setUpdateCheckHandler(new WallUpdateCheckHandler() {
|
||||
|
||||
@Override
|
||||
public boolean check(String table, String column, Object setValue, List<Object> filterValues) {
|
||||
return true;
|
||||
}
|
||||
});
|
||||
{
|
||||
WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status = 4");
|
||||
assertTrue(result.getViolations().size() > 0);
|
||||
}
|
||||
assertEquals(0, wallProvider.getWhiteListHitCount());
|
||||
assertEquals(0, wallProvider.getBlackListHitCount());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user