From bac7e8621727fd0f23bb81fb4cf3d852e7d15a9c Mon Sep 17 00:00:00 2001
From: creeper123123321 <7974274+creeper123123321@users.noreply.github.com>
Date: Sat, 5 Dec 2020 16:06:58 -0300
Subject: [PATCH] Fix #80 #81

---
 .../clientside/ProtocolDetectionHandler.java       | 12 +++++++-----
 .../viafabric/service/ProtocolAutoDetector.java    | 14 ++++++++++----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/main/java/com/github/creeper123123321/viafabric/handler/clientside/ProtocolDetectionHandler.java b/src/main/java/com/github/creeper123123321/viafabric/handler/clientside/ProtocolDetectionHandler.java
index f35e307..52d5124 100644
--- a/src/main/java/com/github/creeper123123321/viafabric/handler/clientside/ProtocolDetectionHandler.java
+++ b/src/main/java/com/github/creeper123123321/viafabric/handler/clientside/ProtocolDetectionHandler.java
@@ -30,7 +30,6 @@ import com.github.creeper123123321.viafabric.service.ProtocolAutoDetector;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
-import io.netty.util.ReferenceCountUtil;
 import us.myles.ViaVersion.api.Pair;
 
 import java.net.InetSocketAddress;
@@ -59,8 +58,6 @@ public class ProtocolDetectionHandler extends ChannelDuplexHandler {
                 }, 10, TimeUnit.SECONDS);
                 ProtocolAutoDetector.SERVER_VER.get(((InetSocketAddress) ctx.channel().remoteAddress()))
                         .whenComplete((obj, ex) -> {
-                            hold = false;
-                            drainQueue(ctx);
                             ctx.pipeline().remove(this);
                             timeoutRun.cancel(false);
                         });
@@ -91,6 +88,12 @@ public class ProtocolDetectionHandler extends ChannelDuplexHandler {
         }
     }
 
+    @Override
+    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+        drainQueue(ctx);
+        super.channelInactive(ctx);
+    }
+
     private void drainQueue(ChannelHandlerContext ctx) {
         queuedMessages.forEach(it -> ctx.write(it.getKey(), it.getValue()));
         queuedMessages.clear();
@@ -100,8 +103,7 @@ public class ProtocolDetectionHandler extends ChannelDuplexHandler {
 
     @Override
     public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
-        queuedMessages.forEach(ReferenceCountUtil::release);
-        queuedMessages.clear();
+        drainQueue(ctx);
         super.handlerRemoved(ctx);
     }
 }
diff --git a/src/main/java/com/github/creeper123123321/viafabric/service/ProtocolAutoDetector.java b/src/main/java/com/github/creeper123123321/viafabric/service/ProtocolAutoDetector.java
index 9934400..362f9f7 100644
--- a/src/main/java/com/github/creeper123123321/viafabric/service/ProtocolAutoDetector.java
+++ b/src/main/java/com/github/creeper123123321/viafabric/service/ProtocolAutoDetector.java
@@ -42,6 +42,7 @@ import net.minecraft.network.packet.c2s.handshake.HandshakeC2SPacket;
 import net.minecraft.network.packet.c2s.query.QueryRequestC2SPacket;
 import net.minecraft.network.packet.s2c.query.QueryPongS2CPacket;
 import net.minecraft.network.packet.s2c.query.QueryResponseS2CPacket;
+import net.minecraft.server.ServerMetadata;
 import net.minecraft.text.LiteralText;
 import net.minecraft.text.Text;
 import us.myles.ViaVersion.api.protocol.ProtocolVersion;
@@ -92,10 +93,15 @@ public class ProtocolAutoDetector {
                                 clientConnection.setPacketListener(new ClientQueryPacketListener() {
                                     @Override
                                     public void onResponse(QueryResponseS2CPacket packet) {
-                                        ProtocolVersion ver = ProtocolVersion.getProtocol(packet.getServerMetadata().getVersion()
-                                                .getProtocolVersion());
-                                        future.complete(ver);
-                                        ViaFabric.JLOGGER.info("Auto-detected " + ver + " for " + address);
+                                        ServerMetadata meta = packet.getServerMetadata();
+                                        ServerMetadata.Version version;
+                                        if (meta != null && (version = meta.getVersion()) != null) {
+                                            ProtocolVersion ver = ProtocolVersion.getProtocol(version.getProtocolVersion());
+                                            future.complete(ver);
+                                            ViaFabric.JLOGGER.info("Auto-detected " + ver + " for " + address);
+                                        } else {
+                                            future.completeExceptionally(new IllegalArgumentException("Null version in query response"));
+                                        }
                                         clientConnection.disconnect(LiteralText.EMPTY);
                                     }