Support splitting vanilla packets that are too large on forge<->forge connections (#7485)

Currently only for SUpdateRecipesPacket and STagsListPacket that blow up with huge datapacks

src/main/java/net/minecraftforge/common/ForgeMod.java

@@ -47,6 +47,7 @@ import net.minecraftforge.fml.event.server.FMLServerStoppingEvent;
 import net.minecraftforge.fml.javafmlmod.FMLJavaModLoadingContext;
 import net.minecraftforge.fml.loading.FMLEnvironment;
 import net.minecraftforge.fml.loading.progress.StartupMessageManager;
+import net.minecraftforge.network.VanillaPacketSplitter;
 import net.minecraftforge.registries.DeferredRegister;
 import net.minecraftforge.server.command.EnumArgument;
 import net.minecraftforge.server.command.ModIdArgument;
@@ -171,6 +172,7 @@ public class ForgeMod implements WorldPersistenceHooks.WorldPersistenceHook
         VersionChecker.startVersionCheck();
 
         registerArgumentTypes();
+        VanillaPacketSplitter.register();
     }
 
     @SuppressWarnings({"unchecked", "rawtypes"})

src/main/java/net/minecraftforge/fml/server/ServerLifecycleHooks.java

@@ -68,7 +68,7 @@ import net.minecraftforge.fml.loading.moddiscovery.ModFile;
 import net.minecraftforge.fml.packs.ModFileResourcePack;
 import net.minecraftforge.fml.packs.ResourcePackLoader;
 import net.minecraftforge.forgespi.language.IModInfo;
-import net.minecraftforge.network.VanillaConnectionNetworkFilter;
+import net.minecraftforge.network.NetworkFilters;
 import net.minecraftforge.registries.GameData;
 
 public class ServerLifecycleHooks
@@ -169,7 +169,7 @@ public class ServerLifecycleHooks
         if (packet.getRequestedState() == ProtocolType.STATUS) return true;
 
         NetworkHooks.registerServerLoginChannel(manager, packet);
-        VanillaConnectionNetworkFilter.injectIfNecessary(manager);
+        NetworkFilters.injectIfNecessary(manager);
         return true;
 
     }

src/main/java/net/minecraftforge/network/ForgeConnectionNetworkFilter.java

@@ -0,0 +1,47 @@
+package net.minecraftforge.network;
+
+import java.util.List;
+
+import io.netty.channel.ChannelHandler;
+import net.minecraft.network.IPacket;
+import net.minecraft.network.NetworkManager;
+import net.minecraft.network.PacketDirection;
+import net.minecraft.network.ProtocolType;
+import net.minecraft.network.play.server.STagsListPacket;
+import net.minecraft.network.play.server.SUpdateRecipesPacket;
+import net.minecraftforge.fml.network.NetworkHooks;
+
+import com.google.common.collect.ImmutableMap;
+
+/**
+ * Network filter for forge-forge connections.
+ */
+@ChannelHandler.Sharable
+public class ForgeConnectionNetworkFilter extends VanillaPacketFilter
+{
+
+    public ForgeConnectionNetworkFilter()
+    {
+        super(
+                ImmutableMap.of(
+                        SUpdateRecipesPacket.class, ForgeConnectionNetworkFilter::splitPacket,
+                        STagsListPacket.class, ForgeConnectionNetworkFilter::splitPacket
+                )
+        );
+    }
+
+    @Override
+    protected boolean isNecessary(NetworkManager manager)
+    {
+        // not needed on local connections, because packets are not encoded to bytes there
+        return !manager.isLocalChannel() && !NetworkHooks.isVanillaConnection(manager);
+    }
+
+    private static void splitPacket(IPacket<?> packet, List<? super IPacket<?>> out)
+    {
+        VanillaPacketSplitter.appendPackets(
+                ProtocolType.PLAY, PacketDirection.CLIENTBOUND, packet, out
+        );
+    }
+
+}

src/main/java/net/minecraftforge/network/NetworkFilters.java

@@ -0,0 +1,37 @@
+package net.minecraftforge.network;
+
+import java.util.Map;
+
+import net.minecraft.network.NetworkManager;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import com.google.common.collect.ImmutableMap;
+
+public class NetworkFilters
+{
+
+    private static final Logger LOGGER = LogManager.getLogger();
+
+    private static final Map<String, VanillaPacketFilter> instances = ImmutableMap.of(
+            "forge:vanilla_filter", new VanillaConnectionNetworkFilter(),
+            "forge:forge_fixes", new ForgeConnectionNetworkFilter()
+    );
+
+    public static void injectIfNecessary(NetworkManager manager)
+    {
+        instances.forEach((key, filter) -> {
+            if (filter.isNecessary(manager))
+            {
+                manager.channel().pipeline().addBefore("packet_handler", key, filter);
+                LOGGER.debug("Injected {} into {}", filter, manager);
+            }
+        });
+    }
+
+    private NetworkFilters()
+    {
+    }
+
+}

src/main/java/net/minecraftforge/network/VanillaConnectionNetworkFilter.java

@@ -22,13 +22,11 @@ package net.minecraftforge.network;
 import java.util.AbstractMap;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import java.util.function.Function;
+import java.util.function.BiConsumer;
 
 import javax.annotation.Nonnull;
 
-import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.MessageToMessageEncoder;
+import io.netty.channel.ChannelHandler;
 import net.minecraft.command.ISuggestionProvider;
 import net.minecraft.command.arguments.ArgumentTypes;
 import net.minecraft.network.IPacket;
@@ -49,33 +47,26 @@ import com.mojang.brigadier.tree.RootCommandNode;
  * A filter for network packets, used to filter/modify parts of vanilla network messages that
  * will cause errors or warnings on vanilla clients, for example entity attributes that are added by Forge or mods.
  */
-public class VanillaConnectionNetworkFilter extends MessageToMessageEncoder<IPacket<?>>
+@ChannelHandler.Sharable
+public class VanillaConnectionNetworkFilter extends VanillaPacketFilter
 {
 
     private static final Logger LOGGER = LogManager.getLogger();
 
-    private static final Map<Class<? extends IPacket<?>>, Function<IPacket<?>, ? extends IPacket<?>>> handlers = ImmutableMap.<Class<? extends IPacket<?>>, Function<IPacket<?>, ? extends IPacket<?>>>builder()
-            .put(handler(SEntityPropertiesPacket.class, VanillaConnectionNetworkFilter::filterEntityProperties))
-            .put(handler(SCommandListPacket.class, VanillaConnectionNetworkFilter::filterCommandList))
-            .build();
-
-
-    public static void injectIfNecessary(NetworkManager manager)
+    public VanillaConnectionNetworkFilter()
     {
-        if (NetworkHooks.isVanillaConnection(manager))
-        {
-            manager.channel().pipeline().addBefore("packet_handler", "forge:vanilla_filter", new VanillaConnectionNetworkFilter());
-            LOGGER.debug("Injected into {}", manager);
-        }
+        super(
+                ImmutableMap.<Class<? extends IPacket<?>>, BiConsumer<IPacket<?>, List<? super IPacket<?>>>>builder()
+                .put(handler(SEntityPropertiesPacket.class, VanillaConnectionNetworkFilter::filterEntityProperties))
+                .put(handler(SCommandListPacket.class, VanillaConnectionNetworkFilter::filterCommandList))
+                .build()
+        );
     }
 
-    /**
-     * Helper function for building the handler map.
-     */
-    @Nonnull
-    private static <T extends IPacket<?>> Map.Entry<Class<? extends IPacket<?>>, Function<IPacket<?>, ? extends IPacket<?>>> handler(Class<T> cls, Function<T, ? extends IPacket<?>> function)
+    @Override
+    protected boolean isNecessary(NetworkManager manager)
     {
-        return new AbstractMap.SimpleEntry<>(cls, function.compose(cls::cast));
+        return NetworkHooks.isVanillaConnection(manager);
     }
 
     /**
@@ -109,11 +100,4 @@ public class VanillaConnectionNetworkFilter extends MessageToMessageEncoder<IPac
         });
         return new SCommandListPacket(newRoot);
     }
-
-    @Override
-    protected void encode(ChannelHandlerContext ctx, IPacket<?> msg, List<Object> out)
-    {
-        Function<IPacket<?>, ? extends IPacket<?>> function = handlers.getOrDefault(msg.getClass(), Function.identity());
-        out.add(function.apply(msg));
-    }
 }

src/main/java/net/minecraftforge/network/VanillaPacketFilter.java

@@ -0,0 +1,59 @@
+package net.minecraftforge.network;
+
+import java.util.AbstractMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+
+import javax.annotation.Nonnull;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.MessageToMessageEncoder;
+import net.minecraft.network.IPacket;
+import net.minecraft.network.NetworkManager;
+
+/**
+ * A filter for vanilla network packets.
+ */
+public abstract class VanillaPacketFilter extends MessageToMessageEncoder<IPacket<?>>
+{
+
+    protected final Map<Class<? extends IPacket<?>>, BiConsumer<IPacket<?>, List<? super IPacket<?>>>> handlers;
+
+    protected VanillaPacketFilter(Map<Class<? extends IPacket<?>>, BiConsumer<IPacket<?>, List<? super IPacket<?>>>> handlers)
+    {
+        this.handlers = handlers;
+    }
+
+    /**
+     * Helper function for building the handler map.
+     */
+    @Nonnull
+    protected static <T extends IPacket<?>> Map.Entry<Class<? extends IPacket<?>>, BiConsumer<IPacket<?>, List<? super IPacket<?>>>> handler(Class<T> cls, Function<T, ? extends IPacket<?>> function)
+    {
+        return handler(cls, (pkt, list) -> list.add(function.apply(cls.cast(pkt))));
+    }
+
+    /**
+     * Helper function for building the handler map.
+     */
+    @Nonnull
+    protected static <T extends IPacket<?>> Map.Entry<Class<? extends IPacket<?>>, BiConsumer<IPacket<?>, List<? super IPacket<?>>>> handler(Class<T> cls, BiConsumer<IPacket<?>, List<? super IPacket<?>>> consumer)
+    {
+        return new AbstractMap.SimpleEntry<>(cls, consumer);
+    }
+
+    /**
+     * Whether this filter is necessary on the given connection.
+     */
+    protected abstract boolean isNecessary(NetworkManager manager);
+
+    @Override
+    protected void encode(ChannelHandlerContext ctx, IPacket<?> msg, List<Object> out)
+    {
+        BiConsumer<IPacket<?>, List<? super IPacket<?>>> consumer = handlers.getOrDefault(msg.getClass(), (pkt, list) -> list.add(pkt));
+        consumer.accept(msg, out);
+    }
+
+}

src/main/java/net/minecraftforge/network/VanillaPacketSplitter.java

@@ -0,0 +1,188 @@
+package net.minecraftforge.network;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import net.minecraft.client.Minecraft;
+import net.minecraft.client.network.play.ClientPlayNetHandler;
+import net.minecraft.network.IPacket;
+import net.minecraft.network.PacketBuffer;
+import net.minecraft.network.PacketDirection;
+import net.minecraft.network.ProtocolType;
+import net.minecraft.network.play.server.SCustomPayloadPlayPacket;
+import net.minecraft.network.play.server.STagsListPacket;
+import net.minecraft.network.play.server.SUpdateRecipesPacket;
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.fml.network.NetworkDirection;
+import net.minecraftforge.fml.network.NetworkEvent;
+import net.minecraftforge.fml.network.NetworkRegistry;
+import net.minecraftforge.fml.network.event.EventNetworkChannel;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import com.google.common.collect.ImmutableSet;
+
+/**
+ * A custom payload channel that allows sending vanilla server-to-client packets, even if they would normally
+ * be too large for the vanilla protocol. This is achieved by splitting them into multiple custom payload packets.
+ */
+public class VanillaPacketSplitter
+{
+
+    private static final Logger LOGGER = LogManager.getLogger();
+
+    private static final ResourceLocation CHANNEL = new ResourceLocation("forge", "split");
+    private static final String VERSION = "1.0";
+
+    private static final Set<Class<? extends IPacket<?>>> ALLOWED_PACKETS = ImmutableSet.of(
+            SUpdateRecipesPacket.class,
+            STagsListPacket.class
+    );
+
+    private static final int PROTOCOL_MAX = 2097152;
+
+    private static final int PAYLOAD_TO_CLIENT_MAX = 1048576;
+    // 1 byte for state, 5 byte for VarInt PacketID
+    private static final int PART_SIZE = PAYLOAD_TO_CLIENT_MAX - 1 - 5;
+
+    private static final byte STATE_FIRST = 1;
+    private static final byte STATE_LAST = 2;
+
+    public static void register()
+    {
+        EventNetworkChannel channel = NetworkRegistry.newEventChannel(CHANNEL, () -> VERSION, VERSION::equals, VERSION::equals);
+        channel.addListener(VanillaPacketSplitter::onClientPacket);
+    }
+
+    /**
+     * Append the given packet to the given list. If the packet needs to be split, multiple packets will be appened.
+     * Otherwise only the packet itself.
+     */
+    public static void appendPackets(ProtocolType protocol, PacketDirection direction, IPacket<?> packet, List<? super IPacket<?>> out)
+    {
+        if (heuristicIsDefinitelySmallEnough(packet))
+        {
+            out.add(packet);
+        }
+        else
+        {
+            PacketBuffer buf = new PacketBuffer(Unpooled.buffer());
+            try
+            {
+                packet.writePacketData(buf);
+            }
+            catch (IOException e)
+            {
+                throw new UncheckedIOException(e);
+            }
+            if (buf.readableBytes() <= PROTOCOL_MAX)
+            {
+                buf.release();
+                out.add(packet);
+            }
+            else
+            {
+                int parts = (int)Math.ceil(((double)buf.readableBytes()) / PART_SIZE);
+                if (parts == 1)
+                {
+                    buf.release();
+                    out.add(packet);
+                }
+                else
+                {
+                    for (int part = 0; part < parts; part++)
+                    {
+                        ByteBuf partPrefix;
+                        if (part == 0)
+                        {
+                            partPrefix = Unpooled.buffer(5);
+                            partPrefix.writeByte(STATE_FIRST);
+                            new PacketBuffer(partPrefix).writeVarInt(protocol.getPacketId(direction, packet));
+                        }
+                        else
+                        {
+                            partPrefix = Unpooled.buffer(1);
+                            partPrefix.writeByte(part == parts - 1 ? STATE_LAST : 0);
+                        }
+                        int partSize = Math.min(PART_SIZE, buf.readableBytes());
+                        ByteBuf partBuf = Unpooled.wrappedBuffer(
+                                partPrefix,
+                                buf.retainedSlice(buf.readerIndex(), partSize)
+                        );
+                        buf.skipBytes(partSize);
+                        out.add(new SCustomPayloadPlayPacket(CHANNEL, new PacketBuffer(partBuf)));
+                    }
+                    // we retained all the slices, so we do not need this one anymore
+                    buf.release();
+                }
+            }
+        }
+    }
+
+    private static boolean heuristicIsDefinitelySmallEnough(IPacket<?> packet)
+    {
+        return false;
+    }
+
+    private static final List<PacketBuffer> receivedBuffers = new ArrayList<>();
+
+    @SuppressWarnings("unchecked")
+    private static void onClientPacket(NetworkEvent.ServerCustomPayloadEvent event)
+    {
+        NetworkEvent.Context ctx = event.getSource().get();
+        PacketDirection direction = ctx.getDirection() == NetworkDirection.PLAY_TO_CLIENT ? PacketDirection.CLIENTBOUND : PacketDirection.SERVERBOUND;
+        ProtocolType protocol = ProtocolType.PLAY;
+
+        ctx.setPacketHandled(true);
+
+        PacketBuffer buf = event.getPayload();
+
+        byte state = buf.readByte();
+        if (state == STATE_FIRST)
+        {
+            if (!receivedBuffers.isEmpty())
+            {
+                LOGGER.warn("forge:split received out of order - inbound buffer not empty when receiving first");
+                receivedBuffers.clear();
+            }
+        }
+        buf.retain(); // retain the buffer, it is released after this handler otherwise
+        receivedBuffers.add(buf);
+        if (state == STATE_LAST)
+        {
+            PacketBuffer full = new PacketBuffer(Unpooled.wrappedBuffer(receivedBuffers.toArray(new PacketBuffer[0])));
+            int packetId = full.readVarInt();
+            IPacket<?> packet = protocol.getPacket(direction, packetId);
+            if (packet == null)
+            {
+                LOGGER.error("Received invalid packet ID {} in forge:split", packetId);
+            }
+            else if (!ALLOWED_PACKETS.contains(packet.getClass()))
+            {
+                LOGGER.error("Received not allowed packet type {} in forge:split", packet);
+            }
+            else
+            {
+                try
+                {
+                    packet.readPacketData(full);
+                }
+                catch (IOException e)
+                {
+                    throw new UncheckedIOException(e);
+                } finally {
+                    receivedBuffers.clear();
+                    full.release();
+                }
+                ctx.enqueueWork(() -> ((IPacket<ClientPlayNetHandler>)packet).processPacket(Minecraft.getInstance().getConnection()));
+            }
+        }
+    }
+
+}