Track plugin message channels that registered by clients (#1276)

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

@@ -311,6 +311,7 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
             + "ready. Channel: {}. Packet discarded.", packet.getChannel());
       } else if (PluginMessageUtil.isRegister(packet)) {
         List<String> channels = PluginMessageUtil.getChannels(packet);
+        player.getClientsideChannels().addAll(channels);
         List<ChannelIdentifier> channelIdentifiers = new ArrayList<>();
         for (String channel : channels) {
           try {
@@ -324,6 +325,7 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
                 new PlayerChannelRegisterEvent(player, ImmutableList.copyOf(channelIdentifiers)));
         backendConn.write(packet.retain());
       } else if (PluginMessageUtil.isUnregister(packet)) {
+        player.getClientsideChannels().removeAll(PluginMessageUtil.getChannels(packet));
         backendConn.write(packet.retain());
       } else if (PluginMessageUtil.isMcBrand(packet)) {
         String brand = PluginMessageUtil.readBrandMessage(packet.content());
@@ -592,6 +594,10 @@ public class ClientPlaySessionHandler implements MinecraftSessionHandler {
     if (!channels.isEmpty()) {
       serverMc.delayedWrite(constructChannelsPacket(serverVersion, channels));
     }
+    // Tell the server about this client's plugin message channels.
+    if (!player.getClientsideChannels().isEmpty()) {
+      serverMc.delayedWrite(constructChannelsPacket(serverVersion, player.getClientsideChannels()));
+    }
 
     // If we had plugin messages queued during login/FML handshake, send them now.
     PluginMessagePacket pm;

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

@@ -99,6 +99,7 @@ import com.velocitypowered.proxy.tablist.VelocityTabListLegacy;
 import com.velocitypowered.proxy.util.ClosestLocaleMatcher;
 import com.velocitypowered.proxy.util.DurationUtils;
 import com.velocitypowered.proxy.util.TranslatableMapper;
+import com.velocitypowered.proxy.util.collect.CappedSet;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import java.net.InetSocketAddress;
@@ -144,6 +145,7 @@ import org.jetbrains.annotations.NotNull;
 public class ConnectedPlayer implements MinecraftConnectionAssociation, Player, KeyIdentifiable,
     VelocityInboundConnection {
 
+  private static final int MAX_CLIENTSIDE_PLUGIN_CHANNELS = 1024;
   private static final PlainTextComponentSerializer PASS_THRU_TRANSLATE =
       PlainTextComponentSerializer.builder().flattener(TranslatableMapper.FLATTENER).build();
   static final PermissionProvider DEFAULT_PERMISSIONS = s -> PermissionFunction.ALWAYS_UNDEFINED;
@@ -173,6 +175,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
   private final InternalTabList tabList;
   private final VelocityServer server;
   private ClientConnectionPhase connectionPhase;
+  private final Collection<String> clientsideChannels;
   private final CompletableFuture<Void> teardownFuture = new CompletableFuture<>();
   private @MonotonicNonNull List<String> serversToTry = null;
   private final ResourcePackHandler resourcePackHandler;
@@ -205,6 +208,7 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
     this.permissionFunction = PermissionFunction.ALWAYS_UNDEFINED;
     this.connectionPhase = connection.getType().getInitialClientPhase();
     this.onlineMode = onlineMode;
+    this.clientsideChannels = CappedSet.create(MAX_CLIENTSIDE_PLUGIN_CHANNELS);
 
     if (connection.getProtocolVersion().noLessThan(ProtocolVersion.MINECRAFT_1_19_3)) {
       this.tabList = new VelocityTabList(this);
@@ -1342,6 +1346,15 @@ public class ConnectedPlayer implements MinecraftConnectionAssociation, Player,
     this.connectionPhase = connectionPhase;
   }
 
+  /**
+   * Return all the plugin message channels that registered by client.
+   *
+   * @return the channels
+   */
+  public Collection<String> getClientsideChannels() {
+    return clientsideChannels;
+  }
+
   @Override
   public @Nullable IdentifiedKey getIdentifiedKey() {
     return playerKey;

proxy/src/main/java/com/velocitypowered/proxy/util/collect/CappedSet.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2019-2023 Velocity Contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.velocitypowered.proxy.util.collect;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ForwardingSet;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * An unsynchronized collection that puts an upper bound on the size of the collection.
+ */
+public final class CappedSet<T> extends ForwardingSet<T> {
+
+  private final Set<T> delegate;
+  private final int upperSize;
+
+  private CappedSet(Set<T> delegate, int upperSize) {
+    this.delegate = delegate;
+    this.upperSize = upperSize;
+  }
+
+  /**
+   * Creates a capped collection backed by a {@link HashSet}.
+   *
+   * @param maxSize the maximum size of the collection
+   * @param <T>     the type of elements in the collection
+   * @return the new collection
+   */
+  public static <T> Set<T> create(int maxSize) {
+    return new CappedSet<>(new HashSet<>(), maxSize);
+  }
+
+  @Override
+  protected Set<T> delegate() {
+    return delegate;
+  }
+
+  @Override
+  public boolean add(T element) {
+    if (this.delegate.size() >= upperSize) {
+      Preconditions.checkState(this.delegate.contains(element),
+          "collection is too large (%s >= %s)",
+          this.delegate.size(), this.upperSize);
+      return false;
+    }
+    return this.delegate.add(element);
+  }
+
+  @Override
+  public boolean addAll(Collection<? extends T> collection) {
+    return this.standardAddAll(collection);
+  }
+}

proxy/src/test/java/com/velocitypowered/proxy/util/collect/CappedSetTest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2019-2021 Velocity Contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.velocitypowered.proxy.util.collect;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.google.common.collect.ImmutableSet;
+import java.util.Collection;
+import java.util.Set;
+import org.junit.jupiter.api.Test;
+
+class CappedSetTest {
+
+  @Test
+  void basicVerification() {
+    Collection<String> coll = CappedSet.create(1);
+    assertTrue(coll.add("coffee"), "did not add single item");
+    assertThrows(IllegalStateException.class, () -> coll.add("tea"),
+        "item was added to collection although it is too full");
+    assertEquals(1, coll.size(), "collection grew in size unexpectedly");
+  }
+
+  @Test
+  void testAddAll() {
+    Set<String> doesFill1 = ImmutableSet.of("coffee", "tea");
+    Set<String> doesFill2 = ImmutableSet.of("chocolate");
+    Set<String> overfill = ImmutableSet.of("Coke", "Pepsi");
+
+    Collection<String> coll = CappedSet.create(3);
+    assertTrue(coll.addAll(doesFill1), "did not add items");
+    assertTrue(coll.addAll(doesFill2), "did not add items");
+    assertThrows(IllegalStateException.class, () -> coll.addAll(overfill),
+        "items added to collection although it is too full");
+    assertEquals(3, coll.size(), "collection grew in size unexpectedly");
+  }
+
+  @Test
+  void handlesSetBehaviorCorrectly() {
+    Set<String> doesFill1 = ImmutableSet.of("coffee", "tea");
+    Set<String> doesFill2 = ImmutableSet.of("coffee", "chocolate");
+    Set<String> overfill = ImmutableSet.of("coffee", "Coke", "Pepsi");
+
+    Collection<String> coll = CappedSet.create(3);
+    assertTrue(coll.addAll(doesFill1), "did not add items");
+    assertTrue(coll.addAll(doesFill2), "did not add items");
+    assertThrows(IllegalStateException.class, () -> coll.addAll(overfill),
+        "items added to collection although it is too full");
+
+    assertFalse(coll.addAll(doesFill1), "added items?!?");
+
+    assertEquals(3, coll.size(), "collection grew in size unexpectedly");
+  }
+}