From 2aa0859066a023f3560f211eee197db6b68bfe21 Mon Sep 17 00:00:00 2001 From: oopswoo3 <102633566+oopswoo3@users.noreply.github.com> Date: Thu, 18 Jun 2026 14:56:10 +0800 Subject: [PATCH] fix(auth): make Redis rate-limit counters atomic Run INCR and first-hit PEXPIRE in one Lua script for auth rate limits and failed-login counters, matching EdgeFunctionRateLimiter and preventing TTL-less keys that never reset. --- .../auth/service/RateLimiterService.java | 37 +++++++--- .../auth/service/RateLimiterServiceTest.java | 72 +++++++++++++++++++ 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/src/main/java/ai/nubase/auth/service/RateLimiterService.java b/src/main/java/ai/nubase/auth/service/RateLimiterService.java index 25adcb8..b0c39b8 100644 --- a/src/main/java/ai/nubase/auth/service/RateLimiterService.java +++ b/src/main/java/ai/nubase/auth/service/RateLimiterService.java @@ -5,11 +5,13 @@ import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.data.redis.core.StringRedisTemplate; +import org.springframework.data.redis.core.script.RedisScript; import org.springframework.stereotype.Service; import java.time.Duration; import java.util.ArrayDeque; import java.util.Deque; +import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -26,6 +28,20 @@ @Slf4j public class RateLimiterService { + /** + * INCR 与首次 EXPIRE 必须在同一 Lua 脚本里完成：分两次调用时，若在 EXPIRE 前进程崩溃， + * 会留下无 TTL 的计数 key，窗口永不重置（永久 429）或失败计数永不衰减。 + * 与 {@link ai.nubase.functions.service.EdgeFunctionRateLimiter} 使用相同模式。 + */ + private static final RedisScript INCREMENT_WITH_EXPIRE_SCRIPT = RedisScript.of( + """ + local count = redis.call('INCR', KEYS[1]) + if count == 1 then + redis.call('PEXPIRE', KEYS[1], ARGV[1]) + end + return count + """, Long.class); + private final EffectiveAuthConfig effectiveAuthConfig; /** Optional — present only when Redis is configured. Null → in-process fallback. */ @@ -129,11 +145,9 @@ public void recordSuccess(String identifier) { // ---------------------------------------------------------------- Redis backend private void checkRateRedis(AuthConfig.RateLimitSettings cfg, String action, String identifier) { - String key = "rl:" + key(action, identifier); - Long n = redisTemplate.opsForValue().increment(key); - if (n != null && n == 1L) { - redisTemplate.expire(key, Duration.ofSeconds(cfg.getWindowSeconds())); - } + String redisKey = "rl:" + key(action, identifier); + long ttlMillis = cfg.getWindowSeconds() * 1000L; + Long n = redisIncrementWithExpire(redisKey, ttlMillis); if (n != null && n > cfg.getMaxRequests()) { throw new RateLimitExceededException( "Rate limit exceeded for " + action + ". Please try again later."); @@ -142,10 +156,8 @@ private void checkRateRedis(AuthConfig.RateLimitSettings cfg, String action, Str private void recordFailureRedis(AuthConfig.RateLimitSettings cfg, String identifier) { String countKey = "rl:fail:" + tenant() + ":" + identifier; - Long n = redisTemplate.opsForValue().increment(countKey); - if (n != null && n == 1L) { - redisTemplate.expire(countKey, Duration.ofSeconds(cfg.getLockoutSeconds())); - } + long ttlMillis = cfg.getLockoutSeconds() * 1000L; + Long n = redisIncrementWithExpire(countKey, ttlMillis); if (n != null && n >= cfg.getMaxFailedLogins()) { redisTemplate.opsForValue().set("rl:lock:" + tenant() + ":" + identifier, "1", Duration.ofSeconds(cfg.getLockoutSeconds())); @@ -153,6 +165,13 @@ private void recordFailureRedis(AuthConfig.RateLimitSettings cfg, String identif } } + private Long redisIncrementWithExpire(String redisKey, long ttlMillis) { + return redisTemplate.execute( + INCREMENT_WITH_EXPIRE_SCRIPT, + List.of(redisKey), + String.valueOf(ttlMillis)); + } + private String key(String action, String identifier) { return tenant() + ":" + action + ":" + identifier; } diff --git a/src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java b/src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java index 2d38f99..6bbc68c 100644 --- a/src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java +++ b/src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java @@ -3,9 +3,27 @@ import ai.nubase.common.config.AuthConfig; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.ArgumentMatchers; +import org.springframework.data.redis.core.StringRedisTemplate; +import org.springframework.data.redis.core.ValueOperations; +import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.test.util.ReflectionTestUtils; +import java.time.Duration; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; /** * Pure unit tests for {@link RateLimiterService} (sliding-window cap + failed-login lockout). @@ -23,6 +41,13 @@ private RateLimiterService limiter(int maxReq, int window, int maxFail, int lock return new RateLimiterService(new EffectiveAuthConfig(cfg)); } + private RateLimiterService redisLimiter(int maxReq, int window, int maxFail, int lockout) { + RateLimiterService rl = limiter(maxReq, window, maxFail, lockout); + StringRedisTemplate redisTemplate = mock(StringRedisTemplate.class); + ReflectionTestUtils.setField(rl, "redisTemplate", redisTemplate); + return rl; + } + @Test @DisplayName("checkRate allows up to the cap, then throws") void rateCap() { @@ -81,4 +106,51 @@ void disabled() { } assertThatCode(() -> rl.assertNotLockedOut("a@x.com")).doesNotThrowAnyException(); } + + @Test + @DisplayName("Redis checkRate uses atomic INCR+PEXPIRE script and enforces cap") + void redisRateCapUsesAtomicScript() { + RateLimiterService rl = redisLimiter(2, 300, 5, 900); + StringRedisTemplate redisTemplate = + (StringRedisTemplate) ReflectionTestUtils.getField(rl, "redisTemplate"); + when(redisTemplate.execute(ArgumentMatchers.>any(), anyList(), any())) + .thenReturn(1L, 2L, 3L); + + assertThatCode(() -> rl.checkRate("otp", "a@x.com")).doesNotThrowAnyException(); + assertThatCode(() -> rl.checkRate("otp", "a@x.com")).doesNotThrowAnyException(); + assertThatThrownBy(() -> rl.checkRate("otp", "a@x.com")) + .isInstanceOf(RateLimiterService.RateLimitExceededException.class); + + @SuppressWarnings("unchecked") + ArgumentCaptor> keysCaptor = ArgumentCaptor.forClass(List.class); + verify(redisTemplate, times(3)) + .execute(ArgumentMatchers.>any(), keysCaptor.capture(), eq("300000")); + assertThat(keysCaptor.getValue()).containsExactly("rl:_:otp:a@x.com"); + verify(redisTemplate, never()).opsForValue(); + } + + @Test + @DisplayName("Redis recordFailure uses atomic INCR+PEXPIRE script before lockout") + void redisFailureCountUsesAtomicScript() { + RateLimiterService rl = redisLimiter(100, 300, 3, 900); + StringRedisTemplate redisTemplate = + (StringRedisTemplate) ReflectionTestUtils.getField(rl, "redisTemplate"); + ValueOperations valueOps = mock(ValueOperations.class); + when(redisTemplate.opsForValue()).thenReturn(valueOps); + when(redisTemplate.execute(ArgumentMatchers.>any(), anyList(), any())) + .thenReturn(1L, 2L, 3L); + when(redisTemplate.hasKey("rl:lock:_:victim@x.com")).thenReturn(false, true); + + String id = "victim@x.com"; + rl.recordFailure(id); + rl.recordFailure(id); + assertThatCode(() -> rl.assertNotLockedOut(id)).doesNotThrowAnyException(); + rl.recordFailure(id); + assertThatThrownBy(() -> rl.assertNotLockedOut(id)) + .isInstanceOf(RateLimiterService.RateLimitExceededException.class); + + verify(redisTemplate, times(3)) + .execute(ArgumentMatchers.>any(), anyList(), eq("900000")); + verify(valueOps).set(eq("rl:lock:_:victim@x.com"), eq("1"), eq(Duration.ofSeconds(900))); + } }