防止SQL注入笔记类1

import com.alibaba.druid.wall.Violation;

import com.alibaba.druid.wall.WallCheckResult;
import com.alibaba.druid.wall.WallConfig;
import com.alibaba.druid.wall.WallProvider;
import com.alibaba.druid.wall.spi.MySqlWallProvider;
import com.alibaba.druid.wall.violation.SyntaxErrorViolation;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;

import java.sql.SQLException;
import java.util.List;
import java.util.Properties;
import java.util.regex.Pattern;

/**
* @date 2023-08-16 sql 注入检测
* @author tom
*/

@Component
@Intercepts({@Signature(
type = Executor.class,
method = "query",
args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
), @Signature(
type = Executor.class,
method = "query",
args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}
)})
public class SqlInjInterceptor implements Interceptor {

private static final Pattern PATTERN = Pattern.compile("[\t\r\n]");
private final static WallProvider PROVIDER = new MySqlWallProvider(new WallConfig(MySqlWallProvider.DEFAULT_CONFIG_DIR));

public SqlInjInterceptor() {
}

public Object intercept(Invocation invocation) throws Throwable {
Object[] args = invocation.getArgs();
MappedStatement ms = (MappedStatement)args[0];
Object parameter = args[1];
BoundSql boundSql;
if (args.length == 4) {
boundSql = ms.getBoundSql(parameter);
} else {
boundSql = (BoundSql)args[5];
}

sqlInj(boundSql.getSql());
return invocation.proceed();
}

public Object plugin(Object target) {
return Plugin.wrap(target, this);
}

public void setProperties(Properties properties) {
}

private static void sqlInj(String sql) throws SQLException {
sql = removeDataAuthFlag(sql);
WallCheckResult checkResult = PROVIDER.check(sql);
List<Violation> violations = checkResult.getViolations();

if (violations.size() == 0){
return;
}
Violation firstViolation = violations.get(0);
if (violations.get(0) instanceof SyntaxErrorViolation) {
SyntaxErrorViolation violation = (SyntaxErrorViolation) violations.get(0);
throw new SQLException("sql injection violation: " + firstViolation.getMessage(),
violation.getException());
} else {
throw new SQLException("sql injection violation: " + firstViolation.getMessage());
}
}

// 数据权限拦截误判
private static String removeDataAuthFlag(String sql) {
if (StringUtils.isBlank(sql)) {
return sql;
}
if (!sql.contains("(@") && !sql.contains("@)")) {
return sql;
}
sql = sql.toLowerCase();
for (;;) {
if (sql.contains("(@") && sql.contains("@)")) {
String front = sql.substring(0, sql.indexOf("(@"));
String after = sql.substring(sql.indexOf("@)") + 2);
// and 处理：
// 若 (@ 前面是 where, @) 后方是 and, 要去除 后方的 and
// 若 (@ 前面是 and, 移除前方的 and
if (front.trim().endsWith("where") && after.trim().startsWith("and")) {
after = after.substring(after.indexOf("and") + 3);
}
if (front.trim().endsWith("and")) {
front = front.substring(0, front.lastIndexOf("and"));
}
sql = front + after;
} else {
break;
}
}
return sql;
}

public static void main(String[] args) {

String sql = "select 1 from dual\n" +
"where (@{\"columnCode\":\"wh_code\",\"tableName\":\"cd_warehouse\",\"custom\":\"rh.wh_code\",\"dataCode\":\"114444411\"}@) " +
"AND 1=2 \r\n\t and (@{\"columnCode\":\"wh_code\",\"tableName\":\"cd_warehouse\",\"custom\":\"rh.wh_code\",\"dataCode\":\"22882222\"}@) " +
" and 1 = 2 \n order by =sleep(1)";

System.out.println(sql);
sql = removeDataAuthFlag(sql);
System.out.println(sql);
}

}