Add queue length check

Author: SupremeMortal

transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakDisconnectReason.java

@@ -26,5 +26,6 @@ public enum RakDisconnectReason {
     NO_FREE_INCOMING_CONNECTIONS,
     INCOMPATIBLE_PROTOCOL_VERSION,
     IP_RECENTLY_CONNECTED,
-    BAD_PACKET
+    BAD_PACKET,
+    QUEUE_TOO_LONG
 }

transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/config/DefaultRakSessionConfig.java

@@ -38,6 +38,7 @@ public class DefaultRakSessionConfig extends DefaultChannelConfig implements Rak
     private volatile long sessionTimeout = SESSION_TIMEOUT_MS;
     private volatile boolean autoFlush = true;
     private volatile int flushInterval = 10;
+    private volatile int maxQueuedBytes = 64 * 1024 * 1024; // 64 MB
 
     public DefaultRakSessionConfig(Channel channel) {
         super(channel);
@@ -78,6 +79,9 @@ public <T> T getOption(ChannelOption<T> option) {
         if (option == RakChannelOption.RAK_FLUSH_INTERVAL) {
             return (T) Integer.valueOf(this.getFlushInterval());
         }
+        if (option == RakChannelOption.RAK_MAX_QUEUED_BYTES) {
+            return (T) Integer.valueOf(this.getMaxQueuedBytes());
+        }
         return this.channel.parent().config().getOption(option);
     }
 
@@ -102,6 +106,8 @@ public <T> boolean setOption(ChannelOption<T> option, T value) {
             this.setAutoFlush((Boolean) value);
         } else if (option == RakChannelOption.RAK_FLUSH_INTERVAL) {
             this.setFlushInterval((Integer) value);
+        } else if (option == RakChannelOption.RAK_MAX_QUEUED_BYTES) {
+            this.setMaxQueuedBytes((Integer) value);
         } else {
             return this.channel.parent().config().setOption(option, value);
         }
@@ -194,4 +200,14 @@ public int getFlushInterval() {
     public void setFlushInterval(int flushInterval) {
         this.flushInterval = flushInterval;
     }
+
+    @Override
+    public void setMaxQueuedBytes(int maxQueuedBytes) {
+        this.maxQueuedBytes = maxQueuedBytes;
+    }
+
+    @Override
+    public int getMaxQueuedBytes() {
+        return maxQueuedBytes;
+    }
 }

transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/config/RakChannelConfig.java

@@ -51,4 +51,8 @@ public interface RakChannelConfig extends ChannelConfig {
     int getFlushInterval();
 
     void setFlushInterval(int intervalMillis);
+
+    void setMaxQueuedBytes(int maxQueuedBytes);
+
+    int getMaxQueuedBytes();
 }

transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/config/RakChannelMetrics.java

@@ -55,4 +55,7 @@ default void nackIn(int count) {
 
     default void stateChange(RakState state) {
     }
+
+    default void queuedPacketBytes(int count) {
+    }
 }

transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/config/RakChannelOption.java

@@ -189,6 +189,12 @@ public class RakChannelOption<T> extends ChannelOption<T> {
     public static final ChannelOption<Integer> RAK_TIME_BETWEEN_SEND_CONNECTION_ATTEMPTS_MS =
             valueOf(RakChannelOption.class, "RAK_TIME_BETWEEN_SEND_CONNECTION_ATTEMPTS_MS");
 
+    /**
+     * Maximum amount of bytes that can be queued in a single RakNet session.
+     */
+    public static final ChannelOption<Integer> RAK_MAX_QUEUED_BYTES =
+            valueOf(RakChannelOption.class, "RAK_MAX_QUEUED_BYTES");
+
     @SuppressWarnings("deprecation")
     protected RakChannelOption() {
         super(null);

transport-raknet/src/main/java/org/cloudburstmc/netty/handler/codec/raknet/common/RakSessionCodec.java

@@ -38,6 +38,8 @@
 import java.util.ArrayDeque;
 import java.util.Queue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import static org.cloudburstmc.netty.channel.raknet.RakConstants.*;
 
@@ -410,6 +412,23 @@ private void tryTick() {
     private void onTick() {
         long curTime = System.currentTimeMillis();
 
+        int maxQueuedBytes = this.channel.config().getOption(RakChannelOption.RAK_MAX_QUEUED_BYTES);
+
+        if (maxQueuedBytes > 0) {
+            int queuedBytes = 0;
+            try {
+                for (EncapsulatedPacket packet : this.outgoingPackets) {
+                    queuedBytes += packet.getBuffer().readableBytes();
+                    if (queuedBytes > maxQueuedBytes) {
+                        this.disconnect(RakDisconnectReason.QUEUE_TOO_LONG);
+                        return;
+                    }
+                }
+            } finally {
+                getMetrics().queuedPacketBytes(queuedBytes);
+            }
+        }
+
         if (this.state == RakState.UNCONNECTED) {
             if (this.isTimedOut(curTime)) {
                 this.close(RakDisconnectReason.TIMED_OUT);

transport-raknet/src/main/java/org/cloudburstmc/netty/util/FastBinaryMinHeap.java

@@ -20,10 +20,12 @@
 import io.netty.util.internal.ObjectPool;
 
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.NoSuchElementException;
 import java.util.Objects;
+import java.util.function.Consumer;
 
-public class FastBinaryMinHeap<E> extends AbstractReferenceCounted {
+public class FastBinaryMinHeap<E> extends AbstractReferenceCounted implements Iterable<E> {
 
     private static final Entry INFIMUM = new Entry(Long.MAX_VALUE);
     private static final Entry SUPREMUM = new Entry(Long.MIN_VALUE);
@@ -205,6 +207,11 @@ public FastBinaryMinHeap<E> touch(Object hint) {
         return this;
     }
 
+    @Override
+    public Iterator<E> iterator() {
+        return new SimpleIterator();
+    }
+
     private static class Entry extends AbstractReferenceCounted {
         private final ObjectPool.Handle<Entry> handle;
         private Object element;
@@ -233,4 +240,18 @@ public Entry touch(Object hint) {
             return this;
         }
     }
+
+    private class SimpleIterator implements Iterator<E> {
+        private int index = 1;
+
+        @Override
+        public boolean hasNext() {
+            return heap.length > index && heap[index] != INFIMUM;
+        }
+
+        @Override
+        public E next() {
+            return (E) heap[index++].element;
+        }
+    }
 }