improved WallFilter.

This commit is contained in:
wenshao
2017-12-01 18:32:55 +08:00
parent ba07df232b
commit 2180b61dd1
2 changed files with 57 additions and 8 deletions
@@ -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());
}
}