From ec2d75d0f6360c1bed4bba10e2f8291eb3d5bb31 Mon Sep 17 00:00:00 2001
From: ME1312 <admin@me1312.net>
Date: Fri, 3 Jul 2020 19:05:02 -0400
Subject: [PATCH] Modify pings for DNS-Forward servers

---
 .../ME1312/SubServers/Bungee/SubProxy.java    | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/SubServers.Bungee/src/net/ME1312/SubServers/Bungee/SubProxy.java b/SubServers.Bungee/src/net/ME1312/SubServers/Bungee/SubProxy.java
index daf05c62..3fb3c23e 100644
--- a/SubServers.Bungee/src/net/ME1312/SubServers/Bungee/SubProxy.java
+++ b/SubServers.Bungee/src/net/ME1312/SubServers/Bungee/SubProxy.java
@@ -3,6 +3,7 @@ package net.ME1312.SubServers.Bungee;
 import com.dosse.upnp.UPnP;
 import com.google.common.collect.Range;
 import com.google.gson.Gson;
+import net.ME1312.Galaxi.Library.Container.PrimitiveContainer;
 import net.ME1312.Galaxi.Library.Map.ObjectMap;
 import net.ME1312.Galaxi.Library.Container.NamedContainer;
 import net.ME1312.Galaxi.Library.UniversalFile;
@@ -28,6 +29,7 @@ import net.ME1312.Galaxi.Library.Version.Version;
 import net.ME1312.SubServers.Bungee.Network.Packet.PacketOutExReload;
 import net.ME1312.SubServers.Bungee.Network.SubProtocol;
 import net.md_5.bungee.BungeeCord;
+import net.md_5.bungee.BungeeServerInfo;
 import net.md_5.bungee.UserConnection;
 import net.md_5.bungee.api.ChatColor;
 import net.md_5.bungee.api.ServerPing;
@@ -37,6 +39,7 @@ import net.md_5.bungee.api.config.ListenerInfo;
 import net.md_5.bungee.api.config.ServerInfo;
 import net.md_5.bungee.api.event.*;
 import net.md_5.bungee.api.plugin.Listener;
+import net.md_5.bungee.connection.InitialHandler;
 import net.md_5.bungee.event.EventHandler;
 
 import java.io.*;
@@ -843,15 +846,36 @@ public final class SubProxy extends BungeeCord implements Listener {
 
     @EventHandler(priority = Byte.MAX_VALUE)
     public void ping(ProxyPingEvent e) {
-        int offline = 0;
-        for (String name : e.getConnection().getListener().getServerPriority()) {
-            ServerInfo server = api.getServer(name.toLowerCase());
-            if (server == null) server = getServerInfo(name);
-            if (server == null || (server instanceof SubServer && !((SubServer) server).isRunning())) offline++;
-        }
+        boolean forced;
+        ServerInfo override;
+        if ((forced = (override = SmartFallback.getForcedHost(e.getConnection())) != null) || (override = SmartFallback.getDNS(e.getConnection())) != null) {
+            if (override instanceof SubServer && !((SubServer) override).isRunning()) {
+                e.setResponse(new ServerPing(e.getResponse().getVersion(), e.getResponse().getPlayers(), new TextComponent(api.getLang("SubServers", "Bungee.Ping.Offline")), null));
+            } else if (!forced) { // "forced_hosts" already have "ping_passthrough"
+                if (!e.getConnection().getListener().isPingPassthrough()) {
+                    e.setResponse(new ServerPing(e.getResponse().getVersion(), e.getResponse().getPlayers(), new TextComponent(override.getMotd()), null));
+                } else {
+                    PrimitiveContainer<Boolean> lock = new PrimitiveContainer<>(true);
+                    ((BungeeServerInfo) override).ping((ping, error) -> {
+                        if (error != null) {
+                            e.setResponse(new ServerPing(e.getResponse().getVersion(), e.getResponse().getPlayers(), new TextComponent(getTranslation( "ping_cannot_connect" )), null));
+                        } else e.setResponse(ping);
+                        lock.value = false;
+                    }, ((InitialHandler) e.getConnection()).getHandshake().getProtocolVersion());
+                    while (lock.value) Util.isException(() -> Thread.sleep(4));
+                }
+            }
+        } else {
+            int offline = 0;
+            for (String name : e.getConnection().getListener().getServerPriority()) {
+                ServerInfo server = api.getServer(name.toLowerCase());
+                if (server == null) server = getServerInfo(name);
+                if (server == null || (server instanceof SubServer && !((SubServer) server).isRunning())) offline++;
+            }
 
-        if (offline >= e.getConnection().getListener().getServerPriority().size()) {
-            e.setResponse(new ServerPing(e.getResponse().getVersion(), e.getResponse().getPlayers(), new TextComponent(api.getLang("SubServers", "Bungee.Ping.Offline")), null));
+            if (offline >= e.getConnection().getListener().getServerPriority().size()) {
+                e.setResponse(new ServerPing(e.getResponse().getVersion(), e.getResponse().getPlayers(), new TextComponent(api.getLang("SubServers", "Bungee.Ping.Offline")), null));
+            }
         }
     }