improve wall filter

src/main/java/com/alibaba/druid/wall/spi/MySqlWallVisitor.java

@@ -63,7 +63,6 @@ import com.alibaba.druid.wall.WallContext;
 import com.alibaba.druid.wall.WallProvider;
 import com.alibaba.druid.wall.WallSqlTableStat;
 import com.alibaba.druid.wall.WallVisitor;
-import com.alibaba.druid.wall.spi.WallVisitorUtils.WallConditionContext;
 import com.alibaba.druid.wall.violation.ErrorCode;
 import com.alibaba.druid.wall.violation.IllegalSQLObjectViolation;
 
@@ -178,13 +177,13 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
     public boolean visit(Limit x) {
         if (x.getRowCount() instanceof SQLNumericLiteralExpr) {
             WallContext context = WallContext.current();
-            
+
             int rowCount = ((SQLNumericLiteralExpr) x.getRowCount()).getNumber().intValue();
             if (rowCount == 0) {
                 if (context != null) {
                     context.incrementWarnnings();
                 }
-                
+
                 if (!provider.getConfig().isLimitZeroAllow()) {
                     this.getViolations().add(new IllegalSQLObjectViolation(ErrorCode.LIMIT_ZERO, "limit row 0",
                                                                            this.toSQL(x)));
@@ -209,8 +208,21 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
                 if (!checkVar(x.getParent(), x.getName())) {
                     boolean isTop = WallVisitorUtils.isTopNoneFromSelect(this, x);
                     if (!isTop) {
-                        violations.add(new IllegalSQLObjectViolation(ErrorCode.VARIANT_DENY, "variable not allow : "
-                                                                                             + x.getName(), toSQL(x)));
+                        boolean allow = true;
+                        if (isDeny(varName)) {
+                            allow = false;
+                        }
+
+                        if (allow) {
+                            if (WallVisitorUtils.isWhereOrHaving(x)) {
+                                allow = false;
+                            }
+                        }
+
+                        if (!allow) {
+                            violations.add(new IllegalSQLObjectViolation(ErrorCode.VARIANT_DENY, "variable not allow : "
+                                                                                                 + x.getName(), toSQL(x)));
+                        }
                     }
                 }
                 return false;
@@ -248,6 +260,14 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
 
         return false;
     }
+    
+    public boolean isDeny(String varName) {
+        if (varName.startsWith("@@")) {
+            varName = varName.substring(2);
+        }
+        
+        return config.getDenyVariants().contains(varName);
+    }
 
     public boolean visit(SQLVariantRefExpr x) {
         String varName = x.getName();
@@ -256,10 +276,23 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
         }
 
         if (varName.startsWith("@@") && !checkVar(x.getParent(), x.getName())) {
-            boolean noneCondition = WallVisitorUtils.getWallConditionContext() == null;
-            if (!noneCondition) {
-                violations.add(new IllegalSQLObjectViolation(ErrorCode.VARIANT_DENY, "variable not allow : "
-                                                                                     + x.getName(), toSQL(x)));
+            boolean isTop = WallVisitorUtils.isTopNoneFromSelect(this, x);
+            if (!isTop) {
+                boolean allow = true;
+                if (isDeny(varName)) {
+                    allow = false;
+                }
+
+                if (allow) {
+                    if (WallVisitorUtils.isWhereOrHaving(x)) {
+                        allow = false;
+                    }
+                }
+
+                if (!allow) {
+                    violations.add(new IllegalSQLObjectViolation(ErrorCode.VARIANT_DENY, "variable not allow : "
+                                                                                         + x.getName(), toSQL(x)));
+                }
             }
         }
 

src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java

@@ -60,6 +60,7 @@ import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLRollbackStatement;
 import com.alibaba.druid.sql.ast.statement.SQLSelect;
+import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause;
 import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
@@ -622,6 +623,40 @@ public class WallVisitorUtils {
         return false;
     }
 
+    public static boolean isWhereOrHaving(SQLObject x) {
+        if (x == null) {
+            return false;
+        }
+        
+        for (;;) {
+            SQLObject parent = x.getParent();
+
+            if (parent == null) {
+                return false;
+            }
+            
+            if (parent instanceof SQLSelectQueryBlock) {
+                SQLSelectQueryBlock query = (SQLSelectQueryBlock) parent;
+                if (query.getWhere() == x) {
+                    return true;
+                } else {
+                    return false;
+                }
+            }
+            
+            if (parent instanceof SQLSelectGroupByClause) {
+                SQLSelectGroupByClause groupBy = (SQLSelectGroupByClause) parent;
+                if (x == groupBy.getHaving()) {
+                    return true;
+                } else {
+                    return false;
+                }
+            }
+            
+            x = parent;
+        }
+    }
+
     public static class WallConditionContext {
 
         private boolean partAlwayrue;