diff --git a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java index 1e21e376d..e63e47b1d 100644 --- a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java +++ b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java @@ -3,83 +3,54 @@ package com.ruoyi.framework.aspectj; import com.ruoyi.common.annotation.RateLimiter; import com.ruoyi.common.enums.LimitType; import com.ruoyi.common.exception.ServiceException; -import com.ruoyi.common.utils.ServletUtils; -import com.ruoyi.common.utils.StringUtils; +import com.ruoyi.common.utils.RedisUtils; +import lombok.extern.slf4j.Slf4j; 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.redisson.api.RateType; import org.springframework.stereotype.Component; import java.lang.reflect.Method; -import java.util.Collections; -import java.util.List; /** * 限流处理 * - * @author ruoyi + * @author Lion Li */ +@Slf4j @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; - } +public class RateLimiterAspect { // 配置织入点 @Pointcut("@annotation(com.ruoyi.common.annotation.RateLimiter)") - public void rateLimiterPointCut() - { + public void rateLimiterPointCut() { } @Before("rateLimiterPointCut()") - public void doBefore(JoinPoint point) throws Throwable - { + 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) - { + try { + RateType rateType = RateType.OVERALL; + if (rateLimiter.limitType() == LimitType.IP) { + rateType = RateType.PER_CLIENT; + } + // 返回 false 说明 获取令牌失败 + if (!RedisUtils.rateLimiter(key, rateType, count, time)) { throw new ServiceException("访问过于频繁,请稍后再试"); } - log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key); - } - catch (ServiceException e) - { + log.info("限制请求'{}',缓存key'{}'", count, key); + } catch (ServiceException e) { throw e; - } - catch (Exception e) - { + } catch (Exception e) { throw new RuntimeException("服务器限流异常,请稍后再试"); } } @@ -87,30 +58,15 @@ public class RateLimiterAspect /** * 是否存在注解,如果存在就获取 */ - private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint) - { + private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint) { Signature signature = joinPoint.getSignature(); MethodSignature methodSignature = (MethodSignature) signature; Method method = methodSignature.getMethod(); - if (method != null) - { + 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(ServletUtils.getClientIP()); - } - 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/ruoyi-framework/src/main/java/com/ruoyi/framework/config/RedisConfig.java b/ruoyi-framework/src/main/java/com/ruoyi/framework/config/RedisConfig.java index 59f9f977f..c1c303a26 100644 --- a/ruoyi-framework/src/main/java/com/ruoyi/framework/config/RedisConfig.java +++ b/ruoyi-framework/src/main/java/com/ruoyi/framework/config/RedisConfig.java @@ -1,7 +1,6 @@ package com.ruoyi.framework.config; import cn.hutool.core.util.ObjectUtil; -import cn.hutool.core.util.StrUtil; import com.ruoyi.common.utils.StringUtils; import com.ruoyi.framework.config.properties.RedissonProperties; import lombok.extern.slf4j.Slf4j; @@ -19,7 +18,6 @@ import org.springframework.cache.annotation.CachingConfigurerSupport; import org.springframework.cache.annotation.EnableCaching; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; -import org.springframework.data.redis.core.script.DefaultRedisScript; import java.io.IOException; import java.util.HashMap; @@ -125,32 +123,4 @@ public class RedisConfig extends CachingConfigurerSupport { return new RedissonSpringCacheManager(redissonClient, config, JsonJacksonCodec.INSTANCE); } - @Bean - public DefaultRedisScript limitScript() { - DefaultRedisScript redisScript = new DefaultRedisScript<>(); - redisScript.setScriptText(limitScriptText()); - redisScript.setResultType(Long.class); - return redisScript; - } - - /** - * 限流脚本 - */ - private String limitScriptText() { - return StrUtil.builder() - .append("local key = KEYS[1]\n") - .append("local count = tonumber(ARGV[1])\n") - .append("local time = tonumber(ARGV[2])\n") - .append("local current = redis.call('get', key);\n") - .append("if current and tonumber(current) > count then\n") - .append(" return current;\n") - .append("end\n") - .append("current = redis.call('incr', key)\n") - .append("if tonumber(current) == 1 then\n") - .append(" redis.call('expire', key, time)\n") - .append("end\n") - .append("return current;") - .toString(); - } - }