Better-Whitelist/src/main/java/io/gitlab/jfronny/betterwhitelist/server/BetterWhitelistServer.java

package io.gitlab.jfronny.betterwhitelist.server;

import com.mojang.authlib.GameProfile;
import com.mojang.brigadier.context.CommandContext;
import io.gitlab.jfronny.betterwhitelist.BetterWhitelist;
import io.gitlab.jfronny.betterwhitelist.DSerializer;
import io.gitlab.jfronny.betterwhitelist.packet.ChallengeResponsePacket;
import io.gitlab.jfronny.betterwhitelist.packet.HandshakePacket;
import io.gitlab.jfronny.betterwhitelist.server.mixin.ServerLoginNetworkHandlerAccessor;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.muscript.ast.context.Script;
import io.gitlab.jfronny.muscript.core.LocationalException;
import io.gitlab.jfronny.muscript.core.MuScriptVersion;
import io.gitlab.jfronny.muscript.parser.Parser;
import net.fabricmc.api.DedicatedServerModInitializer;
import net.fabricmc.fabric.api.command.v2.CommandRegistrationCallback;
import net.fabricmc.fabric.api.networking.v1.*;
import net.fabricmc.loader.api.FabricLoader;
import net.minecraft.network.PacketByteBuf;
import net.minecraft.network.packet.c2s.handshake.HandshakeC2SPacket;
import net.minecraft.server.command.ServerCommandSource;
import net.minecraft.server.network.ServerLoginNetworkHandler;
import net.minecraft.text.Text;
import net.minecraft.util.Util;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.*;

import static net.minecraft.server.command.CommandManager.literal;

public class BetterWhitelistServer implements DedicatedServerModInitializer {
    private String scriptSource;
    private Script script;
    private final Map<GameProfile, Challenge> challenges = new HashMap<>();

    @Override
    public void onInitializeServer() {
        BetterWhitelist.initialize();

        CommandRegistrationCallback.EVENT.register((dispatcher, registryAccess, environment) -> {
            dispatcher.register(
                    literal(BetterWhitelist.MOD_ID)
                            .requires(source -> source.hasPermissionLevel(4))
                            .executes(this::printVersion)
                            .then(literal("reload").executes(this::reloadScript))
            );
        });

        ServerLoginConnectionEvents.QUERY_START.register((handler, server, sender, synchronizer) -> {
            GameProfile gp = profile(handler);
            Challenge challenge = new Challenge(gp, sender);
            challenges.put(gp, challenge);
            PacketByteBuf buf = PacketByteBufs.create();
            HandshakePacket.CODEC.encode(buf, new HandshakePacket(BetterWhitelist.PROTOCOL_VERSION));
            challenge.sender.sendPacket(BetterWhitelist.HANDSHAKE_CHANNEL, buf);
            synchronizer.waitFor(challenge.challengeCompleted);
        });

        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;
            }
            LoginPacketSender rns = (LoginPacketSender) responseSender;
            HandshakePacket packet = HandshakePacket.CODEC.decode(buf);
            try {
                if (packet.version() != 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 = rns;
                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().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;
            }
            LoginPacketSender rns = (LoginPacketSender) responseSender;
            ChallengeResponsePacket packet = ChallengeResponsePacket.CODEC.decode(buf);
            try {
                Challenge ch = challenges.get(profile(handler));
                ch.sender = rns;
                BetterWhitelist.LOG.info("Got response from " + ch.profile.getName() + ": " + packet.message());
                ch.response.complete(DSerializer.deserialize(packet.message()));
            } catch (Throwable t) {
                BetterWhitelist.LOG.error("Failed login", t);
                handler.disconnect(Text.literal("Invalid dynamic"));
            }
        });

        try {
            reloadScript();
        } catch (IOException e) {
            throw new RuntimeException("Could not load whitelist script", e);
        }
    }

    private GameProfile profile(ServerLoginNetworkHandler handler) {
        GameProfile gp = ((ServerLoginNetworkHandlerAccessor) handler).getProfile();
        if (gp == null) throw new NullPointerException("Missing GameProfile");
        return gp;
    }

    private int printVersion(CommandContext<ServerCommandSource> context) {
        context.getSource().sendMessage(Text.literal("Loaded " + BetterWhitelist.MOD_METADATA.getName() + " " + BetterWhitelist.MOD_METADATA.getVersion()));
        return 1;
    }

    private int reloadScript(CommandContext<ServerCommandSource> context) {
        try {
            reloadScript();
            context.getSource().sendMessage(Text.literal("Successfully reloaded script"));
        } catch (Throwable t) {
            BetterWhitelist.LOG.error("Could not reload script", t);
            context.getSource().sendError(Text.literal("Could not reload script, check server log for details"));
        }
        return 1;
    }

    private void reloadScript() throws IOException {
        Path scriptPath = FabricLoader.getInstance()
                .getConfigDir()
                .resolve(BetterWhitelist.MOD_ID + ".mu");
        if (!Files.exists(scriptPath)) Files.writeString(scriptPath, """
                // Use this method to execute a closure on the client and get back the result
                clientVersion = challenge({ ->
                  // Note that closures sent to the client do not have access to things you declare elsewhere
                  mod('better-whitelist').version
                })
                
                println("You can, of course, use println-debugging")
                
                // You have access to the same methods on the server as you do on the client
                // You may use the assert method to short-circuit if you encounter a case where the client should not be allowed access
                // Assert can also have a second argument for the message to send if the assertion fails
                assert(mod('better-whitelist').version == clientVersion, 'You have the wrong mod version')
                
                // you can also send server-evaluated parameters with your challenge
                assert(challenge({ arg ->
                  arg::allMatch({ v -> mods::values()::anyMatch({ m -> v.id == m.id & v.version == m.version }) })
                }, mods::values()::filter({ v -> v.environment != 'server' })::map({ v -> { id = v.id, version = v.version } })))
                """);
        String s = Files.readString(scriptPath);
        this.script = Parser.parseScript(MuScriptVersion.DEFAULT, s);
        this.scriptSource = s;
    }
}