111

src/main/java/com/fdkankan/cloud/acl/annotation/RedisLimit.java

@@ -0,0 +1,37 @@
+package com.fdkankan.cloud.acl.annotation;
+
+import com.fdkankan.cloud.acl.constant.LimitType;
+
+import java.lang.annotation.*;
+import java.util.concurrent.TimeUnit;
+
+@Target({ElementType.METHOD,ElementType.TYPE})
+@Retention(RetentionPolicy.RUNTIME)
+@Inherited
+@Documented
+public @interface RedisLimit {
+
+    // 资源名称
+    String name() default "";
+
+    // 资源key
+    String key() default "";
+
+    /**
+     * 滑动窗口时间单位，默认 分钟
+     */
+    TimeUnit timeUnit() default TimeUnit.SECONDS;
+
+    // 时间
+    int period();
+
+    // 最多访问次数
+    int limitCount();
+
+    // 类型
+    LimitType limitType() default LimitType.APP_KEY;
+
+    // 提示信息
+    String msg() default "亲，限流了，稍后再试！";
+
+}

src/main/java/com/fdkankan/cloud/acl/aop/RedisLimitAspect.java

@@ -0,0 +1,111 @@
+package com.fdkankan.cloud.acl.aop;
+
+import cn.hutool.core.util.StrUtil;
+import com.fdkankan.cloud.acl.annotation.RedisLimit;
+import com.fdkankan.cloud.acl.constant.LimitType;
+import com.fdkankan.common.constant.ErrorCode;
+import com.fdkankan.common.exception.BusinessException;
+import com.fdkankan.web.util.WebUtil;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
+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.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.io.ClassPathResource;
+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.scripting.support.ResourceScriptSource;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
+
+import javax.annotation.Resource;
+import javax.servlet.http.HttpServletRequest;
+import java.lang.reflect.Method;
+import java.time.Instant;
+import java.util.Collections;
+import java.util.concurrent.TimeUnit;
+
+@Slf4j
+@Aspect
+@Configuration
+public class RedisLimitAspect {
+
+    private final static String REDIS_LIMIT_KEY_PREFIX = "openapi:limit:key:resource:%s:%s";
+
+    @Autowired
+    private RedisTemplate redisTemplate;
+    @Resource
+    private RedisScript<Long> limitRedisScript;
+
+    @Before("@annotation(com.fdkankan.openApi.aop.RedisLimit)")
+    public void before(JoinPoint joinPoint){
+        MethodSignature methodSignature = (MethodSignature)joinPoint.getSignature();
+        Method method = methodSignature.getMethod();
+        RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);
+        LimitType limitType = redisLimit.limitType();//限流方式
+        String resource = redisLimit.name();//资源名
+
+        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
+        HttpServletRequest request = requestAttributes.getRequest();
+        int period = redisLimit.period();
+        int limitCount = redisLimit.limitCount();
+        TimeUnit timeUnit = redisLimit.timeUnit();
+        String key = null;
+        switch (limitType){
+            case APP_KEY:
+                key = request.getHeader("Authorization");
+                if (StrUtil.isEmpty(key)) {
+                    throw new BusinessException(ErrorCode.AUTH_FAIL);
+                }
+                break;
+            case IP:
+                key = WebUtil.getIpAddress(request);
+                break;
+            case CUSTOMER:
+                key = redisLimit.key();
+                break;
+            default:
+                key = StringUtils.upperCase(method.getName());
+        }
+        boolean limited = this.shouldLimited(key, resource, limitCount, period, timeUnit);
+        if(limited){
+            throw new BusinessException(ErrorCode.SYSTEM_BUSY.code(), redisLimit.msg());
+        }
+    }
+
+    private boolean shouldLimited(String key, String resource, long max, long timeout, TimeUnit timeUnit) {
+        key = String.format(REDIS_LIMIT_KEY_PREFIX, key, resource);
+        // 统一使用单位毫秒
+        long ttl = timeUnit.toMillis(timeout);
+        // 当前时间毫秒数
+        long now = Instant.now().toEpochMilli();
+        long expired = now - ttl;
+        /**
+         * 注意这里必须转为 String,否则会报错 java.lang.Long cannot be cast to java.lang.String
+         * stringRedisTemplate.execute(RedisScript<T> script, List<K> keys, Object... args)
+         */
+        Long executeTimes = (Long)redisTemplate.execute(limitRedisScript, Collections.singletonList(key), String.valueOf(now), String.valueOf(ttl), String.valueOf(expired), String.valueOf(max));
+        if (executeTimes != null) {
+            if (executeTimes == 0) {
+                return true;
+            } else {
+                return false;
+            }
+        }
+        return false;
+    }
+
+    @Bean("limitRedisScript")
+    public RedisScript<Long> limitRedisScript() {
+        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
+        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redisScript/limit.lua")));
+        redisScript.setResultType(Long.class);
+        return redisScript;
+    }
+
+}

src/main/java/com/fdkankan/cloud/acl/constant/LimitType.java

@@ -0,0 +1,27 @@
+package com.fdkankan.cloud.acl.constant;
+
+import lombok.AllArgsConstructor;
+
+@AllArgsConstructor
+public enum LimitType {
+
+    /**
+     * 应用维度
+     */
+    APP_KEY("appKey"),
+    /**
+     * 自定义维度
+     */
+    CUSTOMER("customer"),
+
+    /**
+     * ip维度
+     */
+    IP("ip"),
+    ;
+
+    private String code;
+
+
+
+}

src/main/java/com/fdkankan/cloud/acl/controller/StsController.java

@@ -1,8 +1,10 @@
 package com.fdkankan.cloud.acl.controller;
 
 import com.fdkankan.cloud.acl.annotation.CheckSignature;
+import com.fdkankan.cloud.acl.annotation.RedisLimit;
 import com.fdkankan.cloud.acl.annotation.SaveLog;
 import com.fdkankan.cloud.acl.bean.StsBean;
+import com.fdkankan.cloud.acl.constant.LimitType;
 import com.fdkankan.cloud.acl.service.IStsService;
 import com.fdkankan.web.response.ResultData;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -25,6 +27,7 @@ public class StsController {
      */
     @CheckSignature
     @SaveLog
+    @RedisLimit(name = "sts/assumeRole", limitType = LimitType.IP, limitCount = 1, period = 1)
     @GetMapping("/assumeRole")
     public ResultData assumeRole(String appCode, Long timestamp){
         return ResultData.ok(stsService.getAssumeRole(appCode));

src/main/java/com/fdkankan/cloud/acl/service/impl/StsServiceImpl.java

@@ -57,63 +57,6 @@ public class StsServiceImpl implements IStsService {
     @Autowired
     private RedisUtil redisUtil;
 
-    public static void main(String[] args) throws Exception {
-
-        // 工程代码泄露可能会导致 AccessKey 泄露，并威胁账号下所有资源的安全性。以下代码示例仅供参考，建议使用更安全的 STS 方式，更多鉴权访问方式请参见：https://help.aliyun.com/document_detail/378657.html。
-        com.aliyun.teaopenapi.models.Config config =
-                new com.aliyun.teaopenapi.models.Config()
-                        // 必填，请确保代码运行环境设置了环境变量 ALIBABA_CLOUD_ACCESS_KEY_ID。
-                        .setAccessKeyId("LTAI5tJnqRzcURptuSVnG57Z")
-                        // 必填，请确保代码运行环境设置了环境变量 ALIBABA_CLOUD_ACCESS_KEY_SECRET。
-                        .setAccessKeySecret("1RomdGaoay2ABIUVSHOMG5HclPWw8D")
-                        .setEndpoint("sts.cn-shenzhen.aliyuncs.com");
-        // Endpoint 请参考 https://api.aliyun.com/product/Sts
-        Client stsClient = new Client(config);
-
-        com.aliyun.sts20150401.models.AssumeRoleRequest assumeRoleRequest = new com.aliyun.sts20150401.models.AssumeRoleRequest()
-                .setDurationSeconds(3600L)
-                .setRoleArn("acs:ram::1899912233141089:role/server-sts")
-                .setRoleSessionName("app_user");
-        //设置权限，如果为空，则默认跟随rolearn的权限
-//        if(StrUtil.isNotEmpty(appConfig.getPolicy())){
-//            assumeRoleRequest.setPolicy(appConfig.getPolicy());
-//        }
-
-        for (int i = 0; i < 200; i++){
-            new Thread(new Runnable() {
-                @Override
-                public void run() {
-                    RuntimeOptions runtime = new RuntimeOptions();
-                    AssumeRoleResponse resp = null;
-
-                    while (Objects.isNull(resp)){
-                        resp = get(assumeRoleRequest, runtime);
-                        if(Objects.isNull(resp)){
-                            try {
-                                Thread.sleep(500L);
-                            } catch (InterruptedException e) {
-                                throw new RuntimeException(e);
-                            }
-                        }
-                    }
-                    System.out.println(JSON.toJSONString(resp));
-                }
-                public AssumeRoleResponse get(AssumeRoleRequest assumeRoleRequest, RuntimeOptions runtime){
-                    AssumeRoleResponse assumeRoleResponse = null;
-                    try {
-                        assumeRoleResponse = stsClient.assumeRoleWithOptions(assumeRoleRequest, runtime);
-                    } catch (Exception e) {
-                        System.out.println("请求失败");
-                    }
-                    return assumeRoleResponse;
-                }
-
-            }).start();
-        }
-    }
-
-
-
     @Override
     public StsBean genAssumeRole(String appCode) {
 

src/main/resources/redisScript/limit.lua

@@ -0,0 +1,25 @@
+--local key = KEYS[1]
+local now = tonumber(ARGV[1]) -- 当前时时间戳
+local ttl = tonumber(ARGV[2]) -- key的过期时间
+local expired = tonumber(ARGV[3]) -- 元素过期分数上限（用于移除过期时间窗口元素）
+local limitCount = tonumber(ARGV[4]) --限制访问次数
+
+-- 清除过期的数据
+-- 移除指定分数区间内的所有元素，expired 即已经过期的 score
+-- 根据当前时间毫秒数 - 超时毫秒数，得到过期时间 expired
+redis.call('zremrangebyscore', KEYS[1], 0, expired)
+
+-- 获取 zset 中的当前元素个数
+local current = tonumber(redis.call('zcard', KEYS[1]))
+local next = current + 1
+
+if next > limitCount then
+  -- 达到限流大小 返回 0
+  return 0;
+else
+  -- 往 zset 中添加一个[value,score]均为当前时间戳的元素，[value,score]
+  redis.call("zadd", KEYS[1], now, now)
+  -- 每次访问均重新设置 zset 的过期时间，单位毫秒
+  redis.call("pexpire", KEYS[1], ttl)
+  return next
+end