Merge pull request #687 from yakolee/master

improve sqlwall and parser

src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java

@@ -110,6 +110,8 @@ public class MySqlLexer extends Lexer {
         if (commentHandler != null && commentHandler.handle(lastToken, stringVal)) {
             return;
         }
+        
+        endOfComment = isEOF();
 
         if (!isAllowComment() && (isEOF() || !isSafeComment(stringVal))) {
             throw new NotAllowCommentException();
@@ -448,6 +450,8 @@ public class MySqlLexer extends Lexer {
                 return;
             }
 
+            endOfComment = isEOF();
+            
             if (!isHint && !isAllowComment() && (isEOF() || !isSafeComment(stringVal))) {
                 throw new NotAllowCommentException();
             }
@@ -489,6 +493,8 @@ public class MySqlLexer extends Lexer {
                 return;
             }
 
+            endOfComment = isEOF();
+            
             if (!isAllowComment() && (isEOF() || !isSafeComment(stringVal))) {
                 throw new NotAllowCommentException();
             }

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

@@ -1752,11 +1752,6 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV
 
     @Override
     public boolean visit(MySqlSetTransactionIsolationLevelStatement x) {
-        return false;
-    }
-
-    @Override
-    public void endVisit(MySqlSetTransactionIsolationLevelStatement x) {
         if (x.getGlobal() == null) {
             print("SET TRANSACTION ISOLATION LEVEL ");
         } else if (x.getGlobal().booleanValue()) {
@@ -1765,6 +1760,12 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV
             print("SET SESSION TRANSACTION ISOLATION LEVEL ");
         }
         print(x.getLevel());
+        return false;
+    }
+
+    @Override
+    public void endVisit(MySqlSetTransactionIsolationLevelStatement x) {
+        
     }
 
     @Override

src/main/java/com/alibaba/druid/sql/dialect/odps/parser/OdpsLexer.java

@@ -137,6 +137,7 @@ public class OdpsLexer extends Lexer {
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            endOfComment = isEOF();
             return;
         }
     }

src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java

@@ -274,6 +274,7 @@ public class OracleLexer extends Lexer {
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            endOfComment = isEOF();
             return;
         }
     }

src/main/java/com/alibaba/druid/sql/dialect/sqlserver/parser/SQLServerLexer.java

@@ -149,6 +149,7 @@ public class SQLServerLexer extends Lexer {
             stringVal = subString(mark + 1, bufPos);
             token = Token.LINE_COMMENT;
             hasComment = true;
+            endOfComment = isEOF();
             return;
         }
     }

src/main/java/com/alibaba/druid/sql/parser/Lexer.java

@@ -64,13 +64,15 @@ public class Lexer {
     /*
      * anti sql injection
      */
-    private boolean        allowComment = true;
+    private boolean          allowComment = true;
 
-    private int            varIndex     = -1;
+    private int              varIndex     = -1;
 
     protected CommentHandler commentHandler;
-    
-    protected boolean        hasComment = false;
+
+    protected boolean        hasComment   = false;
+
+    protected boolean        endOfComment = false;
 
     public Lexer(String input){
         this(input, true);
@@ -1058,4 +1060,9 @@ public class Lexer {
         pos = text.length();
         this.token = Token.EOF;
     }
+
+    public boolean isEndOfComment() {
+        return endOfComment;
+    }
+
 }

src/main/java/com/alibaba/druid/sql/parser/SQLParser.java

@@ -184,7 +184,10 @@ public class SQLParser {
         } else if (lexer.token() == Token.USER) {
             alias = lexer.stringVal();
             lexer.nextToken();
-        }
+        } else if (lexer.token() == Token.END) {
+            alias = lexer.stringVal();
+            lexer.nextToken();
+        } 
 
         switch (lexer.token()) {
             case KEY:

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

@@ -32,7 +32,6 @@ import com.alibaba.druid.sql.ast.expr.SQLNCharExpr;
 import com.alibaba.druid.sql.ast.expr.SQLNullExpr;
 import com.alibaba.druid.sql.ast.expr.SQLNumberExpr;
 import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
-import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
 import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
 import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem;
 import com.alibaba.druid.sql.dialect.db2.visitor.DB2ParameterizedOutputVisitor;
@@ -165,7 +164,7 @@ public class ParameterizedOutputVisitorUtils {
         if (parent instanceof SQLDataType //
             || parent instanceof SQLColumnDefinition //
             || parent instanceof SQLServerTop //
-            || parent instanceof SQLAssignItem //
+            //|| parent instanceof SQLAssignItem //
             || parent instanceof SQLSelectOrderByItem //
         ) {
             return false;

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

@@ -604,6 +604,7 @@ public abstract class WallProvider {
         final List<Violation> violations = new ArrayList<Violation>();
         List<SQLStatement> statementList = new ArrayList<SQLStatement>();
         boolean syntaxError = false;
+        boolean endOfComment = false;
         try {
             SQLStatementParser parser = createParser(sql);
             parser.getLexer().setCommentHandler(WallCommentHandler.instance);
@@ -619,10 +620,11 @@ public abstract class WallProvider {
             parser.parseStatementList(statementList);
 
             final Token lastToken = parser.getLexer().token();
-            if (lastToken != Token.EOF) {
+            if (lastToken != Token.EOF && config.isStrictSyntaxCheck()) {
                 violations.add(new IllegalSQLObjectViolation(ErrorCode.SYNTAX_ERROR, "not terminal sql, token "
                                                                                      + lastToken, sql));
             }
+            endOfComment = parser.getLexer().isEndOfComment();
         } catch (NotAllowCommentException e) {
             violations.add(new IllegalSQLObjectViolation(ErrorCode.COMMENT_STATEMENT_NOT_ALLOW, "comment not allow", sql));
             incrementCommentDeniedCount();
@@ -643,6 +645,9 @@ public abstract class WallProvider {
         }
 
         WallVisitor visitor = createWallVisitor();
+        if (!endOfComment) {
+            visitor.getConfig().setSelectUnionCheck(false);
+        }
 
         if (statementList.size() > 0) {
             boolean lastIsHint = false;

src/test/java/com/alibaba/druid/bvt/filter/wall/MustParameterizedTest3.java

@@ -39,7 +39,7 @@ public class MustParameterizedTest3 extends TestCase {
         Assert.assertTrue(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ?", config));
         Assert.assertFalse(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = 7", config));
         Assert.assertTrue(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ? union select * from t", config));
-        Assert.assertFalse(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ? union select 1, 2, 3", config));
+        Assert.assertFalse(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ? union select 1, 2, 3 --", config));
         Assert.assertFalse(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ? union select * from t fid = fid", config));
         Assert.assertFalse(WallUtils.isValidateMySql("select 1, 2, 3 from t where id  = ? union select * from t fid > 5", config));
     }

src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest.java

@@ -44,7 +44,7 @@ public class MySqlWallTest extends TestCase {
         Assert.assertFalse(WallUtils.isValidateMySql("select * FROM T WHERE id = 1 AND select 0x3C3F706870206576616C28245F504F53545B2763275D293F3E into outfile '\\www\\edu\\1.php'"));
         Assert.assertTrue(WallUtils.isValidateMySql("select 'outfile'"));
 
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where c1=1 union select 1, 2"));
+        //Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where c1=1 union select 1, 2"));
         
         Assert.assertTrue(WallUtils.isValidateMySql("select c1 from t where 1=1 or id =1"));
         Assert.assertFalse(WallUtils.isValidateMySql("select c1 from t where id =1 or 1=1"));

src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest134.java

@@ -26,9 +26,9 @@ public class MySqlWallTest134 extends TestCase {
 
     public void test_false() throws Exception {
         WallProvider provider = new MySqlWallProvider();
-
+        provider.getConfig().setSelectUnionCheck(true);
         Assert.assertTrue(provider.checkValid("SELECT 1 FROM dual WHERE 1 = '1'''''''''''''UNION SELECT '2';"));
-        Assert.assertFalse(provider.checkValid("SELECT 1 FROM t WHERE 1 = '1'''''''''''''UNION SELECT '2';"));
+        Assert.assertFalse(provider.checkValid("SELECT 1 FROM t WHERE 1 = '1'''''''''''''UNION SELECT '2' --;")); //end of comment --> false
     }
 
 }

src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest62.java

@@ -54,12 +54,13 @@ public class MySqlWallTest62 extends TestCase {
     public void test_false() throws Exception {
         WallProvider provider = new MySqlWallProvider();
         provider.getConfig().setSchemaCheck(true);
+        provider.getConfig().setSelectUnionCheck(true);
         String sql = "SELECT 1, 2, 3" + //
                      " UNION ALL SELECT  a  from tt where c=1" + //
-                     " UNION ALL SELECT 2 FROM dual ";
+                     " UNION ALL SELECT 2 FROM dual --";
         Assert.assertFalse(provider.checkValid(sql));
         
-        sql = "SELECT a from t where c=1 UNION ALL SELECT 2 FROM dual ";
+        sql = "SELECT a from t where c=1 UNION ALL SELECT 2 FROM dual --";
         Assert.assertFalse(provider.checkValid(sql));
     }
 

src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest63.java

@@ -34,12 +34,17 @@ public class MySqlWallTest63 extends TestCase {
     public void test_true() throws Exception {
         WallProvider provider = new MySqlWallProvider();
         provider.getConfig().setSchemaCheck(true);
+        provider.getConfig().setSelectUnionCheck(true);
+        provider.setBlackListEnable(false);
+        provider.setWhiteListEnable(false);
+
+        Assert.assertTrue(provider.checkValid(//
+        "SELECT FID, FNAME FROM T WHERE C=1 UNION SELECT 1, 'AAA'"));
 
         Assert.assertFalse(provider.checkValid(//
-        "SELECT FID, FNAME FROM T WHERE C=1 UNION SELECT 1, 'AAA'"));
+        "SELECT FID, FNAME FROM T WHERE C=1 UNION SELECT 1, 'AAA' --"));
 
         Assert.assertEquals(1, provider.getTableStats().size());
     }
 
-
 }

src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest_union.java

@@ -28,22 +28,20 @@ public class MySqlWallTest_union extends TestCase {
         WallConfig config = new WallConfig();
         config.setSelectUnionCheck(true);
 
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select 1, 2", config));
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select 1, 2", config)); // not end of comment
+        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select 1, 2 --", config));
 
-        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t union select 1, 2", config));
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t union select 1, 2", config)); // no where
 
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select null, '1', 2",
-                                                     config));
+        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select null, '1', 2 --", config));
 
-        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select c1, c2", config));
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select c1, c2", config)); //union select item is not const
 
         Assert.assertTrue(WallUtils.isValidateMySql("SELECT typeid, typename FROM (SELECT typeid, typename FROM materialtype UNION ALL SELECT ? AS typeid, ? AS typename) a ORDER BY typeid",
-                                                    config));
+                                                    config)); // union select item has alias
 
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from (select 1 as f1, 2 as f2) t union select 'u1', 'u2'",
-                                                     config));
-        
-        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select 'u1' as u1, 'u2' as u2",
-                                                     config));
+        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from (select 1 as f1, 2 as f2) t union select 'u1', 'u2' --", config)); // from is subQuery
+
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where id=1 union select 'u1' as u1, 'u2' as u2",  config)); // union select item has alias
     }
 }

src/test/java/com/alibaba/druid/bvt/filter/wall/OracleWallTest.java

@@ -19,12 +19,16 @@ import junit.framework.TestCase;
 
 import org.junit.Assert;
 
+import com.alibaba.druid.wall.WallConfig;
 import com.alibaba.druid.wall.WallUtils;
 
 public class OracleWallTest extends TestCase {
 
     public void testWall() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where c=1 union select 1, 2"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateOracle("select f1, f2 from t where c=1 union select 1, 2",config));
+        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where c=1 union select 1, 2 --",config));
         
         Assert.assertFalse(WallUtils.isValidateOracle("SELECT * FROM T UNION select * from TAB"));
         Assert.assertFalse(WallUtils.isValidateOracle("SELECT * FROM T UNION select * from ALL_TABLES where (1=1 or (1+1)=2) and (4=8 or 1=1)"));

src/test/java/com/alibaba/druid/bvt/filter/wall/WallUnionTest.java

@@ -19,20 +19,27 @@ import junit.framework.TestCase;
 
 import org.junit.Assert;
 
+import com.alibaba.druid.wall.WallConfig;
 import com.alibaba.druid.wall.WallUtils;
 
 /**
  * 这个场景，被攻击者用于测试当前SQL拥有多少字段
+ * 
  * @author wenshao
- *
  */
 public class WallUnionTest extends TestCase {
 
     public void testMySql() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where f1=1 union select 1, 2"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where f1=1 union select 1, 2", config));
+        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where f1=1 union select 1, 2 --", config));
     }
 
     public void testOracle() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where f1=1 union select 1, 2"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateOracle("select f1, f2 from t where f1=1 union select 1, 2", config));
+        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where f1=1 union select 1, 2 --", config));
     }
 }

src/test/java/com/alibaba/druid/bvt/filter/wall/WallUnionTest2.java

@@ -19,20 +19,31 @@ import junit.framework.TestCase;
 
 import org.junit.Assert;
 
+import com.alibaba.druid.wall.WallConfig;
 import com.alibaba.druid.wall.WallUtils;
 
 /**
  * 这个场景，被攻击者用于测试当前SQL拥有多少字段
+ * 
  * @author wenshao
- *
  */
 public class WallUnionTest2 extends TestCase {
 
     public void testMySql() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateMySql("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1",
+                                                    config));
+        Assert.assertFalse(WallUtils.isValidateMySql("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1 --",
+                                                     config));
     }
 
     public void testOracle() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateOracle("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1",
+                                                     config));
+        Assert.assertFalse(WallUtils.isValidateOracle("select f1, f2 from t where f1 = 1 union select 1, 2 where 1 = 1 --",
+                                                      config));
     }
 }

src/test/java/com/alibaba/druid/bvt/filter/wall/WallUnionTest4.java

@@ -19,20 +19,31 @@ import junit.framework.TestCase;
 
 import org.junit.Assert;
 
+import com.alibaba.druid.wall.WallConfig;
 import com.alibaba.druid.wall.WallUtils;
 
 /**
  * 这个场景，被攻击者用于测试当前SQL拥有多少字段
+ * 
  * @author wenshao
- *
  */
 public class WallUnionTest4 extends TestCase {
 
     public void testMySql() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateMySql("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X)"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateMySql("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X)",
+                                                    config));
+        Assert.assertFalse(WallUtils.isValidateMySql("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X) -- ",
+                                                     config));
     }
-    
+
     public void testORACLE() throws Exception {
-        Assert.assertFalse(WallUtils.isValidateOracle("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X)"));
+        WallConfig config = new WallConfig();
+        config.setSelectUnionCheck(true);
+        Assert.assertTrue(WallUtils.isValidateOracle("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X)",
+                                                     config));
+        Assert.assertFalse(WallUtils.isValidateOracle("SELECT id, product FROM T1 t where id=1 UNION (SELECT * FROM (SELECT 1,'x') X) -- ",
+                                                      config));
     }
 }

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

@@ -26,11 +26,13 @@ public class MySqlParameterizedOutputVisitorTest_6 extends TestCase {
 
     public void test_0() throws Exception {
         String sql = "SET autocommit=1";
-        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        String paramSql = "SET @@autocommit = ?";
+        Assert.assertEquals(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), paramSql);
     }
-    
+
     public void test_1() throws Exception {
         String sql = "SET GLOBAL sort_buffer_size=1000000, SESSION sort_buffer_size=1000000;";
-        Assert.assertSame(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), sql);
+        String paramSql = "SET @@global.sort_buffer_size = ?, @@sort_buffer_size = ?";
+        Assert.assertEquals(ParameterizedOutputVisitorUtils.parameterize(sql, JdbcConstants.MYSQL), paramSql);
     }
 }