SpringCloud/spring-cloud-gateway
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make RedisRateLimiter omit headers if property is set

Browse Source 
Fixes #855
This commit is contained in:
Ronny Bräunlich
2019-02-27 08:15:52 +01:00
parent 9f9454b91f
commit eb0b10ccc6
2 changed files with 36 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java
View File

@@ -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<RedisRateLimiter.Config> implements ApplicationContextAware {
@@ -64,7 +66,7 @@ public class RedisRateLimiter extends AbstractRateLimiter<RedisRateLimiter.Confi
private String burstCapacityHeader = BURST_CAPACITY_HEADER;
public RedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate,
RedisScript<List<Long>> script, Validator validator) {
RedisScript<List<Long>> 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<RedisRateLimiter.Confi
Instant.now().getEpochSecond() + "", "1");
// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
Flux<List<Long>> 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<Long>(), (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<RedisRateLimiter.Confi
}
@NotNull
public HashMap<String, String> getHeaders(Config config, Long tokensLeft) {
HashMap<String, String> 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<String, String> getHeaders(Config config, Long tokensLeft) {
Map<String, String> 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;
}

26
spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java
View File

@@ -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)