增加对多条件的SQL支持状态机&增加对in语句的SQL状态机支持

This commit is contained in:
黑昼
2018-03-16 21:25:03 +08:00
parent b63a97c02f
commit 1a149d1fa3
2 changed files with 37 additions and 7 deletions
@@ -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<SQLExpr> conditions;
if (where instanceof SQLBinaryOpExpr) {
conditions = SQLBinaryOpExpr.split((SQLBinaryOpExpr) where, SQLBinaryOperator.BooleanAnd);
} else if (where instanceof SQLBinaryOpExprGroup) {
conditions = new ArrayList<SQLExpr>();
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<SQLExpr>();
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)));
@@ -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<Object> 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);
}
}
}