From cb0c301bafc4da3e6344720c735a5bd53355d177 Mon Sep 17 00:00:00 2001
From: RK_01 <50594595+RaphiMC@users.noreply.github.com>
Date: Tue, 29 Oct 2024 12:51:47 +0100
Subject: [PATCH] Fix 1.21.1 -> 1.21.2 chunk loading edge case (#4212)

Will usually happen on 1.8 servers or older
---
 .../v1_21to1_21_2/Protocol1_21To1_21_2.java   |  2 +
 .../BlockItemPacketRewriter1_21_2.java        | 51 ++++++++++++++++++-
 .../rewriter/EntityPacketRewriter1_21_2.java  |  8 +++
 .../storage/ChunkLoadTracker.java             | 45 ++++++++++++++++
 4 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/storage/ChunkLoadTracker.java

diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/Protocol1_21To1_21_2.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/Protocol1_21To1_21_2.java
index 768e29cb7..cf3f6a1cc 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/Protocol1_21To1_21_2.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/Protocol1_21To1_21_2.java
@@ -48,6 +48,7 @@ import com.viaversion.viaversion.protocols.v1_21to1_21_2.rewriter.ComponentRewri
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.rewriter.EntityPacketRewriter1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.rewriter.ParticleRewriter1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.BundleStateTracker;
+import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.ChunkLoadTracker;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.PlayerPositionStorage;
 import com.viaversion.viaversion.rewriter.AttributeRewriter;
 import com.viaversion.viaversion.rewriter.SoundRewriter;
@@ -231,6 +232,7 @@ public final class Protocol1_21To1_21_2 extends AbstractProtocol<ClientboundPack
         addEntityTracker(connection, new EntityTrackerBase(connection, EntityTypes1_21_2.PLAYER));
         connection.put(new BundleStateTracker());
         connection.put(new PlayerPositionStorage());
+        connection.put(new ChunkLoadTracker());
     }
 
     @Override
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 239f7aeb6..8e90a472f 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
@@ -25,11 +25,15 @@ import com.viaversion.nbt.tag.Tag;
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.data.FullMappings;
 import com.viaversion.viaversion.api.data.MappingData;
+import com.viaversion.viaversion.api.data.Mappings;
 import com.viaversion.viaversion.api.minecraft.BlockPosition;
+import com.viaversion.viaversion.api.minecraft.ChunkPosition;
 import com.viaversion.viaversion.api.minecraft.Holder;
 import com.viaversion.viaversion.api.minecraft.HolderSet;
 import com.viaversion.viaversion.api.minecraft.Particle;
 import com.viaversion.viaversion.api.minecraft.SoundEvent;
+import com.viaversion.viaversion.api.minecraft.blockentity.BlockEntity;
+import com.viaversion.viaversion.api.minecraft.chunks.Chunk;
 import com.viaversion.viaversion.api.minecraft.data.StructuredDataContainer;
 import com.viaversion.viaversion.api.minecraft.data.StructuredDataKey;
 import com.viaversion.viaversion.api.minecraft.item.Item;
@@ -51,6 +55,8 @@ import com.viaversion.viaversion.protocols.v1_21to1_21_2.Protocol1_21To1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ClientboundPackets1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ServerboundPacket1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ServerboundPackets1_21_2;
+import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.BundleStateTracker;
+import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.ChunkLoadTracker;
 import com.viaversion.viaversion.rewriter.BlockRewriter;
 import com.viaversion.viaversion.rewriter.SoundRewriter;
 import com.viaversion.viaversion.rewriter.StructuredItemRewriter;
@@ -85,7 +91,6 @@ public final class BlockItemPacketRewriter1_21_2 extends StructuredItemRewriter<
         blockRewriter.registerBlockUpdate(ClientboundPackets1_21.BLOCK_UPDATE);
         blockRewriter.registerSectionBlocksUpdate1_20(ClientboundPackets1_21.SECTION_BLOCKS_UPDATE);
         blockRewriter.registerLevelEvent1_21(ClientboundPackets1_21.LEVEL_EVENT, 2001);
-        blockRewriter.registerLevelChunk1_19(ClientboundPackets1_21.LEVEL_CHUNK_WITH_LIGHT, ChunkType1_20_2::new);
         blockRewriter.registerBlockEntityData(ClientboundPackets1_21.BLOCK_ENTITY_DATA);
 
         registerAdvancements1_20_3(ClientboundPackets1_21.UPDATE_ADVANCEMENTS);
@@ -346,6 +351,50 @@ public final class BlockItemPacketRewriter1_21_2 extends StructuredItemRewriter<
 
             wrapper.write(Types.BOOLEAN, state == RECIPE_INIT); // Replace
         });
+
+        protocol.registerClientbound(ClientboundPackets1_21.LEVEL_CHUNK_WITH_LIGHT, wrapper -> {
+            final Chunk chunk = blockRewriter.handleChunk1_19(wrapper, ChunkType1_20_2::new);
+            final Mappings blockEntityMappings = protocol.getMappingData().getBlockEntityMappings();
+            if (blockEntityMappings != null) {
+                final List<BlockEntity> blockEntities = chunk.blockEntities();
+                for (int i = 0; i < blockEntities.size(); i++) {
+                    final BlockEntity blockEntity = blockEntities.get(i);
+                    final int id = blockEntity.typeId();
+                    final int mappedId = blockEntityMappings.getNewIdOrDefault(id, id);
+                    if (id != mappedId) {
+                        blockEntities.set(i, blockEntity.withTypeId(mappedId));
+                    }
+                }
+            }
+
+            final ChunkLoadTracker chunkLoadTracker = wrapper.user().get(ChunkLoadTracker.class);
+            if (chunkLoadTracker.isChunkLoaded(chunk.getX(), chunk.getZ())) {
+                // Unload the old chunk, so the new one can be loaded without graphical glitches
+                // Bundling it prevents the client from falling through the world during the chunk swap
+                final boolean isBundling = wrapper.user().get(BundleStateTracker.class).isBundling();
+                if (!isBundling) {
+                    final PacketWrapper bundleStart = wrapper.create(ClientboundPackets1_21_2.BUNDLE_DELIMITER);
+                    bundleStart.send(Protocol1_21To1_21_2.class);
+                }
+
+                final PacketWrapper forgetLevelChunk = wrapper.create(ClientboundPackets1_21_2.FORGET_LEVEL_CHUNK);
+                forgetLevelChunk.write(Types.CHUNK_POSITION, new ChunkPosition(chunk.getX(), chunk.getZ()));
+                forgetLevelChunk.send(Protocol1_21To1_21_2.class);
+                wrapper.send(Protocol1_21To1_21_2.class);
+                wrapper.cancel();
+
+                if (!isBundling) {
+                    final PacketWrapper bundleEnd = wrapper.create(ClientboundPackets1_21_2.BUNDLE_DELIMITER);
+                    bundleEnd.send(Protocol1_21To1_21_2.class);
+                }
+            } else {
+                chunkLoadTracker.addChunk(chunk.getX(), chunk.getZ());
+            }
+        });
+        protocol.registerClientbound(ClientboundPackets1_21.FORGET_LEVEL_CHUNK, wrapper -> {
+            final ChunkPosition chunkPosition = wrapper.passthrough(Types.CHUNK_POSITION);
+            wrapper.user().get(ChunkLoadTracker.class).removeChunk(chunkPosition.chunkX(), chunkPosition.chunkZ());
+        });
     }
 
     private void convertServerboundRecipeDisplayId(final PacketWrapper wrapper) {
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/EntityPacketRewriter1_21_2.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/EntityPacketRewriter1_21_2.java
index ebb4e35eb..8c156304f 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/EntityPacketRewriter1_21_2.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/EntityPacketRewriter1_21_2.java
@@ -18,6 +18,7 @@
 package com.viaversion.viaversion.protocols.v1_21to1_21_2.rewriter;
 
 import com.viaversion.nbt.tag.CompoundTag;
+import com.viaversion.viaversion.api.data.entity.EntityTracker;
 import com.viaversion.viaversion.api.minecraft.RegistryEntry;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_21_2;
@@ -33,6 +34,7 @@ import com.viaversion.viaversion.protocols.v1_21to1_21_2.Protocol1_21To1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ClientboundPackets1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ServerboundPackets1_21_2;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.BundleStateTracker;
+import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.ChunkLoadTracker;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.ClientVehicleStorage;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.storage.PlayerPositionStorage;
 import com.viaversion.viaversion.rewriter.EntityRewriter;
@@ -121,6 +123,12 @@ public final class EntityPacketRewriter1_21_2 extends EntityRewriter<Clientbound
             wrapper.passthrough(Types.VAR_INT); // Portal cooldown
 
             wrapper.write(Types.VAR_INT, 64); // Sea level
+
+            final EntityTracker entityTracker = tracker(wrapper.user());
+            if (entityTracker.currentWorld() != null && !entityTracker.currentWorld().equals(world)) {
+                wrapper.user().get(ChunkLoadTracker.class).clear();
+            }
+
             trackWorldDataByKey1_20_5(wrapper.user(), dimensionId, world);
 
             wrapper.user().remove(ClientVehicleStorage.class);
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/storage/ChunkLoadTracker.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/storage/ChunkLoadTracker.java
new file mode 100644
index 000000000..c1f8cbb00
--- /dev/null
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/storage/ChunkLoadTracker.java
@@ -0,0 +1,45 @@
+/*
+ * This file is part of ViaVersion - https://github.com/ViaVersion/ViaVersion
+ * Copyright (C) 2016-2024 ViaVersion and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package com.viaversion.viaversion.protocols.v1_21to1_21_2.storage;
+
+import com.viaversion.viaversion.api.connection.StorableObject;
+import com.viaversion.viaversion.api.minecraft.ChunkPosition;
+import java.util.HashSet;
+import java.util.Set;
+
+public class ChunkLoadTracker implements StorableObject {
+
+    private final Set<Long> loadedChunks = new HashSet<>();
+
+    public void addChunk(final int x, final int z) {
+        this.loadedChunks.add(ChunkPosition.chunkKey(x, z));
+    }
+
+    public void removeChunk(final int x, final int z) {
+        this.loadedChunks.remove(ChunkPosition.chunkKey(x, z));
+    }
+
+    public boolean isChunkLoaded(final int x, final int z) {
+        return this.loadedChunks.contains(ChunkPosition.chunkKey(x, z));
+    }
+
+    public void clear() {
+        this.loadedChunks.clear();
+    }
+
+}