From 0d1e2b3f92760888a1d56a0ed366e44168c899f0 Mon Sep 17 00:00:00 2001
From: TheFreaky <maksoh88@gmail.com>
Date: Wed, 13 Nov 2024 18:03:05 +0100
Subject: [PATCH] Allow larger rate limits for extended time intervals

---
 .../filter/ratelimit/RedisRateLimiter.java    | 21 +++++++++++++------
 .../RedisRateLimiterConfigTests.java          |  8 +++++++
 .../RedisRateLimiterLuaScriptTests.java       | 17 +++++++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java
index 6ae0c46766..1acb7d5a75 100644
--- a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java
+++ b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiter.java
@@ -83,6 +83,13 @@ public class RedisRateLimiter extends AbstractRateLimiter<RedisRateLimiter.Confi
 	 */
 	public static final String REQUESTED_TOKENS_HEADER = "X-RateLimit-Requested-Tokens";
 
+	/**
+	 * Constant representing the maximum number that can be safely used in Redis scripts.
+	 * This is the largest number that can be represented accurately in Java's `long` type
+	 * and is equal to 2^53 - 1, which is the largest number that Lua can reliably handle.
+	 */
+	private static final Long REDIS_LUA_MAX_SAFE_INTEGER = 9007199254740991L;
+
 	private Log log = LogFactory.getLog(getClass());
 
 	private ReactiveStringRedisTemplate redisTemplate;
@@ -129,7 +136,7 @@ public RedisRateLimiter(ReactiveStringRedisTemplate redisTemplate, RedisScript<L
 	 * @param defaultBurstCapacity how many tokens the bucket can hold in token-bucket
 	 * algorithm.
 	 */
-	public RedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
+	public RedisRateLimiter(int defaultReplenishRate, long defaultBurstCapacity) {
 		super(Config.class, CONFIGURATION_PROPERTY_NAME, (ConfigurationService) null);
 		this.defaultConfig = new Config().setReplenishRate(defaultReplenishRate).setBurstCapacity(defaultBurstCapacity);
 	}
@@ -141,7 +148,7 @@ public RedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
 	 * algorithm.
 	 * @param defaultRequestedTokens how many tokens are requested per request.
 	 */
-	public RedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity, int defaultRequestedTokens) {
+	public RedisRateLimiter(int defaultReplenishRate, long defaultBurstCapacity, int defaultRequestedTokens) {
 		this(defaultReplenishRate, defaultBurstCapacity);
 		this.defaultConfig.setRequestedTokens(defaultRequestedTokens);
 	}
@@ -240,7 +247,7 @@ public Mono<Response> isAllowed(String routeId, String id) {
 		int replenishRate = routeConfig.getReplenishRate();
 
 		// How much bursting do you want to allow?
-		int burstCapacity = routeConfig.getBurstCapacity();
+		long burstCapacity = routeConfig.getBurstCapacity();
 
 		// How many tokens are requested per request?
 		int requestedTokens = routeConfig.getRequestedTokens();
@@ -313,7 +320,7 @@ public static class Config {
 		private int replenishRate;
 
 		@Min(0)
-		private int burstCapacity = 1;
+		private long burstCapacity = 1;
 
 		@Min(1)
 		private int requestedTokens = 1;
@@ -327,13 +334,15 @@ public Config setReplenishRate(int replenishRate) {
 			return this;
 		}
 
-		public int getBurstCapacity() {
+		public long getBurstCapacity() {
 			return burstCapacity;
 		}
 
-		public Config setBurstCapacity(int burstCapacity) {
+		public Config setBurstCapacity(long burstCapacity) {
 			Assert.isTrue(burstCapacity >= this.replenishRate, "BurstCapacity(" + burstCapacity
 					+ ") must be greater than or equal than replenishRate(" + this.replenishRate + ")");
+			Assert.isTrue(burstCapacity <= REDIS_LUA_MAX_SAFE_INTEGER, "BurstCapacity(" + burstCapacity
+					+ ") must not exceed the maximum allowed value of " + REDIS_LUA_MAX_SAFE_INTEGER);
 			this.burstCapacity = burstCapacity;
 			return this;
 		}
diff --git a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterConfigTests.java b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterConfigTests.java
index 708c039781..049426897d 100644
--- a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterConfigTests.java
+++ b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterConfigTests.java
@@ -58,6 +58,14 @@ public void shouldThrowAnErrorWhenReplenishRateIsHigherThanBurstCapacity() {
 		Assertions.assertThatThrownBy(() -> new RedisRateLimiter(10, 5)).isInstanceOf(IllegalArgumentException.class);
 	}
 
+	@Test
+	public void shouldThrowAnErrorWhenBurstCapacityExceedsMaxAllowed() {
+		long burstCapacity = Long.MAX_VALUE;
+
+		Assertions.assertThatThrownBy(() -> new RedisRateLimiter(10, burstCapacity))
+			.isInstanceOf(IllegalArgumentException.class);
+	}
+
 	@Test
 	public void redisRateConfiguredFromEnvironment() {
 		assertFilter("redis_rate_limiter_config_test", 10, 20, 1, false);
diff --git a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java
index dc36f9291b..a303ce911e 100644
--- a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java
+++ b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java
@@ -47,6 +47,7 @@
 public class RedisRateLimiterLuaScriptTests {
 
 	static final String KEY_PREFIX = "redis-rate-limiter-lua-script-tests";
+	static final Long REDIS_LUA_MAX_SAFE_INTEGER = 9007199254740991L;
 
 	@Container
 	public static GenericContainer redis = new GenericContainer<>("redis:5.0.14-alpine").withExposedPorts(6379);
@@ -147,6 +148,22 @@ public void testTokensNotEnough() {
 		assertThat(result.get(1)).isEqualTo(10);
 	}
 
+	@Test
+	public void testCapacityExceedsMaxInt() {
+		long rate = 1;
+		long capacity = REDIS_LUA_MAX_SAFE_INTEGER;
+		long now = System.currentTimeMillis();
+		long requested = 1;
+
+		List<String> keys = getKeys("capacity_exceeds_max_int");
+		List<String> args = getArgs(rate, capacity, now, requested);
+
+		List<Long> result = redisTemplate.execute(redisScript, keys, args).blockFirst();
+
+		assertThat(result.get(0)).isEqualTo(1);
+		assertThat(result.get(1)).isEqualTo(REDIS_LUA_MAX_SAFE_INTEGER - 1);
+	}
+
 	@EnableAutoConfiguration
 	@SpringBootConfiguration
 	public static class TestConfig {