From 707fe5214ece8d9bd7b2ac45a012959743070a55 Mon Sep 17 00:00:00 2001 From: wenshao Date: Fri, 5 Apr 2013 22:32:54 +0800 Subject: [PATCH] improve wall filter --- .../druid/stat/DruidStatManagerFacade.java | 1 + .../com/alibaba/druid/wall/WallContext.java | 4 ++ .../com/alibaba/druid/wall/WallFilter.java | 57 ++++++++++--------- .../com/alibaba/druid/wall/WallProvider.java | 10 +--- 4 files changed, 37 insertions(+), 35 deletions(-) diff --git a/src/main/java/com/alibaba/druid/stat/DruidStatManagerFacade.java b/src/main/java/com/alibaba/druid/stat/DruidStatManagerFacade.java index 1eee75d9f0..49d2f60f80 100644 --- a/src/main/java/com/alibaba/druid/stat/DruidStatManagerFacade.java +++ b/src/main/java/com/alibaba/druid/stat/DruidStatManagerFacade.java @@ -232,6 +232,7 @@ public final class DruidStatManagerFacade { return newMap; } + @SuppressWarnings({ "rawtypes", "unchecked" }) private static List> mergeNamedList(List listA, List listB) { Map> mapped = new HashMap>(); for (Object item : (List) listA) { diff --git a/src/main/java/com/alibaba/druid/wall/WallContext.java b/src/main/java/com/alibaba/druid/wall/WallContext.java index d0cf3e3910..2a064cb7e5 100644 --- a/src/main/java/com/alibaba/druid/wall/WallContext.java +++ b/src/main/java/com/alibaba/druid/wall/WallContext.java @@ -92,6 +92,10 @@ public class WallContext { public static void clearContext() { contextLocal.remove(); } + + public static void setContext(WallContext context) { + contextLocal.set(context); + } public WallSqlStat getSqlState() { return sqlState; diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java index dd87bc00d6..85d06c6095 100644 --- a/src/main/java/com/alibaba/druid/wall/WallFilter.java +++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java @@ -42,18 +42,20 @@ import com.alibaba.druid.wall.violation.SyntaxErrorViolation; public class WallFilter extends FilterAdapter implements WallFilterMBean { - private final static Log LOG = LogFactory.getLog(WallFilter.class); + private final static Log LOG = LogFactory.getLog(WallFilter.class); - private boolean inited = false; + private boolean inited = false; - private WallProvider provider; + private WallProvider provider; - private String dbType; + private String dbType; - private WallConfig config; + private WallConfig config; - private volatile boolean logViolation = false; - private volatile boolean throwException = true; + private volatile boolean logViolation = false; + private volatile boolean throwException = true; + + public final static String ATTR_CONTEXT = "wall.ctx"; @Override public void init(DataSourceProxy dataSource) { @@ -162,8 +164,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public void statement_addBatch(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); chain.statement_addBatch(statement, sql); @@ -300,20 +301,23 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql) throws SQLException { String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + + WallContext originalContext = WallContext.current(); try { + WallContext.create(dbType); sql = check(sql); return chain.statement_execute(statement, sql); } finally { - WallContext.clearContext(); + if (originalContext != null) { + WallContext.setContext(originalContext); + } } } @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, int autoGeneratedKeys) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_execute(statement, sql, autoGeneratedKeys); @@ -325,8 +329,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, int columnIndexes[]) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_execute(statement, sql, columnIndexes); @@ -338,8 +341,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, String columnNames[]) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_execute(statement, sql, columnNames); @@ -351,8 +353,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public ResultSetProxy statement_executeQuery(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_executeQuery(statement, sql); @@ -363,8 +364,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_executeUpdate(statement, sql); @@ -376,8 +376,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, int autoGeneratedKeys) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_executeUpdate(statement, sql, autoGeneratedKeys); @@ -389,8 +388,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, int columnIndexes[]) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_executeUpdate(statement, sql, columnIndexes); @@ -399,11 +397,16 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { } } + private void setWallContext(StatementProxy statement) { + String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); + WallContext context = WallContext.create(dbType); + statement.getAttributes().put(ATTR_CONTEXT, context); + } + @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, String columnNames[]) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + setWallContext(statement); try { sql = check(sql); return chain.statement_executeUpdate(statement, sql, columnNames); diff --git a/src/main/java/com/alibaba/druid/wall/WallProvider.java b/src/main/java/com/alibaba/druid/wall/WallProvider.java index 703f0c35cd..cc8681001a 100644 --- a/src/main/java/com/alibaba/druid/wall/WallProvider.java +++ b/src/main/java/com/alibaba/druid/wall/WallProvider.java @@ -355,7 +355,7 @@ public abstract class WallProvider { WallContext originalContext = WallContext.current(); try { - WallContext.create(dbType); + WallContext.createIfNotExists(dbType); return checkInternal(sql); } finally { if (originalContext == null) { @@ -364,17 +364,11 @@ public abstract class WallProvider { } } - public WallCheckResult checkInternal(String sql) { + private WallCheckResult checkInternal(String sql) { checkCount.incrementAndGet(); WallContext context = WallContext.current(); - if (context != null && !this.dbType.equals(context.getDbType())) { - WallContext.clearContext(); - } - - WallContext.createIfNotExists(this.dbType); - if (config.isDoPrivilegedAllow() && ispPivileged()) { return new WallCheckResult(); }