From 9cf35793f2b449f859b27b6afe37bb4807e7592e Mon Sep 17 00:00:00 2001
From: mkwcat <cat@mkw.re>
Date: Tue, 2 Jan 2024 18:49:55 -0500
Subject: [PATCH] QR2: Implement group host migration

---
 qr2/group.go     | 42 +++++++++++++++++++++++++++++++++++++++---
 qr2/heartbeat.go | 11 +++++++++++
 qr2/session.go   |  2 +-
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/qr2/group.go b/qr2/group.go
index be1e8bc..345042b 100644
--- a/qr2/group.go
+++ b/qr2/group.go
@@ -181,14 +181,49 @@ func ProcessGPStatusUpdate(profileID uint32, senderIP uint64, status string) {
 			delete(groups, session.GroupPointer.GroupName)
 		} else if session.GroupPointer.Server == session {
 			logging.Notice("QR2", "Server down in group", aurora.Cyan(session.GroupPointer.GroupName))
-			session.GroupPointer.Server = nil
-			// TODO: Search for new host via dwc_hoststate
+			session.GroupPointer.findNewServer()
 		}
 
 		session.GroupPointer = nil
 	}
 }
 
+// findNewServer attempts to find the new server/host in the group when the current server goes down.
+// If no server is found, the group's server pointer is set to nil.
+// Expects the mutex to be locked.
+func (g *Group) findNewServer() {
+	server := (*Session)(nil)
+	serverJoinIndex := -1
+	for session := range g.Players {
+		if session.Data["dwc_hoststate"] != "2" {
+			continue
+		}
+
+		joinIndex, err := strconv.Atoi(session.Data["+joinindex"])
+		if err != nil {
+			continue
+		}
+
+		if server == nil || joinIndex < serverJoinIndex {
+			server = session
+			serverJoinIndex = joinIndex
+		}
+	}
+
+	g.Server = server
+	g.updateMatchType()
+}
+
+// updateMatchType updates the match type of the group based on the host's dwc_mtype value.
+// Expects the mutex to be locked.
+func (g *Group) updateMatchType() {
+	if g.Server == nil || g.Server.Data["dwc_mtype"] == "" {
+		return
+	}
+
+	g.MatchType = g.Server.Data["dwc_mtype"]
+}
+
 type GroupInfo struct {
 	GroupName   string                       `json:"id"`
 	GameName    string                       `json:"game"`
@@ -212,6 +247,7 @@ func GetGroups(gameName string) []GroupInfo {
 		groupInfo := GroupInfo{
 			GroupName:   group.GroupName,
 			GameName:    group.GameName,
+			MatchType:   "",
 			Suspend:     true,
 			ServerIndex: "",
 			MKWRegion:   "",
@@ -226,7 +262,7 @@ func GetGroups(gameName string) []GroupInfo {
 			groupInfo.MatchType = "unknown"
 		}
 
-		if groupInfo.MatchType == "private" && group.Server != nil {
+		if group.Server != nil {
 			groupInfo.ServerIndex = group.Server.Data["+joinindex"]
 		}
 
diff --git a/qr2/heartbeat.go b/qr2/heartbeat.go
index 20d261a..e68e86a 100644
--- a/qr2/heartbeat.go
+++ b/qr2/heartbeat.go
@@ -97,5 +97,16 @@ func heartbeat(moduleName string, conn net.PacketConn, addr net.Addr, buffer []b
 		sendClientExploit(moduleName, session)
 	}
 
+	mutex.Lock()
+	if session.GroupPointer != nil {
+		if session.GroupPointer.Server == nil {
+			session.GroupPointer.findNewServer()
+		} else {
+			// Update the match type if needed
+			session.GroupPointer.updateMatchType()
+		}
+	}
+	mutex.Unlock()
+
 	logging.Info(moduleName, "Heartbeat ok")
 }
diff --git a/qr2/session.go b/qr2/session.go
index b47ee3d..6650069 100644
--- a/qr2/session.go
+++ b/qr2/session.go
@@ -58,7 +58,7 @@ func removeSession(addr uint64) {
 		} else if session.GroupPointer.Server == session {
 			logging.Notice("QR2", "Server down in group", aurora.Cyan(session.GroupPointer.GroupName))
 			session.GroupPointer.Server = nil
-			// TODO: Search for new host via dwc_hoststate
+			session.GroupPointer.findNewServer()
 		}
 
 		session.GroupPointer = nil