JavaでRate Limitをする方法(bucket4j)

java

Rate Limitとは、一定時間あたりにアクセスできる回数(クライアントがリクエストする回数)に制限をかけることです。

過剰にアクセスしようとする悪意のあるユーザーに制限をかけることで、DOS対策になります。

ここでは、javaのspring bootを使っている時にrate limitを実装する方法を紹介します。　　

色々ライブラリーがあるようですが、ここではbucket4jというライブラリーを使いました。

github.com


ゲストとloginしているユーザー、またuserのroleに応じて制限する回数を変える仕様にしました。
そして、クライアントごとに制限をかけました。


Mavenの設定
<dependency>
        <groupId>com.giffing.bucket4j.spring.boot.starter</groupId>
        <artifactId>bucket4j-spring-boot-starter</artifactId>
        <version>0.2.0</version>
</dependency>


Interceptorクラス


インターセプタとは、メソッドの前後で任意の処理を実行させることができる機能のことで、HandlerInterceptorクラスを実装したクラスはこの機能を持ちます。
リクエストがある度に、メソッドの実行前後でログを取りたい時などにも使われるようです。
ここでは、リクエストがある度にrate Limitをしたいので、このクラスを用います。

preHandleメソッドは、リクエストが実行される前に行われる処理で、falseを返すとリクエストが実行できない(制限される)ようになります。


public class PerClientRateLimitInterceptor implements HandlerInterceptor {

    private Integer anonLimit = 10;   

    private Integer commonLimit = 40;

    private Integer vipLimit = 100;

    private final Map<String, Bucket> buckets = new ConcurrentHashMap<>();   //(1)

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {

        String path = request.getServletPath();
　　
　　 String remoteIp = getRemoteAddr(request);

        //loginしているユーザーの情報を取得。loginしていなければnull
        User currentUser = ShiroUtils.getUserDO();

        Bucket requestBucket;

        if (currentUser != null) {
            if (ShiroUtils.hasRole("vip")) {
                requestBucket = this.buckets.computeIfAbsent(currentUser.getAccount(), key -> vipBucket());  //(2)
            } else {
                requestBucket = this.buckets.computeIfAbsent(currentUser.getAccount(), key -> commonBucket());
            }
        } else {
            //ログインしていないユーザーはIPアドレス毎に制限
            requestBucket = this.buckets.computeIfAbsent(remoteIp, key -> anonBucket());
        }

        ConsumptionProbe probe = requestBucket.tryConsumeAndReturnRemaining(1);　//(3)
        if (probe.isConsumed()) {       //(4)
            response.addHeader("X-Rate-Limit-Remaining",
                    Long.toString(probe.getRemainingTokens()));
            return true;
        }

        // 429
        response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());   //(5)
        response.addHeader("X-Rate-Limit-Retry-After-Milliseconds",
                Long.toString(TimeUnit.NANOSECONDS.toMillis(probe.getNanosToWaitForRefill())));
        return false;
    }

    //anon用
    private Bucket anonBucket() {
        return Bucket4j.builder()
                .addLimit(Bandwidth.classic(anonLimit, Refill.intervally(anonLimit, Duration.ofMinutes(1))))     //(6)
                .build();
    }

    //common用
     private Bucket commonBucket() {
        return Bucket4j.builder()
                .addLimit(Bandwidth.classic(commonLimit, Refill.intervally(commonLimit, Duration.ofMinutes(1))))
                .build();
    }

    //vip用
     private Bucket vipBucket() {
        return Bucket4j.builder()
                .addLimit(Bandwidth.classic(vipLimit, Refill.intervally(vipLimit, Duration.ofMinutes(1))))
                .build();
    }

    //IPアドレス取得    
    private String getRemoteAddr(HttpServletRequest request) {     //(7)
        String xForwardedFor =  request.getHeader("X-Forwarded-For");
        //ELB等を経由していたらxForwardedForを返す
        if (xForwardedFor != null) {
            return xForwardedFor;
        }
        return request.getRemoteAddr();
    }

}


loginしているユーザーはaccount名ごとに、loginしていない場合はIPアドレスごとに制限をかけました。
また、ユーザーのroleに応じて制限する回数を変えました。loginしていない場合のroleはanonです。


(1)識別するクライアントがkey、制限する回数を数えるbucketがvalueです。

(2)ここではvipユーザーなのでvipBucketが使われています。

(3)bucketのカウントが1増えます。

(4)bucketのカウントが制限以内であればpreHandleメソッドはtrueを返します。

(5)probe.isConsumed()がfalseであれば、制限を超えたことになるので、responseのstatusを429に設定しpreHandleメソッドはfalseを返します。

(6)1分間でanonLimit回の制限をかけるbucketです。1分が経過するとbucket内のカウントがゼロになり、anonLimitを超えるとリクエストの制限が発動します。もちろん、1分間の部分や回数は自由に変えられます。

(7)クライアントのIPアドレスの取得方法です。詳しくはこちらで。

yu-memorandum.hatenablog.com


Configクラス

このクラスを作らないと、rate limitは機能しないので注意。

@Configuration
public class InterceptorConfig implements WebMvcConfigurer {
    @Bean
    public PerClientRateLimitInterceptor perClientRateLimitInterceptor() {
        return new PerClientRateLimitInterceptor();
    }

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(perClientRateLimitInterceptor()).addPathPatterns("/**"); 

    }

}


addInterceptorsメソッドで、用いるInterceptorクラスを指定しています。
また、addPathPatternsの部分でInterceptを適応するパスを指定できます。
ここでは全てのパスで適応する仕様。 　 　

これでクライアントごとにRate Limitをかけられるはずです。
ここではクライアント毎に、そしてroleに応じて制限する回数を変えたりと複雑になりましたが、色々と要望に応じてRate Limitを実装出来るようです。


実装するに当たって、この記事を参考にしました。
実装方法だけでなく、アルゴリズムも書いてあります。(英語ですが、google翻訳を使えば問題なく読めます、、、)

golb.hplar.ch