From 8b7b8cdffdb7aba97ef817037859d1151f4b2665 Mon Sep 17 00:00:00 2001
From: onebeastchris <github@onechris.mozmail.com>
Date: Mon, 29 Apr 2024 16:08:14 +0200
Subject: [PATCH] Properly shutdown LocalSession's, ensure transferring works
 properly regardless if we're injected or not

---
 bootstrap/mod/fabric/build.gradle.kts               |  2 --
 .../platform/spigot/GeyserSpigotInjector.java       |  2 +-
 .../geysermc/geyser/network/netty/LocalSession.java | 13 +++++++------
 .../org/geysermc/geyser/session/GeyserSession.java  |  4 ++--
 4 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/bootstrap/mod/fabric/build.gradle.kts b/bootstrap/mod/fabric/build.gradle.kts
index 53e4dfe53..cd513c1e4 100644
--- a/bootstrap/mod/fabric/build.gradle.kts
+++ b/bootstrap/mod/fabric/build.gradle.kts
@@ -49,9 +49,7 @@ application {
 
 relocate("org.cloudburstmc.netty")
 relocate("org.cloudburstmc.protocol")
-relocate("com.github.steveice10.mc.protocol")
 relocate("com.github.steveice10.mc.auth")
-relocate("com.github.steveice10.packetlib")
 
 tasks {
     remapJar {
diff --git a/bootstrap/spigot/src/main/java/org/geysermc/geyser/platform/spigot/GeyserSpigotInjector.java b/bootstrap/spigot/src/main/java/org/geysermc/geyser/platform/spigot/GeyserSpigotInjector.java
index 6d22a77ae..5dcfbd0f8 100644
--- a/bootstrap/spigot/src/main/java/org/geysermc/geyser/platform/spigot/GeyserSpigotInjector.java
+++ b/bootstrap/spigot/src/main/java/org/geysermc/geyser/platform/spigot/GeyserSpigotInjector.java
@@ -178,7 +178,7 @@ public class GeyserSpigotInjector extends GeyserInjector {
         MinecraftProtocol protocol = new MinecraftProtocol();
         LocalSession session = new LocalSession(bootstrap.getGeyserConfig().getRemote().address(),
                 bootstrap.getGeyserConfig().getRemote().port(), this.serverSocketAddress,
-                InetAddress.getLoopbackAddress().getHostAddress(), protocol, protocol.createHelper(), false);
+                InetAddress.getLoopbackAddress().getHostAddress(), protocol, protocol.createHelper());
         session.connect();
         session.disconnect("");
     }
diff --git a/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java b/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
index b5598d063..958e88288 100644
--- a/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
+++ b/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
@@ -30,6 +30,7 @@ import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.*;
 import io.netty.channel.unix.PreferredDirectByteBufAllocator;
 import io.netty.handler.codec.haproxy.*;
+import io.netty.util.concurrent.DefaultThreadFactory;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.geysermc.mcprotocollib.network.BuiltinFlags;
 import org.geysermc.mcprotocollib.network.codec.PacketCodecHelper;
@@ -42,6 +43,7 @@ import org.geysermc.mcprotocollib.protocol.codec.MinecraftCodecHelper;
 import java.net.Inet4Address;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Manages a Minecraft Java session over our LocalChannel implementations.
@@ -54,24 +56,23 @@ public final class LocalSession extends TcpSession {
     private final String clientIp;
     private final PacketCodecHelper codecHelper;
 
-    private final boolean transferring;
-
-    public LocalSession(String host, int port, SocketAddress targetAddress, String clientIp, PacketProtocol protocol, MinecraftCodecHelper codecHelper, boolean transferring) {
+    public LocalSession(String host, int port, SocketAddress targetAddress, String clientIp, PacketProtocol protocol, MinecraftCodecHelper codecHelper) {
         super(host, port, protocol);
         this.targetAddress = targetAddress;
         this.clientIp = clientIp;
         this.codecHelper = codecHelper;
-        this.transferring = transferring;
     }
 
     @Override
-    public void connect(boolean wait) {
+    public void connect(boolean wait, boolean transferring) {
         if (this.disconnected) {
             throw new IllegalStateException("Connection has already been disconnected.");
         }
 
         if (DEFAULT_EVENT_LOOP_GROUP == null) {
-            DEFAULT_EVENT_LOOP_GROUP = new DefaultEventLoopGroup();
+            DEFAULT_EVENT_LOOP_GROUP = new DefaultEventLoopGroup(new DefaultThreadFactory(this.getClass(), true));
+            Runtime.getRuntime().addShutdownHook(new Thread(
+                    () -> DEFAULT_EVENT_LOOP_GROUP.shutdownGracefully(100, 500, TimeUnit.MILLISECONDS)));
         }
 
         try {
diff --git a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
index 869999357..d10a20b3d 100644
--- a/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
+++ b/core/src/main/java/org/geysermc/geyser/session/GeyserSession.java
@@ -880,7 +880,7 @@ public class GeyserSession implements GeyserConnection, GeyserCommandSource {
             // We're going to connect through the JVM and not through TCP
             downstream = new LocalSession(this.remoteServer.address(), this.remoteServer.port(),
                     geyser.getBootstrap().getSocketAddress(), upstream.getAddress().getAddress().getHostAddress(),
-                    this.protocol, this.protocol.createHelper(), loginEvent.transferring());
+                    this.protocol, this.protocol.createHelper());
             this.downstream = new DownstreamSession(downstream);
         } else {
             downstream = new TcpClientSession(this.remoteServer.address(), this.remoteServer.port(), this.protocol);
@@ -1070,7 +1070,7 @@ public class GeyserSession implements GeyserConnection, GeyserCommandSource {
             setDaylightCycle(true);
         }
 
-        downstream.connect(false);
+        downstream.connect(false, loginEvent.transferring());
     }
 
     public void disconnect(String reason) {