update 同步 ruoyi 自定义限流 新功能

This commit is contained in:
疯狂的狮子li 2021-08-17 14:39:38 +08:00
parent c271ffc066
commit 8cc673ec95
1 changed files with 116 additions and 116 deletions

View File

@ -1,116 +1,116 @@
package com.ruoyi.framework.aspectj; package com.ruoyi.framework.aspectj;
import java.lang.reflect.Method; import com.ruoyi.common.annotation.RateLimiter;
import java.util.Collections; import com.ruoyi.common.enums.LimitType;
import java.util.List; import com.ruoyi.common.exception.ServiceException;
import org.aspectj.lang.JoinPoint; import com.ruoyi.common.utils.ServletUtils;
import org.aspectj.lang.Signature; import com.ruoyi.common.utils.StringUtils;
import org.aspectj.lang.annotation.Aspect; import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Before; import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Pointcut; import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature; import org.aspectj.lang.annotation.Before;
import org.slf4j.Logger; import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.LoggerFactory; import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired; import org.slf4j.Logger;
import org.springframework.data.redis.core.RedisTemplate; import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.script.RedisScript; import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component; import org.springframework.data.redis.core.RedisTemplate;
import com.ruoyi.common.annotation.RateLimiter; import org.springframework.data.redis.core.script.RedisScript;
import com.ruoyi.common.enums.LimitType; import org.springframework.stereotype.Component;
import com.ruoyi.common.exception.ServiceException;
import com.ruoyi.common.utils.ServletUtils; import java.lang.reflect.Method;
import com.ruoyi.common.utils.StringUtils; import java.util.Collections;
import com.ruoyi.common.utils.ip.IpUtils; import java.util.List;
/** /**
* 限流处理 * 限流处理
* *
* @author ruoyi * @author ruoyi
*/ */
@Aspect @Aspect
@Component @Component
public class RateLimiterAspect public class RateLimiterAspect
{ {
private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class); private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);
private RedisTemplate<Object, Object> redisTemplate; private RedisTemplate<Object, Object> redisTemplate;
private RedisScript<Long> limitScript; private RedisScript<Long> limitScript;
@Autowired @Autowired
public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate) public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate)
{ {
this.redisTemplate = redisTemplate; this.redisTemplate = redisTemplate;
} }
@Autowired @Autowired
public void setLimitScript(RedisScript<Long> limitScript) public void setLimitScript(RedisScript<Long> limitScript)
{ {
this.limitScript = limitScript; this.limitScript = limitScript;
} }
// 配置织入点 // 配置织入点
@Pointcut("@annotation(com.ruoyi.common.annotation.RateLimiter)") @Pointcut("@annotation(com.ruoyi.common.annotation.RateLimiter)")
public void rateLimiterPointCut() public void rateLimiterPointCut()
{ {
} }
@Before("rateLimiterPointCut()") @Before("rateLimiterPointCut()")
public void doBefore(JoinPoint point) throws Throwable public void doBefore(JoinPoint point) throws Throwable
{ {
RateLimiter rateLimiter = getAnnotationRateLimiter(point); RateLimiter rateLimiter = getAnnotationRateLimiter(point);
String key = rateLimiter.key(); String key = rateLimiter.key();
int time = rateLimiter.time(); int time = rateLimiter.time();
int count = rateLimiter.count(); int count = rateLimiter.count();
String combineKey = getCombineKey(rateLimiter, point); String combineKey = getCombineKey(rateLimiter, point);
List<Object> keys = Collections.singletonList(combineKey); List<Object> keys = Collections.singletonList(combineKey);
try try
{ {
Long number = redisTemplate.execute(limitScript, keys, count, time); Long number = redisTemplate.execute(limitScript, keys, count, time);
if (StringUtils.isNull(number) || number.intValue() > count) if (StringUtils.isNull(number) || number.intValue() > count)
{ {
throw new ServiceException("访问过于频繁,请稍后再试"); throw new ServiceException("访问过于频繁,请稍后再试");
} }
log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key); log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key);
} }
catch (ServiceException e) catch (ServiceException e)
{ {
throw e; throw e;
} }
catch (Exception e) catch (Exception e)
{ {
throw new RuntimeException("服务器限流异常,请稍后再试"); throw new RuntimeException("服务器限流异常,请稍后再试");
} }
} }
/** /**
* 是否存在注解如果存在就获取 * 是否存在注解如果存在就获取
*/ */
private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint) private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint)
{ {
Signature signature = joinPoint.getSignature(); Signature signature = joinPoint.getSignature();
MethodSignature methodSignature = (MethodSignature) signature; MethodSignature methodSignature = (MethodSignature) signature;
Method method = methodSignature.getMethod(); Method method = methodSignature.getMethod();
if (method != null) if (method != null)
{ {
return method.getAnnotation(RateLimiter.class); return method.getAnnotation(RateLimiter.class);
} }
return null; return null;
} }
public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
{ {
StringBuffer stringBuffer = new StringBuffer(rateLimiter.key()); StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
if (rateLimiter.limitType() == LimitType.IP) if (rateLimiter.limitType() == LimitType.IP)
{ {
stringBuffer.append(IpUtils.getIpAddr(ServletUtils.getRequest())); stringBuffer.append(ServletUtils.getClientIP());
} }
MethodSignature signature = (MethodSignature) point.getSignature(); MethodSignature signature = (MethodSignature) point.getSignature();
Method method = signature.getMethod(); Method method = signature.getMethod();
Class<?> targetClass = method.getDeclaringClass(); Class<?> targetClass = method.getDeclaringClass();
stringBuffer.append("-").append(targetClass.getName()).append("- ").append(method.getName()); stringBuffer.append("-").append(targetClass.getName()).append("- ").append(method.getName());
return stringBuffer.toString(); return stringBuffer.toString();
} }
} }