接口限制

src/main/java/com/sqx/common/utils/ApiAccessLimitUtil.java

@@ -0,0 +1,147 @@
+package com.sqx.common.utils;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+
+import java.time.DayOfWeek;
+import java.time.LocalDate;
+import java.time.ZoneOffset;
+import java.time.temporal.TemporalAdjusters;
+import java.util.Objects;
+
+
+/**
+ * 接口访问次数 月为单位
+ */
+@Component
+public class ApiAccessLimitUtil {
+
+    private static final String ACCESS_COUNT_KEY_PREFIX = "sys:limit:";
+
+    private static RedisUtils redisUtils;
+
+    private static final int DEFAULT_ACCESS_COUNT = 5;
+    private static final String DATE_TIME_FORMAT = "month";
+
+
+    @Autowired
+    public void setRedisUtils(RedisUtils redisUtils) {
+        ApiAccessLimitUtil.redisUtils = redisUtils;
+    }
+
+    /**
+     * 默认 当月5次
+     * @param id         唯一值
+     * @param key        接口名称  sys:limit:接口名称
+     * @return
+     */
+    public static boolean isAccessAllowed(String id, String key) {
+        String redisKey = generateRedisKey(key, id);
+        Object countObj = redisUtils.get(redisKey);
+        if (countObj == null) {
+            // 根据不同时间周期设置过期时间并初始化访问次数为1
+            long expireAt = calculateExpireAt(DATE_TIME_FORMAT);
+            redisUtils.set(redisKey, 1, expireAt);
+            return true;
+        }
+        if ((int) countObj < DEFAULT_ACCESS_COUNT) {
+            // 访问次数未达上限，次数加1
+            redisUtils.incr(redisKey);
+            return true;
+        }
+        return false;
+    }
+
+
+    /**
+     * 默认月 month/月/自然月
+     * @param id         唯一值
+     * @param key        接口名称  sys:limit:接口名称
+     * @param count      次数限制
+     * @return
+     */
+    public static boolean isAccessAllowed(String id, String key, Integer count) {
+        String redisKey = generateRedisKey(key, id);
+        Object countObj = redisUtils.get(redisKey);
+        if (countObj == null) {
+            // 根据不同时间周期设置过期时间并初始化访问次数为1
+            long expireAt = calculateExpireAt(DATE_TIME_FORMAT);
+            redisUtils.set(redisKey, 1, expireAt);
+            return true;
+        }
+        if ((int) countObj < count) {
+            // 访问次数未达上限，次数加1
+            redisUtils.incr(redisKey);
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * 默认 5次
+     * @param id         唯一值
+     * @param key        接口名称  sys:limit:接口名称
+     * @param timeFormat  day/天/自然天 week/周/本周日 month/月/自然月 year/年/自然年
+     * @return
+     */
+    public static boolean isAccessAllowed(String id, String key, String timeFormat) {
+        String redisKey = generateRedisKey(key, id);
+        Object countObj = redisUtils.get(redisKey);
+        if (countObj == null) {
+            // 根据不同时间周期设置过期时间并初始化访问次数为1
+            long expireAt = calculateExpireAt(timeFormat);
+            redisUtils.set(redisKey, 1, expireAt);
+            return true;
+        }
+        if ((int) countObj < DEFAULT_ACCESS_COUNT) {
+            // 访问次数未达上限，次数加1
+            redisUtils.incr(redisKey);
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * @param id         唯一值
+     * @param key        接口名称  sys:limit:接口名称
+     * @param count      次数限制
+     * @param timeFormat  day/天/自然天 week/周/本周日 month/月/自然月 year/年/自然年
+     * @return
+     */
+    public static boolean isAccessAllowed(String id, String key, Integer count, String timeFormat) {
+        String redisKey = generateRedisKey(key, id);
+        Object countObj = redisUtils.get(redisKey);
+        if (countObj == null) {
+            // 根据不同时间周期设置过期时间并初始化访问次数为1
+            long expireAt = calculateExpireAt(timeFormat);
+            redisUtils.set(redisKey, 1, expireAt);
+            return true;
+        }
+        if ((int) countObj < count) {
+            // 访问次数未达上限，次数加1
+            redisUtils.incr(redisKey);
+            return true;
+        }
+        return false;
+    }
+
+
+    private static String generateRedisKey(String key, String id) {
+        return ACCESS_COUNT_KEY_PREFIX + key + ":" + id;
+    }
+
+    private static long calculateExpireAt(String timePeriod) {
+        LocalDate now = LocalDate.now();
+        LocalDate expireDate = null;
+        if ("day".equals(timePeriod)) {
+            expireDate = now.plusDays(1).atStartOfDay().toLocalDate();
+        } else if ("week".equals(timePeriod)) {
+            expireDate = now.plusWeeks(0).with(TemporalAdjusters.nextOrSame(java.time.DayOfWeek.SUNDAY));
+        } else if ("month".equals(timePeriod)) {
+            expireDate = now.plusMonths(1).withDayOfMonth(1).minusDays(1);
+        } else if ("year".equals(timePeriod)) {
+            expireDate = now.plusYears(1).withDayOfYear(1).minusDays(1);
+        }
+        return Objects.requireNonNull(expireDate).atTime(23, 59, 59).toEpochSecond(ZoneOffset.UTC);
+    }
+}

src/main/java/com/sqx/common/utils/RedisUtils.java

@@ -128,6 +128,21 @@ public class RedisUtils {
         return get(key, NOT_EXPIRE);
     }
 
+    /**
+     * 对指定的键执行自增操作，返回自增后的值
+     *
+     * @param key 要自增的键
+     * @return 自增后的值，如果键不存在则初始化为1后返回1，如果出现异常返回null
+     */
+    public Long incr(String key) {
+        try {
+            return redisTemplate.opsForValue().increment(key);
+        } catch (Exception e) {
+            e.printStackTrace();
+            return null;
+        }
+    }
+
     public void delete(String key) {
         redisTemplate.delete(key);
     }