spring boot实战分享-自定义注解做接口限流

利用自定义注解做界面实现对于controller的接口访问限制 废话不多说直接看看最终的自定义注解长什么样

// 限制查询次数 每分钟只能查询10次
@AccessLimit(second = 60,max = 10)
@GetMapping("/getProjectList")
public R getProjectList(String codes) {
    return R.data();
}

我们只需要在需要限流的controller上面加上注解 就能够实现对于接口请求访问的限制，那么具体如何编写这个注解呢？代码如下

首先我们需要先创建一个注解类 AccessLimit.java

import java.lang.annotation.*;

/**
 * @program: carbon-offset
 * @description: 用户限流注解
 * @author: yyge
 * @create: 2023-11-15 14:13
 **/

@Documented
@Target(value = ElementType.METHOD) //可以注解在方法上面 用在controller层
@Retention(RetentionPolicy.RUNTIME)
public @interface AccessLimit {
   // 这里只是注解 现在还需要写一个切面
    int second() default 60; // 默认情况下 每1分钟 可以访问10次
    int max() default 10; // 默认情况下 每1分钟 可以访问10次
}

然后我们还需要去编写一个切面类AccessLimitAspect.java

/**
 * @program: carbon-offset
 * @description: 用户限流 切面
 * @author: yyf
 * @create: 2023-11-15 14:17
 **/

import com.chy.core.log.exception.ServiceException;
import com.chy.core.redis.cache.ChyRedis;
import com.chy.core.tool.utils.Func;
import com.chy.modules.auth.enums.I18nInfoEnum;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.time.Duration;
import java.util.Objects;


@Aspect
@Slf4j
@Component
@AllArgsConstructor
public class AccessLimitAspect {

    // 注入Redis服务 (redisTemplate)
    private final ChyRedis chyRedis;

    // 使用一个自定义的注解 定义为切点
    @Pointcut("@annotation(com.chy.modules.co.aspect.AccessLimit)")
    private void annotationPointCut() {
    }

    // 定义切面函数
    @Around("annotationPointCut() &&  @annotation(accessLimit)")
    public Object before(ProceedingJoinPoint joinPoint,com.chy.modules.co.aspect.AccessLimit accessLimit) throws Throwable {
        // 获取客户端的IP地址
        // 获取当前的HttpServletRequest对象
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();

        // 初始化变量addr和clientIp
        String addr  = null;
        String clientIp = null;

        // 定义需要检查的请求头名称数组
        String[] headerNames = {
          "X-Real-IP",
          "Client-Ip",
          "Proxy-Client-IP",
          "WL-Proxy-Client-IP",
          "HTTP_CLIENT_IP"
        };
        for (String header : headerNames) {
          String value = request.getHeader(header);
          if (value != null && value.length() != 0 && !"unknown".equalsIgnoreCase(value)) {
            clientIp = value;
            break;
          }
        }

        // 若clientIp依然为null，则使用getRemoteAddr()方法作为备选方案
        if(clientIp == null){
           addr  = request.getHeader("x-forwarded-for");
           if (addr==null){
              addr = request.getRemoteAddr();
           }
        }else{
           addr=clientIp;
        }

        
        //先获取一下方法名称 全路径
        String class_name = joinPoint.getTarget().getClass().getName(); //获取注解所在类的类名
        String method_name = joinPoint.getSignature().getName(); //获取注解所在类的方法名
        String redisMark="limit:"+addr+":"+class_name+"."+method_name; //生成redis的标识 方便查询次数

        // 现在是要根据ID去做标识
        //判断注解是否开启限制-如果开启限制则从Redis中查询数据 查看登记的次数
         Integer key = chyRedis.get(redisMark);
         if(Func.isNull(key)){
            long maxNumber= accessLimit.second();
            chyRedis.setEx(String.valueOf(redisMark),1, Duration.ofSeconds(maxNumber));
         }else{
             // 如果max是10次 则第11次就会提示错误
             if(key >= accessLimit.max()){
                 throw new ServiceException(I18nInfoEnum.SYSTEM_LIMIT_ERROR.getKey()); //错误提示
             }else{
                 //获取剩余过期时间
                 Long ttl = chyRedis.ttl(redisMark);
                 //value值+1
                 key++;
                 //重新设置值以及过期时间
                 chyRedis.setEx(String.valueOf(redisMark),key, Duration.ofSeconds(ttl));
             }
         }
        //返回并执行正常的controller 内的内容
        return joinPoint.proceed();
    }
}