From 671df77c1f5f2328a7359b9ef2f4254287c2e184 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 12 Aug 2018 08:06:50 -0400
Subject: [PATCH] Forward the keep-alive packet directly onto the client.

This improves anti-cheat compatibility.
---
 .../backend/BackendPlaySessionHandler.java    |  5 +--
 .../client/ClientPlaySessionHandler.java      | 32 +++++--------------
 2 files changed, 11 insertions(+), 26 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 b0784b179..bd6ad4d77 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
@@ -27,8 +27,9 @@ public class BackendPlaySessionHandler implements MinecraftSessionHandler {
         ClientPlaySessionHandler playerHandler =
                 (ClientPlaySessionHandler) connection.getProxyPlayer().getConnection().getSessionHandler();
         if (packet instanceof KeepAlive) {
-            // Forward onto the server
-            connection.getMinecraftConnection().write(packet);
+            // Forward onto the player
+            playerHandler.setLastPing(((KeepAlive) packet).getRandomId());
+            connection.getProxyPlayer().getConnection().write(packet);
         } else if (packet instanceof Disconnect) {
             Disconnect original = (Disconnect) packet;
             connection.getProxyPlayer().handleConnectionException(connection.getServerInfo(), original);
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 c1451ef9e..549c8fb46 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
@@ -25,7 +25,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     private static final int MAX_PLUGIN_CHANNELS = 128;
 
     private final ConnectedPlayer player;
-    private ScheduledFuture<?> pingTask;
     private long lastPing = -1;
     private boolean spawned = false;
     private final List<UUID> serverBossBars = new ArrayList<>();
@@ -36,30 +35,15 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
         this.player = player;
     }
 
-    @Override
-    public void activated() {
-        EventLoop loop = player.getConnection().getChannel().eventLoop();
-        pingTask = loop.scheduleAtFixedRate(this::ping, 5, 15, TimeUnit.SECONDS);
-    }
-
-    private void ping() {
-        long randomId = ThreadLocalRandom.current().nextInt();
-        lastPing = randomId;
-        KeepAlive keepAlive = new KeepAlive();
-        keepAlive.setRandomId(randomId);
-        player.getConnection().write(keepAlive);
-    }
-
     @Override
     public void handle(MinecraftPacket packet) {
         if (packet instanceof KeepAlive) {
             KeepAlive keepAlive = (KeepAlive) packet;
             if (keepAlive.getRandomId() != lastPing) {
-                throw new IllegalStateException("Client sent invalid keepAlive; expected " + lastPing + ", got " + keepAlive.getRandomId());
+                // The last keep alive we got was probably from a different server. Let's ignore it, and hope the next
+                // ping is alright.
+                return;
             }
-
-            // Do not forward the packet to the player's server, because we handle pings for all servers already.
-            return;
         }
 
         if (packet instanceof ClientSettings) {
@@ -132,11 +116,6 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     @Override
     public void disconnected() {
         player.teardown();
-
-        if (pingTask != null && !pingTask.isCancelled()) {
-            pingTask.cancel(false);
-            pingTask = null;
-        }
     }
 
     @Override
@@ -149,6 +128,7 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     }
 
     public void handleBackendJoinGame(JoinGame joinGame) {
+        lastPing = Long.MIN_VALUE; // reset last ping
         if (!spawned) {
             // nothing special to do here
             spawned = true;
@@ -244,4 +224,8 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     public EntityIdRemapper getIdRemapper() {
         return idRemapper;
     }
+
+    public void setLastPing(long lastPing) {
+        this.lastPing = lastPing;
+    }
 }