分钟、小时、天限流

ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/config/GatewayConfig.java

@@ -7,17 +7,13 @@ import org.redisson.api.RedissonClient;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.cloud.gateway.filter.GlobalFilter;
 import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
-import org.springframework.cloud.gateway.route.RouteDefinition;
 import org.springframework.cloud.gateway.route.RouteDefinitionRepository;
 import org.springframework.cloud.gateway.support.ConfigurationService;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-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.script.DefaultRedisScript;
 import org.springframework.data.redis.core.script.RedisScript;
-import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
-import org.springframework.data.redis.serializer.StringRedisSerializer;
 
 import javax.annotation.Resource;
 import java.util.List;
@@ -31,26 +27,9 @@ import java.util.List;
 @Configuration
 public class GatewayConfig
 {
-
-    @Resource
-    private RedisTemplate<String, RouteDefinition> redisTemplate;
     @Resource
     private RedissonClient redissonClient;
 
-    @Bean
-    public RedisTemplate<String, RouteDefinition> redisTemplate(RedisConnectionFactory redisConnectionFactory)
-    {
-        final RedisTemplate<String, RouteDefinition> template = new RedisTemplate<>();
-        template.setConnectionFactory(redisConnectionFactory);
-        Jackson2JsonRedisSerializer<RouteDefinition> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(RouteDefinition.class);
-        template.setValueSerializer(jackson2JsonRedisSerializer);
-        template.setHashValueSerializer(jackson2JsonRedisSerializer);
-        final StringRedisSerializer stringRedisSerializer = new StringRedisSerializer();
-        template.setKeySerializer(stringRedisSerializer);
-        template.setHashKeySerializer(stringRedisSerializer);
-        return template;
-    }
-
     /**
      * Redis 路由仓库
      *
@@ -65,8 +44,8 @@ public class GatewayConfig
     @Bean
     public CustomerRedisRateLimiter customerRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate,
              @Qualifier(RedisRateLimiter.REDIS_SCRIPT_NAME) RedisScript<List<Long>> redisScript,
-             ConfigurationService configurationService) {
+             ConfigurationService configurationService, DefaultRedisScript<Long> timeRedisScript) {
-        return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService);
+        return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService, timeRedisScript);
     }
 
     @Bean

ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java

@@ -9,7 +9,6 @@ import com.ruoyi.gateway.utils.GatewayUtils;
 import com.ruoyi.gateway.utils.beans.IscRule;
 import org.springframework.cloud.gateway.filter.GatewayFilterChain;
 import org.springframework.cloud.gateway.filter.GlobalFilter;
-import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
 import org.springframework.cloud.gateway.route.Route;
 import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
 import org.springframework.core.Ordered;
@@ -28,6 +27,7 @@ import reactor.core.publisher.Mono;
 
 import java.net.URI;
 import java.util.*;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -48,10 +48,11 @@ import static org.springframework.util.CollectionUtils.unmodifiableMultiValueMap
  * @date 2021-10-15
  */
 public class CustomerGlobalFilter implements GlobalFilter, Ordered {
-    private CustomerRedisRateLimiter rateLimiter;
+    private final CustomerRedisRateLimiter rateLimiter;
     public CustomerGlobalFilter(CustomerRedisRateLimiter rateLimiter) {
         this.rateLimiter = rateLimiter;
     }
+    public static final TimeUnit[] TIME_UNITS = {TimeUnit.SECONDS, TimeUnit.MINUTES, TimeUnit.HOURS, TimeUnit.DAYS};
 
     @Override
     public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
@@ -65,7 +66,7 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
         if (HttpMethod.GET.equals(httpMethod)) {
             final MultiValueMap<String, String> queryParams = new LinkedMultiValueMap<>(request.getQueryParams());
             final IscRule rule = handleRule(headerAk, () -> queryParams.get(accessKeyName), route);
-            Supplier<Mono<Void>> rateLimiterAftersupplier = () -> {
+            Supplier<Mono<Void>> rateLimiterAfterSupplier = () -> {
                 removeParam(headerAk, request, accessKeyName, () -> queryParams.remove(accessKeyName));
                 handleHiddenParams(route, queryParams, (next, map) -> {
                     Object value;
@@ -83,7 +84,7 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                 return chain.filter(exchange.mutate().request(updatedRequest).build());
             };
 
-            return rateLimiter(exchange, rule, route, rateLimiterAftersupplier);
+            return rateLimiter(exchange, rule, route, 0, rateLimiterAfterSupplier);
         } else if (HttpMethod.POST.equals(httpMethod)) {
             final ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
             final Mono<String> modifiedBody = serverRequest.bodyToMono(String.class);
@@ -93,7 +94,7 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                     JSONObject jsonObj = JSONUtil.parseObj(body);
                     final IscRule rule = handleRule(headerAk, () ->
                             Arrays.asList(jsonObj.get(accessKeyName, String.class, true)), route);
-                    Supplier<Mono<String>> rateLimiterAftersupplier = () -> {
+                    Supplier<Mono<String>> rateLimiterAfterSupplier = () -> {
                         removeParam(headerAk, request, accessKeyName, () -> jsonObj.remove(accessKeyName));
                         handleHiddenParams(route, jsonObj, (next, map) -> {
                             map.set(next.getKey(), next.getValue());
@@ -101,13 +102,13 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                         return Mono.just(jsonObj.toString());
                     };
 
-                    return rateLimiter(exchange, rule, route, rateLimiterAftersupplier);
+                    return rateLimiter(exchange, rule, route, 0, rateLimiterAfterSupplier);
                 } else if (MediaType.APPLICATION_FORM_URLENCODED.equals(mediaType)) {
                     if (StringUtils.hasText(body)) {
                         final Stream<String[]> stream = Arrays.stream(body.split("&")).map(param -> param.split("="));
                         final IscRule rule = handleRule(headerAk, () -> stream.filter(param -> param.length > 0 &&
                                         accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()), route);
-                        Supplier<Mono<String>> rateLimiterAftersupplier = () -> {
+                        Supplier<Mono<String>> rateLimiterAfterSupplier = () -> {
                             removeParam(headerAk, request, accessKeyName, null);
                             final List<String[]> params = stream.filter(param -> !accessKeyName.equals(param[0])).collect(Collectors.toList());
                             handleHiddenParams(route, params, (next, list) -> {
@@ -125,7 +126,7 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                             return Mono.just(params.stream().map(param -> param[0] + '=' + param[1]).collect(Collectors.joining("&")));
                         };
 
-                        return rateLimiter(exchange, rule, route, rateLimiterAftersupplier);
+                        return rateLimiter(exchange, rule, route, 0, rateLimiterAfterSupplier);
                     }
                 }
                 return Mono.empty();
@@ -207,20 +208,32 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
      * 限流
      * @param exchange
      * @param rule
-     * @Param route
+     * @param route
-     * @param rateLimiterAftersupplier 限流后操作(删除参数、添加隐藏参数，跳转)
+     * @param rateLimiterAfterSupplier 限流后操作(删除参数、添加隐藏参数，跳转)
      * @param <T>
      * @return
      */
-    private <T extends Object> Mono rateLimiter(ServerWebExchange exchange, IscRule rule, Route route,
+    private <T extends Object> Mono rateLimiter(ServerWebExchange exchange, IscRule rule, Route route, final int index,
-                                                Supplier<Mono<T>> rateLimiterAftersupplier) {
+                                                Supplier<Mono<T>> rateLimiterAfterSupplier) {
-        final RedisRateLimiter.Config config = new RedisRateLimiter.Config().setReplenishRate(1);
+        final TimeUnit timeUnit = TIME_UNITS[index];
-        return rateLimiter.isAllowed(config, rule.getId() + ':' + route.getId()).flatMap(response -> {
+        final Long limit = TimeUnit.SECONDS.equals(timeUnit) ? rule.getSecondsLimit() : TimeUnit.MINUTES.equals(timeUnit)
+                ? rule.getMinutesLimit() : TimeUnit.HOURS.equals(timeUnit) ? rule.getHoursLimit() : rule.getDaysLimit();
+        if(Objects.isNull(limit) || limit <= 0L) {
+            if(TimeUnit.DAYS.equals(timeUnit)) {
+                return rateLimiterAfterSupplier.get();
+            }
+            return rateLimiter(exchange, rule, route, index + 1, rateLimiterAfterSupplier);
+        }
+        return rateLimiter.isAllowed(route.getId(), rule.getId(), limit, timeUnit).flatMap(response -> {
             for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                 exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
             }
             if (response.isAllowed()) {
-                return rateLimiterAftersupplier.get();
+                if(TimeUnit.DAYS.equals(timeUnit)) {
+                    return rateLimiterAfterSupplier.get();
+                }
+                return rateLimiter(exchange, rule, route, index + 1, rateLimiterAfterSupplier);
+
             }
             setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
             return exchange.getResponse().setComplete();

ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java

@@ -1,17 +1,23 @@
 package com.ruoyi.gateway.ratelimit;
 
+import cn.hutool.core.date.LocalDateTimeUtil;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
 import org.springframework.cloud.gateway.support.ConfigurationService;
 import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
+import org.springframework.data.redis.core.script.DefaultRedisScript;
 import org.springframework.data.redis.core.script.RedisScript;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -24,20 +30,47 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
     private ReactiveStringRedisTemplate redisTemplate;
 
     private RedisScript<List<Long>> script;
+    private RedisScript<Long> timeRedisScript;
 
     private AtomicBoolean initialized = new AtomicBoolean(false);
 
     public CustomerRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate, RedisScript<List<Long>> script,
-                                    ConfigurationService configurationService) {
+                                    ConfigurationService configurationService, DefaultRedisScript<Long> timeRedisScript) {
         super(redisTemplate, script, configurationService);
         this.redisTemplate = redisTemplate;
-        this.redisTemplate = redisTemplate;
         this.script = script;
+        this.timeRedisScript = timeRedisScript;
         this.initialized.compareAndSet(false, true);
     }
 
-    public Mono<Response> isAllowed(Config routeConfig, String id)
+    public Mono<Response> isAllowed(String routeId, String ak, Long replenishRate, TimeUnit timeUnit) {
-    {
+        Config routeConfig = new Config().setReplenishRate(replenishRate.intValue());
+        String id = ak + ':' + routeId;
+        int time = 60;
+        LocalDateTime now = LocalDateTime.now();
+        List<String> keys;
+        switch (timeUnit) {
+            case DAYS:
+                time += LocalDateTimeUtil.between(now, LocalDateTime.now().with(LocalTime.MAX), ChronoUnit.SECONDS);
+                keys = getKeys(id, now.getDayOfMonth(), timeUnit);
+                break;
+            case HOURS:
+                time += (60 - now.getMinute()) * 60 - now.getSecond();
+                keys = getKeys(id, now.getHour(), timeUnit);
+                break;
+            case MINUTES:
+                time += 60 - now.getSecond();
+                keys = getKeys(id, now.getMinute(), timeUnit);
+                break;
+            case SECONDS:
+                return isAllowed(routeConfig, id);
+            default:
+                return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
+        }
+        return isAllowed(routeConfig, keys, time);
+    }
+
+    public Mono<Response> isAllowed(Config routeConfig, String id) {
         if (!this.initialized.get()) {
             throw new IllegalStateException("RedisRateLimiter is not initialized");
         }
@@ -91,16 +124,66 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
         return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
     }
 
+    public Mono<Response> isAllowed(Config routeConfig, List<String> keys, int time) {
+        if (!this.initialized.get()) {
+            throw new IllegalStateException("RedisRateLimiter is not initialized");
+        }
+
+        // How many requests per second do you want a user to be allowed to do?
+        int replenishRate = routeConfig.getReplenishRate();
+
+        // How much bursting do you want to allow?
+        int burstCapacity = routeConfig.getBurstCapacity();
+
+        // How many tokens are requested per request?
+        int requestedTokens = routeConfig.getRequestedTokens();
+
+        try {
+            // The arguments to the LUA script. time() returns unixtime in seconds.
+            List<String> scriptArgs = Arrays.asList(replenishRate + "", time + "");
+            // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
+            Flux<Long> flux = this.redisTemplate.execute(this.timeRedisScript, keys, scriptArgs);
+            // .log("redisratelimiter", Level.FINER);
+            return flux.next().onErrorResume(throwable -> {
+                if (log.isDebugEnabled()) {
+                    log.debug("Error calling rate limiter lua", throwable);
+                }
+                return Mono.just(-1L);
+            }).map(tokensLeft -> {
+                boolean allowed = tokensLeft < replenishRate;
+                Response response = new Response(allowed, getHeaders(routeConfig, tokensLeft));
+                if (log.isDebugEnabled()) {
+                    log.debug("response: " + response);
+                }
+                return response;
+            });
+        }
+        catch (Exception e) {
+            /*
+             * We don't want a hard dependency on Redis to allow traffic. Make sure to set
+             * an alert so you know if this is happening too much. Stripe's observed
+             * failure rate is 0.01%.
+             */
+            log.error("Error determining if user allowed from redis", e);
+        }
+        return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
+    }
+
     static List<String> getKeys(String id) {
         // use `{}` around keys to use Redis Key hash tags
         // this allows for using redis cluster
 
         // Make a unique key per user.
-        String prefix = "request_rate_limiter.{" + id;
+        String prefix = "rate_limiter.{" + id;
 
         // You need two Redis keys for Token Bucket.
         String tokenKey = prefix + "}.tokens";
         String timestampKey = prefix + "}.timestamp";
         return Arrays.asList(tokenKey, timestampKey);
     }
+
+    static List<String> getKeys(String id, int now, TimeUnit timeUnit) {
+        String tokenKey = "rate_limiter.{" + id + ':' + now + "}." + timeUnit.name();
+        return Arrays.asList(tokenKey);
+    }
 }