springboot 自定义注解+拦截器+Redis实现ip限流

springboot 自定义注解+拦截器+Redis实现ip限流

自定义注解（具体频次可以根据具体使用场景调整）

import java.lang.annotation.*;
/**
 * IPAccessLimit.
 *
 * @date 2022/11/15
 * @since 1.0
 */
@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface IPAccessLimit {
    /**
     * 指定时间 单位：秒
     */
    int seconds() default 60;
    /**
     * 指定时间内API请求次数
     */
    int maxCount() default 50;
}

RequestInterceptor拦截器

import com.fasterxml.jackson.databind.ObjectMapper;
import com.kosun.ks.common.dto.RespDTO;
import com.kosun.ks.common.exception.ErrorCode;
import com.kosun.ks.common.util.RedisUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * RequestInterceptor拦截器.
 * @date 2022/11/2
 * @since 1.0
 */
@Slf4j
@Component
public class RequestInterceptor implements HandlerInterceptor {
    @Autowired
    private RedisUtil redisUtil;
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        try {
            if (handler instanceof HandlerMethod) {
                HandlerMethod handlerMethod = (HandlerMethod) handler;
                IPAccessLimit ipAccessLimit = handlerMethod.getMethodAnnotation(IPAccessLimit.class);
                if (null == ipAccessLimit) {
                    return true;
                }
                int seconds = ipAccessLimit.seconds();
                int maxCount = ipAccessLimit.maxCount();
                String ip = getIpAddress(request);
                String servletPath = request.getServletPath();
                String key = "IPAccessLimit:"+ip + ":" + request.getContextPath() + ":" + servletPath;
                // 已经访问的次数
                String count = redisUtil.get(key) == null ? null : redisUtil.get(key).toString();
                if (null != count) {
                    log.info("(ip限流请求次数) ip:{} 接口名:{} 访问次数:{}", ip, servletPath, count);
                }
                if (null == count || -1 == Integer.parseInt(count)) {
                    redisUtil.set(key, 1, seconds);
                    return true;
                }
                if (Integer.parseInt(count) < maxCount) {
                    redisUtil.incr(key, 1);
                    return true;
                }
                _response(response);
                return false;
            }
        } catch (Exception e) {
            log.error("(ip限流请求次数) 请求异常 ex:{}", e.getMessage());
        }
        return true;
    }

    /**
     * 拦截器异常响应
     *
     * @param response
     * @throws IOException
     */
    public void _response(HttpServletResponse response) throws IOException {
        response.setCharacterEncoding("UTF-8");
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        ObjectMapper objectMapper = new ObjectMapper();
        // R.fail 是接口响应统一封装返回
        response.getWriter().println(objectMapper.writeValueAsString(RespDTO.onFail(ErrorCode.TOO_MANY_REQUESTS)));
        return;
    }

    public static String getIpAddress(HttpServletRequest request) {
        String ip;
        try {
            ip = request.getHeader("x-forwarded-for");
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("Proxy-Client-IP");
            }
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_CLIENT_IP");
            }
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_X_FORWARDED_FOR");
            }
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getRemoteAddr();
            }
        } catch (Exception e) {
            ip = "未知IP";
        }
        return ip;
    }
}

注册 RequestInterceptor自定义拦截器

import com.kosun.ks.common.annotation.RequestInterceptor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurationSupport;
/**
 * 注册 RequestInterceptor自定义拦截器.
 * @date 2022/11/2
 * @since 1.0
 */
@Configuration
public class WebMvcConfig extends WebMvcConfigurationSupport {
    @Autowired
    RequestInterceptor requestInterceptor;
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry
                .addInterceptor(requestInterceptor)
                .addPathPatterns("/translate","/xxx");
//                .addPathPatterns("/**");
    }
}

最后再接口上加上自定义注解就可以使用了。