improve sql-wall

src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlParameterizedOutputVisitor.java

@@ -15,6 +15,8 @@
  */
 package com.alibaba.druid.sql.dialect.mysql.visitor;
 
+import java.util.List;
+
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
@@ -24,6 +26,7 @@ import com.alibaba.druid.sql.ast.expr.SQLNCharExpr;
 import com.alibaba.druid.sql.ast.expr.SQLNumberExpr;
 import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLInsertStatement.ValuesClause;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
 import com.alibaba.druid.sql.visitor.ParameterizedOutputVisitorUtils;
 import com.alibaba.druid.sql.visitor.ParameterizedVisitor;
@@ -134,9 +137,13 @@ public class MySqlParameterizedOutputVisitor extends MySqlOutputVisitor implemen
     }
 
     protected void printValuesList(MySqlInsertStatement x) {
+        List<ValuesClause> valuesList = x.getValuesList();
         print("VALUES ");
         incrementIndent();
-        x.getValuesList().get(0).accept(this);
+        valuesList.get(0).accept(this);
         decrementIndent();
+        if (valuesList.size() > 1) {
+            this.incrementReplaceCunt();
+        }
     }
 }

src/main/java/com/alibaba/druid/sql/visitor/ParameterizedVisitor.java

@@ -1,3 +1,18 @@
+/*
+ * Copyright 1999-2011 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package com.alibaba.druid.sql.visitor;
 
 public interface ParameterizedVisitor extends SQLASTVisitor {

src/main/java/com/alibaba/druid/wall/WallCheckResult.java

@@ -34,6 +34,8 @@ public class WallCheckResult {
 
     private final WallSqlStat                      sqlStat;
 
+    private String                                 sql;
+
     public WallCheckResult(){
         this(null);
     }
@@ -70,6 +72,14 @@ public class WallCheckResult {
         this.syntaxError = syntaxError;
     }
 
+    public String getSql() {
+        return sql;
+    }
+
+    public void setSql(String sql) {
+        this.sql = sql;
+    }
+
     public List<Violation> getViolations() {
         return violations;
     }
@@ -90,7 +100,6 @@ public class WallCheckResult {
         return syntaxError;
     }
 
-    
     public WallSqlStat getSqlStat() {
         return sqlStat;
     }

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

@@ -582,8 +582,8 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
             }
         }
 
-        // TODO
-        return sql;
+        String resultSql = checkResult.getSql();
+        return resultSql;
     }
 
     @Override

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

@@ -28,6 +28,7 @@ import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
+import com.alibaba.druid.sql.SQLUtils;
 import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.parser.Lexer;
 import com.alibaba.druid.sql.parser.NotAllowCommentException;
@@ -532,7 +533,9 @@ public abstract class WallProvider {
         WallContext context = WallContext.current();
 
         if (config.isDoPrivilegedAllow() && ispPivileged()) {
-            return new WallCheckResult();
+            WallCheckResult checkResult = new WallCheckResult();
+            checkResult.setSql(sql);
+            return checkResult;
         }
 
         // first step, check whiteList
@@ -540,6 +543,7 @@ public abstract class WallProvider {
         if (!mulltiTenant) {
             WallCheckResult checkResult = checkWhiteAndBlackList(sql);
             if (checkResult != null) {
+                checkResult.setSql(sql);
                 return checkResult;
             }
         }
@@ -630,12 +634,23 @@ public abstract class WallProvider {
             recordStats(tableStats, functionStats);
         }
 
+        WallCheckResult result;
         if (sqlStat != null) {
             context.setSqlStat(sqlStat);
-            return new WallCheckResult(sqlStat, statementList);
+            result = new WallCheckResult(sqlStat, statementList);
+        } else {
+            result = new WallCheckResult(sqlStat, violations, tableStats, functionStats, statementList, syntaxError);
         }
 
-        return new WallCheckResult(sqlStat, violations, tableStats, functionStats, statementList, syntaxError);
+        String resultSql;
+        if (visitor.isSqlModified()) {
+            resultSql = SQLUtils.toSQLString(statementList, dbType);
+        } else {
+            resultSql = sql;
+        }
+        result.setSql(resultSql);
+
+        return result;
     }
 
     private WallCheckResult checkWhiteAndBlackList(String sql) {

src/main/java/com/alibaba/druid/wall/WallVisitor.java

@@ -23,14 +23,18 @@ import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 public interface WallVisitor extends SQLASTVisitor {
 
     WallConfig getConfig();
-    
+
     WallProvider getProvider();
 
     List<Violation> getViolations();
-    
+
     void addViolation(Violation violation);
-    
+
     boolean isDenyTable(String name);
 
     String toSQL(SQLObject obj);
+
+    boolean isSqlModified();
+
+    void setSqlModified(boolean sqlModified);
 }

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

@@ -53,21 +53,35 @@ public class DB2WallVisitor extends DB2ASTVisitorAdapter implements WallVisitor
 
     private final WallConfig      config;
     private final WallProvider    provider;
-    private final List<Violation> violations = new ArrayList<Violation>();
+    private final List<Violation> violations  = new ArrayList<Violation>();
+    private boolean               sqlModified = false;
 
     public DB2WallVisitor(WallProvider provider){
         this.config = provider.getConfig();
         this.provider = provider;
     }
 
+    @Override
+    public boolean isSqlModified() {
+        return sqlModified;
+    }
+
+    @Override
+    public void setSqlModified(boolean sqlModified) {
+        this.sqlModified = sqlModified;
+    }
+
+    @Override
     public WallProvider getProvider() {
         return provider;
     }
 
+    @Override
     public WallConfig getConfig() {
         return config;
     }
 
+    @Override
     public void addViolation(Violation violation) {
         this.violations.add(violation);
     }
@@ -129,11 +143,11 @@ public class DB2WallVisitor extends DB2ASTVisitorAdapter implements WallVisitor
 
         return true;
     }
-    
+
     @Override
     public boolean visit(DB2SelectQueryBlock x) {
         WallVisitorUtils.checkSelelct(this, x);
-        
+
         return true;
     }
 

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

@@ -70,25 +70,40 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi
 
     private final WallConfig      config;
     private final WallProvider    provider;
-    private final List<Violation> violations = new ArrayList<Violation>();
+    private final List<Violation> violations  = new ArrayList<Violation>();
+    private boolean               sqlModified = false;
 
     public MySqlWallVisitor(WallProvider provider){
         this.config = provider.getConfig();
         this.provider = provider;
     }
 
+    @Override
+    public boolean isSqlModified() {
+        return sqlModified;
+    }
+
+    @Override
+    public void setSqlModified(boolean sqlModified) {
+        this.sqlModified = sqlModified;
+    }
+
+    @Override
     public WallProvider getProvider() {
         return provider;
     }
 
+    @Override
     public WallConfig getConfig() {
         return config;
     }
 
+    @Override
     public void addViolation(Violation violation) {
         this.violations.add(violation);
     }
 
+    @Override
     public List<Violation> getViolations() {
         return violations;
     }

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

@@ -61,25 +61,40 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
 
     private final WallConfig      config;
     private final WallProvider    provider;
-    private final List<Violation> violations = new ArrayList<Violation>();
+    private final List<Violation> violations  = new ArrayList<Violation>();
+    private boolean               sqlModified = false;
 
     public OracleWallVisitor(WallProvider provider){
         this.config = provider.getConfig();
         this.provider = provider;
     }
+    
+    @Override
+    public boolean isSqlModified() {
+        return sqlModified;
+    }
 
+    @Override
+    public void setSqlModified(boolean sqlModified) {
+        this.sqlModified = sqlModified;
+    }
+
+    @Override
     public WallProvider getProvider() {
         return provider;
     }
 
+    @Override
     public WallConfig getConfig() {
         return config;
     }
 
+    @Override
     public void addViolation(Violation violation) {
         this.violations.add(violation);
     }
 
+    @Override
     public List<Violation> getViolations() {
         return violations;
     }

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

@@ -53,25 +53,40 @@ public class PGWallVisitor extends PGASTVisitorAdapter implements WallVisitor {
 
     private final WallConfig      config;
     private final WallProvider    provider;
-    private final List<Violation> violations = new ArrayList<Violation>();
+    private final List<Violation> violations  = new ArrayList<Violation>();
+    private boolean               sqlModified = false;
 
     public PGWallVisitor(WallProvider provider){
         this.config = provider.getConfig();
         this.provider = provider;
     }
+    
+    @Override
+    public boolean isSqlModified() {
+        return sqlModified;
+    }
 
+    @Override
+    public void setSqlModified(boolean sqlModified) {
+        this.sqlModified = sqlModified;
+    }
+
+    @Override
     public WallProvider getProvider() {
         return provider;
     }
 
+    @Override
     public WallConfig getConfig() {
         return config;
     }
 
+    @Override
     public void addViolation(Violation violation) {
         this.violations.add(violation);
     }
 
+    @Override
     public List<Violation> getViolations() {
         return violations;
     }

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

@@ -57,13 +57,25 @@ public class SQLServerWallVisitor extends SQLServerASTVisitorAdapter implements
 
     private final WallConfig      config;
     private final WallProvider    provider;
-    private final List<Violation> violations = new ArrayList<Violation>();
+    private final List<Violation> violations  = new ArrayList<Violation>();
+    private boolean               sqlModified = false;
 
     public SQLServerWallVisitor(WallProvider provider){
         this.config = provider.getConfig();
         this.provider = provider;
     }
+    
+    @Override
+    public boolean isSqlModified() {
+        return sqlModified;
+    }
 
+    @Override
+    public void setSqlModified(boolean sqlModified) {
+        this.sqlModified = sqlModified;
+    }
+
+    @Override
     public WallProvider getProvider() {
         return provider;
     }
@@ -73,6 +85,7 @@ public class SQLServerWallVisitor extends SQLServerASTVisitorAdapter implements
         return this.config;
     }
 
+    @Override
     public void addViolation(Violation violation) {
         this.violations.add(violation);
     }

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

@@ -406,11 +406,17 @@ public class WallVisitorUtils {
         }
 
         if (parent instanceof SQLDeleteStatement) {
-            ((SQLDeleteStatement) parent).setWhere(condition);
+            SQLDeleteStatement deleteStmt = (SQLDeleteStatement) parent;
+            deleteStmt.setWhere(condition);
+            visitor.setSqlModified(true);
         } else if (parent instanceof SQLUpdateStatement) {
-            ((SQLUpdateStatement) parent).setWhere(condition);
+            SQLUpdateStatement updateStmt = (SQLUpdateStatement) parent;
+            updateStmt.setWhere(condition);
+            visitor.setSqlModified(true);
         } else if (parent instanceof SQLSelectQueryBlock) {
-            ((SQLSelectQueryBlock) parent).setWhere(condition);
+            SQLSelectQueryBlock queryBlock = (SQLSelectQueryBlock) parent;
+            queryBlock.setWhere(condition);
+            visitor.setSqlModified(true);
         }
     }
 
@@ -422,10 +428,6 @@ public class WallVisitorUtils {
 
         SQLExpr condition = join.getCondition();
 
-        if (checkLeft) {
-
-        }
-
         SQLTableSource right = join.getRight();
         if (right instanceof SQLExprTableSource) {
             SQLExpr tableExpr = ((SQLExprTableSource) right).getExpr();
@@ -450,6 +452,7 @@ public class WallVisitorUtils {
 
         if (condition != join.getCondition()) {
             join.setCondition(condition);
+            visitor.setSqlModified(true);
         }
     }
 

src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest.java

@@ -51,13 +51,15 @@ public class MySqlParameterizedOutputVisitorTest extends TestCase {
 
         MySqlStatementParser parser = new MySqlStatementParser(sql);
         List<SQLStatement> statementList = parser.parseStatementList();
-        SQLStatement statemen = statementList.get(0);
+        SQLStatement stmt = statementList.get(0);
 
         Assert.assertEquals(1, statementList.size());
 
         StringBuilder out = new StringBuilder();
         MySqlParameterizedOutputVisitor visitor = new MySqlParameterizedOutputVisitor(out);
-        statemen.accept(visitor);
+        stmt.accept(visitor);
+        
+        Assert.assertTrue(visitor.getReplaceCount() > 0);
 
         Assert.assertEquals(expect, out.toString());
     }
@@ -66,13 +68,15 @@ public class MySqlParameterizedOutputVisitorTest extends TestCase {
 
         OracleStatementParser parser = new OracleStatementParser(sql);
         List<SQLStatement> statementList = parser.parseStatementList();
-        SQLStatement statemen = statementList.get(0);
+        SQLStatement stmt = statementList.get(0);
 
         Assert.assertEquals(1, statementList.size());
 
         StringBuilder out = new StringBuilder();
         OracleParameterizedOutputVisitor visitor = new OracleParameterizedOutputVisitor(out, false);
-        statemen.accept(visitor);
+        stmt.accept(visitor);
+        
+        Assert.assertTrue(visitor.getReplaceCount() > 0);
 
         Assert.assertEquals(expect, out.toString());
     }

src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest_1.java

@@ -36,4 +36,56 @@ public class MySqlParameterizedOutputVisitorTest_1 extends TestCase {
         Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.HSQL), sql);
     }
 
+    public void test_1() throws Exception {
+        String sql = "SELECT * FROM T WHERE ID = ?";
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.ORACLE), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DB2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.SQL_SERVER), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.POSTGRESQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.H2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DERBY), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.HSQL), sql);
+    }
+    
+    public void test_2() throws Exception {
+        String sql = "SELECT * FROM T WHERE ID = ? AND Name = ?";
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.ORACLE), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DB2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.SQL_SERVER), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.POSTGRESQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.H2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DERBY), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.HSQL), sql);
+    }
+    
+    public void test_3() throws Exception {
+        String sql = "SELECT * FROM T WHERE ID IS NULL";
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.ORACLE), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DB2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.SQL_SERVER), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.POSTGRESQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.H2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DERBY), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.HSQL), sql);
+    }
+    
+    public void test_4() throws Exception {
+        String sql = "INSERT INTO T (FID, FNAME) VALUES(?, ?)";
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.ORACLE), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DB2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.SQL_SERVER), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.POSTGRESQL), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.H2), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.DERBY), sql);
+        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.HSQL), sql);
+    }
+    
+    public void test_mysql() throws Exception {
+        String sql = "INSERT INTO T (FID, FNAME) VALUES(?, ?), (?, ?)";
+        Assert.assertNotSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+    }
 }