From 3436764820a9666cab665086dc43fe895fa7c034 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B8=A9=E5=B0=91?= <szujobs@hotmail.com>
Date: Mon, 19 Mar 2012 07:11:46 +0800
Subject: [PATCH] bug fixed for sql wall

---
 .../druid/wall/spi/SQLServerWallVisitor.java       | 14 ++++++++++++++
 .../alibaba/druid/wall/spi/WallVisitorUtils.java   |  5 ++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java
index 6553c1ab1a..39502db3bb 100644
--- a/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java
+++ b/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java
@@ -11,6 +11,7 @@ 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.SQLDeleteStatement;
 import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
@@ -180,6 +181,19 @@ public class SQLServerWallVisitor extends SQLServerASTVisitorAdapter implements
         return true;
     }
     
+    public boolean visit(SQLVariantRefExpr x) {
+        String varName = x.getName();
+        if (varName == null) {
+            return false;
+        }
+
+        if (config.isVariantCheck() && varName.startsWith("@@")) {
+            violations.add(new IllegalSQLObjectViolation(toSQL(x)));
+        }
+
+        return false;
+    }
+    
     @Override
     public boolean visit(SQLServerObjectReferenceExpr x) {
         if (x.getSchema() != null && config.isPermitSchema(x.getSchema())) {
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 e494aecaf0..e2a53c8932 100644
--- a/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java
+++ b/src/main/java/com/alibaba/druid/wall/spi/WallVisitorUtils.java
@@ -18,7 +18,6 @@ import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
-import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
 import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr;
 import com.alibaba.druid.sql.ast.expr.SQLNCharExpr;
@@ -450,6 +449,10 @@ public class WallVisitorUtils {
             if (visitor.getConfig().isPermitSchema(owner)) {
                 addViolation(visitor, x);
             }
+            
+            if (visitor.getConfig().isPermitObjects(owner)) {
+                addViolation(visitor, x);
+            }
         }
         
         //if (ownerExpr instanceof SQLPropertyExpr) {