diff --git a/src/client/java/io/gitlab/jfronny/betterwhitelist/client/BetterWhitelistClient.java b/src/client/java/io/gitlab/jfronny/betterwhitelist/client/BetterWhitelistClient.java index 3634170..c037f56 100644 --- a/src/client/java/io/gitlab/jfronny/betterwhitelist/client/BetterWhitelistClient.java +++ b/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() diff --git a/src/main/java/io/gitlab/jfronny/betterwhitelist/BetterWhitelist.java b/src/main/java/io/gitlab/jfronny/betterwhitelist/BetterWhitelist.java index 38647f3..94f44f1 100644 --- a/src/main/java/io/gitlab/jfronny/betterwhitelist/BetterWhitelist.java +++ b/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(); diff --git a/src/main/java/io/gitlab/jfronny/betterwhitelist/server/BetterWhitelistServer.java b/src/main/java/io/gitlab/jfronny/betterwhitelist/server/BetterWhitelistServer.java index 0d84f46..8fb9c9c 100644 --- a/src/main/java/io/gitlab/jfronny/betterwhitelist/server/BetterWhitelistServer.java +++ b/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 challenges = new HashMap<>(); - private static class Challenge { - public final Queue> 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 = challenges.get(gp); - - Scope fork = BetterWhitelist.SCOPE.fork(); - 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 " + gp.getName() + ": " + challengeString); - buf.writeString(challengeString); - List> 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); + Challenge challenge = new Challenge(gp, sender); + challenges.put(gp, challenge); + PacketByteBuf handshake = PacketByteBufs.create(); + handshake.writeInt(BetterWhitelist.PROTOCOL_VERSION); + challenge.sender.sendPacket(BetterWhitelist.HANDSHAKE_CHANNEL, handshake); + synchronizer.waitFor(challenge.challengeCompleted); }); - 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 context) { context.getSource().sendMessage(Text.literal("Loaded " + BetterWhitelist.MOD_METADATA.getName() + " " + BetterWhitelist.MOD_METADATA.getVersion())); return 1; diff --git a/src/main/java/io/gitlab/jfronny/betterwhitelist/server/Challenge.java b/src/main/java/io/gitlab/jfronny/betterwhitelist/server/Challenge.java new file mode 100644 index 0000000..9875e9e --- /dev/null +++ b/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> response = new ManualFuture<>(); + public final GameProfile profile; + public PacketSender sender; + + public Challenge(GameProfile profile, PacketSender sender) { + this.profile = profile; + this.sender = sender; + } +} diff --git a/src/main/java/io/gitlab/jfronny/betterwhitelist/server/ManualFuture.java b/src/main/java/io/gitlab/jfronny/betterwhitelist/server/ManualFuture.java new file mode 100644 index 0000000..466b24c --- /dev/null +++ b/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 implements Future { + 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 + } +} diff --git a/src/main/java/io/gitlab/jfronny/betterwhitelist/server/ServerScope.java b/src/main/java/io/gitlab/jfronny/betterwhitelist/server/ServerScope.java new file mode 100644 index 0000000..281c2b8 --- /dev/null +++ b/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> 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); + } + } +}