mirror of
https://gitee.com/dromara/RuoYi-Vue-Plus.git
synced 2026-09-20 02:08:15 +08:00
分钟、小时、天限流
This commit is contained in:
Wenchao Gong
parent
0271887c6b
commit
3930b70a2b
@ -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) {
|
||||
return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService);
|
||||
ConfigurationService configurationService, DefaultRedisScript<Long> timeRedisScript) {
|
||||
return new CustomerRedisRateLimiter(redisTemplate, redisScript, configurationService, timeRedisScript);
|
||||
}
|
||||
|
||||
@Bean
|
||||
|
||||
@ -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 rateLimiterAftersupplier 限流后操作(删除参数、添加隐藏参数，跳转)
|
||||
* @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 -> {
|
||||
private <T extends Object> Mono rateLimiter(ServerWebExchange exchange, IscRule rule, Route route, final int index,
|
||||
Supplier<Mono<T>> rateLimiterAfterSupplier) {
|
||||
final TimeUnit timeUnit = TIME_UNITS[index];
|
||||
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();
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user