From 9e999e1e5a3f709f57d0a747aef86ae31f8ca212 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Fri, 21 Sep 2018 23:18:58 -0400
Subject: [PATCH] Maintain keep alives on the server connection, not on the
 client.

---
 .../backend/BackendPlaySessionHandler.java     |  2 +-
 .../backend/VelocityServerConnection.java      | 15 +++++++++++++++
 .../client/ClientPlaySessionHandler.java       | 18 ++----------------
 3 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
index 88cba258d..35ef63d92 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BackendPlaySessionHandler.java
@@ -42,7 +42,7 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
                 (ClientPlaySessionHandler) connection.getPlayer().getConnection().getSessionHandler();
         if (packet instanceof KeepAlive) {
             // Forward onto the player
-            playerHandler.setLastPing(((KeepAlive) packet).getRandomId());
+            connection.setLastPingId(((KeepAlive) packet).getRandomId());
             connection.getPlayer().getConnection().write(packet);
         } else if (packet instanceof Disconnect) {
             Disconnect original = (Disconnect) packet;
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
index 80210e33a..5b0a8cf3d 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/VelocityServerConnection.java
@@ -44,6 +44,8 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
     private boolean legacyForge = false;
     private boolean hasCompletedJoin = false;
     private boolean gracefulDisconnect = false;
+    private long lastPingId;
+    private long lastPingSent;
 
     public VelocityServerConnection(VelocityRegisteredServer registeredServer, ConnectedPlayer proxyPlayer, VelocityServer server) {
         this.registeredServer = registeredServer;
@@ -193,4 +195,17 @@ public class VelocityServerConnection implements MinecraftConnectionAssociation,
     public boolean isGracefulDisconnect() {
         return gracefulDisconnect;
     }
+
+    public long getLastPingId() {
+        return lastPingId;
+    }
+
+    public long getLastPingSent() {
+        return lastPingSent;
+    }
+
+    public void setLastPingId(long lastPingId) {
+        this.lastPingId = lastPingId;
+        this.lastPingSent = System.currentTimeMillis();
+    }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
index cc50e7877..81471ffa3 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
@@ -29,8 +29,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     private static final int MAX_PLUGIN_CHANNELS = 1024;
 
     private final ConnectedPlayer player;
-    private long lastPingID = -1;
-    private long lastPingSent = -1;
     private boolean spawned = false;
     private final List<UUID> serverBossBars = new ArrayList<>();
     private final Set<String> clientPluginMsgChannels = new HashSet<>();
@@ -64,13 +62,12 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
 
         if (packet instanceof KeepAlive) {
             KeepAlive keepAlive = (KeepAlive) packet;
-            if (keepAlive.getRandomId() != lastPingID) {
+            if (keepAlive.getRandomId() != serverConnection.getLastPingId()) {
                 // The last keep alive we got was probably from a different server. Let's ignore it, and hope the next
                 // ping is alright.
                 return;
             }
-            player.setPing(System.currentTimeMillis() - lastPingSent);
-            resetPingData();
+            player.setPing(System.currentTimeMillis() - serverConnection.getLastPingSent());
             serverConnection.getMinecraftConnection().write(packet);
             return;
         }
@@ -155,7 +152,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     }
 
     public void handleBackendJoinGame(JoinGame joinGame) {
-        resetPingData(); // reset ping data
         if (!spawned) {
             // Nothing special to do with regards to spawning the player
             spawned = true;
@@ -300,16 +296,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
         return clientPluginMsgChannels;
     }
 
-    public void setLastPing(long lastPing) {
-        this.lastPingID = lastPing;
-        this.lastPingSent = System.currentTimeMillis();
-    }
-    
-    private void resetPingData() {
-        this.lastPingID = -1;
-        this.lastPingSent = -1;
-    }
-
     public void handleTabCompleteResponse(TabCompleteResponse response) {
         if (outstandingTabComplete != null) {
             if (!outstandingTabComplete.isAssumeCommand()) {