redis限流增加

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);
+    }
 }

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("请求过于频繁,请稍后再试");
+            }
         }
+
+
     }
 
 

src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java

@@ -20,7 +20,7 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
 
     private Map<String, LongAdder> countHolder = new ConcurrentHashMap();
     /**
-     * 总数限流到0后,标记
+     * 总数限流到0后,标记,会有些许误判,不在乎内存的话,可用hashset存
      */
     private BloomFilter<CharSequence> TOTAL_LIMIT_ZERO_FLAG = BloomFilter.create(Funnels.stringFunnel(Charset.defaultCharset()), MAX_HOLDER_SIZE * 20);
 

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<String> getKeys(String key) {

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()));
+    }
+
+
+}

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

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

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<batch;i++){
             System.out.println(bussinessController.rateLimitKey());
         }
     }
 
+    @Test
+    public void reateLimitCount() throws Exception {
+        int batch=101;
+        for(int i=0;i<batch;i++){
+           bussinessController.rateCount(i);
+        }
+    }
+
+
+
 }