SpringBoot - 流量控制(二)：基于REDIS+LUA脚本的限流

电玩女神 2023-10-08 01:18 36阅读 0赞

#### 写在前面 ####

公司近期在做对外接口的限流，由于是轻量级的处理，所以采用“基于REDIS+LUA脚本”的方式进行访问控制。  
当然也可以使用RequestRateLimiterGatewayFilterFactory或者阿里巴巴的sentinel。

[SpringCloud - 流量控制(一)：基于RequestRateLimiterGatewayFilterFactory的限流][SpringCloud - _RequestRateLimiterGatewayFilterFactory]

#### 如何编写？ ####

###### ① 定义注解 ######

/**
     * 限流注解
     * @author ROCKY
     */
    @Target(ElementType.METHOD)
    @Retention(RetentionPolicy.RUNTIME)
    @Documented
    public @interface Limiter {
        
    
        // 限流KEY,用于REDIS的存储
        public String key()；
    
        // 限流时间(秒)
        public int time() default 60;
    
        // 限流次数
        public int count() default 100;
    
        // 限流类型
        public LimitType limitType() default LimitType.DEFAULT;
    }

###### ② 切面处理 ######

/**
     * 限流AOP
     * @author ROCKY
     */
    @Aspect
    @Component
    @Slf4j
    public class LimiterAspect {
        
    
        @Autowired
        private RedisTemplate<Object, Object> redisTemplate;
    
        @Autowired
        private RedisScript<Long> limitScript;
    
        @Before("@annotation(limiter)")
        public void doBefore(JoinPoint point, Limiter limiter) throws Throwable {
        
    	
            // 可以使用指定的KEY，也可以自己封装
            // final String limiterKey = limiter.key();
            final String limiterKey = getLimiterKey(limiter, point);
    
            int time = limiter.time();
            int count = limiter.count();
    
            List<Object> keys = Collections.singletonList(limiterKey);
            try {
        
                Long number = redisTemplate.execute(limitScript, keys, count, time);
                if (StringUtils.isEmpty(number) || number.intValue() > count) {
        
                    log.error("限流请求总次数为: '{}', 当前为第'{}'请求, 缓存KEY为: '{}'", count, number.intValue(), limiterKey);
                    throw new MyException(MyException.LIMITER_IN_THE_CONTROL);
                }
            } catch (MyException me) {
        
                throw me;
            } catch (Exception e) {
        
                throw new RuntimeException("限流控制内部异常，请稍候再试");
            }
        }
    
        // 获取限流数据的KEY(小写)
        public String getLimiterKey(Limiter limiter, JoinPoint point) {
        
            StringBuffer stringBuffer = new StringBuffer(Constants.LIMITER_IN_REDIS_PREFIX_KEY);
            if (limiter.limitType() == LimitType.IP) {
        
                stringBuffer.append(IpUtils.getIpAddr(ServletUtils.getRequest())).append("-");
            }
            if (limiter.limitType() == LimitType.USER) {
        
                // 从RequestContextHolder中获取用户名
                ...
            }
            if (limiter.limitType() == LimitType.TOKEN) {
        
                // 从RequestContextHolder中获取TOKEN
                ...
            }
    
            MethodSignature signature = (MethodSignature) point.getSignature();
            Method method = signature.getMethod();
            Class<?> targetClass = method.getDeclaringClass();
            stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
            return stringBuffer.toString().toLowerCase();
        }
    }

###### ③ 定义脚本 ######

/**
     * LUA限流脚本
     * @author ROCKY
     */
    @Configuration
    @EnableCaching
    public class RedisScriptConfig extends CachingConfigurerSupport {
        
        @Bean
        public DefaultRedisScript<Long> limitScript() {
        
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptText(limitScriptText());
            redisScript.setResultType(Long.class);
            return redisScript;
        }
    
        /**
         * 限流脚本
         */
        private String limitScriptText() {
        
            return "local key = KEYS[1]\n" +
                    "local count = tonumber(ARGV[1])\n" +
                    "local time = tonumber(ARGV[2])\n" +
                    "local current = redis.call('get', key);\n" +
                    "if current and tonumber(current) > count then\n" +
                    "    return tonumber(current);\n" +
                    "end\n" +
                    "current = redis.call('incr', key)\n" +
                    "if tonumber(current) == 1 then\n" +
                    "    redis.call('expire', key, time)\n" +
                    "end\n" +
                    "return tonumber(current);";
        }
    }

###### ④ 其他定义 ######

/**
     * 限流类型
     * @author ROCKY
     */
    
    public enum LimitType {
        
        // 默认策略全局限流
        DEFAULT,
    
        // 根据请求者IP进行限流
        IP,
    
        // 根据请求者进行限流
        USER,
    
        // 根据TOKEN进行限流
        TOKEN
    }

/**
     * SERVLET工具类
     * @author ROCKY
     */
    public class ServletUtils {
        
    
        // 获取String参数
        public static String getParameter(String name) {
        
            return getRequest().getParameter(name);
        }
    
        // 获取request
        public static HttpServletRequest getRequest() {
        
            try {
        
                return getRequestAttributes().getRequest();
            } catch (Exception e) {
        
                return null;
            }
        }
    
        // 获取response
        public static HttpServletResponse getResponse() {
        
            try {
        
                return getRequestAttributes().getResponse();
            } catch (Exception e) {
        
                return null;
            }
        }
    
        // 获取session
        public static HttpSession getSession() {
        
            return getRequest().getSession();
        }
    
        public static ServletRequestAttributes getRequestAttributes() {
        
            try {
        
                RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
                return (ServletRequestAttributes) attributes;
            } catch (Exception e) {
        
                return null;
            }
        }
    
        public static Map<String, String> getHeaders(HttpServletRequest request) {
        
            Map<String, String> map = new LinkedHashMap<>();
            Enumeration<String> enumeration = request.getHeaderNames();
            if (enumeration != null) {
        
                while (enumeration.hasMoreElements()) {
        
                    String key = enumeration.nextElement();
                    String value = request.getHeader(key);
                    map.put(key, value);
                }
            }
            return map;
        }
    
        /**
         * 将字符串渲染到客户端
         */
        public static String renderString(HttpServletResponse response, String string) {
        
            try {
        
                response.setStatus(200);
                response.setContentType("application/json");
                response.setCharacterEncoding("utf-8");
                response.getWriter().print(string);
            } catch (IOException e) {
        
                e.printStackTrace();
            }
            return null;
        }
    
        /**
         * 是否包含字符串
         */
        private static  boolean inStringIgnoreCase(String str, String... strs) {
        
            if (str != null && strs != null) {
        
                for (String s : strs) {
        
                    if (str.equalsIgnoreCase(StringUtils.trim(s))) {
        
                        return true;
                    }
                }
            }
            return false;
        }
    }

/**
     * 获取IP工具类
     * @author ROCKY
     */
    public class IpUtils {
        
        public static String getIpAddr(HttpServletRequest request) {
        
            if (request == null) {
        
                return null;
            }
    
            String ip = null;
    
            // X-Forwarded-For：Squid 服务代理
            String ipAddresses = request.getHeader("X-Forwarded-For");
            if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
        
                // Proxy-Client-IP：apache 服务代理
                ipAddresses = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
        
                // WL-Proxy-Client-IP：weblogic 服务代理
                ipAddresses = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
        
                // HTTP_CLIENT_IP：有些代理服务器
                ipAddresses = request.getHeader("HTTP_CLIENT_IP");
            }
            if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
        
                // X-Real-IP：nginx服务代理
                ipAddresses = request.getHeader("X-Real-IP");
            }
    
            // 有些网络通过多层代理，那么获取到的ip就会有多个，一般都是通过逗号（,）分割开来，并且第一个ip为客户端的真实IP
            if (ipAddresses != null && ipAddresses.length() != 0) {
        
                ip = ipAddresses.split(",")[0];
            }
    
            // 还是不能获取到，最后再通过request.getRemoteAddr();获取
            if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
        
                ip = request.getRemoteAddr();
            }
            return ip.equals("0:0:0:0:0:0:0:1") ? "127.0.0.1" : ip;
        }
    
        public static boolean internalIp(String ip) {
        
            byte[] addr = textToNumericFormatV4(ip);
            return internalIp(addr) || "127.0.0.1".equals(ip);
        }
    
        private static boolean internalIp(byte[] addr) {
        
            if (null == addr || addr.length < 2) {
        
                return true;
            }
            final byte b0 = addr[0];
            final byte b1 = addr[1];
            // 10.x.x.x/8
            final byte SECTION_1 = 0x0A;
            // 172.16.x.x/12
            final byte SECTION_2 = (byte) 0xAC;
            final byte SECTION_3 = (byte) 0x10;
            final byte SECTION_4 = (byte) 0x1F;
            // 192.168.x.x/16
            final byte SECTION_5 = (byte) 0xC0;
            final byte SECTION_6 = (byte) 0xA8;
            switch (b0) {
        
                case SECTION_1:
                    return true;
                case SECTION_2:
                    if (b1 >= SECTION_3 && b1 <= SECTION_4) {
        
                        return true;
                    }
                case SECTION_5:
                    switch (b1) {
        
                        case SECTION_6:
                            return true;
                    }
                default:
                    return false;
            }
        }
    
        // 将IPv4地址转换成字节
        public static byte[] textToNumericFormatV4(String text) {
        
            if (text.length() == 0) {
        
                return null;
            }
    
            byte[] bytes = new byte[4];
            String[] elements = text.split("\\.", -1);
            try {
        
                long l;
                int i;
                switch (elements.length) {
        
                    case 1:
                        l = Long.parseLong(elements[0]);
                        if ((l < 0L) || (l > 4294967295L)) {
        
                            return null;
                        }
                        bytes[0] = (byte) (int) (l >> 24 & 0xFF);
                        bytes[1] = (byte) (int) ((l & 0xFFFFFF) >> 16 & 0xFF);
                        bytes[2] = (byte) (int) ((l & 0xFFFF) >> 8 & 0xFF);
                        bytes[3] = (byte) (int) (l & 0xFF);
                        break;
                    case 2:
                        l = Integer.parseInt(elements[0]);
                        if ((l < 0L) || (l > 255L)) {
        
                            return null;
                        }
                        bytes[0] = (byte) (int) (l & 0xFF);
                        l = Integer.parseInt(elements[1]);
                        if ((l < 0L) || (l > 16777215L)) {
        
                            return null;
                        }
                        bytes[1] = (byte) (int) (l >> 16 & 0xFF);
                        bytes[2] = (byte) (int) ((l & 0xFFFF) >> 8 & 0xFF);
                        bytes[3] = (byte) (int) (l & 0xFF);
                        break;
                    case 3:
                        for (i = 0; i < 2; ++i) {
        
                            l = Integer.parseInt(elements[i]);
                            if ((l < 0L) || (l > 255L)) {
        
                                return null;
                            }
                            bytes[i] = (byte) (int) (l & 0xFF);
                        }
                        l = Integer.parseInt(elements[2]);
                        if ((l < 0L) || (l > 65535L)) {
        
                            return null;
                        }
                        bytes[2] = (byte) (int) (l >> 8 & 0xFF);
                        bytes[3] = (byte) (int) (l & 0xFF);
                        break;
                    case 4:
                        for (i = 0; i < 4; ++i) {
        
                            l = Integer.parseInt(elements[i]);
                            if ((l < 0L) || (l > 255L)) {
        
                                return null;
                            }
                            bytes[i] = (byte) (int) (l & 0xFF);
                        }
                        break;
                    default:
                        return null;
                }
            } catch (NumberFormatException e) {
        
                return null;
            }
            return bytes;
        }
    
        public static String getHostIp() {
        
            try {
        
                return InetAddress.getLocalHost().getHostAddress();
            } catch (UnknownHostException e) {
        
            }
            return "127.0.0.1";
        }
        public static String getHostPort(HttpServletRequest request) {
        
            return  String.valueOf(request.getLocalPort());
        }
    
        public static String getHostName() {
        
            try {
        
                return InetAddress.getLocalHost().getHostName();
            } catch (UnknownHostException e) {
        
            }
            return "未知";
        }
    }

#### 如何调试？ ####

[LUA调试器][LUA]

[SpringCloud - _RequestRateLimiterGatewayFilterFactory]: https://blog.csdn.net/goodjava2007/article/details/125859107?spm=1001.2014.3001.5501
[LUA]: https://redis.io/docs/manual/programmability/lua-debugging/