Merge pull request #4072 from jmrt47/add-throw-on-limit-option

Add Exception on Limit Exceed Switch to RequestRateLimiterGatewayFilterFactory

Author: ryanjbaxter

docs/modules/ROOT/pages/spring-cloud-gateway-server-webflux/gatewayfilter-factories/requestratelimiter-factory.adoc

@@ -2,9 +2,13 @@
 = `RequestRateLimiter` `GatewayFilter` Factory
 
 The `RequestRateLimiter` `GatewayFilter` factory uses a `RateLimiter` implementation to determine if the current request is allowed to proceed. If it is not, a status of `HTTP 429 - Too Many Requests` (by default) is returned.
+The default status code can be configured using the `StatusCode` property of the `RequestRateLimiter` filter.
 
 This filter takes an optional `keyResolver` parameter and parameters specific to the rate limiter (described xref:spring-cloud-gateway-server-webflux/gatewayfilter-factories/requestratelimiter-factory.adoc#key-resolver-section[later in this section]).
 
+It also supports the `throwOnLimit` option, which is `false` by default. When set to `true`, the filter will throw a `HttpClientErrorException` when the request is denied by the rate limiter, instead of just setting the response status.
+The `HttpClientErrorException` is created with the configured status code from the `StatusCode` property. Depending on the status code, Spring may return a more specific subclass (for example, if `StatusCode` is set to `HTTP 429 - Too Many Requests`, the exception will be a `HttpClientErrorException.TooManyRequests`).
+
 `keyResolver` is a bean that implements the `KeyResolver` interface.
 In configuration, reference the bean by name using SpEL.
 `#{@myKeyResolver}` is a SpEL expression that references a bean named `myKeyResolver`.

spring-cloud-gateway-server-webflux/src/main/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactory.java

@@ -20,6 +20,7 @@
 import java.util.Objects;
 
 import org.jspecify.annotations.Nullable;
+import reactor.core.publisher.Mono;
 
 import org.springframework.boot.context.properties.ConfigurationProperties;
 import org.springframework.cloud.gateway.filter.GatewayFilter;
@@ -30,6 +31,7 @@
 import org.springframework.cloud.gateway.support.HttpStatusHolder;
 import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
 import org.springframework.http.HttpStatus;
+import org.springframework.web.client.HttpClientErrorException;
 
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;
 
@@ -60,6 +62,12 @@ public class RequestRateLimiterGatewayFilterFactory
 	/** HttpStatus to return when denyEmptyKey is true, defaults to FORBIDDEN. */
 	private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();
 
+	/**
+	 * Switch to throw a {@link HttpClientErrorException} when the request is denied by
+	 * the RateLimiter, defaults to false.
+	 */
+	private boolean throwOnLimit = false;
+
 	public RequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
 		super(Config.class);
 		this.defaultRateLimiter = defaultRateLimiter;
@@ -90,6 +98,14 @@ public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
 		this.emptyKeyStatusCode = emptyKeyStatusCode;
 	}
 
+	public boolean isThrowOnLimit() {
+		return throwOnLimit;
+	}
+
+	public void setThrowOnLimit(boolean throwOnLimit) {
+		this.throwOnLimit = throwOnLimit;
+	}
+
 	@SuppressWarnings("unchecked")
 	@Override
 	public GatewayFilter apply(Config config) {
@@ -98,6 +114,7 @@ public GatewayFilter apply(Config config) {
 		boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
 		HttpStatusHolder emptyKeyStatus = HttpStatusHolder
 			.parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));
+		boolean throwLimit = getOrDefault(config.throwOnLimit, this.throwOnLimit);
 
 		return (exchange, chain) -> resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
 			if (EMPTY_KEY.equals(key)) {
@@ -122,6 +139,11 @@ public GatewayFilter apply(Config config) {
 					return chain.filter(exchange);
 				}
 
+				if (throwLimit) {
+					return Mono.error(HttpClientErrorException.create(config.getStatusCode(), "Too Many Requests",
+							exchange.getResponse().getHeaders(), null, null));
+				}
+
 				setResponseStatus(exchange, config.getStatusCode());
 				return exchange.getResponse().setComplete();
 			});
@@ -144,6 +166,8 @@ public static class Config implements HasRouteId {
 
 		private @Nullable String emptyKeyStatus;
 
+		private @Nullable Boolean throwOnLimit;
+
 		private @Nullable String routeId;
 
 		public @Nullable KeyResolver getKeyResolver() {
@@ -191,6 +215,15 @@ public Config setEmptyKeyStatus(String emptyKeyStatus) {
 			return this;
 		}
 
+		public @Nullable Boolean getThrowOnLimit() {
+			return throwOnLimit;
+		}
+
+		public Config setThrowOnLimit(Boolean throwOnLimit) {
+			this.throwOnLimit = throwOnLimit;
+			return this;
+		}
+
 		@Override
 		public void setRouteId(String routeId) {
 			this.routeId = routeId;

spring-cloud-gateway-server-webflux/src/test/java/org/springframework/cloud/gateway/filter/factory/RequestRateLimiterGatewayFilterFactoryTests.java

@@ -21,6 +21,7 @@
 
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -44,6 +45,7 @@
 import org.springframework.mock.web.server.MockServerWebExchange;
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.bean.override.mockito.MockitoBean;
+import org.springframework.web.client.HttpClientErrorException;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
@@ -83,22 +85,32 @@ public void notAllowedWorks() {
 		assertFilterFactory(resolver2, "notallowedkey", false, HttpStatus.TOO_MANY_REQUESTS);
 	}
 
+	@Test
+	public void notAllowedThrows() {
+		assertFilterFactory(null, "allowedkey", false, HttpStatus.TOO_MANY_REQUESTS, null, true);
+	}
+
 	@Test
 	public void emptyKeyDenied() {
 		assertFilterFactory(exchange -> Mono.empty(), null, true, HttpStatus.FORBIDDEN);
 	}
 
 	@Test
 	public void emptyKeyAllowed() {
-		assertFilterFactory(exchange -> Mono.empty(), null, true, HttpStatus.OK, false);
+		assertFilterFactory(exchange -> Mono.empty(), null, true, HttpStatus.OK, false, false);
+	}
+
+	@Test
+	public void emptyKeyDeniedWithThrowOnLimit() {
+		assertFilterFactory(exchange -> Mono.empty(), null, true, HttpStatus.FORBIDDEN, true, true);
 	}
 
 	private void assertFilterFactory(KeyResolver keyResolver, String key, boolean allowed, HttpStatus expectedStatus) {
-		assertFilterFactory(keyResolver, key, allowed, expectedStatus, null);
+		assertFilterFactory(keyResolver, key, allowed, expectedStatus, null, false);
 	}
 
 	private void assertFilterFactory(KeyResolver keyResolver, String key, boolean allowed, HttpStatus expectedStatus,
-			Boolean denyEmptyKey) {
+			Boolean denyEmptyKey, Boolean throwOnLimit) {
 
 		String tokensRemaining = allowed ? "1" : "0";
 
@@ -122,18 +134,40 @@ private void assertFilterFactory(KeyResolver keyResolver, String key, boolean al
 		if (denyEmptyKey != null) {
 			factory.setDenyEmptyKey(denyEmptyKey);
 		}
+		if (throwOnLimit != null) {
+			factory.setThrowOnLimit(throwOnLimit);
+		}
 		GatewayFilter filter = factory.apply(config -> {
 			config.setRouteId("myroute");
 			config.setKeyResolver(keyResolver);
 		});
 
-		Mono<Void> response = filter.filter(exchange, this.filterChain);
-		response.subscribe(aVoid -> {
-			assertThat(exchange.getResponse().getStatusCode()).isEqualTo(expectedStatus);
-			assertThat(exchange.getResponse().getHeaders().asMultiValueMap()).containsEntry("X-Tokens-Remaining",
-					Collections.singletonList(tokensRemaining));
-		});
+		StepVerifier.FirstStep<Void> voidFirstStep = StepVerifier.create(filter.filter(exchange, this.filterChain));
 
+		if (key == null) {
+			// empty key case won't contain the header, so we just check the status code
+			voidFirstStep.consumeSubscriptionWith(aVoid -> {
+				assertThat(exchange.getResponse().getStatusCode()).isEqualTo(expectedStatus);
+			}).expectComplete().verify();
+		}
+		else if (throwOnLimit != null && throwOnLimit && !allowed) {
+			// if throwOnLimit is true and the request is denied, we expect an error instead of a complete signal
+			voidFirstStep.consumeErrorWith(throwable -> {
+				assertThat(throwable).isInstanceOf(HttpClientErrorException.class);
+				HttpClientErrorException ex = (HttpClientErrorException) throwable;
+				assertThat(ex.getStatusCode()).isEqualTo(expectedStatus);
+				assertThat(ex.getResponseHeaders().toSingleValueMap()).containsEntry("X-Tokens-Remaining",
+						tokensRemaining);
+			}).verify();
+		}
+		else {
+			// allowed and denied, we expect a complete signal with status and headers
+			voidFirstStep.consumeSubscriptionWith(aVoid -> {
+				assertThat(exchange.getResponse().getStatusCode()).isEqualTo(expectedStatus);
+				assertThat(exchange.getResponse().getHeaders().toSingleValueMap()).containsEntry("X-Tokens-Remaining",
+						tokensRemaining);
+			}).expectComplete().verify();
+		}
 	}
 
 	@EnableAutoConfiguration