io.github.dengchen2020.ratelimiter.local.LocalRateLimiterInterceptor Maven / Gradle / Ivy
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)) return true;
RateLimit rateLimit = handlerMethod.getMethod().getAnnotation(RateLimit.class);
if (rateLimit == null) rateLimit = handlerMethod.getBeanType().getAnnotation(RateLimit.class);
if (rateLimit == null) return true;
RateLimitStrategy strategy = rateLimit.strategy();
String limitKey;
switch (strategy) {
case userAndUri -> {
Principal principal = request.getUserPrincipal();
if (principal == null) {
limitKey = IPUtils.getIpAddr(request) + request.getRequestURI() + request.getMethod();
} else {
limitKey = principal.getName() + request.getRequestURI() + request.getMethod();
}
}
case ip -> limitKey = IPUtils.getIpAddr(request);
case ipAndUri -> limitKey = IPUtils.getIpAddr(request) + request.getRequestURI() + request.getMethod();
case user -> {
Principal principal = request.getUserPrincipal();
if (principal == null) {
limitKey = IPUtils.getIpAddr(request);
} else {
limitKey = principal.getName();
}
}
case uri -> limitKey = request.getRequestURI() + request.getMethod();
case null, default -> limitKey = IPUtils.getIpAddr(request) + request.getRequestURI() + request.getMethod();
}
RateLimiter rateLimiter;
if (rateLimit.timeUnit() == TimeUnit.MINUTES) {
rateLimiter = minuteRateLimiter;
} else {
rateLimiter = secondRateLimiter;
}
if (rateLimiter.limit(limitKey, rateLimit.value())) throw new RateLimitException(rateLimit.errorMsg());
return true;
}
}