自定义限流注解
文章目录
-
- :apple: 准备
- :watermelon: 准备
- :banana: 测试
使用 redis 工具，整合 aop 做接口限流。
准备
- 添加依赖
web 依赖 aop 依赖 redis 依赖
- 添加
redis配置
spring.redis.host=localhost
spring.redis.port=6379
spring.redis.database=10
准备
定义常量类：
/**
* @author: yueLQ
* @date: 2022-08-17 19:47
*
* 枚举类：限流的类型
*/
public enum LimitType {
/**
* 默认的限流策略，针对某一个接口进行限流，例：一分钟限流一百次
*/
DEFAULT,
/**
* 针对某一个 ip 地址限流
*/
IP;
}
自定义限流注解 @RateLimiter :
package org.javaboy.rate_limit.annotation;
import org.javaboy.rate_limit.enums.LimitType;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
/**
* @author: yueLQ
* @date: 2022-08-17 19:50
*
* 限流注解
*/
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiter {
/**
* 限流的 key, 主要是指前缀
* @return
*/
String key() default "rate_limite:";
/**
* 限流时间窗，默认 60 秒
* @return
*/
int time() default 60;
/**
* 在时间窗内的限流次数
* @return
*/
int count() default 100;
LimitType limitType() default LimitType.DEFAULT;
}
在 resources 文件夹下，新建 lua 目录，编写 lua 脚本，进行 redis 数据库的操作：
-- 定义 key
local key=KEYS[1]
-- redis 调用时候的传值
local time= tonumber(ARGV[1])
local count=tonumber(ARGV[2])
-- 调用 get 方法
local current=redis.call('get',key)
if current and tonumber(current)>count then
return tonumber(current)
end
-- 进行自增加
current = redis.call('incr', key)
if tonumber(current)==1 then
redis.call('expire',key,time)
end
-- 返回最终结果
return tonumber(current)
redis 配置：
package org.javaboy.rate_limit.config;
import org.springframework.beans.factory.annotation.Autowired;
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: yueLQ
* @date: 2022-08-17 19:57
*
*
*/
@Configuration
public class RedisConfig {
/**
* 使用 RedisTemplate 存储JDK 序列化方式 （默认），存入到数据库中会乱码
* 换一种序列化的方案，(JSON 的序列化方案)
*/
@Bean
RedisTemplate<Object,Object> redisTemplate(RedisConnectionFactory redisConnectionFactory){
RedisTemplate<Object, Object> template = new RedisTemplate<>();
template.setConnectionFactory(redisConnectionFactory);
Jackson2JsonRedisSerializer<Object> serializer = new Jackson2JsonRedisSerializer<>(Object.class);
template.setKeySerializer(serializer);
template.setHashKeySerializer(serializer);
template.setValueSerializer(serializer);
template.setHashValueSerializer(serializer);
return template;
}
/**
* 配置加载，lua 脚本
* 泛型，就是我们 lua 脚本中返回值的数据类型
*/
@Bean
DefaultRedisScript<Long> limitScript(){
DefaultRedisScript<Long> script = new DefaultRedisScript<>();
script.setResultType(Long.class);
script.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
return script;
}
}
根据 ip 进行限流，创建 IpUtils 工具类：
package org.javaboy.rate_limit.utils;
import javax.servlet.http.HttpServletRequest;
import java.net.InetAddress;
import java.net.UnknownHostException;
/**
* @author: yueLQ
* @date: 2022-08-22 19:38
*/
public class IpUtils {
/**
* 获取客户端IP
*
* @param request 请求对象
* @return IP地址
*/
public static String getIpAddr(HttpServletRequest request) {
if (request == null) {
return "unknown";
}
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("X-Forwarded-For");
}
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 "0:0:0:0:0:0:0:1".equals(ip) ? "127.0.0.1" : getMultistageReverseProxyIp(ip);
}
/**
* 检查是否为内部IP地址
*
* @param ip IP地址
* @return 结果
*/
public static boolean internalIp(String ip) {
byte[] addr = textToNumericFormatV4(ip);
return internalIp(addr) || "127.0.0.1".equals(ip);
}
/**
* 检查是否为内部IP地址
*
* @param addr byte地址
* @return 结果
*/
private static boolean internalIp(byte[] addr) {
if (addr == null || 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地址转换成字节
*
* @param text IPv4地址
* @return byte 字节
*/
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;
}
/**
* 获取IP地址
*
* @return 本地IP地址
*/
public static String getHostIp() {
try {
return InetAddress.getLocalHost().getHostAddress();
} catch (UnknownHostException e) {
}
return "127.0.0.1";
}
/**
* 获取主机名
*
* @return 本地主机名
*/
public static String getHostName() {
try {
return InetAddress.getLocalHost().getHostName();
} catch (UnknownHostException e) {
}
return "未知";
}
/**
* 从多级反向代理中获得第一个非unknown IP地址
*
* @param ip 获得的IP地址
* @return 第一个非unknown IP地址
*/
public static String getMultistageReverseProxyIp(String ip) {
// 多级反向代理检测
if (ip != null && ip.indexOf(",") > 0) {
final String[] ips = ip.trim().split(",");
for (String subIp : ips) {
if (false == isUnknown(subIp)) {
ip = subIp;
break;
}
}
}
return ip;
}
/**
* 检测给定字符串是否为未知，多用于检测HTTP请求相关
*
* @param checkString 被检测的字符串
* @return 是否未知
*/
public static boolean isUnknown(String checkString) {
return checkString == null || checkString.length() == 0 || "unknown".equalsIgnoreCase(checkString);
}
}
自定义业务异常类 RateLimitException：
package org.javaboy.rate_limit.exception;
/**
* @author: yueLQ
* @date: 2022-08-22 19:09
*/
public class RateLimitException extends Exception {
public RateLimitException(String message) {
super(message);
}
}
全局异常类GlobalException 统一处理：
package org.javaboy.rate_limit.exception;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import java.util.HashMap;
import java.util.Map;
/**
* @author: yueLQ
* @date: 2022-08-22 19:10
*/
@RestControllerAdvice
public class GlobalException {
@ExceptionHandler(RateLimitException.class)
public Map<String,String> rateLimitException(RateLimitException e){
HashMap<String, String> map = new HashMap<>();
map.put("state","500");
map.put("msg",e.getMessage());
return map;
}
}
定义切面，进行限流操作：
package org.javaboy.rate_limit.aspectJ;
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.javaboy.rate_limit.annotation.RateLimiter;
import org.javaboy.rate_limit.enums.LimitType;
import org.javaboy.rate_limit.exception.RateLimitException;
import org.javaboy.rate_limit.utils.IpUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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 java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
/**
* @author: yueLQ
* @date: 2022-08-22 19:19
*/
@Aspect
@Component
public class RateLimiterAspect {
private final static Logger LOGGER = LoggerFactory.getLogger(RateLimiterAspect.class);
@Autowired
private RedisTemplate<Object, Object> template;
@Autowired
private RedisScript<Long> script;
/**
* 限流的话我们只需要一个前置通知即可
*/
@Before("@annotation(rateLimiter)")
public void before(JoinPoint joinPoint, RateLimiter rateLimiter) throws RateLimitException {
int time = rateLimiter.time();
int count = rateLimiter.count();
String combinKey = getCombinKey(rateLimiter, joinPoint);
try {
// 执行 lua 表达式
List<Object> keys = Collections.singletonList(combinKey);
Long num = template.execute(script, keys, time, count);
if (num == null || num.intValue() > count) {
// 超过限流了
LOGGER.info("当前接口，已达到最大限流次数!");
throw new RateLimitException("访问过于频繁，请稍后访问");
}
LOGGER.info("一个事件窗内，请求次数是{},当前请求次数{},缓存的 key 为{}",count,num,combinKey);
} catch (Exception e) {
e.printStackTrace();
throw e;
}
}
/**
* 这个 key 就是接口调用的次数，存储在 redis 中的 key
* key:127.0.0.1-org.javaboy.rate_limit.controller.HelloController-hello // 基于 ip 限流
* key:org.javaboy.rate_limit.controller.HelloController-hello // 基于接口限流
*
* @param rateLimiter
* @param joinPoint
* @return
*/
private String getCombinKey(RateLimiter rateLimiter, JoinPoint joinPoint) {
StringBuffer buffer = new StringBuffer(rateLimiter.key());
if (rateLimiter.limitType() == LimitType.IP) {
buffer.append(
IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest())
).append("-");
}
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method method = signature.getMethod();
// 获取类的路径
buffer.append(method.getDeclaringClass().getName())
.append("-")
.append(method.getName());
return buffer.toString();
}
}
创建接口进行测试：
package org.javaboy.rate_limit.controller;
import org.javaboy.rate_limit.annotation.RateLimiter;
import org.javaboy.rate_limit.enums.LimitType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
/**
* @author: yueLQ
* @date: 2022-08-22 20:12
*/
@RestController
public class HelloController {
@GetMapping("/hello")
/**
* 限流： 5秒之内，接口可以访问三次
*/
@RateLimiter(time = 10,count = 3,limitType = LimitType.IP)
public String hello(){
return "hello";
}
}
测试
