처리율 제한장치 적용 (bucket4j-redis)

이세원·2024년 7월 17일

배경

운영 서버에서 DOS 등 악의적이거나 과도한 트래픽을 받게되면 서버에 예상치 못한 방향으로 동작하거나 다운될 위험이 있다.

이를 고려하여 프로젝트에서도 이를 방지하기위한 기능인 처리율 제한장치를 구현해보려 한다.

처리율 제한장치 자체는 다양한 위치 및 구현방식이 존재하지만,
나는 서버측에서 처리율 제한장치를 구현해보려고 한다.

java+springboot로 서버를 운영중이기에 구현에 용이한 방식인 Bucket4j-redis를 사용하여 처리율 장치를 구현해보겠다.
(Api gateway를 사용해도 되지만 나중에 많이 벌어서 사용하겠다,,)

ref) https://bucket4j.com/8.9.0/toc.html#bucket4j-redis
공식문서만 달달 보면서 구현해본건 처음인 것 같다,,

갑자기 Redis?

라이브러리 이름을 보게되면 bucket4j-redis라고 되어있다.
갑자기 redis가 왜 튀어나온 것일까?

기본적으로 Bucket4j는 서버에서 인메모리로 동작하도록 되어있다.
하지만 이는 분산환경에서 큰 문제가 될 수도 있다.
(원하는 방향이 아닐 것이다.)

분산환경에서도 어떤 서버로 접근하든 동일한 처리율 제한 정책을 구현하기 위해서는
Redis와 같이 외부 스토리지를 이용해서 WAS와 분리해야한다.
(jwt토큰을 외부 스토리지에 저장하는것과 같은 맥락이다.)

Bucket4j도 역시 다른 많은 저장소를 지원하지만, Redis 사용경험이 있기에 이를 적용하기로 했다.

이제 코드를 보면서 살펴보자

구현

대부분 공식문서를 바탕으로 구현하였다.

build.gradle

    //Bucket4j
    implementation 'com.bucket4j:bucket4j-core:8.9.0'
    implementation 'com.bucket4j:bucket4j-redis:8.9.0'

RedisConfig

import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
import io.github.bucket4j.distributed.proxy.ProxyManager;
import io.github.bucket4j.redis.lettuce.cas.LettuceBasedProxyManager;
import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisURI;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.codec.ByteArrayCodec;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.codec.StringCodec;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.beans.factory.annotation.Value;
import java.time.Duration;

@Configuration
public class RedisConfig {

    @Value("${spring.data.redis.port}")
    private int port;
    @Value("${spring.data.redis.host}")
    private String host;

    private RedisClient redisClient() {
        return RedisClient.create(RedisURI.builder()
                .withHost(host)
                .withPort(port)
                .build());
    }

    /**
     * Reference
     * https://bucket4j.com/8.9.0/toc.html#bucket4j-redis
     */

    // proxy manager 가 redis 사용한다.
    // LettuceBasedProxyManager 는 비동기 고성능 처리 지원
    @Bean
    public ProxyManager<String> lettuceBasedProxyManager() {
        RedisClient redisClient = redisClient();
        StatefulRedisConnection<String, byte[]> redisConnection = redisClient
                .connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE)); // 키는 UTF-8 문자로, 값은 바이트 배열로 직렬화

        // Expiration 전략 설정
        return LettuceBasedProxyManager.builderFor(redisConnection)
                .withExpirationStrategy(
                        ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofMinutes(1L)))
                .build();
    }
}

RateLimitFilter

import io.github.bucket4j.*;
import io.github.bucket4j.distributed.proxy.ProxyManager;
import jakarta.servlet.Filter;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.ServletRequest;
import jakarta.servlet.ServletResponse;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

import static java.time.Duration.ofMinutes;

@Component
@Slf4j
public class RateLimitFilter implements Filter {

    /**
     * Reference
     * https://bucket4j.com/8.9.0/toc.html#bucket4j-redis
     */

    private ProxyManager<String> proxyManager;

    @Autowired
    public RateLimitFilter(ProxyManager<String> proxyManager) {
        this.proxyManager = proxyManager;
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse,
                         FilterChain filterChain) throws IOException, ServletException {

        HttpServletRequest httpRequest = (HttpServletRequest) servletRequest;
        String key = httpRequest.getRemoteAddr(); // ip 기반 처리율 제한, 키에 따라서 기반이 달라진다.
        Supplier<BucketConfiguration> bucketConfigurationSupplier = getConfigSupplier();
        Bucket bucket = proxyManager.builder().build(key, bucketConfigurationSupplier); // 버킷 생성
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1); // 이 메소드는 요청을 1개 소비 시도하고, 소비가 성공했는지 및 남아 있는 토큰 수를 반환합니다.

        if (probe.isConsumed()) {
            filterChain.doFilter(servletRequest, servletResponse);
        } else {
            // fail case
//            log.info("수용 불가능한 처리율" + "\trequest ip : " + key + "\turi : " + httpRequest.getRequestURI() + "\tremainingToken : " + probe.getRemainingTokens());
            HttpServletResponse httpServletResponse = makeRateLimitResponse(servletResponse, probe);
        }
    }

    // 공식문서상 여기 위치가 올바름
    // 파라미터로 인자 받아서 커스텀 가능 ex) id
    public Supplier<BucketConfiguration> getConfigSupplier() {

        // 토큰 개수 상황에 맞게 custom
        return () ->
                BucketConfiguration.builder()
                        .addLimit(limit -> limit.capacity(60).refillGreedy(60, ofMinutes(1))) // 토큰 리필 전략, 필요에 따라 메소드 선택
                        .build();
    }

    private HttpServletResponse makeRateLimitResponse(ServletResponse servletResponse, ConsumptionProbe probe) throws IOException {

        HttpServletResponse httpResponse = (HttpServletResponse) servletResponse;
        httpResponse.setContentType("text/plain");
        httpResponse.setHeader("X-Rate-Limit-Retry-After-Seconds", "" +
                TimeUnit.NANOSECONDS.toSeconds(probe.getNanosToWaitForRefill()));
        httpResponse.setStatus(429);
        httpResponse.getWriter().append("Too many requests");

        return httpResponse;
    }

}

FilterConfig

import lombok.RequiredArgsConstructor;
import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Arrays;

@Configuration
@RequiredArgsConstructor
public class FilterConfig {

    private final RateLimitFilter rateLimitFilter;
    private static final String[] INCLUDE_PATHS = {
            "/book/*"
    };

    @Bean
    public FilterRegistrationBean<RateLimitFilter> filterBean() {

        FilterRegistrationBean<RateLimitFilter> registrationBean
                = new FilterRegistrationBean<>();

        registrationBean.setFilter(rateLimitFilter);
        registrationBean.setOrder(1);
        registrationBean.setUrlPatterns(Arrays.asList(INCLUDE_PATHS));

        return registrationBean;
    }
}

키값을 무엇으로 주어야할까?

이부분은 본인이 어떤식으로 처리율을 제한하느냐에 따라서 커스텀 가능한 영역이다.
나는 request ip를 기준으로 제한을 두기로 하여서 위와 같이 설정하였다.
(api endpoint 마다 처리율을 제한하고 싶다면 key를 endpoint로 설정해야겠죠,,?)

토큰값은 얼마나 주어야할까?

기본적으로 bucket4j는 토큰버킷 알고리즘으로 동작한다.
토큰값을 튜닝하는것이 토큰버킷 알고리즘에서 까다로운 부분이다.
이 부분에 대해서는 각자의 프로젝트의 요구사항에 맞게 토큰 전략을 잘 세울필요가 있다.
(부하테스트를 진행하면서 원하는 요구사항에 맞는 적정량으로 커스텀하면 될 것 같다.)

동작결과

다음과 같이 처리율 제한장치의 동작을 로그로 확인해볼 수 있다.
로그를 잘 설정해둔다면 악의적인 요청들도 잘 파악할 수 있다고 생각된다.

(토큰 리필 전략에 따라서 한번에 60개가 공급되는것이 아닌 1분동안 수시로 공급된다.)

그리고 처리율 제한이 걸렸을때는 코드상에서도 확인할 수 있지만 status 429 및 Retry-After-Seconds 정보가 헤더에 담겨서 클라이언트에 전달된다.

아래는 influxdb + grafana로 서버의 요청처리를 시각화한 결과이다.
처리율을 벗어나는 요청의 경우 429 Status를 반환하며 운영되고있음을 확인할 수 있다.
(DOS 상황처럼 k6로 가상 유저를 생성하여 계속 요청해보았다.)

분산환경에서의 결과

아까 분산 환경을 위해 redis를 사용했다고 했는데 못 믿겠을 당신을 위해 결과를 첨부한다.
잘 동작한다.


0개의 댓글