add limit tag for controller

增加限流注解， 控制接口访问

pom.xml

@@ -75,7 +75,10 @@
             <artifactId>commons-pool2</artifactId>
             <version>2.5.0</version>
         </dependency>
-
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+        </dependency>
         <!--jwt-->
         <dependency>
             <groupId>io.jsonwebtoken</groupId>

src/main/java/me/zhengjie/common/aop/limit/Limit.java

+package me.zhengjie.common.aop.limit;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * @author jacky
+ */
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface Limit {
+
+    // 资源名称，用于描述接口功能
+    String name() default "";
+
+    // 资源 key
+    String key() default "";
+
+    // key prefix
+    String prefix() default "";
+
+    // 时间的，单位秒
+    int period();
+
+    // 限制访问次数
+    int count();
+
+    // 限制类型
+    LimitType limitType() default LimitType.CUSTOMER;
+
+}

src/main/java/me/zhengjie/common/aop/limit/LimitAspect.java

+package me.zhengjie.common.aop.limit;
+
+import com.google.common.collect.ImmutableList;
+import me.zhengjie.common.exception.BadRequestException;
+import me.zhengjie.common.utils.IpUtil;
+import me.zhengjie.common.utils.RequestHolder;
+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.annotation.Pointcut;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+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 {
+    @Autowired
+    private RedisTemplate redisTemplate;
+    private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);
+
+
+    @Pointcut("@annotation(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 name = limit.name();
+        String key = limit.key();
+        if (StringUtils.isEmpty(key)) {
+            switch (limitType) {
+                case IP:
+                    key = IpUtil.getIP(request);
+                    break;
+                default:
+                    key = signatureMethod.getName();
+            }
+        }
+
+        ImmutableList keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));
+
+        String luaScript = buildLuaScript();
+        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
+        Number count = (Number) 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 BadRequestException("访问次数受限制");
+        }
+    }
+
+    /**
+     * 限流脚本
+     */
+    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;";
+    }
+}

src/main/java/me/zhengjie/common/aop/limit/LimitType.java

+package me.zhengjie.common.aop.limit;
+
+public enum LimitType {
+    CUSTOMER,
+//     by ip addr
+    IP;
+}

src/main/java/me/zhengjie/core/config/WebSecurityConfig.java

@@ -78,7 +78,6 @@ public class WebSecurityConfig extends WebSecurityConfigurerAdapter {
 
                 .antMatchers("/auth/**").permitAll()
                 .antMatchers("/websocket/**").permitAll()
-
                 .antMatchers("/druid/**").anonymous()
                 // swagger start
                 .antMatchers("/swagger-ui.html").anonymous()
@@ -86,7 +85,7 @@ public class WebSecurityConfig extends WebSecurityConfigurerAdapter {
                 .antMatchers("/webjars/**").anonymous()
                 .antMatchers("/*/api-docs").anonymous()
                 // swagger end
-
+                .antMatchers("/test/**").anonymous()
                 .antMatchers(HttpMethod.OPTIONS, "/**").anonymous()
 
                 // 所有请求都需要认证

src/main/java/me/zhengjie/tools/rest/TestController.java

+package me.zhengjie.tools.rest;
+
+import me.zhengjie.common.aop.limit.Limit;
+import org.springframework.security.access.prepost.PreAuthorize;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+@RestController
+@RequestMapping("test")
+public class TestController {
+    private static final AtomicInteger ATOMIC_INTEGER = new AtomicInteger();
+
+    /**
+     * 测试限流注解，下面配置说明该接口 60秒内最多只能访问 10次，保存到redis的键名为 limit_test，
+     */
+    @Limit(key = "test", period = 60, count = 10, name = "testLimit", prefix = "limit")
+    @GetMapping("limit")
+    public int testLimit() {
+        return ATOMIC_INTEGER.incrementAndGet();
+    }
+}