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

@@ -0,0 +1,33 @@
+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

@@ -0,0 +1,85 @@
+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

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

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

@@ -46,8 +46,8 @@ public class WebSecurityConfig extends WebSecurityConfigurerAdapter {
     @Autowired
     public void configureGlobal(AuthenticationManagerBuilder auth) throws Exception {
         auth
-            .userDetailsService(jwtUserDetailsService)
-            .passwordEncoder(passwordEncoderBean());
+                .userDetailsService(jwtUserDetailsService)
+                .passwordEncoder(passwordEncoderBean());
     }
 
     @Bean
@@ -65,56 +65,55 @@ public AuthenticationManager authenticationManagerBean() throws Exception {
     protected void configure(HttpSecurity httpSecurity) throws Exception {
         httpSecurity
 
-            // 禁用 CSRF
-            .csrf().disable()
+                // 禁用 CSRF
+                .csrf().disable()
 
-            // 授权异常
-            .exceptionHandling().authenticationEntryPoint(unauthorizedHandler).and()
+                // 授权异常
+                .exceptionHandling().authenticationEntryPoint(unauthorizedHandler).and()
 
-            // 不创建会话
-            .sessionManagement().sessionCreationPolicy(SessionCreationPolicy.STATELESS).and()
+                // 不创建会话
+                .sessionManagement().sessionCreationPolicy(SessionCreationPolicy.STATELESS).and()
 
-            .authorizeRequests()
+                .authorizeRequests()
 
-            .antMatchers("/auth/**").permitAll()
-            .antMatchers("/websocket/**").permitAll()
+                .antMatchers("/auth/**").permitAll()
+                .antMatchers("/websocket/**").permitAll()
+                .antMatchers("/druid/**").anonymous()
+                // swagger start
+                .antMatchers("/swagger-ui.html").anonymous()
+                .antMatchers("/swagger-resources/**").anonymous()
+                .antMatchers("/webjars/**").anonymous()
+                .antMatchers("/*/api-docs").anonymous()
+                // swagger end
+                .antMatchers("/test/**").anonymous()
+                .antMatchers(HttpMethod.OPTIONS, "/**").anonymous()
 
-            .antMatchers("/druid/**").anonymous()
-            // swagger start
-            .antMatchers("/swagger-ui.html").anonymous()
-            .antMatchers("/swagger-resources/**").anonymous()
-            .antMatchers("/webjars/**").anonymous()
-            .antMatchers("/*/api-docs").anonymous()
-            // swagger end
-
-            .antMatchers(HttpMethod.OPTIONS, "/**").anonymous()
-
-            // 所有请求都需要认证
-            .anyRequest().authenticated();
+                // 所有请求都需要认证
+                .anyRequest().authenticated();
 
         httpSecurity
-            .addFilterBefore(authenticationTokenFilter, UsernamePasswordAuthenticationFilter.class);
+                .addFilterBefore(authenticationTokenFilter, UsernamePasswordAuthenticationFilter.class);
     }
 
     @Override
     public void configure(WebSecurity web) throws Exception {
         // AuthenticationTokenFilter will ignore the below paths
         web
-            .ignoring()
-            .antMatchers(
-                HttpMethod.POST,
-                authenticationPath
-            )
-
-            // allow anonymous resource requests
-            .and()
-            .ignoring()
-            .antMatchers(
-                HttpMethod.GET,
-                "/*.html",
-                "/**/*.html",
-                "/**/*.css",
-                "/**/*.js"
-            );
+                .ignoring()
+                .antMatchers(
+                        HttpMethod.POST,
+                        authenticationPath
+                )
+
+                // allow anonymous resource requests
+                .and()
+                .ignoring()
+                .antMatchers(
+                        HttpMethod.GET,
+                        "/*.html",
+                        "/**/*.html",
+                        "/**/*.css",
+                        "/**/*.js"
+                );
     }
 }

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

@@ -0,0 +1,24 @@
+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();
+    }
+}