From 5f80737ead01c9a0f1c035f06eb575b5de66c222 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B8=A9=E5=B0=91?= <szujobs@hotmail.com>
Date: Thu, 15 Mar 2012 13:11:53 +0800
Subject: [PATCH] add wall support

---
 .../alibaba/druid/filter/wall/WallFilter.java |  1 +
 .../com/alibaba/druid/sql/parser/Lexer.java   | 27 +++++++++++++++----
 .../sql/parser/NotAllowCommentException.java  | 19 +++++++++++++
 3 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 src/main/java/com/alibaba/druid/sql/parser/NotAllowCommentException.java

diff --git a/src/main/java/com/alibaba/druid/filter/wall/WallFilter.java b/src/main/java/com/alibaba/druid/filter/wall/WallFilter.java
index 975909bbc3..0711d01f55 100644
--- a/src/main/java/com/alibaba/druid/filter/wall/WallFilter.java
+++ b/src/main/java/com/alibaba/druid/filter/wall/WallFilter.java
@@ -97,6 +97,7 @@ public class WallFilter extends FilterAdapter {
         String dbType = connection.getDirectDataSource().getDbType();
 
         SQLStatementParser parser = createParser(dbType, sql);
+        parser.getLexer().setAllowComment(false); // permit comment
 
         List<SQLStatement> statementList = parser.parseStatementList();
         if (statementList.size() > 1) {
diff --git a/src/main/java/com/alibaba/druid/sql/parser/Lexer.java b/src/main/java/com/alibaba/druid/sql/parser/Lexer.java
index 3af57c7386..504dfc55e7 100644
--- a/src/main/java/com/alibaba/druid/sql/parser/Lexer.java
+++ b/src/main/java/com/alibaba/druid/sql/parser/Lexer.java
@@ -53,7 +53,7 @@ public class Lexer {
     protected int                              tokenPos;
 
     /** A character buffer for literals. */
-    protected final static ThreadLocal<char[]> sbufRef     = new ThreadLocal<char[]>();
+    protected final static ThreadLocal<char[]> sbufRef      = new ThreadLocal<char[]>();
     protected char[]                           sbuf;
 
     /** string point as size */
@@ -62,25 +62,38 @@ public class Lexer {
     /** string point as offset */
     protected int                              np;
 
-    protected SymbolTable                      symbolTable = new SymbolTable();
+    protected SymbolTable                      symbolTable  = new SymbolTable();
 
     /**
      * The token, set by nextToken().
      */
     protected Token                            token;
 
-    protected Keywords                         keywods     = Keywords.DEFAULT_KEYWORDS;
+    protected Keywords                         keywods      = Keywords.DEFAULT_KEYWORDS;
 
     protected String                           stringVal;
 
-    protected boolean                          skipComment = true;
+    protected boolean                          skipComment  = true;
 
-    private SavePoint                          savePoint   = null;
+    private SavePoint                          savePoint    = null;
+
+    /*
+     * anti sql injection
+     */
+    private boolean                            allowComment = true;
 
     public Lexer(String input){
         this(input, true);
     }
 
+    public boolean isAllowComment() {
+        return allowComment;
+    }
+
+    public void setAllowComment(boolean allowComment) {
+        this.allowComment = allowComment;
+    }
+
     private static class SavePoint {
 
         int   bp;
@@ -575,6 +588,10 @@ public class Lexer {
     }
 
     public void scanComment() {
+        if (!allowComment) {
+            throw new NotAllowCommentException();
+        }
+
         if (ch != '/') {
             throw new IllegalStateException();
         }
diff --git a/src/main/java/com/alibaba/druid/sql/parser/NotAllowCommentException.java b/src/main/java/com/alibaba/druid/sql/parser/NotAllowCommentException.java
new file mode 100644
index 0000000000..e9702a995f
--- /dev/null
+++ b/src/main/java/com/alibaba/druid/sql/parser/NotAllowCommentException.java
@@ -0,0 +1,19 @@
+package com.alibaba.druid.sql.parser;
+
+public class NotAllowCommentException extends ParserException {
+
+    private static final long serialVersionUID = 1L;
+
+    public NotAllowCommentException(){
+        super();
+    }
+
+    public NotAllowCommentException(String message, Throwable e){
+        super(message, e);
+    }
+
+    public NotAllowCommentException(String message){
+        super(message);
+    }
+
+}