Skip to content

Commit 6f6a602

Browse files
oopswoo3oopswoo3
authored
fix(auth): make Redis rate-limit counters atomic (#9)
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.
1 parent 23a3f98 commit 6f6a602

2 files changed

Lines changed: 100 additions & 9 deletions

File tree

‎src/main/java/ai/nubase/auth/service/RateLimiterService.java‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,13 @@
55
import lombok.extern.slf4j.Slf4j;
66
import org.springframework.beans.factory.annotation.Autowired;
77
import org.springframework.data.redis.core.StringRedisTemplate;
8+
import org.springframework.data.redis.core.script.RedisScript;
89
import org.springframework.stereotype.Service;
910

1011
import java.time.Duration;
1112
import java.util.ArrayDeque;
1213
import java.util.Deque;
14+
import java.util.List;
1315
import java.util.Map;
1416
import java.util.concurrent.ConcurrentHashMap;
1517

@@ -26,6 +28,20 @@
2628
@Slf4j
2729
public class RateLimiterService {
2830

31+
/**
32+
* INCR 与首次 EXPIRE 必须在同一 Lua 脚本里完成：分两次调用时，若在 EXPIRE 前进程崩溃，
33+
* 会留下无 TTL 的计数 key，窗口永不重置（永久 429）或失败计数永不衰减。
34+
* 与 {@link ai.nubase.functions.service.EdgeFunctionRateLimiter} 使用相同模式。
35+
*/
36+
private static final RedisScript<Long> INCREMENT_WITH_EXPIRE_SCRIPT = RedisScript.of(
37+
"""
38+
local count = redis.call('INCR', KEYS[1])
39+
if count == 1 then
40+
redis.call('PEXPIRE', KEYS[1], ARGV[1])
41+
end
42+
return count
43+
""", Long.class);
44+
2945
private final EffectiveAuthConfig effectiveAuthConfig;
3046

3147
/** Optional — present only when Redis is configured. Null → in-process fallback. */
@@ -129,11 +145,9 @@ public void recordSuccess(String identifier) {
129145
// ---------------------------------------------------------------- Redis backend
130146

131147
private void checkRateRedis(AuthConfig.RateLimitSettings cfg, String action, String identifier) {
132-
String key = "rl:" + key(action, identifier);
133-
Long n = redisTemplate.opsForValue().increment(key);
134-
if (n != null && n == 1L) {
135-
redisTemplate.expire(key, Duration.ofSeconds(cfg.getWindowSeconds()));
136-
}
148+
String redisKey = "rl:" + key(action, identifier);
149+
long ttlMillis = cfg.getWindowSeconds() * 1000L;
150+
Long n = redisIncrementWithExpire(redisKey, ttlMillis);
137151
if (n != null && n > cfg.getMaxRequests()) {
138152
throw new RateLimitExceededException(
139153
"Rate limit exceeded for " + action + ". Please try again later.");
@@ -142,17 +156,22 @@ private void checkRateRedis(AuthConfig.RateLimitSettings cfg, String action, Str
142156

143157
private void recordFailureRedis(AuthConfig.RateLimitSettings cfg, String identifier) {
144158
String countKey = "rl:fail:" + tenant() + ":" + identifier;
145-
Long n = redisTemplate.opsForValue().increment(countKey);
146-
if (n != null && n == 1L) {
147-
redisTemplate.expire(countKey, Duration.ofSeconds(cfg.getLockoutSeconds()));
148-
}
159+
long ttlMillis = cfg.getLockoutSeconds() * 1000L;
160+
Long n = redisIncrementWithExpire(countKey, ttlMillis);
149161
if (n != null && n >= cfg.getMaxFailedLogins()) {
150162
redisTemplate.opsForValue().set("rl:lock:" + tenant() + ":" + identifier, "1",
151163
Duration.ofSeconds(cfg.getLockoutSeconds()));
152164
log.warn("Identity '{}' locked out after {} failed attempts (redis)", identifier, n);
153165
}
154166
}
155167

168+
private Long redisIncrementWithExpire(String redisKey, long ttlMillis) {
169+
return redisTemplate.execute(
170+
INCREMENT_WITH_EXPIRE_SCRIPT,
171+
List.of(redisKey),
172+
String.valueOf(ttlMillis));
173+
}
174+
156175
private String key(String action, String identifier) {
157176
return tenant() + ":" + action + ":" + identifier;
158177
}

‎src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,27 @@
33
import ai.nubase.common.config.AuthConfig;
44
import org.junit.jupiter.api.DisplayName;
55
import org.junit.jupiter.api.Test;
6+
import org.mockito.ArgumentCaptor;
7+
import org.mockito.ArgumentMatchers;
8+
import org.springframework.data.redis.core.StringRedisTemplate;
9+
import org.springframework.data.redis.core.ValueOperations;
10+
import org.springframework.data.redis.core.script.RedisScript;
11+
import org.springframework.test.util.ReflectionTestUtils;
612

13+
import java.time.Duration;
14+
import java.util.List;
15+
16+
import static org.assertj.core.api.Assertions.assertThat;
717
import static org.assertj.core.api.Assertions.assertThatCode;
818
import static org.assertj.core.api.Assertions.assertThatThrownBy;
19+
import static org.mockito.ArgumentMatchers.any;
20+
import static org.mockito.ArgumentMatchers.anyList;
21+
import static org.mockito.ArgumentMatchers.eq;
22+
import static org.mockito.Mockito.mock;
23+
import static org.mockito.Mockito.never;
24+
import static org.mockito.Mockito.times;
25+
import static org.mockito.Mockito.verify;
26+
import static org.mockito.Mockito.when;
927

1028
/**
1129
* 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
2341
return new RateLimiterService(new EffectiveAuthConfig(cfg));
2442
}
2543

44+
private RateLimiterService redisLimiter(int maxReq, int window, int maxFail, int lockout) {
45+
RateLimiterService rl = limiter(maxReq, window, maxFail, lockout);
46+
StringRedisTemplate redisTemplate = mock(StringRedisTemplate.class);
47+
ReflectionTestUtils.setField(rl, "redisTemplate", redisTemplate);
48+
return rl;
49+
}
50+
2651
@Test
2752
@DisplayName("checkRate allows up to the cap, then throws")
2853
void rateCap() {
@@ -81,4 +106,51 @@ void disabled() {
81106
}
82107
assertThatCode(() -> rl.assertNotLockedOut("a@x.com")).doesNotThrowAnyException();
83108
}
109+
110+
@Test
111+
@DisplayName("Redis checkRate uses atomic INCR+PEXPIRE script and enforces cap")
112+
void redisRateCapUsesAtomicScript() {
113+
RateLimiterService rl = redisLimiter(2, 300, 5, 900);
114+
StringRedisTemplate redisTemplate =
115+
(StringRedisTemplate) ReflectionTestUtils.getField(rl, "redisTemplate");
116+
when(redisTemplate.execute(ArgumentMatchers.<RedisScript<Long>>any(), anyList(), any()))
117+
.thenReturn(1L, 2L, 3L);
118+
119+
assertThatCode(() -> rl.checkRate("otp", "a@x.com")).doesNotThrowAnyException();
120+
assertThatCode(() -> rl.checkRate("otp", "a@x.com")).doesNotThrowAnyException();
121+
assertThatThrownBy(() -> rl.checkRate("otp", "a@x.com"))
122+
.isInstanceOf(RateLimiterService.RateLimitExceededException.class);
123+
124+
@SuppressWarnings("unchecked")
125+
ArgumentCaptor<List<String>> keysCaptor = ArgumentCaptor.forClass(List.class);
126+
verify(redisTemplate, times(3))
127+
.execute(ArgumentMatchers.<RedisScript<Long>>any(), keysCaptor.capture(), eq("300000"));
128+
assertThat(keysCaptor.getValue()).containsExactly("rl:_:otp:a@x.com");
129+
verify(redisTemplate, never()).opsForValue();
130+
}
131+
132+
@Test
133+
@DisplayName("Redis recordFailure uses atomic INCR+PEXPIRE script before lockout")
134+
void redisFailureCountUsesAtomicScript() {
135+
RateLimiterService rl = redisLimiter(100, 300, 3, 900);
136+
StringRedisTemplate redisTemplate =
137+
(StringRedisTemplate) ReflectionTestUtils.getField(rl, "redisTemplate");
138+
ValueOperations<String, String> valueOps = mock(ValueOperations.class);
139+
when(redisTemplate.opsForValue()).thenReturn(valueOps);
140+
when(redisTemplate.execute(ArgumentMatchers.<RedisScript<Long>>any(), anyList(), any()))
141+
.thenReturn(1L, 2L, 3L);
142+
when(redisTemplate.hasKey("rl:lock:_:victim@x.com")).thenReturn(false, true);
143+
144+
String id = "victim@x.com";
145+
rl.recordFailure(id);
146+
rl.recordFailure(id);
147+
assertThatCode(() -> rl.assertNotLockedOut(id)).doesNotThrowAnyException();
148+
rl.recordFailure(id);
149+
assertThatThrownBy(() -> rl.assertNotLockedOut(id))
150+
.isInstanceOf(RateLimiterService.RateLimitExceededException.class);
151+
152+
verify(redisTemplate, times(3))
153+
.execute(ArgumentMatchers.<RedisScript<Long>>any(), anyList(), eq("900000"));
154+
verify(valueOps).set(eq("rl:lock:_:victim@x.com"), eq("1"), eq(Duration.ofSeconds(900)));
155+
}
84156
}

0 commit comments

Comments
 (0)