io.github.dengchen2020.ratelimiter.local.LocalRateLimiterInterceptor Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of dc-spring-boot-starter-ratelimiter Show documentation
限流
There is a newer version: 0.0.28
Show newest version
package io.github.dengchen2020.ratelimiter.local;

import io.github.dengchen2020.core.utils.IPUtils;
import io.github.dengchen2020.ratelimiter.annotation.RateLimit;
import io.github.dengchen2020.ratelimiter.annotation.RateLimitStrategy;
import io.github.dengchen2020.ratelimiter.exception.RateLimitException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import java.security.Principal;
import java.time.Duration;
import java.util.concurrent.TimeUnit;

public class LocalRateLimiterInterceptor implements HandlerInterceptor {

    RateLimiter minuteRateLimiter = new RateLimiter(Duration.ofMinutes(1));

    RateLimiter secondRateLimiter = new RateLimiter(Duration.ofSeconds(1));

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (handler instanceof HandlerMethod handlerMethod) {
            RateLimit rateLimit = handlerMethod.getMethod().getAnnotation(RateLimit.class);
            if (rateLimit == null) rateLimit = handlerMethod.getBeanType().getAnnotation(RateLimit.class);
            if (rateLimit == null) return true;
            String uri = request.getRequestURI();
            RateLimiter rateLimiter;
            if (rateLimit.timeUnit() == TimeUnit.MINUTES) {
                rateLimiter = minuteRateLimiter;
            } else {
                rateLimiter = secondRateLimiter;
            }
            RateLimitStrategy strategy = rateLimit.strategy();
            String limitKey;
            if (strategy == RateLimitStrategy.userAndUri) {
                Principal principal = request.getUserPrincipal();
                if (principal == null) {
                    limitKey = IPUtils.getIpAddr(request) + uri;
                } else {
                    limitKey = principal.getName() + uri;
                }
            } else if (strategy == RateLimitStrategy.ip) {
                limitKey = IPUtils.getIpAddr(request);
            } else if (strategy == RateLimitStrategy.ipAndUri) {
                limitKey = IPUtils.getIpAddr(request) + uri;
            } else if (strategy == RateLimitStrategy.user) {
                Principal principal = request.getUserPrincipal();
                if (principal == null) {
                    limitKey = IPUtils.getIpAddr(request);
                } else {
                    limitKey = principal.getName();
                }
            } else if (strategy == RateLimitStrategy.uri) {
                limitKey = uri;
            } else {
                limitKey = IPUtils.getIpAddr(request) + uri;
            }
            if (rateLimiter.limit(limitKey, rateLimit.value())) {
                throw new RateLimitException(rateLimit.errorMsg());
            }
        }
        return true;
    }
}