diff --git a/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java b/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java index 75242f0d79..d2df0e8c94 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java +++ b/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java @@ -57,12 +57,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.net.URL; -import java.util.ArrayList; -import java.util.Collections; -import java.util.Enumeration; -import java.util.List; -import java.util.Set; -import java.util.Stack; +import java.util.*; import static com.alibaba.druid.sql.visitor.SQLEvalVisitor.EVAL_VALUE; @@ -899,6 +894,15 @@ public class WallVisitorUtils { List conditions; if (where instanceof SQLBinaryOpExpr) { conditions = SQLBinaryOpExpr.split((SQLBinaryOpExpr) where, SQLBinaryOperator.BooleanAnd); + } else if (where instanceof SQLBinaryOpExprGroup) { + conditions = new ArrayList(); + for (SQLExpr each : ((SQLBinaryOpExprGroup) where).getItems()) { + if (each instanceof SQLBinaryOpExpr) { + conditions.addAll(SQLBinaryOpExpr.split((SQLBinaryOpExpr) each, SQLBinaryOperator.BooleanAnd)); + } else if (each instanceof SQLInListExpr) { + conditions.add(each); + } + } } else { conditions = new ArrayList(); conditions.add(where); @@ -919,6 +923,14 @@ public class WallVisitorUtils { filterValueExprList.add(right); } } + } else if (condition instanceof SQLInListExpr) { + SQLInListExpr listExpr = (SQLInListExpr) condition; + if (listExpr.getExpr() instanceof SQLIdentifierExpr) { + SQLIdentifierExpr nameExpr = (SQLIdentifierExpr) listExpr.getExpr(); + if (nameExpr.getName().equals(checkColumn)) { + filterValueExprList.addAll(((SQLInListExpr) condition).getTargetList()); + } + } } } @@ -938,6 +950,7 @@ public class WallVisitorUtils { for (SQLExpr expr : filterValueExprList) { filterValues.add(((SQLValuableExpr) expr).getValue()); } + filterValues = new ArrayList(new HashSet(filterValues)); boolean validate = updateCheckHandler.check(tableName, checkColumn, setValue, filterValues); if (!validate) { visitor.addViolation(new IllegalSQLObjectViolation(ErrorCode.UPDATE_CHECK_FAIL, "update check failed.", visitor.toSQL(x))); 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 71a548ef39..23f66ecf8e 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,8 +19,8 @@ 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"); @@ -50,5 +50,22 @@ public class WallUpdateCheckTest extends TestCase { } assertEquals(0, wallProvider.getWhiteListHitCount()); assertEquals(0, wallProvider.getBlackListHitCount()); + wallProvider.getConfig().setUpdateCheckHandler(new WallUpdateCheckHandler() { + + @Override + public boolean check(String table, String column, Object setValue, List filterValues) { + //增加对in语句的支持, status in (1, 2) 应该返回的filterValue为1和2 + assertTrue(filterValues.size() == 2); + return true; + } + }); + { + WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status in (1, 2)"); + assertTrue(result.getViolations().size() == 0); + } + { + WallCheckResult result = wallProvider.check("update t_orders set status = 3 where id = 3 and status = 3 and status in (3, 4)"); + assertTrue(result.getViolations().size() == 0); + } } }