From 0271887c6bc4ddc0af07ed1ffa3a35ebdfedfb24 Mon Sep 17 00:00:00 2001 From: Wenchao Gong Date: Sun, 17 Oct 2021 01:24:30 +0800 Subject: [PATCH] =?UTF-8?q?QPS=20=E9=99=90=E6=B5=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ruoyi/gateway/config/GatewayConfig.java | 18 ++- .../gateway/filter/CustomerGlobalFilter.java | 148 ++++++++++++------ .../ratelimit/CustomerRedisRateLimiter.java | 106 +++++++++++++ 3 files changed, 223 insertions(+), 49 deletions(-) create mode 100644 ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java diff --git a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/config/GatewayConfig.java b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/config/GatewayConfig.java index 0836e7c70..d6ee44530 100644 --- a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/config/GatewayConfig.java +++ b/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> redisScript, + ConfigurationService configurationService) { + return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService); + } + + @Bean + public GlobalFilter customerGlobalFilter(CustomerRedisRateLimiter customerRedisRateLimiter) { - return new CustomerGlobalFilter(); + return new CustomerGlobalFilter(customerRedisRateLimiter); } } diff --git a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java index b71b86929..5bb256884 100644 --- a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java +++ b/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 filter(ServerWebExchange exchange, GatewayFilterChain chain) { final Route route = exchange.getRequiredAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR); @@ -56,22 +64,26 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered { String headerAk = GatewayUtils.getValue(null, () -> request.getHeaders().get(accessKeyName)); if (HttpMethod.GET.equals(httpMethod)) { final MultiValueMap queryParams = new LinkedMultiValueMap<>(request.getQueryParams()); - handleRule(headerAk, () -> queryParams.get(accessKeyName), route, request, accessKeyName); - queryParams.remove(accessKeyName); - handleHiddenParams(route, queryParams, (next, map) -> { - Object value; - if (Objects.isNull(value = next.getValue())) { - map.add(next.getKey(), StrUtil.EMPTY); - } else if(value instanceof JSONArray) { - map.put(next.getKey(), ((JSONArray)value).toList(String.class)); - } else { - map.add(next.getKey(), value.toString()); - } - }); - URI newUri = UriComponentsBuilder.fromUri(request.getURI()) - .replaceQueryParams(unmodifiableMultiValueMap(queryParams)).build().toUri(); - ServerHttpRequest updatedRequest = exchange.getRequest().mutate().uri(newUri).build(); - return chain.filter(exchange.mutate().request(updatedRequest).build()); + final IscRule rule = handleRule(headerAk, () -> queryParams.get(accessKeyName), route); + Supplier> rateLimiterAftersupplier = () -> { + removeParam(headerAk, request, accessKeyName, () -> queryParams.remove(accessKeyName)); + handleHiddenParams(route, queryParams, (next, map) -> { + Object value; + if (Objects.isNull(value = next.getValue())) { + map.add(next.getKey(), StrUtil.EMPTY); + } else if (value instanceof JSONArray) { + map.put(next.getKey(), ((JSONArray) value).toList(String.class)); + } else { + map.add(next.getKey(), value.toString()); + } + }); + URI newUri = UriComponentsBuilder.fromUri(request.getURI()) + .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 modifiedBody = serverRequest.bodyToMono(String.class); @@ -79,34 +91,41 @@ 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)), - route, request, accessKeyName); - jsonObj.remove(accessKeyName); - handleHiddenParams(route, jsonObj, (next, map) -> { - map.set(next.getKey(), next.getValue()); - }); - return Mono.just(jsonObj.toString()); + final IscRule rule = handleRule(headerAk, () -> + Arrays.asList(jsonObj.get(accessKeyName, String.class, true)), route); + Supplier> 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 stream = Arrays.stream(body.split("&")).map(param -> param.split("=")); - handleRule(headerAk, () -> stream.filter(param -> param.length > 0 && - accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()), - route, request, accessKeyName); - final List params = stream.filter(param -> !accessKeyName.equals(param[0])).collect(Collectors.toList()); - handleHiddenParams(route, params, (next, list) -> { - Object value; - if (Objects.isNull(value = next.getValue())) { - list.add(new String[]{next.getKey(), StrUtil.EMPTY}); - } else if(value instanceof JSONArray) { - ((JSONArray)value).stream().map(o -> String.valueOf(o)).forEach(v -> { - list.add(new String[]{next.getKey(), v}); - }); - } else { - list.add(new String[]{next.getKey(), next.getValue().toString()}); - } + final IscRule rule = handleRule(headerAk, () -> stream.filter(param -> param.length > 0 && + accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()), route); + Supplier> rateLimiterAftersupplier = () -> { + removeParam(headerAk, request, accessKeyName, null); + final List params = stream.filter(param -> !accessKeyName.equals(param[0])).collect(Collectors.toList()); + handleHiddenParams(route, params, (next, list) -> { + Object value; + if (Objects.isNull(value = next.getValue())) { + list.add(new String[]{next.getKey(), StrUtil.EMPTY}); + } else if (value instanceof JSONArray) { + ((JSONArray) value).stream().map(o -> String.valueOf(o)).forEach(v -> { + list.add(new String[]{next.getKey(), v}); + }); + } 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 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> valueSupplier, Route route, - ServerHttpRequest request, String accessKeyName) { + private IscRule handleRule(String headerAk, Supplier> valueSupplier, Route route) { //获取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 + */ + private void removeParam(String headerAk, ServerHttpRequest request, String accessKeyName, Supplier 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 + * @return + */ + private Mono rateLimiter(ServerWebExchange exchange, IscRule rule, Route route, + Supplier> rateLimiterAftersupplier) { + final RedisRateLimiter.Config config = new RedisRateLimiter.Config().setReplenishRate(1); + return rateLimiter.isAllowed(config, rule.getId() + ':' + route.getId()).flatMap(response -> { + for (Map.Entry 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(); + }); + } } diff --git a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java new file mode 100644 index 000000000..cd388ffca --- /dev/null +++ b/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> script; + + private AtomicBoolean initialized = new AtomicBoolean(false); + + public CustomerRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate, RedisScript> script, + ConfigurationService configurationService) { + super(redisTemplate, script, configurationService); + this.redisTemplate = redisTemplate; + this.redisTemplate = redisTemplate; + this.script = script; + this.initialized.compareAndSet(false, true); + } + + public Mono 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 keys = getKeys(id); + + // The arguments to the LUA script. time() returns unixtime in seconds. + List scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", + Instant.now().getEpochSecond() + "", requestedTokens + ""); + // allowed, tokens_left = redis.eval(SCRIPT, keys, args) + Flux> 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(), (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 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); + } +}