readd itemstack data sanitization (#10454)

* readd itemstack data sanitization

* use autocloseable

2 files changed, 228 insertions, 88 deletions

diff --git a/patches/server/1048-Prevent-sending-oversized-item-data-in-equipment-and.patch b/patches/server/1048-Prevent-sending-oversized-item-data-in-equipment-and.patch
new file mode 100644
index 0000000000..a395e17562
--- /dev/null
+++ b/patches/server/1048-Prevent-sending-oversized-item-data-in-equipment-and.patch
@@ -0,0 +1,228 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: Nassim Jahnke <[email protected]>
+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 <[email protected]>
+
+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..e9436f8a73ee0a02096d66e14d73edaae28d5a41
+--- /dev/null
++++ b/src/main/java/io/papermc/paper/util/DataSanitizationUtil.java
+@@ -0,0 +1,100 @@
++package io.papermc.paper.util;
++
++import java.util.ArrayList;
++import java.util.List;
++import java.util.concurrent.atomic.AtomicBoolean;
++import java.util.function.UnaryOperator;
++import net.minecraft.network.RegistryFriendlyByteBuf;
++import net.minecraft.network.codec.StreamCodec;
++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.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;
++        }
++        final List<ItemStack> items = projectiles.getItems();
++        final List<ItemStack> sanitized = new ArrayList<>();
++        for (int i = 0; i < Math.min(items.size(), 3); i++) {
++            // we want to preserve item type as vanilla client can change visuals based on type
++            sanitized.add(new ItemStack(items.get(i).getItemHolder()));
++        }
++        return ChargedProjectiles.of(sanitized);
++    }
++
++    private static BundleContents sanitizeBundleContents(final BundleContents contents) {
++        // Bundles change their texture based on their fullness.
++        int sizeUsed = 0;
++        for (final ItemStack item : contents.items()) {
++            final int scale = 64 / item.getMaxStackSize();
++            sizeUsed += scale * item.getCount();
++        }
++        // Now we add a single fake item that uses the same amount of slots as all other items.
++        final List<ItemStack> items = new ArrayList<>();
++        items.add(new ItemStack(Items.PAPER, sizeUsed));
++        return new BundleContents(items);
++    }
++
++    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 5632974af9c603d333ffc30a5a1b1e851821a3bb..9b2a209cda955ef3e5d8ff3ed1b2249888c7d139 100644
+--- a/src/main/java/net/minecraft/core/component/DataComponents.java
++++ b/src/main/java/net/minecraft/core/component/DataComponents.java
+@@ -139,10 +139,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()
+@@ -206,7 +206,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 e092a486c4041ab1cfe9e29c88d0d94528a6e9a6..3945ca04ede578121b370592482ac917f2d4cf96 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 id, List<Pair<EquipmentSlot, ItemStack>> equipmentList) {
++        // Paper start - data sanitization
++        this(id, equipmentList, false);
++    }
++    private boolean sanitize = false;
++    public ClientboundSetEquipmentPacket(int id, List<Pair<EquipmentSlot, ItemStack>> equipmentList, boolean sanitize) {
++        this.sanitize = sanitize;
++        // Paper end - data sanitization
+         this.entity = id;
+         this.slots = equipmentList;
+     }
+@@ -41,6 +48,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();
+@@ -49,6 +57,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 f16a69775332a08ed0e87d27acd0fc959359694c..a2279262c93408c11f5d2290b48fd794975e8cfe 100644
+--- a/src/main/java/net/minecraft/server/level/ServerEntity.java
++++ b/src/main/java/net/minecraft/server/level/ServerEntity.java
+@@ -361,7 +361,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 b9b3277c8ed94e0cd30b20b9c00a33eaad48e5ac..c450447585af4c8cdc87abe871c229ff895c3e53 100644
+--- a/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java
++++ b/src/main/java/net/minecraft/server/network/ServerGamePacketListenerImpl.java
+@@ -2788,7 +2788,7 @@ public class ServerGamePacketListenerImpl extends ServerCommonPacketListenerImpl
+                                 entity.refreshEntityData(ServerGamePacketListenerImpl.this.player);
+                                 // SPIGOT-7136 - Allays
+                                 if (entity instanceof Allay) {
+-                                    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();
+                                 }
+                             }
+diff --git a/src/main/java/net/minecraft/world/entity/LivingEntity.java b/src/main/java/net/minecraft/world/entity/LivingEntity.java
+index ef0f118aecf0893e45cb9423a677d7e42496324b..2c2aefa1d57a939c3ac6c3d6f3a22de848cabf14 100644
+--- a/src/main/java/net/minecraft/world/entity/LivingEntity.java
++++ b/src/main/java/net/minecraft/world/entity/LivingEntity.java
+@@ -3318,7 +3318,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) {
diff --git a/removed-patches-1-20-5/0675-Prevent-sending-oversized-item-data-in-equipment-and.patch b/removed-patches-1-20-5/0675-Prevent-sending-oversized-item-data-in-equipment-and.patch
deleted file mode 100644
index 5e19c6e46b..0000000000
--- a/removed-patches-1-20-5/0675-Prevent-sending-oversized-item-data-in-equipment-and.patch
+++ /dev/null
@@ -1,88 +0,0 @@
-From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
-From: Nassim Jahnke <[email protected]>
-Date: Wed, 1 Dec 2021 12:36:25 +0100
-Subject: [PATCH] Prevent sending oversized item data in equipment and metadata
-
-TODO: Check if still needed with compacted items over the network and limits
-
-
-diff --git a/src/main/java/net/minecraft/network/syncher/EntityDataSerializers.java b/src/main/java/net/minecraft/network/syncher/EntityDataSerializers.java
-index 3ee4c01e4505241e575b8b2e96338ba27b793a2b..d70f77e4d5ca59c9f802ecb354fa8d53d1e10134 100644
---- a/src/main/java/net/minecraft/network/syncher/EntityDataSerializers.java
-+++ b/src/main/java/net/minecraft/network/syncher/EntityDataSerializers.java
-@@ -44,7 +44,7 @@ public class EntityDataSerializers {
-     public static final EntityDataSerializer<ItemStack> ITEM_STACK = new EntityDataSerializer<ItemStack>() {
-         @Override
-         public void write(FriendlyByteBuf buf, ItemStack value) {
--            buf.writeItem(value);
-+            buf.writeItem(net.minecraft.world.entity.LivingEntity.sanitizeItemStack(value, true)); // Paper - prevent oversized data
-         }
- 
-         @Override
-diff --git a/src/main/java/net/minecraft/server/level/ServerEntity.java b/src/main/java/net/minecraft/server/level/ServerEntity.java
-index 91563f903834c1a0636dc087f8c6376815165b6c..a51564e8dce3c125ed5f05cc23548a05c1e79a95 100644
---- a/src/main/java/net/minecraft/server/level/ServerEntity.java
-+++ b/src/main/java/net/minecraft/server/level/ServerEntity.java
-@@ -335,7 +335,10 @@ public class ServerEntity {
-                 ItemStack itemstack = ((LivingEntity) this.entity).getItemBySlot(enumitemslot);
- 
-                 if (!itemstack.isEmpty()) {
--                    list.add(Pair.of(enumitemslot, itemstack.copy()));
-+                    // Paper start - prevent oversized data
-+                    final ItemStack sanitized = LivingEntity.sanitizeItemStack(itemstack.copy(), false);
-+                    list.add(Pair.of(enumitemslot, sanitized));
-+                    // Paper end - prevent oversized data
-                 }
-             }
- 
-diff --git a/src/main/java/net/minecraft/world/entity/LivingEntity.java b/src/main/java/net/minecraft/world/entity/LivingEntity.java
-index ed6fd20ff608e764d6b0f517f6c9c85c533f1646..8025e351fb3e24aa67b31eca74be1bc368592851 100644
---- a/src/main/java/net/minecraft/world/entity/LivingEntity.java
-+++ b/src/main/java/net/minecraft/world/entity/LivingEntity.java
-@@ -3208,7 +3208,10 @@ public abstract class LivingEntity extends Entity implements Attackable {
-         equipmentChanges.forEach((enumitemslot, itemstack) -> {
-             ItemStack itemstack1 = itemstack.copy();
- 
--            list.add(Pair.of(enumitemslot, itemstack1));
-+            // Paper start - prevent oversized data
-+            ItemStack toSend = sanitizeItemStack(itemstack1, true);
-+            list.add(Pair.of(enumitemslot, toSend));
-+            // Paper end - prevent oversized data
-             switch (enumitemslot.getType()) {
-                 case HAND:
-                     this.setLastHandItem(enumitemslot, itemstack1);
-@@ -3221,6 +3224,34 @@ public abstract class LivingEntity extends Entity implements Attackable {
-         ((ServerLevel) this.level()).getChunkSource().broadcast(this, new ClientboundSetEquipmentPacket(this.getId(), list));
-     }
- 
-+    // Paper start - prevent oversized data
-+    public static ItemStack sanitizeItemStack(final ItemStack itemStack, final boolean copyItemStack) {
-+        if (itemStack.isEmpty() || !itemStack.hasTag()) {
-+            return itemStack;
-+        }
-+
-+        final ItemStack copy = copyItemStack ? itemStack.copy() : itemStack;
-+        final CompoundTag tag = copy.getTag();
-+        if (copy.is(Items.BUNDLE) && tag.get("Items") instanceof ListTag oldItems && !oldItems.isEmpty()) {
-+            // Bundles change their texture based on their fullness.
-+            org.bukkit.inventory.meta.BundleMeta bundleMeta = (org.bukkit.inventory.meta.BundleMeta) copy.asBukkitMirror().getItemMeta();
-+            int sizeUsed = 0;
-+            for (org.bukkit.inventory.ItemStack item : bundleMeta.getItems()) {
-+                int scale = 64 / item.getMaxStackSize();
-+                sizeUsed += scale * item.getAmount();
-+            }
-+            // Now we add a single fake item that uses the same amount of slots as all other items.
-+            ListTag items = new ListTag();
-+            items.add(new ItemStack(Items.PAPER, sizeUsed).save(new CompoundTag()));
-+            tag.put("Items", items);
-+        }
-+        if (tag.get("BlockEntityTag") instanceof CompoundTag blockEntityTag) {
-+            blockEntityTag.remove("Items");
-+        }
-+        return copy;
-+    }
-+    // Paper end - prevent oversized data
-+
-     private ItemStack getLastArmorItem(EquipmentSlot slot) {
-         return (ItemStack) this.lastArmorItemStacks.get(slot.getIndex());
-     }