From 7515cabfdbfca3e801106ba71bc53a493e5eb2b2 Mon Sep 17 00:00:00 2001 From: dushitaoyuan Date: Fri, 6 Sep 2019 14:07:45 +0800 Subject: [PATCH] =?UTF-8?q?redis=E9=99=90=E6=B5=81=E5=A2=9E=E5=8A=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../controller/BussinessController.java | 5 +++ .../security/RateLimitAspect.java | 14 +++++++-- .../security/ratelimit/GuavaRateLimiter.java | 2 +- .../security/ratelimit/RedisRateLimiter.java | 9 ++++-- .../securitydemo/utils/StringIntUtil.java | 22 +++++++++++++ .../resources/META-INF/rate_limiter_count.lua | 31 ++++++++++++++----- src/main/resources/application.properties | 1 + .../SecurityDemoApplicationTests.java | 15 ++++++++- 8 files changed, 83 insertions(+), 16 deletions(-) create mode 100644 src/main/java/com/taoyuanx/securitydemo/utils/StringIntUtil.java diff --git a/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java b/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java index cf1df85..a859409 100644 --- a/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java +++ b/src/main/java/com/taoyuanx/securitydemo/controller/BussinessController.java @@ -206,4 +206,9 @@ public class BussinessController { public void upload(HttpServletRequest request, HttpServletResponse response) throws Exception { fileHandler.handleFile(response, request); } + @RateLimit(type = RateLimitType.TOTAL_COUNT,key = "rate_count",totalCount = 100) + @GetMapping("rate_count") + public void rateCount(int index) throws Exception { + System.out.println("rate_count\t"+index); + } } diff --git a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java index 5bfe731..e873b2f 100644 --- a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java +++ b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java @@ -87,11 +87,19 @@ public class RateLimitAspect { } } if (LOG.isDebugEnabled()) { - LOG.debug("采用[{}]限流策略,限流key:{}", RateLimitType.IP, key); + LOG.debug("采用[{}]限流策略,限流key:{}", type, key); } - if (!rateLimiter.tryAcquire(key, rateLimit.limit())) { - throw new LimitException("请求过于频繁,请稍后再试"); + if (type.equals(RateLimitType.TOTAL_COUNT)) { + if (!rateLimiter.tryCount(1, key, rateLimit.totalCount())) { + throw new LimitException("访问字数已达最大限制:" + rateLimit.totalCount() + ",请稍后再试"); + } + } else { + if (!rateLimiter.tryAcquire(key, rateLimit.limit())) { + throw new LimitException("请求过于频繁,请稍后再试"); + } } + + } diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java index f4339e8..d3b017e 100644 --- a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java +++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java @@ -20,7 +20,7 @@ public class GuavaRateLimiter extends AbstractRateLimiter { private Map countHolder = new ConcurrentHashMap(); /** - * 总数限流到0后,标记 + * 总数限流到0后,标记,会有些许误判,不在乎内存的话,可用hashset存 */ private BloomFilter TOTAL_LIMIT_ZERO_FLAG = BloomFilter.create(Funnels.stringFunnel(Charset.defaultCharset()), MAX_HOLDER_SIZE * 20); diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java index e4b6ecb..d8f8be7 100644 --- a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java +++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/RedisRateLimiter.java @@ -1,5 +1,7 @@ package com.taoyuanx.securitydemo.security.ratelimit; +import com.sun.org.apache.xml.internal.utils.StringToIntTable; +import com.taoyuanx.securitydemo.utils.StringIntUtil; import org.springframework.core.io.ClassPathResource; import org.springframework.data.redis.core.StringRedisTemplate; import org.springframework.data.redis.core.script.DefaultRedisScript; @@ -45,9 +47,10 @@ public class RedisRateLimiter extends AbstractRateLimiter { @Override public boolean tryCount(int count, String key, Long totalCount) { - String[] scriptArgs = {count + "", totalCount + ""}; - Long result = redisTemplate.execute(this.countScript, Arrays.asList(key), scriptArgs); - return result == 1L; + + String[] scriptArgs = {count + "", totalCount + "", StringIntUtil.toInt(key) +""}; + Long result = redisTemplate.execute(this.countScript, Arrays.asList(key,"zero_flag"), scriptArgs); + return result >= 0; } private List getKeys(String key) { diff --git a/src/main/java/com/taoyuanx/securitydemo/utils/StringIntUtil.java b/src/main/java/com/taoyuanx/securitydemo/utils/StringIntUtil.java new file mode 100644 index 0000000..7313d17 --- /dev/null +++ b/src/main/java/com/taoyuanx/securitydemo/utils/StringIntUtil.java @@ -0,0 +1,22 @@ +package com.taoyuanx.securitydemo.utils; + +import org.apache.commons.codec.binary.Hex; +import org.apache.commons.codec.digest.DigestUtils; + +/** + * @author dushitaoyuan + * @desc 长字符转long + * @date 2019/9/6 + */ +public class StringIntUtil { + + public static int toInt(String str) { + return Math.abs(digest(str).hashCode()); + } + + private static String digest(String str) { + return Hex.encodeHexString(DigestUtils.getMd5Digest().digest(str.getBytes())); + } + + +} diff --git a/src/main/resources/META-INF/rate_limiter_count.lua b/src/main/resources/META-INF/rate_limiter_count.lua index 406a432..74768b7 100644 --- a/src/main/resources/META-INF/rate_limiter_count.lua +++ b/src/main/resources/META-INF/rate_limiter_count.lua @@ -1,24 +1,39 @@ ---如果等于0说明超时,其他则是当前资源的访问数量 - +--判断资源归0标记是否存在 +--标记实现,可通过 布隆过滤和bitmap实现 注意redis支持情况 +local bit_key_offset=tonumber(ARGV[3]) +redis.log(redis.LOG_WARNING, "bit_key_offset " .. ARGV[3]) +local zero_flag =redis.call("GETBIT", KEYS[2],bit_key_offset) +if zero_flag == 1 then + return -1 +end --申请资源数量 local count= tonumber(ARGV[1]) +--redis.log(redis.LOG_WARNING, "key " .. KEYS[1]) + +--redis.log(redis.LOG_WARNING, "count " .. count) +--redis.log(redis.LOG_WARNING, "total_count " .. ARGV[2]) if count == nil then count = 1 end - -- 获取剩余资源数量 local last_count = tonumber(redis.call("get", KEYS[1])) if last_count == nil then - last_count = ARGV[2]; + last_count = tonumber(ARGV[2]); + redis.call("set", KEYS[1],last_count); end ---计数减少 -if last_count > count then +--计数减少,资源归0时,标记 +if last_count >= count then redis.call("DECRBY", KEYS[1], count) - return 1 + last_count=tonumber(redis.call("get", KEYS[1])) +-- redis.log(redis.LOG_WARNING, "last_count " .. last_count) + return last_count else - return 0 + redis.call("SETBIT", KEYS[2],bit_key_offset,1) + redis.call("DEL", KEYS[2]) + return -1 end + --todo bitmap 标记 资源为0 \ No newline at end of file diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties index 859e814..7cb8779 100644 --- a/src/main/resources/application.properties +++ b/src/main/resources/application.properties @@ -41,3 +41,4 @@ application.config.file-storage-dir=G:/temp application.config.systemFileFormat=http://localhost:${server.port}/api/file +logging.level.org.springframework=INFO diff --git a/src/test/java/com/taoyuanx/securitydemo/SecurityDemoApplicationTests.java b/src/test/java/com/taoyuanx/securitydemo/SecurityDemoApplicationTests.java index 113cc45..25ec02e 100644 --- a/src/test/java/com/taoyuanx/securitydemo/SecurityDemoApplicationTests.java +++ b/src/test/java/com/taoyuanx/securitydemo/SecurityDemoApplicationTests.java @@ -1,12 +1,15 @@ package com.taoyuanx.securitydemo; import com.taoyuanx.securitydemo.controller.BussinessController; +import com.taoyuanx.securitydemo.utils.StringIntUtil; import org.junit.Test; import org.junit.runner.RunWith; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.test.context.junit4.SpringRunner; +import java.util.BitSet; + @RunWith(SpringRunner.class) @SpringBootTest public class SecurityDemoApplicationTests { @@ -21,10 +24,20 @@ public class SecurityDemoApplicationTests { @Test public void reateLimit() { - int batch=100; + int batch=101; for(int i=0;i