diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java index 8c947ed28e..82fefd0713 100644 --- a/src/main/java/com/alibaba/druid/wall/WallFilter.java +++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java @@ -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 sqlTableStats = sqlStat.getTableStats(); - if (sqlTableStats == null) { - return; - } - - for (Map.Entry 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 sqlTableStats = sqlStat.getTableStats(); - if (sqlTableStats == null) { - return; - } - - for (Map.Entry 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() { diff --git a/src/main/java/com/alibaba/druid/wall/WallProvider.java b/src/main/java/com/alibaba/druid/wall/WallProvider.java index 269dfd6789..31daf86e33 100644 --- a/src/main/java/com/alibaba/druid/wall/WallProvider.java +++ b/src/main/java/com/alibaba/druid/wall/WallProvider.java @@ -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 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 blackList; - private int blackSqlMaxSize = 100; // 1k + private int blackSqlMaxSize = 100; // 1k + + private LRUCache sqlList; protected final WallConfig config; @@ -59,8 +62,14 @@ public abstract class WallProvider { private static final ThreadLocal privileged = new ThreadLocal(); - private final ConcurrentMap functionStats = new ConcurrentHashMap(16, 0.75f, 1); - private final ConcurrentMap tableStats = new ConcurrentHashMap(16, 0.75f, 1); + private final ConcurrentMap functionStats = new ConcurrentHashMap( + 16, + 0.75f, + 1); + private final ConcurrentMap tableStats = new ConcurrentHashMap( + 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 sqlTableStats = sqlStat.getTableStats(); + if (sqlTableStats == null) { + return; + } + + for (Map.Entry 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 sqlTableStats = sqlStat.getTableStats(); + if (sqlTableStats == null) { + return; + } + + for (Map.Entry 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(whiteSqlMaxSize); } + if (sqlList == null) { + sqlList = new LRUCache(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(blackSqlMaxSize); } + if (sqlList == null) { + sqlList = new LRUCache(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. unmodifiableSet(hashSet); } + + public Set getSqlList() { + Set hashSet = new HashSet(); + lock.readLock().lock(); + try { + if (sqlList != null) { + hashSet.addAll(sqlList.keySet()); + } + } finally { + lock.readLock().unlock(); + } + + return Collections. unmodifiableSet(hashSet); + } public Set getBlackList() { Set hashSet = new HashSet(); @@ -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); } diff --git a/src/main/java/com/alibaba/druid/wall/WallSqlTableStat.java b/src/main/java/com/alibaba/druid/wall/WallSqlTableStat.java index 3452766a78..ac9854838b 100644 --- a/src/main/java/com/alibaba/druid/wall/WallSqlTableStat.java +++ b/src/main/java/com/alibaba/druid/wall/WallSqlTableStat.java @@ -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++; } diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallProviderTest.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallProviderTest.java new file mode 100644 index 0000000000..87a64cf63a --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallProviderTest.java @@ -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()); + } +}