package com.xxx.demo.aspect;
 
import java.lang.annotation.*;
 
 
/**
 * 接口訪(fǎng)問(wèn)頻率注解，默認(rèn)一分鐘只能訪(fǎng)問(wèn)10次
 */
@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestLimit {
    // 限制時(shí)間 單位：秒(默認(rèn)值：一分鐘）
    long period() default 60;
    // 允許請(qǐng)求的次數(shù)(默認(rèn)值：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 {
        // 從注解中獲取限制次數(shù)和窗口時(shí)間
        long period = requestLimit.period();
        long limitCount = requestLimit.count();
 
        // 請(qǐng)求
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        assert attributes != null;
        HttpServletRequest request = attributes.getRequest();
        String ip = IpUtil.getIpFromRequest(request);
        String uri = request.getRequestURI();
        //設(shè)置客戶(hù)端訪(fǎng)問(wèn)的key
        String key = "req_limit_".concat(uri).concat(ip);
 
        ZSetOperations zSetOperations = redisTemplate.opsForZSet();
        // 添加當(dāng)前時(shí)間戳，分?jǐn)?shù)為當(dāng)前時(shí)間戳
        long currentMs = System.currentTimeMillis();
        zSetOperations.add(key, currentMs, currentMs);
        // 設(shè)置窗口時(shí)間作為過(guò)期時(shí)間
        redisTemplate.expire(key, period, TimeUnit.SECONDS);
        // 移除掉不在窗口里的數(shù)據(jù)
        zSetOperations.removeRangeByScore(key, 0, currentMs - period * 1000);
        // 查詢(xún)窗口內(nèi)已經(jīng)訪(fǎng)問(wèn)過(guò)的次數(shù)
        Long count = zSetOperations.zCard(key);
        if (count > limitCount) {
            log.error("接口攔截：{} 請(qǐng)求超過(guò)限制頻率【{}次/{}s】,IP為{}", uri, limitCount, period, ip);
            throw new BusinessException(ErrorCode.REQUEST_LIMITED.getCode(), ErrorCode.REQUEST_LIMITED.getMessage());
        }
 
        // 繼續(xù)執(zhí)行請(qǐng)求
        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";
 
    /**
     * 從請(qǐng)求中獲取IP
     *
     * @return IP；當(dāng)獲取不到時(shí)，返回null
     */
    public static String getIpFromRequest(HttpServletRequest request ) {
        return getRealIp(request);
    }
 
    /**
     * 獲取請(qǐng)求的真實(shí)IP，優(yōu)先級(jí)從高到低為：<br/>
     * 1.從請(qǐng)求頭X-Forwarded-For中獲取ip，并且只獲取第一個(gè)ip（從左到右） <br/>
     * 2.從請(qǐng)求頭X-Real-IP中獲取ip <br/>
     * 3.使用{@link HttpServletRequest#getRemoteAddr()}方法獲取ip
     *
     * @param request 請(qǐng)求對(duì)象，必須不能為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) {
                // 如果存在多個(gè)ip，則取第一個(gè)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();
        }
    }
}