From 8753c6f8c791ff01de2976dd725c39deeaee2872 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E5=8A=9B?= <soldierliuli@hotmail.com>
Date: Thu, 12 May 2022 13:28:10 +0800
Subject: [PATCH] fix bug

---
 .../com/storeroom/aspect/LimitAspect.java     | 82 +++++++++++++++++++
 .../com/storeroom/config/SwaggerConfig.java   |  6 +-
 .../com/storeroom/utils/RequestHolder.java    | 14 ++++
 3 files changed, 99 insertions(+), 3 deletions(-)
 create mode 100644 common/src/main/java/com/storeroom/aspect/LimitAspect.java
 create mode 100644 common/src/main/java/com/storeroom/utils/RequestHolder.java

diff --git a/common/src/main/java/com/storeroom/aspect/LimitAspect.java b/common/src/main/java/com/storeroom/aspect/LimitAspect.java
new file mode 100644
index 0000000..0f908f9
--- /dev/null
+++ b/common/src/main/java/com/storeroom/aspect/LimitAspect.java
@@ -0,0 +1,82 @@
+package com.storeroom.aspect;
+
+import com.google.common.collect.ImmutableList;
+import com.storeroom.annotaion.Limit;
+import com.storeroom.exception.BaseException;
+import com.storeroom.utils.RequestHolder;
+import com.storeroom.utils.StringUtils;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+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.stereotype.Component;
+
+import javax.servlet.http.HttpServletRequest;
+import java.lang.reflect.Method;
+
+@Aspect
+@Component
+public class LimitAspect {
+
+    private final RedisTemplate<Object,Object> redisTemplate;
+    private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);
+
+    public LimitAspect(RedisTemplate<Object,Object> redisTemplate) {
+        this.redisTemplate = redisTemplate;
+    }
+
+    @Pointcut("@annotation(com.storeroom.annotaion.Limit)")
+    public void pointcut() {
+    }
+
+    @Around("pointcut()")
+    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
+        HttpServletRequest request = RequestHolder.getHttpServletRequest();
+        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
+        Method signatureMethod = signature.getMethod();
+        Limit limit = signatureMethod.getAnnotation(Limit.class);
+        LimitType limitType = limit.limitType();
+        String key = limit.key();
+        if (StringUtils.isEmpty(key)) {
+            if (limitType == LimitType.IP) {
+                key = StringUtils.getIp(request);
+            } else {
+                key = signatureMethod.getName();
+            }
+        }
+
+        ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));
+
+        String luaScript = buildLuaScript();
+        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
+        Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
+        if (null != count && count.intValue() <= limit.count()) {
+            logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name());
+            return joinPoint.proceed();
+        } else {
+            throw new BaseException("访问次数受限制");
+        }
+    }
+
+    /**
+     * 限流脚本
+     */
+    private String buildLuaScript() {
+        return "local c" +
+                "\nc = redis.call('get',KEYS[1])" +
+                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
+                "\nreturn c;" +
+                "\nend" +
+                "\nc = redis.call('incr',KEYS[1])" +
+                "\nif tonumber(c) == 1 then" +
+                "\nredis.call('expire',KEYS[1],ARGV[2])" +
+                "\nend" +
+                "\nreturn c;";
+    }
+}
diff --git a/common/src/main/java/com/storeroom/config/SwaggerConfig.java b/common/src/main/java/com/storeroom/config/SwaggerConfig.java
index 1fee6e6..9623cf2 100644
--- a/common/src/main/java/com/storeroom/config/SwaggerConfig.java
+++ b/common/src/main/java/com/storeroom/config/SwaggerConfig.java
@@ -65,9 +65,9 @@ public class SwaggerConfig {
 
     private ApiInfo apiInfo() {
         return new ApiInfoBuilder()
-                .description("阅行客后台管理系统")
-                .title("YXK-ADMIN 接口文档")
-                .version("2.4")
+                .description("智能库房综合管理系统")
+                .title("接口文档")
+                .version("1.0")
                 .build();
     }
 
diff --git a/common/src/main/java/com/storeroom/utils/RequestHolder.java b/common/src/main/java/com/storeroom/utils/RequestHolder.java
new file mode 100644
index 0000000..27cb1a3
--- /dev/null
+++ b/common/src/main/java/com/storeroom/utils/RequestHolder.java
@@ -0,0 +1,14 @@
+package com.storeroom.utils;
+
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
+
+import javax.servlet.http.HttpServletRequest;
+import java.util.Objects;
+
+public class RequestHolder {
+
+    public static HttpServletRequest getHttpServletRequest() {
+        return ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
+    }
+}