diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSetStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSetStatement.java index 2f7074c2ec..ab6adab03b 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSetStatement.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSetStatement.java @@ -18,6 +18,7 @@ package com.alibaba.druid.sql.ast.statement; import java.util.ArrayList; import java.util.List; +import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLStatementImpl; import com.alibaba.druid.sql.visitor.SQLASTVisitor; @@ -25,6 +26,8 @@ import com.alibaba.druid.sql.visitor.SQLASTVisitor; public class SQLSetStatement extends SQLStatementImpl { private List items = new ArrayList(); + + private List hints; public SQLSetStatement(){ } @@ -41,10 +44,19 @@ public class SQLSetStatement extends SQLStatementImpl { this.items = items; } + public List getHints() { + return hints; + } + + public void setHints(List hints) { + this.hints = hints; + } + @Override protected void accept0(SQLASTVisitor visitor) { if (visitor.visit(this)) { acceptChild(visitor, this.items); + acceptChild(visitor, this.hints); } visitor.endVisit(this); } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlLockTableStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlLockTableStatement.java index eefa25e75b..2bc8510805 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlLockTableStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlLockTableStatement.java @@ -15,6 +15,9 @@ */ package com.alibaba.druid.sql.dialect.mysql.ast.statement; +import java.util.List; + +import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor; @@ -24,6 +27,8 @@ public class MySqlLockTableStatement extends MySqlStatementImpl { private SQLExprTableSource tableSource; private LockType lockType; + + private List hints; public SQLExprTableSource getTableSource() { return tableSource; @@ -64,4 +69,12 @@ public class MySqlLockTableStatement extends MySqlStatementImpl { this.name = name; } } + + public List getHints() { + return hints; + } + + public void setHints(List hints) { + this.hints = hints; + } } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlShowIndexesStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlShowIndexesStatement.java index a6febc2969..670f2079fb 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlShowIndexesStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlShowIndexesStatement.java @@ -15,6 +15,9 @@ */ package com.alibaba.druid.sql.dialect.mysql.ast.statement; +import java.util.List; + +import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; @@ -22,8 +25,9 @@ import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor; public class MySqlShowIndexesStatement extends MySqlStatementImpl implements MySqlShowStatement { - private SQLName table; - private SQLName database; + private SQLName table; + private SQLName database; + private List hints; public SQLName getTable() { return table; @@ -54,4 +58,12 @@ public class MySqlShowIndexesStatement extends MySqlStatementImpl implements MyS } visitor.endVisit(this); } + + public List getHints() { + return hints; + } + + public void setHints(List hints) { + this.hints = hints; + } } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlStartTransactionStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlStartTransactionStatement.java index fc1431ece4..4edaf85dd8 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlStartTransactionStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlStartTransactionStatement.java @@ -15,14 +15,19 @@ */ package com.alibaba.druid.sql.dialect.mysql.ast.statement; +import java.util.List; + +import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor; public class MySqlStartTransactionStatement extends MySqlStatementImpl { - private boolean consistentSnapshot = false; + private boolean consistentSnapshot = false; - private boolean begin = false; - private boolean work = false; + private boolean begin = false; + private boolean work = false; + + private List hints; public boolean isConsistentSnapshot() { return consistentSnapshot; @@ -53,4 +58,12 @@ public class MySqlStartTransactionStatement extends MySqlStatementImpl { visitor.endVisit(this); } + + public List getHints() { + return hints; + } + + public void setHints(List hints) { + this.hints = hints; + } } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java index 6674793ac5..db27a022b2 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java @@ -654,6 +654,9 @@ public class MySqlStatementParser extends SQLStatementParser { stmt.setLockType(LockType.LOW_PRIORITY_WRITE); } + if(lexer.token() == Token.HINT) { + stmt.setHints(this.exprParser.parseHints()); + } statementList.add(stmt); return true; } @@ -1101,6 +1104,10 @@ public class MySqlStatementParser extends SQLStatementParser { stmt.setDatabase(database); } } + + if (lexer.token() == Token.HINT) { + stmt.setHints(this.exprParser.parseHints()); + } return stmt; } @@ -1492,6 +1499,10 @@ public class MySqlStatementParser extends SQLStatementParser { stmt.setWork(true); } } + + if(lexer.token() == Token.HINT) { + stmt.setHints(this.exprParser.parseHints()); + } return stmt; } @@ -2123,6 +2134,10 @@ public class MySqlStatementParser extends SQLStatementParser { SQLVariantRefExpr varRef = (SQLVariantRefExpr) stmt.getItems().get(0).getTarget(); varRef.setGlobal(true); } + + if(lexer.token() == Token.HINT) { + stmt.setHints(this.exprParser.parseHints()); + } return stmt; } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java index ca66aaeaaa..0fc377d48e 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java @@ -1118,6 +1118,14 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV if (x.isConsistentSnapshot()) { print(" WITH CONSISTENT SNAPSHOT"); } + + List hints = x.getHints(); + if (hints != null && !hints.isEmpty()) { + print(" "); + for (SQLCommentHint hint : hints) { + hint.accept(this); + } + } if (x.isBegin()) { print(" BEGIN"); @@ -2130,6 +2138,14 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV } x.getTable().accept(this); } + + List hints = x.getHints(); + if (hints != null && !hints.isEmpty()) { + print(" "); + for (SQLCommentHint hint : hints) { + hint.accept(this); + } + } return false; } @@ -2664,6 +2680,14 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV print(' '); print(x.getLockType().name); } + + List hints = x.getHints(); + if (hints != null && !hints.isEmpty()) { + print(" "); + for (SQLCommentHint hint : hints) { + hint.accept(this); + } + } return false; } diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java index c58e34e3f5..535a48c133 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java @@ -1148,6 +1148,14 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Printab public boolean visit(SQLSetStatement x) { print("SET "); printAndAccept(x.getItems(), ", "); + + List hints = x.getHints(); + if (hints != null && !hints.isEmpty()) { + print(" "); + for (SQLCommentHint hint : hints) { + hint.accept(this); + } + } return false; } diff --git a/src/main/java/com/alibaba/druid/wall/WallConfig.java b/src/main/java/com/alibaba/druid/wall/WallConfig.java index 1d67b4f236..c6ddc04353 100644 --- a/src/main/java/com/alibaba/druid/wall/WallConfig.java +++ b/src/main/java/com/alibaba/druid/wall/WallConfig.java @@ -42,6 +42,7 @@ public class WallConfig implements WallConfigMBean { private boolean renameTableAllow = true; private boolean hintAllow = true; private boolean lockTableAllow = true; + private boolean startTransactionAllow = true; private boolean conditionAndAlwayTrueAllow = false; private boolean conditionAndAlwayFalseAllow = false; @@ -761,4 +762,12 @@ public class WallConfig implements WallConfigMBean { this.lockTableAllow = lockTableAllow; } + public boolean isStartTransactionAllow() { + return startTransactionAllow; + } + + public void setStartTransactionAllow(boolean startTransactionAllow) { + this.startTransactionAllow = startTransactionAllow; + } + } diff --git a/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java b/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java index 929c14973b..b5f9cb496f 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java +++ b/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java @@ -107,6 +107,7 @@ import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSetCharSetStatemen import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSetNamesStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlShowGrantsStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlShowStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlStartTransactionStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; import com.alibaba.druid.sql.dialect.oracle.ast.stmt.OracleCreateSequenceStatement; @@ -2416,6 +2417,10 @@ public class WallVisitorUtils { allow = config.isLockTableAllow(); denyMessage = "lock table not allow"; errorCode = ErrorCode.LOCK_TABLE_NOT_ALLOW; + } else if (x instanceof MySqlStartTransactionStatement) { + allow = config.isStartTransactionAllow(); + denyMessage = "start transaction not allow"; + errorCode = ErrorCode.START_TRANSACTION_NOT_ALLOW; } else { allow = config.isNoneBaseStatementAllow(); errorCode = ErrorCode.NONE_BASE_STATEMENT_NOT_ALLOW; diff --git a/src/main/java/com/alibaba/druid/wall/violation/ErrorCode.java b/src/main/java/com/alibaba/druid/wall/violation/ErrorCode.java index f6d0582e61..3dfef0bd83 100644 --- a/src/main/java/com/alibaba/druid/wall/violation/ErrorCode.java +++ b/src/main/java/com/alibaba/druid/wall/violation/ErrorCode.java @@ -33,6 +33,7 @@ public interface ErrorCode { public final static int CALL_NOT_ALLOW = 1300; public final static int COMMIT_NOT_ALLOW = 1301; public final static int ROLLBACK_NOT_ALLOW = 1302; + public final static int START_TRANSACTION_NOT_ALLOW = 1303; public final static int SET_NOT_ALLOW = 1200; public final static int DESC_NOT_ALLOW = 1201; @@ -47,7 +48,7 @@ public interface ErrorCode { public final static int DROP_TABLE_NOT_ALLOW = 1103; public final static int COMMENT_STATEMENT_NOT_ALLOW = 1104; public final static int RENAME_TABLE_NOT_ALLOW = 1105; - public final static int LOCK_TABLE_NOT_ALLOW = 1105; + public final static int LOCK_TABLE_NOT_ALLOW = 1106; public final static int LIMIT_ZERO = 2200; public final static int MULTI_STATEMENT = 2201; diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest_hint.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest_hint.java index fd0762bf73..96e067daf8 100644 --- a/src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest_hint.java +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/MySqlWallTest_hint.java @@ -17,8 +17,11 @@ public class MySqlWallTest_hint extends TestCase { } public void test_false_1() throws Exception { + WallConfig config = new WallConfig(); + config.setHintAllow(true); + config.setMultiStatementAllow(true); String sql = "select * from person where id = '3'/**/union select 0,1,v from (select 1,2,user/*!() as v*/) a where '1'<>''"; - Assert.assertFalse(WallUtils.isValidateMySql(sql)); // + Assert.assertFalse(WallUtils.isValidateMySql(sql, config)); // } public void test_true() throws Exception { @@ -77,4 +80,19 @@ public class MySqlWallTest_hint extends TestCase { Assert.assertTrue(WallUtils.isValidateMySql(sql, config)); // } + public void test_true_6() throws Exception { + String sql = "START TRANSACTION /*!40100 WITH CONSISTENT SNAPSHOT */"; + Assert.assertTrue(WallUtils.isValidateMySql(sql)); // + } + + public void test_true_7() throws Exception { + String sql = "LOCK TABLES `m_rpt_adgroupeffect` READ /*!32311 LOCAL */"; + Assert.assertTrue(WallUtils.isValidateMySql(sql)); // + } + + public void test_true_8() throws Exception { + String sql = "SET SQL_QUOTE_SHOW_CREATE=1/*!40102 ,SQL_MODE=concat(@@sql_mode, _utf8 ',NO_KEY_OPTIONS,NO_TABLE_OPTIONS,NO_FIELD_OPTIONS') */"; + Assert.assertTrue(WallUtils.isValidateMySql(sql)); // + } + }