| | |
| | | package com.ruoyi.common.ratelimiter.aspectj; |
| | | |
| | | import cn.hutool.core.util.ArrayUtil; |
| | | import com.ruoyi.common.core.constant.GlobalConstants; |
| | | import com.ruoyi.common.core.exception.ServiceException; |
| | | import com.ruoyi.common.core.utils.MessageUtils; |
| | | import com.ruoyi.common.core.utils.ServletUtils; |
| | | import com.ruoyi.common.core.utils.StringUtils; |
| | | import com.ruoyi.common.ratelimiter.annotation.RateLimiter; |
| | | import com.ruoyi.common.ratelimiter.enums.LimitType; |
| | | import com.ruoyi.common.redis.utils.RedisUtils; |
| | |
| | | import org.aspectj.lang.annotation.Before; |
| | | import org.aspectj.lang.reflect.MethodSignature; |
| | | import org.redisson.api.RateType; |
| | | import org.springframework.core.DefaultParameterNameDiscoverer; |
| | | import org.springframework.core.ParameterNameDiscoverer; |
| | | import org.springframework.expression.EvaluationContext; |
| | | import org.springframework.expression.ExpressionParser; |
| | | import org.springframework.expression.ParserContext; |
| | | import org.springframework.expression.common.TemplateParserContext; |
| | | import org.springframework.expression.spel.standard.SpelExpressionParser; |
| | | import org.springframework.expression.spel.support.StandardEvaluationContext; |
| | | import org.springframework.stereotype.Component; |
| | | |
| | | import java.lang.reflect.Method; |
| | | |
| | |
| | | */ |
| | | @Slf4j |
| | | @Aspect |
| | | @Component |
| | | public class RateLimiterAspect { |
| | | |
| | | /** |
| | | * 定义spel表达式解析器 |
| | | */ |
| | | private final ExpressionParser parser = new SpelExpressionParser(); |
| | | /** |
| | | * 定义spel解析模版 |
| | | */ |
| | | private final ParserContext parserContext = new TemplateParserContext(); |
| | | /** |
| | | * 定义spel上下文对象进行解析 |
| | | */ |
| | | private final EvaluationContext context = new StandardEvaluationContext(); |
| | | /** |
| | | * 方法参数解析器 |
| | | */ |
| | | private final ParameterNameDiscoverer pnd = new DefaultParameterNameDiscoverer(); |
| | | |
| | | @Before("@annotation(rateLimiter)") |
| | | public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable { |
| | |
| | | } |
| | | long number = RedisUtils.rateLimiter(combineKey, rateType, count, time); |
| | | if (number == -1) { |
| | | throw new ServiceException(MessageUtils.message("rate.limiter.message")); |
| | | String message = rateLimiter.message(); |
| | | if (StringUtils.startsWith(message, "{") && StringUtils.endsWith(message, "}")) { |
| | | message = MessageUtils.message(StringUtils.substring(message, 1, message.length() - 1)); |
| | | } |
| | | throw new ServiceException(message); |
| | | } |
| | | log.info("限制令牌 => {}, 剩余令牌 => {}, 缓存key => '{}'", count, number, combineKey); |
| | | } catch (ServiceException e) { |
| | | throw e; |
| | | } catch (Exception e) { |
| | | throw new RuntimeException("服务器限流异常,请稍候再试"); |
| | | if (e instanceof ServiceException) { |
| | | throw e; |
| | | } else { |
| | | throw new RuntimeException("服务器限流异常,请稍候再试"); |
| | | } |
| | | } |
| | | } |
| | | |
| | | public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) { |
| | | StringBuilder stringBuffer = new StringBuilder(rateLimiter.key()); |
| | | if (rateLimiter.limitType() == LimitType.IP) { |
| | | // 获取请求ip |
| | | stringBuffer.append(ServletUtils.getClientIP()).append("-"); |
| | | } else if (rateLimiter.limitType() == LimitType.CLUSTER) { |
| | | // 获取客户端实例id |
| | | stringBuffer.append(RedisUtils.getClient().getId()).append("-"); |
| | | } |
| | | String key = rateLimiter.key(); |
| | | // 获取方法(通过方法签名来获取) |
| | | MethodSignature signature = (MethodSignature) point.getSignature(); |
| | | Method method = signature.getMethod(); |
| | | Class<?> targetClass = method.getDeclaringClass(); |
| | | stringBuffer.append(targetClass.getName()).append("-").append(method.getName()); |
| | | return stringBuffer.toString(); |
| | | // 判断是否是spel格式 |
| | | if (StringUtils.containsAny(key, "#")) { |
| | | // 获取参数值 |
| | | Object[] args = point.getArgs(); |
| | | // 获取方法上参数的名称 |
| | | String[] parameterNames = pnd.getParameterNames(method); |
| | | if (ArrayUtil.isEmpty(parameterNames)) { |
| | | throw new ServiceException("限流key解析异常!请联系管理员!"); |
| | | } |
| | | for (int i = 0; i < parameterNames.length; i++) { |
| | | context.setVariable(parameterNames[i], args[i]); |
| | | } |
| | | // 解析返回给key |
| | | try { |
| | | key = parser.parseExpression(key, parserContext).getValue(context, String.class) + ":"; |
| | | } catch (Exception e) { |
| | | throw new ServiceException("限流key解析异常!请联系管理员!"); |
| | | } |
| | | } |
| | | StringBuilder stringBuffer = new StringBuilder(GlobalConstants.RATE_LIMIT_KEY); |
| | | stringBuffer.append(ServletUtils.getRequest().getRequestURI()).append(":"); |
| | | if (rateLimiter.limitType() == LimitType.IP) { |
| | | // 获取请求ip |
| | | stringBuffer.append(ServletUtils.getClientIP()).append(":"); |
| | | } else if (rateLimiter.limitType() == LimitType.CLUSTER) { |
| | | // 获取客户端实例id |
| | | stringBuffer.append(RedisUtils.getClient().getId()).append(":"); |
| | | } |
| | | return stringBuffer.append(key).toString(); |
| | | } |
| | | } |