From f6078e9b74fac3f7fabf5060e4303143b08da59b Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sat, 26 Dec 2020 20:52:12 -0500
Subject: [PATCH] Fix several problems and clean up the BungeeCord plugin
 messaging support.

Fixes #402
---
 .../backend/BungeeCordMessageResponder.java   | 65 ++++++-------------
 .../server/VelocityRegisteredServer.java      |  7 +-
 2 files changed, 26 insertions(+), 46 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BungeeCordMessageResponder.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BungeeCordMessageResponder.java
index 1450564a6..1a0289051 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BungeeCordMessageResponder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/backend/BungeeCordMessageResponder.java
@@ -236,75 +236,52 @@ public class BungeeCordMessageResponder {
     });
   }
 
-  private ByteBuf prepareForwardMessage(ByteBufDataInput in) {
-    String channel = in.readUTF();
-    short messageLength = in.readShort();
-
-    ByteBuf buf = Unpooled.buffer();
-    ByteBufDataOutput forwarded = new ByteBufDataOutput(buf);
-    forwarded.writeUTF(channel);
-    forwarded.writeShort(messageLength);
-    buf.writeBytes(in.unwrap().readSlice(messageLength));
-    return buf;
-  }
-
   private void processForwardToPlayer(ByteBufDataInput in) {
-    proxy.getPlayer(in.readUTF())
-        .ifPresent(foundPlayer -> sendServerResponse((ConnectedPlayer) foundPlayer,
-            prepareForwardMessage(in)));
+    Optional<Player> player = proxy.getPlayer(in.readUTF());
+    if (player.isPresent()) {
+      ByteBuf toForward = in.unwrap().copy();
+      sendServerResponse((ConnectedPlayer) player.get(), toForward);
+    }
   }
 
   private void processForwardToServer(ByteBufDataInput in) {
     String target = in.readUTF();
-    ByteBuf toForward = prepareForwardMessage(in);
+    ByteBuf toForward = in.unwrap().copy();
     if (target.equals("ALL")) {
-      ByteBuf unreleasableForward = Unpooled.unreleasableBuffer(toForward);
       try {
         for (RegisteredServer rs : proxy.getAllServers()) {
-          ((VelocityRegisteredServer) rs).sendPluginMessage(LEGACY_CHANNEL, unreleasableForward);
+          ((VelocityRegisteredServer) rs).sendPluginMessage(LEGACY_CHANNEL,
+              toForward.retainedSlice());
         }
       } finally {
         toForward.release();
       }
     } else {
-      proxy.getServer(target).ifPresent(rs -> ((VelocityRegisteredServer) rs)
-          .sendPluginMessage(LEGACY_CHANNEL, toForward));
+      Optional<RegisteredServer> server = proxy.getServer(target);
+      if (server.isPresent()) {
+        ((VelocityRegisteredServer) server.get()).sendPluginMessage(LEGACY_CHANNEL, toForward);
+      } else {
+        toForward.release();
+      }
     }
   }
 
-  // Note: this method will always release the buffer!
-  private void sendResponseOnConnection(ByteBuf buf) {
-    sendServerResponse(this.player, buf);
-  }
-
   static String getBungeeCordChannel(ProtocolVersion version) {
     return version.compareTo(ProtocolVersion.MINECRAFT_1_13) >= 0 ? MODERN_CHANNEL.getId()
         : LEGACY_CHANNEL.getId();
   }
 
+  // Note: this method will always release the buffer!
+  private void sendResponseOnConnection(ByteBuf buf) {
+    sendServerResponse(this.player, buf);
+  }
+
   // Note: this method will always release the buffer!
   private static void sendServerResponse(ConnectedPlayer player, ByteBuf buf) {
     MinecraftConnection serverConnection = player.ensureAndGetCurrentServer().ensureConnected();
     String chan = getBungeeCordChannel(serverConnection.getProtocolVersion());
-
-    PluginMessage msg = null;
-    boolean released = false;
-
-    try {
-      VelocityServerConnection vsc = player.getConnectedServer();
-      if (vsc == null) {
-        return;
-      }
-
-      MinecraftConnection serverConn = vsc.ensureConnected();
-      msg = new PluginMessage(chan, buf);
-      serverConn.write(msg);
-      released = true;
-    } finally {
-      if (!released && msg != null) {
-        msg.release();
-      }
-    }
+    PluginMessage msg = new PluginMessage(chan, buf);
+    serverConnection.write(msg);
   }
 
   boolean process(PluginMessage message) {
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/server/VelocityRegisteredServer.java b/proxy/src/main/java/com/velocitypowered/proxy/server/VelocityRegisteredServer.java
index 5987311e1..3cf0313d9 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/server/VelocityRegisteredServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/server/VelocityRegisteredServer.java
@@ -125,7 +125,9 @@ public class VelocityRegisteredServer implements RegisteredServer, ForwardingAud
   }
 
   /**
-   * Sends a plugin message to the server through this connection.
+   * Sends a plugin message to the server through this connection. The message will be released
+   * afterwards.
+   *
    * @param identifier the channel ID to use
    * @param data the data
    * @return whether or not the message was sent
@@ -133,11 +135,12 @@ public class VelocityRegisteredServer implements RegisteredServer, ForwardingAud
   public boolean sendPluginMessage(ChannelIdentifier identifier, ByteBuf data) {
     for (ConnectedPlayer player : players.values()) {
       VelocityServerConnection connection = player.getConnectedServer();
-      if (connection != null && connection.getServerInfo().equals(serverInfo)) {
+      if (connection != null && connection.getServer() == this) {
         return connection.sendPluginMessage(identifier, data);
       }
     }
 
+    data.release();
     return false;
   }