Add option to exclude addresses from being blocked

CloudburstMC · Mar 30, 2024 · 94f729d · 94f729d
1 parent 2f77498
commit 94f729d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakServerChannel.java b/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakServerChannel.java
@@ -99,8 +99,7 @@ public void onCloseTriggered(ChannelPromise promise) {
     public boolean tryBlockAddress(InetAddress address, long time, TimeUnit unit) {
         RakServerRateLimiter rateLimiter = this.pipeline().get(RakServerRateLimiter.class);
         if (rateLimiter != null) {
-            rateLimiter.blockAddress(address, time, unit);
-            return true;
+            return rateLimiter.blockAddress(address, time, unit);
         }
         return false;
     }

diff --git a/...rc/main/java/org/cloudburstmc/netty/handler/codec/raknet/server/RakServerRateLimiter.java b/...rc/main/java/org/cloudburstmc/netty/handler/codec/raknet/server/RakServerRateLimiter.java
@@ -25,8 +25,7 @@
 import org.cloudburstmc.netty.channel.raknet.RakServerChannel;
 
 import java.net.InetAddress;
-import java.util.Iterator;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -41,6 +40,8 @@ public class RakServerRateLimiter extends SimpleChannelInboundHandler<DatagramPa
     private final ConcurrentHashMap<InetAddress, AtomicInteger> rateLimitMap = new ConcurrentHashMap<>();
     private final Map<InetAddress, Long> blockedConnections = new ConcurrentHashMap<>();
 
+    private final Collection<InetAddress> exceptions = Collections.synchronizedCollection(new HashSet<>());
+
     private final AtomicLong globalCounter = new AtomicLong(0);
 
     private ScheduledFuture<?> tickFuture;
@@ -81,9 +82,14 @@ private void onBlockedTick() {
         }
     }
 
-    public void blockAddress(InetAddress address, long time, TimeUnit unit) {
+    public boolean blockAddress(InetAddress address, long time, TimeUnit unit) {
+        if (this.exceptions.contains(address)) {
+            return false;
+        }
+
         long millis = unit.toMillis(time);
         this.blockedConnections.put(address, System.currentTimeMillis() + millis);
+        return true;
     }
 
     public void unblockAddress(InetAddress address) {
@@ -96,6 +102,14 @@ public boolean isAddressBlocked(InetAddress address) {
         return this.blockedConnections.containsKey(address);
     }
 
+    public void addException(InetAddress address) {
+        this.exceptions.add(address);
+    }
+
+    public void removeException(InetAddress address) {
+        this.exceptions.remove(address);
+    }
+
     @Override
     protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket datagram) throws Exception {
         if (this.globalCounter.incrementAndGet() > this.channel.config().getGlobalPacketLimit()) {
@@ -111,8 +125,8 @@ protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket datagram)
         }
 
         AtomicInteger counter = this.rateLimitMap.computeIfAbsent(address, a -> new AtomicInteger());
-        if (counter.incrementAndGet() > this.channel.config().getPacketLimit()) {
-            this.blockAddress(address, 10, TimeUnit.SECONDS);
+        if (counter.incrementAndGet() > this.channel.config().getPacketLimit() &&
+                this.blockAddress(address, 10, TimeUnit.SECONDS)) {
             log.warn("[{}] Blocked because packet limit was reached");
         } else {
             ctx.fireChannelRead(datagram.retain());