From daaa682e7f92d1bebbc4e5534a83915cccf04c0d Mon Sep 17 00:00:00 2001
From: maxomatic458 <104733404+maxomatic458@users.noreply.github.com>
Date: Wed, 23 Oct 2024 09:12:42 +0200
Subject: [PATCH] Fix desync when modifying player inventory (from valence)
 when player is viewing an open inventory  (#667)

# Objective

- fixes #666
# Solution

- also check for changes in the player's `Inventory` when viewing an
open inventory
---
 crates/valence_inventory/src/lib.rs | 38 ++++++++++++++++---
 src/tests/inventory.rs              | 57 +++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 5 deletions(-)

diff --git a/crates/valence_inventory/src/lib.rs b/crates/valence_inventory/src/lib.rs
index 19841bfa0..42a20e28c 100644
--- a/crates/valence_inventory/src/lib.rs
+++ b/crates/valence_inventory/src/lib.rs
@@ -679,7 +679,9 @@ fn update_open_inventories(
     for (client_entity, mut client, mut inv_state, cursor_item, mut open_inventory) in &mut clients
     {
         // Validate that the inventory exists.
-        let Ok(inventory) = inventories.get_mut(open_inventory.entity) else {
+        let Ok([inventory, player_inventory]) =
+            inventories.get_many_mut([open_inventory.entity, client_entity])
+        else {
             // The inventory no longer exists, so close the inventory.
             commands.entity(client_entity).remove::<OpenInventory>();
 
@@ -725,12 +727,34 @@ fn update_open_inventories(
                 // Send the changed slots.
 
                 // The slots that were NOT changed by this client, and they need to be sent.
-                let changed_filtered = inventory.changed & !open_inventory.client_changed;
+                let changed_filtered =
+                    u128::from(inventory.changed & !open_inventory.client_changed);
 
-                if changed_filtered != 0 {
-                    inv_state.state_id += 1;
+                // The slots changed in the player inventory (e.g by calling
+                // `inventory.set_slot` while the player is viewing the inventory).
+                let mut player_inventory_changed = u128::from(player_inventory.changed);
+
+                // Ignore the armor and crafting grid slots because they are not part of
+                // the open inventory.
+                player_inventory_changed >>= *PlayerInventory::SLOTS_MAIN.start();
+                // "Append" the player inventory to the end of the slots belonging to the opened
+                // inventory.
+                player_inventory_changed <<= inventory.slot_count();
 
-                    for (i, slot) in inventory.slots.iter().enumerate() {
+                let changed_filtered = changed_filtered | player_inventory_changed;
+
+                if changed_filtered != 0 {
+                    for (i, slot) in inventory
+                        .slots
+                        .iter()
+                        .chain(
+                            player_inventory
+                                .slots
+                                .iter()
+                                .skip(*PlayerInventory::SLOTS_MAIN.start() as usize),
+                        )
+                        .enumerate()
+                    {
                         if (changed_filtered >> i) & 1 == 1 {
                             client.write_packet(&ScreenHandlerSlotUpdateS2c {
                                 window_id: inv_state.window_id as i8,
@@ -740,6 +764,10 @@ fn update_open_inventories(
                             });
                         }
                     }
+
+                    player_inventory
+                        .map_unchanged(|f| &mut f.changed)
+                        .set_if_neq(0);
                 }
             }
         }
diff --git a/src/tests/inventory.rs b/src/tests/inventory.rs
index 042811066..e17f29947 100644
--- a/src/tests/inventory.rs
+++ b/src/tests/inventory.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use bevy_app::prelude::*;
 use bevy_ecs::prelude::*;
 
@@ -409,6 +411,61 @@ fn test_should_modify_open_inventory_click_slot() {
     assert_eq!(cursor_item.0, ItemStack::new(ItemKind::Diamond, 2, None));
 }
 
+#[test]
+fn test_sync_inventory_change_made_from_valence_while_inventory_is_opened() {
+    let ScenarioSingleClient {
+        mut app,
+        client,
+        mut helper,
+        ..
+    } = ScenarioSingleClient::new();
+
+    // Process a tick to get past the "on join" logic.
+    app.update();
+    helper.clear_received();
+
+    // Open a inventory for the player. (27 slots)
+    set_up_open_inventory(&mut app, client);
+    app.update();
+    helper.clear_received();
+
+    // Modify the player's inventory.
+    let client_inv_state = app
+        .world_mut()
+        .get::<ClientInventoryState>(client)
+        .expect("could not find client inventory state");
+
+    let inv_state_window_id = client_inv_state.window_id();
+    let inv_state_state_id = client_inv_state.state_id();
+
+    let mut inventory = app
+        .world_mut()
+        .get_mut::<Inventory>(client)
+        .expect("could not find inventory for client");
+
+    inventory.set_slot(9, ItemStack::new(ItemKind::Diamond, 2, None));
+
+    app.update();
+
+    // Make assertions
+    let sent_packets = helper.collect_received();
+
+    let received = sent_packets.0[0]
+        .decode::<ScreenHandlerSlotUpdateS2c>()
+        .unwrap();
+
+    assert_eq!(received.window_id, inv_state_window_id as i8,);
+
+    assert_eq!(received.slot_idx, 27);
+
+    assert_eq!(
+        received.slot_data,
+        Cow::Borrowed(&ItemStack::new(ItemKind::Diamond, 2, None))
+    );
+
+    assert_eq!(received.state_id, VarInt(inv_state_state_id.0));
+}
+
 #[test]
 fn test_prevent_modify_open_inventory_click_slot_readonly_inventory() {
     let ScenarioSingleClient {