From b5b677da716a0fb61cb3a2bd49446bcb58d97d3d Mon Sep 17 00:00:00 2001
From: FlorianMichael <60033407+FlorianMichael@users.noreply.github.com>
Date: Mon, 29 Jan 2024 13:57:34 +0100
Subject: [PATCH] Don't reset selected version when playing singleplayer

Closes https://github.com/ViaVersion/ViaFabricPlus/issues/337
---
 .../fixes/data/ResourcePackHeaderDiff.java         |  3 ++-
 .../base/integration/MixinMinecraftClient.java     | 14 +++++++++-----
 .../minecraft/MixinServerResourcePackLoader_4.java |  9 ++++-----
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/main/java/de/florianmichael/viafabricplus/fixes/data/ResourcePackHeaderDiff.java b/src/main/java/de/florianmichael/viafabricplus/fixes/data/ResourcePackHeaderDiff.java
index be5447d1f..04acb56b0 100644
--- a/src/main/java/de/florianmichael/viafabricplus/fixes/data/ResourcePackHeaderDiff.java
+++ b/src/main/java/de/florianmichael/viafabricplus/fixes/data/ResourcePackHeaderDiff.java
@@ -98,8 +98,9 @@ public static void checkOutdated() {
     public static GameVersion get(final VersionEnum version) {
         if (!GAME_VERSION_DIFF.containsKey(version)) {
             return SharedConstants.getGameVersion();
+        } else {
+            return GAME_VERSION_DIFF.get(version);
         }
-        return GAME_VERSION_DIFF.get(version);
     }
 
     private static void registerVersion(final VersionEnum version, final int packFormat, final String name) {
diff --git a/src/main/java/de/florianmichael/viafabricplus/injection/mixin/base/integration/MixinMinecraftClient.java b/src/main/java/de/florianmichael/viafabricplus/injection/mixin/base/integration/MixinMinecraftClient.java
index dd228e8b7..7d2ff8544 100644
--- a/src/main/java/de/florianmichael/viafabricplus/injection/mixin/base/integration/MixinMinecraftClient.java
+++ b/src/main/java/de/florianmichael/viafabricplus/injection/mixin/base/integration/MixinMinecraftClient.java
@@ -19,10 +19,15 @@
 
 package de.florianmichael.viafabricplus.injection.mixin.base.integration;
 
+import com.llamalad7.mixinextras.sugar.Local;
 import de.florianmichael.viafabricplus.event.PostGameLoadCallback;
 import de.florianmichael.viafabricplus.protocolhack.ProtocolHack;
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.RunArgs;
+import net.minecraft.network.ClientConnection;
+import net.minecraft.resource.ResourcePackManager;
+import net.minecraft.server.SaveLoader;
+import net.minecraft.world.level.storage.LevelStorage;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
@@ -36,11 +41,10 @@ private void callPostGameLoadEvent(RunArgs args, CallbackInfo ci) {
         PostGameLoadCallback.EVENT.invoker().postGameLoad();
     }
 
-    @Inject(method = "startIntegratedServer", at = @At("HEAD"))
-    private void disableProtocolHack(CallbackInfo ci) {
-        // Set the target version to the native version when starting a singleplayer world
-        // This will automatically reload all the mappings and reset the target version to the forced version
-        ProtocolHack.setTargetVersion(ProtocolHack.NATIVE_VERSION);
+    @Inject(method = "startIntegratedServer", at = @At(value = "INVOKE", target = "Lnet/minecraft/network/ClientConnection;connect(Ljava/lang/String;ILnet/minecraft/network/listener/ClientLoginPacketListener;)V", shift = At.Shift.BEFORE))
+    private void disableProtocolHack(LevelStorage.Session session, ResourcePackManager dataPackManager, SaveLoader saveLoader, boolean newWorld, CallbackInfo ci, @Local ClientConnection clientConnection) {
+        ProtocolHack.setTargetVersion(ProtocolHack.NATIVE_VERSION, true);
+        clientConnection.channel.closeFuture().addListener(channel -> ProtocolHack.resetPreviousVersion());
     }
 
 }
diff --git a/src/main/java/de/florianmichael/viafabricplus/injection/mixin/fixes/minecraft/MixinServerResourcePackLoader_4.java b/src/main/java/de/florianmichael/viafabricplus/injection/mixin/fixes/minecraft/MixinServerResourcePackLoader_4.java
index bdccb25cd..89afb9059 100644
--- a/src/main/java/de/florianmichael/viafabricplus/injection/mixin/fixes/minecraft/MixinServerResourcePackLoader_4.java
+++ b/src/main/java/de/florianmichael/viafabricplus/injection/mixin/fixes/minecraft/MixinServerResourcePackLoader_4.java
@@ -45,12 +45,11 @@ private void removeHeaders(CallbackInfoReturnable<Map<String, String>> cir) {
         final LinkedHashMap<String, String> modifiableMap = new LinkedHashMap<>(cir.getReturnValue());
         if (ProtocolHack.getTargetVersion().isOlderThanOrEqualTo(VersionEnum.r1_14_3)) {
             modifiableMap.remove("X-Minecraft-Version-ID");
+            if (ProtocolHack.getTargetVersion().isOlderThanOrEqualTo(VersionEnum.r1_12_2)) {
+                modifiableMap.remove("X-Minecraft-Pack-Format");
+                modifiableMap.remove("User-Agent");
+            }
         }
-        if (ProtocolHack.getTargetVersion().isOlderThanOrEqualTo(VersionEnum.r1_12_2)) {
-            modifiableMap.remove("X-Minecraft-Pack-Format");
-            modifiableMap.remove("User-Agent");
-        }
-
         cir.setReturnValue(modifiableMap);
     }