Bucket4j - API Rate Limit 적용

현재 사이드프로젝트로 앱을 하나 출시했다. DDOS 공격이나, 악의적인 트래픽 공격으로부터 보호하기 위해 API Rate Limit을 적용해야 겠다고 생각이 들었다. Naver Cloud Platform에서 제공하는 서버를 사용하고 있었는데 혹시 자체적으로 Rate Limit 기능이 있나 찾아봤지만 제공을 하지 않는건지, 내가 못찾는건지? 정보를 찾을 수 없었다. 그래서 직접 애플리케이션단에 적용을 해야 겠다고 생각이 들었다.

API Rate Limit 알고리즘에는 여러가지가 있었지만 나는 Token Bucket 방식을 사용하기로 결정했다.

Token Bucket 방식이란 하나의 통안에 토큰을 정해진 수만큼 넣고 요청을 하려면 통에서 토큰을 꺼낸뒤 사용처리를 하고 요청을 진행하는 방식이다. 토큰이 없다면 제한된 수량을 다 넣었다고 판단하고 요청을 제한하게 된다. 물론 특정 시간후 다시 토큰을 채우게 하는 기능도 가능하다.

그중 대표적인 라이브러리인 Bucket4j를 사용하기로 결정했다.

우선 gradle에 의존성을 추가해준다.

// bucket4j
    implementation("com.giffing.bucket4j.spring.boot.starter:bucket4j-spring-boot-starter:0.5.2"){
        exclude(group= "org.ehcache", module= "ehcache")
    }

그리고 각 서비스 정책에 맞게 제한 정책을 정해줘야한다. 아직 특정 서비스 유저가 없어서 일반, 관리자로 나누어 정책을 설정했다.

@Service
class RateLimitService(
){

    init {
        for (value in LimitStrategy.values()) {
            rateLimitMap[value.name] = Bucket.builder().addLimit(value.getBucket()).build()
        }
    }

    fun resolveBucket(name: String): Bucket {
        return rateLimitMap[name]?: throw CustomException(CustomErrorCode.UNDEFINED_ERROR)
    }

    companion object{
        private val rateLimitMap = ConcurrentHashMap<String, Bucket>()
    }

    // policy
    private enum class LimitStrategy{
        ROLE_USER {
            override fun getBucket(): Bandwidth {
                return Bandwidth.classic(3000, Refill.intervally(3000, Duration.ofMinutes(60)))
            }
        },
        ROLE_ADMIN{
            override fun getBucket(): Bandwidth {
                return Bandwidth.classic(10000, Refill.intervally(10000, Duration.ofMinutes(60)))
            }
        };

        abstract fun getBucket(): Bandwidth

        companion object {
            private val strategyMap = values().associateBy(LimitStrategy::name)
            fun getStrategy(name: String) = strategyMap[name.uppercase()]?.let { it }?: throw CustomException(CustomErrorCode.UNDEFINED_ERROR)
        }
    }
}

그리고 controller단에 지정을 하거나 filter에서 rate limit를 체크하도록 설정할 수 있는데 저는 필터에서 처리하도록 결정하고 필터에 설정을 해줬습니다.

// check api rate limit
  val userRole = tokenService.getUserRole(token)
  val bucket = rateLimitService.resolveBucket(userRole)
  if (bucket.tryConsume(1)) {
   } else {
      throw CustomException(CustomErrorCode.TOO_MANY_REQUESTS)
  }

설정 위치는 고민후 서비스 성격에 맞게 위치를 지정해주면 됩니다.