fix: add rateLimiter wansenai#286

core/api/src/main/java/com/wansenai/api/RateLimitException.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2023-2033 WanSen AI Team, Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://opensource.wansenai.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package com.wansenai.api;
+
+
+public class RateLimitException extends RuntimeException {
+    private String errorCode;
+    private String errorMessage;
+
+    public RateLimitException(String errorCode, String errorMessage) {
+        super(errorMessage);
+        this.errorCode = errorCode;
+        this.errorMessage = errorMessage;
+    }
+}

core/api/src/main/java/com/wansenai/api/config/RateLimiter.java

@@ -0,0 +1,30 @@
+package com.wansenai.api.config;
+
+import com.wansenai.utils.enums.LimitType;
+
+import java.lang.annotation.*;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface RateLimiter {
+    /**
+     * 限流key
+     */
+    String key() default "rate_limit:";
+
+    /**
+     * 限流时间,单位秒
+     */
+    int time() default 60;
+
+    /**
+     * 限流次数
+     */
+    int count() default 100;
+
+    /**
+     * 限流类型
+     */
+    LimitType limitType() default LimitType.DEFAULT;
+}

core/api/src/main/java/com/wansenai/api/config/RateLimiterAspect.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2023-2033 WanSen AI Team, Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://opensource.wansenai.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package com.wansenai.api.config;
+
+import com.wansenai.api.RateLimitException;
+import com.wansenai.utils.IpUtils;
+import com.wansenai.utils.enums.BaseCodeEnum;
+import com.wansenai.utils.enums.LimitType;
+import com.wansenai.utils.redis.RedisUtil;
+import lombok.extern.slf4j.Slf4j;
+import org.aspectj.lang.JoinPoint;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Before;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.springframework.stereotype.Component;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
+
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.List;
+
+@Component
+@Aspect
+@Slf4j
+public class RateLimiterAspect {
+
+    private final RedisUtil redisUtil;
+
+    public RateLimiterAspect(RedisUtil redisUtil) {
+        this.redisUtil = redisUtil;
+    }
+
+
+    @Before("@annotation(rateLimiter)")
+    public void doBefore(JoinPoint point, RateLimiter rateLimiter) {
+        String key = rateLimiter.key();
+        int time = rateLimiter.time();
+        int count = rateLimiter.count();
+
+        String combineKey = getCombineKey(rateLimiter, point);
+        List<String> keys = Collections.singletonList(combineKey);
+        try {
+            if (redisUtil.get(keys.get(0)) != null) {
+                long number = redisUtil.incr(keys.get(0), 1);
+                if ((int) number > count) {
+                    throw new RateLimitException(BaseCodeEnum.FREQUENT_SYSTEM_ACCESS.getCode(), BaseCodeEnum.FREQUENT_SYSTEM_ACCESS.getMsg());
+                }
+                log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, (int) number, keys.get(0));
+            } else {
+                redisUtil.set(keys.get(0), 1, time);
+            }
+        } catch (Exception e) {
+            throw new RateLimitException(BaseCodeEnum.SYSTEM_BUSY.getCode(), BaseCodeEnum.SYSTEM_BUSY.getMsg());
+        }
+    }
+
+    /**
+     * 获取ip为key
+     * @param rateLimiter
+     * @param point
+     * @return
+     */
+    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
+        StringBuilder stringBuffer = new StringBuilder(rateLimiter.key());
+        if (rateLimiter.limitType() == LimitType.IP) {
+            stringBuffer.append(
+                            IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes())
+                                    .getRequest()))
+                    .append("-");
+        }
+        MethodSignature signature = (MethodSignature) point.getSignature();
+        Method method = signature.getMethod();
+        Class<?> targetClass = method.getDeclaringClass();
+        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
+        return stringBuffer.toString();
+    }
+}