Improvements

5 files changed, 130 insertions, 62 deletions

diff --git a/src/main/java/gq/malwarefight/nosession/mixin/InitialTweaker.java b/src/main/java/gq/malwarefight/nosession/mixin/InitialTweaker.java
index 7d9c308..51425fc 100644
--- a/src/main/java/gq/malwarefight/nosession/mixin/InitialTweaker.java
+++ b/src/main/java/gq/malwarefight/nosession/mixin/InitialTweaker.java
@@ -21,10 +21,17 @@ import org.spongepowered.asm.mixin.Mixins;
 
 import java.io.File;
 import java.io.IOException;
+import java.net.Socket;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Random;
 
 @SuppressWarnings({"unused", "unchecked"})
 public class InitialTweaker implements ITweaker {
@@ -56,16 +63,59 @@ public class InitialTweaker implements ITweaker {
         );
     }
 
+    public static boolean createLockFile(long value) {
+        Path path = Paths.get(System.getProperty("java.io.tmpdir"), "NoSessionLock" + value);
+        try {
+            Files.createFile(path);
+        } catch (FileAlreadyExistsException e) {
+            System.out.println("You won the lottery! Two NoSession instances used the same ID. The chance that a new NoSession instance uses the same ID as an old one is (NoSession instances) / 2^64");
+            return false;
+        } catch (IOException e) {
+            LogManager.getLogger().warn("Failed to create lockfile " + e.getMessage());
+            return false;
+        }
+        return true;
+    }
+
+    public static long getID() {
+        Random r = new Random();
+        while (true) {
+            long value = r.nextLong();
+            if (createLockFile(value)) {
+                Runtime.getRuntime().addShutdownHook(
+                        new Thread(() -> {
+                            try {
+                                Files.delete(Paths.get(System.getProperty("java.io.tmpdir"), "NoSessionLock" + value));
+                            } catch (Exception ignored) {}
+                        })
+                );
+                return value;
+            }
+        }
+    }
 
     public void setToken(String token) throws IOException, URISyntaxException {
+        long value = getID();
+        Utils.ID = value;
         File javaExecutable = Paths.get(System.getProperty("java.home"), "bin", "java").toFile();
         ProcessBuilder processBuilder = new ProcessBuilder(
-                javaExecutable.getAbsolutePath(), "-cp", getClasspath(), "gq.malwarefight.tokenapp.Main"
+                javaExecutable.getAbsolutePath(), "-cp", getClasspath(), "gq.malwarefight.tokenapp.Main", Long.toString(value)
         );
         processBuilder.redirectOutput(ProcessBuilder.Redirect.INHERIT).redirectError(ProcessBuilder.Redirect.INHERIT);
         Process c = processBuilder.start();
         c.getOutputStream().write((token + "\n").getBytes(StandardCharsets.UTF_8));
         c.getOutputStream().flush();
+        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();
+                    }
+                })
+        );
     }
 
     /**
@@ -81,6 +131,11 @@ public class InitialTweaker implements ITweaker {
         ArrayList<String> argsCopy = new ArrayList<>(args);
         for (int i = 0; i < argsCopy.size(); i++) {
             if (argsCopy.get(i).equals("--accessToken")) {
+                try {
+                    setToken(args.get(i + 1));
+                } catch (IOException | URISyntaxException e) {
+                    throw new RuntimeException(e);
+                }
                 argsCopy.set(i + 1, "<noSessionAccessToken>");
             }
         }
@@ -141,11 +196,6 @@ public class InitialTweaker implements ITweaker {
         ArrayList<String> launchArgs = (ArrayList<String>) Launch.blackboard.get("ArgumentList");
         for (int i = 0; i < launchArgs.size(); i++) {
             if (launchArgs.get(i).equals("--accessToken")) {
-                try {
-                    setToken(launchArgs.get(i + 1));
-                } catch (Exception e) {
-                    throw new RuntimeException(e);
-                }
                 launchArgs.set(i + 1, "<noSessionAccessToken>");
             }
         }
diff --git a/src/main/java/gq/malwarefight/nosession/mixin/Utils.java b/src/main/java/gq/malwarefight/nosession/mixin/Utils.java
index e67bddc..2aaec37 100644
--- a/src/main/java/gq/malwarefight/nosession/mixin/Utils.java
+++ b/src/main/java/gq/malwarefight/nosession/mixin/Utils.java
@@ -1,21 +1,24 @@
 package gq.malwarefight.nosession.mixin;
 
 import gq.malwarefight.nosession.mixin.asm.ReplacingMethodVisitor;
-import net.minecraft.launchwrapper.ITweaker;
 import net.minecraft.launchwrapper.Launch;
 import org.objectweb.asm.*;
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
-import java.util.AbstractList;
-import java.util.ArrayList;
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 
 public class Utils {
     static int num = 0;
+    static int PORT = -1;
+    static long ID = -1;
+    private static final int BASE_PORT = 47777;
 
     public static byte[] read(InputStream i, Character delimiter) throws IOException {
         byte[] buffer = new byte[512];
@@ -39,7 +42,7 @@ public class Utils {
     }
 
     public static String readString(InputStream i, Character delimiter) throws IOException {
-        return new String(read(i, delimiter));
+        return new String(read(i, delimiter), StandardCharsets.UTF_8);
     }
 
     public static void createClass(byte[] classArray, String name) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
@@ -86,4 +89,34 @@ public class Utils {
         num++;
         return className;
     }
+
+    public static Socket getProperSocket() {
+        if (PORT == -1) {
+            Socket socket = null;
+            int port = 0;
+            for (int i = BASE_PORT; i < BASE_PORT + 10; i++) {
+                try {
+                    socket = new Socket();
+                    socket.connect(new InetSocketAddress(InetAddress.getLoopbackAddress(), i));
+                    socket.getOutputStream().write("id\n".getBytes(StandardCharsets.UTF_8));
+                    String value = readString(socket.getInputStream(), '\n');
+                    if (value.equals(Long.toString(ID))) {
+                        port = i;
+                        break;
+                    }
+                } catch (Exception ignored) {}
+            }
+            PORT = port;
+            return socket;
+        } else {
+            try {
+                Socket socket = new Socket();
+                socket.connect(new InetSocketAddress(InetAddress.getLoopbackAddress(), PORT));
+                return socket;
+            } catch (IOException e) {
+                PORT = -1;
+                return getProperSocket();
+            }
+        }
+    }
 }
diff --git a/src/main/java/gq/malwarefight/nosession/mixin/client/YggdrasilSessionMixin.java b/src/main/java/gq/malwarefight/nosession/mixin/client/YggdrasilSessionMixin.java
index 440ca17..48327fe 100644
--- a/src/main/java/gq/malwarefight/nosession/mixin/client/YggdrasilSessionMixin.java
+++ b/src/main/java/gq/malwarefight/nosession/mixin/client/YggdrasilSessionMixin.java
@@ -3,40 +3,29 @@ package gq.malwarefight.nosession.mixin.client;
 import com.mojang.authlib.GameProfile;
 import com.mojang.authlib.exceptions.AuthenticationException;
 import com.mojang.authlib.yggdrasil.YggdrasilMinecraftSessionService;
+import gq.malwarefight.nosession.mixin.Utils;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 import java.io.IOException;
-import java.net.InetAddress;
-import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.nio.charset.StandardCharsets;
 import java.util.Objects;
-import java.util.Scanner;
 
 
 @Mixin(YggdrasilMinecraftSessionService.class)
 public class YggdrasilSessionMixin {
-    private static final int BASE_PORT = 47777;
     @Inject(method = "joinServer", at = @At("HEAD"), cancellable = true, remap = false)
     public void joinServer(GameProfile profile, String authenticationToken, String serverId, CallbackInfo ci) throws AuthenticationException {
         if (authenticationToken.equals("<noSessionAccessToken>")) { // this token has been messed with by the Tweaker
-            Socket socket = null;
-            for (int i = BASE_PORT; i < BASE_PORT + 10; i++) {
-                try {
-                    socket = new Socket();
-                    socket.connect(new InetSocketAddress(InetAddress.getLoopbackAddress(), i));
-                    break;
-                } catch (Exception ignored) {}
-            }
+            Socket socket = Utils.getProperSocket();
             try {
                 assert socket != null;
                 socket.getOutputStream().write(("login " + serverId + "\n").getBytes(StandardCharsets.UTF_8));
                 socket.getOutputStream().flush();
-                Scanner result = new Scanner(socket.getInputStream(), StandardCharsets.UTF_8.name());
-                String line = result.nextLine();
+                String line = Utils.readString(socket.getInputStream(), '\n');
                 String[] parts = line.split(" ");
                 if (Objects.equals(parts[0], "200")) {
                     ci.cancel();
@@ -52,7 +41,7 @@ public class YggdrasilSessionMixin {
             }
 
             try {
-                socket.getOutputStream().write("disconnect".getBytes(StandardCharsets.UTF_8));
+                socket.getOutputStream().write("disconnect\n".getBytes(StandardCharsets.UTF_8));
                 socket.close();
             } catch (IOException ignored) {}
         }
diff --git a/src/main/java/gq/malwarefight/tokenapp/Main.java b/src/main/java/gq/malwarefight/tokenapp/Main.java
index 750ae83..66afad1 100644
--- a/src/main/java/gq/malwarefight/tokenapp/Main.java
+++ b/src/main/java/gq/malwarefight/tokenapp/Main.java
@@ -15,7 +15,7 @@ import java.util.UUID;
 
 public class Main {
     public static final int BASE_PORT = 47777;
-    public static volatile boolean running = true;
+    public static long ID;
     public static YggdrasilMinecraftSessionService sessionService = null;
     public static YggdrasilAuthenticationService authenticationService = null;
     public static GameProfile gameProfile = null;
@@ -40,6 +40,7 @@ public class Main {
     }
 
     public static void main(String[] args) {
+        ID = Long.parseLong(args[0]);
         try {
             setup();
         } catch (Exception e) {
@@ -49,25 +50,22 @@ public class Main {
         ServerSocket sock = null;
         for (int i = BASE_PORT; i < BASE_PORT + 10; i++) {
             try {
-                System.out.println(i);
                 //noinspection resource
                 sock = new ServerSocket(i, 50, InetAddress.getLoopbackAddress());
-                System.out.println(i);
                 break;
-            } catch (Exception ignored) {
-            }
+            } catch (Exception ignored) {}
         }
         if (sock == null) {
             System.err.println("Could not bind to any valid port");
             System.exit(1);
         }
-        while (running) {
+        while (true) {
             try {
                 Socket connection = sock.accept();
                 Thread t = new SocketThread(connection);
                 t.start();
             } catch (IOException ignored) {
-                running = false;
+                System.exit(0);
             }
         }
     }
diff --git a/src/main/java/gq/malwarefight/tokenapp/SocketThread.java b/src/main/java/gq/malwarefight/tokenapp/SocketThread.java
index 572ca57..05db6f0 100644
--- a/src/main/java/gq/malwarefight/tokenapp/SocketThread.java
+++ b/src/main/java/gq/malwarefight/tokenapp/SocketThread.java
@@ -1,13 +1,13 @@
 package gq.malwarefight.tokenapp;
 
 import com.mojang.authlib.exceptions.AuthenticationException;
+import gq.malwarefight.nosession.mixin.Utils;
 
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.Socket;
 import java.nio.charset.StandardCharsets;
-import java.util.Scanner;
 
 public class SocketThread extends Thread {
     private final Socket sock;
@@ -27,37 +27,35 @@ public class SocketThread extends Thread {
         try {
             InputStream in = sock.getInputStream();
             OutputStream out = sock.getOutputStream();
-            Scanner scanner = new Scanner(in, StandardCharsets.UTF_8.name());
             while (true) {
-                if (scanner.hasNextLine()) {
-                    String line = scanner.nextLine();
-                    String[] parts = line.split(" ");
-                    if ("login".equals(parts[0])) {
-                        try {
-                            Main.sessionService.joinServer(
-                                    Main.gameProfile, Main.authenticationService.getClientToken(), parts[1]
-                            );
-                        } catch (Exception e) {
-                            if (e instanceof ArrayIndexOutOfBoundsException) {
-                                writeResponse("400 Bad Request", out);
-                            } else if (e instanceof AuthenticationException) {
-                                writeResponse("401 Unauthorized", out);
-                            } else {
-                                writeResponse("500 Internal Server Error", out);
-                            }
-                            continue;
+                String line = Utils.readString(in, '\n');
+                String[] parts = line.split(" ");
+                if ("login".equals(parts[0])) {
+                    try {
+                        Main.sessionService.joinServer(
+                                Main.gameProfile, Main.authenticationService.getClientToken(), parts[1]
+                        );
+                    } catch (Exception e) {
+                        if (e instanceof ArrayIndexOutOfBoundsException) {
+                            writeResponse("400 Bad Request", out);
+                        } else if (e instanceof AuthenticationException) {
+                            writeResponse("401 Unauthorized", out);
+                        } else {
+                            writeResponse("500 Internal Server Error", out);
                         }
-                        writeResponse("200 OK", out);
-                    } else if ("disconnect".equals(parts[0])) {
-                        sock.close();
-                    } else if ("fullquit".equals(parts[0])) {
-                        Main.running = false;
-                        sock.close();
-                        return;
-                    } else {
-                        writeResponse("418 I'm a teapot", out);
+                        continue;
                     }
+                    writeResponse("200 OK", out);
+                } else if ("disconnect".equals(parts[0])) {
+                    sock.close();
+                } else if ("fullquit".equals(parts[0])) {
+                    System.exit(0);
+                } else if ("id".equals(parts[0])) {
+                    writeResponse(Long.toString(Main.ID), out);
+                } else {
+                    writeResponse("418 I'm a teapot", out);
                 }
+
             }
         } catch (IOException ignored) {}
     }