Negotiate protocol version and refactor BetterWhitelistServer

src/client/java/io/gitlab/jfronny/betterwhitelist/client/BetterWhitelistClient.java

@@ -20,7 +20,14 @@ import java.util.concurrent.CompletableFuture;
 public class BetterWhitelistClient implements ClientModInitializer {
     @Override
     public void onInitializeClient() {
-        ClientLoginNetworking.registerGlobalReceiver(BetterWhitelist.LOGIN_CHANNEL, (client, handler, buf, listenerAdder) -> {
+        ClientLoginNetworking.registerGlobalReceiver(BetterWhitelist.HANDSHAKE_CHANNEL, (client, handler, buf, listenerAdder) -> {
+            // buf also contains the server version
+            PacketByteBuf resultBuf = PacketByteBufs.create();
+            resultBuf.writeInt(BetterWhitelist.PROTOCOL_VERSION);
+            return CompletableFuture.completedFuture(resultBuf);
+        });
+
+        ClientLoginNetworking.registerGlobalReceiver(BetterWhitelist.CHALLENGE_CHANNEL, (client, handler, buf, listenerAdder) -> {
             Scope fork = BetterWhitelist.SCOPE.fork();
             fork.set("resourcePacks", MinecraftClient.getInstance()
                     .getResourcePackManager()

src/main/java/io/gitlab/jfronny/betterwhitelist/BetterWhitelist.java

@@ -19,7 +19,9 @@ import static io.gitlab.jfronny.muscript.data.dynamic.additional.DFinal.of;
 public class BetterWhitelist {
     public static final String MOD_ID = "better-whitelist";
     public static final Logger LOG = Logger.forName(MOD_ID);
-    public static final Identifier LOGIN_CHANNEL = new Identifier(MOD_ID, "challenge");
+    public static final Identifier HANDSHAKE_CHANNEL = new Identifier(MOD_ID, "handshake");
+    public static final Identifier CHALLENGE_CHANNEL = new Identifier(MOD_ID, "challenge");
+    public static final int PROTOCOL_VERSION = 1;
     public static final ModMetadata MOD_METADATA = FabricLoader.getInstance().getModContainer(MOD_ID).orElseThrow().getMetadata();
     public static final Scope SCOPE = StandardLib.createScope();
 

src/main/java/io/gitlab/jfronny/betterwhitelist/server/BetterWhitelistServer.java

@@ -6,11 +6,8 @@ import io.gitlab.jfronny.betterwhitelist.BetterWhitelist;
 import io.gitlab.jfronny.betterwhitelist.server.mixin.ServerLoginNetworkHandlerAccessor;
 import io.gitlab.jfronny.commons.StringFormatter;
 import io.gitlab.jfronny.muscript.compiler.Parser;
-import io.gitlab.jfronny.muscript.data.Scope;
 import io.gitlab.jfronny.muscript.data.Script;
-import io.gitlab.jfronny.muscript.data.dynamic.DNull;
 import io.gitlab.jfronny.muscript.data.dynamic.Dynamic;
-import io.gitlab.jfronny.muscript.data.dynamic.additional.DFinal;
 import io.gitlab.jfronny.muscript.error.LocationalException;
 import net.fabricmc.api.DedicatedServerModInitializer;
 import net.fabricmc.fabric.api.command.v2.CommandRegistrationCallback;
@@ -26,7 +23,7 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.*;
-import java.util.concurrent.FutureTask;
+import java.util.concurrent.*;
 
 import static net.minecraft.server.command.CommandManager.literal;
 
@@ -35,17 +32,6 @@ public class BetterWhitelistServer implements DedicatedServerModInitializer {
     private Script script;
     private final Map<GameProfile, Challenge> challenges = new HashMap<>();
 
-    private static class Challenge {
-        public final Queue<Dynamic<?>> responses = new LinkedList<>();
-        public final GameProfile profile;
-        public PacketSender response;
-
-        public Challenge(GameProfile profile, PacketSender response) {
-            this.profile = profile;
-            this.response = response;
-        }
-    }
-
     @Override
     public void onInitializeServer() {
         CommandRegistrationCallback.EVENT.register((dispatcher, registryAccess, environment) -> {
@@ -59,75 +45,64 @@ public class BetterWhitelistServer implements DedicatedServerModInitializer {
 
         ServerLoginConnectionEvents.QUERY_START.register((handler, server, sender, synchronizer) -> {
             GameProfile gp = profile(handler);
-            challenges.put(gp, new Challenge(gp, sender));
+            Challenge challenge = new Challenge(gp, sender);
-            Challenge challenge = challenges.get(gp);
+            challenges.put(gp, challenge);
-
+            PacketByteBuf handshake = PacketByteBufs.create();
-            Scope fork = BetterWhitelist.SCOPE.fork();
+            handshake.writeInt(BetterWhitelist.PROTOCOL_VERSION);
-            fork.set("assert", args -> {
+            challenge.sender.sendPacket(BetterWhitelist.HANDSHAKE_CHANNEL, handshake);
-                if (args.size() != 1 && args.size() != 2) throw new IllegalArgumentException("Invalid number of arguments for assert: expected 1 or 2 but got " + args.size());
+            synchronizer.waitFor(challenge.challengeCompleted);
-                if (!args.get(0).asBool().getValue()) throw new AssertFail(args.size() > 1 ? args.get(1).asString().getValue() : "Failed Whitelist Check");
-                return new DNull();
-            }).set("challenge", args -> {
-                if (args.size() == 0) throw new IllegalArgumentException("Invalid number of arguments for challenge: expected 1 or more but got 0");
-                PacketByteBuf buf = PacketByteBufs.create();
-                String challengeString = Dynamic.serialize(args.get(0).asCallable());
-                BetterWhitelist.LOG.info("Sending challenge to " + gp.getName() + ": " + challengeString);
-                buf.writeString(challengeString);
-                List<Dynamic<?>> params = args.getValue().subList(1, args.size());
-                buf.writeInt(params.size());
-                params.forEach(p -> buf.writeString(Dynamic.serialize(p)));
-                challenge.response.sendPacket(BetterWhitelist.LOGIN_CHANNEL, buf);
-                int i = 0;
-                while (challenge.responses.isEmpty() && handler.isConnectionOpen() && i++ < 100) sleep(10);
-                if (!challenge.responses.isEmpty()) return challenge.responses.remove();
-                else throw new AssertFail("Took too long to respond");
-            }).set("user", Map.of(
-                    "id", gp.getId() == null ? new DNull() : DFinal.of(gp.getId().toString()),
-                    "name", DFinal.of(gp.getName())
-            ));
-
-            FutureTask<?> future = new FutureTask<>(() -> {
-                try {
-                    try {
-                        script.run(fork);
-                        BetterWhitelist.LOG.info("Completed challenge for " + gp.getName());
-                    } catch (LocationalException le) {
-                        for (Throwable t = le; t != null; t = t.getCause()) {
-                            if (t instanceof AssertFail af) throw af;
-                        }
-                        throw le;
-                    }
-                } catch (AssertFail fail) {
-                    BetterWhitelist.LOG.warn("Failed challenge for " + gp.getName() + ": " + fail.getMessage());
-                    if (handler.isConnectionOpen()) handler.disconnect(Text.literal(fail.getMessage()));
-                } catch (Throwable t) {
-                    BetterWhitelist.LOG.error("Something went wrong while trying to execute a challenge\n"
-                            + StringFormatter.toString(t, e ->
-                            e instanceof LocationalException le
-                                    ? le.asPrintable(scriptSource).toString()
-                                    : e.toString()
-                    ));
-                    if (handler.isConnectionOpen()) handler.disconnect(Text.literal("Something went wrong"));
-                }
-                challenges.remove(gp);
-                return null;
-            });
-
-            Util.getMainWorkerExecutor().execute(future);
-            synchronizer.waitFor(future);
         });
 
-        ServerLoginNetworking.registerGlobalReceiver(BetterWhitelist.LOGIN_CHANNEL, (server, handler, understood, buf, synchronizer, responseSender) -> {
+        ServerLoginNetworking.registerGlobalReceiver(BetterWhitelist.HANDSHAKE_CHANNEL, (server, handler, understood, buf, synchronizer, responseSender) -> {
+            if (!understood) {
+                handler.disconnect(Text.literal("This server requires better-whitelist to be installed"));
+                return;
+            }
+            try {
+                if (buf.readInt() != BetterWhitelist.PROTOCOL_VERSION) {
+                    handler.disconnect(Text.literal("This server requires a version of better-whitelist supporting the protocol version " + BetterWhitelist.PROTOCOL_VERSION));
+                    return;
+                }
+                Challenge challenge = challenges.get(profile(handler));
+                challenge.sender = responseSender;
+                Util.getMainWorkerExecutor().execute(new FutureTask<>(() -> {
+                    try {
+                        ServerScope.run(script, challenge);
+                        challenge.challengeCompleted.complete(null);
+                        BetterWhitelist.LOG.info("Completed challenge for " + challenge.profile.getName());
+                    } catch (ServerScope.AssertFail fail) {
+                        BetterWhitelist.LOG.warn("Failed challenge for " + challenge.profile.getName() + ": " + fail.getMessage());
+                        if (handler.isConnectionOpen()) handler.disconnect(Text.literal(fail.getMessage()));
+                        challenge.challengeCompleted.cancel();
+                    } catch (Throwable t) {
+                        BetterWhitelist.LOG.error("Something went wrong while trying to execute a challenge\n"
+                                + StringFormatter.toString(t, e ->
+                                e instanceof LocationalException le
+                                        ? le.asPrintable(scriptSource).toString()
+                                        : e.toString()
+                        ));
+                        if (handler.isConnectionOpen()) handler.disconnect(Text.literal("Something went wrong"));
+                        challenge.challengeCompleted.cancel();
+                    }
+                    challenges.remove(challenge.profile);
+                    return null;
+                }));
+            } catch (Throwable t) {
+                handler.disconnect(Text.literal("Handshake failed"));
+            }
+        });
+
+        ServerLoginNetworking.registerGlobalReceiver(BetterWhitelist.CHALLENGE_CHANNEL, (server, handler, understood, buf, synchronizer, responseSender) -> {
             if (!understood) {
                 handler.disconnect(Text.literal("This server requires better-whitelist to be installed"));
                 return;
             }
             try {
                 Challenge ch = challenges.get(profile(handler));
-                ch.response = responseSender;
+                ch.sender = responseSender;
                 String response = buf.readString();
                 BetterWhitelist.LOG.info("Got response from " + ch.profile.getName() + ": " + response);
-                ch.responses.add(Dynamic.deserialize(response));
+                ch.response.complete(Dynamic.deserialize(response));
             } catch (Throwable t) {
                 BetterWhitelist.LOG.error("Failed login", t);
                 handler.disconnect(Text.literal("Invalid dynamic"));
@@ -147,19 +122,6 @@ public class BetterWhitelistServer implements DedicatedServerModInitializer {
         return gp;
     }
 
-    private static void sleep(long millis) {
-        try {
-            Thread.sleep(millis);
-        } catch (InterruptedException ignored) {
-        }
-    }
-
-    private static class AssertFail extends RuntimeException {
-        public AssertFail(String message) {
-            super(message);
-        }
-    }
-
     private int printVersion(CommandContext<ServerCommandSource> context) {
         context.getSource().sendMessage(Text.literal("Loaded " + BetterWhitelist.MOD_METADATA.getName() + " " + BetterWhitelist.MOD_METADATA.getVersion()));
         return 1;

src/main/java/io/gitlab/jfronny/betterwhitelist/server/Challenge.java

@@ -0,0 +1,17 @@
+package io.gitlab.jfronny.betterwhitelist.server;
+
+import com.mojang.authlib.GameProfile;
+import io.gitlab.jfronny.muscript.data.dynamic.Dynamic;
+import net.fabricmc.fabric.api.networking.v1.PacketSender;
+
+class Challenge {
+    public final ManualFuture<?> challengeCompleted = new ManualFuture<>();
+    public final ManualFuture<Dynamic<?>> response = new ManualFuture<>();
+    public final GameProfile profile;
+    public PacketSender sender;
+
+    public Challenge(GameProfile profile, PacketSender sender) {
+        this.profile = profile;
+        this.sender = sender;
+    }
+}

src/main/java/io/gitlab/jfronny/betterwhitelist/server/ManualFuture.java

@@ -0,0 +1,66 @@
+package io.gitlab.jfronny.betterwhitelist.server;
+
+import org.jetbrains.annotations.NotNull;
+
+import java.util.concurrent.*;
+
+public class ManualFuture<T> implements Future<T> {
+    private Result state = Result.RUNNING;
+    private T result;
+
+    @Override
+    public boolean cancel(boolean b) {
+        if (state != Result.RUNNING) return false;
+        cancel();
+        return true;
+    }
+
+    @Override
+    public boolean isCancelled() {
+        return state == Result.CANCELLED;
+    }
+
+    @Override
+    public boolean isDone() {
+        return state == Result.DONE;
+    }
+
+    @Override
+    public T get() throws InterruptedException, ExecutionException {
+        while (state == Result.RUNNING) Thread.sleep(10);
+        if (state == Result.CANCELLED) throw new CancellationException();
+        return result;
+    }
+
+    @Override
+    public T get(long l, @NotNull TimeUnit timeUnit) throws TimeoutException {
+        long millis = timeUnit.toMillis(l) / 10;
+        while (state == Result.RUNNING && millis-- > 0) {
+            try {
+                Thread.sleep(10);
+            } catch (InterruptedException ignored) {
+            }
+        }
+        if (millis <= 0) throw new TimeoutException();
+        if (state == Result.CANCELLED) throw new CancellationException();
+        return result;
+    }
+
+    public void cancel() {
+        this.state = Result.CANCELLED;
+    }
+
+    public void complete(T result) {
+        if (state != Result.RUNNING) throw new IllegalStateException("Attempted to complete non-running future");
+        this.state = Result.DONE;
+        this.result = result;
+    }
+
+    public void reset() {
+        this.state = Result.RUNNING;
+    }
+
+    private enum Result {
+        RUNNING, CANCELLED, DONE
+    }
+}

src/main/java/io/gitlab/jfronny/betterwhitelist/server/ServerScope.java

@@ -0,0 +1,62 @@
+package io.gitlab.jfronny.betterwhitelist.server;
+
+import io.gitlab.jfronny.betterwhitelist.BetterWhitelist;
+import io.gitlab.jfronny.muscript.data.Scope;
+import io.gitlab.jfronny.muscript.data.Script;
+import io.gitlab.jfronny.muscript.data.dynamic.DNull;
+import io.gitlab.jfronny.muscript.data.dynamic.Dynamic;
+import io.gitlab.jfronny.muscript.data.dynamic.additional.DFinal;
+import io.gitlab.jfronny.muscript.error.LocationalException;
+import net.fabricmc.fabric.api.networking.v1.PacketByteBufs;
+import net.minecraft.network.PacketByteBuf;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class ServerScope {
+    public static Dynamic<?> run(Script script, Challenge challenge) {
+        try {
+            return script.run(fork(challenge));
+        } catch (LocationalException le) {
+            for (Throwable t = le; t != null; t = t.getCause()) {
+                if (t instanceof ServerScope.AssertFail af) throw af;
+            }
+            throw le;
+        }
+    }
+
+    private static Scope fork(Challenge challenge) {
+        return BetterWhitelist.SCOPE.fork().set("assert", args -> {
+            if (args.size() != 1 && args.size() != 2) throw new IllegalArgumentException("Invalid number of arguments for assert: expected 1 or 2 but got " + args.size());
+            if (!args.get(0).asBool().getValue()) throw new AssertFail(args.size() > 1 ? args.get(1).asString().getValue() : "Failed Whitelist Check");
+            return new DNull();
+        }).set("challenge", args -> {
+            if (args.size() == 0) throw new IllegalArgumentException("Invalid number of arguments for challenge: expected 1 or more but got 0");
+            PacketByteBuf buf = PacketByteBufs.create();
+            String challengeString = Dynamic.serialize(args.get(0).asCallable());
+            BetterWhitelist.LOG.info("Sending challenge to " + challenge.profile.getName() + ": " + challengeString);
+            buf.writeString(challengeString);
+            List<Dynamic<?>> params = args.getValue().subList(1, args.size());
+            buf.writeInt(params.size());
+            params.forEach(p -> buf.writeString(Dynamic.serialize(p)));
+            challenge.response.reset();
+            challenge.sender.sendPacket(BetterWhitelist.CHALLENGE_CHANNEL, buf);
+            try {
+                return challenge.response.get(1, TimeUnit.SECONDS);
+            } catch (TimeoutException e) {
+                throw new AssertFail("Took too long to respond");
+            }
+        }).set("user", Map.of(
+                "id", challenge.profile.getId() == null ? new DNull() : DFinal.of(challenge.profile.getId().toString()),
+                "name", DFinal.of(challenge.profile.getName())
+        ));
+    }
+
+    public static class AssertFail extends RuntimeException {
+        public AssertFail(String message) {
+            super(message);
+        }
+    }
+}