From 86b0403dafc800dd83e13d41b3ceb5bef9136b5d Mon Sep 17 00:00:00 2001 From: mkwcat Date: Sat, 20 Jan 2024 18:05:32 -0500 Subject: [PATCH] QR2: Resend client messages if no ack received --- go.mod | 3 +- go.sum | 2 + qr2/main.go | 1 + qr2/message.go | 182 ++++++++++++++++++++++++---------------- qr2/session.go | 27 +++--- serverbrowser/server.go | 5 +- 6 files changed, 135 insertions(+), 85 deletions(-) diff --git a/go.mod b/go.mod index 1c3f75d..0818363 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index 972db14..c883923 100644 --- a/go.sum +++ b/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= diff --git a/qr2/main.go b/qr2/main.go index 866df92..ac8057d 100644 --- a/qr2/main.go +++ b/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: diff --git a/qr2/message.go b/qr2/message.go index c374ee0..c653a27 100644 --- a/qr2/message.go +++ b/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 = "" - } - - 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 = "" + } + + 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)) diff --git a/qr2/session.go b/qr2/session.go index e6880cf..0e92a67 100644 --- a/qr2/session.go +++ b/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{}, } } diff --git a/serverbrowser/server.go b/serverbrowser/server.go index 25026ed..c9e8c47 100644 --- a/serverbrowser/server.go +++ b/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:]) }