add sql wall support

src/main/java/com/alibaba/druid/filter/wall/IllegalSQLObjectViolation.java

@@ -1,14 +1,14 @@
 package com.alibaba.druid.filter.wall;
 
-public class IllegalStatementViolation implements Violation {
+public class IllegalSQLObjectViolation implements Violation {
 
     private String condition;
 
-    public IllegalStatementViolation(){
+    public IllegalSQLObjectViolation(){
 
     }
 
-    public IllegalStatementViolation(String condition){
+    public IllegalSQLObjectViolation(String condition){
         this.condition = condition;
     }
 

src/main/java/com/alibaba/druid/filter/wall/WallFilter.java

@@ -1,25 +1,30 @@
 package com.alibaba.druid.filter.wall;
 
 import java.sql.SQLException;
-import java.util.List;
 import java.util.Properties;
 
 import com.alibaba.druid.filter.FilterAdapter;
 import com.alibaba.druid.filter.FilterChain;
+import com.alibaba.druid.filter.wall.spi.MySqlWallProvider;
+import com.alibaba.druid.filter.wall.spi.MySqlWallVisitor;
+import com.alibaba.druid.filter.wall.spi.OracleWallVisitor;
 import com.alibaba.druid.proxy.jdbc.CallableStatementProxy;
 import com.alibaba.druid.proxy.jdbc.ConnectionProxy;
+import com.alibaba.druid.proxy.jdbc.DataSourceProxy;
 import com.alibaba.druid.proxy.jdbc.PreparedStatementProxy;
-import com.alibaba.druid.sql.ast.SQLStatement;
-import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
-import com.alibaba.druid.sql.dialect.oracle.parser.OracleStatementParser;
-import com.alibaba.druid.sql.dialect.postgresql.parser.PGSQLStatementParser;
-import com.alibaba.druid.sql.dialect.sqlserver.parser.SQLServerStatementParser;
-import com.alibaba.druid.sql.parser.SQLStatementParser;
-import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 import com.alibaba.druid.util.JdbcUtils;
 
 public class WallFilter extends FilterAdapter {
 
+    private WallProvider provider;
+
+    @Override
+    public void init(DataSourceProxy dataSource) {
+        this.dataSource = dataSource;
+
+        provider = createWallProvider(dataSource.getDbType());
+    }
+
     @Override
     public ConnectionProxy connection_connect(FilterChain chain, Properties info) throws SQLException {
         return chain.connection_connect(info);
@@ -94,44 +99,27 @@ public class WallFilter extends FilterAdapter {
     }
 
     public void check(ConnectionProxy connection, String sql) throws SQLException {
-        String dbType = connection.getDirectDataSource().getDbType();
+        provider.check(sql, true);
+    }
 
-        SQLStatementParser parser = createParser(dbType, sql);
-        parser.getLexer().setAllowComment(false); // permit comment
-
-        List<SQLStatement> statementList = parser.parseStatementList();
-        if (statementList.size() > 1) {
-            throw new SQLException("multi-statement : " + sql);
+    public WallProvider createWallProvider(String dbType) {
+        if (JdbcUtils.MYSQL.equals(dbType)) {
+            return new MySqlWallProvider();
         }
 
-        SQLStatement stmt = statementList.get(0);
-
-        SQLASTVisitor visitor = createWallVisitor(dbType);
-
-        stmt.accept(visitor);
+        throw new IllegalStateException("dbType not support : " + dbType);
     }
 
-    public SQLASTVisitor createWallVisitor(String dbType) throws SQLException {
-        throw new SQLException("dbType not support : " + dbType);
-    }
-
-    public SQLStatementParser createParser(String dbType, String sql) throws SQLException {
+    public WallVisitor createWallVisitor(String dbType) throws SQLException {
         if (JdbcUtils.ORACLE.equals(dbType)) {
-            return new OracleStatementParser(sql);
+            return new OracleWallVisitor();
         }
 
         if (JdbcUtils.MYSQL.equals(dbType)) {
-            return new MySqlStatementParser(sql);
-        }
-
-        if (JdbcUtils.POSTGRESQL.equals(dbType)) {
-            return new PGSQLStatementParser(sql);
-        }
-
-        if (JdbcUtils.SQL_SERVER.equals(dbType)) {
-            return new SQLServerStatementParser(sql);
+            return new MySqlWallVisitor();
         }
 
         throw new SQLException("dbType not support : " + dbType);
     }
+
 }

src/main/java/com/alibaba/druid/filter/wall/WallProvider.java

@@ -0,0 +1,79 @@
+package com.alibaba.druid.filter.wall;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+
+public abstract class WallProvider {
+
+    // Dummy value to associate with an Object in the backing Map
+    private static final Object           PRESENT           = new Object();
+
+    private LinkedHashMap<String, Object> whiteList;
+
+    private int                           whileListMaxSize  = 1024;
+
+    private int                           whiteSqlMaxLength = 1024;        // 1k
+
+    public synchronized void addWhiteSql(String sql) {
+        if (whiteList == null) {
+            whiteList = new LinkedHashMap<String, Object>(whileListMaxSize, 0.75f, true);
+        }
+
+        whiteList.put(sql, PRESENT);
+    }
+
+    public synchronized boolean whiteContains(String sql) {
+        if (whiteList == null) {
+            return false;
+        }
+        
+        return whiteList.get(sql) != null;
+    }
+
+    public abstract SQLStatementParser createParser(String sql);
+
+    public abstract WallVisitor createWallVisitor();
+
+    public boolean check(String sql, boolean throwException) throws WallRuntimeException {
+        // first step, check whiteList
+        boolean isWhite = whiteContains(sql);
+        if (isWhite) {
+            return true;
+        }
+
+        SQLStatementParser parser = createParser(sql);
+        parser.getLexer().setAllowComment(false); // permit comment
+
+        List<SQLStatement> statementList = parser.parseStatementList();
+        if (statementList.size() > 1) {
+            if (throwException) {
+                throw new WallRuntimeException("multi-statement : " + sql);
+            }
+
+            return false;
+        }
+
+        SQLStatement stmt = statementList.get(0);
+
+        WallVisitor visitor = createWallVisitor();
+
+        stmt.accept(visitor);
+
+        if (visitor.getViolations().size() > 0) {
+            if (throwException) {
+                throw new WallRuntimeException();
+            } else {
+                return false;
+            }
+        }
+
+        if (sql.length() < whiteSqlMaxLength) {
+            this.addWhiteSql(sql);
+        }
+
+        return true;
+    }
+}

src/main/java/com/alibaba/druid/filter/wall/WallRuntimeException.java

@@ -0,0 +1,24 @@
+package com.alibaba.druid.filter.wall;
+
+
+public class WallRuntimeException extends RuntimeException {
+
+    private static final long serialVersionUID = 1L;
+
+    public WallRuntimeException(){
+        super();
+    }
+
+    public WallRuntimeException(String reason, Throwable cause){
+        super(reason, cause);
+    }
+
+    public WallRuntimeException(String reason){
+        super(reason);
+    }
+
+    public WallRuntimeException(Throwable cause){
+        super(cause);
+    }
+
+}

src/main/java/com/alibaba/druid/filter/wall/WallSQLException.java

@@ -0,0 +1,25 @@
+package com.alibaba.druid.filter.wall;
+
+import java.sql.SQLException;
+
+public class WallSQLException extends SQLException {
+
+    private static final long serialVersionUID = 1L;
+
+    public WallSQLException(){
+        super();
+    }
+
+    public WallSQLException(String reason, Throwable cause){
+        super(reason, cause);
+    }
+
+    public WallSQLException(String reason){
+        super(reason);
+    }
+
+    public WallSQLException(Throwable cause){
+        super(cause);
+    }
+
+}

src/main/java/com/alibaba/druid/filter/wall/WallUtils.java

@@ -0,0 +1,11 @@
+package com.alibaba.druid.filter.wall;
+
+import com.alibaba.druid.filter.wall.spi.MySqlWallProvider;
+
+public class WallUtils {
+
+    public static boolean isValidateMySql(String sql) {
+        MySqlWallProvider provider = new MySqlWallProvider();
+        return provider.check(sql, false);
+    }
+}

src/main/java/com/alibaba/druid/filter/wall/spi/MySqlWallProvider.java

@@ -0,0 +1,21 @@
+package com.alibaba.druid.filter.wall.spi;
+
+import com.alibaba.druid.filter.wall.WallProvider;
+import com.alibaba.druid.filter.wall.WallVisitor;
+import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+
+
+public class MySqlWallProvider extends WallProvider {
+
+    @Override
+    public SQLStatementParser createParser(String sql) {
+        return new MySqlStatementParser(sql);
+    }
+
+    @Override
+    public WallVisitor createWallVisitor() {
+        return new MySqlWallVisitor();
+    }
+
+}

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

@@ -1,14 +1,22 @@
 package com.alibaba.druid.filter.wall.spi;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
-import com.alibaba.druid.filter.wall.IllegalStatementViolation;
+import com.alibaba.druid.filter.wall.IllegalSQLObjectViolation;
 import com.alibaba.druid.filter.wall.Violation;
 import com.alibaba.druid.filter.wall.WallVisitor;
 import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
+import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
+import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr;
+import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
+import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
 import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement;
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlExecuteStatement;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
@@ -19,14 +27,21 @@ import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
 
 public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisitor {
 
+    private Set<String>           permitFunctions = new HashSet<String>();
+    private Set<String>           permitSchemas   = new HashSet<String>();
+
     private final List<Violation> violations;
 
-    public MySqlWallVisitor() {
+    public MySqlWallVisitor(){
         this(new ArrayList<Violation>());
     }
 
-    public MySqlWallVisitor(List<Violation> violations) {
+    public MySqlWallVisitor(List<Violation> violations){
         this.violations = violations;
+
+permitFunctions.add("load_file");
+
+permitSchemas.add("information_schema");
     }
 
     public List<Violation> getViolations() {
@@ -34,50 +49,92 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
     }
 
     public boolean visit(MySqlExecuteStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toMySqlString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
         return false;
     }
-    
+
     @Override
     public boolean visit(MySqlShowTablesStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toMySqlString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
         return false;
     }
-    
+
     @Override
     public boolean visit(MySqlShowDatabasesStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toMySqlString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
         return false;
     }
-    
+
     @Override
     public boolean visit(MySqlShowColumnsStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toMySqlString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
         return false;
     }
-    
+
     public boolean visit(SQLBinaryOpExpr x) {
         WallVisitorUtils.check(this, x);
 
         return true;
     }
-    
+
     @Override
     public boolean visit(SQLSelectQueryBlock x) {
         if (x.getWhere() != null) {
             x.getWhere().setParent(x);
         }
-        
+
         return true;
     }
-    
+
     @Override
     public boolean visit(MySqlSelectQueryBlock x) {
         return visit((SQLSelectQueryBlock) x);
     }
-    
+
     public boolean visit(SQLDropTableStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toMySqlString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
         return false;
     }
+
+    public boolean visit(SQLVariantRefExpr x) {
+        String varName = x.getName();
+        if (varName == null) {
+            return false;
+        }
+
+        if (varName.startsWith("@@")) {
+            violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
+        }
+
+        return false;
+    }
+
+    @Override
+    public boolean visit(SQLMethodInvokeExpr x) {
+        String methodName = x.getMethodName();
+
+        if (permitFunctions.contains(methodName.toLowerCase())) {
+            violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
+        }
+
+        return true;
+    }
+
+    public boolean visit(SQLExprTableSource x) {
+        SQLExpr expr = x.getExpr();
+
+        if (expr instanceof SQLPropertyExpr) {
+            SQLPropertyExpr propExpr = (SQLPropertyExpr) expr;
+
+            if (propExpr.getOwner() instanceof SQLIdentifierExpr) {
+                String ownerName = ((SQLIdentifierExpr) propExpr.getOwner()).getName();
+
+                if (permitSchemas.contains(ownerName.toLowerCase())) {
+                    violations.add(new IllegalSQLObjectViolation(SQLUtils.toMySqlString(x)));
+                }
+            }
+        }
+
+        return true;
+    }
 }

src/main/java/com/alibaba/druid/filter/wall/spi/OracleWallVisitor.java

@@ -3,7 +3,7 @@ package com.alibaba.druid.filter.wall.spi;
 import java.util.ArrayList;
 import java.util.List;
 
-import com.alibaba.druid.filter.wall.IllegalStatementViolation;
+import com.alibaba.druid.filter.wall.IllegalSQLObjectViolation;
 import com.alibaba.druid.filter.wall.Violation;
 import com.alibaba.druid.filter.wall.WallVisitor;
 import com.alibaba.druid.sql.SQLUtils;
@@ -34,7 +34,7 @@ public class OracleWallVisitor extends OracleASTVIsitorAdapter implements WallVi
     }
     
     public boolean visit(SQLDropTableStatement x) {
-        violations.add(new IllegalStatementViolation(SQLUtils.toOracleString(x)));
+        violations.add(new IllegalSQLObjectViolation(SQLUtils.toOracleString(x)));
         return false;
     }
 }