Redis Lua Rate Limiter

애플리케이션 검사만으로 막을 수 없던 경쟁

CTF 답안 제출은 팀마다 허용 횟수가 정해져 있었습니다. 제출 횟수는 DB에 저장했지만, 제출 로직 안에서 조회·증가·한도 판정이 나뉘어 있었습니다. 두 요청이 거의 같은 시각에 들어오면 같은 횟수를 기준으로 둘 다 한도 안이라고 판단해 정해진 횟수보다 많은 제출이 통과할 수 있었습니다.

검사와 증가를 한 번에 실행하기

읽기와 쓰기 사이에 틈이 있는 한 애플리케이션에서 아무리 검사해도 같은 문제가 남습니다. 그래서 판단 자체를 Redis 안으로 옮겼습니다. Redis는 script를 실행하는 동안 다른 명령을 끼워 넣지 않으므로, 카운터를 올리고 한도와 비교하는 일이 하나의 실행 안에서 끝납니다.

-- KEYS[1] 팀 단위 카운터, ARGV[1] 윈도 길이(초), ARGV[2] 윈도당 허용 횟수
local count = redis.call('INCR', KEYS[1])

if count == 1 then
  redis.call('EXPIRE', KEYS[1], ARGV[1])
end

if count > tonumber(ARGV[2]) then
  return 0
end

return 1

카운터가 처음 만들어질 때만 TTL을 겁니다. 윈도가 지나면 key가 스스로 사라지므로 따로 정리하는 작업이 필요하지 않습니다.

제한 규칙을 한곳에 모으기

제한이 필요한 자리마다 Redis 호출을 직접 넣으면 규칙이 코드 곳곳에 흩어집니다. 어디에 걸리는지는 annotation으로 선언하고, 실제 검사는 AOP가 가로채도록 나눴습니다.

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimited {
    String action();
    int windowSeconds();
    int limit();
}

@Aspect
@Component
public class RateLimitAspect {

    private static final RedisScript<Long> FIXED_WINDOW = RedisScript.of(
            new ClassPathResource("lua/fixed-window.lua"), Long.class);

    private final StringRedisTemplate redis;

    public RateLimitAspect(StringRedisTemplate redis) {
        this.redis = redis;
    }

    @Before("@annotation(rateLimited)")
    public void check(RateLimited rateLimited) {
        String key = "rate:%s:%s"
                .formatted(rateLimited.action(), currentTeamId());

        Long allowed = redis.execute(
                FIXED_WINDOW,
                List.of(key),
                String.valueOf(rateLimited.windowSeconds()),
                String.valueOf(rateLimited.limit()));

        if (!Long.valueOf(1L).equals(allowed)) {
            throw new RateLimitExceededException(key);
        }
    }
}

제출 처리에는 @RateLimited(action = "submit", windowSeconds = 60, limit = 3) 한 줄만 붙습니다. 다른 기능에 같은 제한이 필요해지면 값만 바꿔 붙이면 됩니다.

동시 요청을 테스트 코드로 확인하기

Spring Boot 테스트에서는 여러 작업이 같은 시점에 Lua script를 실행하도록 CountDownLatch로 시작 시점을 맞췄습니다.

@SpringBootTest
class RedisLuaRateLimiterConcurrencyTest {

    @Autowired
    private StringRedisTemplate redis;

    private final RedisScript<Long> script = RedisScript.of(
            new ClassPathResource("lua/fixed-window.lua"), Long.class);

    @Test
    void concurrentRequestsDoNotExceedLimit() throws Exception {
        int requests = 100;
        int concurrency = 20;
        int limit = 10;
        String key = "rate:test:" + UUID.randomUUID();

        ExecutorService executor = Executors.newFixedThreadPool(concurrency);
        CountDownLatch ready = new CountDownLatch(concurrency);
        CountDownLatch start = new CountDownLatch(1);

        try {
            List<Future<Long>> results = IntStream.range(0, requests)
                    .mapToObj(ignored -> executor.submit(() -> {
                        ready.countDown();
                        start.await();
                        return redis.execute(
                                script,
                                List.of(key),
                                "60",
                                String.valueOf(limit));
                    }))
                    .toList();

            assertThat(ready.await(3, TimeUnit.SECONDS)).isTrue();
            start.countDown();

            long allowed = 0;
            for (Future<Long> result : results) {
                if (Long.valueOf(1L).equals(result.get(3, TimeUnit.SECONDS))) {
                    allowed++;
                }
            }

            assertThat(allowed).isEqualTo(limit);
        } finally {
            start.countDown();
            executor.shutdownNow();
            executor.awaitTermination(3, TimeUnit.SECONDS);
            redis.delete(key);
        }
    }
}

테스트 결과

예시 시나리오는 요청 100건을 최대 20개 작업 스레드로 실행하고 허용 한도를 10건으로 둡니다. 허용 결과가 10건과 다르면 테스트가 실패합니다.

테스트 결과: 허용 10건, 차단 90건.