ViaVersion support

src/main/java/com/moulberry/axiom/AxiomPaper.java

@@ -51,6 +51,7 @@ public class AxiomPaper extends JavaPlugin implements Listener {
     public final Map<UUID, RateLimiter> playerBlockBufferRateLimiters = new ConcurrentHashMap<>();
     public final Map<UUID, Restrictions> playerRestrictions = new ConcurrentHashMap<>();
     public final Map<UUID, IdMapper<BlockState>> playerBlockRegistry = new ConcurrentHashMap<>();
+    public final Set<UUID> mismatchedDataVersionUsingViaVersion = Collections.newSetFromMap(new ConcurrentHashMap<>());
     public Configuration configuration;
 
     public IdMapper<BlockState> allowedBlockRegistry = null;
@@ -268,6 +269,7 @@ public class AxiomPaper extends JavaPlugin implements Listener {
             playerBlockBufferRateLimiters.keySet().retainAll(stillActiveAxiomPlayers);
             playerRestrictions.keySet().retainAll(stillActiveAxiomPlayers);
             playerBlockRegistry.keySet().retainAll(stillActiveAxiomPlayers);
+            mismatchedDataVersionUsingViaVersion.retainAll(stillActiveAxiomPlayers);
         }, 20, 20);
 
         boolean sendMarkers = configuration.getBoolean("send-markers");
@@ -295,8 +297,8 @@ public class AxiomPaper extends JavaPlugin implements Listener {
         return this.playerBlockBufferRateLimiters.get(uuid);
     }
 
-    public boolean hasCustomBlockRegistry(UUID uuid) {
+    public boolean isMismatchedDataVersion(UUID uuid) {
-        return this.playerBlockRegistry.containsKey(uuid);
+        return this.mismatchedDataVersionUsingViaVersion.contains(uuid);
     }
 
     public IdMapper<BlockState> getBlockRegistry(UUID uuid) {

src/main/java/com/moulberry/axiom/packet/AxiomBigPayloadHandler.java

@@ -11,6 +11,7 @@ import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.level.ServerPlayer;
 
+import java.io.IOException;
 import java.util.List;
 
 public class AxiomBigPayloadHandler extends ByteToMessageDecoder {
@@ -35,14 +36,31 @@ public class AxiomBigPayloadHandler extends ByteToMessageDecoder {
 
     @Override
     protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
-        // Don't process if channel isn't active
+        // No bytes to read?! Go away
-        if (!ctx.channel().isActive()) {
+        if (in.readableBytes() == 0) {
-            in.skipBytes(in.readableBytes());
+            return;
+        }
+
+        // Don't handle if player doesn't have permission to use Axiom
+        ServerPlayer player = connection.getPlayer();
+        if (player == null || !AxiomPaper.PLUGIN.canUseAxiom(player.getBukkitEntity())) {
+            ctx.fireChannelRead(in.retain());
+
+            // Skip remaining bytes
+            if (in.readableBytes() > 0) {
+                in.skipBytes(in.readableBytes());
+            }
+            return;
+        }
+
+        // Don't process if channel isn't active
+        if (!ctx.channel().isActive()) {
+            if (in.readableBytes() > 0) {
+                in.skipBytes(in.readableBytes());
+            }
             return;
         }
 
-        int i = in.readableBytes();
-        if (i != 0) {
         int readerIndex = in.readerIndex();
         boolean success = false;
         try {
@@ -52,24 +70,22 @@ public class AxiomBigPayloadHandler extends ByteToMessageDecoder {
             if (packetId == payloadId) {
                 ResourceLocation identifier = buf.readResourceLocation();
                 if (identifier.equals(SET_BUFFER)) {
-                        ServerPlayer player = connection.getPlayer();
-                        if (AxiomPaper.PLUGIN.canUseAxiom(player.getBukkitEntity())) {
                     setBlockBuffer.onReceive(player, buf);
                     success = true;
-                            in.skipBytes(in.readableBytes());
+                    if (in.readableBytes() > 0) {
-                            return;
+                        throw new IOException("Axiom packet " + identifier + " was larger than I expected, found " + in.readableBytes() +
+                            " bytes extra whilst reading packet");
                     }
+                    return;
                 } else if (identifier.equals(UPLOAD_BLUEPRINT)) {
-                        ServerPlayer player = connection.getPlayer();
-                        if (AxiomPaper.PLUGIN.canUseAxiom(player.getBukkitEntity())) {
                     uploadBlueprint.onReceive(player, buf);
                     success = true;
-                            in.skipBytes(in.readableBytes());
+                    if (in.readableBytes() > 0) {
-                            return;
+                        throw new IOException("Axiom packet " + identifier + " was larger than I expected, found " + in.readableBytes() +
+                            " bytes extra whilst reading packet");
                     }
+                    return;
                 } else if (requestChunkDataPacketListener != null && identifier.equals(REQUEST_CHUNK_DATA)) {
-                        ServerPlayer player = connection.getPlayer();
-                        if (AxiomPaper.PLUGIN.canUseAxiom(player.getBukkitEntity())) {
                     byte[] bytes = new byte[buf.writerIndex() - buf.readerIndex()];
                     buf.getBytes(buf.readerIndex(), bytes);
 
@@ -84,18 +100,28 @@ public class AxiomBigPayloadHandler extends ByteToMessageDecoder {
                     });
 
                     success = true;
+                    if (in.readableBytes() > 0) {
                         in.skipBytes(in.readableBytes());
+                    }
                     return;
                 }
             }
+        } catch (Throwable t) {
+            if (!(t instanceof IndexOutOfBoundsException)) {
+                // Skip remaining bytes
+                success = true;
+                if (in.readableBytes() > 0) {
+                    in.skipBytes(in.readableBytes());
+                }
+
+                // Throw error, will disconnect client
+                throw t;
             }
-            } catch (Throwable ignored) {
         } finally {
             if (!success) {
                 in.readerIndex(readerIndex);
             }
         }
-        }
 
         ctx.fireChannelRead(in.retain());
 

src/main/java/com/moulberry/axiom/packet/BlueprintRequestPacketListener.java

@@ -6,6 +6,8 @@ import com.moulberry.axiom.blueprint.RawBlueprint;
 import com.moulberry.axiom.blueprint.ServerBlueprintManager;
 import com.moulberry.axiom.blueprint.ServerBlueprintRegistry;
 import io.netty.buffer.Unpooled;
+import net.kyori.adventure.text.Component;
+import net.minecraft.SharedConstants;
 import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.resources.ResourceLocation;
 import org.bukkit.craftbukkit.v1_20_R1.entity.CraftPlayer;
@@ -30,6 +32,11 @@ public class BlueprintRequestPacketListener implements PluginMessageListener {
             return;
         }
 
+        if (this.plugin.isMismatchedDataVersion(player.getUniqueId())) {
+            player.sendMessage(Component.text("Axiom+ViaVersion: This feature isn't supported. Switch your client version to " + SharedConstants.VERSION_STRING + " to use this"));
+            return;
+        }
+
         FriendlyByteBuf friendlyByteBuf = new FriendlyByteBuf(Unpooled.wrappedBuffer(message));
         String path = friendlyByteBuf.readUtf();
 

src/main/java/com/moulberry/axiom/packet/HelloPacketListener.java

@@ -10,13 +10,16 @@ import com.moulberry.axiom.viaversion.ViaVersionHelper;
 import com.moulberry.axiom.world_properties.server.ServerWorldPropertiesRegistry;
 import com.viaversion.viaversion.api.Via;
 import com.viaversion.viaversion.api.data.*;
+import com.viaversion.viaversion.api.protocol.ProtocolPathEntry;
 import com.viaversion.viaversion.api.protocol.version.ProtocolVersion;
 import com.viaversion.viaversion.libs.opennbt.tag.builtin.CompoundTag;
 import io.netty.buffer.Unpooled;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
 import net.minecraft.SharedConstants;
+import net.minecraft.core.IdMapper;
 import net.minecraft.network.FriendlyByteBuf;
+import net.minecraft.world.level.block.state.BlockState;
 import org.bukkit.Bukkit;
 import org.bukkit.NamespacedKey;
 import org.bukkit.World;
@@ -54,12 +57,13 @@ public class HelloPacketListener implements PluginMessageListener {
 
         int serverDataVersion = SharedConstants.getCurrentVersion().getDataVersion().getVersion();
         if (dataVersion != serverDataVersion) {
+            String incompatibleDataVersion = plugin.configuration.getString("incompatible-data-version");
+            if (incompatibleDataVersion == null) incompatibleDataVersion = "warn";
+
             if (!Bukkit.getPluginManager().isPluginEnabled("ViaVersion")) {
                 Component text = Component.text("Axiom: Incompatible data version detected (client " + dataVersion +
                     ", server " + serverDataVersion  + ")");
 
-                String incompatibleDataVersion = plugin.configuration.getString("incompatible-data-version");
-                if (incompatibleDataVersion == null) incompatibleDataVersion = "kick";
                 if (incompatibleDataVersion.equals("warn")) {
                     player.sendMessage(text.color(NamedTextColor.RED));
                     return;
@@ -68,38 +72,32 @@ public class HelloPacketListener implements PluginMessageListener {
                     return;
                 }
             } else {
-//                int playerVersion = Via.getAPI().getPlayerVersion(player.getUniqueId());
+                int playerVersion = Via.getAPI().getPlayerVersion(player.getUniqueId());
-//                if (ProtocolVersion.isRegistered(playerVersion)) {
-//                    ProtocolVersion version = ProtocolVersion.getProtocol(playerVersion);
-//                    String name = version.getName().split("/")[0];
-//
-//
-//                }
 
-                CompoundTag tag = MappingDataLoader.loadNBT("mappings-1.20.2to1.20.3.nbt");
+                IdMapper<BlockState> mapper;
+                try {
+                    mapper = ViaVersionHelper.getBlockRegistryForVersion(this.plugin.allowedBlockRegistry, playerVersion);
+                } catch (Exception e) {
+                    String clientDescription = "client: " + ProtocolVersion.getProtocol(playerVersion);
+                    String serverDescription = "server: " + ProtocolVersion.getProtocol(SharedConstants.getProtocolVersion());
+                    String description = clientDescription + " <-> " + serverDescription;
+                    Component text = Component.text("Axiom+ViaVersion: " + e.getMessage() + " (" + description + ")");
 
-                if (tag == null) {
+                    if (incompatibleDataVersion.equals("warn")) {
-                    player.kick(Component.text("Axiom+ViaVersion: Failed to load mappings (1.20.2 <-> 1.20.3)"));
+                        player.sendMessage(text.color(NamedTextColor.RED));
+                    } else {
+                        player.kick(text);
+                    }
                     return;
                 }
 
-                Mappings mappings = MappingDataLoader.loadMappings(tag, "blockstates");
+                this.plugin.playerBlockRegistry.put(player.getUniqueId(), mapper);
-
+                this.plugin.mismatchedDataVersionUsingViaVersion.add(player.getUniqueId());
-                if (mappings == null) {
-                    player.kick(Component.text("Axiom+ViaVersion: Failed to load mapped blockstates (1.20.2 <-> 1.20.3"));
-                    return;
-                }
-
-                this.plugin.playerBlockRegistry.put(player.getUniqueId(), ViaVersionHelper.applyMappings(this.plugin.allowedBlockRegistry,
-                        BiMappings.of(mappings).inverse()));
 
                 Component text = Component.text("Axiom: Warning, client and server versions don't match. " +
                         "Axiom will try to use ViaVersion conversions, but this process may cause problems");
                 player.sendMessage(text.color(NamedTextColor.RED));
             }
-//            inverse.getNewIdOrDefault()
-
-
         }
 
         if (apiVersion != AxiomConstants.API_VERSION) {
@@ -121,7 +119,7 @@ public class HelloPacketListener implements PluginMessageListener {
             Component text = Component.text("This server requires the use of Axiom 2.3 or later. Contact the server administrator if you believe this is unintentional");
 
             String unsupportedRestrictions = plugin.configuration.getString("client-doesnt-support-restrictions");
-            if (unsupportedRestrictions == null) unsupportedRestrictions = "kick";
+            if (unsupportedRestrictions == null) unsupportedRestrictions = "warn";
             if (unsupportedRestrictions.equals("warn")) {
                 player.sendMessage(text.color(NamedTextColor.RED));
                 return;
@@ -161,6 +159,7 @@ public class HelloPacketListener implements PluginMessageListener {
 
         // Initialize Hotbars
         PersistentDataContainer container = player.getPersistentDataContainer();
+        if (!this.plugin.isMismatchedDataVersion(player.getUniqueId())) {
             int activeHotbarIndex = container.getOrDefault(AxiomConstants.ACTIVE_HOTBAR_INDEX, PersistentDataType.BYTE, (byte) 0);
             PersistentDataContainer hotbarItems = container.get(AxiomConstants.HOTBAR_DATA, PersistentDataType.TAG_CONTAINER);
             if (hotbarItems != null) {
@@ -180,6 +179,7 @@ public class HelloPacketListener implements PluginMessageListener {
                 buf.getBytes(0, bytes);
                 player.sendPluginMessage(this.plugin, "axiom:initialize_hotbars", bytes);
             }
+        }
 
         // Initialize Views
         UUID activeView = container.get(AxiomConstants.ACTIVE_VIEW, UUIDDataType.INSTANCE);

src/main/java/com/moulberry/axiom/packet/RequestChunkDataPacketListener.java

@@ -46,7 +46,7 @@ public class RequestChunkDataPacketListener implements PluginMessageListener {
         FriendlyByteBuf friendlyByteBuf = new FriendlyByteBuf(Unpooled.wrappedBuffer(message));
         long id = friendlyByteBuf.readLong();
 
-        if (!this.plugin.canUseAxiom(bukkitPlayer) || this.plugin.hasCustomBlockRegistry(bukkitPlayer.getUniqueId())) {
+        if (!this.plugin.canUseAxiom(bukkitPlayer) || this.plugin.isMismatchedDataVersion(bukkitPlayer.getUniqueId())) {
             // We always send an 'empty' response in order to make the client happy
             sendEmptyResponse(player, id);
             return;

src/main/java/com/moulberry/axiom/packet/SetBlockBufferPacketListener.java

@@ -340,14 +340,14 @@ public class SetBlockBufferPacketListener {
                         return;
                     }
 
+                    var holder = registry.getHolder(biome);
+                    if (holder.isPresent()) {
                         var chunk = (LevelChunk) world.getChunk(x >> 2, z >> 2, ChunkStatus.FULL, false);
                         if (chunk == null) return;
 
                         var section = chunk.getSection(cy - minSection);
                         PalettedContainer<Holder<Biome>> container = (PalettedContainer<Holder<Biome>>) section.getBiomes();
 
-                    var holder = registry.getHolder(biome);
-                    if (holder.isPresent()) {
                         if (!Integration.canPlaceBlock(player.getBukkitEntity(),
                             new Location(player.getBukkitEntity().getWorld(), x+1, y+1, z+1))) return;
 

src/main/java/com/moulberry/axiom/packet/SetHotbarSlotPacketListener.java

@@ -3,6 +3,7 @@ package com.moulberry.axiom.packet;
 import com.moulberry.axiom.AxiomConstants;
 import com.moulberry.axiom.AxiomPaper;
 import com.moulberry.axiom.persistence.ItemStackDataType;
+import com.viaversion.viaversion.api.Via;
 import io.netty.buffer.Unpooled;
 import net.minecraft.network.FriendlyByteBuf;
 import org.bukkit.NamespacedKey;
@@ -22,7 +23,7 @@ public class SetHotbarSlotPacketListener implements PluginMessageListener {
 
     @Override
     public void onPluginMessageReceived(@NotNull String channel, @NotNull Player player, @NotNull byte[] message) {
-        if (!this.plugin.canUseAxiom(player)) {
+        if (!this.plugin.canUseAxiom(player) || this.plugin.isMismatchedDataVersion(player.getUniqueId())) {
             return;
         }
 

src/main/java/com/moulberry/axiom/packet/SwitchActiveHotbarPacketListener.java

@@ -25,7 +25,7 @@ public class SwitchActiveHotbarPacketListener implements PluginMessageListener {
 
     @Override
     public void onPluginMessageReceived(@NotNull String channel, @NotNull Player player, @NotNull byte[] message) {
-        if (!this.plugin.canUseAxiom(player)) {
+        if (!this.plugin.canUseAxiom(player) || this.plugin.isMismatchedDataVersion(player.getUniqueId())) {
             return;
         }
 

src/main/java/com/moulberry/axiom/packet/UploadBlueprintPacketListener.java

@@ -5,7 +5,9 @@ import com.moulberry.axiom.blueprint.BlueprintIo;
 import com.moulberry.axiom.blueprint.RawBlueprint;
 import com.moulberry.axiom.blueprint.ServerBlueprintManager;
 import com.moulberry.axiom.blueprint.ServerBlueprintRegistry;
+import net.minecraft.SharedConstants;
 import net.minecraft.network.FriendlyByteBuf;
+import net.minecraft.network.chat.Component;
 import net.minecraft.server.level.ServerPlayer;
 
 import java.io.BufferedOutputStream;
@@ -26,6 +28,11 @@ public class UploadBlueprintPacketListener  {
             return;
         }
 
+        if (this.plugin.isMismatchedDataVersion(serverPlayer.getUUID())) {
+            serverPlayer.sendSystemMessage(Component.literal("Axiom+ViaVersion: This feature isn't supported. Switch your client version to " + SharedConstants.VERSION_STRING + " to use this"));
+            return;
+        }
+
         ServerBlueprintRegistry registry = ServerBlueprintManager.getRegistry();
         if (registry == null || this.plugin.blueprintFolder == null) {
             return;

src/main/java/com/moulberry/axiom/viaversion/ViaVersionHelper.java

@@ -1,15 +1,68 @@
 package com.moulberry.axiom.viaversion;
 
 import com.moulberry.axiom.buffer.BlockBuffer;
+import com.viaversion.viaversion.api.Via;
 import com.viaversion.viaversion.api.data.BiMappings;
+import com.viaversion.viaversion.api.data.MappingData;
 import com.viaversion.viaversion.api.data.Mappings;
+import com.viaversion.viaversion.api.protocol.ProtocolPathEntry;
+import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import net.kyori.adventure.text.Component;
+import net.kyori.adventure.text.format.NamedTextColor;
+import net.minecraft.SharedConstants;
 import net.minecraft.core.IdMapper;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.state.BlockState;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
 
 public class ViaVersionHelper {
 
-    public static IdMapper<BlockState> applyMappings(IdMapper<BlockState> registry, BiMappings mappings) {
+    private static final Int2ObjectOpenHashMap<IdMapper<BlockState>> blockRegistryCache = new Int2ObjectOpenHashMap<>();
+    private static final Int2ObjectOpenHashMap<String> blockRegistryErrorCache = new Int2ObjectOpenHashMap<>();
+
+    public static IdMapper<BlockState> getBlockRegistryForVersion(IdMapper<BlockState> mapper, int playerVersion) {
+        if (blockRegistryErrorCache.containsKey(playerVersion)) {
+            throw new RuntimeException(blockRegistryErrorCache.get(playerVersion));
+        }
+        if (blockRegistryCache.containsKey(playerVersion)) {
+            return blockRegistryCache.get(playerVersion);
+        }
+
+        List<ProtocolPathEntry> path = Via.getManager().getProtocolManager().getProtocolPath(playerVersion,
+            SharedConstants.getProtocolVersion());
+
+        if (path == null) {
+            blockRegistryErrorCache.put(playerVersion, "Failed to find protocol path");
+            throw new RuntimeException("Failed to find protocol path");
+        }
+
+        for (int i = path.size()-1; i >= 0; i--) {
+            ProtocolPathEntry protocolPathEntry = path.get(i);
+
+            MappingData mappingData = protocolPathEntry.protocol().getMappingData();
+
+            if (mappingData == null) {
+                blockRegistryErrorCache.put(playerVersion, "Failed to load mapping data (" + protocolPathEntry + ")");
+                throw new RuntimeException("Failed to load mapping data (" + protocolPathEntry + ")");
+            }
+
+            Mappings blockStateMappings = mappingData.getBlockStateMappings();
+
+            if (blockStateMappings == null) {
+                blockRegistryErrorCache.put(playerVersion, "Failed to load BlockState mappings (" + protocolPathEntry + ")");
+                throw new RuntimeException("Failed to load BlockState mappings (" + protocolPathEntry + ")");
+            }
+
+            mapper = ViaVersionHelper.applyMappings(mapper, blockStateMappings);
+        }
+
+        blockRegistryCache.put(playerVersion, mapper);
+        return mapper;
+    }
+
+    public static IdMapper<BlockState> applyMappings(IdMapper<BlockState> registry, Mappings mappings) {
         IdMapper<BlockState> newBlockRegistry = new IdMapper<>();
 
         // Add empty mappings for non-existent blocks