Skip to content

fix(auth): make Redis rate-limit counters atomic #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Chat2DB-Pro merged 1 commit into from
Jun 18, 2026
Merged

fix(auth): make Redis rate-limit counters atomic #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 28 additions & 9 deletions src/main/java/ai/nubase/auth/service/RateLimiterService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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<Long> 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. */
Expand Down Expand Up @@ -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.");
Expand All @@ -142,17 +156,22 @@ 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()));
log.warn("Identity '{}' locked out after {} failed attempts (redis)", identifier, n);
}
}

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;
}
Expand Down
72 changes: 72 additions & 0 deletions src/test/java/ai/nubase/auth/service/RateLimiterServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand All @@ -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() {
Expand Down Expand Up @@ -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.<RedisScript<Long>>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<List<String>> keysCaptor = ArgumentCaptor.forClass(List.class);
verify(redisTemplate, times(3))
.execute(ArgumentMatchers.<RedisScript<Long>>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<String, String> valueOps = mock(ValueOperations.class);
when(redisTemplate.opsForValue()).thenReturn(valueOps);
when(redisTemplate.execute(ArgumentMatchers.<RedisScript<Long>>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.<RedisScript<Long>>any(), anyList(), eq("900000"));
verify(valueOps).set(eq("rl:lock:_:victim@x.com"), eq("1"), eq(Duration.ofSeconds(900)));
}
}
Loading