apply rate limiter per ip

Author: rfresh2

src/main/java/dev/zenith/web/api/WebServer.java

@@ -1,6 +1,8 @@
 package dev.zenith.web.api;
 
 import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import com.zenith.Globals;
 import com.zenith.command.api.CommandContext;
 import com.zenith.discord.EmbedSerializer;
@@ -9,7 +11,6 @@
 import dev.zenith.web.api.model.CommandRequest;
 import dev.zenith.web.api.model.CommandResponse;
 import io.javalin.Javalin;
-import io.javalin.http.util.NaiveRateLimit;
 import io.javalin.json.JavalinJackson;
 import org.eclipse.jetty.util.thread.ExecutorThreadPool;
 
@@ -21,6 +22,9 @@
 
 public class WebServer {
     private Javalin server;
+    private final Cache<String, Integer> rateLimitCache = CacheBuilder.newBuilder()
+        .expireAfterWrite(1, TimeUnit.MINUTES)
+        .build();
 
     public synchronized void start() {
         if (server != null) {
@@ -56,6 +60,20 @@ private Javalin createServer() {
                 config.jsonMapper(new JavalinJackson(objectMapper, false));
             })
             .beforeMatched(ctx -> {
+                if (PLUGIN_CONFIG.rateLimiter) {
+                    String ip = ctx.ip();
+                    synchronized (this) {
+                        int reqCount = rateLimitCache.get(ip, () -> 0);
+                        rateLimitCache.put(ip, reqCount + 1);
+                        if (reqCount >= PLUGIN_CONFIG.rateLimitRequestsPerMinute) {
+                            ctx.status(429);
+                            ctx.json(new AuthErrorResponse("Rate limit exceeded"));
+                            ctx.skipRemainingHandlers();
+                            LOG.warn("Rate limit exceeded for IP: {}", ip);
+                            return;
+                        }
+                    }
+                }
                 var authHeaderValue = ctx.header("Authorization");
                 if (authHeaderValue != null) {
                     var expectedHeaderValue = PLUGIN_CONFIG.authToken;
@@ -73,9 +91,6 @@ private Javalin createServer() {
                 LOG.warn("Denied request from {}: {}", ctx.ip(), reason);
             })
             .post("/command", ctx -> {
-                if (PLUGIN_CONFIG.rateLimiter) {
-                    NaiveRateLimit.requestPerTimeUnit(ctx, PLUGIN_CONFIG.rateLimitRequestsPerMinute, TimeUnit.MINUTES);
-                }
                 var req = ctx.bodyAsClass(CommandRequest.class);
                 var command = req.command();
                 var context = CommandContext.create(command, WebAPICommandSource.INSTANCE);

src/main/java/dev/zenith/web/command/WebAPICommand.java

@@ -29,7 +29,9 @@ public CommandUsage commandUsage() {
                 "on/off",
                 "port <port>",
                 "auth <token>",
-                "commandsAccountOwnerPerms on/off"
+                "commandsAccountOwnerPerms on/off",
+                "rateLimiter on/off",
+                "rateLimiter requestsPerMinute <requestCount>"
             )
             .build();
     }
@@ -64,7 +66,18 @@ public LiteralArgumentBuilder<CommandContext> register() {
                 PLUGIN_CONFIG.commandsAccountOwnerPerms = getToggle(c, "toggle");
                 c.getSource().getEmbed()
                     .title("Commands Account Owner Perms " + toggleStrCaps(PLUGIN_CONFIG.commandsAccountOwnerPerms));
-            })));
+            })))
+            .then(literal("rateLimiter")
+                .then(argument("toggle", toggle()).executes(c -> {
+                    PLUGIN_CONFIG.rateLimiter = getToggle(c, "toggle");
+                    c.getSource().getEmbed()
+                        .title("Rate Limiter " + toggleStrCaps(PLUGIN_CONFIG.rateLimiter));
+                }))
+                .then(literal("requestsPerMinute").then(argument("requests", integer(1)).executes(c -> {
+                    PLUGIN_CONFIG.rateLimitRequestsPerMinute = getInteger(c, "requests");
+                    c.getSource().getEmbed()
+                        .title("Rate Limiter Requests Per Minute Set");
+                }))));
     }
 
     @Override
@@ -74,6 +87,8 @@ public void defaultEmbed(Embed embed) {
             .addField("Port", PLUGIN_CONFIG.port)
             .addField("Auth Token", PLUGIN_CONFIG.authToken)
             .addField("Commands Account Owner Perms", PLUGIN_CONFIG.commandsAccountOwnerPerms)
+            .addField("Rate Limiter", toggleStr(PLUGIN_CONFIG.rateLimiter))
+            .addField("Rate Limit Requests Per Minute", PLUGIN_CONFIG.rateLimitRequestsPerMinute)
             .primaryColor();
     }
 }