springboot通过注解和redis实现接口限流

说明

本文通过 自定义注解 + Redis 实现接口限流。

实现

• 自定义注解

@Inherited
@Documented
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
// 默认就是60秒内可以访问2次
public @interface RequestLimit {
    /**
     * 时间，秒为单位
     */
    int second() default 60;
    /**
     * second时间内允许访问的次数
     */
    int maxCount() default 2;
}

• 通过AOP实现限流

@Configuration
@EnableAspectJAutoProxy
@Aspect
@Slf4j
@Order(-5) // 如果有多个AOP，需要将Order设置到最小，优先判断限流
public class RequestLimitAspect {
    // 这里通过将访问信息保存在redis中，判断是否限流
    @Autowired
    private RedisTemplate<String, String> redisTemplate;
    // 定义切点，所有添加RequestLimit注解的
    @Pointcut(value = "@annotation(org.xxx.xxx.annotation.RequestLimit)")
    private void handleRequestLimit() {
    }

    @Around("handleRequestLimit()")
    public Object handleResponseEncrypt(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RequestLimit requestLimit = method.getAnnotation(RequestLimit.class);
        // 没有添加注解直接放行
        if (null == requestLimit) {
            return pjp.proceed();
        }
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        // 获取HttpServletRequest
        HttpServletRequest request = (HttpServletRequest) requestAttributes.resolveReference(RequestAttributes.REFERENCE_REQUEST);
        // 获取公网IP
        String ip = CommonUtils.getPublicIp(request);
        // 实现同一个公网IP 限流
        String key = "limit:" + ip + ":" + request.getServletPath();
        String count = redisTemplate.opsForValue().get(key);
        if (StringUtils.isEmpty(count)) {
            count = "1";
            redisTemplate.opsForValue().set(key, count, requestLimit.second(), TimeUnit.SECONDS);
            return pjp.proceed();
        }
        if (Integer.valueOf(count) < requestLimit.maxCount()) {
            redisTemplate.opsForValue().increment(key);
            return pjp.proceed();
        }
        log.info("RemoteAddr:{} 请求接口:{} 请求过于频繁, 已拒绝请求.", ip, request.getServletPath());
        HttpServletResponse response = requestAttributes.getResponse();
        response.setCharacterEncoding("UTF-8");
        response.setContentType("application/json; charset=utf-8");
        response.getWriter()
                // 触发限流自定义返回内容
                .println(JSON.toJSONString(new Response(Constant.ResponseCode.RequestLimit, Constant.ResponseMsg.RequestLimit)));
        return null;

    }
}

• 获取公网IP

public static String getPublicIp(HttpServletRequest request) {
    String ip = null;
    try {
        ip = request.getHeader("x-forwarded-for");
        if (StringUtils.isEmpty(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (StringUtils.isEmpty(ip) || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (StringUtils.isEmpty(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_CLIENT_IP");
        }
        if (StringUtils.isEmpty(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_X_FORWARDED_FOR");
        }
        if (StringUtils.isEmpty(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
    } catch (Exception e) {
        log.error("getPublicIp ERROR ", e);
    }
    //使用代理，则获取第一个IP地址
    if (StringUtils.isEmpty(ip) && ip.length() > 15) {
        if (ip.indexOf(",") > 0) {
            ip = ip.substring(0, ip.indexOf(","));
        }
    }
    return ip;
}