diff --git a/src/main/java/com/ruoyi/common/constant/Constants.java b/src/main/java/com/ruoyi/common/constant/Constants.java index 7fdf163..7a57fc0 100644 --- a/src/main/java/com/ruoyi/common/constant/Constants.java +++ b/src/main/java/com/ruoyi/common/constant/Constants.java @@ -74,6 +74,11 @@ public class Constants */ public static final String REPEAT_SUBMIT_KEY = "repeat_submit:"; + /** + * 限流 redis key + */ + public static final String RATE_LIMIT_KEY = "rate_limit:"; + /** * 验证码有效期(分钟) */ diff --git a/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java b/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java new file mode 100644 index 0000000..c31a67e --- /dev/null +++ b/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java @@ -0,0 +1,116 @@ +package com.ruoyi.framework.aspectj; + +import java.lang.reflect.Method; +import java.util.Collections; +import java.util.List; +import org.aspectj.lang.JoinPoint; +import org.aspectj.lang.Signature; +import org.aspectj.lang.annotation.Aspect; +import org.aspectj.lang.annotation.Before; +import org.aspectj.lang.annotation.Pointcut; +import org.aspectj.lang.reflect.MethodSignature; +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 com.ruoyi.common.exception.ServiceException; +import com.ruoyi.common.utils.ServletUtils; +import com.ruoyi.common.utils.StringUtils; +import com.ruoyi.common.utils.ip.IpUtils; +import com.ruoyi.framework.aspectj.lang.annotation.RateLimiter; +import com.ruoyi.framework.aspectj.lang.enums.LimitType; + +/** + * 限流处理 + * + * @author ruoyi + */ +@Aspect +@Component +public class RateLimiterAspect +{ + private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class); + + private RedisTemplate redisTemplate; + + private RedisScript limitScript; + + @Autowired + public void setRedisTemplate1(RedisTemplate redisTemplate) + { + this.redisTemplate = redisTemplate; + } + + @Autowired + public void setLimitScript(RedisScript limitScript) + { + this.limitScript = limitScript; + } + + // 配置织入点 + @Pointcut("@annotation(com.ruoyi.framework.aspectj.lang.annotation.RateLimiter)") + public void rateLimiterPointCut() + { + } + + @Before("rateLimiterPointCut()") + public void doBefore(JoinPoint point) throws Throwable + { + RateLimiter rateLimiter = getAnnotationRateLimiter(point); + String key = rateLimiter.key(); + int time = rateLimiter.time(); + int count = rateLimiter.count(); + + String combineKey = getCombineKey(rateLimiter, point); + List keys = Collections.singletonList(combineKey); + try + { + Long number = redisTemplate.execute(limitScript, keys, count, time); + if (StringUtils.isNull(number) || number.intValue() > count) + { + throw new ServiceException("访问过于频繁,请稍后再试"); + } + log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key); + } + catch (ServiceException e) + { + throw e; + } + catch (Exception e) + { + throw new RuntimeException("服务器限流异常,请稍后再试"); + } + } + + /** + * 是否存在注解,如果存在就获取 + */ + private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint) + { + Signature signature = joinPoint.getSignature(); + MethodSignature methodSignature = (MethodSignature) signature; + Method method = methodSignature.getMethod(); + + if (method != null) + { + return method.getAnnotation(RateLimiter.class); + } + return null; + } + + public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) + { + StringBuffer stringBuffer = new StringBuffer(rateLimiter.key()); + if (rateLimiter.limitType() == LimitType.IP) + { + stringBuffer.append(IpUtils.getIpAddr(ServletUtils.getRequest())); + } + MethodSignature signature = (MethodSignature) point.getSignature(); + Method method = signature.getMethod(); + Class targetClass = method.getDeclaringClass(); + stringBuffer.append("-").append(targetClass.getName()).append("- ").append(method.getName()); + return stringBuffer.toString(); + } +} diff --git a/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/Excels.java b/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/Excels.java index d9645a0..d2e5d74 100644 --- a/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/Excels.java +++ b/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/Excels.java @@ -14,5 +14,5 @@ import java.lang.annotation.Target; @Retention(RetentionPolicy.RUNTIME) public @interface Excels { - Excel[] value(); -} \ No newline at end of file + public Excel[] value(); +} diff --git a/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/RateLimiter.java b/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/RateLimiter.java new file mode 100644 index 0000000..722451e --- /dev/null +++ b/src/main/java/com/ruoyi/framework/aspectj/lang/annotation/RateLimiter.java @@ -0,0 +1,40 @@ +package com.ruoyi.framework.aspectj.lang.annotation; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; +import com.ruoyi.common.constant.Constants; +import com.ruoyi.framework.aspectj.lang.enums.LimitType; + +/** + * 限流注解 + * + * @author ruoyi + */ +@Target(ElementType.METHOD) +@Retention(RetentionPolicy.RUNTIME) +@Documented +public @interface RateLimiter +{ + /** + * 限流key + */ + public String key() default Constants.RATE_LIMIT_KEY; + + /** + * 限流时间,单位秒 + */ + public int time() default 60; + + /** + * 限流次数 + */ + public int count() default 100; + + /** + * 限流类型 + */ + public LimitType limitType() default LimitType.DEFAULT; +} diff --git a/src/main/java/com/ruoyi/framework/aspectj/lang/enums/LimitType.java b/src/main/java/com/ruoyi/framework/aspectj/lang/enums/LimitType.java new file mode 100644 index 0000000..12d9bb9 --- /dev/null +++ b/src/main/java/com/ruoyi/framework/aspectj/lang/enums/LimitType.java @@ -0,0 +1,20 @@ +package com.ruoyi.framework.aspectj.lang.enums; + +/** + * 限流类型 + * + * @author ruoyi + */ + +public enum LimitType +{ + /** + * 默认策略全局限流 + */ + DEFAULT, + + /** + * 根据请求者IP进行限流 + */ + IP +} diff --git a/src/main/java/com/ruoyi/framework/config/RedisConfig.java b/src/main/java/com/ruoyi/framework/config/RedisConfig.java index 7422f3f..322d6e3 100644 --- a/src/main/java/com/ruoyi/framework/config/RedisConfig.java +++ b/src/main/java/com/ruoyi/framework/config/RedisConfig.java @@ -6,6 +6,7 @@ import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; 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.StringRedisSerializer; import com.fasterxml.jackson.annotation.JsonAutoDetect; import com.fasterxml.jackson.annotation.JsonTypeInfo; @@ -47,4 +48,32 @@ public class RedisConfig extends CachingConfigurerSupport template.afterPropertiesSet(); return template; } + + @Bean + public DefaultRedisScript limitScript() + { + DefaultRedisScript 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 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 current;"; + } }