From 632d0b6ba8a91a2c856b10a83a2a8d80bb58416d Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Wed, 24 Apr 2019 21:28:43 -0400
Subject: [PATCH] Handle plugin channel registrations during initial connect as
 well.

Fixes #193
---
 .../client/ClientPlaySessionHandler.java      |  8 +++-
 .../client/InitialConnectSessionHandler.java  | 40 ++++++++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
index 74b1d6730..030de405d 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/ClientPlaySessionHandler.java
@@ -50,7 +50,7 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 public class ClientPlaySessionHandler implements MinecraftSessionHandler {
 
   private static final Logger logger = LogManager.getLogger(ClientPlaySessionHandler.class);
-  private static final int MAX_PLUGIN_CHANNELS = 1024;
+  static final int MAX_PLUGIN_CHANNELS = 1024;
 
   private final ConnectedPlayer player;
   private boolean spawned = false;
@@ -63,6 +63,12 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
   public ClientPlaySessionHandler(VelocityServer server, ConnectedPlayer player) {
     this.player = player;
     this.server = server;
+
+    if (player.getMinecraftConnection().getSessionHandler()
+        instanceof InitialConnectSessionHandler) {
+      this.knownChannels.addAll(((InitialConnectSessionHandler) player.getMinecraftConnection()
+          .getSessionHandler()).getKnownChannels());
+    }
   }
 
   @Override
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java
index 998e01cc5..33d78cd30 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/InitialConnectSessionHandler.java
@@ -1,22 +1,56 @@
 package com.velocitypowered.proxy.connection.client;
 
+import static com.velocitypowered.proxy.connection.client.ClientPlaySessionHandler.MAX_PLUGIN_CHANNELS;
+
 import com.velocitypowered.proxy.connection.MinecraftSessionHandler;
 import com.velocitypowered.proxy.connection.backend.VelocityServerConnection;
 import com.velocitypowered.proxy.protocol.packet.PluginMessage;
+import com.velocitypowered.proxy.protocol.util.PluginMessageUtil;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 
 public class InitialConnectSessionHandler implements MinecraftSessionHandler {
 
   private final ConnectedPlayer player;
+  private final Set<String> knownChannels;
 
   InitialConnectSessionHandler(ConnectedPlayer player) {
     this.player = player;
+    this.knownChannels = new HashSet<>();
   }
 
   @Override
   public boolean handle(PluginMessage packet) {
     VelocityServerConnection serverConn = player.getConnectionInFlight();
     if (serverConn != null) {
-      if (!player.getPhase().handle(player, packet, serverConn)) {
+      if (player.getPhase().handle(player, packet, serverConn)) {
+        return true;
+      }
+
+      if (PluginMessageUtil.isRegister(packet)) {
+        List<String> actuallyRegistered = new ArrayList<>();
+        List<String> channels = PluginMessageUtil.getChannels(packet);
+        for (String channel : channels) {
+          if (knownChannels.size() >= MAX_PLUGIN_CHANNELS && !knownChannels.contains(channel)) {
+            throw new IllegalStateException("Too many plugin message channels registered");
+          }
+          if (knownChannels.add(channel)) {
+            actuallyRegistered.add(channel);
+          }
+        }
+
+        if (!actuallyRegistered.isEmpty()) {
+          PluginMessage newRegisterPacket = PluginMessageUtil.constructChannelsPacket(serverConn
+              .ensureConnected().getProtocolVersion(), actuallyRegistered);
+          serverConn.ensureConnected().write(newRegisterPacket);
+        }
+      } else if (PluginMessageUtil.isUnregister(packet)) {
+        List<String> channels = PluginMessageUtil.getChannels(packet);
+        knownChannels.removeAll(channels);
+        serverConn.ensureConnected().write(packet);
+      } else {
         serverConn.ensureConnected().write(packet);
       }
     }
@@ -28,4 +62,8 @@ public class InitialConnectSessionHandler implements MinecraftSessionHandler {
     // the user cancelled the login process
     player.teardown();
   }
+
+  Set<String> getKnownChannels() {
+    return knownChannels;
+  }
 }