From 2e38e0e1cb62b42bfe456355f379b6577f12c07f Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sat, 11 Apr 2020 21:52:01 -0400
Subject: [PATCH] Properly fix the previous patch

---
 .../proxy/connection/MinecraftConnection.java | 23 ++++++++++++++++---
 .../proxy/network/netty/DiscardHandler.java   | 17 ++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 proxy/src/main/java/com/velocitypowered/proxy/network/netty/DiscardHandler.java

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
index beefa6b05..d36846cc7 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
@@ -16,6 +16,7 @@ import com.velocitypowered.natives.encryption.VelocityCipher;
 import com.velocitypowered.natives.encryption.VelocityCipherFactory;
 import com.velocitypowered.natives.util.Natives;
 import com.velocitypowered.proxy.VelocityServer;
+import com.velocitypowered.proxy.network.netty.DiscardHandler;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCipherDecoder;
@@ -131,8 +132,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
         try {
           sessionHandler.exception(cause);
         } catch (Exception ex) {
-          logger.error("{}: exception handling exception", (association != null ? association :
-              channel.remoteAddress()), cause);
+          logger.error("{}: exception handling exception in {}",
+              (association != null ? association : channel.remoteAddress()), sessionHandler, cause);
         }
       }
 
@@ -140,10 +141,11 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
         if (cause instanceof ReadTimeoutException) {
           logger.error("{}: read timed out", association);
         } else {
-          logger.error("{}: exception encountered", association, cause);
+          logger.error("{}: exception encountered in {}", association, sessionHandler, cause);
         }
       }
 
+      ctx.pipeline().addBefore(MINECRAFT_DECODER, "discard", DiscardHandler.HANDLER);
       ctx.close();
     }
   }
@@ -155,6 +157,10 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
     }
   }
 
+  private void ensureInEventLoop() {
+    Preconditions.checkState(this.channel.eventLoop().inEventLoop(), "Not in event loop");
+  }
+
   public EventLoop eventLoop() {
     return channel.eventLoop();
   }
@@ -233,6 +239,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    * @param autoReading whether or not we should read data automatically
    */
   public void setAutoReading(boolean autoReading) {
+    ensureInEventLoop();
+
     channel.config().setAutoRead(autoReading);
     if (autoReading) {
       // For some reason, the channel may not completely read its queued contents once autoread
@@ -249,6 +257,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    * @param state the new state
    */
   public void setState(StateRegistry state) {
+    ensureInEventLoop();
+
     this.state = state;
     this.channel.pipeline().get(MinecraftEncoder.class).setState(state);
     this.channel.pipeline().get(MinecraftDecoder.class).setState(state);
@@ -263,6 +273,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    * @param protocolVersion the protocol version to use
    */
   public void setProtocolVersion(ProtocolVersion protocolVersion) {
+    ensureInEventLoop();
+
     this.protocolVersion = protocolVersion;
     this.nextProtocolVersion = protocolVersion;
     if (protocolVersion != ProtocolVersion.LEGACY) {
@@ -284,6 +296,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    * @param sessionHandler the handler to use
    */
   public void setSessionHandler(MinecraftSessionHandler sessionHandler) {
+    ensureInEventLoop();
+
     if (this.sessionHandler != null) {
       this.sessionHandler.deactivated();
     }
@@ -302,6 +316,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    */
   public void setCompressionThreshold(int threshold) {
     ensureOpen();
+    ensureInEventLoop();
 
     if (threshold == -1) {
       channel.pipeline().remove(COMPRESSION_DECODER);
@@ -325,6 +340,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
    */
   public void enableEncryption(byte[] secret) throws GeneralSecurityException {
     ensureOpen();
+    ensureInEventLoop();
 
     SecretKey key = new SecretKeySpec(secret, "AES");
 
@@ -342,6 +358,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
   }
 
   public void setAssociation(MinecraftConnectionAssociation association) {
+    ensureInEventLoop();
     this.association = association;
   }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/network/netty/DiscardHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/network/netty/DiscardHandler.java
new file mode 100644
index 000000000..b4421fd96
--- /dev/null
+++ b/proxy/src/main/java/com/velocitypowered/proxy/network/netty/DiscardHandler.java
@@ -0,0 +1,17 @@
+package com.velocitypowered.proxy.network.netty;
+
+import io.netty.channel.ChannelHandler.Sharable;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.util.ReferenceCountUtil;
+
+@Sharable
+public class DiscardHandler extends ChannelInboundHandlerAdapter {
+
+  public static final DiscardHandler HANDLER = new DiscardHandler();
+
+  @Override
+  public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+    ReferenceCountUtil.release(msg);
+  }
+}
\ No newline at end of file