mirror of
https://gitee.com/dromara/RuoYi-Vue-Plus.git
synced 2026-09-20 02:08:15 +08:00
QPS 限流
This commit is contained in:
Wenchao Gong
parent
71ee37556d
commit
0271887c6b
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -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,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<String, String> 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<Mono<Void>> 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<String> 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<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 &&
|
||||
accessKeyName.equals(param[0])).map(param -> param[1]).collect(Collectors.toList()),
|
||||
route, request, accessKeyName);
|
||||
final List<String[]> 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<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;
|
||||
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<List<String>> valueSupplier, Route route,
|
||||
ServerHttpRequest request, String accessKeyName) {
|
||||
private IscRule handleRule(String headerAk, Supplier<List<String>> 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 <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();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user