Port as int, default host in config and separate implementation of handler

src/main/java/me/uniodex/velocityrcon/commandsource/IRconCommandSource.java → src/main/java/me/uniodex/velocityrcon/RconCommandSource.java

@@ -1,5 +1,6 @@
-package me.uniodex.velocityrcon.commandsource;
+package me.uniodex.velocityrcon;
 
+import com.velocitypowered.api.command.CommandSource;
 import com.velocitypowered.api.permission.PermissionFunction;
 import com.velocitypowered.api.permission.Tristate;
 import com.velocitypowered.api.proxy.ProxyServer;
@@ -13,15 +14,15 @@
 
 import static com.velocitypowered.api.permission.PermissionFunction.ALWAYS_TRUE;
 
-public class IRconCommandSource implements RconCommandSource {
+public class RconCommandSource implements CommandSource {
 
     private final StringBuffer buffer = new StringBuffer();
-    private PermissionFunction permissionFunction = ALWAYS_TRUE;
+    private final PermissionFunction permissionFunction = ALWAYS_TRUE;
 
     @Getter
-    private ProxyServer server;
+    private final ProxyServer server;
 
-    public IRconCommandSource(ProxyServer server) {
+    public RconCommandSource(ProxyServer server) {
         this.server = server;
     }
 

src/main/java/me/uniodex/velocityrcon/VelocityRcon.java

@@ -12,7 +12,6 @@
 import io.netty.channel.ChannelFuture;
 import lombok.Getter;
 import me.uniodex.velocityrcon.server.RconServer;
-import me.uniodex.velocityrcon.utils.Utils;
 import org.slf4j.Logger;
 
 import java.io.File;
@@ -29,8 +28,12 @@ public class VelocityRcon {
     @Getter
     private final Logger logger;
 
+    private static final String DEFAULT_HOST = "127.0.0.1";
+
+    private final Toml toml;
+
     @Getter
-    private String rconHost = "127.0.0.1";
+    private String rconHost = DEFAULT_HOST;
     @Getter
     private int rconPort = 1337;
     @Getter
@@ -47,25 +50,13 @@ public VelocityRcon(ProxyServer server, Logger logger, @DataDirectory final Path
         this.logger = logger;
         instance = this;
 
-        Toml toml = loadConfig(folder);
+        toml = loadToml(folder);
         if (toml == null) {
             logger.warn("Failed to load rcon.toml. Shutting down.");
             return;
         }
 
-        if (Utils.isInteger(toml.getString("rcon-port"))) {
-            rconPort = Integer.valueOf(toml.getString("rcon-port"));
-        } else {
-            logger.warn("Invalid rcon port. Shutting down.");
-            return;
-        }
-        rconHost = toml.getString("rcon-host");
-        if (rconHost == null) {
-            logger.warn("rcon-host is not specified in the config! 127.0.0.1 will be used.");
-            rconHost = "127.0.0.1";
-        }
-        rconPassword = toml.getString("rcon-password");
-        rconColored = toml.getBoolean("rcon-colored");
+        loadDataFromConfig();
     }
 
     @Subscribe
@@ -80,10 +71,9 @@ public void onShutdown(ProxyShutdownEvent event) {
 
     private void startListener() {
         InetSocketAddress address = new InetSocketAddress(rconHost, rconPort);
-        rconServer = new RconServer(server, rconPassword);
-        logger.info("Binding rcon to address: /" + address.getHostName() + ":" + address.getPort());
+        rconServer = new RconServer(address, rconPassword, new VelocityRconCommandHandler(server, logger));
 
-        ChannelFuture future = rconServer.bind(address);
+        ChannelFuture future = rconServer.bind();
         Channel channel = future.awaitUninterruptibly().channel();
 
         if (!channel.isActive()) {
@@ -93,13 +83,28 @@ private void startListener() {
 
     private void stopListener() {
         if (rconServer != null) {
-            logger.info("Trying to stop RCON listener");
-
             rconServer.shutdown();
         }
     }
 
-    private Toml loadConfig(Path path) {
+    private void loadDataFromConfig() {
+        try {
+            rconPort = toml.getLong("rcon-port").intValue();
+        } catch (ClassCastException ignored) {
+            try {
+                rconPort = Integer.parseInt(toml.getString("rcon-port"));
+            } catch (ClassCastException ignored2) {
+                logger.warn("Invalid rcon port. Shutting down.");
+                return;
+            }
+        }
+
+        rconHost = toml.getString("rcon-host", DEFAULT_HOST);
+        rconPassword = toml.getString("rcon-password");
+        rconColored = toml.getBoolean("rcon-colored", false);
+    }
+
+    private Toml loadToml(Path path) {
         File folder = path.toFile();
         File file = new File(folder, "rcon.toml");
         if (!file.getParentFile().exists()) {

src/main/java/me/uniodex/velocityrcon/VelocityRconCommandHandler.java

@@ -0,0 +1,78 @@
+package me.uniodex.velocityrcon;
+
+import com.velocitypowered.api.proxy.ProxyServer;
+import me.uniodex.velocityrcon.server.RconHandler;
+import me.uniodex.velocityrcon.server.RconServer;
+import me.uniodex.velocityrcon.utils.Utils;
+import org.slf4j.Logger;
+
+import java.net.SocketAddress;
+
+public class VelocityRconCommandHandler implements RconHandler {
+    private final Logger logger;
+    private final ProxyServer proxyServer;
+
+    private final RconCommandSource commandSender;
+
+
+    public VelocityRconCommandHandler(ProxyServer proxyServer, Logger logger) {
+        this.proxyServer = proxyServer;
+        this.logger = logger;
+        this.commandSender = new RconCommandSource(proxyServer);
+    }
+
+    @Override
+    public void onBind(RconServer server) {
+        logger.info("Binding RCON to address: /" + server.getAddress().getHostName() + ":" + server.getAddress().getPort());
+    }
+
+    @Override
+    public void onClientConnected(RconServer server, SocketAddress clientAddress) {
+        logger.info("RCON connection from [{}]", clientAddress);
+    }
+
+    @Override
+    public void onClientDisconnected(RconServer server, SocketAddress clientAddress) {
+        logger.info("RCON client [{}] disconnected ", clientAddress);
+    }
+
+    @Override
+    public String processData(RconServer server, String payload, SocketAddress clientAddress) {
+        boolean success;
+        String message;
+
+        if (payload.equalsIgnoreCase("end") || payload.equalsIgnoreCase("stop")) {
+            success = true;
+            message = "Shutting down the proxy...";
+            proxyServer.shutdown();
+        } else {
+            try {
+                success = proxyServer.getCommandManager().executeAsync(commandSender, payload).join();
+                if (success) {
+                    message = commandSender.flush();
+                } else {
+                    message = "No such command";
+                }
+            } catch (Exception e) {
+                e.printStackTrace();
+                success = false;
+                message = "Unknown error";
+            }
+        }
+
+        if (!success) {
+            message = String.format("Error executing: %s (%s)", payload, message);
+        }
+
+        if (!VelocityRcon.getInstance().isRconColored()) {
+            message = Utils.stripColor(message);
+        }
+
+        return message;
+    }
+
+    @Override
+    public void onShutdown(RconServer server) {
+        logger.info("Stopping RCON listener");
+    }
+}

src/main/java/me/uniodex/velocityrcon/server/RconChannelHandler.java

@@ -0,0 +1,108 @@
+package me.uniodex.velocityrcon.server;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.SimpleChannelInboundHandler;
+
+import java.nio.ByteOrder;
+import java.nio.charset.StandardCharsets;
+
+public class RconChannelHandler extends SimpleChannelInboundHandler<ByteBuf> {
+
+    private static final byte FAILURE = -1;
+    private static final byte TYPE_RESPONSE = 0;
+    private static final byte TYPE_COMMAND = 2;
+    private static final byte TYPE_LOGIN = 3;
+
+    private final String password;
+
+    private boolean loggedIn = false;
+
+    /**
+     * The {@link RconServer} this handler belongs to.
+     */
+    private final RconServer rconServer;
+
+    @Override
+    public void channelUnregistered(ChannelHandlerContext ctx) {
+        rconServer.getHandler().onClientDisconnected(rconServer, ctx.channel().remoteAddress());
+    }
+
+    public RconChannelHandler(RconServer rconServer, String password) {
+        this.rconServer = rconServer;
+        this.password = password;
+    }
+
+    @Override
+    protected void channelRead0(ChannelHandlerContext ctx, ByteBuf buf) {
+        buf = buf.order(ByteOrder.LITTLE_ENDIAN);
+        if (buf.readableBytes() < 8) {
+            return;
+        }
+
+        int requestId = buf.readInt();
+        int type = buf.readInt();
+
+        byte[] payloadData = new byte[buf.readableBytes() - 2];
+        buf.readBytes(payloadData);
+        String payload = new String(payloadData, StandardCharsets.UTF_8);
+
+        buf.readBytes(2); // two byte padding
+
+        if (type == TYPE_LOGIN) {
+            handleLogin(ctx, payload, requestId);
+        } else if (type == TYPE_COMMAND) {
+            handleCommand(ctx, payload, requestId);
+        } else {
+            sendLargeResponse(ctx, requestId, "Unknown request " + Integer.toHexString(type));
+        }
+    }
+
+    private void handleLogin(ChannelHandlerContext ctx, String payload, int requestId) {
+        if (password.equals(payload)) {
+            loggedIn = true;
+
+            sendResponse(ctx, requestId, TYPE_COMMAND, "");
+
+            rconServer.getHandler().onClientConnected(rconServer, ctx.channel().remoteAddress());
+        } else {
+            loggedIn = false;
+            sendResponse(ctx, FAILURE, TYPE_COMMAND, "");
+        }
+    }
+
+    private void handleCommand(ChannelHandlerContext ctx, String payload, int requestId) {
+        if (!loggedIn) {
+            sendResponse(ctx, FAILURE, TYPE_COMMAND, "");
+            return;
+        }
+        String message = rconServer.getHandler().processData(rconServer, payload, ctx.channel().remoteAddress());
+        sendLargeResponse(ctx, requestId, message);
+    }
+
+    private void sendResponse(ChannelHandlerContext ctx, int requestId, int type, String payload) {
+        ByteBuf buf = ctx.alloc().buffer().order(ByteOrder.LITTLE_ENDIAN);
+        buf.writeInt(requestId);
+        buf.writeInt(type);
+        buf.writeBytes(payload.getBytes(StandardCharsets.UTF_8));
+        buf.writeByte(0);
+        buf.writeByte(0);
+        ctx.write(buf);
+    }
+
+    private void sendLargeResponse(ChannelHandlerContext ctx, int requestId, String payload) {
+        if (payload.length() == 0) {
+            sendResponse(ctx, requestId, TYPE_RESPONSE, "");
+            return;
+        }
+
+        int start = 0;
+        while (start < payload.length()) {
+            int length = payload.length() - start;
+            int truncated = Math.min(length, 2048);
+
+            sendResponse(ctx, requestId, TYPE_RESPONSE, payload.substring(start, truncated));
+            start += truncated;
+        }
+    }
+}