根据ip限制接口访问次数

前言

我们利用redis去实现这个功能，redis的天然高并发和内存单线程速度拉满，非常适合做这个场景。为了可用性，我们把它封装成注解形式，哪个接口想被根据ip限制接口访问次数，直接标注上注解即可。

一、添加配置

在yaml文件中添加如下配置：

spring.redis.host: 172.xx.xx.xx
spring.redis.port: 6379
spring.redis.database: 1
spring.redis.password: tenxcloud

二、封装注解

封装一个注解使用，并且给一个默认值，防止空指针异常。

package com.xxx.ai.intelligentqa.annotate;

import java.lang.annotation.*;

/**
 * 接口访问频率注解，默认一分钟只能访问20次
 */

@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestLimit {

    // 限制时间 单位：秒(默认值：一分钟）
    long period() default 60;

    // 允许请求的次数(默认值：20次）
    long count() default 20;
}

三、主体逻辑切面实现

我们利用AOP的切面，来配合注解实现限流逻辑：

package com.xxx.ai.intelligentqa.aop;

import com.xxx.ai.intelligentqa.annotate.RequestLimit;
import com.xxx.ai.intelligentqa.tools.RequestUtil;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.BoundSetOperations;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletRequest;
import java.util.concurrent.TimeUnit;


@Aspect
@Component
public class RequestLimitAspect {

    @Autowired
    StringRedisTemplate redisTemplate;

    private static final Logger logger = LoggerFactory.getLogger(RequestLimitAspect.class);

    private static final String blackListKey = "ai:black_list";


    // 切点
    @Pointcut("@annotation(requestLimit)")
    public void controllerAspect(RequestLimit requestLimit) {}

    @Around("controllerAspect(requestLimit)")
    public Object doAround(ProceedingJoinPoint joinPoint, RequestLimit requestLimit) throws Throwable {
        SseEmitter emitter = new SseEmitter(0L);
        //获取当前请求request对象
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            throw new IllegalStateException("在上下文中没有请求的属性，可能是非web程序访问。");
        }
        HttpServletRequest request = attributes.getRequest();
        long period = requestLimit.period();
        long limitCount = requestLimit.count();

        String ip = RequestUtil.getIpAdrress(request); //解析出来真是请求的ip，防止多重代理ip攻击
        String uri = request.getRequestURI();
        String key = "ai:req_limit_".concat(uri).concat(":").concat(ip);
        // 检查用户IP是否在黑名单中
        BoundSetOperations<String, String> blackListOperations = redisTemplate.boundSetOps(blackListKey);
        if (blackListOperations.isMember(ip)) {
            logger.error("接口拦截：真实IP为{},已经在黑名单中", ip);
            //这里被我aop环绕的接口返回值是sse类型，所以此处我也需要使用sse形式返回。根据你接口返回值来
            SseEmitter emitter = new SseEmitter(0L);
            emitter.send(SseEmitter.event().name(AppConsts.EVENT_ERROR).data("您的请求过于频繁，请于5分钟后再次访问"));
            emitter.complete();
            return emitter;
        }

        ZSetOperations zSetOperations = redisTemplate.opsForZSet();

        // 添加当前时间戳
        long currentMs = System.currentTimeMillis();
        //zSetOperations.add(key, currentMs, currentMs);
        zSetOperations.add(key, String.valueOf(currentMs), currentMs);

        // 设置用户的过期时间
        redisTemplate.expire(key, period, TimeUnit.SECONDS);

        // 删除当前窗口之外的值(如果时间窗口是60秒，那么在60秒内的同一IP请求会被计数，超过60秒的请求就不应该再被计数了，因为它们已经滑出时间窗口了)
        Long aLong = zSetOperations.removeRangeByScore(key, 0, currentMs - period * 1000);

        // 检查所有可用计数
        Long count = zSetOperations.zCard(key);

        if (count > limitCount) {
            logger.error("接口拦截：{} 请求超过限制频率【{}次/{}s】,真实IP为{}", uri, limitCount, period, ip);

            // 将用户IP添加到黑名单并设置过期时间为5分钟（300秒）
            blackListOperations.add(ip);
            redisTemplate.expire(blackListKey, 300, TimeUnit.SECONDS);
            SseEmitter emitter = new SseEmitter(0L);
            emitter.send(SseEmitter.event().name(AppConsts.EVENT_ERROR).data("系统繁忙，请稍后重试"));
            emitter.complete();
            return emitter;
        }

        // 如果条件不成立，将继续执行 controller 层的方法
        return  joinPoint.proceed();
    }
}

四、把注解加在想要限流的接口上

@RequestLimit(count = 30)
    public Result knowledgeConverseEvents(QAAppDto dto, FacadeBase FacadeBase) {
        return appIntelligentAS.knowledgeConverseStream(dto);
    }