From 80477f6a0e3c2e7e14e78b11d8d07b4e0f5e4a3c Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Thu, 26 Jul 2018 14:50:42 -0400
Subject: [PATCH] Properly handle pings.

---
 .../proxy/connection/MinecraftConnection.java |  1 +
 .../connection/MinecraftSessionHandler.java   |  4 +++
 .../backend/PlaySessionHandler.java           |  2 +-
 .../connection/backend/ServerConnection.java  |  1 +
 .../connection/client/PlaySessionHandler.java | 35 +++++++++++++++++--
 5 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java b/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
index 5bc3e8f7d..17c8f9c06 100644
--- a/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
+++ b/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
@@ -133,6 +133,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
 
     public void setSessionHandler(MinecraftSessionHandler sessionHandler) {
         this.sessionHandler = sessionHandler;
+        sessionHandler.activated();
     }
 
     private void ensureOpen() {
diff --git a/src/main/java/com/velocitypowered/proxy/connection/MinecraftSessionHandler.java b/src/main/java/com/velocitypowered/proxy/connection/MinecraftSessionHandler.java
index 976c62b2f..b2eee3ff7 100644
--- a/src/main/java/com/velocitypowered/proxy/connection/MinecraftSessionHandler.java
+++ b/src/main/java/com/velocitypowered/proxy/connection/MinecraftSessionHandler.java
@@ -18,6 +18,10 @@ public interface MinecraftSessionHandler {
 
     }
 
+    default void activated() {
+
+    }
+
     default void exception(Throwable throwable) {
 
     }
diff --git a/src/main/java/com/velocitypowered/proxy/connection/backend/PlaySessionHandler.java b/src/main/java/com/velocitypowered/proxy/connection/backend/PlaySessionHandler.java
index 1cf2636ea..46e7e5dd3 100644
--- a/src/main/java/com/velocitypowered/proxy/connection/backend/PlaySessionHandler.java
+++ b/src/main/java/com/velocitypowered/proxy/connection/backend/PlaySessionHandler.java
@@ -19,7 +19,7 @@ public class PlaySessionHandler implements MinecraftSessionHandler {
     @Override
     public void handle(MinecraftPacket packet) {
         if (packet instanceof Ping) {
-            // Make sure to reply back to the server so it doesn't think we're gone.
+            // Forward onto the server
             connection.getChannel().write(packet);
         } else if (packet instanceof Disconnect) {
             // The server wants to disconnect us. TODO fallback handling
diff --git a/src/main/java/com/velocitypowered/proxy/connection/backend/ServerConnection.java b/src/main/java/com/velocitypowered/proxy/connection/backend/ServerConnection.java
index b4ecdfa60..0bea9ca94 100644
--- a/src/main/java/com/velocitypowered/proxy/connection/backend/ServerConnection.java
+++ b/src/main/java/com/velocitypowered/proxy/connection/backend/ServerConnection.java
@@ -72,6 +72,7 @@ public class ServerConnection {
         handshake.setPort(serverInfo.getAddress().getPort());
         channel.write(handshake);
 
+        channel.setProtocolVersion(proxyPlayer.getConnection().getProtocolVersion());
         channel.setState(StateRegistry.LOGIN);
 
         // Login
diff --git a/src/main/java/com/velocitypowered/proxy/connection/client/PlaySessionHandler.java b/src/main/java/com/velocitypowered/proxy/connection/client/PlaySessionHandler.java
index ce5aec44f..357921da6 100644
--- a/src/main/java/com/velocitypowered/proxy/connection/client/PlaySessionHandler.java
+++ b/src/main/java/com/velocitypowered/proxy/connection/client/PlaySessionHandler.java
@@ -4,22 +4,46 @@ import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.packets.Ping;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import io.netty.buffer.ByteBuf;
+import io.netty.channel.EventLoop;
+
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
 
 public class PlaySessionHandler implements MinecraftSessionHandler {
     private final ConnectedPlayer player;
+    private ScheduledFuture<?> pingTask;
+    private long lastPing = -1;
 
     public PlaySessionHandler(ConnectedPlayer player) {
         this.player = player;
     }
 
+    @Override
+    public void activated() {
+        EventLoop loop = player.getConnection().getChannel().eventLoop();
+        loop.scheduleAtFixedRate(this::ping, 5, 15, TimeUnit.SECONDS);
+    }
+
+    private void ping() {
+        long randomId = ThreadLocalRandom.current().nextLong();
+        lastPing = randomId;
+        Ping ping = new Ping();
+        ping.setRandomId(randomId);
+        player.getConnection().write(ping);
+    }
+
     @Override
     public void handle(MinecraftPacket packet) {
         if (packet instanceof Ping) {
-            // Handle the ping.
-            player.getConnection().write(packet);
+            Ping ping = (Ping) packet;
+            if (ping.getRandomId() != lastPing) {
+                throw new IllegalStateException("Client sent invalid ping; expected " + lastPing + ", got " + ping.getRandomId());
+            }
+
+            // Do not forward the packet to the player's server, because we handle pings for all servers already.
             return;
         }
-
         // If we don't want to handle this packet, just forward it on.
         player.getConnectedServer().getChannel().write(packet);
     }
@@ -32,5 +56,10 @@ public class PlaySessionHandler implements MinecraftSessionHandler {
     @Override
     public void disconnected() {
         player.getConnectedServer().disconnect();
+
+        if (pingTask != null && !pingTask.isCancelled()) {
+            pingTask.cancel(false);
+            pingTask = null;
+        }
     }
 }