From bc6ad16d401b83fc94278faa983cfaa99bde00d5 Mon Sep 17 00:00:00 2001
From: EnZaXD <florian.michael07@gmail.com>
Date: Fri, 1 Nov 2024 18:50:54 +0100
Subject: [PATCH] Properly remove custom_data component if we created it in
 1.20.5+ protocols (#4229)

---
 .../rewriter/BlockItemPacketRewriter1_21.java | 33 ++++++++-----------
 .../BlockItemPacketRewriter1_21_2.java        |  8 ++---
 .../rewriter/StructuredItemRewriter.java      | 13 ++++++++
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
index d95204ab3..ed21ef5ec 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
@@ -136,14 +136,13 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
         super.handleItemToClient(connection, item);
         updateItemData(item);
 
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        if (dataContainer.has(StructuredDataKey.RARITY)) {
+        final StructuredDataContainer data = item.dataContainer();
+        if (data.has(StructuredDataKey.RARITY)) {
             return item;
         }
-
         // Change rarity of trident and piglin banner pattern
         if (item.identifier() == 1188 || item.identifier() == 1200) {
-            dataContainer.set(StructuredDataKey.RARITY, 0); // Common
+            data.set(StructuredDataKey.RARITY, 0); // Common
             saveTag(createCustomTag(item), new ByteTag(true), "rarity");
         }
         return item;
@@ -174,7 +173,16 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
 
         super.handleItemToServer(connection, item);
         downgradeItemData(item);
-        resetRarityValues(item, nbtTagName("rarity"));
+
+        final StructuredDataContainer data = item.dataContainer();
+        final CompoundTag customData = data.get(StructuredDataKey.CUSTOM_DATA);
+        if (customData == null) {
+            return item;
+        }
+        if (customData.remove(nbtTagName("rarity")) != null) {
+            data.remove(StructuredDataKey.RARITY);
+            removeCustomTag(data, customData);
+        }
         return item;
     }
 
@@ -206,21 +214,6 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
         });
     }
 
-    public static void resetRarityValues(final Item item, final String tagName) {
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        final CompoundTag customData = dataContainer.get(StructuredDataKey.CUSTOM_DATA);
-        if (customData == null) {
-            return;
-        }
-
-        if (customData.remove(tagName) != null) {
-            dataContainer.remove(StructuredDataKey.RARITY);
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
-        }
-    }
-
     private int itemToJubeboxSong(final int id) {
         String identifier = Protocol1_20_5To1_21.MAPPINGS.getFullItemMappings().identifier(id);
         if (!identifier.contains("music_disc_")) {
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/BlockItemPacketRewriter1_21_2.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/BlockItemPacketRewriter1_21_2.java
index 52431b954..58319b783 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/BlockItemPacketRewriter1_21_2.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/BlockItemPacketRewriter1_21_2.java
@@ -484,9 +484,7 @@ public final class BlockItemPacketRewriter1_21_2 extends StructuredItemRewriter<
 
         if (customData.remove(nbtTagName("remove_custom_name")) != null) {
             dataContainer.remove(StructuredDataKey.CUSTOM_NAME);
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
+            removeCustomTag(dataContainer, customData);
         }
 
         final IntArrayTag emptyEnchantments = customData.getIntArrayTag(nbtTagName("0_enchants"));
@@ -504,9 +502,7 @@ public final class BlockItemPacketRewriter1_21_2 extends StructuredItemRewriter<
             if (customData.remove(nbtTagName("remove_glint")) != null) {
                 dataContainer.remove(StructuredDataKey.ENCHANTMENT_GLINT_OVERRIDE);
             }
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
+            removeCustomTag(dataContainer, customData);
         }
         return item;
     }
diff --git a/common/src/main/java/com/viaversion/viaversion/rewriter/StructuredItemRewriter.java b/common/src/main/java/com/viaversion/viaversion/rewriter/StructuredItemRewriter.java
index 27da771d2..89a073891 100644
--- a/common/src/main/java/com/viaversion/viaversion/rewriter/StructuredItemRewriter.java
+++ b/common/src/main/java/com/viaversion/viaversion/rewriter/StructuredItemRewriter.java
@@ -41,6 +41,8 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 public class StructuredItemRewriter<C extends ClientboundPacketType, S extends ServerboundPacketType,
     T extends Protocol<C, ?, ?, S>> extends ItemRewriter<C, S, T> {
 
+    public static final String MARKER_KEY = "VV|custom_data";
+
     public StructuredItemRewriter(
         T protocol,
         Type<Item> itemType, Type<Item[]> itemArrayType, Type<Item> mappedItemType, Type<Item[]> mappedItemArrayType,
@@ -205,15 +207,18 @@ public class StructuredItemRewriter<C extends ClientboundPacketType, S extends S
         // Remove custom name
         if (customData.remove(nbtTagName("added_custom_name")) != null) {
             data.remove(StructuredDataKey.CUSTOM_NAME);
+            removeCustomTag(data, customData);
         } else {
             final Tag customName = removeBackupTag(customData, "custom_name");
             if (customName != null) {
                 data.set(StructuredDataKey.CUSTOM_NAME, customName);
+                removeCustomTag(data, customData);
             }
 
             final Tag itemName = removeBackupTag(customData, "item_name");
             if (itemName != null) {
                 data.set(StructuredDataKey.ITEM_NAME, itemName);
+                removeCustomTag(data, customData);
             }
         }
     }
@@ -223,6 +228,7 @@ public class StructuredItemRewriter<C extends ClientboundPacketType, S extends S
         CompoundTag customData = data.get(StructuredDataKey.CUSTOM_DATA);
         if (customData == null) {
             customData = new CompoundTag();
+            customData.putBoolean(MARKER_KEY, true);
             data.set(StructuredDataKey.CUSTOM_DATA, customData);
         }
         return customData;
@@ -239,6 +245,13 @@ public class StructuredItemRewriter<C extends ClientboundPacketType, S extends S
         return customData.remove(nbtTagName(tagName));
     }
 
+    protected void removeCustomTag(final StructuredDataContainer data, final CompoundTag customData) {
+        // Only remove if we initially added it and only the marker is left
+        if (customData.contains(MARKER_KEY) && customData.size() == 1) {
+            data.remove(StructuredDataKey.CUSTOM_DATA);
+        }
+    }
+
     @FunctionalInterface
     private interface ItemHandler {