1 files changed, 24 insertions, 13 deletions
diff --git a/src/main/java/gq/malwarefight/nosession/tweaks/initial/InitialTweaker.java b/src/main/java/gq/malwarefight/nosession/tweaks/initial/InitialTweaker.java
index e2ca1ae..21b2646 100644
--- a/src/main/java/gq/malwarefight/nosession/tweaks/initial/InitialTweaker.java
+++ b/src/main/java/gq/malwarefight/nosession/tweaks/initial/InitialTweaker.java
@@ -9,11 +9,11 @@ import org.spongepowered.asm.mixin.MixinEnvironment;
 import org.spongepowered.asm.mixin.Mixins;
 
 import java.io.File;
-import java.io.IOException;
 import java.net.Socket;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.UUID;
 
 public class InitialTweaker implements ITweaker {
 
@@ -30,22 +30,13 @@ public class InitialTweaker implements ITweaker {
      */
     @Override
     public final void acceptOptions(List<String> args, File gameDir, File assetsDir, String version) {
+        boolean isRelaunch = false; // is this invocation from a relaunch?
+        String uuid = null;
         ArrayList<String> argsCopy = new ArrayList<>(args);
         for (int i = 0; i < argsCopy.size(); i++) {
             if (argsCopy.get(i).equals("--accessToken")) {
                 if (argsCopy.get(i + 1).equals("<noSessionAccessToken>")) {
-                    Runtime.getRuntime().addShutdownHook(
-                            new Thread(() -> {
-                                try {
-                                    Socket socket = Utils.getProperSocket();
-                                    socket.getOutputStream().write("fullquit\n".getBytes(StandardCharsets.UTF_8));
-                                    socket.close();
-                                } catch (IOException e) {
-                                    e.printStackTrace();
-                                }
-                            })
-                    );
-                    return; // don't do anything, the change has already been made
+                    isRelaunch = true;
                 }
                 try {
                     Utils.setToken(args.get(i + 1));
@@ -53,8 +44,28 @@ public class InitialTweaker implements ITweaker {
                     throw new RuntimeException(e);
                 }
                 argsCopy.set(i + 1, "<noSessionAccessToken>");
+            } else if (argsCopy.get(i).equals("--uuid")) {
+                uuid = argsCopy.get(i + 1);
             }
         }
+        if (isRelaunch) {
+            String finalUuid = uuid;
+            if (finalUuid == null) {
+                return; // if we don't have a uuid, quit
+            }
+            Runtime.getRuntime().addShutdownHook(
+                    new Thread(() -> {
+                        try {
+                            Socket socket = Utils.getProperSocket(UUID.fromString(finalUuid));
+                            socket.getOutputStream().write("fullquit\n".getBytes(StandardCharsets.UTF_8));
+                            socket.close();
+                        } catch (Exception e) {
+                            e.printStackTrace();
+                        }
+                    })
+            );
+            return; // don't do anything, the change has already been made
+        }
         System.out.println("=======================");
         System.out.println("NoSession: relaunching without the token");
         System.out.println("=======================");