bbps-backend/RateLimitingFilter.java

package com.billpaymentservice.bbps.filter;

import java.io.IOException;
import java.time.Instant;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.web.filter.OncePerRequestFilter;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.RedisConnectionFailureException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.billpaymentservice.bbps.ratelimit.LocalBucketManager;
import com.billpaymentservice.bbps.util.Helper;

import io.github.bucket4j.Bucket;


public class RateLimitingFilter extends OncePerRequestFilter {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitingFilter.class);
   
    private static final long RATE_PER_WINDOW = 1L;
    private static final long BURST = 0L;
    private static final long WINDOW_MILLIS = 60_000L;

    private final StringRedisTemplate redisTemplate;
    private final LocalBucketManager localBucketManager;
    private final ObjectMapper mapper = new ObjectMapper();

    private final DefaultRedisScript<Long> incrWithExpireScript;

    public RateLimitingFilter(StringRedisTemplate redisTemplate, LocalBucketManager localBucketManager) {
        this.redisTemplate = redisTemplate;
        this.localBucketManager = localBucketManager;
        String script =
                "local v = redis.call('INCR', KEYS[1]);" +
                "if v == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end;" +
                "return v";

        this.incrWithExpireScript = new DefaultRedisScript<>(script, Long.class);
    }

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

        String mid = validateAndExtractMid(request, response);
        if (mid == null) return;

        LOGGER.info("RateLimiter activated → path={} mid={}", request.getRequestURI(), mid);
        response.setHeader("X-RateLimit-Filter", "invoked");

        long limit = RATE_PER_WINDOW + BURST;
        String redisKey = buildRedisKey(mid);

        // 1️⃣ Local Bucket Check
        if (!consumeLocalToken(mid, response)) {
            long resetAt = calculateResetTime();
            setRateHeaders(response, limit, 0, resetAt);
            writeErrorResponse(request, response, 429, "Too Many Requests", "Rate limit exceeded (local)", mid);
            return;
        }

        // 2️⃣ Redis Counter
        Long count = incrementRedisCounter(redisKey);
        long resetAt = calculateResetTime();
        long remaining = Math.max(0, limit - count);

        setRateHeaders(response, limit, remaining, resetAt);
        response.setHeader("X-RateLimit-Count", String.valueOf(count));

        if (count > limit) {
            handleGlobalLimitExceeded(mid, response, request, resetAt);
            return;
        }

        filterChain.doFilter(request, response);
       
    }     
    
    //  -------------------------HELPERS-------------------------
    
    
    private String validateAndExtractMid(HttpServletRequest request, HttpServletResponse response) throws IOException 
    { String mid = request.getHeader("mid"); 
    if (mid == null || mid.trim().isEmpty()) 
    { writeErrorResponse(request, response, 400, "Bad Request", "Required header 'mid' is missing or empty", null); 
    return null; } 
    return mid.trim(); }
    
    private boolean consumeLocalToken(String mid, HttpServletResponse response) {
        Bucket bucket = localBucketManager.getBucket("local:mid:" + mid);
        boolean consumed = bucket.tryConsume(1);
        response.setHeader("X-RateLimit-Local-Remaining", String.valueOf(bucket.getAvailableTokens()));
        return consumed;
    }

    private Long incrementRedisCounter(String key) {
        long ttlSeconds = (WINDOW_MILLIS / 1000L) + 2;

        try {
            Long result = redisTemplate.execute(
                    incrWithExpireScript,
                    Collections.singletonList(key),
                    String.valueOf(ttlSeconds)
            );
            return (result == null ? 1 : result);
        } catch (Exception e) {
            LOGGER.warn("Redis unavailable → fallback mode");
            return 1L;
        }
    }

    private long calculateResetTime() {
        long nowSec = System.currentTimeMillis() / 1000L;
        long win = WINDOW_MILLIS / 1000L;
        return ((nowSec / win) * win) + win;
    }

    private String buildRedisKey(String mid) {
        long index = System.currentTimeMillis() / WINDOW_MILLIS;
        return "ratelimit:mid:" + mid + ":ts:" + index;
    }

    private void setRateHeaders(HttpServletResponse response, long limit, long remaining, long resetAt) {
        response.setHeader("X-RateLimit-Limit", String.valueOf(limit));
        response.setHeader("X-RateLimit-Remaining", String.valueOf(remaining));
        response.setHeader("X-RateLimit-Reset", String.valueOf(resetAt));
    }

    private void handleGlobalLimitExceeded(String mid,
                                           HttpServletResponse res,
                                           HttpServletRequest req,
                                           long resetAt) throws IOException {

        long now = System.currentTimeMillis() / 1000L;
        res.setHeader("Retry-After", String.valueOf(Math.max(1, resetAt - now)));

        writeErrorResponse(req, res, 429,
                "Too Many Requests", "Rate limit exceeded (global)", mid);
    }

    private void writeErrorResponse(HttpServletRequest request,
                                    HttpServletResponse response,
                                    int status, String error,
                                    String message, String mid) throws IOException {

        if (response.isCommitted()) return;

        Map<String, Object> body = new HashMap<>();
        body.put("timestamp", Instant.now().toString());
        body.put("status", status);
        body.put("error", error);
        body.put("message", message);
        body.put("path", request.getRequestURI());

        if (mid != null) body.put("mid", mid);

        response.setStatus(status);
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        mapper.writeValue(response.getWriter(), body);
    }
}