diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java index 2593ccf7..6e74dc29 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayAutoConfiguration.java @@ -19,6 +19,7 @@ package org.springframework.cloud.gateway.config; import java.util.List; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.boot.actuate.endpoint.Endpoint; import org.springframework.boot.autoconfigure.AutoConfigureAfter; import org.springframework.boot.autoconfigure.AutoConfigureBefore; @@ -41,6 +42,7 @@ import org.springframework.cloud.gateway.filter.factory.RedirectToWebFilterFacto import org.springframework.cloud.gateway.filter.factory.RemoveNonProxyHeadersWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.RemoveRequestHeaderWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.RemoveResponseHeaderWebFilterFactory; +import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.RewritePathWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.SecureHeadersProperties; import org.springframework.cloud.gateway.filter.factory.SecureHeadersWebFilterFactory; @@ -48,6 +50,8 @@ import org.springframework.cloud.gateway.filter.factory.SetPathWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.SetResponseHeaderWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.SetStatusWebFilterFactory; import org.springframework.cloud.gateway.filter.factory.WebFilterFactory; +import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter; +import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter; import org.springframework.cloud.gateway.handler.FilteringWebHandler; import org.springframework.cloud.gateway.handler.NettyProxyWebHandler; import org.springframework.cloud.gateway.handler.RoutePredicateHandlerMapping; @@ -77,6 +81,12 @@ import org.springframework.context.annotation.Primary; import com.netflix.hystrix.HystrixObservableCommand; +import org.springframework.core.io.ClassPathResource; +import org.springframework.data.redis.core.RedisTemplate; +import org.springframework.data.redis.core.StringRedisTemplate; +import org.springframework.data.redis.core.script.DefaultRedisScript; +import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.scripting.support.ResourceScriptSource; import reactor.core.publisher.Flux; import reactor.ipc.netty.http.client.HttpClient; import rx.RxReactiveStreams; @@ -180,7 +190,7 @@ public class GatewayAutoConfiguration { return new WriteResponseFilter(); } - // Request Predicate beans + // Predicate Factory beans @Bean public AfterRoutePredicateFactory afterRoutePredicateFactory() { @@ -232,7 +242,7 @@ public class GatewayAutoConfiguration { return new RemoteAddrRoutePredicateFactory(); } - // Filter Factory beans + // WebFilter Factory beans @Bean public AddRequestHeaderWebFilterFactory addRequestHeaderWebFilterFactory() { @@ -283,6 +293,12 @@ public class GatewayAutoConfiguration { return new RemoveResponseHeaderWebFilterFactory(); } + @Bean + @ConditionalOnBean(RateLimiter.class) + public RequestRateLimiterWebFilterFactory requestRateLimiterWebFilterFactory(RateLimiter rateLimiter) { + return new RequestRateLimiterWebFilterFactory(rateLimiter); + } + @Bean public RewritePathWebFilterFactory rewritePathWebFilterFactory() { return new RewritePathWebFilterFactory(); @@ -309,6 +325,23 @@ public class GatewayAutoConfiguration { } + @ConditionalOnClass(RedisTemplate.class) + protected static class GatewayRedisConfiguration { + @Bean + public RedisScript redistRequestRateLimiterScript() { + DefaultRedisScript redisScript = new DefaultRedisScript<>(); + redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/scripts/request_rate_limiter.lua"))); + redisScript.setResultType(List.class); + return redisScript; + } + + @Bean + public RedisRateLimiter redisRateLimiter(StringRedisTemplate redisTemplate, + @Qualifier("redistRequestRateLimiterScript") RedisScript redisScript) { + return new RedisRateLimiter(redisTemplate, redisScript); + } + } + /*@Bean public RouterFunction test() { RouterFunction route = RouterFunctions.route( diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java index 7b5b3f47..32ce3489 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactory.java @@ -17,32 +17,22 @@ package org.springframework.cloud.gateway.filter.factory; -import java.time.Instant; -import java.util.Arrays; -import java.util.List; - -import org.apache.commons.logging.Log; -import org.apache.commons.logging.LogFactory; -import org.springframework.data.redis.core.StringRedisTemplate; -import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter; +import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter.Response; import org.springframework.http.HttpStatus; import org.springframework.tuple.Tuple; import org.springframework.web.server.WebFilter; /** - * Sample User Request Rate Throttle filter. + * User Request Rate Limiter filter. * See https://stripe.com/blog/rate-limiters and - * https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L11-L34 */ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory { - private Log log = LogFactory.getLog(getClass()); - private final StringRedisTemplate redisTemplate; - private final RedisScript script; + private final RateLimiter rateLimiter; - public RequestRateLimiterWebFilterFactory(StringRedisTemplate redisTemplate, RedisScript script) { - this.redisTemplate = redisTemplate; - this.script = script; + public RequestRateLimiterWebFilterFactory(RateLimiter rateLimiter) { + this.rateLimiter = rateLimiter; } @SuppressWarnings("unchecked") @@ -57,11 +47,11 @@ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory { return (exchange, chain) -> { // exchange.getPrincipal().flatMap(principal -> {}) //TODO: get user from request, maybe a KeyResolutionStrategy.resolve(exchange). Lookup strategy bean via arg - boolean allowed = isAllowed(replenishRate, capacity, "me"); + Response response = rateLimiter.isAllowed("me", replenishRate, capacity); //TODO: set some headers for rate, tokens left - if (allowed) { + if (response.isAllowed()) { return chain.filter(exchange); } exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS); @@ -69,33 +59,4 @@ public class RequestRateLimiterWebFilterFactory implements WebFilterFactory { }; } - //TODO: move to interface - //TODO: use tuple args except for id - /* for testing */ boolean isAllowed(int replenishRate, int capacity, String id) { - boolean allowed = false; - - try { - // # Make a unique key per user. - String prefix = "request_rate_limiter." + id; - - // # You need two Redis keys for Token Bucket. - List keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp"); - - // The arguments to the LUA script. time() returns unixtime in seconds. - String[] args = new String[]{ replenishRate+"", capacity+"", Instant.now().getEpochSecond()+"", "1"}; - // allowed, tokens_left = redis.eval(SCRIPT, keys, args) - List results = this.redisTemplate.execute(this.script, keys, args); - - allowed = new Long(1L).equals(results.get(0)); - Long tokensLeft = (Long) results.get(1); - - if (log.isDebugEnabled()) { - log.debug("isAllowed("+id+")=" + allowed + ", tokensLeft: "+tokensLeft); - } - - } catch (Exception e) { - log.error("Error determining if user allowed from redis", e); - } - return allowed; - } } 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 new file mode 100644 index 00000000..9a5644ed --- /dev/null +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RateLimiter.java @@ -0,0 +1,35 @@ +package org.springframework.cloud.gateway.filter.ratelimit; + +/** + * @author Spencer Gibb + */ +public interface RateLimiter { + Response isAllowed(String id, int replenishRate, int capacity); + + class Response { + private final boolean allowed; + private final long tokensRemaining; + + public Response(boolean allowed, long tokensRemaining) { + this.allowed = allowed; + this.tokensRemaining = tokensRemaining; + } + + public boolean isAllowed() { + return allowed; + } + + public long getTokensRemaining() { + return tokensRemaining; + } + + @Override + public String toString() { + final StringBuffer sb = new StringBuffer("Response{"); + sb.append("allowed=").append(allowed); + sb.append(", tokensRemaining=").append(tokensRemaining); + sb.append('}'); + return sb.toString(); + } + } +} 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 new file mode 100644 index 00000000..03b83949 --- /dev/null +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java @@ -0,0 +1,71 @@ +package org.springframework.cloud.gateway.filter.ratelimit; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.springframework.data.redis.core.StringRedisTemplate; +import org.springframework.data.redis.core.script.RedisScript; + +import java.time.Instant; +import java.util.Arrays; +import java.util.List; + +/** + * See https://stripe.com/blog/rate-limiters and + * https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L11-L34 + * + * @author Spencer Gibb + */ +public class RedisRateLimiter implements RateLimiter { + private Log log = LogFactory.getLog(getClass()); + + private final StringRedisTemplate redisTemplate; + private final RedisScript script; + + public RedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript script) { + this.redisTemplate = redisTemplate; + this.script = script; + } + + /** + * 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 capacity + * @param id + * @return + */ + @Override + //TODO: signature? params (tuple?). Return type, tokens left? + public Response isAllowed(String id, int replenishRate, int capacity) { + + try { + // Make a unique key per user. + String prefix = "request_rate_limiter." + id; + + // You need two Redis keys for Token Bucket. + List keys = Arrays.asList(prefix + ".tokens", prefix + ".timestamp"); + + // The arguments to the LUA script. time() returns unixtime in seconds. + String[] args = new String[]{ replenishRate+"", capacity+"", Instant.now().getEpochSecond()+"", "1"}; + // allowed, tokens_left = redis.eval(SCRIPT, keys, args) + List results = this.redisTemplate.execute(this.script, keys, args); + + boolean allowed = new Long(1L).equals(results.get(0)); + Long tokensLeft = (Long) results.get(1); + + Response response = new Response(allowed, 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 new Response(true, -1); + } +} diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java index e9bb2534..c50f6284 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterWebFilterFactoryTests.java @@ -1,26 +1,15 @@ package org.springframework.cloud.gateway.filter.factory; -import java.util.List; -import java.util.UUID; - import org.junit.Test; import org.junit.runner.RunWith; -import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.cloud.gateway.test.BaseWebClientTests; -import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Import; -import org.springframework.core.io.ClassPathResource; -import org.springframework.data.redis.core.StringRedisTemplate; -import org.springframework.data.redis.core.script.DefaultRedisScript; -import org.springframework.data.redis.core.script.RedisScript; -import org.springframework.scripting.support.ResourceScriptSource; import org.springframework.test.annotation.DirtiesContext; import org.springframework.test.context.junit4.SpringRunner; -import static org.assertj.core.api.Assertions.assertThat; import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT; /** @@ -28,25 +17,24 @@ import static org.springframework.boot.test.context.SpringBootTest.WebEnvironmen * @author Spencer Gibb */ @RunWith(SpringRunner.class) -@SpringBootTest(properties = "logging.level.org.springframework.cloud.gateway.filter.factory=DEBUG", - webEnvironment = RANDOM_PORT) +@SpringBootTest(webEnvironment = RANDOM_PORT) @DirtiesContext public class RequestRateLimiterWebFilterFactoryTests extends BaseWebClientTests { - @Autowired + /*@Autowired private StringRedisTemplate redisTemplate; @Autowired - private RedisScript script; + private RedisScript script;*/ @Test public void requestRateLimiterWebFilterFactoryWorks() throws Exception { - String id = UUID.randomUUID().toString(); + /*String id = UUID.randomUUID().toString(); RequestRateLimiterWebFilterFactory filterFactory = new RequestRateLimiterWebFilterFactory(this.redisTemplate, this.script); int replenishRate = 10; - int capacity = 5 * replenishRate; + int capacity = 2 * replenishRate; // Bursts work for (int i = 0; i < capacity; i++) { @@ -66,14 +54,14 @@ public class RequestRateLimiterWebFilterFactoryTests extends BaseWebClientTests } allowed = filterFactory.isAllowed(replenishRate, capacity, id); - assertThat(allowed).isFalse(); + assertThat(allowed).isFalse();*/ } @EnableAutoConfiguration @SpringBootConfiguration @Import(BaseWebClientTests.DefaultTestConfig.class) public static class TestConfig { - @Bean + /*@Bean public RedisScript requestRateLimiterScript() { DefaultRedisScript redisScript = new DefaultRedisScript<>(); redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/scripts/request_rate_limiter.lua"))); @@ -84,6 +72,6 @@ public class RequestRateLimiterWebFilterFactoryTests extends BaseWebClientTests @Bean public RequestRateLimiterWebFilterFactory requestRateLimiterWebFilterFactory(StringRedisTemplate redisTemplate) { return new RequestRateLimiterWebFilterFactory(redisTemplate, requestRateLimiterScript()); - } + }*/ } } 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 new file mode 100644 index 00000000..b3512d67 --- /dev/null +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java @@ -0,0 +1,66 @@ +package org.springframework.cloud.gateway.filter.ratelimit; + +import java.util.UUID; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter.Response; +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 static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT; + +/** + * see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L36-L62 + * @author Spencer Gibb + */ +@RunWith(SpringRunner.class) +@SpringBootTest(webEnvironment = RANDOM_PORT) +@DirtiesContext +public class RedisRateLimiterTests extends BaseWebClientTests { + + @Autowired + private RedisRateLimiter rateLimiter; + + @Test + public void requestRateLimiterWebFilterFactoryWorks() throws Exception { + String id = UUID.randomUUID().toString(); + + int replenishRate = 10; + int capacity = 2 * replenishRate; + + // Bursts work + for (int i = 0; i < capacity; i++) { + Response response = rateLimiter.isAllowed(id, replenishRate, capacity); + assertThat(response.isAllowed()).as("Burst # %s is allowed", i).isTrue(); + } + + Response response = rateLimiter.isAllowed(id, replenishRate, capacity); + assertThat(response.isAllowed()).as("Burst # %s is not allowed", capacity).isFalse(); + + Thread.sleep(1000); + + // # After the burst is done, check the steady state + for (int i = 0; i < replenishRate; i++) { + response = rateLimiter.isAllowed(id, replenishRate, capacity); + assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue(); + } + + response = rateLimiter.isAllowed(id, replenishRate, capacity); + assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse(); + } + + @EnableAutoConfiguration + @SpringBootConfiguration + @Import(BaseWebClientTests.DefaultTestConfig.class) + public static class TestConfig { + + } +} diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/AdhocTestSuite.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/AdhocTestSuite.java index 955bfc01..16ef7002 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/AdhocTestSuite.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/AdhocTestSuite.java @@ -37,6 +37,7 @@ import org.springframework.cloud.gateway.filter.factory.SetPathWebFilterFactoryI import org.springframework.cloud.gateway.filter.factory.SetPathWebFilterFactoryTests; import org.springframework.cloud.gateway.filter.factory.SetResponseWebFilterFactoryTests; import org.springframework.cloud.gateway.filter.factory.SetStatusWebFilterFactoryTests; +import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiterTests; import org.springframework.cloud.gateway.handler.predicate.AfterRoutePredicateFactoryTests; import org.springframework.cloud.gateway.handler.predicate.BeforeRoutePredicateFactoryTests; import org.springframework.cloud.gateway.handler.predicate.BetweenRoutePredicateFactoryTests; @@ -56,6 +57,7 @@ import static org.junit.Assume.assumeThat; @SuiteClasses({GatewayIntegrationTests.class, FormIntegrationTests.class, PostTests.class, + RedisRateLimiterTests.class, // route filter tests AddRequestHeaderWebFilterFactoryTests.class, AddRequestParameterWebFilterFactoryTests.class,