Fixed not taking screenshot when disconnected by the server (#24)

Minenash · Mar 29, 2021 · 6d45537 · 6d45537
1 parent 2f559b1
commit 6d45537
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/src/main/java/com/minenash/seamless_loading_screen/FinishQuit.java b/src/main/java/com/minenash/seamless_loading_screen/FinishQuit.java
@@ -7,11 +7,13 @@
 import net.minecraft.client.gui.screen.Screen;
 import net.minecraft.client.gui.screen.TitleScreen;
 import net.minecraft.client.gui.screen.multiplayer.MultiplayerScreen;
+import net.minecraft.client.network.ClientPlayNetworkHandler;
 import net.minecraft.client.realms.gui.screen.RealmsBridgeScreen;
 import net.minecraft.client.realms.gui.screen.RealmsMainScreen;
 import net.minecraft.client.texture.NativeImage;
 import net.minecraft.client.util.ScreenshotUtils;
 import net.minecraft.client.util.math.MatrixStack;
+import net.minecraft.text.Text;
 import net.minecraft.text.TranslatableText;
 import net.minecraft.util.Util;
 
@@ -28,7 +30,14 @@ public FinishQuit() {
 
     private static boolean hudHidden = false;
     private static boolean stop = false;
+    private static ClientPlayNetworkHandler serverOrderedDisconnectHandler = null;
+    private static Text serverOrderedDisconnectReason = null;
 
+    public static void run(ClientPlayNetworkHandler screen, Text reason) {
+        FinishQuit.serverOrderedDisconnectHandler = screen;
+        FinishQuit.serverOrderedDisconnectReason = reason;
+        run(false);
+    }
 
     public static void run(boolean stop) {
         MinecraftClient client = MinecraftClient.getInstance();
@@ -37,6 +46,8 @@ public static void run(boolean stop) {
             ScreenshotLoader.allowCustomScreenshot = false;
             if (stop)
                 client.scheduleStop();
+            else if (serverOrderedDisconnectHandler != null)
+                serverOrderedDisconnectHandler.onDisconnected(serverOrderedDisconnectReason);
             else
                 quit(client);
             return;
@@ -77,8 +88,12 @@ public void render(MatrixStack matrices, int _mouseX, int _mouseY, float _delta)
         if(Config.resolution != Config.ScreenshotResolution.Native)
             resizeScreen(client, client.getWindow().getWidth(), client.getWindow().getHeight());
 
+        SeamlessLoadingScreen.isDisconnecting = true;  //Fapi 0.30.0 compat
+
         if (stop)
             client.scheduleStop();
+        else if (serverOrderedDisconnectHandler != null)
+            serverOrderedDisconnectHandler.onDisconnected(serverOrderedDisconnectReason);
         else
             quit(client);
 
@@ -97,7 +112,6 @@ private static void quit(MinecraftClient client) {
         boolean isSinglePlayer = client.isInSingleplayer();
         boolean isRealms = client.isConnectedToRealms();
 
-        SeamlessLoadingScreen.isDisconnecting = true;  //Fapi 0.30.0 compat
         client.world.disconnect();
         if (isSinglePlayer)
             client.disconnect(new SaveLevelScreen(new TranslatableText("menu.savingLevel")));

diff --git a/src/main/java/com/minenash/seamless_loading_screen/mixin/ClientPlayNetworkHandlerMixin.java b/src/main/java/com/minenash/seamless_loading_screen/mixin/ClientPlayNetworkHandlerMixin.java
@@ -1,22 +1,48 @@
 package com.minenash.seamless_loading_screen.mixin;
 
+import com.minenash.seamless_loading_screen.FinishQuit;
 import com.minenash.seamless_loading_screen.ScreenshotLoader;
 import com.minenash.seamless_loading_screen.SeamlessLoadingScreen;
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.client.gui.screen.DisconnectedScreen;
+import net.minecraft.client.gui.screen.Screen;
+import net.minecraft.client.gui.screen.TitleScreen;
+import net.minecraft.client.gui.screen.multiplayer.MultiplayerScreen;
 import net.minecraft.client.network.ClientPlayNetworkHandler;
+import net.minecraft.client.realms.gui.screen.DisconnectedRealmsScreen;
+import net.minecraft.client.realms.gui.screen.RealmsScreen;
 import net.minecraft.client.world.ClientWorld;
+import net.minecraft.text.Text;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 @Mixin(ClientPlayNetworkHandler.class)
 public class ClientPlayNetworkHandlerMixin {
 
+    @Shadow @Final private Screen loginScreen;
+
+    @Shadow @Final private static Text field_26620;
+
     @Redirect(method = "onGameJoin", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/MinecraftClient;joinWorld(Lnet/minecraft/client/world/ClientWorld;)V"))
     private void setChangeWorldJoinScreen(MinecraftClient client, ClientWorld world) {
         if (ScreenshotLoader.loaded) {
             SeamlessLoadingScreen.changeWorldJoinScreen = true;
         }
         client.joinWorld(world);
     }
+
+    private boolean stopDisconnect = true;
+    @Inject(method = "onDisconnected", at = @At("HEAD"), cancellable = true)
+    private void onServerOrderedDisconnect(Text reason, CallbackInfo info) {
+        if (stopDisconnect) {
+            info.cancel();
+            FinishQuit.run((ClientPlayNetworkHandler) (Object) this, reason);
+        }
+        stopDisconnect = !stopDisconnect;
+    }
 }