springboot自定义注解+aop实现接口限流

oNuoyi

已于 2023-09-07 16:59:04 修改

阅读量210

点赞数

文章标签： spring boot java

于 2023-09-07 16:56:42 首次发布

本文链接：https://blog.csdn.net/qq_41973632/article/details/132692984

版权

先自定义一个限流注解RequestLimit

package com.nuoyi.limitrequest.common.annotation;

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

/**
 * @desc: 默认30/60s
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RequestLimit {

    int count() default 30;//次数

    long expire() default 60;//限定时长

    TimeUnit timeUnit() default TimeUnit.SECONDS;//时长单位


}

然后再通过aop切面进行限制需要限流的接口
aop主要做了：
1、拿到访问的ip判断是否开启了ip黑名单，开启了ip黑名单且ip在黑名单中则直接返回异常无权限
2、拿到请求ip+账号+接口的请求次数和限制的次数对比，不超过则+1放行
3、拿到请求ip+账号+接口的请求次数和限制的次数对比，超过则判断是否开启黑名单，开启则给ip添加进黑名单，然后禁用账号并返回异常请求频繁
超过请求限制的账号处理可自行根据需求调整

package com.nuoyi.limitrequest.aop;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.nuoyi.limitrequest.common.annotation.RequestLimit;
import com.nuoyi.limitrequest.common.config.AppConfiguration;
import com.nuoyi.limitrequest.dao.mapper.UserMapper;
import com.nuoyi.limitrequest.common.enums.ApiCodeEnum;
import com.nuoyi.limitrequest.common.enums.RedisKeyEnum;
import com.nuoyi.limitrequest.common.exception.ServiceException;
import com.nuoyi.limitrequest.common.util.IPUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;

/**
 * @desc: 请求限制aop
 * @Author: nuoyi
 */
@Component
@Aspect
@Slf4j
@RequiredArgsConstructor
public class ReqLimitAop {

    private final RedisTemplate<String, String> redisTemplate;
    private final UserMapper userMapper;
    private final AppConfiguration configuration;

    /**
     * 接口请求限制
     */
    @Pointcut("execution(* com.nuoyi.limitrequest.controller.*.*(..))")
    public void reqLimitPointcut() {
    }

    /**
     * 限流监控
     */
    @Before(value = "reqLimitPointcut() && @annotation(requestLimit)")
    public void reqLimitPointcutExecute(JoinPoint joinPoint, RequestLimit requestLimit) {

        String key = RedisKeyEnum.IpLimit.getKey();
        String blackKey = RedisKeyEnum.IpBlack.getKey();
        HttpServletRequest httpServletRequest = getHttpServletRequest();
        String ip = IPUtil.getIpAddr(getHttpServletRequest());
        String requestURI = httpServletRequest.getRequestURI();
        Integer accountId = 1;//动态获取 来自解析token后存储本地线程中获取


        //开启ip黑名单则优先判断ip是否存在黑名单中
        if (configuration.getIsOpenIpBlack() && StrUtil.isNotBlank(ip) && redisTemplate.opsForHash().hasKey(blackKey, ip)) {
            throw new ServiceException(ApiCodeEnum.NoAuthority.getCode(), ApiCodeEnum.NoAuthority.getDescription());
        }

        //缓存key+ip+账号id+请求地址  （可针对ip、账号做限制）
        String allKey = key.concat(":").concat(ip).concat(":").concat(Convert.toStr(accountId)).concat(":").concat(requestURI.replace("/", "_"));
        Integer count = StrUtil.isBlank(ip) ? 0 : StrUtil.isBlank(redisTemplate.opsForValue().get(allKey)) ? 0 : Convert.toInt(redisTemplate.opsForValue().get(allKey));
        if (count >= requestLimit.count()) {
            if (configuration.getIsOpenIpBlack()) {
                redisTemplate.opsForHash().put(blackKey, ip, JSON.toJSONString(accountId));
            }
            //账号禁用
            userMapper.updateAccountState(accountId);
            throw new ServiceException(ApiCodeEnum.RequestFrequent.getCode(), ApiCodeEnum.RequestFrequent.getDescription());
        } else if (count == 0) {
            redisTemplate.opsForValue().set(allKey, Convert.toStr(count + 1), requestLimit.expire(), requestLimit.timeUnit());
        } else {
            Long expire = redisTemplate.getExpire(allKey);
            redisTemplate.opsForValue().set(allKey, Convert.toStr(count + 1), null == expire ? requestLimit.expire() : expire, requestLimit.timeUnit());
        }

    }


    /**
     * 获取request对象
     *
     * @return request对象
     */
    private HttpServletRequest getHttpServletRequest() {
        RequestAttributes ra = RequestContextHolder.getRequestAttributes();
        ServletRequestAttributes sra = (ServletRequestAttributes) ra;
        if (ObjectUtil.isNull(sra)) {
            return null;
        }
        return sra.getRequest();
    }


}

AppConfiguration类是动态配置类，可自行调整
在这里插入图片描述

package com.nuoyi.limitrequest.common.config;

import lombok.Data;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

/**
 * @desc:
 * @Author: nuoyi
 */
@Configuration
@ConfigurationProperties("nuoyi")
@Data
public class AppConfiguration {
    /**
     * mysql链接地址
     */
    private String mysqlUrl;

    /**
     * mysql密码
     */
    private String mysqlPassword;

    /**
     * mysql用户名
     */
    private String mysqlUsername;

    /**
     * mysql端口号
     */
    private String mysqlPort;

    /**
     * redis主机地址
     */
    private String redisHost;

    /**
     * redis链接密码
     */
    private String redisPassword;

    /**
     * redis库
     */
    private Integer redisDataBase;

    /**
     * redis链接端口号
     */
    private Integer redisPort;

    /**
     * 是否开启ip黑名单
     */
    private Boolean isOpenIpBlack;

    /**
     * 项目数据源初始化失败是否自动退出
     */
    private Boolean isAutoExit;
}

自定义异常码枚举类ApiCodeEnum

package com.nuoyi.limitrequest.common.enums;

import lombok.AllArgsConstructor;
import lombok.Getter;

/**
 * @desc:
 * @Author: nuoyi
 */
@Getter
@AllArgsConstructor
public enum ApiCodeEnum {

    /**
     * 请求频繁
     */
    RequestFrequent(3008, "Request Frequent, Please wait a moment and visit again"),

    /**
     * 暂无权限 或 暂时没有操作机会
     */
    NoAuthority(3007, "no authority");

    private final Integer code;

    private final String description;
}

自定义异常类ServiceException

package com.nuoyi.limitrequest.common.exception;

import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;

/**
 * 服务异常
 *
 * @author nuoyi
 */
@Data
@NoArgsConstructor
@EqualsAndHashCode(callSuper = true)
public class ServiceException extends RuntimeException {

    private static final long serialVersionUID = -8032230285855412076L;

    /**
     * 错误码
     */
    private Integer code;

    /**
     * 错误消息
     */
    private String message;

    /**
     * 错误对象
     */
    private Exception exception;

    public ServiceException(Integer code, String message) {
        this.code = code;
        this.message = message;
    }

    public ServiceException(String message, Integer code) {
        this.code = code;
        this.message = message;
    }
    public ServiceException(Exception e) {
        this.exception = e;
    }

    public ServiceException(Integer code, Exception e) {
        this.code = code;
        this.exception = e;
    }
}

解析IP工具类IPUtil

package com.nuoyi.limitrequest.common.util;

import cn.hutool.core.util.ObjectUtil;

import javax.servlet.http.HttpServletRequest;

/**
 * @desc: ip工具类
 * @Author: nuoyi
 */
public class IPUtil {
    /**
     * 获取IP
     *
     * @param request 请求
     * @return ip
     */
    public static String getIpAddr(HttpServletRequest request) {
        if (ObjectUtil.isNull(request)) {
            return "No IP";
        }
        String ip = request.getHeader("x-forwarded-for");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            // 多次反向代理后会有多个ip值，第一个ip才是真实ip
            if (ip.contains(",")) {
                ip = ip.split(",")[0];
            }
        }
        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.getHeader("X-Real-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip) || "0:0:0:0:0:0:0:1".equals(ip)) {
            return "No IP";
        }
        return ip;
    }

}

redis的key枚举类

package com.nuoyi.limitrequest.common.enums;

import lombok.Getter;

/**
 * @desc:
 * @Author: nuoyi
 */
public enum RedisKeyEnum {

    /**
     * 项目缓存前缀
     */
    CACHE_KEY_PREFIX("Nuoyi:"),

    /**
     * token
     * */
    Token("Token"),

    /**
     * ip请求
     */
    IpLimit("Ip:Limit"),

    /**
     * ip黑名单地址
     */
    IpBlack("Ip:Black"),

    ;

    @Getter
    private String name;

    RedisKeyEnum(String name) {
        this.name = name;
    }

    /**
     * 获取全称缓存路径
     */
    public String getKey() {
        return RedisKeyEnum.CACHE_KEY_PREFIX.name + this.getName();
    }

}

最后在需要的接口上添加限流注解就可以进行测试了

package com.nuoyi.limitrequest.controller;

import com.nuoyi.limitrequest.common.annotation.RequestLimit;
import lombok.RequiredArgsConstructor;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.concurrent.TimeUnit;

/**
 * @desc:
 * @Author: nuoyi
 */
@RestController
@RequestMapping("test")
@RequiredArgsConstructor
@Validated
public class TestController {

    /**
     * 首页
     */
    @GetMapping("index")
    @RequestLimit
    public String index() {
        return "访问主页";
    }

    /**
     * 首页1
     */
    @GetMapping("index1")
    @RequestLimit(count = 20, expire = 60, timeUnit = TimeUnit.SECONDS)
    public String index1() {
        return "访问主页";
    }
}