springboot实现自定义注解限流-Peak-Gao

最近搭建的博客网站，详情被人刷了，特意以此来提醒该加限流处理了

引入依赖

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>

自定义注解实现，默认10秒内只能请求5次，当然这个是根据自己的实际情况修改

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    //5次
    int count() default 5
    //10秒
    int second() default 10
}

aop

/**
 * 限流注解
 * Created by PeakGao on 2023/3/2.
 */
@Aspect
@Component
public class IpLimitAspect extends AccessLimitIntercept {
    @Pointcut("@annotation(com.fyg.common.annotation.RateLimit)")
    public void rateLimit() {

    }

    @Before("rateLimit()")
    public void before(JoinPoint point) throws IOException, InterruptedException {
        long beginTime = System.currentTimeMillis()
        MethodSignature signature = (MethodSignature) point.getSignature()
        Method method = signature.getMethod()
        SysLog sysLog = new SysLog()
        RateLimit ipLimit = method.getAnnotation(RateLimit.class)
        if (ipLimit != null) {
            //注解上的描述
            sysLog.setOperation(String.valueOf(ipLimit.count()) + String.valueOf(ipLimit.second()))
        }

        //请求的参数
        Object[] args = point.getArgs()
        try {
            String params = new Gson().toJson(args)
            sysLog.setParams(params)
        } catch (Exception e) {

        }
        //请求的方法名
        String className = point.getTarget().getClass().getName()
        String methodName = signature.getName()
        sysLog.setMethod(className + "." + methodName + "()")
        //获取request
        HttpServletRequest request = HttpContextUtils.getHttpServletRequest()
        HttpServletResponse response = HttpContextUtils.getHttpServletResponse()
//        try {
        this.preHandle(request, response, ipLimit.count(), ipLimit.second())
//        } catch (Exception e) {
//            logger.error("限流内部程序出错：{}", e.getMessage());
//        }
        // 执行时长(毫秒)
        long time = System.currentTimeMillis() - beginTime
        logger.info("API：{}，限流注解程序执行：{} 毫秒", request.getRequestURI(), time)
    }

具体实现逻辑，采用ip限制控流

/**
 * 接口限流
 * 同一个接口10s内请求超过5次进行限流
 * Created by PeakGao on 2023/3/2.
 */
public class AccessLimitIntercept extends BaseController {

    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, int count, int second) throws InterruptedException, IOException {
        //文章搜索则不进行限流,如需部分接口地址限流可自定义注解实现
        // 拼接redis key = IP + Api限流
        String prefix = Constant.REDIS_KEY_PREFIX + Constant.RATE_LIMIT
        String key =  ipUtils.getIpAddr(request) + request.getRequestURI()
        // 获取redis的value
        Integer maxTimes = null
        Object value = redisUtil.get(prefix+key)
        if (value != null) {
            maxTimes = (Integer) value
        }
        if (maxTimes == null) {
            // 如果redis中没有该ip对应的时间则表示第一次调用，保存key到redis
            redisUtil.set(Constant.REDIS_KEY_PREFIX + Constant.RATE_LIMIT, key, 1, second)
        } else if (maxTimes < count) {
            // 如果redis中的时间比注解上的时间小则表示可以允许访问,这是修改redis的value时间
            redisUtil.set(Constant.REDIS_KEY_PREFIX + Constant.RATE_LIMIT, key, maxTimes + 1, second)
        } else {
            // 请求过于频繁
            output(response, "{"code":"8002","message":"请求过于频繁，请稍后再试"}")
            throw new RuntimeException("API请求限流拦截启动,当前接口：【" + key + "】请求过于频繁")
        }
        return true
    }


    public void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8")
        ServletOutputStream outputStream = null
        try {
            outputStream = response.getOutputStream()
            outputStream.write(msg.getBytes(StandardCharsets.UTF_8))
        } catch (IOException e) {
            e.printStackTrace()
        } finally {
            if (ObjectUtils.isNotEmpty(outputStream)) {
                outputStream.flush()
                outputStream.close()
            }
        }
    }

}