From eb0b10ccc6dba136f2080c8caa90f05a9b9d49b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ronny=20Br=C3=A4unlich?= Date: Wed, 27 Feb 2019 08:15:52 +0100 Subject: [PATCH] Make RedisRateLimiter omit headers if property is set Fixes #855 --- .../filter/ratelimit/RedisRateLimiter.java | 20 ++++++++------ .../ratelimit/RedisRateLimiterTests.java | 26 +++++++++++++++++-- 2 files changed, 36 insertions(+), 10 deletions(-) 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 791718545..7fb2e2895 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 @@ -5,6 +5,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; import javax.validation.constraints.Min; @@ -29,6 +30,7 @@ import org.springframework.validation.annotation.Validated; * https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L11-L34 * * @author Spencer Gibb + * @author Ronny Bräunlich */ @ConfigurationProperties("spring.cloud.gateway.redis-rate-limiter") public class RedisRateLimiter extends AbstractRateLimiter implements ApplicationContextAware { @@ -64,7 +66,7 @@ public class RedisRateLimiter extends AbstractRateLimiter redisTemplate, - RedisScript> script, Validator validator) { + RedisScript> script, Validator validator) { super(Config.class, CONFIGURATION_PROPERTY_NAME, validator); this.redisTemplate = redisTemplate; this.script = script; @@ -159,12 +161,12 @@ public class RedisRateLimiter extends AbstractRateLimiter> flux = this.redisTemplate.execute(this.script, keys, scriptArgs); - // .log("redisratelimiter", Level.FINER); + // .log("redisratelimiter", Level.FINER); return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L))) .reduce(new ArrayList(), (longs, l) -> { longs.addAll(l); return longs; - }) .map(results -> { + }).map(results -> { boolean allowed = results.get(0) == 1L; Long tokensLeft = results.get(1); @@ -188,11 +190,13 @@ public class RedisRateLimiter extends AbstractRateLimiter getHeaders(Config config, Long tokensLeft) { - HashMap headers = new HashMap<>(); - headers.put(this.remainingHeader, tokensLeft.toString()); - headers.put(this.replenishRateHeader, String.valueOf(config.getReplenishRate())); - headers.put(this.burstCapacityHeader, String.valueOf(config.getBurstCapacity())); + public Map getHeaders(Config config, Long tokensLeft) { + Map headers = new HashMap<>(); + if (isIncludeHeaders()) { + headers.put(this.remainingHeader, tokensLeft.toString()); + headers.put(this.replenishRateHeader, String.valueOf(config.getReplenishRate())); + headers.put(this.burstCapacityHeader, String.valueOf(config.getBurstCapacity())); + } return headers; } 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 5b6b4a0fd..6608da176 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 @@ -5,6 +5,7 @@ import java.util.UUID; import org.junit.Rule; 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; @@ -25,6 +26,7 @@ import static org.springframework.boot.test.context.SpringBootTest.WebEnvironmen /** * see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#file-1-check_request_rate_limiter-rb-L36-L62 * @author Spencer Gibb + * @author Ronny Bräunlich */ @RunWith(SpringRunner.class) @SpringBootTest(webEnvironment = RANDOM_PORT) @@ -71,7 +73,7 @@ public class RedisRateLimiterTests extends BaseWebClientTests { Thread.sleep(1000); - // # After the burst is done, check the steady state + // # After the burst is done, check the steady state for (int i = 0; i < replenishRate; i++) { response = rateLimiter.isAllowed(routeId, id).block(); assertThat(response.isAllowed()).as("steady state # %s is allowed", i).isTrue(); @@ -80,13 +82,33 @@ public class RedisRateLimiterTests extends BaseWebClientTests { response = rateLimiter.isAllowed(routeId, id).block(); assertThat(response.isAllowed()).as("steady state # %s is allowed", replenishRate).isFalse(); } - + @Test public void keysUseRedisKeyHashTags() { assertThat(RedisRateLimiter.getKeys("1")) .containsExactly("request_rate_limiter.{1}.tokens", "request_rate_limiter.{1}.timestamp"); } + @Test + public void redisRateLimiterDoesNotSendHeadersIfDeactivated() throws Exception { + assumeThat("Ignore on Circle", + System.getenv("CIRCLECI"), is(nullValue())); + + String id = UUID.randomUUID().toString(); + String routeId = "myroute"; + + rateLimiter.setIncludeHeaders(false); + + Response response = rateLimiter.isAllowed(routeId, id).block(); + assertThat(response.isAllowed()).isTrue(); + assertThat(response.getHeaders()) + .doesNotContainKey(RedisRateLimiter.REMAINING_HEADER); + assertThat(response.getHeaders()). + doesNotContainKey(RedisRateLimiter.REPLENISH_RATE_HEADER); + assertThat(response.getHeaders()). + doesNotContainKey(RedisRateLimiter.BURST_CAPACITY_HEADER); + } + @EnableAutoConfiguration @SpringBootConfiguration @Import(BaseWebClientTests.DefaultTestConfig.class)