From f965c5b9c91e058e9f9383a03654b406d1529620 Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Tue, 17 Nov 2020 05:27:05 +0100
Subject: [PATCH] Added write and read limit

---
 .../server/network/netty/NettyServer.java     | 69 +++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/src/main/java/net/minestom/server/network/netty/NettyServer.java b/src/main/java/net/minestom/server/network/netty/NettyServer.java
index d4314ef80..920ca098f 100644
--- a/src/main/java/net/minestom/server/network/netty/NettyServer.java
+++ b/src/main/java/net/minestom/server/network/netty/NettyServer.java
@@ -12,6 +12,7 @@ import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.ServerSocketChannel;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.channel.socket.nio.NioServerSocketChannel;
+import io.netty.handler.traffic.ChannelTrafficShapingHandler;
 import net.minestom.server.network.PacketProcessor;
 import net.minestom.server.network.netty.channel.ClientChannel;
 import net.minestom.server.network.netty.codec.LegacyPingHandler;
@@ -32,6 +33,10 @@ public class NettyServer {
     private String address;
     private int port;
 
+    // Options
+    private long writeLimit;
+    private long readLimit;
+
     public NettyServer(@NotNull PacketProcessor packetProcessor) {
         Class<? extends ServerChannel> channel;
 
@@ -63,6 +68,11 @@ public class NettyServer {
 
                 ChannelPipeline pipeline = ch.pipeline();
 
+                ChannelTrafficShapingHandler channelTrafficShapingHandler =
+                        new ChannelTrafficShapingHandler(writeLimit, readLimit, 200);
+
+                pipeline.addLast("traffic-limiter", channelTrafficShapingHandler);
+
                 // First check should verify if the packet is a legacy ping (from 1.6 version and earlier)
                 pipeline.addLast("legacy-ping", new LegacyPingHandler());
 
@@ -97,14 +107,73 @@ public class NettyServer {
         }
     }
 
+    /**
+     * Gets the address of the server.
+     *
+     * @return the server address
+     */
     public String getAddress() {
         return address;
     }
 
+    /**
+     * Gets the port used by the server.
+     *
+     * @return the server port
+     */
     public int getPort() {
         return port;
     }
 
+    /**
+     * Gets the server write limit.
+     * <p>
+     * Used when you want to limit the bandwidth used by a single connection.
+     * Can also prevent the networking threads from being unresponsive.
+     *
+     * @return the write limit in bytes
+     */
+    public long getWriteLimit() {
+        return writeLimit;
+    }
+
+    /**
+     * Changes the server write limit
+     * <p>
+     * WARNING: the change will only apply to new connections, the current ones will not be updated.
+     *
+     * @param writeLimit the new write limit in bytes
+     * @see #getWriteLimit()
+     */
+    public void setWriteLimit(long writeLimit) {
+        this.writeLimit = writeLimit;
+    }
+
+
+    /**
+     * Gets the server read limit.
+     * <p>
+     * Used when you want to limit the bandwidth used by a single connection.
+     * Can also prevent the networking threads from being unresponsive.
+     *
+     * @return the read limit in bytes
+     */
+    public long getReadLimit() {
+        return readLimit;
+    }
+
+    /**
+     * Changes the server read limit
+     * <p>
+     * WARNING: the change will only apply to new connections, the current ones will not be updated.
+     *
+     * @param readLimit the new read limit in bytes
+     * @see #getWriteLimit()
+     */
+    public void setReadLimit(long readLimit) {
+        this.readLimit = readLimit;
+    }
+
     public void stop() {
         serverChannel.close();