SpringBoot使用Redis对用户IP进行接口限流

使用接口限流的主要目的在于提高系统的稳定性，防止接口被恶意打击（短时间内大量请求）。

一、创建限流注解

引入redis依赖

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

创建一个自定义限流注解

package com.example.vuespringboot.annotation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * @author qx
 * @date 2023/7/27
 * @des 自定义限流注解
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiter {
    /**
     * 限流时间，单位秒
     *
     * @return
     */
    int time() default 5;

    /**
     * 限流次数
     *
     * @return
     */
    int count() default 10;
}

二、定义lua脚本

在resources/lua下新建limit.lua：

-- 获取redis键
local key = KEYS[1]
-- 获取第一个参数（次数）
local count = tonumber(ARGV[1])
-- 获取第二个参数（时间）
local time = tonumber(ARGV[2])
-- 获取当前流量
local current = redis.call('get', key);
-- 如果current值存在，且值大于规定的次数，则拒绝放行（直接返回当前流量）
if current and tonumber(current) > count then
    return tonumber(current)
end
-- 如果值小于规定次数，或值不存在，则允许放行，当前流量数+1  (值不存在情况下，可以自增变为1)
current = redis.call('incr', key);
-- 如果是第一次进来，那么开始设置键的过期时间。
if tonumber(current) == 1 then
    redis.call('expire', key, time);
end
-- 返回当前流量
return tonumber(current)

三、注入Lua执行脚本

package com.example.vuespringboot.config;

import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.PropertyAccessor;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
import org.springframework.scripting.support.ResourceScriptSource;

/**
 * @author qx
 * @date 2023/7/27
 * @des 自定义限流配置类
 */
@Configuration
public class RedisConfig {

    @Bean
    public RedisTemplate<Object, Object> redisTemplate(RedisConnectionFactory connectionFactory) {
        RedisTemplate<Object, Object> redisTemplate = new RedisTemplate<>();
        redisTemplate.setConnectionFactory(connectionFactory);
        // 使用Jackson2JsonRedisSerialize 替换默认序列化(默认采用的是JDK序列化)
        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(Object.class);
        ObjectMapper om = new ObjectMapper();
        om.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        om.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        jackson2JsonRedisSerializer.setObjectMapper(om);
        redisTemplate.setKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashValueSerializer(jackson2JsonRedisSerializer);
        return redisTemplate;
    }


    /**
     * 解析lua脚本的bean
     */
    @Bean("limitScript")
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
        redisScript.setResultType(Long.class);
        return redisScript;
    }

}

四、定义Aop切面

先定义一个获取IP地址的工具类

package com.example.vuespringboot.util;

import javax.servlet.http.HttpServletRequest;

/**
 * @author zhuxiaomeng
 * @date 2017/12/28
 */
public class IpUtil {

    public static String getIp(HttpServletRequest request) {
        String ip = request.getHeader("x-forwarded-for");
        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("X-Real-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}

接下来定义aop切面类

package com.example.vuespringboot.aop;

import cn.hutool.core.net.Ipv4Util;
import com.example.vuespringboot.annotation.RateLimiter;
import com.example.vuespringboot.util.IpUtil;
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.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.util.WebUtils;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * @author qx
 * @date 2023/7/27
 * @des 自定义限流aop切面类
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {

    @Autowired
    private RedisTemplate redisTemplate;
    @Autowired
    private RedisScript<Long> limitScript;

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable {
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(point);
        List<String> keys = Collections.singletonList(combineKey);
        try {
            Long number = (Long) redisTemplate.execute(limitScript, keys, count, time);
            // 当前流量number已超过限制，则抛出异常
            if (number == null || number.intValue() > count) {
                throw new RuntimeException("访问过于频繁，请稍后再试");
            }
            log.info("[limit] 限制请求数'{}',当前请求数'{}',缓存key'{}'", count, number.intValue(), combineKey);
        } catch (Exception ex) {
            ex.printStackTrace();
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }

    /**
     * 把用户IP和接口方法名拼接成 redis 的 key
     *
     * @param point 切入点
     * @return 组合key
     */
    private String getCombineKey(JoinPoint point) {
        StringBuilder sb = new StringBuilder("rate_limit:");
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();
        sb.append(IpUtil.getIp(request));


        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        // keyPrefix + "-" + class + "-" + method
        return sb.append("-").append(targetClass.getName())
                .append("-").append(method.getName()).toString();
    }
}

五、测试

package com.example.vuespringboot.controller;

import com.example.vuespringboot.annotation.RateLimiter;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;


/**
 * @author qx
 * @date 2023/7/26
 * @des 测试
 */
@RestController
public class IndexController {

    /**
     * 使用限流注解 5秒内只允许请求两次
     *
     * @return
     */
    @RateLimiter(time = 5, count = 2)
    @GetMapping("/")
    public String showData() {
        return "hello from spring boot";
    }

}

启动项目，在浏览器进行接口测试

第一次访问

SpringBoot使用Redis对用户IP进行接口限流-LMLPHP

2023-07-27 15:09:39.688  INFO 1456 --- [nio-8081-exec-6] c.e.vuespringboot.aop.RateLimiterAspect  : [limit] 限制请求数'2',当前请求数'1',缓存key'rate_limit:0:0:0:0:0:0:0:1-com.example.vuespringboot.controller.IndexController-showData'

第二次访问也正常

如果是第三次访问的话就会出现限流的消息提示