[Spring] Bucket4j를 이용하여 트래픽 제어하기

이재민·2024년 5월 16일
0

트러블슈팅&개선

목록 보기
2/5

개요

현재 사내 서비스는 모바일 웹 서비스를 제공하고 있습니다. 모바일 웹 서비스로 무작위 호출을 하는 등의 공격이 잦아지고 있었습니다.
그러다 보니 모니터링 알림에도 알림이 누적되고 구성원에게 좋지 못한 경험을 주고 있었습니다.
게다가, 이러한 공격이 계속되다가 혹여나 공격이 성공하게 된다면 내부 자원이 보호되지 못할 수도 있게 됩니다.
이슈가 발생할때마다 모바일 웹에서는 해당 ip들의 요청을 차단하고 있었습니다. 쉬는날에도 팀원들이 이러한 대응을 하고 있었기에 이를 해결하고자 하였습니다.
대비책을 세워두고자 방안을 모색하였고 그 중 Bucket4j를 이용하여 트래픽을 제어하고자 하였습니다.

Bucket4j란

Bucket4j는 token-bucket 알고리즘을 기반의 자바 처리율 제한기 라이브러리이다.

Bucket4j에서 Api Throttling을 구현 방법

  1. 먼저, 각 사용자 또는 API 키에 대한 버킷을 생성한다. 버킷은 특정 시간 동안 허용되는 요청의 최대 수를 나타낸다
  2. 사용자 또는 API 키가 API를 호출할 때마다, 해당 버킷에서 토큰을 하나 사용한다.
  3. 버킷에 토큰이 없으면 요청은 거부된다. 이는 사용자가 허용된 요청 수를 초과했음을 의미.
  4. 버킷에는 일정 시간마다 토큰이 자동으로 추가된다. 이를 통해 사용자는 일정 시간이 지나면 다시 API를 호출할 수 있다.

Bucket4j를 사용하면 이러한 과정을 쉽게 구현할 수 있다.
Bucket4j는 버킷 생성, 토큰 사용, 토큰 재충전 등의 기능을 제공하므로, 개발자는 API Throttling 로직에 집중할 수 있다.

API Throttling

간단하게 API Throttling은 API 요청에 속도와 횟수를 제한하는 것을 말합니다. 이를 속도 제한(Rate Limit) 기능으로 부르기도 합니다.
API Throttling은 이번 포스팅의 목적과 같이 악의적인 사용자가 Dos와 같은 공격을 시도해서 요청이 폭증하거나 응답 레이턴시가 증가하는 것을 막고자
즉, 시스템 보안성 강화를 위해서 사용하곤 합니다.
또한 과금의 조율을 위해서도 사용합니다. Baeldung 가이드를 보시면 Pricing Plan을 도입하여 업체별로 제한을 할수도 있습니다.

API 스로틀링을 지탱하는 알고리즘은 다음과 같습니다.
1. Leaky Bucket
2. Token Bucket
3. Fixed Window, Sliding Window
위 내용들에 대해서는 나중에 자세히 다뤄보도록 하겠습니다.

Bucket4j 구현

현재 서비스에서는 분산 서버를 운용중에 있고 Redis Cluster를 구성해서 사용중에 있다.
Bucket4j는 로컬 메모리 외에도 JDBC, Redis 등 다양한 분산 환경에 대한 지원을 하고 있다.
필자는 Bucket4j를 Redis에 적용해서 구현하였다.
전체 코드는 Github을 참고해주세요.

Bucket4j-redis

Bucket4j는 Redis Jedis, Lettuce, Redisson 클라이언트를 모두 지원하며 Lettuce를 사용하여 적용하였다.

1. 의존성 추가

...
implementation 'com.bucket4j:bucket4j-redis:8.7.0'

2. RateLimiterConfig

@Slf4j
@RequiredArgsConstructor
@Configuration
public class RateLimiterConfig {

    // 버킷에 담길 수 있는 토큰의 최대 수
    private static final int CAPACITY = 20;

    // REFILL_DURATION 으로 지정된 시간 동안 버킷에 추가될 토큰의 수
    private static final int REFILL_TOKEN_AMOUNT = 3;

    // 토큰이 재충전되는 빈도
    private static final Duration REFILL_DURATION = Duration.ofSeconds(5);

    private final RedisClient lettuceRedisClient;

@Bean
    public LettuceBasedProxyManager lettuceBasedProxyManager() {
        StatefulRedisConnection<String, byte[]> connect = lettuceRedisClient.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));

        return LettuceBasedProxyManager.builderFor(connect)
            // 버킷 만료 정책(최대 용량으로 채워지는데 필요한 60초 동안 버킷으로 유지)
            .withExpirationStrategy(ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofSeconds(60)))
            .build();
    }

    @Bean
    public BucketConfiguration bucketConfiguration() {
        return BucketConfiguration.builder()
            // REFILL_DURATION 동안 REFILL_TOKEN_AMOUNT 만큼의 토근이 버킷에 추가(5초동안 3개의 토큰이 추가됨)
            .addLimit(Bandwidth.classic(CAPACITY, Refill.intervally(REFILL_TOKEN_AMOUNT, REFILL_DURATION)))
            .build();
    }
}
  • withExpirationStrategy : 각 버킷의 만료 정책을 설정할 수 있다. 각 버킷이 최대 용량으로 채워지는 동안 정해진 시간 동안 버킷을 Redis에 유지한다.
    만약, 그동안 요청이 없다면 해당 버킷은 Redis에서 제거한다.
  • StatefulRedisConnection : StatefulRedisConnection 를 사용해서 여러 스레드에서 동시에 사용할 수 있으며 Thread-safe하다.
    또한, 여러 스레드가 하나의 StatefulRedisConnection 를 공유할 수 있어. 매번 connection을 생성 및 종료하는 오버헤드를 줄일 수 있다.
@Slf4j
@RequiredArgsConstructor
@Service
public class RateLimiterService {

    private final RateLimiterConfig rateLimiterConfig;

    public boolean tryConsume(String remoteAddrKey) {
        Bucket bucket = getOrCreateBucket(remoteAddrKey);

        ConsumptionProbe probe = consumeToken(bucket);

        logConsumption(remoteAddrKey, probe);

        handleNotConsumed(probe);

        return probe.isConsumed();
    }

    private Bucket getOrCreateBucket(String apiKey) {
        return rateLimiterConfig.lettuceBasedProxyManager().builder()
            .build(apiKey, () -> rateLimiterConfig.bucketConfiguration());
    }

    private ConsumptionProbe consumeToken(Bucket bucket) {
        return bucket.tryConsumeAndReturnRemaining(1);
    }

    private void logConsumption(String remoteAddrKey, ConsumptionProbe probe) {
        log.info("API Key: {}, RemoteAddress: {}, tryConsume: {}, remainToken: {}, tryTime: {}",
            remoteAddrKey, remoteAddrKey, probe.isConsumed(), probe.getRemainingTokens(), LocalDateTime.now());
    }

    private void handleNotConsumed(ConsumptionProbe probe) {
        if (!probe.isConsumed()) {
            throw new RateLimiterException(RateLimiterException.TOO_MANY_REQUEST);
        }
    }

    public long getRemainToken(String apiKey) {
        Bucket bucket = getOrCreateBucket(apiKey);
        return bucket.getAvailableTokens();
    }
}
  • tryConsume : 해당 메소드는 크게 5개의 단계로 나뉘어져있다.
  1. 버킷가져오기
  2. 토큰 소모
  3. 로깅
  4. 잔여 토큰 체크
  5. 토큰 소모 여부 반환
  • getOrCreateBucket() : 해당 메소드를 통해 Bucket을 Redis에서 관리하도록 한다.
    이를 통해 분산 환경에서 동기화를 가져갈 수 있다.

3. 테스트 및 검증

@SpringBootTest
class RateLimiterServiceTest {

    @Autowired
    private RateLimiterService rateLimiterService;

    @Autowired
    private RedisTemplate redisTemplate;

    @BeforeEach
    void setUp() {
        redisTemplate.getConnectionFactory().getConnection().flushAll();
    }

    @DisplayName("처리율 제한에 걸리지 않으면 true를 반환한다.")
    @Test
    void return_true_if_not_rate_limit() {
        // given
        String key = "127.0.0.1";

        // when
        boolean result = rateLimiterService.tryConsume(key);

        // then
        Assertions.assertThat(result).isTrue();
    }

    @DisplayName("남은 토큰 개수를 반환한다.")
    @Test
    void return_remain_token() {
        // given
        String key = "127.0.0.1";

        // when
        long remainToken = rateLimiterService.getRemainToken(key);

        // then
        Assertions.assertThat(remainToken).isEqualTo(20);
    }

    @DisplayName("처리율 제한에 걸리면 예외를 발생시킨다.")
    @Test
    void throw_exception_if_rate_limit() {
        // given
        String key = "127.0.0.1";

        // when
        for (int i = 0; i < 20; i++) {
            rateLimiterService.tryConsume(key);
        }

        // then
        Assertions.assertThatThrownBy(() -> rateLimiterService.tryConsume(key))
            .isInstanceOf(RateLimiterException.class)
            .hasMessage(RateLimiterException.TOO_MANY_REQUEST);
    }
}
  • 총 3개의 테스트를 진행하였다.

적용 방안

이번 학습을 통해 고민하고 자료를 찾아본 결과 처리율 제한을 구현할 수 있는 방법은 많은 것 같다.

  1. 어노테이션을 이용해 특정 엔드포인트를 중점으로 처리율 제한 -> 깃헙 코드를 확인하면 된다.

  2. Interceptor를 통해 전반적인 요청에 대한 처리율 제한 -> Pricing Plan 을 사용할 때 조금 더 적절해보인다.

  3. API Gateway에서 처리

마치며

처리율 제한과 관련된 다른 라이브러리들도 많지만, 단순히 처리율 제한을 위해서만 사용할 것이기에 Bucket4j를 선택하게 되었다.
나의 velog에 포스팅 된 Resilience4j를 사용할 수 있지만 단순 처리율 제한을 사용하기엔 Bucket4j보단 복잡하고 해줘야 할 설정이 많긴하다.
처리율 제한을 적용하여 외부로 부터의 무작위 공격에 대응하며 서버를 보호할 수 있게 되었다.

참고
https://github.com/bucket4j/bucket4j/tree/8.0?tab=readme-ov-file#java-compatibility-matrix
https://www.baeldung.com/spring-bucket4j
https://dkswnkk.tistory.com/732

profile
문제 해결과 개선 과제를 수행하며 성장을 추구하는 것을 좋아합니다.

0개의 댓글