Rate Limiting với Bucket4j trong Spring Boot

<dependencies>
    <!-- Bucket4j Core & Redis Backend -->
    <dependency>
        <groupId>com.bucket4j</groupId>
        <artifactId>bucket4j-redis</artifactId>
        <version>8.10.1</version>
    </dependency>
    <!-- Lettuce Driver cho Redis -->
    <dependency>
        <groupId>io.lettuce</groupId>
        <artifactId>lettuce-core</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-redis</artifactId>
    </dependency>
</dependencies>

package com.tungdadev.config;
 
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.BucketConfiguration;
import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
import io.github.bucket4j.distributed.proxy.ClientSideConfig;
import io.github.bucket4j.distributed.proxy.ProxyManager;
import io.github.bucket4j.redis.lettuce.cas.LettuceBasedProxyManager;
import io.lettuce.core.RedisClient;
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.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
 
import java.time.Duration;
 
@Configuration
public class RateLimitConfig {
 
    // Spring Boot không expose sẵn bean RedisClient của Lettuce, nên tự tạo
    @Bean(destroyMethod = "shutdown")
    public RedisClient redisClient(@Value("${spring.data.redis.url:redis://localhost:6379}") String redisUrl) {
        return RedisClient.create(redisUrl);
    }
 
    @Bean(destroyMethod = "close")
    public StatefulRedisConnection<String, byte[]> rateLimitRedisConnection(RedisClient redisClient) {
        // Key là String, value là byte[] (trạng thái bucket đã serialize)
        return redisClient.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));
    }
 
    @Bean
    public ProxyManager<String> proxyManager(StatefulRedisConnection<String, byte[]> connection) {
        // builderFor(RedisClient) chỉ cho ProxyManager<byte[]>; muốn key String phải truyền connection có codec String
        return LettuceBasedProxyManager.builderFor(connection)
                .withClientSideConfig(ClientSideConfig.getDefault()
                        .withExpirationAfterWriteStrategy(
                                ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofMinutes(10))))
                .build();
    }
 
    // Định nghĩa hạn mức: 10 request / phút, nạp đều đặn
    public static BucketConfiguration standardTierConfig() {
        return BucketConfiguration.builder()
                .addLimit(Bandwidth.builder()
                        .capacity(10)
                        .refillGreedy(10, Duration.ofMinutes(1))
                        .build())
                .build();
    }
}

package com.tungdadev.annotation;
 
import java.lang.annotation.*;
 
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimited {
    String key() default "client_ip"; // Hoặc "api_key", "user_id"
    int capacity() default 10;
    int refillMinutes() default 1;
}

package com.tungdadev.interceptor;
 
import com.tungdadev.annotation.RateLimited;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.ConsumptionProbe;
import io.github.bucket4j.distributed.proxy.ProxyManager;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
 
import java.time.Duration;
 
@Component
@RequiredArgsConstructor
public class RateLimitInterceptor implements HandlerInterceptor {
 
    private final ProxyManager<String> proxyManager;
 
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (!(handler instanceof HandlerMethod handlerMethod)) {
            return true;
        }
 
        RateLimited rateLimited = handlerMethod.getMethodAnnotation(RateLimited.class);
        if (rateLimited == null) {
            return true;
        }
 
        // Lấy định danh: ví dụ API Key trong header hoặc Bearer token
        String apiKey = request.getHeader("X-API-KEY");
        // Gắn thêm tên method: mỗi endpoint một bucket riêng. Nếu dùng chung key, cấu hình (capacity)
        // chỉ được áp lúc bucket tạo lần đầu, endpoint khác sẽ dùng nhầm hạn mức của endpoint trước.
        String endpoint = handlerMethod.getBeanType().getSimpleName() + "." + handlerMethod.getMethod().getName();
        String limitKey = "rate_limit:" + endpoint + ":" + (apiKey != null ? apiKey : request.getRemoteAddr());
 
        // Lấy hoặc khởi tạo Bucket phân tán trong Redis
        Bucket bucket = proxyManager.builder().build(limitKey, () ->
            io.github.bucket4j.BucketConfiguration.builder()
                .addLimit(io.github.bucket4j.Bandwidth.builder()
                    .capacity(rateLimited.capacity())
                    .refillGreedy(rateLimited.capacity(), Duration.ofMinutes(rateLimited.refillMinutes()))
                    .build())
                .build()
        );
 
        // Thử tiêu thụ 1 token
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
 
        if (probe.isConsumed()) {
            // Cho phép đi tiếp, kèm theo Header thông báo hạn mức còn lại
            response.addHeader("X-RateLimit-Limit", String.valueOf(rateLimited.capacity()));
            response.addHeader("X-RateLimit-Remaining", String.valueOf(probe.getRemainingTokens()));
            return true;
        } else {
            // Hết lượt! Trả về HTTP 429 và thời gian cần chờ
            long waitForRefillNanos = probe.getNanosToWaitForRefill();
            // Làm tròn lên, làm tròn xuống thì client quay lại sớm quá và lại bị 429
            long waitForRefillSeconds = Math.max(1, (waitForRefillNanos + 999_999_999) / 1_000_000_000);
 
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            response.addHeader(HttpHeaders.RETRY_AFTER, String.valueOf(waitForRefillSeconds));
            response.setContentType("application/json");
            response.setCharacterEncoding("UTF-8"); // message có tiếng Việt
            response.getWriter().write(String.format(
                "{\"error\": \"Too Many Requests\", \"message\": \"Vượt quá giới hạn gọi API. Vui lòng thử lại sau %d giây.\"}",
                waitForRefillSeconds
            ));
            return false;
        }
    }
}

@Configuration
@RequiredArgsConstructor
public class WebMvcConfig implements WebMvcConfigurer {
 
    private final RateLimitInterceptor rateLimitInterceptor;
 
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(rateLimitInterceptor).addPathPatterns("/api/**");
    }
}

@RestController
@RequestMapping("/api/v1/orders")
public class OrderController {
 
    @PostMapping
    @RateLimited(capacity = 5, refillMinutes = 1) // Tối đa 5 lượt đặt hàng mỗi phút
    public ResponseEntity<String> placeOrder(@RequestBody OrderRequest request) {
        return ResponseEntity.ok("Đặt hàng thành công!");
    }
}