ServerBrowser: Match self lookup with public IP

common/encoding.go

@@ -1,6 +1,10 @@
 package common
 
-import "fmt"
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
 
 func Base32Encode(value int64) string {
 	alpha := "0123456789abcdefghijklmnopqrstuv"
@@ -26,3 +30,39 @@ func reverse(s string) string {
 
 	return string(rns)
 }
+
+func IPFormatToInt(ip string) (int32, uint16) {
+	port := 0
+
+	if strings.Contains(ip, ":") {
+		ipSplit := strings.Split(ip, ":")
+
+		var err error
+		port, err = strconv.Atoi(ipSplit[1])
+		if err != nil {
+			panic(err)
+		}
+
+		ip = ipSplit[0]
+	}
+
+	var intIP int
+	for i, s := range strings.Split(ip, ".") {
+		val, err := strconv.Atoi(s)
+		if err != nil {
+			panic(err)
+		}
+
+		intIP |= val << (24 - i*8)
+	}
+
+	// TODO: Check if this handles negative numbers properly
+	return int32(intIP), uint16(port)
+}
+
+func IPFormatToString(ip string) (string, string) {
+	intIP, intPort := IPFormatToInt(ip)
+
+	// TODO: Check if this handles negative numbers properly
+	return strconv.FormatInt(int64(intIP), 10), strconv.FormatUint(uint64(intPort), 10)
+}

serverbrowser/filter.go

@@ -20,20 +20,37 @@ var (
 	regexMatchmaking = regexp.MustCompile(`^dwc_mver = -?(\d{1,10}) and dwc_pid != (\d{1,10}) and maxplayers = -?(\d{1,10}) and numplayers < -?(\d{1,10}) and dwc_mtype = -?(\d{1,10}) and dwc_hoststate = (\d{1,10}) and dwc_suspend = (\d{1,10}) and \((.*)\)$`)
 )
 
-func FilterServers(servers []map[string]string, queryGame string, filter string) []map[string]string {
+func filterServers(servers []map[string]string, queryGame string, filter string, publicIP string) []map[string]string {
 	if match := regexSelfLookup.FindStringSubmatch(filter); match != nil {
 		dwc_pid := match[1]
 
+		filtered := []map[string]string{}
+
 		// Search for where the profile ID matches
 		for _, server := range servers {
 			if server["dwc_pid"] == dwc_pid {
 				logging.Info(ModuleName, "Self lookup from", aurora.Cyan(dwc_pid), "ok")
 				return []map[string]string{server}
 			}
+
+			// Alternatively, if the server hasn't set its dwc_pid field yet, we return servers matching the request's public IP.
+			// If multiple servers exist with the same public IP then the client will use the one with the matching port.
+			// This is a bit of a hack to speed up server creation.
+			if _, ok := server["dwc_pid"]; !ok && server["publicip"] == publicIP {
+				server["dwc_pid"] = dwc_pid
+				server["dwc_mtype"] = "0"
+				server["dwc_mver"] = "0"
+				filtered = append(filtered, server)
+			}
 		}
 
-		logging.Error(ModuleName, "Could not find server with dwc_pid", aurora.Cyan(dwc_pid))
-		return []map[string]string{}
+		if len(filtered) == 0 {
+			logging.Error(ModuleName, "Could not find server with dwc_pid", aurora.Cyan(dwc_pid))
+			return []map[string]string{}
+		}
+
+		logging.Info(ModuleName, "Self lookup for", aurora.Cyan(dwc_pid), "matched", aurora.BrightCyan(len(filtered)), "servers via public IP")
+		return filtered
 	}
 
 	if match := regexMatchmaking.FindStringSubmatch(filter); match != nil {

serverbrowser/server.go

@@ -38,11 +38,6 @@ const (
 	LimitResultCountOption  = 1 << 7 // 0x80 / 128
 )
 
-func FindServers(queryGame string, filter string) ([]map[string]string, error) {
-	// TODO: Handle gueryGame, filter
-	return FilterServers(qr2.GetSessionServers(), queryGame, filter), nil
-}
-
 func popString(buffer []byte, index int) (string, int) {
 	str := common.GetString(buffer[index:])
 	return str, index + len(str) + 1
@@ -94,6 +89,7 @@ func handleServerListRequest(conn net.Conn, buffer []byte) {
 		output = binary.BigEndian.AppendUint16(output, 6500)
 
 		// Write the encrypted reply
+		// TODO: Key is currently hardcoded to Mario Kart Wii
 		conn.Write(common.EncryptTypeX([]byte("9r3Rmy"), challenge, output))
 		return
 	}
@@ -115,10 +111,8 @@ func handleServerListRequest(conn net.Conn, buffer []byte) {
 	}
 	output = append(output, 0x00) // Zero length string to end the list
 
-	servers, err := FindServers(queryGame, filter)
-	if err != nil {
-		panic(err)
-	}
+	publicIP, _ := common.IPFormatToString(conn.RemoteAddr().String())
+	servers := filterServers(qr2.GetSessionServers(), queryGame, filter, publicIP)
 
 	for _, server := range servers {
 		var flags byte