Fire PluginMessageEvent correctly if it's sent when a connection is established

Fixes #774

proxy/src/main/java/com/velocitypowered/proxy/connection/client/AuthSessionHandler.java

@@ -193,7 +193,7 @@ public class AuthSessionHandler implements MinecraftSessionHandler {
               return;
             }
 
-            mcConnection.setSessionHandler(new InitialConnectSessionHandler(player));
+            mcConnection.setSessionHandler(new InitialConnectSessionHandler(player, server));
             server.getEventManager().fire(new PostLoginEvent(player))
                 .thenCompose((ignored) -> connectToInitialServer(player))
                 .exceptionally((ex) -> {

proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java

@@ -17,11 +17,18 @@
 
 package com.velocitypowered.proxy.connection.client;
 
+import com.velocitypowered.api.event.connection.PluginMessageEvent;
+import com.velocitypowered.api.proxy.messages.ChannelIdentifier;
+import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import com.velocitypowered.proxy.connection.backend.BungeeCordMessageResponder;
 import com.velocitypowered.proxy.connection.backend.VelocityServerConnection;
 import com.velocitypowered.proxy.protocol.packet.PluginMessage;
 import com.velocitypowered.proxy.protocol.util.PluginMessageUtil;
+import io.netty.buffer.ByteBufUtil;
+import io.netty.buffer.Unpooled;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 /**
  * Handles the play state between exiting the login phase and establishing the first connection
@@ -29,10 +36,15 @@ import com.velocitypowered.proxy.protocol.util.PluginMessageUtil;
  */
 public class InitialConnectSessionHandler implements MinecraftSessionHandler {
 
+  private static final Logger logger = LogManager.getLogger(InitialConnectSessionHandler.class);
+
   private final ConnectedPlayer player;
 
-  InitialConnectSessionHandler(ConnectedPlayer player) {
+  private final VelocityServer server;
+
+  InitialConnectSessionHandler(ConnectedPlayer player, VelocityServer server) {
     this.player = player;
+    this.server = server;
   }
 
   @Override
@@ -45,12 +57,37 @@ public class InitialConnectSessionHandler implements MinecraftSessionHandler {
 
       if (PluginMessageUtil.isRegister(packet)) {
         player.getKnownChannels().addAll(PluginMessageUtil.getChannels(packet));
+        serverConn.ensureConnected().write(packet.retain());
+        return true;
       } else if (PluginMessageUtil.isUnregister(packet)) {
         player.getKnownChannels().removeAll(PluginMessageUtil.getChannels(packet));
+        serverConn.ensureConnected().write(packet.retain());
+        return true;
       } else if (BungeeCordMessageResponder.isBungeeCordMessage(packet)) {
         return true;
       }
-      serverConn.ensureConnected().write(packet.retain());
+
+      ChannelIdentifier id = server.getChannelRegistrar().getFromId(packet.getChannel());
+      if (id == null) {
+        serverConn.ensureConnected().write(packet.retain());
+        return true;
+      }
+
+      byte[] copy = ByteBufUtil.getBytes(packet.content());
+      PluginMessageEvent event = new PluginMessageEvent(serverConn, serverConn.getPlayer(), id,
+          copy);
+      server.getEventManager().fire(event)
+          .thenAcceptAsync(pme -> {
+            if (pme.getResult().isAllowed() && serverConn.isActive()) {
+              PluginMessage copied = new PluginMessage(packet.getChannel(),
+                  Unpooled.wrappedBuffer(copy));
+              serverConn.ensureConnected().write(copied);
+            }
+          }, player.getConnection().eventLoop())
+          .exceptionally((ex) -> {
+            logger.error("Exception while handling plugin message {}", packet, ex);
+            return null;
+          });
     }
     return true;
   }