QPS 限流

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

@@ -2,18 +2,25 @@ package com.ruoyi.gateway.config;
 
 import com.ruoyi.gateway.config.provider.RedisRouteDefinitionRepository;
 import com.ruoyi.gateway.filter.CustomerGlobalFilter;
+import com.ruoyi.gateway.ratelimit.CustomerRedisRateLimiter;
 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.RedisScript;
 import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
 import org.springframework.data.redis.serializer.StringRedisSerializer;
 
 import javax.annotation.Resource;
+import java.util.List;
 
 /**
  * Gateway 配置文件
@@ -56,8 +63,15 @@ public class GatewayConfig
     }
 
     @Bean
-    public GlobalFilter customerGlobalFilter()
+    public CustomerRedisRateLimiter customerRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate,
+             @Qualifier(RedisRateLimiter.REDIS_SCRIPT_NAME) RedisScript<List<Long>> redisScript,
+             ConfigurationService configurationService) {
+        return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService);
+    }
+
+    @Bean
+    public GlobalFilter customerGlobalFilter(CustomerRedisRateLimiter customerRedisRateLimiter)
     {
-        return new CustomerGlobalFilter();
+        return new CustomerGlobalFilter(customerRedisRateLimiter);
     }
 }

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

@@ -4,19 +4,21 @@ import cn.hutool.core.util.StrUtil;
 import cn.hutool.json.JSONArray;
 import cn.hutool.json.JSONObject;
 import cn.hutool.json.JSONUtil;
+import com.ruoyi.gateway.ratelimit.CustomerRedisRateLimiter;
 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;
 import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
-import org.springframework.util.MultiValueMapAdapter;
 import org.springframework.util.StringUtils;
 import org.springframework.web.reactive.function.server.HandlerStrategies;
 import org.springframework.web.reactive.function.server.ServerRequest;
@@ -31,6 +33,7 @@ import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;
 import static org.springframework.util.CollectionUtils.unmodifiableMultiValueMap;
 
 /**
@@ -45,6 +48,11 @@ import static org.springframework.util.CollectionUtils.unmodifiableMultiValueMap
  * @date 2021-10-15
  */
 public class CustomerGlobalFilter implements GlobalFilter, Ordered {
+    private CustomerRedisRateLimiter rateLimiter;
+    public CustomerGlobalFilter(CustomerRedisRateLimiter rateLimiter) {
+        this.rateLimiter = rateLimiter;
+    }
+
     @Override
     public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
         final Route route = exchange.getRequiredAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
@@ -56,8 +64,9 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
         String headerAk = GatewayUtils.getValue(null, () -> request.getHeaders().get(accessKeyName));
         if (HttpMethod.GET.equals(httpMethod)) {
             final MultiValueMap<String, String> queryParams = new LinkedMultiValueMap<>(request.getQueryParams());
-            handleRule(headerAk, () -> queryParams.get(accessKeyName), route, request, accessKeyName);
+            final IscRule rule = handleRule(headerAk, () -> queryParams.get(accessKeyName), route);
-            queryParams.remove(accessKeyName);
+            Supplier<Mono<Void>> rateLimiterAftersupplier = () -> {
+                removeParam(headerAk, request, accessKeyName, () -> queryParams.remove(accessKeyName));
                 handleHiddenParams(route, queryParams, (next, map) -> {
                     Object value;
                     if (Objects.isNull(value = next.getValue())) {
@@ -72,6 +81,9 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                         .replaceQueryParams(unmodifiableMultiValueMap(queryParams)).build().toUri();
                 ServerHttpRequest updatedRequest = exchange.getRequest().mutate().uri(newUri).build();
                 return chain.filter(exchange.mutate().request(updatedRequest).build());
+            };
+
+            return rateLimiter(exchange, rule, route, rateLimiterAftersupplier);
         } else if (HttpMethod.POST.equals(httpMethod)) {
             final ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
             final Mono<String> modifiedBody = serverRequest.bodyToMono(String.class);
@@ -79,19 +91,24 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                 MediaType mediaType = request.getHeaders().getContentType();
                 if (MediaType.APPLICATION_JSON.equals(mediaType)) {
                     JSONObject jsonObj = JSONUtil.parseObj(body);
-                    handleRule(headerAk, () -> Arrays.asList(jsonObj.get(accessKeyName, String.class, true)),
+                    final IscRule rule = handleRule(headerAk, () ->
-                            route, request, accessKeyName);
+                            Arrays.asList(jsonObj.get(accessKeyName, String.class, true)), route);
-                    jsonObj.remove(accessKeyName);
+                    Supplier<Mono<String>> rateLimiterAftersupplier = () -> {
+                        removeParam(headerAk, request, accessKeyName, () -> jsonObj.remove(accessKeyName));
                         handleHiddenParams(route, jsonObj, (next, map) -> {
                             map.set(next.getKey(), next.getValue());
                         });
                         return Mono.just(jsonObj.toString());
+                    };
+
+                    return rateLimiter(exchange, rule, route, 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("="));
-                        handleRule(headerAk, () -> stream.filter(param -> param.length > 0 &&
+                        final IscRule rule = handleRule(headerAk, () -> stream.filter(param -> param.length > 0 &&
-                                        accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()),
+                                        accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()), route);
-                                route, request, accessKeyName);
+                        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) -> {
                                 Object value;
@@ -104,9 +121,11 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                                 } else {
                                     list.add(new String[]{next.getKey(), next.getValue().toString()});
                                 }
-
                             });
                             return Mono.just(params.stream().map(param -> param[0] + '=' + param[1]).collect(Collectors.joining("&")));
+                        };
+
+                        return rateLimiter(exchange, rule, route, rateLimiterAftersupplier);
                     }
                 }
                 return Mono.empty();
@@ -127,12 +146,9 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
      * @param headerAk
      * @param valueSupplier
      * @param route
-     * @param request
-     * @param accessKeyName
      * @return
      */
-    private String handleRule(String headerAk, Supplier<List<String>> valueSupplier, Route route,
+    private IscRule handleRule(String headerAk, Supplier<List<String>> valueSupplier, Route route) {
-                              ServerHttpRequest request, String accessKeyName) {
         //获取AK
         final String ak = GatewayUtils.getValue(headerAk, valueSupplier, () -> new RuntimeException("AK 不存在"));
         //获取规则
@@ -140,13 +156,27 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                 () -> new RuntimeException("AK异常"));
         //是否到期
         GatewayUtils.isBefore(rule, () -> new RuntimeException("AK已过期"));
-        //TODO 限流
+        //设置AK 到ID 为了传参方便
+        rule.setId(ak);
+        return rule;
+    }
 
+    /**
+     * 删除参数
+     * @param headerAk
+     * @param request
+     * @param accessKeyName
+     * @param removeSupplier
+     * @param <T>
+     */
+    private <T> void removeParam(String headerAk, ServerHttpRequest request, String accessKeyName, Supplier<T> removeSupplier) {
         //如果header中有AK,则删除
         if (Objects.nonNull(headerAk)) {
             request.getHeaders().remove(accessKeyName);
         }
-        return ak;
+        if(Objects.nonNull(removeSupplier)) {
+            removeSupplier.get();
+        }
     }
 
     /**
@@ -172,4 +202,28 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
             mapper.accept(next, result);
         }
     }
+
+    /**
+     * 限流
+     * @param exchange
+     * @param rule
+     * @Param route
+     * @param rateLimiterAftersupplier 限流后操作(删除参数、添加隐藏参数，跳转)
+     * @param <T>
+     * @return
+     */
+    private <T extends Object> Mono rateLimiter(ServerWebExchange exchange, IscRule rule, Route route,
+                                                Supplier<Mono<T>> rateLimiterAftersupplier) {
+        final RedisRateLimiter.Config config = new RedisRateLimiter.Config().setReplenishRate(1);
+        return rateLimiter.isAllowed(config, rule.getId() + ':' + route.getId()).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();
+            }
+            setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
+            return exchange.getResponse().setComplete();
+        });
+    }
 }

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

@@ -0,0 +1,106 @@
+package com.ruoyi.gateway.ratelimit;
+
+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.RedisScript;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * @author Wenchao Gong
+ * @date 2021-10-16
+ */
+@Slf4j
+public class CustomerRedisRateLimiter extends RedisRateLimiter {
+
+    private ReactiveStringRedisTemplate redisTemplate;
+
+    private RedisScript<List<Long>> script;
+
+    private AtomicBoolean initialized = new AtomicBoolean(false);
+
+    public CustomerRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate, RedisScript<List<Long>> script,
+                                    ConfigurationService configurationService) {
+        super(redisTemplate, script, configurationService);
+        this.redisTemplate = redisTemplate;
+        this.redisTemplate = redisTemplate;
+        this.script = script;
+        this.initialized.compareAndSet(false, true);
+    }
+
+    public Mono<Response> isAllowed(Config routeConfig, String id)
+    {
+        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 {
+            List<String> keys = getKeys(id);
+
+            // The arguments to the LUA script. time() returns unixtime in seconds.
+            List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "",
+                    Instant.now().getEpochSecond() + "", requestedTokens + "");
+            // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
+            Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
+            // .log("redisratelimiter", Level.FINER);
+            return flux.onErrorResume(throwable -> {
+                if (log.isDebugEnabled()) {
+                    log.debug("Error calling rate limiter lua", throwable);
+                }
+                return Flux.just(Arrays.asList(1L, -1L));
+            }).reduce(new ArrayList<Long>(), (longs, l) -> {
+                longs.addAll(l);
+                return longs;
+            }).map(results -> {
+                boolean allowed = results.get(0) == 1L;
+                Long tokensLeft = results.get(1);
+
+                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;
+
+        // You need two Redis keys for Token Bucket.
+        String tokenKey = prefix + "}.tokens";
+        String timestampKey = prefix + "}.timestamp";
+        return Arrays.asList(tokenKey, timestampKey);
+    }
+}