diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java index 123d5301..41e43370 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java @@ -17,9 +17,6 @@ package org.springframework.cloud.gateway.filter.factory; -import java.util.Arrays; -import java.util.List; - import org.springframework.cloud.gateway.filter.GatewayFilter; import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver; import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter; @@ -27,69 +24,44 @@ import org.springframework.http.HttpStatus; import org.springframework.tuple.Tuple; /** - * User Request Rate Limiter filter. - * See https://stripe.com/blog/rate-limiters and + * User Request Rate Limiter filter. See https://stripe.com/blog/rate-limiters and */ public class RequestRateLimiterGatewayFilterFactory implements GatewayFilterFactory { - public static final String REPLENISH_RATE_KEY = "replenishRate"; - public static final String BURST_CAPACITY_KEY = "burstCapacity"; public static final String KEY_RESOLVER_KEY = "keyResolver"; private final RateLimiter rateLimiter; private final KeyResolver defaultKeyResolver; - public RequestRateLimiterGatewayFilterFactory(RateLimiter rateLimiter, KeyResolver defaultKeyResolver) { + public RequestRateLimiterGatewayFilterFactory(RateLimiter rateLimiter, + KeyResolver defaultKeyResolver) { this.rateLimiter = rateLimiter; this.defaultKeyResolver = defaultKeyResolver; } - @Override - public List argNames() { - return Arrays.asList(REPLENISH_RATE_KEY, BURST_CAPACITY_KEY, KEY_RESOLVER_KEY); - } - - public GatewayFilter apply(int replenishRate, int burstCapacity) { - return apply(replenishRate, burstCapacity, this.defaultKeyResolver); - } - - public GatewayFilter apply(int replenishRate, int burstCapacity, KeyResolver keyResolver) { - return (exchange, chain) -> - keyResolver.resolve(exchange).flatMap(key -> - //TODO: if key is empty? - rateLimiter.isAllowed(key, replenishRate, burstCapacity).flatMap(response -> { - //TODO: set some headers for rate, tokens left - - if (response.isAllowed()) { - return chain.filter(exchange); - } - exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS); - return exchange.getResponse().setComplete(); - })); - } - @SuppressWarnings("unchecked") @Override public GatewayFilter apply(Tuple args) { - // How many requests per second do you want a user to be allowed to do? - int replenishRate = args.getInt(REPLENISH_RATE_KEY); - - // How much bursting do you want to allow? - int burstCapacity; - if (args.hasFieldName(BURST_CAPACITY_KEY)) { - burstCapacity = args.getInt(BURST_CAPACITY_KEY); - } else { - burstCapacity = 0; - } KeyResolver keyResolver; if (args.hasFieldName(KEY_RESOLVER_KEY)) { keyResolver = args.getValue(KEY_RESOLVER_KEY, KeyResolver.class); - } else { + } + else { keyResolver = defaultKeyResolver; } - return apply(replenishRate, burstCapacity, keyResolver); + return (exchange, chain) -> keyResolver.resolve(exchange).flatMap(key -> + // TODO: if key is empty? + rateLimiter.isAllowed(key, args).flatMap(response -> { + // TODO: set some headers for rate, tokens left + + if (response.isAllowed()) { + return chain.filter(exchange); + } + exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS); + return exchange.getResponse().setComplete(); + })); } } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java index a45c49c1..874b0737 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java @@ -1,13 +1,14 @@ package org.springframework.cloud.gateway.filter.ratelimit; +import org.springframework.tuple.Tuple; import reactor.core.publisher.Mono; /** * @author Spencer Gibb */ public interface RateLimiter { - //TODO: move ints to tuple - Mono isAllowed(String id, int replenishRate, int burstCapacity); + + Mono isAllowed(String id, Tuple args); class Response { private final boolean allowed; diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java index 9852fba5..62d13714 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java @@ -4,13 +4,13 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.logging.Level; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.springframework.data.redis.core.ReactiveRedisTemplate; import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.tuple.Tuple; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -21,6 +21,9 @@ import reactor.core.publisher.Mono; * @author Spencer Gibb */ public class RedisRateLimiter implements RateLimiter { + public static final String REPLENISH_RATE_KEY = "replenishRate"; + public static final String BURST_CAPACITY_KEY = "burstCapacity"; + private Log log = LogFactory.getLog(getClass()); private final ReactiveRedisTemplate redisTemplate; @@ -36,15 +39,23 @@ public class RedisRateLimiter implements RateLimiter { * This uses a basic token bucket algorithm and relies on the fact that Redis scripts * execute atomically. No other operations can run between fetching the count and * writing the new count. - * @param replenishRate - * @param burstCapacity * @param id + * @param args * @return */ @Override - // TODO: signature? params (tuple?). @SuppressWarnings("unchecked") - public Mono isAllowed(String id, int replenishRate, int burstCapacity) { + public Mono isAllowed(String id, Tuple args) { + // How many requests per second do you want a user to be allowed to do? + int replenishRate = args.getInt(REPLENISH_RATE_KEY); + + // How much bursting do you want to allow? + int burstCapacity; + if (args.hasFieldName(BURST_CAPACITY_KEY)) { + burstCapacity = args.getInt(BURST_CAPACITY_KEY); + } else { + burstCapacity = 0; + } try { // Make a unique key per user. @@ -54,10 +65,10 @@ public class RedisRateLimiter implements RateLimiter { List keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp"); // The arguments to the LUA script. time() returns unixtime in seconds. - List args = Arrays.asList(replenishRate + "", burstCapacity + "", + List scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1"); // allowed, tokens_left = redis.eval(SCRIPT, keys, args) - Flux> flux = this.redisTemplate.execute(this.script, keys, args); + Flux> flux = this.redisTemplate.execute(this.script, keys, scriptArgs); // .log("redisratelimiter", Level.FINER); return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L))) .reduce(new ArrayList(), (longs, l) -> { diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java index 9e2d6493..949703bd 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java @@ -20,10 +20,14 @@ import org.springframework.mock.http.server.reactive.MockServerHttpRequest; import org.springframework.mock.web.server.MockServerWebExchange; import org.springframework.test.annotation.DirtiesContext; import org.springframework.test.context.junit4.SpringRunner; +import org.springframework.tuple.Tuple; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.when; import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY; +import static org.springframework.tuple.TupleBuilder.tuple; import reactor.core.publisher.Mono; @@ -67,7 +71,8 @@ public class RequestRateLimiterGatewayFilterFactoryTests extends BaseWebClientTe int replenishRate = 10; int burstCapacity = 2 * replenishRate; - when(rateLimiter.isAllowed(key, replenishRate, burstCapacity)) + Tuple args = tuple().of(REPLENISH_RATE_KEY, replenishRate, BURST_CAPACITY_KEY, burstCapacity); + when(rateLimiter.isAllowed(key, args)) .thenReturn(Mono.just(new Response(allowed, 1))); @@ -77,10 +82,9 @@ public class RequestRateLimiterGatewayFilterFactoryTests extends BaseWebClientTe when(this.filterChain.filter(exchange)).thenReturn(Mono.empty()); - Mono response = filterFactory.apply(replenishRate, burstCapacity, keyResolver).filter(exchange, this.filterChain); - response.subscribe(aVoid -> { - assertThat(exchange.getResponse().getStatusCode()).isEqualTo(expectedStatus); - }); + Mono response = filterFactory.apply(args).filter(exchange, this.filterChain); + response.subscribe(aVoid -> assertThat(exchange.getResponse().getStatusCode()) + .isEqualTo(expectedStatus)); } diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/PrincipalNameKeyResolverIntegrationTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/PrincipalNameKeyResolverIntegrationTests.java index 35f1ade2..305ee82c 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/PrincipalNameKeyResolverIntegrationTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/PrincipalNameKeyResolverIntegrationTests.java @@ -33,8 +33,8 @@ import org.springframework.web.bind.annotation.RestController; import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.DEFINED_PORT; import static org.springframework.cloud.gateway.filter.factory.GatewayFilters.prefixPath; -import static org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory.BURST_CAPACITY_KEY; -import static org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory.REPLENISH_RATE_KEY; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY; import static org.springframework.cloud.gateway.handler.predicate.RoutePredicates.path; import static org.springframework.tuple.TupleBuilder.tuple; import static org.springframework.web.reactive.function.client.ExchangeFilterFunctions.basicAuthentication; @@ -108,7 +108,7 @@ public class PrincipalNameKeyResolverIntegrationTests { @Bean RateLimiter rateLimiter() { - return (id, replenishRate, burstCapacity) -> Mono.just(new RateLimiter.Response(true, Long.MAX_VALUE)); + return (id, args) -> Mono.just(new RateLimiter.Response(true, Long.MAX_VALUE)); } @Bean diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java index e78b0013..9de79910 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java @@ -13,9 +13,13 @@ import org.springframework.cloud.gateway.test.BaseWebClientTests; import org.springframework.context.annotation.Import; import org.springframework.test.annotation.DirtiesContext; import org.springframework.test.context.junit4.SpringRunner; +import org.springframework.tuple.Tuple; import static org.assertj.core.api.Assertions.assertThat; import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.BURST_CAPACITY_KEY; +import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.REPLENISH_RATE_KEY; +import static org.springframework.tuple.TupleBuilder.tuple; /** * see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L36-L62 @@ -36,15 +40,17 @@ public class RedisRateLimiterTests extends BaseWebClientTests { int replenishRate = 10; int burstCapacity = 2 * replenishRate; + Tuple args = tuple().of(REPLENISH_RATE_KEY, replenishRate, BURST_CAPACITY_KEY, burstCapacity); + // Bursts work for (int i = 0; i < burstCapacity; i++) { - Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block(); + Response response = rateLimiter.isAllowed(id, args).block(); assertThat(response.isAllowed()).as("Burst # %s is allowed", i).isTrue(); } - Response response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block(); + Response response = rateLimiter.isAllowed(id, args).block(); if (response.isAllowed()) { //TODO: sometimes there is an off by one error - response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block(); + response = rateLimiter.isAllowed(id, args).block(); } assertThat(response.isAllowed()).as("Burst # %s is not allowed", burstCapacity).isFalse(); @@ -52,11 +58,11 @@ public class RedisRateLimiterTests extends BaseWebClientTests { // # After the burst is done, check the steady state for (int i = 0; i < replenishRate; i++) { - response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block(); + response = rateLimiter.isAllowed(id, args).block(); assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue(); } - response = rateLimiter.isAllowed(id, replenishRate, burstCapacity).block(); + response = rateLimiter.isAllowed(id, args).block(); assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse(); }