From ab55050d11d3c449ed6329494a758fbd0c69551b Mon Sep 17 00:00:00 2001
From: SupremeMortal <6178101+SupremeMortal@users.noreply.github.com>
Date: Fri, 13 Sep 2019 14:37:31 +0100
Subject: [PATCH] Close remote java connection when bedrock player leaves.

---
 api/src/main/java/org/geysermc/api/Player.java       |  9 +++++++++
 .../java/org/geysermc/connector/GeyserConnector.java |  2 ++
 .../network/ConnectorServerEventHandler.java         | 10 +++++++++-
 .../connector/network/session/GeyserSession.java     | 12 +++++++++---
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/api/src/main/java/org/geysermc/api/Player.java b/api/src/main/java/org/geysermc/api/Player.java
index 7df58f46..f8107cb1 100644
--- a/api/src/main/java/org/geysermc/api/Player.java
+++ b/api/src/main/java/org/geysermc/api/Player.java
@@ -29,6 +29,8 @@ import org.geysermc.api.command.CommandSender;
 import org.geysermc.api.session.AuthData;
 import org.geysermc.api.window.FormWindow;
 
+import java.net.InetSocketAddress;
+
 public interface Player extends CommandSender {
 
     /**
@@ -68,4 +70,11 @@ public interface Player extends CommandSender {
      * @param id the id of the window
      */
     void sendForm(FormWindow window, int id);
+
+    /**
+     * Returns the current hostname and port the player is connected with.
+     *
+     * @return player's socket address.
+     */
+    InetSocketAddress getSocketAddress();
 }
diff --git a/connector/src/main/java/org/geysermc/connector/GeyserConnector.java b/connector/src/main/java/org/geysermc/connector/GeyserConnector.java
index 63bdc8bf..28ab191a 100644
--- a/connector/src/main/java/org/geysermc/connector/GeyserConnector.java
+++ b/connector/src/main/java/org/geysermc/connector/GeyserConnector.java
@@ -173,10 +173,12 @@ public class GeyserConnector implements Connector {
     public void addPlayer(Player player) {
         players.put(player.getAuthenticationData().getName(), player);
         players.put(player.getAuthenticationData().getUUID(), player);
+        players.put(player.getSocketAddress(), player);
     }
 
     public void removePlayer(Player player) {
         players.remove(player.getAuthenticationData().getName());
         players.remove(player.getAuthenticationData().getUUID());
+        players.remove(player.getSocketAddress());
     }
 }
diff --git a/connector/src/main/java/org/geysermc/connector/network/ConnectorServerEventHandler.java b/connector/src/main/java/org/geysermc/connector/network/ConnectorServerEventHandler.java
index 0ac0f997..d89151bc 100644
--- a/connector/src/main/java/org/geysermc/connector/network/ConnectorServerEventHandler.java
+++ b/connector/src/main/java/org/geysermc/connector/network/ConnectorServerEventHandler.java
@@ -29,6 +29,7 @@ import com.github.steveice10.mc.protocol.data.status.ServerStatusInfo;
 import com.nukkitx.protocol.bedrock.BedrockPong;
 import com.nukkitx.protocol.bedrock.BedrockServerEventHandler;
 import com.nukkitx.protocol.bedrock.BedrockServerSession;
+import org.geysermc.api.Player;
 import org.geysermc.api.events.PingEvent;
 import org.geysermc.connector.GeyserConnector;
 import org.geysermc.connector.configuration.GeyserConfiguration;
@@ -102,7 +103,14 @@ public class ConnectorServerEventHandler implements BedrockServerEventHandler {
     public void onSessionCreation(BedrockServerSession bedrockServerSession) {
         bedrockServerSession.setLogging(true);
         bedrockServerSession.setPacketHandler(new UpstreamPacketHandler(connector, new GeyserSession(connector, bedrockServerSession)));
-        bedrockServerSession.addDisconnectHandler((x) -> GeyserLogger.DEFAULT.warning("Bedrock user with ip: " + bedrockServerSession.getAddress().getAddress() + " has disconnected for reason " + x));
+        bedrockServerSession.addDisconnectHandler(disconnectReason -> {
+            GeyserLogger.DEFAULT.info("Bedrock user with ip: " + bedrockServerSession.getAddress().getAddress() + " has disconnected for reason " + disconnectReason);
+
+            Player player = connector.getPlayers().get(bedrockServerSession.getAddress());
+            if (player != null) {
+                player.disconnect(disconnectReason.name());
+            }
+        });
         bedrockServerSession.setPacketCodec(GeyserConnector.BEDROCK_PACKET_CODEC);
     }
 }
\ No newline at end of file
diff --git a/connector/src/main/java/org/geysermc/connector/network/session/GeyserSession.java b/connector/src/main/java/org/geysermc/connector/network/session/GeyserSession.java
index 3eed5ada..8503bb52 100644
--- a/connector/src/main/java/org/geysermc/connector/network/session/GeyserSession.java
+++ b/connector/src/main/java/org/geysermc/connector/network/session/GeyserSession.java
@@ -60,14 +60,15 @@ import org.geysermc.connector.network.session.cache.*;
 import org.geysermc.connector.network.translators.Registry;
 import org.geysermc.connector.utils.Toolbox;
 
+import java.net.InetSocketAddress;
 import java.util.UUID;
 
 @Getter
 public class GeyserSession implements PlayerSession, Player {
 
-    private GeyserConnector connector;
+    private final GeyserConnector connector;
+    private final BedrockServerSession upstream;
     private RemoteServer remoteServer;
-    private BedrockServerSession upstream;
 
     private Client downstream;
 
@@ -179,7 +180,7 @@ public class GeyserSession implements PlayerSession, Player {
             if (downstream != null && downstream.getSession() != null) {
                 downstream.getSession().disconnect(reason);
             }
-            if (upstream != null) {
+            if (upstream != null && !upstream.isClosed()) {
                 upstream.disconnect(reason);
             }
         }
@@ -241,6 +242,11 @@ public class GeyserSession implements PlayerSession, Player {
         windowCache.showWindow(window, id);
     }
 
+    @Override
+    public InetSocketAddress getSocketAddress() {
+        return this.upstream.getAddress();
+    }
+
     public void sendForm(FormWindow window) {
         windowCache.showWindow(window);
     }