QR2: Resend client messages if no ack received

go.mod

@@ -1,6 +1,6 @@
 module wwfc
 
-go 1.21
+go 1.21.1
 
 require (
 	github.com/jackc/pgx/v4 v4.18.1
@@ -21,4 +21,5 @@ require (
 	github.com/petermattis/goid v0.0.0-20180202154549-b0b1615b78e5 // indirect
 	golang.org/x/crypto v0.17.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
+	gvisor.dev/gvisor v0.0.0-20240119232905-7b151e25d076
 )

go.sum

@@ -206,4 +206,6 @@ gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gvisor.dev/gvisor v0.0.0-20240119232905-7b151e25d076 h1:LbaTr9qML03qYVNb18i2L5QYAf5Go7BoFILZQ3KXETs=
+gvisor.dev/gvisor v0.0.0-20240119232905-7b151e25d076/go.mod h1:10sU+Uh5KKNv1+2x2A0Gvzt8FjD3ASIhorV3YsauXhk=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=

qr2/main.go

@@ -123,6 +123,7 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) {
 		// In case ClientExploitReply is lost, this can be checked as well
 		// This would be sent either after the payload is downloaded, or the client is already patched
 		session.ExploitReceived = true
+		session.MessageAckWaker.Assert()
 		return
 
 	case KeepAliveRequest:

qr2/message.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+	"net"
 	"os"
 	"strconv"
 	"time"
@@ -11,6 +12,7 @@ import (
 	"wwfc/logging"
 
 	"github.com/logrusorgru/aurora/v3"
+	"gvisor.dev/gvisor/pkg/sleep"
 )
 
 func printHex(data []byte) string {
@@ -224,74 +226,7 @@ func SendClientMessage(senderIP string, destSearchID uint64, message []byte) {
 		logging.Error(moduleName, "Invalid message:", aurora.Cyan(printHex(message)))
 	}
 
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	destPid, ok := receiver.Data["dwc_pid"]
-	if !ok || destPid == "" {
-		destPid = "<UNKNOWN>"
-	}
-
-	destSessionID := receiver.SessionID
-	packetCount := receiver.PacketCount + 1
-	receiver.PacketCount = packetCount
-	destAddr := receiver.Addr
-
-	if isNatnegPacket {
-		cookie := binary.BigEndian.Uint32(message[0x2:0x6])
-		logging.Notice(moduleName, "Send NN cookie", aurora.Cyan(cookie), "to", aurora.BrightCyan(destPid))
-	} else {
-		cmd := message[8]
-		common.LogMatchCommand(moduleName, destPid, cmd, matchData)
-
-		if cmd == common.MatchReservation {
-			resvError := checkReservationAllowed(moduleName, sender, receiver, matchData.Reservation.MatchType)
-			if resvError != "ok" {
-				if resvError == "restricted" || resvError == "restricted_join" {
-					logging.Error(moduleName, "RESERVATION: Restricted player attempted to join a public match")
-
-					if sender.Login != nil && sender.Login.Restricted {
-						callback := sender.Login.GPErrorCallback
-						profileId := sender.Login.ProfileID
-
-						mutex.Unlock()
-						callback(profileId, resvError)
-						mutex.Lock()
-					}
-					if receiver.Login != nil && receiver.Login.Restricted {
-						callback := receiver.Login.GPErrorCallback
-						profileId := receiver.Login.ProfileID
-
-						mutex.Unlock()
-						callback(profileId, resvError)
-						mutex.Lock()
-					}
-				}
-				return
-			}
-
-			sender.Reservation = matchData
-			sender.ReservationID = receiver.SearchID
-		} else if cmd == common.MatchResvOK || cmd == common.MatchResvDeny || cmd == common.MatchResvWait {
-			if receiver.ReservationID != sender.SearchID || receiver.Reservation.Reservation == nil {
-				logging.Error(moduleName, "Destination has no reservation with the sender")
-				return
-			}
-
-			if receiver.Reservation.Version != matchData.Version {
-				logging.Error(moduleName, "Reservation version mismatch")
-				return
-			}
-
-			if cmd == common.MatchResvOK {
-				if !processResvOK(moduleName, matchData.Version, *receiver.Reservation.Reservation, *matchData.ResvOK, sender, receiver) {
-					return
-				}
-			} else {
-				receiver.ReservationID = 0
-			}
-		}
-	}
+	destSessionID, packetCount, destAddr := processClientMessage(moduleName, sender, receiver, message, isNatnegPacket, matchData)
 
 	payload := createResponseHeader(ClientMessageRequest, destSessionID)
 
@@ -299,12 +234,117 @@ func SendClientMessage(senderIP string, destSearchID uint64, message []byte) {
 	binary.BigEndian.PutUint32(payload[len(payload)-4:], packetCount)
 	payload = append(payload, message...)
 
-	_, err := masterConn.WriteTo(payload, destAddr)
-	if err != nil {
-		logging.Error(moduleName, "Error sending message:", err.Error())
+	receiver.MessageMutex.Lock()
+
+	s := sleep.Sleeper{}
+	s.AddWaker(&receiver.MessageAckWaker)
+	timeWaker := sleep.Waker{}
+	s.AddWaker(&timeWaker)
+
+	receiver.MessageAckWaker.Clear()
+	timeOutCount := 0
+	for {
+		time.AfterFunc(1*time.Second, func() {
+			timeWaker.Assert()
+		})
+
+		_, err := masterConn.WriteTo(payload, destAddr)
+		if err != nil {
+			logging.Error(moduleName, "Error sending message:", err.Error())
+		}
+
+		// Wait for an ack or timeout
+		switch s.Fetch(true) {
+		case &timeWaker:
+			timeOutCount++
+			if timeOutCount > 8 {
+				logging.Error(moduleName, "Timed out waiting for ack")
+				receiver.MessageMutex.Unlock()
+				return
+			}
+			break
+
+		default:
+			receiver.MessageMutex.Unlock()
+			return
+		}
 	}
 }
 
+func processClientMessage(moduleName string, sender, receiver *Session, message []byte, isNatnegPacket bool, matchData common.MatchCommandData) (destSessionID uint32, packetCount uint32, destAddr net.Addr) {
+	mutex.Lock()
+
+	destPid, ok := receiver.Data["dwc_pid"]
+	if !ok || destPid == "" {
+		destPid = "<UNKNOWN>"
+	}
+
+	destSessionID = receiver.SessionID
+	packetCount = receiver.PacketCount + 1
+	receiver.PacketCount = packetCount
+	destAddr = receiver.Addr
+
+	if isNatnegPacket {
+		mutex.Unlock()
+		cookie := binary.BigEndian.Uint32(message[0x2:0x6])
+		logging.Notice(moduleName, "Send NN cookie", aurora.Cyan(cookie), "to", aurora.BrightCyan(destPid))
+		return
+	}
+	defer mutex.Unlock()
+
+	cmd := message[8]
+	common.LogMatchCommand(moduleName, destPid, cmd, matchData)
+
+	if cmd == common.MatchReservation {
+		resvError := checkReservationAllowed(moduleName, sender, receiver, matchData.Reservation.MatchType)
+		if resvError != "ok" {
+			if resvError == "restricted" || resvError == "restricted_join" {
+				logging.Error(moduleName, "RESERVATION: Restricted player attempted to join a public match")
+
+				if sender.Login != nil && sender.Login.Restricted {
+					callback := sender.Login.GPErrorCallback
+					profileId := sender.Login.ProfileID
+
+					mutex.Unlock()
+					callback(profileId, resvError)
+					mutex.Lock()
+				}
+				if receiver.Login != nil && receiver.Login.Restricted {
+					callback := receiver.Login.GPErrorCallback
+					profileId := receiver.Login.ProfileID
+
+					mutex.Unlock()
+					callback(profileId, resvError)
+					mutex.Lock()
+				}
+			}
+			return
+		}
+
+		sender.Reservation = matchData
+		sender.ReservationID = receiver.SearchID
+	} else if cmd == common.MatchResvOK || cmd == common.MatchResvDeny || cmd == common.MatchResvWait {
+		if receiver.ReservationID != sender.SearchID || receiver.Reservation.Reservation == nil {
+			logging.Error(moduleName, "Destination has no reservation with the sender")
+			return
+		}
+
+		if receiver.Reservation.Version != matchData.Version {
+			logging.Error(moduleName, "Reservation version mismatch")
+			return
+		}
+
+		if cmd == common.MatchResvOK {
+			if !processResvOK(moduleName, matchData.Version, *receiver.Reservation.Reservation, *matchData.ResvOK, sender, receiver) {
+				return
+			}
+		} else {
+			receiver.ReservationID = 0
+		}
+	}
+	return
+}
+
 func sendClientExploit(moduleName string, sessionCopy Session) {
 	if len(sessionCopy.Login.GameCode) != 4 || !common.IsUppercaseAlphanumeric(sessionCopy.Login.GameCode) {
 		logging.Error(moduleName, "Invalid game code:", aurora.Cyan(sessionCopy.Login.GameCode))

qr2/session.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/logrusorgru/aurora/v3"
 	"github.com/sasha-s/go-deadlock"
+	"gvisor.dev/gvisor/pkg/sleep"
 )
 
 const (
@@ -31,8 +32,10 @@ type Session struct {
 	Endianness      byte // Some fields depend on the client's endianness
 	Data            map[string]string
 	PacketCount     uint32
-	ReservationID   uint64
 	Reservation     common.MatchCommandData
+	ReservationID   uint64
+	MessageMutex    deadlock.Mutex
+	MessageAckWaker sleep.Waker
 	GroupPointer    *Group
 }
 
@@ -103,16 +106,18 @@ func setSessionData(moduleName string, addr net.Addr, sessionId uint32, payload
 
 	if !sessionExists {
 		session = &Session{
-			SessionID:     sessionId,
-			Addr:          addr,
-			Challenge:     "",
-			Authenticated: false,
-			LastKeepAlive: time.Now().Unix(),
-			Endianness:    ClientNoEndian,
-			Data:          payload,
-			PacketCount:   0,
-			Reservation:   common.MatchCommandData{},
-			ReservationID: 0,
+			SessionID:       sessionId,
+			Addr:            addr,
+			Challenge:       "",
+			Authenticated:   false,
+			LastKeepAlive:   time.Now().Unix(),
+			Endianness:      ClientNoEndian,
+			Data:            payload,
+			PacketCount:     0,
+			Reservation:     common.MatchCommandData{},
+			ReservationID:   0,
+			MessageMutex:    deadlock.Mutex{},
+			MessageAckWaker: sleep.Waker{},
 		}
 	}
 

serverbrowser/server.go

@@ -4,7 +4,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"github.com/logrusorgru/aurora/v3"
 	"net"
 	"regexp"
 	"strconv"
@@ -12,6 +11,8 @@ import (
 	"wwfc/common"
 	"wwfc/logging"
 	"wwfc/qr2"
+
+	"github.com/logrusorgru/aurora/v3"
 )
 
 const (
@@ -339,5 +340,5 @@ func handleSendMessageRequest(conn net.Conn, buffer []byte) {
 
 	logging.Notice(ModuleName, "Send message from", aurora.BrightCyan(conn.RemoteAddr()), "to", aurora.Cyan(fmt.Sprintf("%012x", searchID)))
 
-	qr2.SendClientMessage(conn.RemoteAddr().String(), searchID, buffer[9:])
+	go qr2.SendClientMessage(conn.RemoteAddr().String(), searchID, buffer[9:])
 }