com.github.dennisit.vplus.data.interceptor.LimitInterceptor Maven / Gradle / Ivy
package com.github.dennisit.vplus.data.interceptor;
import com.github.dennisit.vplus.data.utils.WebUtils;
import com.google.common.collect.ImmutableList;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import java.io.Serializable;
import java.lang.reflect.Method;
import javax.servlet.http.HttpServletRequest;
/**
* // 限流 10秒内最多允许1000次访问
* // @Limit(key = "USER:LOGIN", period = 10, count = 1000)
* 使用aop拦截注解实现应用服务器分布式限流的目的,相同key共享限制次数
*/
@Slf4j
@Aspect
@Configuration
public class LimitInterceptor {
private final RedisTemplate redisTemplate;
@Autowired
public LimitInterceptor(RedisTemplate redisTemplate) {
this.redisTemplate = redisTemplate;
}
@Around("execution(public * *(..)) && @annotation(com.github.dennisit.vplus.data.interceptor.Limit)")
public Object interceptor(ProceedingJoinPoint pjp) {
MethodSignature signature = (MethodSignature) pjp.getSignature();
Method method = signature.getMethod();
Limit limitAnnotation = method.getAnnotation(Limit.class);
LimitType limitType = limitAnnotation.limitType();
String name = limitAnnotation.name();
String key;
int limitPeriod = limitAnnotation.period();
int limitCount = limitAnnotation.count();
switch (limitType) {
case IP:
key = WebUtils.getIP(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest());
break;
case CUSTOMER:
key = limitAnnotation.key();
break;
default:
key = StringUtils.upperCase(method.getName());
}
ImmutableList keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(), key));
try {
String luaScript = buildLuaScript();
RedisScript redisScript = new DefaultRedisScript<>(luaScript, Number.class);
Number count = redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
log.info("Access try count is {} for name={} and key = {}", count, name, key);
if (count != null && count.intValue() <= limitCount) {
return pjp.proceed();
} else {
throw new RuntimeException("You have been dragged into the blacklist");
}
} catch (Throwable e) {
if (e instanceof RuntimeException) {
throw new RuntimeException(e.getLocalizedMessage());
}
throw new RuntimeException("server exception");
}
}
/**
* 限流 脚本
*
* @return lua脚本
*/
public String buildLuaScript() {
StringBuilder lua = new StringBuilder();
lua.append("local c");
lua.append("\nc = redis.call('get',KEYS[1])");
// 调用不超过最大值，则直接返回
lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
lua.append("\nreturn c;");
lua.append("\nend");
// 执行计算器自加
lua.append("\nc = redis.call('incr',KEYS[1])");
lua.append("\nif tonumber(c) == 1 then");
// 从第一次调用开始限流，设置对应键值的过期
lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
lua.append("\nend");
lua.append("\nreturn c;");
return lua.toString();
}
private static final String UNKNOWN = "unknown";
}