From 11806b9166eb09794f5c7c0025edb23ac10465bd Mon Sep 17 00:00:00 2001
From: TheMode <themode@outlook.fr>
Date: Wed, 6 Oct 2021 21:14:08 +0200
Subject: [PATCH] Ensure no packet corruption due to encryption. A proxy is
 still recommended for optimal performance

Signed-off-by: TheMode <themode@outlook.fr>
---
 .../player/PlayerSocketConnection.java        | 68 ++++++++-----------
 1 file changed, 30 insertions(+), 38 deletions(-)

diff --git a/src/main/java/net/minestom/server/network/player/PlayerSocketConnection.java b/src/main/java/net/minestom/server/network/player/PlayerSocketConnection.java
index ee514a957..84eb58990 100644
--- a/src/main/java/net/minestom/server/network/player/PlayerSocketConnection.java
+++ b/src/main/java/net/minestom/server/network/player/PlayerSocketConnection.java
@@ -231,6 +231,17 @@ public class PlayerSocketConnection extends PlayerConnection {
     @ApiStatus.Internal
     public void write(@NotNull ByteBuffer buffer) {
         synchronized (bufferLock) {
+            if (encrypted) { // Encryption support
+                ByteBuffer output = PacketUtils.localBuffer();
+                try {
+                    this.encryptCipher.update(buffer, output);
+                    buffer = output.flip();
+                } catch (ShortBufferException e) {
+                    MinecraftServer.getExceptionManager().handleException(e);
+                    return;
+                }
+            }
+
             BinaryBuffer localBuffer = tickBuffer.getPlain();
             final int capacity = localBuffer.capacity();
             final int size = buffer.remaining();
@@ -273,47 +284,28 @@ public class PlayerSocketConnection extends PlayerConnection {
                 final BinaryBuffer localBuffer = tickBuffer.getPlain();
                 final boolean emptyWaitingList = waitingBuffers.isEmpty();
                 if (localBuffer.readableBytes() == 0 && emptyWaitingList) return;
-                // Update tick buffer
                 try {
-                    // Encryption support
-                    if (encrypted) {
-                        ByteBuffer cipherInput = localBuffer.asByteBuffer(0, localBuffer.writerOffset());
-                        try {
-                            this.encryptCipher.update(cipherInput, cipherInput.duplicate());
-                        } catch (ShortBufferException e) {
-                            MinecraftServer.getExceptionManager().handleException(e);
-                        }
+                    // Try to write the current buffer
+                    if (emptyWaitingList && localBuffer.writeChannel(channel)) {
+                        // Can reuse buffer
+                        localBuffer.clear();
+                        return;
                     }
-
-                    try {
-                        // Try to write the current buffer
-                        if (emptyWaitingList && localBuffer.writeChannel(channel)) {
-                            // Can reuse buffer
-                            localBuffer.clear();
-                            return;
-                        }
-                        this.tickBuffer.setPlain(PooledBuffers.get());
-                        this.waitingBuffers.add(localBuffer);
-                        if (emptyWaitingList) return;
-                        // Write as much as possible from the waiting list
-                        Iterator<BinaryBuffer> iterator = waitingBuffers.iterator();
-                        while (iterator.hasNext()) {
-                            BinaryBuffer waitingBuffer = iterator.next();
-                            try {
-                                if (!waitingBuffer.writeChannel(channel)) break;
-                                iterator.remove();
-                                PooledBuffers.add(waitingBuffer);
-                            } catch (IOException e) {
-                                MinecraftServer.getExceptionManager().handleException(e);
-                                throw new ClosedChannelException();
-                            }
-                        }
-                    } catch (IOException e) {
-                        MinecraftServer.getExceptionManager().handleException(e);
-                        throw new ClosedChannelException();
+                    this.tickBuffer.setPlain(PooledBuffers.get());
+                    this.waitingBuffers.add(localBuffer);
+                    if (emptyWaitingList) return;
+                    // Write as much as possible from the waiting list
+                    Iterator<BinaryBuffer> iterator = waitingBuffers.iterator();
+                    while (iterator.hasNext()) {
+                        BinaryBuffer waitingBuffer = iterator.next();
+                        if (!waitingBuffer.writeChannel(channel)) break;
+                        iterator.remove();
+                        PooledBuffers.add(waitingBuffer);
                     }
-
-                } catch (OutOfMemoryError e) {
+                } catch (IOException e) { // Couldn't write to the socket
+                    MinecraftServer.getExceptionManager().handleException(e);
+                    throw new ClosedChannelException();
+                } catch (OutOfMemoryError e) { // Couldn't allocate a pooled buffer
                     this.waitingBuffers.clear();
                     System.gc(); // Explicit gc forcing buffers to be collected
                     throw new ClosedChannelException();