GS4 Query handler (#20)

proxy/src/main/java/com/velocitypowered/network/ConnectionManager.java

@@ -5,14 +5,13 @@ import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.client.HandshakeSessionHandler;
 import com.velocitypowered.proxy.protocol.ProtocolConstants;
 import com.velocitypowered.proxy.protocol.StateRegistry;
+import com.velocitypowered.proxy.protocol.netty.GS4QueryHandler;
 import com.velocitypowered.proxy.protocol.netty.LegacyPingDecoder;
 import com.velocitypowered.proxy.protocol.netty.LegacyPingEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftVarintFrameDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftVarintLengthEncoder;
-import org.apache.logging.log4j.Logger;
-import org.apache.logging.log4j.LogManager;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.bootstrap.ServerBootstrap;
 import io.netty.channel.Channel;
@@ -21,15 +20,20 @@ import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.EventLoopGroup;
 import io.netty.channel.epoll.Epoll;
+import io.netty.channel.epoll.EpollDatagramChannel;
 import io.netty.channel.epoll.EpollEventLoopGroup;
 import io.netty.channel.epoll.EpollServerSocketChannel;
 import io.netty.channel.epoll.EpollSocketChannel;
 import io.netty.channel.nio.NioEventLoopGroup;
+import io.netty.channel.socket.DatagramChannel;
 import io.netty.channel.socket.ServerSocketChannel;
 import io.netty.channel.socket.SocketChannel;
+import io.netty.channel.socket.nio.NioDatagramChannel;
 import io.netty.channel.socket.nio.NioServerSocketChannel;
 import io.netty.channel.socket.nio.NioSocketChannel;
 import io.netty.handler.timeout.ReadTimeoutHandler;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 import java.net.InetSocketAddress;
 import java.util.HashSet;
@@ -54,6 +58,7 @@ public final class ConnectionManager {
     private final Set<Channel> endpoints = new HashSet<>();
     private final Class<? extends ServerSocketChannel> serverSocketChannelClass;
     private final Class<? extends SocketChannel> socketChannelClass;
+    private final Class<? extends DatagramChannel> datagramChannelClass;
     private final EventLoopGroup bossGroup;
     private final EventLoopGroup workerGroup;
 
@@ -62,11 +67,13 @@ public final class ConnectionManager {
         if (epoll) {
             this.serverSocketChannelClass = EpollServerSocketChannel.class;
             this.socketChannelClass = EpollSocketChannel.class;
+            this.datagramChannelClass = EpollDatagramChannel.class;
             this.bossGroup = new EpollEventLoopGroup(0, createThreadFactory("Netty Epoll Boss #%d"));
             this.workerGroup = new EpollEventLoopGroup(0, createThreadFactory("Netty Epoll Worker #%d"));
         } else {
             this.serverSocketChannelClass = NioServerSocketChannel.class;
             this.socketChannelClass = NioSocketChannel.class;
+            this.datagramChannelClass = NioDatagramChannel.class;
             this.bossGroup = new NioEventLoopGroup(0, createThreadFactory("Netty Nio Boss #%d"));
             this.workerGroup = new NioEventLoopGroup(0, createThreadFactory("Netty Nio Worker #%d"));
         }
@@ -120,6 +127,24 @@ public final class ConnectionManager {
                 });
     }
 
+    public void queryBind(final String hostname, final int port) {
+        new Bootstrap()
+                .channel(datagramChannelClass)
+                .group(this.workerGroup)
+                .handler(new GS4QueryHandler())
+                .localAddress(hostname, port)
+                .bind()
+                .addListener((ChannelFutureListener) future -> {
+                    final Channel channel = future.channel();
+                    if (future.isSuccess()) {
+                        this.endpoints.add(channel);
+                        logger.info("Listening for GS4 query on {}", channel.localAddress());
+                    } else {
+                        logger.error("Can't bind to {}", channel.localAddress(), future.cause());
+                    }
+                });
+    }
+
     public Bootstrap createWorker() {
         return new Bootstrap()
                 .channel(this.socketChannelClass)

proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java

@@ -9,7 +9,6 @@ import com.velocitypowered.api.proxy.ProxyServer;
 import com.velocitypowered.natives.util.Natives;
 import com.velocitypowered.network.ConnectionManager;
 import com.velocitypowered.proxy.config.VelocityConfiguration;
-import com.velocitypowered.proxy.connection.MinecraftConnection;
 import com.velocitypowered.proxy.connection.client.ConnectedPlayer;
 import com.velocitypowered.proxy.connection.http.NettyHttpClient;
 import com.velocitypowered.api.server.ServerInfo;
@@ -97,6 +96,10 @@ public class VelocityServer implements ProxyServer {
         httpClient = new NettyHttpClient(this);
 
         this.cm.bind(configuration.getBind());
+
+        if (configuration.isQueryEnabled()) {
+            this.cm.queryBind(configuration.getBind().getHostString(), configuration.getQueryPort());
+        }
     }
 
     public ServerMap getServers() {

proxy/src/main/java/com/velocitypowered/proxy/config/VelocityConfiguration.java

@@ -32,12 +32,15 @@ public class VelocityConfiguration {
     private final int compressionThreshold;
     private final int compressionLevel;
 
+    private final boolean queryEnabled;
+    private final int queryPort;
+
     private Component motdAsComponent;
 
     private VelocityConfiguration(String bind, String motd, int showMaxPlayers, boolean onlineMode,
                                   IPForwardingMode ipForwardingMode, Map<String, String> servers,
                                   List<String> attemptConnectionOrder, int compressionThreshold,
-                                  int compressionLevel) {
+                                  int compressionLevel, boolean queryEnabled, int queryPort) {
         this.bind = bind;
         this.motd = motd;
         this.showMaxPlayers = showMaxPlayers;
@@ -47,6 +50,8 @@ public class VelocityConfiguration {
         this.attemptConnectionOrder = attemptConnectionOrder;
         this.compressionThreshold = compressionThreshold;
         this.compressionLevel = compressionLevel;
+        this.queryEnabled = queryEnabled;
+        this.queryPort = queryPort;
     }
 
     public boolean validate() {
@@ -126,6 +131,14 @@ public class VelocityConfiguration {
         return AddressUtil.parseAddress(bind);
     }
 
+    public boolean isQueryEnabled() {
+        return queryEnabled;
+    }
+
+    public int getQueryPort() {
+        return queryPort;
+    }
+
     public String getMotd() {
         return motd;
     }
@@ -181,6 +194,8 @@ public class VelocityConfiguration {
                 ", attemptConnectionOrder=" + attemptConnectionOrder +
                 ", compressionThreshold=" + compressionThreshold +
                 ", compressionLevel=" + compressionLevel +
+                ", queryEnabled=" + queryEnabled +
+                ", queryPort=" + queryPort +
                 ", motdAsComponent=" + motdAsComponent +
                 '}';
     }
@@ -209,7 +224,9 @@ public class VelocityConfiguration {
                     ImmutableMap.copyOf(servers),
                     toml.getTable("servers").getList("try"),
                     toml.getTable("advanced").getLong("compression-threshold", 1024L).intValue(),
-                    toml.getTable("advanced").getLong("compression-level", -1L).intValue());
+                    toml.getTable("advanced").getLong("compression-level", -1L).intValue(),
+                    toml.getTable("query").getBoolean("enabled"),
+                    toml.getTable("query").getLong("port", 25577L).intValue());
         }
     }
 }

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/GS4QueryHandler.java

@@ -0,0 +1,189 @@
+package com.velocitypowered.proxy.protocol.netty;
+
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.collect.ImmutableSet;
+import com.velocitypowered.api.proxy.Player;
+import com.velocitypowered.proxy.VelocityServer;
+import com.velocitypowered.proxy.protocol.ProtocolConstants;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.SimpleChannelInboundHandler;
+import io.netty.channel.socket.DatagramPacket;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.net.InetAddress;
+import java.nio.charset.StandardCharsets;
+import java.util.Collection;
+import java.util.Set;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+
+public class GS4QueryHandler extends SimpleChannelInboundHandler<DatagramPacket> {
+    private static final Logger logger = LogManager.getLogger(GS4QueryHandler.class);
+
+    private final static short QUERY_MAGIC_FIRST = 0xFE;
+    private final static short QUERY_MAGIC_SECOND = 0xFD;
+    private final static byte QUERY_TYPE_HANDSHAKE = 0x09;
+    private final static byte QUERY_TYPE_STAT = 0x00;
+    private final static byte[] QUERY_RESPONSE_FULL_PADDING = new byte[] { 0x73, 0x70, 0x6C, 0x69, 0x74, 0x6E, 0x75, 0x6D, 0x00, (byte) 0x80, 0x00 };
+    private final static byte[] QUERY_RESPONSE_FULL_PADDING2 = new byte[] { 0x01, 0x70, 0x6C, 0x61, 0x79, 0x65, 0x72, 0x5F, 0x00, 0x00 };
+
+    // Contents to add into basic stat response. See ResponseWriter class below
+    private final static Set<String> QUERY_BASIC_RESPONSE_CONTENTS = ImmutableSet.of(
+            "hostname",
+            "gametype",
+            "map",
+            "numplayers",
+            "maxplayers",
+            "hostport",
+            "hostip"
+    );
+
+    private final static Cache<InetAddress, Integer> sessions = CacheBuilder.newBuilder()
+            .expireAfterWrite(30, TimeUnit.SECONDS)
+            .build();
+
+    @Override
+    protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket msg) throws Exception {
+        ByteBuf queryMessage = msg.content();
+        InetAddress senderAddress = msg.sender().getAddress();
+
+        // Verify query packet magic
+        if (queryMessage.readUnsignedByte() != QUERY_MAGIC_FIRST && queryMessage.readUnsignedByte() != QUERY_MAGIC_SECOND) {
+            throw new IllegalStateException("Invalid query packet magic");
+        }
+
+        // Read packet header
+        short type = queryMessage.readUnsignedByte();
+        int sessionId = queryMessage.readInt();
+
+        // Allocate buffer for response
+        ByteBuf queryResponse = ctx.alloc().buffer();
+        DatagramPacket responsePacket = new DatagramPacket(queryResponse, msg.sender());
+
+        switch(type) {
+            case QUERY_TYPE_HANDSHAKE: {
+                // Generate new challenge token and put it into the sessions cache
+                int challengeToken = ThreadLocalRandom.current().nextInt();
+                sessions.put(senderAddress, challengeToken);
+
+                // Respond with challenge token
+                queryResponse.writeByte(QUERY_TYPE_HANDSHAKE);
+                queryResponse.writeInt(sessionId);
+                writeString(queryResponse, Integer.toString(challengeToken));
+                break;
+            }
+
+            case QUERY_TYPE_STAT: {
+                // Check if query was done with session previously generated using a handshake packet
+                int challengeToken = queryMessage.readInt();
+                Integer session = sessions.getIfPresent(senderAddress);
+                if (session == null || session != challengeToken) {
+                    throw new IllegalStateException("Invalid challenge token");
+                }
+
+                // Check which query response client expects
+                if (queryMessage.readableBytes() != 0 && queryMessage.readableBytes() != 4) {
+                    throw new IllegalStateException("Invalid query packet");
+                }
+
+                // Packet header
+                queryResponse.writeByte(QUERY_TYPE_STAT);
+                queryResponse.writeInt(sessionId);
+
+                // Fetch information
+                VelocityServer server = VelocityServer.getServer();
+                Collection<Player> players = server.getAllPlayers();
+
+                // Start writing the response
+                ResponseWriter responseWriter = new ResponseWriter(queryResponse, queryMessage.readableBytes() == 0);
+                responseWriter.write("hostname", server.getConfiguration().getMotd());
+                responseWriter.write("gametype", "SMP");
+
+                responseWriter.write("game_id", "MINECRAFT");
+                responseWriter.write("version", ProtocolConstants.SUPPORTED_GENERIC_VERSION_STRING);
+                responseWriter.write("plugins", "");
+
+                responseWriter.write("map", "Velocity");
+                responseWriter.write("numplayers", players.size());
+                responseWriter.write("maxplayers", server.getConfiguration().getShowMaxPlayers());
+                responseWriter.write("hostport", server.getConfiguration().getBind().getPort());
+                responseWriter.write("hostip", server.getConfiguration().getBind().getHostString());
+
+                responseWriter.writePlayers(players);
+                break;
+            }
+
+            default: {
+                throw new IllegalStateException("Invalid query type: " + type);
+            }
+        }
+
+        // Send the response
+        ctx.writeAndFlush(responsePacket);
+    }
+
+    @Override
+    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
+        logger.warn("Error while trying to handle a query packet from {}", ctx.channel().remoteAddress(), cause);
+    }
+
+    private static void writeString(ByteBuf buf, String string) {
+        buf.writeCharSequence(string, StandardCharsets.ISO_8859_1);
+        buf.writeByte(0x00);
+    }
+
+    private static class ResponseWriter {
+        private final ByteBuf buf;
+        private final boolean isBasic;
+
+        ResponseWriter(ByteBuf buf, boolean isBasic) {
+            this.buf = buf;
+            this.isBasic = isBasic;
+
+            if (!isBasic) {
+                buf.writeBytes(QUERY_RESPONSE_FULL_PADDING);
+            }
+        }
+
+        // Writes k/v to stat packet body if this writer is initialized
+        // for full stat response. Otherwise this follows
+        // GS4QueryHandler#QUERY_BASIC_RESPONSE_CONTENTS to decide what
+        // to write into packet body
+        void write(String key, Object value) {
+            if (isBasic) {
+                // Basic contains only specific set of data
+                if (!QUERY_BASIC_RESPONSE_CONTENTS.contains(key)) {
+                    return;
+                }
+
+                // Special case hostport
+                if (key.equals("hostport")) {
+                    buf.writeShortLE((Integer) value);
+                } else {
+                    writeString(buf, value.toString());
+                }
+            } else {
+                writeString(buf, key);
+                writeString(buf, value.toString());
+            }
+        }
+
+        // Ends packet k/v body writing and writes stat player list to
+        // the packet if this writer is initialized for full stat response
+        void writePlayers(Collection<Player> players) {
+            if (isBasic) {
+                return;
+            }
+
+            // Ends the full stat key-value body with \0
+            buf.writeByte(0x00);
+
+            buf.writeBytes(QUERY_RESPONSE_FULL_PADDING2);
+            players.forEach(player -> writeString(buf, player.getUsername()));
+            buf.writeByte(0x00);
+        }
+    }
+}

proxy/src/main/resources/velocity.toml

@@ -38,4 +38,11 @@ try = [
 compression-threshold = 1024
 
 # How much compression should be done (from 0-9). The default is -1, which uses zlib's default level of 6.
-compression-level = -1
+compression-level = -1
+
+[query]
+# Whether to enable responding to GameSpy 4 query responses or not
+enabled = false
+
+# If query responding is enabled, on what port should query response listener listen on?
+port = 25577