improve wall filter

2023-03-04 19:17:16 +00:00 · 2013-07-30 22:37:11 +08:00
parent 63574c4702
commit 30fcdc8a4e
4 changed files with 216 additions and 79 deletions
@@ -20,7 +20,6 @@ import java.sql.SQLException;
 import java.sql.Wrapper;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;

 import com.alibaba.druid.filter.FilterAdapter;
@@ -102,7 +101,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            if (config == null) {
                config = new WallConfig(DB2WallProvider.DEFAULT_CONFIG_DIR);
            }
-            
+
            provider = new DB2WallProvider(config);
        } else {
            throw new IllegalStateException("dbType not support : " + dbType + ", url " + dataSource.getUrl());
@@ -175,7 +174,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            WallContext.clearContext();
        }
    }
-    
+
    @Override
    public void preparedStatement_addBatch(FilterChain chain, PreparedStatementProxy statement) throws SQLException {
        chain.preparedStatement_addBatch(statement);
@@ -413,7 +412,11 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            for (int i = 0; i < updateCounts.length; ++i) {
                updateCount += updateCounts[i];
            }
-            statExecuteUpdate(sqlStat, updateCount);
+
+            if (sqlStat != null) {
+                provider.addUpdateCount(sqlStat, updateCount);
+            }
+
            return updateCounts;
        } finally {
            WallContext.clearContext();
@@ -504,7 +507,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
        if (!firstResult) {
            WallSqlStat sqlStat = (WallSqlStat) statement.getAttribute(ATTR_SQL_STAT);
            int updateCount = statement.getUpdateCount();
-            statExecuteUpdate(sqlStat, updateCount);
+            if (sqlStat != null) {
+                provider.addUpdateCount(sqlStat, updateCount);
+            }
        }

        return firstResult;
@@ -520,7 +525,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
    public int preparedStatement_executeUpdate(FilterChain chain, PreparedStatementProxy statement) throws SQLException {
        int updateCount = chain.preparedStatement_executeUpdate(statement);
        WallSqlStat sqlStat = (WallSqlStat) statement.getAttribute(ATTR_SQL_STAT);
-        statExecuteUpdate(sqlStat, updateCount);
+        if (sqlStat != null) {
+            provider.addUpdateCount(sqlStat, updateCount);
+        }
        return updateCount;
    }

@@ -549,37 +556,8 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            return;
        }

-        if (updateCount > 0) {
-            statExecuteUpdate(sqlStat, updateCount);
-        }
-    }
-
-    private void statExecuteUpdate(WallSqlStat sqlStat, int updateCount) {
-        if (sqlStat == null) {
-            return;
-        }
-
-        Map<String, WallSqlTableStat> sqlTableStats = sqlStat.getTableStats();
-        if (sqlTableStats == null) {
-            return;
-        }
-
-        for (Map.Entry<String, WallSqlTableStat> entry : sqlTableStats.entrySet()) {
-            String tableName = entry.getKey();
-            WallTableStat tableStat = provider.getTableStat(tableName);
-            if (tableStat == null) {
-                continue;
-            }
-
-            WallSqlTableStat sqlTableStat = entry.getValue();
-
-            if (sqlTableStat.getDeleteCount() > 0) {
-                tableStat.addDeleteDataCount(updateCount);
-            } else if (sqlTableStat.getUpdateCount() > 0) {
-                tableStat.addUpdateDataCount(updateCount);
-            } else if (sqlTableStat.getInsertCount() > 0) {
-                tableStat.addInsertDataCount(updateCount);
-            }
+        if (updateCount > 0 && sqlStat != null) {
+            provider.addUpdateCount(sqlStat, updateCount);
        }
    }

@@ -604,7 +582,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            }
        }

-            //TODO
+        // TODO
        return sql;
    }

@@ -664,24 +642,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
            return;
        }

-        Map<String, WallSqlTableStat> sqlTableStats = sqlStat.getTableStats();
-        if (sqlTableStats == null) {
-            return;
-        }
-
-        for (Map.Entry<String, WallSqlTableStat> entry : sqlTableStats.entrySet()) {
-            String tableName = entry.getKey();
-            WallTableStat tableStat = provider.getTableStat(tableName);
-            if (tableStat == null) {
-                continue;
-            }
-
-            WallSqlTableStat sqlTableStat = entry.getValue();
-
-            if (sqlTableStat.getSelectCount() > 0) {
-                tableStat.addFetchRowCount(fetchRowCount);
-            }
-        }
+        provider.addFetchRowCount(sqlStat, fetchRowCount);
    }

    public long getViolationCount() {
@@ -35,6 +35,7 @@ import com.alibaba.druid.sql.parser.ParserException;
 import com.alibaba.druid.sql.parser.SQLStatementParser;
 import com.alibaba.druid.sql.parser.Token;
 import com.alibaba.druid.sql.visitor.ExportParameterVisitor;
+import com.alibaba.druid.sql.visitor.ParameterizedOutputVisitorUtils;
 import com.alibaba.druid.util.LRUCache;
 import com.alibaba.druid.wall.spi.WallVisitorUtils;
 import com.alibaba.druid.wall.violation.ErrorCode;
@@ -45,13 +46,15 @@ public abstract class WallProvider {

    private LRUCache<String, WallSqlStat>                 whiteList;

-    private int                                           MAX_SQL_LENGTH          = 2048;                                             // 1k
+    private int                                           MAX_SQL_LENGTH          = 2048;                                              // 1k

-    private int                                           whiteSqlMaxSize         = 500;                                              // 1k
+    private int                                           whiteSqlMaxSize         = 500;                                               // 1k

    private LRUCache<String, WallSqlStat>                 blackList;

-    private int                                           blackSqlMaxSize         = 100;                                              // 1k
+    private int                                           blackSqlMaxSize         = 100;                                               // 1k
+
+    private LRUCache<String, WallSqlStat>                 sqlList;

    protected final WallConfig                            config;

@@ -59,8 +62,14 @@ public abstract class WallProvider {

    private static final ThreadLocal<Boolean>             privileged              = new ThreadLocal<Boolean>();

-    private final ConcurrentMap<String, WallFunctionStat> functionStats           = new ConcurrentHashMap<String, WallFunctionStat>(16, 0.75f, 1);
-    private final ConcurrentMap<String, WallTableStat>    tableStats              = new ConcurrentHashMap<String, WallTableStat>(16, 0.75f, 1);
+    private final ConcurrentMap<String, WallFunctionStat> functionStats           = new ConcurrentHashMap<String, WallFunctionStat>(
+                                                                                                                                    16,
+                                                                                                                                    0.75f,
+                                                                                                                                    1);
+    private final ConcurrentMap<String, WallTableStat>    tableStats              = new ConcurrentHashMap<String, WallTableStat>(
+                                                                                                                                 16,
+                                                                                                                                 0.75f,
+                                                                                                                                 1);

    public final WallDenyStat                             commentDeniedStat       = new WallDenyStat();

@@ -102,6 +111,16 @@ public abstract class WallProvider {
        return this.functionStats;
    }

+    public WallSqlStat getSqlStat(String sql) {
+        WallSqlStat sqlStat = this.getWhiteSql(sql);
+
+        if (sqlStat == null) {
+            sqlStat = this.getBlackSql(sql);
+        }
+
+        return sqlStat;
+    }
+
    public WallTableStat getTableStat(String tableName) {
        String lowerCaseName = tableName.toLowerCase();
        if (lowerCaseName.startsWith("`") && lowerCaseName.endsWith("`")) {
@@ -111,6 +130,54 @@ public abstract class WallProvider {
        return getTableStatWithLowerName(lowerCaseName);
    }

+    public void addUpdateCount(WallSqlStat sqlStat, long updateCount) {
+        Map<String, WallSqlTableStat> sqlTableStats = sqlStat.getTableStats();
+        if (sqlTableStats == null) {
+            return;
+        }
+
+        for (Map.Entry<String, WallSqlTableStat> entry : sqlTableStats.entrySet()) {
+            String tableName = entry.getKey();
+            WallTableStat tableStat = this.getTableStat(tableName);
+            if (tableStat == null) {
+                continue;
+            }
+
+            WallSqlTableStat sqlTableStat = entry.getValue();
+
+            if (sqlTableStat.getDeleteCount() > 0) {
+                tableStat.addDeleteDataCount(updateCount);
+            } else if (sqlTableStat.getUpdateCount() > 0) {
+                tableStat.addUpdateDataCount(updateCount);
+            } else if (sqlTableStat.getInsertCount() > 0) {
+                tableStat.addInsertDataCount(updateCount);
+            }
+        }
+    }
+
+    public void addFetchRowCount(WallSqlStat sqlStat, long fetchRowCount) {
+        sqlStat.addAndGetEffectRowCount(fetchRowCount);
+
+        Map<String, WallSqlTableStat> sqlTableStats = sqlStat.getTableStats();
+        if (sqlTableStats == null) {
+            return;
+        }
+
+        for (Map.Entry<String, WallSqlTableStat> entry : sqlTableStats.entrySet()) {
+            String tableName = entry.getKey();
+            WallTableStat tableStat = this.getTableStat(tableName);
+            if (tableStat == null) {
+                continue;
+            }
+
+            WallSqlTableStat sqlTableStat = entry.getValue();
+
+            if (sqlTableStat.getSelectCount() > 0) {
+                tableStat.addFetchRowCount(fetchRowCount);
+            }
+        }
+    }
+
    public WallTableStat getTableStatWithLowerName(String lowerCaseName) {
        WallTableStat stat = tableStats.get(lowerCaseName);
        if (stat == null) {
@@ -154,9 +221,19 @@ public abstract class WallProvider {
                whiteList = new LRUCache<String, WallSqlStat>(whiteSqlMaxSize);
            }

+            if (sqlList == null) {
+                sqlList = new LRUCache<String, WallSqlStat>(blackSqlMaxSize + whiteSqlMaxSize);
+            }
+
            WallSqlStat wallStat = whiteList.get(sql);
            if (wallStat == null) {
-                wallStat = new WallSqlStat(tableStats, functionStats, syntaxError);
+                String mergedSql = ParameterizedOutputVisitorUtils.parameterize(sql, dbType);
+                wallStat = sqlList.get(mergedSql);
+                if (wallStat == null) {
+                    wallStat = new WallSqlStat(tableStats, functionStats, syntaxError);
+                    sqlList.put(mergedSql, wallStat);
+                }
+
                wallStat.incrementAndGetExecuteCount();
                whiteList.put(sql, wallStat);
            }
@@ -176,9 +253,19 @@ public abstract class WallProvider {
                blackList = new LRUCache<String, WallSqlStat>(blackSqlMaxSize);
            }

+            if (sqlList == null) {
+                sqlList = new LRUCache<String, WallSqlStat>(blackSqlMaxSize + whiteSqlMaxSize);
+            }
+
            WallSqlStat wallStat = blackList.get(sql);
            if (wallStat == null) {
-                wallStat = new WallSqlStat(tableStats, functionStats, violations, syntaxError);
+                String mergedSql = ParameterizedOutputVisitorUtils.parameterize(sql, dbType);
+                wallStat = sqlList.get(mergedSql);
+                if (wallStat == null) {
+                    wallStat = new WallSqlStat(tableStats, functionStats, violations, syntaxError);
+                    sqlList.put(mergedSql, wallStat);
+                }
+
                wallStat.incrementAndGetExecuteCount();
                blackList.put(sql, wallStat);
            }
@@ -202,6 +289,20 @@ public abstract class WallProvider {

        return Collections.<String> unmodifiableSet(hashSet);
    }
+    
+    public Set<String> getSqlList() {
+        Set<String> hashSet = new HashSet<String>();
+        lock.readLock().lock();
+        try {
+            if (sqlList != null) {
+                hashSet.addAll(sqlList.keySet());
+            }
+        } finally {
+            lock.readLock().unlock();
+        }
+
+        return Collections.<String> unmodifiableSet(hashSet);
+    }

    public Set<String> getBlackList() {
        Set<String> hashSet = new HashSet<String>();
@@ -246,7 +347,22 @@ public abstract class WallProvider {
    }

    public void clearCache() {
-        clearWhiteList();
+        lock.writeLock().lock();
+        try {
+            if (whiteList != null) {
+                whiteList = null;
+            }
+            
+            if (blackList != null) {
+                blackList = null;
+            }
+            
+            if (sqlList != null) {
+                sqlList = null;
+            }
+        } finally {
+            lock.writeLock().unlock();
+        }
    }

    public void clearWhiteList() {
@@ -627,11 +743,11 @@ public abstract class WallProvider {
    public long getHardCheckCount() {
        return hardCheckCount.get();
    }
-    
+
    public long getViolationEffectRowCount() {
        return violationEffectRowCount.get();
    }
-    
+
    public void addViolationEffectRowCount(long rowCount) {
        violationEffectRowCount.addAndGet(rowCount);
    }
@@ -17,17 +17,27 @@ package com.alibaba.druid.wall;
 */
 public class WallSqlTableStat {

-    private int selectCount;
-    private int selectIntoCount;
-    private int insertCount;
-    private int updateCount;
-    private int deleteCount;
-    private int truncateCount;
-    private int createCount;
-    private int alterCount;
-    private int dropCount;
-    private int replaceCount;
-    private int showCount;
+    private int    selectCount;
+    private int    selectIntoCount;
+    private int    insertCount;
+    private int    updateCount;
+    private int    deleteCount;
+    private int    truncateCount;
+    private int    createCount;
+    private int    alterCount;
+    private int    dropCount;
+    private int    replaceCount;
+    private int    showCount;
+
+    private String sample;
+
+    public String getSample() {
+        return sample;
+    }
+
+    public void setSample(String sample) {
+        this.sample = sample;
+    }

    public int getReplaceCount() {
        return replaceCount;
@@ -149,11 +159,10 @@ public class WallSqlTableStat {
        this.dropCount += value;
    }

-    
    public int getShowCount() {
        return showCount;
    }
-    
+
    public void incrementShowCount() {
        this.showCount++;
    }
@@ -0,0 +1,51 @@
+package com.alibaba.druid.bvt.filter.wall;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.wall.WallSqlStat;
+import com.alibaba.druid.wall.WallTableStat;
+import com.alibaba.druid.wall.spi.MySqlWallProvider;
+
+
+public class WallProviderTest extends TestCase {
+    public void test_getSqlStat() throws Exception {
+        String whiteSql_1 = "select * from t1 where fid = 1";
+        String whiteSql_2 = "select * from t1 where fid = 2";
+        String blackSql_1 = "select * from t2 where fid = 1 OR 1 = 1";
+        String blackSql_2 = "select * from t2 where fid = 2 OR 1 = 1";
+        
+        MySqlWallProvider provider = new MySqlWallProvider();
+        
+        Assert.assertTrue(provider.checkValid(whiteSql_1));
+        Assert.assertTrue(provider.checkValid(whiteSql_2));
+        Assert.assertEquals(1, provider.getSqlList().size());
+        
+        Assert.assertFalse(provider.checkValid(blackSql_1));
+        Assert.assertFalse(provider.checkValid(blackSql_2));
+        Assert.assertEquals(2, provider.getSqlList().size());
+        
+        Assert.assertNotNull(provider.getSqlStat(whiteSql_1));
+        Assert.assertNotNull(provider.getSqlStat(blackSql_1));
+        
+        Assert.assertSame(provider.getSqlStat(blackSql_1), provider.getSqlStat(blackSql_2));
+        
+        WallSqlStat whiteStat_1 = provider.getSqlStat(whiteSql_1);
+        WallSqlStat whiteStat_2 = provider.getSqlStat(whiteSql_2);
+        Assert.assertSame(whiteStat_1, whiteStat_2);
+        provider.addFetchRowCount(whiteStat_1, 3);
+        WallTableStat wallTableStat = provider.getTableStat("t1");
+        Assert.assertNotNull(wallTableStat);
+        Assert.assertEquals(3, wallTableStat.getFetchRowCount());
+        
+        for (int i = 1000; i < 1000 * 2; ++i) {
+            String sql_x = "select * from t1 where fid = " + i;
+            Assert.assertTrue(provider.checkValid(sql_x));
+            Assert.assertSame(whiteStat_1, provider.getSqlStat(sql_x));
+        }
+        
+        Assert.assertEquals(500, provider.getWhiteList().size());
+        Assert.assertEquals(2, provider.getSqlList().size());
+    }
+}