From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Wed, 23 Sep 2020 01:46:46 -0400
Subject: [PATCH] Add Velocity natives for encryption and compression

This patch is licensed under the MIT license. See /licenses/MIT.md.

diff --git a/pom.xml b/pom.xml
index 3c60bf75dacd00cac2e35b0e597992ce0ce9448c..84bedea0b7db6929698503f78ede260f6ea3d56e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -156,6 +156,13 @@
             <version>4.8.47</version>
             <scope>test</scope>
         </dependency>
+        <!-- for optimized protocol handling -->
+        <dependency>
+            <groupId>com.velocitypowered</groupId>
+            <artifactId>velocity-native</artifactId>
+            <version>1.1.0-SNAPSHOT</version>
+            <scope>compile</scope>
+        </dependency>
     </dependencies>
 
     <!-- This builds a completely 'ready to start' jar with all dependencies inside -->
diff --git a/src/main/java/net/minecraft/network/NetworkManager.java b/src/main/java/net/minecraft/network/NetworkManager.java
index 2abf64fda6e504a89b2e443af48894023d446cef..69bfe8d0a9cf745557c3e06141dec1ed3bb7d567 100644
--- a/src/main/java/net/minecraft/network/NetworkManager.java
+++ b/src/main/java/net/minecraft/network/NetworkManager.java
@@ -568,10 +568,16 @@ public class NetworkManager extends SimpleChannelInboundHandler<Packet<?>> {
         return this.channel instanceof LocalChannel || this.channel instanceof LocalServerChannel;
     }
 
-    public void a(Cipher cipher, Cipher cipher1) {
+    public void a(javax.crypto.SecretKey secretkey) { // Tuinity
         this.n = true;
-        this.channel.pipeline().addBefore("splitter", "decrypt", new PacketDecrypter(cipher));
-        this.channel.pipeline().addBefore("prepender", "encrypt", new PacketEncrypter(cipher1));
+        // Tuinity start
+        try {
+        this.channel.pipeline().addBefore("splitter", "decrypt", new PacketDecrypter(/*MinecraftEncryption.a(2, secretkey)*/ secretkey));
+        this.channel.pipeline().addBefore("prepender", "encrypt", new PacketEncrypter(/*MinecraftEncryption.a(1, secretkey)*/ secretkey));
+        } catch (java.security.GeneralSecurityException e) {
+            throw new RuntimeException("Couldn't enable encryption", e);
+        }
+        // Tuinity end
     }
 
     public boolean isConnected() {
diff --git a/src/main/java/net/minecraft/network/PacketCompressor.java b/src/main/java/net/minecraft/network/PacketCompressor.java
index 45b9d3d3c84d11e7f27f699506a1036dff9fdc53..43ea54479ba70a5c3f4a0f8e4e5a0d03e93162c1 100644
--- a/src/main/java/net/minecraft/network/PacketCompressor.java
+++ b/src/main/java/net/minecraft/network/PacketCompressor.java
@@ -7,14 +7,18 @@ import java.util.zip.Deflater;
 
 public class PacketCompressor extends MessageToByteEncoder<ByteBuf> {
 
-    private final byte[] a = new byte[8192];
-    private final Deflater b;
+    // Tuinity start - use Velocity natives
+//    private final byte[] a = new byte[8192];
+//    private final Deflater b;
     private int c;
+    private final com.velocitypowered.natives.compression.VelocityCompressor compressor;
 
     public PacketCompressor(int i) {
         this.c = i;
-        this.b = new Deflater();
+//        this.b = new Deflater();
+        this.compressor = com.velocitypowered.natives.util.Natives.compress.get().create(-1);
     }
+    // Tuinity end
 
     protected void encode(ChannelHandlerContext channelhandlercontext, ByteBuf bytebuf, ByteBuf bytebuf1) throws Exception {
         int i = bytebuf.readableBytes();
@@ -24,24 +28,46 @@ public class PacketCompressor extends MessageToByteEncoder<ByteBuf> {
             packetdataserializer.d(0);
             packetdataserializer.writeBytes(bytebuf);
         } else {
-            byte[] abyte = new byte[i];
-
-            bytebuf.readBytes(abyte);
-            packetdataserializer.d(abyte.length);
-            this.b.setInput(abyte, 0, i);
-            this.b.finish();
-
-            while (!this.b.finished()) {
-                int j = this.b.deflate(this.a);
-
-                packetdataserializer.writeBytes(this.a, 0, j);
+            // Tuinity start - delegate to Velocity natives
+//            byte[] abyte = new byte[i];
+//
+//            bytebuf.readBytes(abyte);
+//            packetdataserializer.d(abyte.length);
+//            this.b.setInput(abyte, 0, i);
+//            this.b.finish();
+//
+//            while (!this.b.finished()) {
+//                int j = this.b.deflate(this.a);
+//
+//                packetdataserializer.writeBytes(this.a, 0, j);
+//            }
+//
+//            this.b.reset();
+            packetdataserializer.d(i);
+            ByteBuf source = com.velocitypowered.natives.util.MoreByteBufUtils.ensureCompatible(channelhandlercontext.alloc(),
+                    this.compressor, bytebuf);
+            try {
+                this.compressor.deflate(source, bytebuf1);
+            } finally {
+                source.release();
             }
-
-            this.b.reset();
+            // Tuinity end
         }
 
     }
 
+    // Tuinity start
+    @Override
+    protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect) throws Exception {
+        return com.velocitypowered.natives.util.MoreByteBufUtils.preferredBuffer(ctx.alloc(), this.compressor, msg.readableBytes() + 1);
+    }
+
+    @Override
+    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
+        this.compressor.close();
+    }
+    // Tuinity end
+
     public void a(int i) {
         this.c = i;
     }
diff --git a/src/main/java/net/minecraft/network/PacketDecompressor.java b/src/main/java/net/minecraft/network/PacketDecompressor.java
index 1932ca55dad37ca773f215eaec23164533d509d3..ce7d7b4e8c1e5de935c83d7db5186ad8a0849e60 100644
--- a/src/main/java/net/minecraft/network/PacketDecompressor.java
+++ b/src/main/java/net/minecraft/network/PacketDecompressor.java
@@ -10,13 +10,17 @@ import java.util.zip.Inflater;
 
 public class PacketDecompressor extends ByteToMessageDecoder {
 
-    private final Inflater a;
+    // Tuinity start - use Velocity natives
+    //private final Inflater a;
+    private final com.velocitypowered.natives.compression.VelocityCompressor compressor;
     private int b;
 
     public PacketDecompressor(int i) {
         this.b = i;
-        this.a = new Inflater();
+        //this.a = new Inflater();
+        this.compressor = com.velocitypowered.natives.util.Natives.compress.get().create(-1);
     }
+    // Tuinity end
 
     protected void decode(ChannelHandlerContext channelhandlercontext, ByteBuf bytebuf, List<Object> list) throws Exception {
         if (bytebuf.readableBytes() != 0) {
@@ -34,20 +38,41 @@ public class PacketDecompressor extends ByteToMessageDecoder {
                     throw new DecoderException("Badly compressed packet - size of " + i + " is larger than protocol maximum of " + 2097152);
                 }
 
-                byte[] abyte = new byte[packetdataserializer.readableBytes()];
-
-                packetdataserializer.readBytes(abyte);
-                this.a.setInput(abyte);
-                byte[] abyte1 = new byte[i];
-
-                this.a.inflate(abyte1);
-                list.add(Unpooled.wrappedBuffer(abyte1));
-                this.a.reset();
+                // Tuinity start
+                ByteBuf compatibleIn = com.velocitypowered.natives.util.MoreByteBufUtils.ensureCompatible(channelhandlercontext.alloc(), compressor, bytebuf);
+                ByteBuf uncompressed = com.velocitypowered.natives.util.MoreByteBufUtils.preferredBuffer(channelhandlercontext.alloc(), compressor, i);
+                try {
+                    compressor.inflate(compatibleIn, uncompressed, i);
+                    list.add(uncompressed);
+                    bytebuf.clear();
+                } catch (Exception e) {
+                    uncompressed.release();
+                    throw e;
+                } finally {
+                    compatibleIn.release();
+                }
+//                byte[] abyte = new byte[packetdataserializer.readableBytes()];
+//
+//                packetdataserializer.readBytes(abyte);
+//                this.a.setInput(abyte);
+//                byte[] abyte1 = new byte[i];
+//
+//                this.a.inflate(abyte1);
+//                list.add(Unpooled.wrappedBuffer(abyte1));
+//                this.a.reset();
+                // Tuinity end
             }
 
         }
     }
 
+    // Tuinity start
+    @Override
+    public void handlerRemoved0(ChannelHandlerContext ctx) throws Exception {
+        this.compressor.close();
+    }
+    // Tuinity end
+
     public void a(int i) {
         this.b = i;
     }
diff --git a/src/main/java/net/minecraft/network/PacketDecrypter.java b/src/main/java/net/minecraft/network/PacketDecrypter.java
index c8aa02b288d67efe1f9e20e038248a4f032f92dd..4bd7c7db965818054929c9222a0968f6d127cc5c 100644
--- a/src/main/java/net/minecraft/network/PacketDecrypter.java
+++ b/src/main/java/net/minecraft/network/PacketDecrypter.java
@@ -8,13 +8,24 @@ import javax.crypto.Cipher;
 
 public class PacketDecrypter extends MessageToMessageDecoder<ByteBuf> {
 
-    private final PacketEncryptionHandler a;
+    // Tuinity start
+    private final com.velocitypowered.natives.encryption.VelocityCipher cipher;
+    //private final PacketEncryptionHandler a;
 
-    public PacketDecrypter(Cipher cipher) {
-        this.a = new PacketEncryptionHandler(cipher);
+    public PacketDecrypter(javax.crypto.SecretKey key /* Cipher cipher */) throws java.security.GeneralSecurityException {
+        //this.a = new PacketEncryptionHandler(cipher);
+        this.cipher = com.velocitypowered.natives.util.Natives.cipher.get().forDecryption(key);
     }
 
     protected void decode(ChannelHandlerContext channelhandlercontext, ByteBuf bytebuf, List<Object> list) throws Exception {
-        list.add(this.a.a(channelhandlercontext, bytebuf));
+        ByteBuf compatible = com.velocitypowered.natives.util.MoreByteBufUtils.ensureCompatible(channelhandlercontext.alloc(), cipher, bytebuf).slice();
+        try {
+            cipher.process(compatible);
+            list.add(compatible);
+        } catch (Exception e) {
+            compatible.release(); // compatible will never be used if we throw an exception
+            throw e;
+        }
     }
+    // Tuinity end
 }
diff --git a/src/main/java/net/minecraft/network/PacketEncrypter.java b/src/main/java/net/minecraft/network/PacketEncrypter.java
index 5759f91d5e9dc52b16c8955b8d318da2b53c7af4..0960f415fd3c6c764cf3dd273bb9e9c59d80ab31 100644
--- a/src/main/java/net/minecraft/network/PacketEncrypter.java
+++ b/src/main/java/net/minecraft/network/PacketEncrypter.java
@@ -5,15 +5,38 @@ import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.MessageToByteEncoder;
 import javax.crypto.Cipher;
 
-public class PacketEncrypter extends MessageToByteEncoder<ByteBuf> {
+// Tuinity start
+// We rewrite this class as the Velocity natives support in-place encryption
+import io.netty.handler.codec.MessageToMessageEncoder; // An unfortunate import, but this is required to fix a compiler error
+public class PacketEncrypter extends MessageToMessageEncoder<ByteBuf> {
 
-    private final PacketEncryptionHandler a;
+    private final com.velocitypowered.natives.encryption.VelocityCipher cipher;
+    //private final PacketEncryptionHandler a;
 
-    public PacketEncrypter(Cipher cipher) {
-        this.a = new PacketEncryptionHandler(cipher);
+    public PacketEncrypter(javax.crypto.SecretKey key /* Cipher cipher */) throws java.security.GeneralSecurityException {
+        // this.a = new PacketEncryptionHandler(cipher);
+        this.cipher = com.velocitypowered.natives.util.Natives.cipher.get().forEncryption(key);
     }
 
-    protected void encode(ChannelHandlerContext channelhandlercontext, ByteBuf bytebuf, ByteBuf bytebuf1) throws Exception {
-        this.a.a(bytebuf, bytebuf1);
+//    protected void encode(ChannelHandlerContext channelhandlercontext, ByteBuf bytebuf, ByteBuf bytebuf1) throws Exception {
+//        this.a.a(bytebuf, bytebuf1);
+//    }
+
+    @Override
+    protected void encode(ChannelHandlerContext ctx, ByteBuf msg, java.util.List<Object> out) throws Exception {
+        ByteBuf compatible = com.velocitypowered.natives.util.MoreByteBufUtils.ensureCompatible(ctx.alloc(), this.cipher, msg);
+        try {
+            this.cipher.process(compatible);
+            out.add(compatible);
+        } catch (Exception e) {
+            compatible.release(); // compatible will never be used if we throw an exception
+            throw e;
+        }
+    }
+
+    @Override
+    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
+        cipher.close();
     }
 }
+// Tuinity end
diff --git a/src/main/java/net/minecraft/server/network/LoginListener.java b/src/main/java/net/minecraft/server/network/LoginListener.java
index 185667110cd6f566b23546728d20fc79223f3c92..dc98ef48a664d9ee2a302fff8c611fd16dc704b6 100644
--- a/src/main/java/net/minecraft/server/network/LoginListener.java
+++ b/src/main/java/net/minecraft/server/network/LoginListener.java
@@ -255,7 +255,7 @@ public class LoginListener implements PacketLoginInListener {
 
             s = (new BigInteger(MinecraftEncryption.a("", this.server.getKeyPair().getPublic(), this.loginKey))).toString(16);
             this.g = LoginListener.EnumProtocolState.AUTHENTICATING;
-            this.networkManager.a(cipher, cipher1);
+            this.networkManager.a(this.loginKey); // Tuinity
         } catch (CryptographyException cryptographyexception) {
             throw new IllegalStateException("Protocol error", cryptographyexception);
         }
diff --git a/src/main/java/net/minecraft/server/network/ServerConnection.java b/src/main/java/net/minecraft/server/network/ServerConnection.java
index dc362724ea0cc1b2f9d9ceffff483217b4356c40..70fde7bad2e0a6d7432d8509fdb7c46d9f020d4c 100644
--- a/src/main/java/net/minecraft/server/network/ServerConnection.java
+++ b/src/main/java/net/minecraft/server/network/ServerConnection.java
@@ -87,6 +87,11 @@ public class ServerConnection {
                 ServerConnection.LOGGER.info("Using default channel type");
             }
 
+            // Tuinity start - indicate Velocity natives in use
+            ServerConnection.LOGGER.info("Tuinity: Using " + com.velocitypowered.natives.util.Natives.compress.getLoadedVariant() + " compression from Velocity.");
+            ServerConnection.LOGGER.info("Tuinity: Using " + com.velocitypowered.natives.util.Natives.cipher.getLoadedVariant() + " cipher from Velocity.");
+            // Tuinity end
+
             this.listeningChannels.add(((ServerBootstrap) ((ServerBootstrap) (new ServerBootstrap()).channel(oclass)).childHandler(new ChannelInitializer<Channel>() {
                 protected void initChannel(Channel channel) throws Exception {
                     try {
