自定义注解配合redis实现接口流量控制

一、自定义注解

seconds为时间，maxCount为最大请求量。表示在该时间内最大的请求数。

/**
 * 接口限流注解
 * @author 陈加炎
 * @date 2022/12/21
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface AccessLimit {
    /**
     * 时间（单位s）
     * @return
     */
    int seconds();

    /**
     * 最大请求量
     * @return
     */
    int maxCount();
}

二、编写拦截器

redisService中封装了一些redis的一些常用操作，大家可以自己编写。

/**
 * 流量控制
 * chenjiayan
 * 2022/12/23
 */
@Slf4j
public class WebSecurityHandler implements HandlerInterceptor {

    @Autowired
    private RedisService redisService = new RedisServiceImpl();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // 如果请求的是方法（接口）
        if(handler instanceof HandlerMethod){
            HandlerMethod hm = (HandlerMethod) handler;
            // 获取方法中的注解
            AccessLimit accessLimit = hm.getMethodAnnotation(AccessLimit.class);
            if(accessLimit!=null){
                // 方法上添加了该注解
                long seconds = accessLimit.seconds(); // 获取时长
                int maxCount = accessLimit.maxCount();//该时长内最大请求次数
                // 生成key 关于key的生成规则可以自己定义 本项目需求是对每个方法都加上限流功能，如果你只是针对ip地址限流，那么key只需要只用ip就好
                String key = IpUtils.getIpAddress(request) + hm.getMethod().getName();
                // 从redis中获取用户访问的次数
                try{
                    Long count = redisService.incrExpire(key, seconds);
                    // 单位时间内超出最大次数
                    if(count>maxCount){
                        // 渲染数据
                        render(response, Result.fail("请求过于频繁，请稍候再试"));
                        log.warn(key+"请求次数超过每"+seconds+"秒"+maxCount+"次");
                        return false;
                    }
                    return true;
                }catch (RedisConnectionFailureException e){
                    log.warn("redis错误："+e.getMessage());
                }
            }
        }
        return true;
    }

    /**
     * 渲染数据（向response响应体中写入数据）
     * @param response
     * @param result
     * @throws IOException
     */
    private void render(HttpServletResponse response, Result<Object> result) throws IOException {
        response.setContentType("application/json;charset=utf-8");
        ServletOutputStream out = response.getOutputStream();
        String strResult = JSON.toJSONString(result);
        out.write(strResult.getBytes(StandardCharsets.UTF_8));
        out.flush();
    }
}

三、配置拦截器

在web mvc配置中添加拦截器

/**
 * web mvc配置
 * chenjiayan
 * 2022/12/22
 */
@Configuration
public class WebMvcConfig implements WebMvcConfigurer { 
    // 添加拦截器
    @Override
    public void addInterceptors(InterceptorRegistry registry){
        // 分页管理
        registry.addInterceptor(new PageableHandlerInterceptor());
        // 流量限制拦截器
        registry.addInterceptor(new WebSecurityHandler());
    }
}

四、使用

使用就非常简单了，只用在方法（接口）上添加@AccessLimit注解就行了。

/**
 * 发送邮箱验证码
 *
 * @param username 用户名
 * @return {@link Result<>}
 */
@AccessLimit(seconds = 60, maxCount = 1)
@ApiOperation(value = "发送邮箱验证码")
@ApiImplicitParam(name = "username", value = "用户名", required = true, dataType = "String")
@GetMapping("/users/code")
public Result<?> sendCode(String username) {
    userAuthService.sendCode(username);
    return Result.ok();
}