forked from dataease/dataease
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
b2a5cfa
commit 904d2c9
Showing
8 changed files
with
263 additions
and
10 deletions.
There are no files selected for viewing
24 changes: 24 additions & 0 deletions
24
backend/src/main/java/io/dataease/auth/annotation/DeRateLimiter.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
package io.dataease.auth.annotation; | ||
|
||
import org.springframework.core.annotation.AliasFor; | ||
|
||
import java.lang.annotation.*; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
@Target(ElementType.METHOD) | ||
@Retention(RetentionPolicy.RUNTIME) | ||
@Documented | ||
public @interface DeRateLimiter { | ||
|
||
long DEFAULT_REQUEST = 2; | ||
|
||
@AliasFor("max") long value() default DEFAULT_REQUEST; | ||
|
||
@AliasFor("value") long max() default DEFAULT_REQUEST; | ||
|
||
String key() default ""; | ||
|
||
long timeout() default 500; | ||
|
||
TimeUnit timeUnit() default TimeUnit.MILLISECONDS; | ||
} |
55 changes: 55 additions & 0 deletions
55
backend/src/main/java/io/dataease/auth/aop/DeRateLimiterHandler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
package io.dataease.auth.aop; | ||
|
||
import cn.hutool.core.util.StrUtil; | ||
import io.dataease.auth.annotation.DeRateLimiter; | ||
import io.dataease.auth.service.DeLimitService; | ||
import io.dataease.commons.utils.IPUtils; | ||
import io.dataease.commons.utils.ServletUtils; | ||
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.springframework.core.annotation.AnnotationUtils; | ||
import org.springframework.stereotype.Component; | ||
|
||
import javax.annotation.Resource; | ||
import java.lang.reflect.Method; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
@Aspect | ||
@Component | ||
public class DeRateLimiterHandler { | ||
|
||
private final static String SEPARATOR = ":"; | ||
|
||
|
||
@Resource | ||
private DeLimitService deLimitService; | ||
|
||
|
||
@Around(value = "@annotation(io.dataease.auth.annotation.DeRateLimiter)") | ||
public Object around(ProceedingJoinPoint point) throws Throwable { | ||
MethodSignature signature = (MethodSignature) point.getSignature(); | ||
Method method = signature.getMethod(); | ||
DeRateLimiter rateLimiter = AnnotationUtils.findAnnotation(method, DeRateLimiter.class); | ||
if (rateLimiter != null) { | ||
String key = rateLimiter.key(); | ||
if (StrUtil.isBlank(key)) { | ||
key = method.getDeclaringClass().getName() + StrUtil.DOT + method.getName(); | ||
} | ||
key = key + SEPARATOR + IPUtils.get(); | ||
|
||
long max = rateLimiter.max(); | ||
long timeout = rateLimiter.timeout(); | ||
TimeUnit timeUnit = rateLimiter.timeUnit(); | ||
Boolean limited = deLimitService.checkRestricted(key, max, timeout, timeUnit); | ||
if (limited) { | ||
String msg = "The current API [%s] is limited, please try again later!"; | ||
String requestURI = ServletUtils.request().getRequestURI(); | ||
throw new RuntimeException(String.format(msg, requestURI)); | ||
} | ||
} | ||
|
||
return point.proceed(); | ||
} | ||
} |
8 changes: 8 additions & 0 deletions
8
backend/src/main/java/io/dataease/auth/service/DeLimitService.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,8 @@ | ||
package io.dataease.auth.service; | ||
|
||
import java.util.concurrent.TimeUnit; | ||
|
||
public interface DeLimitService { | ||
|
||
Boolean checkRestricted(String key, long max, long timeout, TimeUnit timeUnit); | ||
} |
51 changes: 51 additions & 0 deletions
51
backend/src/main/java/io/dataease/auth/service/impl/RedisLimitServiceImpl.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
package io.dataease.auth.service.impl; | ||
|
||
import io.dataease.auth.service.DeLimitService; | ||
import io.dataease.commons.condition.RedisStatusCondition; | ||
import io.dataease.commons.utils.LogUtil; | ||
import org.slf4j.Logger; | ||
import org.springframework.context.annotation.Conditional; | ||
import org.springframework.context.annotation.Primary; | ||
import org.springframework.data.redis.core.StringRedisTemplate; | ||
import org.springframework.data.redis.core.script.RedisScript; | ||
import org.springframework.stereotype.Component; | ||
|
||
import javax.annotation.Resource; | ||
import java.time.Instant; | ||
import java.util.Collections; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
@Conditional({RedisStatusCondition.class}) | ||
@Component | ||
@Primary | ||
public class RedisLimitServiceImpl implements DeLimitService { | ||
|
||
Logger log = LogUtil.getLogger(); | ||
private final static String REDIS_LIMIT_KEY_PREFIX = "limit:"; | ||
@Resource | ||
private RedisScript<Long> limitRedisScript; | ||
|
||
@Resource | ||
private StringRedisTemplate stringRedisTemplate; | ||
|
||
@Override | ||
public Boolean checkRestricted(String key, long max, long timeout, TimeUnit timeUnit) { | ||
key = REDIS_LIMIT_KEY_PREFIX + key; | ||
long ttl = timeUnit.toMillis(timeout); | ||
long now = Instant.now().toEpochMilli(); | ||
long expired = now - ttl; | ||
|
||
Long executeTimes = stringRedisTemplate.execute(limitRedisScript, Collections.singletonList(key), now + "", ttl + "", expired + "", max + ""); | ||
if (executeTimes != null) { | ||
if (executeTimes == 0) { | ||
|
||
log.error("【{}】在单位时间 {} 毫秒内已达到访问上限,当前接口上限 {}", key, ttl, max); | ||
return true; | ||
} else { | ||
log.info("【{}】在单位时间 {} 毫秒内访问 {} 次", key, ttl, executeTimes); | ||
return false; | ||
} | ||
} | ||
return false; | ||
} | ||
} |
26 changes: 26 additions & 0 deletions
26
backend/src/main/java/io/dataease/auth/service/impl/StandaloneLimitServiceImpl.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
package io.dataease.auth.service.impl; | ||
|
||
|
||
import com.google.common.util.concurrent.RateLimiter; | ||
import io.dataease.auth.service.DeLimitService; | ||
import org.springframework.stereotype.Service; | ||
|
||
import java.util.concurrent.ConcurrentHashMap; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
|
||
@Service | ||
public class StandaloneLimitServiceImpl implements DeLimitService { | ||
|
||
private static ConcurrentHashMap<String, RateLimiter> RATE_LIMITER = new ConcurrentHashMap<>(); | ||
|
||
@Override | ||
public Boolean checkRestricted(String key, long max, long timeout, TimeUnit timeUnit) { | ||
RateLimiter rateLimiter = null; | ||
if (!RATE_LIMITER.containsKey(key)) { | ||
RATE_LIMITER.put(key, RateLimiter.create(max)); | ||
} | ||
rateLimiter = RATE_LIMITER.get(key); | ||
return !rateLimiter.tryAcquire(timeout, timeUnit); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
-- 下标从 1 开始 获取key | ||
local key = KEYS[1] | ||
-- 下标从 1 开始 获取参数 | ||
local now = tonumber(ARGV[1]) -- 当前时间错 | ||
local ttl = tonumber(ARGV[2]) -- 有效 | ||
local expired = tonumber(ARGV[3]) -- | ||
local max = tonumber(ARGV[4]) | ||
|
||
-- 清除过期的数据 | ||
-- 移除指定分数区间内的所有元素,expired 即已经过期的 score | ||
-- 根据当前时间毫秒数 - 超时毫秒数,得到过期时间 expired | ||
redis.call('zremrangebyscore', key, 0, expired) | ||
|
||
-- 获取 zset 中的当前元素个数 | ||
local current = tonumber(redis.call('zcard', key)) | ||
local next = current + 1 | ||
|
||
if next > max then | ||
-- 达到限流大小 返回 0 | ||
return 0; | ||
else | ||
-- 往 zset 中添加一个值、得分均为当前时间戳的元素,[value,score] | ||
redis.call("zadd", key, now, now) | ||
-- 每次访问均重新设置 zset 的过期时间,单位毫秒 | ||
redis.call("pexpire", key, ttl) | ||
return next | ||
end |