From 616c088b66830c8640a2b6e9f39f36af9ef14eb5 Mon Sep 17 00:00:00 2001
From: Camotoy <20743703+Camotoy@users.noreply.github.com>
Date: Wed, 20 Jul 2022 18:35:40 -0400
Subject: [PATCH] Fix custom items with ItemTranslator#getBedrockItemMapping

---
 .../geysermc/geyser/inventory/Inventory.java    |  7 +++----
 .../inventory/item/ItemTranslator.java          | 17 ++++++++++++-----
 .../BedrockInventoryTransactionTranslator.java  |  2 +-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/inventory/Inventory.java b/core/src/main/java/org/geysermc/geyser/inventory/Inventory.java
index ca7e90a25..58d551489 100644
--- a/core/src/main/java/org/geysermc/geyser/inventory/Inventory.java
+++ b/core/src/main/java/org/geysermc/geyser/inventory/Inventory.java
@@ -34,7 +34,6 @@ import lombok.Getter;
 import lombok.Setter;
 import lombok.ToString;
 import org.geysermc.geyser.GeyserImpl;
-import org.geysermc.geyser.registry.type.ItemMapping;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.inventory.item.ItemTranslator;
 import org.jetbrains.annotations.Range;
@@ -136,9 +135,9 @@ public abstract class Inventory {
 
     protected void updateItemNetId(GeyserItemStack oldItem, GeyserItemStack newItem, GeyserSession session) {
         if (!newItem.isEmpty()) {
-            ItemMapping oldMapping = ItemTranslator.getBedrockItemMapping(session, oldItem);
-            ItemMapping newMapping = ItemTranslator.getBedrockItemMapping(session, newItem);
-            if (oldMapping.getBedrockId() == newMapping.getBedrockId()) {
+            int oldMapping = ItemTranslator.getBedrockItemId(session, oldItem);
+            int newMapping = ItemTranslator.getBedrockItemId(session, newItem);
+            if (oldMapping == newMapping) {
                 newItem.setNetId(oldItem.getNetId());
             } else {
                 newItem.setNetId(session.getNextItemNetId());
diff --git a/core/src/main/java/org/geysermc/geyser/translator/inventory/item/ItemTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/inventory/item/ItemTranslator.java
index 0a2ab57df..b36833cb1 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/inventory/item/ItemTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/inventory/item/ItemTranslator.java
@@ -266,16 +266,23 @@ public abstract class ItemTranslator {
     }
 
     /**
-     * Given an item stack, determine the item mapping that should be applied to Bedrock players.
+     * Given an item stack, determine the Bedrock item ID that should be applied to Bedrock players.
      */
-    @Nonnull
-    public static ItemMapping getBedrockItemMapping(GeyserSession session, @Nonnull GeyserItemStack itemStack) {
+    public static int getBedrockItemId(GeyserSession session, @Nonnull GeyserItemStack itemStack) {
         if (itemStack.isEmpty()) {
-            return ItemMapping.AIR;
+            return ItemMapping.AIR.getJavaId();
         }
         int javaId = itemStack.getJavaId();
-        return ITEM_STACK_TRANSLATORS.getOrDefault(javaId, DEFAULT_TRANSLATOR)
+        ItemMapping mapping = ITEM_STACK_TRANSLATORS.getOrDefault(javaId, DEFAULT_TRANSLATOR)
                 .getItemMapping(javaId, itemStack.getNbt(), session.getItemMappings());
+
+        int customItemId = getCustomItem(itemStack.getNbt(), mapping);
+        if (customItemId == -1) {
+            // No custom item
+            return mapping.getBedrockId();
+        } else {
+            return customItemId;
+        }
     }
 
     private static final ItemTranslator DEFAULT_TRANSLATOR = new ItemTranslator() {
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockInventoryTransactionTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockInventoryTransactionTranslator.java
index 24c046ef2..815456132 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockInventoryTransactionTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/BedrockInventoryTransactionTranslator.java
@@ -542,7 +542,7 @@ public class BedrockInventoryTransactionTranslator extends PacketTranslator<Inve
 
     private boolean isIncorrectHeldItem(GeyserSession session, InventoryTransactionPacket packet) {
         int javaSlot = session.getPlayerInventory().getOffsetForHotbar(packet.getHotbarSlot());
-        int expectedItemId = ItemTranslator.getBedrockItemMapping(session, session.getPlayerInventory().getItem(javaSlot)).getBedrockId();
+        int expectedItemId = ItemTranslator.getBedrockItemId(session, session.getPlayerInventory().getItem(javaSlot));
         int heldItemId = packet.getItemInHand() == null ? ItemData.AIR.getId() : packet.getItemInHand().getId();
 
         if (expectedItemId != heldItemId) {