基于Redis实现API接口访问次数限制

package com.xxx.demo.aspect;
 
import java.lang.annotation.*;
 
 
/**
 * 接口访问频率注解，默认一分钟只能访问10次
 */
@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestLimit {
    // 限制时间 单位：秒(默认值：一分钟）
    long period() default 60;
    // 允许请求的次数(默认值：10次）
    long count() default 10;
}

package com.xxx.demo.aspect;
 
import com.xgd.demo.commons.ErrorCode;
import com.xgd.demo.handler.BusinessException;
import com.xgd.demo.util.IpUtil;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.log4j.Log4j2;
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.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
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 java.util.concurrent.TimeUnit;
 
/**
 * @date 2024/11/8 上午8:43
 */
@Aspect
@Component
@Log4j2
public class RequestLimitAspect {
    @Autowired
    RedisTemplate redisTemplate;
 
    @Pointcut("@annotation(requestLimit)")
    public void controllerAspect(RequestLimit requestLimit) {}
 
    @Around("controllerAspect(requestLimit)")
    public Object doAround(ProceedingJoinPoint joinPoint, RequestLimit requestLimit) throws Throwable {
        // 从注解中获取限制次数和窗口时间
        long period = requestLimit.period();
        long limitCount = requestLimit.count();
 
        // 请求
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        assert attributes != null;
        HttpServletRequest request = attributes.getRequest();
        String ip = IpUtil.getIpFromRequest(request);
        String uri = request.getRequestURI();
        //设置客户端访问的key
        String key = "req_limit_".concat(uri).concat(ip);
 
        ZSetOperations zSetOperations = redisTemplate.opsForZSet();
        // 添加当前时间戳，分数为当前时间戳
        long currentMs = System.currentTimeMillis();
        zSetOperations.add(key, currentMs, currentMs);
        // 设置窗口时间作为过期时间
        redisTemplate.expire(key, period, TimeUnit.SECONDS);
        // 移除掉不在窗口里的数据
        zSetOperations.removeRangeByScore(key, 0, currentMs - period * 1000);
        // 查询窗口内已经访问过的次数
        Long count = zSetOperations.zCard(key);
        if (count > limitCount) {
            log.error("接口拦截：{} 请求超过限制频率【{}次/{}s】,IP为{}", uri, limitCount, period, ip);
            throw new BusinessException(ErrorCode.REQUEST_LIMITED.getCode(), ErrorCode.REQUEST_LIMITED.getMessage());
        }
 
        // 继续执行请求
        return  joinPoint.proceed();
    }
}

package com.xxx.demo.util;
 
import jakarta.servlet.http.HttpServletRequest;
import java.util.Objects;
 
/**
 * @date 2024/11/8 上午9:06
 */
public class IpUtil {
    private static final String X_FORWARDED_FOR_HEADER = "X-Forwarded-For";
    private static final String X_REAL_IP_HEADER = "X-Real-IP";
 
    /**
     * 从请求中获取IP
     *
     * @return IP；当获取不到时，返回null
     */
    public static String getIpFromRequest(HttpServletRequest request ) {
        return getRealIp(request);
    }
 
    /**
     * 获取请求的真实IP，优先级从高到低为：<br/>
     * 1.从请求头X-Forwarded-For中获取ip，并且只获取第一个ip（从左到右） <br/>
     * 2.从请求头X-Real-IP中获取ip <br/>
     * 3.使用{@link HttpServletRequest#getRemoteAddr()}方法获取ip
     *
     * @param request 请求对象，必须不能为null
     * @return ip
     */
    private static String getRealIp(HttpServletRequest request) {
        Objects.requireNonNull(request, "request must be not null");
 
        String ip = request.getHeader(X_FORWARDED_FOR_HEADER);
        if (ip != null && !ip.isBlank()) {
            int delimiterIndex = ip.indexOf(',');
            if (delimiterIndex != -1) {
                // 如果存在多个ip，则取第一个ip
                ip = ip.substring(0, delimiterIndex);
            }
 
            return ip;
        }
 
        ip = request.getHeader(X_REAL_IP_HEADER);
        if (ip != null && !ip.isBlank()) {
            return ip;
        } else {
            return request.getRemoteAddr();
        }
    }
}