AddressParser-related changes and fixes

build.gradle

@@ -142,6 +142,8 @@ dependencies {
 
     includeJ8("com.viaversion:viaversion:${rootProject.viaver_version}")
     include("com.github.TinfoilMC:ClientCommands:1.1.0")
+
+    testImplementation("org.testng:testng:6.13.1")
 }
 
 remapJar {
@@ -169,6 +171,11 @@ processResources {
     }
 }
 
+test {
+    // Minimal insurance to make sure we get expected results.
+    useTestNG()
+}
+
 List<String> mcReleases = Arrays.stream(rootProject.publish_mc_versions.toString().split(","))
         .map({ it -> it.trim() })
         .collect(Collectors.toList())

src/main/java/com/viaversion/fabric/common/AddressParser.java

@@ -20,85 +20,286 @@
 import com.google.common.collect.Lists;
 import com.google.common.primitives.Ints;
 import com.viaversion.viaversion.api.protocol.version.ProtocolVersion;
+import com.viaversion.viaversion.libs.fastutil.ints.IntArrayList;
+import com.viaversion.viaversion.libs.fastutil.ints.IntImmutableList;
+import com.viaversion.viaversion.libs.fastutil.ints.IntList;
 import org.apache.commons.lang3.StringUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.VisibleForTesting;
 
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 // Based on VIAaaS parser
 public class AddressParser {
-    public Integer protocol;
-    public String viaSuffix;
-    public String serverAddress;
-    public String viaOptions;
+    private static final String VIA = "viafabric";
+    private static final Pattern DOT_SUFFIX = Pattern.compile("\\.$");
 
-    public AddressParser parse(String address) {
-        return parse(address, "viafabric");
+    // Retaining a list for now; not exposing it to the end consumer.
+    // An idea tossed about is to retry the server for the given protocols.
+    @Nullable
+    @VisibleForTesting
+    final IntList protocols;
+    @NotNull
+    private final String serverAddress;
+    @Nullable
+    private final Integer port;
+
+    private AddressParser(@NotNull String address, @Nullable Integer port) {
+        this(address, null, port);
     }
 
-    public String getSuffixWithOptions() {
-        if (viaOptions != null && !viaOptions.isEmpty()) {
-            return viaOptions + "." + viaSuffix;
-        }
-        return viaSuffix;
+    private AddressParser(@NotNull String address, @Nullable IntList protocols, @Nullable Integer port) {
+        this.serverAddress = address;
+        this.port = port;
+        this.protocols = protocols;
+    }
+
+    @NotNull
+    public static AddressParser parse(String address) {
+        return parse(address, VIA);
     }
 
-    public AddressParser parse(String address, String viaHostName) {
+    @NotNull
+    private static AddressParser parse(String address, String viaHostName) {
+        int portIndex = address.lastIndexOf(':');
+        Integer port = portIndex >= 0 ? Ints.tryParse(address.substring(portIndex + 1)) : null;
+
+        if (port != null && port >= 0 && port < 65536) {
+            if (address.charAt(portIndex - 1) == '.') {
+                // Let's not allocate an intermediate string.
+                portIndex -= 1;
+            } else if (port < 10000 &&
+                    // I don't like this but there's not really a better way of doing this
+                    address.lastIndexOf(':', portIndex - 1) > Math.max(
+                            address.lastIndexOf(']', portIndex - 1),
+                            address.lastIndexOf('.', portIndex - 1)
+                    )) {
+                // We parsed an IPv6, a port isn't ideal here.
+                port = null;
+            }
+
+            // Keeping a sane flow control.
+            if (port != null) {
+                // Truncate the port off, as that interferes.
+                address = address.substring(0, portIndex);
+            }
+        }
+
         address = StringUtils.removeEnd(address, ".");
-        String suffixRemoved = StringUtils.removeEnd(address, "." + viaHostName);
 
-        if (suffixRemoved.equals(address)) {
-            serverAddress = address;
-            return this;
+        String truncated = StringUtils.removeEnd(address, '.' + viaHostName);
+        if (!address.equals(truncated)) {
+            return parseSuffix(truncated, port);
+        }
+
+        truncated = StringUtils.removeStart(address, viaHostName + '.');
+        if (!address.equals(truncated)) {
+            final AddressParser addr = parsePrefix(truncated, port);
+            if (addr != null) {
+                return addr;
+            }
         }
 
+        return new AddressParser(address, port);
+    }
+
+    @NotNull
+    private static AddressParser parseSuffix(String address, Integer port) {
         boolean stopOptions = false;
-        List<String> optionsParts = new ArrayList<>();
+        IntList protocolParts = new IntArrayList();
         List<String> serverParts = new ArrayList<>();
 
-        for (String part : Lists.reverse(Arrays.asList(suffixRemoved.split(Pattern.quote("."))))) {
-            if (!stopOptions && parseOption(part)) {
-                optionsParts.add(part);
+        Integer protocol;
+        for (String part : Lists.reverse(Arrays.asList(address.split("\\.")))) {
+            if (!stopOptions && (protocol = parseSuffixOption(part)) != null) {
+                protocolParts.add(protocol.intValue());
                 continue;
             }
             stopOptions = true;
             serverParts.add(part);
         }
 
-        serverAddress = String.join(".", Lists.reverse(serverParts));
-        viaOptions = String.join(".", Lists.reverse(optionsParts));
-        viaSuffix = viaHostName;
+        return new AddressParser(
+                String.join(".", Lists.reverse(serverParts)),
+                new IntImmutableList(Lists.reverse(protocolParts)),
+                port
+        );
+    }
+
+    // Fail condition = returns null; caller must fall through.
+    @Nullable
+    private static AddressParser parsePrefix(String address, Integer port) {
+        IntList protocols = new IntArrayList();
+        int index = 0, lastIndex, colonIndex = address.indexOf(';');
+
+        if (colonIndex < 0) {
+            return null;
+        }
+
+        while ((index = address.indexOf('+', lastIndex = index)) >= 0 && index < colonIndex) {
+            parseAndAdd(address.substring(lastIndex, index), protocols);
+            index++;
+        }
+
+        parseAndAdd(address.substring(lastIndex, colonIndex), protocols);
+
+        return new AddressParser(
+                address.substring(colonIndex + 1),
+                new IntImmutableList(protocols),
+                port
+        );
+    }
+
+    private static void parseAndAdd(String part, IntList protocols) {
+        final Integer protocol = parseSchemeOption(part);
+        if (protocol != null) {
+            protocols.add(protocol.intValue());
+        }
+    }
+
+    public String getSuffixWithOptions() {
+        if (protocols == null) {
+            return "";
+        }
+        if (protocols.isEmpty()) {
+            return VIA;
+        }
+        return protocols.intStream()
+                .mapToObj(AddressParser::toProtocolName)
+                .map(str -> str.replace('.', '_'))
+                .collect(Collectors.joining("._v", "_v", "." + VIA));
+    }
+
+    public String getPrefixWithOptions() {
+        if (protocols == null) {
+            return "";
+        }
+        if (protocols.isEmpty()) {
+            return VIA;
+        }
+        return protocols.intStream()
+                .mapToObj(AddressParser::toProtocolName)
+                .collect(Collectors.joining("+v", VIA + ".v", ""));
+    }
+
+    public boolean hasViaMetadata() {
+        return protocols != null;
+    }
+
+    public boolean hasProtocol() {
+        return protocols != null && !protocols.isEmpty();
+    }
+
+    public Integer protocol() {
+        if (protocols != null && !protocols.isEmpty()) {
+            return protocols.getInt(0);
+        }
+        return null;
+    }
+
+    @NotNull
+    public String serverAddress() {
+        return this.serverAddress;
+    }
+
+    @Nullable
+    public Integer port() {
+        return this.port;
+    }
+
+    public String toAddress() {
+        if (port != null) {
+            return serverAddress + ':' + port;
+        }
+        return serverAddress;
+    }
+
+    public String toSuffixedViaAddress() {
+        final String address = addAddressSuffix(serverAddress);
+        if (port != null) {
+            return address + ':' + port;
+        }
+        return address;
+    }
+
+    public InetAddress resolve() throws UnknownHostException {
+        return this.addAddressSuffix(InetAddress.getByName(serverAddress));
+    }
 
-        return this;
+    public InetSocketAddress addAddressSuffix(InetSocketAddress address) throws UnknownHostException {
+        return new InetSocketAddress(addAddressSuffix(address.getAddress()), address.getPort());
     }
 
-    public boolean parseOption(String part) {
+    public InetAddress addAddressSuffix(InetAddress address) throws UnknownHostException {
+        return InetAddress.getByAddress(addAddressSuffix(address.getHostName()), address.getAddress());
+    }
+
+    public String addAddressSuffix(String input) {
+        if (!this.hasViaMetadata()) {
+            return input;
+        }
+        return DOT_SUFFIX.matcher(input).replaceAll("") + '.' + this.getSuffixWithOptions();
+    }
+
+    private static Integer parseSuffixOption(String part) {
         String option;
         if (part.length() < 2) {
-            return false;
+            return null;
         } else if (part.startsWith("_")) {
             option = String.valueOf(part.charAt(1));
         } else if (part.charAt(1) == '_') {
             option = String.valueOf(part.charAt(0));
         } else {
-            return false;
+            return null;
         }
 
         String arg = part.substring(2);
         if ("v".equals(option)) {
-            parseProtocol(arg);
+            return parseProtocol(arg);
+        }
+
+        return null;
+    }
+
+    private static Integer parseSchemeOption(String part) {
+        if (part.length() < 2) {
+            return null;
         }
+        if (!part.startsWith("v")) {
+            return null;
+        }
+        return parseProtocol(part.substring(1));
+    }
 
-        return true;
+    private static Integer parseProtocol(String arg) {
+        final Integer protocol = Ints.tryParse(arg);
+        if (protocol != null) {
+            return protocol;
+        }
+        ProtocolVersion ver = ProtocolVersion.getClosest(arg.replace('_', '.'));
+        if (ver != null) {
+            return ver.getVersion();
+        }
+        return null;
     }
 
-    public void parseProtocol(String arg) {
-        protocol = Ints.tryParse(arg);
-        if (protocol == null) {
-            ProtocolVersion ver = ProtocolVersion.getClosest(arg.replace("_", "."));
-            if (ver != null) protocol = ver.getVersion();
+    private static String toProtocolName(int protocol) {
+        if (protocol < 0 || !ProtocolVersion.isRegistered(protocol)) {
+            return Integer.toString(protocol);
         }
+        return ProtocolVersion.getProtocol(protocol).getIncludedVersions().iterator().next();
+    }
+
+    @Override
+    public String toString() {
+        return "AddressParser{" + this.toSuffixedViaAddress() + '}';
     }
 }

src/main/java/com/viaversion/fabric/common/protocol/HostnameParserProtocol.java

@@ -20,10 +20,9 @@
 import com.viaversion.fabric.common.AddressParser;
 import com.viaversion.viaversion.api.protocol.AbstractSimpleProtocol;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
+import com.viaversion.viaversion.api.protocol.packet.State;
 import com.viaversion.viaversion.api.protocol.remapper.PacketHandlers;
 import com.viaversion.viaversion.api.protocol.remapper.ValueTransformer;
-import com.viaversion.viaversion.api.type.Type;
-import com.viaversion.viaversion.api.protocol.packet.State;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.protocols.base.ServerboundHandshakePackets;
 
@@ -39,7 +38,7 @@ protected void register() {
                 map(Types.STRING, new ValueTransformer<String, String>(Types.STRING) {
                     @Override
                     public String transform(PacketWrapper packetWrapper, String s) {
-                        return new AddressParser().parse(s).serverAddress;
+                        return AddressParser.parse(s).serverAddress();
                     }
                 });
             }

src/main/java/com/viaversion/fabric/common/provider/AbstractFabricVersionProvider.java

@@ -28,7 +28,6 @@
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
 import com.viaversion.viaversion.api.protocol.packet.State;
 import com.viaversion.viaversion.api.protocol.version.ProtocolVersion;
-import com.viaversion.viaversion.api.type.Type;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.exception.CancelException;
 import com.viaversion.viaversion.protocol.version.BaseVersionProvider;
@@ -42,7 +41,11 @@
 import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
-import java.util.*;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
 import java.util.concurrent.CompletableFuture;
 import java.util.logging.Logger;
 import java.util.stream.IntStream;
@@ -93,8 +96,8 @@ public ProtocolVersion getClosestServerProtocol(UserConnection connection) throw
             SocketAddress addr = connection.getChannel().remoteAddress();
 
             if (addr instanceof InetSocketAddress) {
-                AddressParser parser = new AddressParser();
-                Integer addrVersion = parser.parse(((InetSocketAddress) addr).getHostName()).protocol;
+                AddressParser parser = AddressParser.parse(((InetSocketAddress) addr).getHostName());
+                Integer addrVersion = parser.protocol();
                 if (addrVersion != null) {
                     serverVer = addrVersion;
                 }