From a0739f8deec7fccd6d91c79498461d885ee83d0e Mon Sep 17 00:00:00 2001 From: jacky Date: Fri, 28 Dec 2018 23:32:46 +0800 Subject: [PATCH] =?UTF-8?q?add=20=20limit=20tag=20for=20controller=20?= =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E9=99=90=E6=B5=81=E6=B3=A8=E8=A7=A3=EF=BC=8C?= =?UTF-8?q?=20=E6=8E=A7=E5=88=B6=E6=8E=A5=E5=8F=A3=E8=AE=BF=E9=97=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pom.xml | 5 +- .../me/zhengjie/common/aop/limit/Limit.java | 33 +++++++ .../common/aop/limit/LimitAspect.java | 85 +++++++++++++++++++ .../zhengjie/common/aop/limit/LimitType.java | 7 ++ .../core/config/WebSecurityConfig.java | 79 +++++++++-------- .../zhengjie/tools/rest/TestController.java | 24 ++++++ 6 files changed, 192 insertions(+), 41 deletions(-) create mode 100644 src/main/java/me/zhengjie/common/aop/limit/Limit.java create mode 100644 src/main/java/me/zhengjie/common/aop/limit/LimitAspect.java create mode 100644 src/main/java/me/zhengjie/common/aop/limit/LimitType.java create mode 100644 src/main/java/me/zhengjie/tools/rest/TestController.java diff --git a/pom.xml b/pom.xml index 072cca42e..93c72e072 100644 --- a/pom.xml +++ b/pom.xml @@ -75,7 +75,10 @@ commons-pool2 2.5.0 - + + org.apache.commons + commons-lang3 + io.jsonwebtoken diff --git a/src/main/java/me/zhengjie/common/aop/limit/Limit.java b/src/main/java/me/zhengjie/common/aop/limit/Limit.java new file mode 100644 index 000000000..cc2febd57 --- /dev/null +++ b/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; + +} diff --git a/src/main/java/me/zhengjie/common/aop/limit/LimitAspect.java b/src/main/java/me/zhengjie/common/aop/limit/LimitAspect.java new file mode 100644 index 000000000..86747ea1a --- /dev/null +++ b/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 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;"; + } +} diff --git a/src/main/java/me/zhengjie/common/aop/limit/LimitType.java b/src/main/java/me/zhengjie/common/aop/limit/LimitType.java new file mode 100644 index 000000000..7e8697533 --- /dev/null +++ b/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; +} diff --git a/src/main/java/me/zhengjie/core/config/WebSecurityConfig.java b/src/main/java/me/zhengjie/core/config/WebSecurityConfig.java index ef42476c2..04d748ccd 100644 --- a/src/main/java/me/zhengjie/core/config/WebSecurityConfig.java +++ b/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" + ); } } diff --git a/src/main/java/me/zhengjie/tools/rest/TestController.java b/src/main/java/me/zhengjie/tools/rest/TestController.java new file mode 100644 index 000000000..fd988bf10 --- /dev/null +++ b/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(); + } +}