Paper/patches/server/0954-Prevent-sending-oversized-item-data-in-equipment-and.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Wed, 1 Dec 2021 12:36:25 +0100
Subject: [PATCH] Prevent sending oversized item data in equipment and metadata

Co-authored-by: Jake Potrebic <jake.m.potrebic@gmail.com>

diff --git a/src/main/java/io/papermc/paper/util/DataSanitizationUtil.java b/src/main/java/io/papermc/paper/util/DataSanitizationUtil.java
new file mode 100644
index 0000000000000000000000000000000000000000..72483dedd3b1864fca3463d3ba3f6fad22680b97
--- /dev/null
+++ b/src/main/java/io/papermc/paper/util/DataSanitizationUtil.java
@@ -0,0 +1,108 @@
+package io.papermc.paper.util;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.UnaryOperator;
+import it.unimi.dsi.fastutil.objects.ObjectArrayList;
+import net.minecraft.network.RegistryFriendlyByteBuf;
+import net.minecraft.network.codec.StreamCodec;
+import net.minecraft.util.Mth;
+import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.Items;
+import net.minecraft.world.item.component.BundleContents;
+import net.minecraft.world.item.component.ChargedProjectiles;
+import net.minecraft.world.item.component.ItemContainerContents;
+import org.apache.commons.lang3.math.Fraction;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.framework.qual.DefaultQualifier;
+
+@DefaultQualifier(NonNull.class)
+public final class DataSanitizationUtil {
+
+    private static final ThreadLocal<DataSanitizer> DATA_SANITIZER = ThreadLocal.withInitial(DataSanitizer::new);
+
+    public static DataSanitizer start(final boolean sanitize) {
+        final DataSanitizer sanitizer = DATA_SANITIZER.get();
+        if (sanitize) {
+            sanitizer.start();
+        }
+        return sanitizer;
+    }
+
+    public static final StreamCodec<RegistryFriendlyByteBuf, ChargedProjectiles> CHARGED_PROJECTILES = codec(ChargedProjectiles.STREAM_CODEC, DataSanitizationUtil::sanitizeChargedProjectiles);
+    public static final StreamCodec<RegistryFriendlyByteBuf, BundleContents> BUNDLE_CONTENTS = codec(BundleContents.STREAM_CODEC, DataSanitizationUtil::sanitizeBundleContents);
+    public static final StreamCodec<RegistryFriendlyByteBuf, ItemContainerContents> CONTAINER = codec(ItemContainerContents.STREAM_CODEC, contents -> ItemContainerContents.EMPTY);
+
+    private static ChargedProjectiles sanitizeChargedProjectiles(final ChargedProjectiles projectiles) {
+        if (projectiles.isEmpty()) {
+            return projectiles;
+        }
+
+        return ChargedProjectiles.of(List.of(
+            new ItemStack(projectiles.contains(Items.FIREWORK_ROCKET) ? Items.FIREWORK_ROCKET : Items.ARROW)
+        ));
+    }
+
+    private static BundleContents sanitizeBundleContents(final BundleContents contents) {
+        if (contents.isEmpty()) {
+            return contents;
+        }
+
+        // Bundles change their texture based on their fullness.
+        // A bundles content weight may be anywhere from 0 to, basically, infinity.
+        // A weight of 1 is the usual maximum case
+        int sizeUsed = Mth.mulAndTruncate(contents.weight(), 64);
+        // Early out, *most* bundles should not be overfilled above a weight of one.
+        if (sizeUsed <= 64) return new BundleContents(List.of(new ItemStack(Items.PAPER, Math.max(1, sizeUsed))));
+
+        final List<ItemStack> sanitizedRepresentation = new ObjectArrayList<>(sizeUsed / 64 + 1);
+        while (sizeUsed > 0) {
+            final int stackCount = Math.min(64, sizeUsed);
+            sanitizedRepresentation.add(new ItemStack(Items.PAPER, stackCount));
+            sizeUsed -= stackCount;
+        }
+        // Now we add a single fake item that uses the same amount of slots as all other items.
+        // Ensure that potentially overstacked bundles are not represented by empty (count=0) itemstacks.
+        return new BundleContents(sanitizedRepresentation);
+    }
+
+    private static <B, A> StreamCodec<B, A> codec(final StreamCodec<B, A> delegate, final UnaryOperator<A> sanitizer) {
+        return new DataSanitizationCodec<>(delegate, sanitizer);
+    }
+
+    private record DataSanitizationCodec<B, A>(StreamCodec<B, A> delegate, UnaryOperator<A> sanitizer) implements StreamCodec<B, A> {
+
+        @Override
+        public @NonNull A decode(final @NonNull B buf) {
+            return this.delegate.decode(buf);
+        }
+
+        @Override
+        public void encode(final @NonNull B buf, final @NonNull A value) {
+            if (!DATA_SANITIZER.get().value().get()) {
+                this.delegate.encode(buf, value);
+            } else {
+                this.delegate.encode(buf, this.sanitizer.apply(value));
+            }
+        }
+    }
+
+    public record DataSanitizer(AtomicBoolean value) implements AutoCloseable {
+
+        public DataSanitizer() {
+            this(new AtomicBoolean(false));
+        }
+
+        public void start() {
+            this.value.compareAndSet(false, true);
+        }
+
+        @Override
+        public void close() {
+            this.value.compareAndSet(true, false);
+        }
+    }
+
+    private DataSanitizationUtil() {
+    }
+}
diff --git a/src/main/java/net/minecraft/core/component/DataComponents.java b/src/main/java/net/minecraft/core/component/DataComponents.java
index a7f6a53cd1168bfc53b616cb0586a546fb9793e7..5e3bfd4a06b4a547a5f813f3d0c533fb7f4451f3 100644
--- a/src/main/java/net/minecraft/core/component/DataComponents.java
+++ b/src/main/java/net/minecraft/core/component/DataComponents.java
@@ -180,10 +180,10 @@ public class DataComponents {
         "map_post_processing", builder -> builder.networkSynchronized(MapPostProcessing.STREAM_CODEC)
     );
     public static final DataComponentType<ChargedProjectiles> CHARGED_PROJECTILES = register(
-        "charged_projectiles", builder -> builder.persistent(ChargedProjectiles.CODEC).networkSynchronized(ChargedProjectiles.STREAM_CODEC).cacheEncoding()
+        "charged_projectiles", builder -> builder.persistent(ChargedProjectiles.CODEC).networkSynchronized(io.papermc.paper.util.DataSanitizationUtil.CHARGED_PROJECTILES).cacheEncoding() // Paper - sanitize charged projectiles
     );
     public static final DataComponentType<BundleContents> BUNDLE_CONTENTS = register(
-        "bundle_contents", builder -> builder.persistent(BundleContents.CODEC).networkSynchronized(BundleContents.STREAM_CODEC).cacheEncoding()
+        "bundle_contents", builder -> builder.persistent(BundleContents.CODEC).networkSynchronized(io.papermc.paper.util.DataSanitizationUtil.BUNDLE_CONTENTS).cacheEncoding() // Paper - sanitize bundle contents
     );
     public static final DataComponentType<PotionContents> POTION_CONTENTS = register(
         "potion_contents", builder -> builder.persistent(PotionContents.CODEC).networkSynchronized(PotionContents.STREAM_CODEC).cacheEncoding()
@@ -250,7 +250,7 @@ public class DataComponents {
         "pot_decorations", builder -> builder.persistent(PotDecorations.CODEC).networkSynchronized(PotDecorations.STREAM_CODEC).cacheEncoding()
     );
     public static final DataComponentType<ItemContainerContents> CONTAINER = register(
-        "container", builder -> builder.persistent(ItemContainerContents.CODEC).networkSynchronized(ItemContainerContents.STREAM_CODEC).cacheEncoding()
+        "container", builder -> builder.persistent(ItemContainerContents.CODEC).networkSynchronized(io.papermc.paper.util.DataSanitizationUtil.CONTAINER).cacheEncoding() // Paper - sanitize container contents
     );
     public static final DataComponentType<BlockItemStateProperties> BLOCK_STATE = register(
         "block_state", builder -> builder.persistent(BlockItemStateProperties.CODEC).networkSynchronized(BlockItemStateProperties.STREAM_CODEC).cacheEncoding()
diff --git a/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEntityDataPacket.java b/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEntityDataPacket.java
index 59c1c103545f04fd35e6932df64a9910a1d74cd7..56bde49e6b7790155b032d0be40961d566ab89e9 100644
--- a/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEntityDataPacket.java
+++ b/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEntityDataPacket.java
@@ -19,9 +19,11 @@ public record ClientboundSetEntityDataPacket(int id, List<SynchedEntityData.Data
     }
 
     private static void pack(List<SynchedEntityData.DataValue<?>> trackedValues, RegistryFriendlyByteBuf buf) {
+        try (var ignored = io.papermc.paper.util.DataSanitizationUtil.start(true)) { // Paper - data sanitization
         for (SynchedEntityData.DataValue<?> dataValue : trackedValues) {
             dataValue.write(buf);
         }
+        } // Paper - data sanitization
 
         buf.writeByte(255);
     }
diff --git a/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEquipmentPacket.java b/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEquipmentPacket.java
index e46030178b3b54168a3532def308a08b10888723..830bd76916e26a3a54954d3cf7b7520af52a2258 100644
--- a/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEquipmentPacket.java
+++ b/src/main/java/net/minecraft/network/protocol/game/ClientboundSetEquipmentPacket.java
@@ -19,6 +19,13 @@ public class ClientboundSetEquipmentPacket implements Packet<ClientGamePacketLis
     private final List<Pair<EquipmentSlot, ItemStack>> slots;
 
     public ClientboundSetEquipmentPacket(int entityId, List<Pair<EquipmentSlot, ItemStack>> equipmentList) {
+        // Paper start - data sanitization
+        this(entityId, equipmentList, false);
+    }
+    private boolean sanitize;
+    public ClientboundSetEquipmentPacket(int entityId, List<Pair<EquipmentSlot, ItemStack>> equipmentList, boolean sanitize) {
+        this.sanitize = sanitize;
+        // Paper end - data sanitization
         this.entity = entityId;
         this.slots = equipmentList;
     }
@@ -40,6 +47,7 @@ public class ClientboundSetEquipmentPacket implements Packet<ClientGamePacketLis
         buf.writeVarInt(this.entity);
         int i = this.slots.size();
 
+        try (var ignored = io.papermc.paper.util.DataSanitizationUtil.start(this.sanitize)) {  // Paper - data sanitization
         for (int j = 0; j < i; j++) {
             Pair<EquipmentSlot, ItemStack> pair = this.slots.get(j);
             EquipmentSlot equipmentSlot = pair.getFirst();
@@ -48,6 +56,7 @@ public class ClientboundSetEquipmentPacket implements Packet<ClientGamePacketLis
             buf.writeByte(bl ? k | -128 : k);
             ItemStack.OPTIONAL_STREAM_CODEC.encode(buf, pair.getSecond());
         }
+        } // Paper - data sanitization
     }
 
     @Override
diff --git a/src/main/java/net/minecraft/server/level/ServerEntity.java b/src/main/java/net/minecraft/server/level/ServerEntity.java
index ce3e359ed2b29674aa89a714ee1fb0875d83595d..d5bc702f2676b1b7a32c8f3a4a349fc2710ee825 100644
--- a/src/main/java/net/minecraft/server/level/ServerEntity.java
+++ b/src/main/java/net/minecraft/server/level/ServerEntity.java
@@ -388,7 +388,7 @@ public class ServerEntity {
             }
 
             if (!list.isEmpty()) {
-                sender.accept(new ClientboundSetEquipmentPacket(this.entity.getId(), list));
+                sender.accept(new ClientboundSetEquipmentPacket(this.entity.getId(), list, true)); // Paper - data sanitization
             }
             ((LivingEntity) this.entity).detectEquipmentUpdatesPublic(); // CraftBukkit - SPIGOT-3789: sync again immediately after sending
         }
diff --git a/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java b/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java
index a5ff6d2d4648f3ed4860b8812c6af48cdd6884fa..ed4b69bc1875f1d90b02166d62078137c41e07a5 100644
--- a/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java
+++ b/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java
@@ -2818,7 +2818,7 @@ public class ServerGamePacketListenerImpl extends ServerCommonPacketListenerImpl
                                     entity.refreshEntityData(ServerGamePacketListenerImpl.this.player);
                                     // SPIGOT-7136 - Allays
                                 if (entity instanceof Allay || entity instanceof net.minecraft.world.entity.animal.horse.AbstractHorse) { // Paper - Fix horse armor desync
-                                        ServerGamePacketListenerImpl.this.send(new ClientboundSetEquipmentPacket(entity.getId(), Arrays.stream(net.minecraft.world.entity.EquipmentSlot.values()).map((slot) -> Pair.of(slot, ((LivingEntity) entity).getItemBySlot(slot).copy())).collect(Collectors.toList())));
+                                    ServerGamePacketListenerImpl.this.send(new ClientboundSetEquipmentPacket(entity.getId(), Arrays.stream(net.minecraft.world.entity.EquipmentSlot.values()).map((slot) -> Pair.of(slot, ((LivingEntity) entity).getItemBySlot(slot).copy())).collect(Collectors.toList()), true)); // Paper - sanitize
                                     }
 
                                 ServerGamePacketListenerImpl.this.player.containerMenu.sendAllDataToRemote(); // Paper - fix slot desync - always refresh player inventory
diff --git a/src/main/java/net/minecraft/world/entity/LivingEntity.java b/src/main/java/net/minecraft/world/entity/LivingEntity.java
index 1bb38e6c6a02839fbe72584e91504dfaca3c46e9..696a75ef6695896358ffcbef44a3a59eb97544fb 100644
--- a/src/main/java/net/minecraft/world/entity/LivingEntity.java
+++ b/src/main/java/net/minecraft/world/entity/LivingEntity.java
@@ -3472,7 +3472,7 @@ public abstract class LivingEntity extends Entity implements Attackable {
             }
 
         });
-        ((ServerLevel) this.level()).getChunkSource().broadcast(this, new ClientboundSetEquipmentPacket(this.getId(), list));
+        ((ServerLevel) this.level()).getChunkSource().broadcast(this, new ClientboundSetEquipmentPacket(this.getId(), list, true)); // Paper - data sanitization
     }
 
     private ItemStack getLastArmorItem(EquipmentSlot slot) {