新增全局限速

core/src/main/java/net/momirealms/craftengine/core/pack/host/impl/SelfHost.java

@@ -1,5 +1,6 @@
 package net.momirealms.craftengine.core.pack.host.impl;
 
+import com.google.common.util.concurrent.RateLimiter;
 import io.github.bucket4j.Bandwidth;
 import net.momirealms.craftengine.core.pack.host.ResourcePackDownloadData;
 import net.momirealms.craftengine.core.pack.host.ResourcePackHost;
@@ -18,6 +19,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 
 public class SelfHost implements ResourcePackHost {
     public static final Factory FACTORY = new Factory();
@@ -60,6 +62,7 @@ public class SelfHost implements ResourcePackHost {
 
     public static class Factory implements ResourcePackHostFactory {
 
+        @SuppressWarnings("UnstableApiUsage")
         @Override
         public ResourcePackHost create(Map<String, Object> arguments) {
             SelfHostHttpServer selfHostHttpServer = SelfHostHttpServer.instance();
@@ -78,19 +81,30 @@ public class SelfHost implements ResourcePackHost {
             boolean oneTimeToken = ResourceConfigUtils.getAsBoolean(arguments.getOrDefault("one-time-token", true), "one-time-token");
             String protocol = arguments.getOrDefault("protocol", "http").toString();
             boolean denyNonMinecraftRequest = ResourceConfigUtils.getAsBoolean(arguments.getOrDefault("deny-non-minecraft-request", true), "deny-non-minecraft-request");
-            Map<String, Object> rateMap = MiscUtils.castToMap(ResourceConfigUtils.get(arguments, "rate-map", "rate-limit"), true);
-            int maxRequests = 5;
-            int resetInterval = 20;
-            if (rateMap != null) {
-                maxRequests = Math.max(ResourceConfigUtils.getAsInt(rateMap.getOrDefault("max-requests", 5), "max-requests"), 1);
-                resetInterval = Math.max(ResourceConfigUtils.getAsInt(rateMap.getOrDefault("reset-interval", 20), "reset-interval"), 1);
+            Map<String, Object> rateLimitPerIp = MiscUtils.castToMap(arguments.get("rate-limitation-per-ip"), true);
+            boolean enabledLimitPerIp = false;
+            Bandwidth limit = null;
+            out:
+            if (rateLimitPerIp != null) {
+                enabledLimitPerIp = ResourceConfigUtils.getAsBoolean(rateLimitPerIp.getOrDefault("enable", false), "enable");
+                if (!enabledLimitPerIp) break out;
+                int maxRequests = Math.max(ResourceConfigUtils.getAsInt(rateLimitPerIp.getOrDefault("max-requests", 5), "max-requests"), 1);
+                int resetInterval = Math.max(ResourceConfigUtils.getAsInt(rateLimitPerIp.getOrDefault("reset-interval", 20), "reset-interval"), 1);
+                limit = Bandwidth.builder()
+                        .capacity(maxRequests)
+                        .refillGreedy(maxRequests, Duration.ofSeconds(resetInterval))
+                        .build();
             }
-            Bandwidth limit = Bandwidth.builder()
-                    .capacity(maxRequests)
-                    .refillGreedy(maxRequests, Duration.ofSeconds(resetInterval))
-                    .initialTokens(maxRequests / 2) // 修正首次可以直接突破限制请求 maxRequests * 2 次
-                    .build();
-            selfHostHttpServer.updateProperties(ip, port, url, denyNonMinecraftRequest, protocol, limit, oneTimeToken);
+            Map<String, Object> tokenBucket = MiscUtils.castToMap(arguments.get("token-bucket"), true);
+            boolean enabledTokenBucket = false;
+            RateLimiter globalLimiter = null;
+            out:
+            if (tokenBucket != null) {
+                enabledTokenBucket = ResourceConfigUtils.getAsBoolean(tokenBucket.getOrDefault("enable", false), "enable");
+                if (!enabledTokenBucket) break out;
+                globalLimiter = RateLimiter.create(ResourceConfigUtils.getAsDouble(tokenBucket.getOrDefault("qps", 1000), "qps"));
+            }
+            selfHostHttpServer.updateProperties(ip, port, url, denyNonMinecraftRequest, protocol, limit, enabledLimitPerIp, enabledTokenBucket, globalLimiter, oneTimeToken);
             return INSTANCE;
         }
     }

core/src/main/java/net/momirealms/craftengine/core/pack/host/impl/SelfHostHttpServer.java

@@ -3,6 +3,7 @@ package net.momirealms.craftengine.core.pack.host.impl;
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.github.benmanes.caffeine.cache.Scheduler;
+import com.google.common.util.concurrent.RateLimiter;
 import io.github.bucket4j.Bandwidth;
 import io.github.bucket4j.Bucket;
 import io.netty.bootstrap.ServerBootstrap;
@@ -30,6 +31,7 @@ import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+@SuppressWarnings("UnstableApiUsage")
 public class SelfHostHttpServer {
     private static SelfHostHttpServer instance;
     private final Cache<String, Boolean> oneTimePackUrls = Caffeine.newBuilder()
@@ -46,11 +48,14 @@ public class SelfHostHttpServer {
     private final AtomicLong totalRequests = new AtomicLong();
     private final AtomicLong blockedRequests = new AtomicLong();
 
-    private Bandwidth limit = Bandwidth.builder()
+    private Bandwidth limitPerIp = Bandwidth.builder()
             .capacity(1)
             .refillGreedy(1, Duration.ofSeconds(1))
             .initialTokens(1)
             .build();
+    private RateLimiter globalLimiter = RateLimiter.create(1);
+    private boolean enabledLimitPerIp = false;
+    private boolean enabledGlobalLimit = false;
     private String ip = "localhost";
     private int port = -1;
     private String protocol = "http";
@@ -78,13 +83,19 @@ public class SelfHostHttpServer {
                                  String url,
                                  boolean denyNonMinecraft,
                                  String protocol,
-                                 Bandwidth limit,
+                                 Bandwidth limitPerIp,
+                                 boolean enabledLimitPerIp,
+                                 boolean enabledGlobalLimit,
+                                 RateLimiter globalLimiter,
                                  boolean token) {
         this.ip = ip;
         this.url = url;
         this.denyNonMinecraft = denyNonMinecraft;
         this.protocol = protocol;
-        this.limit = limit;
+        this.limitPerIp = limitPerIp;
+        this.enabledLimitPerIp = enabledLimitPerIp;
+        this.enabledGlobalLimit = enabledGlobalLimit;
+        this.globalLimiter = globalLimiter;
         this.useToken = token;
 
         if (port <= 0 || port > 65535) {
@@ -140,7 +151,13 @@ public class SelfHostHttpServer {
                 String clientIp = ((InetSocketAddress) ctx.channel().remoteAddress())
                         .getAddress().getHostAddress();
 
-                if (checkRateLimit(clientIp)) {
+                if (enabledGlobalLimit && !globalLimiter.tryAcquire()) {
+                    sendError(ctx, HttpResponseStatus.TOO_MANY_REQUESTS, "Rate limit exceeded");
+                    blockedRequests.incrementAndGet();
+                    return;
+                }
+
+                if (enabledLimitPerIp && !checkIpRateLimit(clientIp)) {
                     sendError(ctx, HttpResponseStatus.TOO_MANY_REQUESTS, "Rate limit exceeded");
                     blockedRequests.incrementAndGet();
                     return;
@@ -217,13 +234,12 @@ public class SelfHostHttpServer {
             ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
         }
 
-        private boolean checkRateLimit(String clientIp) {
-            Bucket rateLimiter = ipRateLimiters.get(clientIp, k -> Bucket.builder().addLimit(limit).build());
-            if (rateLimiter == null) { // 怎么可能null?
-                rateLimiter = Bucket.builder().addLimit(limit).build();
-                ipRateLimiters.put(clientIp, rateLimiter);
-            }
-            return !rateLimiter.tryConsume(1);
+        private boolean checkIpRateLimit(String clientIp) {
+            Bucket rateLimiter = ipRateLimiters.get(clientIp, k ->
+                    Bucket.builder().addLimit(limitPerIp).build()
+            );
+            assert rateLimiter != null;
+            return rateLimiter.tryConsume(1);
         }
 
         private boolean validateToken(String token) {

core/src/main/java/net/momirealms/craftengine/core/util/ResourceConfigUtils.java

@@ -218,6 +218,30 @@ public final class ResourceConfigUtils {
         }
     }
 
+    public static long getAsLong(Object o, String option) {
+        switch (o) {
+            case null -> {
+                return 0;
+            }
+            case Long l -> {
+                return l;
+            }
+            case Number number -> {
+                return number.longValue();
+            }
+            case String s -> {
+                try {
+                    return Long.parseLong(s);
+                } catch (NumberFormatException e) {
+                    throw new LocalizedResourceConfigException("warning.config.type.long", e, s, option);
+                }
+            }
+            default -> {
+                throw new LocalizedResourceConfigException("warning.config.type.long", o.toString(), option);
+            }
+        }
+    }
+
     @SuppressWarnings("unchecked")
     public static Map<String, Object> getAsMap(Object obj, String option) {
         if (obj instanceof Map<?, ?> map) {