diff --git a/src/main/java/com/alibaba/druid/wall/WallProvider.java b/src/main/java/com/alibaba/druid/wall/WallProvider.java index 77f7db4d0b..e80e8bd9a3 100644 --- a/src/main/java/com/alibaba/druid/wall/WallProvider.java +++ b/src/main/java/com/alibaba/druid/wall/WallProvider.java @@ -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 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 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); diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallUpdateCheckTest.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallUpdateCheckTest.java index b324a62838..e86de96985 100644 --- a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallUpdateCheckTest.java +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallUpdateCheckTest.java @@ -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 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()); } }