From 23554239d0a804368e273eeeef17246888569cdc Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Sat, 23 Sep 2017 13:18:43 +1000
Subject: [PATCH] #2259: Have proxy handle keepalives

---
 .../main/java/net/md_5/bungee/ServerConnection.java |  6 ++++++
 .../main/java/net/md_5/bungee/UserConnection.java   |  3 ---
 .../md_5/bungee/connection/DownstreamBridge.java    |  2 +-
 .../net/md_5/bungee/connection/UpstreamBridge.java  | 13 ++++++++++++-
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnection.java b/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
index b795bf90..1be12dbc 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
@@ -25,6 +25,12 @@ public class ServerConnection implements Server
     private boolean isObsolete;
     @Getter
     private final boolean forgeServer = false;
+    @Getter
+    @Setter
+    private boolean pingFailed;
+    @Getter
+    @Setter
+    private long sentPingId;
 
     private final Unsafe unsafe = new Unsafe()
     {
diff --git a/proxy/src/main/java/net/md_5/bungee/UserConnection.java b/proxy/src/main/java/net/md_5/bungee/UserConnection.java
index 881a079a..93df9747 100644
--- a/proxy/src/main/java/net/md_5/bungee/UserConnection.java
+++ b/proxy/src/main/java/net/md_5/bungee/UserConnection.java
@@ -91,9 +91,6 @@ public final class UserConnection implements ProxiedPlayer
     /*========================================================================*/
     @Getter
     @Setter
-    private long sentPingId;
-    @Getter
-    @Setter
     private long sentPingTime;
     @Getter
     @Setter
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java b/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
index 4049e815..6b8ae245 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/DownstreamBridge.java
@@ -106,7 +106,7 @@ public class DownstreamBridge extends PacketHandler
     @Override
     public void handle(KeepAlive alive) throws Exception
     {
-        con.setSentPingId( alive.getRandomId() );
+        server.setSentPingId( alive.getRandomId() );
         con.setSentPingTime( System.currentTimeMillis() );
     }
 
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
index 167f1f88..29e763c7 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
@@ -116,11 +116,22 @@ public class UpstreamBridge extends PacketHandler
     @Override
     public void handle(KeepAlive alive) throws Exception
     {
-        if ( alive.getRandomId() == con.getSentPingId() )
+        if ( alive.getRandomId() == con.getServer().getSentPingId() )
         {
             int newPing = (int) ( System.currentTimeMillis() - con.getSentPingTime() );
             con.getTabListHandler().onPingChange( newPing );
             con.setPing( newPing );
+        } else
+        {
+            if ( con.getServer().getSentPingId() != 0 && !con.getServer().isPingFailed() )
+            {
+                alive.setRandomId( con.getServer().getSentPingId() );
+                con.getServer().unsafe().sendPacket( alive );
+
+                con.getServer().setPingFailed( true );
+            }
+
+            throw CancelSendSignal.INSTANCE;
         }
     }