| | |
| | | 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<Object, Object> redisTemplate; |
| | | |
| | | private RedisScript<Long> limitScript; |
| | | |
| | | @Autowired |
| | | public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate) |
| | | { |
| | | this.redisTemplate = redisTemplate; |
| | | } |
| | | |
| | | @Autowired |
| | | public void setLimitScript(RedisScript<Long> 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<Object> 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("服务器限流异常,请稍后再试"); |
| | | } |
| | | } |
| | |
| | | /** |
| | | * 是否存在注解,如果存在就获取 |
| | | */ |
| | | 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(); |
| | | } |
| | | } |