From 17ecf0c73629bae8447f2553e5c5b31a1134c5c7 Mon Sep 17 00:00:00 2001
From: FlorianMichael <florian.michael07@gmail.com>
Date: Sat, 25 May 2024 22:22:44 +0200
Subject: [PATCH] Add ProtocolTranslator#getTargetVersion with channel argument

---
 .../fixes/viaversion/ViaFabricPlusProtocol.java     |  9 +++------
 .../protocoltranslator/ProtocolTranslator.java      | 13 +++++++++++++
 .../ViaFabricPlusBaseVersionProvider.java           |  2 +-
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/main/java/de/florianmichael/viafabricplus/fixes/viaversion/ViaFabricPlusProtocol.java b/src/main/java/de/florianmichael/viafabricplus/fixes/viaversion/ViaFabricPlusProtocol.java
index b2ddb587..258eec32 100644
--- a/src/main/java/de/florianmichael/viafabricplus/fixes/viaversion/ViaFabricPlusProtocol.java
+++ b/src/main/java/de/florianmichael/viafabricplus/fixes/viaversion/ViaFabricPlusProtocol.java
@@ -74,11 +74,8 @@ public class ViaFabricPlusProtocol extends AbstractSimpleProtocol {
 
             final ProtocolVersion version = wrapper.user().getProtocolInfo().serverProtocolVersion();
             if (!payloadDiff.containsKey(channel) || version.olderThan(payloadDiff.get(channel).getLeft())) {
-                // Technically, it's wrong to just drop all payloads.
-                // However, ViaVersion doesn't translate them and the server can't detect if
-                // we handled the payload or not,
-                // so dropping them is easier
-                // than adding a bunch of useless translations for payloads
+                // Technically, it's wrong to just drop all payloads. However, ViaVersion doesn't translate them and the server can't detect if
+                // we handled the payload or not, so dropping them is easier than adding a bunch of useless translations for payloads
                 // which don't do anything on the client anyway.
                 wrapper.cancel();
                 return;
@@ -101,7 +98,7 @@ public class ViaFabricPlusProtocol extends AbstractSimpleProtocol {
     public void init(UserConnection connection) {
         super.init(connection);
 
-        final ProtocolVersion serverVersion = connection.getChannel().attr(ProtocolTranslator.TARGET_VERSION_ATTRIBUTE_KEY).get();
+        final ProtocolVersion serverVersion = ProtocolTranslator.getTargetVersion(connection.getChannel());
 
         // Add storages we need for different fixes here
         if (serverVersion.equals(BedrockProtocolVersion.bedrockLatest)) {
diff --git a/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/ProtocolTranslator.java b/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/ProtocolTranslator.java
index 7883349d..6e075de2 100644
--- a/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/ProtocolTranslator.java
+++ b/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/ProtocolTranslator.java
@@ -158,6 +158,19 @@ public class ProtocolTranslator {
         return targetVersion;
     }
 
+    /**
+     * Gets the target version from the channel attribute, can be used in early stages of the connection
+     *
+     * @param channel the channel
+     * @return the target version
+     */
+    public static ProtocolVersion getTargetVersion(final Channel channel) {
+        if (channel == null || !channel.hasAttr(TARGET_VERSION_ATTRIBUTE_KEY)) {
+            throw new IllegalStateException("Target version attribute not set");
+        }
+        return channel.attr(TARGET_VERSION_ATTRIBUTE_KEY).get();
+    }
+
     /**
      * Sets the target version
      *
diff --git a/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/impl/provider/viaversion/ViaFabricPlusBaseVersionProvider.java b/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/impl/provider/viaversion/ViaFabricPlusBaseVersionProvider.java
index bfe08fae..72faaf8b 100644
--- a/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/impl/provider/viaversion/ViaFabricPlusBaseVersionProvider.java
+++ b/src/main/java/de/florianmichael/viafabricplus/protocoltranslator/impl/provider/viaversion/ViaFabricPlusBaseVersionProvider.java
@@ -29,7 +29,7 @@ public class ViaFabricPlusBaseVersionProvider extends BaseVersionProvider {
     @Override
     public ProtocolVersion getClosestServerProtocol(UserConnection connection) throws Exception {
         if (connection.isClientSide()) {
-            return connection.getChannel().attr(ProtocolTranslator.TARGET_VERSION_ATTRIBUTE_KEY).get();
+            return ProtocolTranslator.getTargetVersion(connection.getChannel());
         } else {
             return super.getClosestServerProtocol(connection);
         }