diff --git a/pom.xml b/pom.xml
index 3920b14..7ce2cb8 100644
--- a/pom.xml
+++ b/pom.xml
@@ -26,11 +26,11 @@
org.springframework.boot
spring-boot-starter
-
+
org.springframework.boot
spring-boot-starter-web
@@ -81,11 +81,13 @@
hutool-core
4.0.3
-
+
+
+
+ org.springframework.boot
+ spring-boot-starter-data-redis
+ true
+
diff --git a/src/main/java/com/taoyuanx/securitydemo/config/GlobalConfig.java b/src/main/java/com/taoyuanx/securitydemo/config/GlobalConfig.java
index 8f9cabf..c9f1e60 100644
--- a/src/main/java/com/taoyuanx/securitydemo/config/GlobalConfig.java
+++ b/src/main/java/com/taoyuanx/securitydemo/config/GlobalConfig.java
@@ -1,11 +1,20 @@
package com.taoyuanx.securitydemo.config;
+import com.taoyuanx.securitydemo.security.ratelimit.AbstractRateLimiter;
+import com.taoyuanx.securitydemo.security.ratelimit.RedisRateLimiter;
import lombok.Data;
import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
+import org.springframework.data.redis.connection.RedisConnectionFactory;
+import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
+import org.springframework.data.redis.core.RedisTemplate;
+import org.springframework.data.redis.core.StringRedisTemplate;
+import java.net.UnknownHostException;
import java.util.List;
/**
@@ -35,5 +44,16 @@ public class GlobalConfig {
return environment.getProperty(configKey);
}
+ /**
+ * 限流实现类
+ * @param redisTemplate
+ * @return
+ */
+ @Bean
+ public AbstractRateLimiter rateLimiter(RedisTemplate redisTemplate){
+ RedisRateLimiter redisRateLimiter=new RedisRateLimiter(redisTemplate);
+ return redisRateLimiter;
+ }
+
}
diff --git a/src/main/java/com/taoyuanx/securitydemo/config/MvcConfig.java b/src/main/java/com/taoyuanx/securitydemo/config/MvcConfig.java
index c206593..6530228 100644
--- a/src/main/java/com/taoyuanx/securitydemo/config/MvcConfig.java
+++ b/src/main/java/com/taoyuanx/securitydemo/config/MvcConfig.java
@@ -159,4 +159,6 @@ public class MvcConfig implements WebMvcConfigurer {
return fileHandler;
}
+
+
}
diff --git a/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java b/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java
index 5be8695..87401f0 100644
--- a/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java
+++ b/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java
@@ -112,7 +112,7 @@ public class BussinessController {
*/
@GetMapping("rateLimit_key")
@ResponseBody
- @RateLimit(type = RateLimitType.SERVICE_KEY, limit = 2, key = "api/rateLimit_key")
+ @RateLimit(type = RateLimitType.SERVICE_KEY, limit = 100, key = "api/rateLimit_key")
public String rateLimitKey() {
return "hello rateLimit!";
}
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/RateLimit.java b/src/main/java/com/taoyuanx/securitydemo/security/RateLimit.java
index 104ed05..e9c578a 100644
--- a/src/main/java/com/taoyuanx/securitydemo/security/RateLimit.java
+++ b/src/main/java/com/taoyuanx/securitydemo/security/RateLimit.java
@@ -11,7 +11,9 @@ import java.lang.annotation.*;
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
/**
- * 限流并发 和限流的key,类型
+ * limit 每秒并发
+ * key 限流的key
+ * type 限流类型 参见:com.taoyuanx.securitydemo.security.RateLimitType
*/
double limit() default 100;
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
index 625481e..5bfe731 100644
--- a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
+++ b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
@@ -1,9 +1,8 @@
package com.taoyuanx.securitydemo.security;
-import com.google.common.collect.Maps;
-import com.google.common.util.concurrent.RateLimiter;
import com.taoyuanx.securitydemo.exception.LimitException;
+import com.taoyuanx.securitydemo.security.ratelimit.AbstractRateLimiter;
import com.taoyuanx.securitydemo.utils.RequestUtil;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
@@ -13,11 +12,11 @@ import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
-import java.util.Map;
/**
* aop限流
@@ -27,8 +26,8 @@ import java.util.Map;
public class RateLimitAspect {
private static final Logger LOG = LoggerFactory.getLogger(RateLimitAspect.class);
- private Map rateHolder = Maps.newConcurrentMap();
- private static final int MAX_HOLDER_SIZE = 50000;
+ @Autowired
+ AbstractRateLimiter rateLimiter;
@Pointcut("execution(* com.taoyuanx.securitydemo.controller..*.*(..))&& (@annotation(RateLimit)||@annotation(Rate))")
public void ratePointCut() {
@@ -46,13 +45,11 @@ public class RateLimitAspect {
handleRateLimit(rateLimit, methodName);
} else {
Rate rate = AnnotationUtils.findAnnotation(currentMethod, Rate.class);
- if (rate != null) {
- RateLimit[] rateLimitArray = rate.rate();
- if (rateLimitArray != null) {
- for (RateLimit limit : rateLimitArray) {
- handleRateLimit(limit, methodName);
- }
+ if (rate != null && rate.rate() != null) {
+ for (RateLimit limit : rate.rate()) {
+ handleRateLimit(limit, methodName);
}
+
}
}
@@ -61,13 +58,6 @@ public class RateLimitAspect {
private void handleRateLimit(RateLimit rateLimit, String methodName) throws Throwable {
- RateLimiter rateLimiter = doGetRateLimiter(rateLimit, methodName);
- if (!rateLimiter.tryAcquire()) {
- throw new LimitException("请求过于频繁,请稍后再试");
- }
- }
-
- private RateLimiter doGetRateLimiter(RateLimit rateLimit, String methodName) {
RateLimitType type = rateLimit.type();
String key = rateLimit.key();
if (type == null) {
@@ -84,40 +74,24 @@ public class RateLimitAspect {
} else {
key = RequestUtil.getRemoteIp() + "_" + serviceKey;
}
- if (LOG.isDebugEnabled()) {
- LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.IP, key);
- }
break;
case METHOD:
key = methodName;
- if (LOG.isDebugEnabled()) {
- LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.METHOD, key);
- }
break;
case SERVICE_KEY:
- if (LOG.isDebugEnabled()) {
- LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.SERVICE_KEY, key);
- }
+
break;
case GLOBAL:
key = "global";
- if (LOG.isDebugEnabled()) {
- LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.GLOBAL, key);
- }
break;
}
}
- //超过固定阈值,清空,重构
- if (rateHolder.size() > MAX_HOLDER_SIZE) {
- rateHolder.clear();
+ if (LOG.isDebugEnabled()) {
+ LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.IP, key);
}
- if (rateHolder.containsKey(key)) {
- return rateHolder.get(key);
+ if (!rateLimiter.tryAcquire(key, rateLimit.limit())) {
+ throw new LimitException("请求过于频繁,请稍后再试");
}
- RateLimiter rateLimiter = RateLimiter.create(rateLimit.limit());
- rateHolder.putIfAbsent(key, rateLimiter);
- return rateHolder.get(key);
-
}
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/AbstractRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/AbstractRateLimiter.java
new file mode 100644
index 0000000..3b897c7
--- /dev/null
+++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/AbstractRateLimiter.java
@@ -0,0 +1,35 @@
+package com.taoyuanx.securitydemo.security.ratelimit;
+
+/**
+ * @author dushitaoyuan
+ * @desc 抽象限流
+ * @date 2019/9/5
+ */
+public abstract class AbstractRateLimiter {
+
+
+ /**
+ * 尝试获取令牌
+ *
+ * @param key 限流标识
+ * @param limit 限流速率
+ * @return
+ */
+ public boolean tryAcquire(String key, Double limit){
+ return doTryAcquire(1,key,limit);
+ }
+
+ /**
+ * 尝试获取令牌
+ *
+ * @param permits 获取令牌数量
+ * @param key 限流标识
+ * @param limit 限流速率
+ * @return
+ */
+ public boolean tryAcquire(int permits, String key, Double limit){
+ return doTryAcquire(permits,key,limit);
+ }
+
+ protected abstract boolean doTryAcquire(int permits, String key, Double limit);
+}
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java
new file mode 100644
index 0000000..af851cc
--- /dev/null
+++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java
@@ -0,0 +1,41 @@
+package com.taoyuanx.securitydemo.security.ratelimit;
+
+import com.google.common.collect.Maps;
+import com.google.common.util.concurrent.RateLimiter;
+
+import java.util.Map;
+
+/**
+ * @author dushitaoyuan
+ * @desc guava限流实现
+ * @date 2019/9/5
+ */
+public class GuavaRateLimiter extends AbstractRateLimiter {
+ private Map rateHolder = Maps.newConcurrentMap();
+ private static final int MAX_HOLDER_SIZE = 50000;
+ /* @Override
+ public boolean tryAcquire(String key, Double limit) {
+ return doTryAcquire(1, key, limit);
+ }
+
+ @Override
+ public boolean tryAcquire(int permits, String key, Double limit) {
+ return doTryAcquire(permits, key, limit);
+ }
+*/
+ protected boolean doTryAcquire(int permits, String key, Double limit) {
+ //超过固定阈值,清空,重构
+ if (rateHolder.size() > MAX_HOLDER_SIZE) {
+ rateHolder.clear();
+ }
+ RateLimiter rateLimiter = null;
+ if (rateHolder.containsKey(key)) {
+ rateLimiter = rateHolder.get(key);
+ }
+ rateLimiter = RateLimiter.create(limit);
+ rateHolder.putIfAbsent(key, rateLimiter);
+ return rateLimiter.tryAcquire(permits);
+ }
+
+
+}
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java
new file mode 100644
index 0000000..0fb0815
--- /dev/null
+++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java
@@ -0,0 +1,62 @@
+package com.taoyuanx.securitydemo.security.ratelimit;
+
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.data.redis.core.RedisTemplate;
+import org.springframework.data.redis.core.script.DefaultRedisScript;
+import org.springframework.data.redis.core.script.RedisScript;
+import org.springframework.scripting.support.ResourceScriptSource;
+
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * @author dushitaoyuan
+ * @desc redis 限流实现
+ * 复制 spring-cloud-gateway,可与spring-boot,spring-cloud等分离使用
+ * @date 2019/9/5
+ */
+public class RedisRateLimiter extends AbstractRateLimiter {
+ private RedisTemplate stringRedisTemplate;
+ private RedisScript> script;
+
+ public RedisRateLimiter(RedisTemplate redisTemplate) {
+ DefaultRedisScript script = new DefaultRedisScript();
+ script.setScriptSource(new ResourceScriptSource(
+ new ClassPathResource("META-INF/demo.lua")));
+ script.setResultType(List.class);
+ this.script = script;
+ this.stringRedisTemplate = redisTemplate;
+ }
+
+ /* @Override
+ public boolean tryAcquire(String key, Double limit) {
+ return doTryAcquire(1, key, limit);
+ }
+
+ @Override
+ public boolean tryAcquire(int permits, String key, Double limit) {
+ return doTryAcquire(permits, key, limit);
+ }*/
+ @Override
+ protected boolean doTryAcquire(int permits, String key, Double limit) {
+ List