make proto proxy great again

MrNavaStar · Aug 13, 2024 · 306288a · 306288a
1 parent e74d1f2
commit 306288a
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 22 deletions.
diff --git a/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/ServerSupplier.java b/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/ServerSupplier.java
@@ -1,9 +1,11 @@
 package me.mrnavastar.protoweaver.proxy;
 
+import me.mrnavastar.protoweaver.proxy.api.ProtoServer;
+
 import java.net.SocketAddress;
 import java.util.List;
 
 public interface ServerSupplier {
 
-    List<SocketAddress> getServers();
+    List<ProtoServer> getServers();
 }
diff --git a/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/Velocity.java b/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/Velocity.java
@@ -4,15 +4,19 @@
 import com.velocitypowered.api.event.Subscribe;
 import com.velocitypowered.api.event.proxy.ProxyInitializeEvent;
 import com.velocitypowered.api.event.proxy.ProxyShutdownEvent;
+import com.velocitypowered.api.event.proxy.server.ServerRegisteredEvent;
+import com.velocitypowered.api.event.proxy.server.ServerUnregisteredEvent;
 import com.velocitypowered.api.plugin.Plugin;
 import com.velocitypowered.api.plugin.annotation.DataDirectory;
 import com.velocitypowered.api.proxy.ProxyServer;
+import com.velocitypowered.api.proxy.server.RegisteredServer;
+import com.velocitypowered.api.proxy.server.ServerInfo;
 import me.mrnavastar.protoweaver.core.util.ProtoConstants;
 import me.mrnavastar.protoweaver.core.util.ProtoLogger;
 import me.mrnavastar.protoweaver.proxy.api.ProtoProxy;
+import me.mrnavastar.protoweaver.proxy.api.ProtoServer;
 import org.slf4j.Logger;
 
-import java.net.SocketAddress;
 import java.nio.file.Path;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -26,13 +30,13 @@
 public class Velocity implements ServerSupplier, ProtoLogger.IProtoLogger {
 
     private final Path dir;
-    private final ProxyServer velocity;
+    private final ProxyServer proxy;
     private final Logger logger;
     private ProtoProxy protoProxy;
 
     @Inject
     public Velocity(ProxyServer proxyServer, Logger logger, @DataDirectory Path dir) {
-        this.velocity = proxyServer;
+        this.proxy = proxyServer;
         this.dir = dir;
         this.logger = logger;
         ProtoLogger.setLogger(this);
@@ -43,15 +47,27 @@ public void onProxyInitialize(ProxyInitializeEvent event) {
         protoProxy = new ProtoProxy(this, dir);
     }
 
+    @Subscribe
+    public void onRegister(ServerRegisteredEvent event) {
+        ServerInfo server = event.registeredServer().getServerInfo();
+        protoProxy.register(new ProtoServer(server.getName(), server.getAddress()));
+    }
+
+    @Subscribe
+    public void onUnregister(ServerUnregisteredEvent event) {
+        ServerInfo server = event.unregisteredServer().getServerInfo();
+        protoProxy.unregister(new ProtoServer(server.getName(), server.getAddress()));
+    }
+
     @Subscribe
     public void onProxyShutdown(ProxyShutdownEvent event) {
         protoProxy.shutdown();
     }
 
     @Override
-    public List<SocketAddress> getServers() {
-        return velocity.getAllServers().stream()
-                .map(server -> server.getServerInfo().getAddress())
+    public List<ProtoServer> getServers() {
+        return proxy.getAllServers().stream()
+                .map(server -> new ProtoServer(server.getServerInfo().getName(), server.getServerInfo().getAddress()))
                 .collect(Collectors.toList());
     }
 

diff --git a/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/Waterfall.java b/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/Waterfall.java
@@ -2,10 +2,9 @@
 
 import me.mrnavastar.protoweaver.core.util.ProtoLogger;
 import me.mrnavastar.protoweaver.proxy.api.ProtoProxy;
-import net.md_5.bungee.api.config.ServerInfo;
+import me.mrnavastar.protoweaver.proxy.api.ProtoServer;
 import net.md_5.bungee.api.plugin.Plugin;
 
-import java.net.SocketAddress;
 import java.util.List;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -31,9 +30,9 @@ public void onDisable() {
     }
 
     @Override
-    public List<SocketAddress> getServers() {
+    public List<ProtoServer> getServers() {
         return getProxy().getServersCopy().values().stream()
-                .map(ServerInfo::getSocketAddress)
+                .map(server -> new ProtoServer(server.getName(), server.getSocketAddress()))
                 .collect(Collectors.toList());
     }
 

diff --git a/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/api/ProtoProxy.java b/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/api/ProtoProxy.java
@@ -16,12 +16,14 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 public class ProtoProxy {
 
-    private static final HashSet<ProtoServer> connectedServers = new HashSet<>();
-    private static final ConcurrentHashMap<SocketAddress, ArrayList<ProtoClient>> servers = new ConcurrentHashMap<>();
+    private static final ConcurrentHashMap<ProtoServer, ArrayList<ProtoClient>> servers = new ConcurrentHashMap<>();
 
     /**
      * Sets the polling rate of servers that are disconnected. Defaults to 5 seconds
@@ -31,35 +33,35 @@ public class ProtoProxy {
     private final String hostsFile;
 
     @ApiStatus.Internal
-    public ProtoProxy(ServerSupplier serverSupplier, Path dir) {
+    public ProtoProxy(ServerSupplier supplier, Path dir) {
         this.hostsFile = dir.toAbsolutePath().toString();
-        serverSupplier.getServers().forEach(server -> servers.put(server, new ArrayList<>()));
+        supplier.getServers().forEach(server -> servers.put(server, new ArrayList<>()));
         ProtoWeaver.PROTOCOL_LOADED.register(this::startProtocol);
         ProtoWeaver.getLoadedProtocols().forEach(this::startProtocol);
     }
 
     private void startProtocol(Protocol protocol) {
         if (protocol.toString().equals("protoweaver:internal")) return;
 
-        servers.forEach((address, clients) -> {
+        servers.forEach((server, clients) -> {
             for (ProtoClient client : clients) {
                 // Don't start a new connection if one already exists for this protocol
                 if (client.getCurrentProtocol().toString().equals(protocol.toString())) return;
             }
-            connectClient(protocol, address, clients);
+            connectClient(protocol, server, clients);
         });
     }
 
-    private void connectClient(Protocol protocol, SocketAddress address, ArrayList<ProtoClient> clients) {
-        ProtoClient client = new ProtoClient((InetSocketAddress) address, hostsFile);
+    private void connectClient(Protocol protocol, ProtoServer server, ArrayList<ProtoClient> clients) {
+        ProtoClient client = new ProtoClient((InetSocketAddress) server.getAddress(), hostsFile);
         client.connect(protocol).onConnectionLost(connection -> {
             clients.remove(client);
 
             if (connection.getDisconnecter().equals(Side.CLIENT)) return;
             Thread.sleep(serverPollRate);
-            connectClient(protocol, address, clients);
+            connectClient(protocol, server, clients);
         }).onConnectionEstablished(connection -> {
-            ProtoLogger.info("Connected to: " + address + " with protocol: " + protocol);
+            ProtoLogger.info("Connected to: " + server + " with protocol: " + protocol);
         });
         clients.add(client);
     }
@@ -70,6 +72,25 @@ public void shutdown() {
         servers.clear();
     }
 
+    @ApiStatus.Internal
+    public void register(ProtoServer server) {
+        if (servers.putIfAbsent(server, new ArrayList<>()) == null)
+            ProtoWeaver.getLoadedProtocols().forEach(this::startProtocol);
+    }
+
+    @ApiStatus.Internal
+    public void unregister(ProtoServer server) {
+        Optional.ofNullable(servers.remove(server)).ifPresent(clients -> clients.forEach(ProtoClient::disconnect));
+    }
+
+    public static ArrayList<ProtoServer> getConnectedServers(@NonNull Protocol protocol) {
+        ArrayList<ProtoServer> connected = new ArrayList<>();
+        servers.forEach((server, clients) -> clients.stream()
+                .filter(c -> protocol.equals(c.getCurrentProtocol()) || c.isConnected())
+                .findFirst().ifPresent(c -> connected.add(server)));
+        return connected;
+    }
+
     /**
      * Sends a packet to every server running protoweaver with the correct protocol.
      */

diff --git a/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/api/ProtoServer.java b/proxy/src/main/java/me/mrnavastar/protoweaver/proxy/api/ProtoServer.java
@@ -1,5 +1,26 @@
 package me.mrnavastar.protoweaver.proxy.api;
 
+import lombok.Getter;
+import lombok.Setter;
 import me.mrnavastar.protoweaver.api.netty.ProtoConnection;
 
-public record ProtoServer(String name, ProtoConnection connection) {}
+import java.net.SocketAddress;
+
+@Setter
+@Getter
+public class ProtoServer {
+
+    private final String name;
+    private final SocketAddress address;
+    private ProtoConnection connection;
+
+    public ProtoServer(String name, SocketAddress address) {
+        this.name = name;
+        this.address = address;
+    }
+
+    @Override
+    public String toString() {
+        return name + " : " + address;
+    }
+}