RateLimitingFilter.java

package com.checkout.payment.gateway.security;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.filter.OncePerRequestFilter;

public class RateLimitingFilter extends OncePerRequestFilter {

    private final boolean rateLimitEnabled;
    private final boolean securityEnabled;
    private final int capacity;
    private final Duration refillPeriod;
    private final Map<String, Bucket> buckets = new ConcurrentHashMap<>();

    public RateLimitingFilter(
            @Value("${gateway.ratelimit.enabled:true}") boolean rateLimitEnabled,
            @Value("${gateway.security.enabled:true}") boolean securityEnabled,
            @Value("${gateway.ratelimit.capacity-per-window:120}") int capacity,
            @Value("${gateway.ratelimit.window:1m}") String window) {
        this.rateLimitEnabled = rateLimitEnabled;
        this.securityEnabled = securityEnabled;
        this.capacity = capacity;
        this.refillPeriod = parseWindow(window);
    }

    private static Duration parseWindow(String window) {
        String w = window.trim().toLowerCase();
        if (w.endsWith("m")) {
            return Duration.ofMinutes(Long.parseLong(w.substring(0, w.length() - 1)));
        }
        if (w.endsWith("s")) {
            return Duration.ofSeconds(Long.parseLong(w.substring(0, w.length() - 1)));
        }
        return Duration.ofMinutes(1);
    }

    @Override
    protected boolean shouldNotFilter(HttpServletRequest request) {
        if (!rateLimitEnabled || !securityEnabled) {
            return true;
        }
        String path = request.getRequestURI();
        if (path.startsWith("/actuator/health")
                || path.startsWith("/swagger-ui")
                || path.startsWith("/v3/api-docs")) {
            return true;
        }
        return false;
    }

    @Override
    protected void doFilterInternal(
            HttpServletRequest request,
            HttpServletResponse response,
            FilterChain filterChain) throws ServletException, IOException {

        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        String key = authentication != null && authentication.isAuthenticated()
                ? authentication.getName()
                : request.getRemoteAddr();

        Bucket bucket = buckets.computeIfAbsent(key, k -> newBucket());
        if (!bucket.tryConsume(1)) {
            response.setStatus(429);
            response.setContentType(MediaType.APPLICATION_JSON_VALUE);
            response.getWriter().write("{\"message\":\"Rate limit exceeded. Try again later.\"}");
            return;
        }

        filterChain.doFilter(request, response);
    }

    private Bucket newBucket() {
        Bandwidth limit = Bandwidth.builder()
                .capacity(capacity)
                .refillGreedy(capacity, refillPeriod)
                .build();
        return Bucket.builder().addLimit(limit).build();
    }
}