From 7ff2bfa4d97f1b80345ca7e0c1ce59431512b5df Mon Sep 17 00:00:00 2001 From: mkwcat Date: Thu, 7 Dec 2023 01:56:59 -0500 Subject: [PATCH] QR2: Session lookup maps for IP and search ID --- qr2/message.go | 206 ++++++++++++++++++++++++++++++++++ qr2/session.go | 291 ++++++------------------------------------------- 2 files changed, 241 insertions(+), 256 deletions(-) create mode 100644 qr2/message.go diff --git a/qr2/message.go b/qr2/message.go new file mode 100644 index 0000000..d313499 --- /dev/null +++ b/qr2/message.go @@ -0,0 +1,206 @@ +package qr2 + +import ( + "bytes" + "encoding/binary" + "github.com/logrusorgru/aurora/v3" + "strconv" + "wwfc/common" + "wwfc/logging" +) + +func SendClientMessage(senderIP string, destSearchID uint64, message []byte) { + moduleName := "QR2/MSG" + + var matchData common.MatchCommandData + var sender *Session + var receiver *Session + senderIPInt, _ := common.IPFormatToInt(senderIP) + + useSearchID := destSearchID < (0x400 << 32) + if useSearchID { + receiver = sessionBySearchID[destSearchID] + } else { + // It's an IP address, used in some circumstances + receiver = sessionByPublicIP[destSearchID] + } + + if receiver == nil { + logging.Error(moduleName, "Destination", aurora.Cyan(destSearchID), "does not exist") + return + } + + // Decode and validate the message + isNatnegPacket := false + if bytes.Equal(message[:2], []byte{0xfd, 0xfc}) { + // Sending natneg cookie + isNatnegPacket = true + if len(message) != 0xA { + logging.Error(moduleName, "Received invalid length NATNEG packet") + return + } + + natnegID := binary.LittleEndian.Uint32(message[0x6:0xA]) + moduleName = "QR2/MSG:s" + strconv.FormatUint(uint64(natnegID), 10) + } else if bytes.Equal(message[:4], []byte{0xbb, 0x49, 0xcc, 0x4d}) { + // DWC match command + if len(message) < 0x14 { + logging.Error(moduleName, "Received invalid length match command packet") + return + } + + senderProfileID := binary.LittleEndian.Uint32(message[0x10:0x14]) + moduleName = "QR2/MSG:p" + strconv.FormatUint(uint64(senderProfileID), 10) + + if (int(message[9]) + 0x14) != len(message) { + logging.Error(moduleName, "Received invalid match command packet header") + return + } + + qr2IP := binary.BigEndian.Uint32(message[0x0C:0x10]) + qr2Port := binary.LittleEndian.Uint16(message[0x0A:0x0C]) + + if senderIPInt != int32(qr2IP) { + logging.Error(moduleName, "Wrong QR2 IP in match command packet header") + } + + sender = sessionByPublicIP[(uint64(qr2Port)<<32)|uint64(qr2IP)] + if sender == nil { + logging.Error(moduleName, "Session does not exist with QR2 IP and port") + return + } + + var ok bool + matchData, ok = common.DecodeMatchCommand(message[8], message[0x14:]) + if !ok { + logging.Error(moduleName, "Received invalid match command:", aurora.Cyan(message[8])) + return + } + + if message[8] == common.MatchReservation { + if matchData.Reservation.MatchType == 3 { + // TODO: Check that this is correct + logging.Error(moduleName, "RESERVATION: Attempt to join a private room over ServerBrowser") + return + } + + if qr2IP != matchData.Reservation.PublicIP { + logging.Error(moduleName, "RESERVATION: Public IP mismatch in header and command") + return + } + + if matchData.Reservation.PublicPort < 1024 { + logging.Error(moduleName, "RESERVATION: Public port is reserved") + return + } + + if matchData.Reservation.LocalPort < 1024 { + logging.Error(moduleName, "RESERVATION: Local port is reserved") + return + } + + if useSearchID { + matchData.Reservation.PublicIP = uint32(sender.SearchID & 0xffffffff) + matchData.Reservation.PublicPort = uint16((sender.SearchID >> 32) & 0xffff) + matchData.Reservation.LocalIP = 0 + matchData.Reservation.LocalPort = 0 + } + } + + if message[8] == common.MatchResvOK { + if qr2IP != matchData.ResvOK.PublicIP { + logging.Error(moduleName, "RESERVATION: Public IP mismatch in header and command") + return + } + + if matchData.ResvOK.PublicPort < 1024 { + logging.Error(moduleName, "RESV_OK: Public port is reserved") + return + } + + if matchData.ResvOK.LocalPort < 1024 { + logging.Error(moduleName, "RESV_OK: Local port is reserved") + return + } + + if matchData.ResvOK.ProfileID != senderProfileID { + logging.Error(moduleName, "RESV_OK: Profile ID mismatch in header") + return + } + + if useSearchID { + matchData.ResvOK.PublicIP = uint32(sender.SearchID & 0xffffffff) + matchData.ResvOK.PublicPort = uint16((sender.SearchID >> 32) & 0xffff) + matchData.ResvOK.LocalIP = 0 + matchData.ResvOK.LocalPort = 0 + } + } + + if message[8] == common.MatchTellAddr { + mutex.Lock() + if sender.Data["publicip"] != receiver.Data["publicip"] { + mutex.Unlock() + logging.Error(moduleName, "TELL_ADDR: Public IP does not match receiver") + return + } + mutex.Unlock() + + if matchData.TellAddr.LocalPort < 1024 { + logging.Error(moduleName, "TELL_ADDR: Local port is reserved") + return + } + } + + if useSearchID { + // Convert public IP to search ID + qr2SearchID := binary.LittleEndian.AppendUint16([]byte{}, uint16((sender.SearchID>>32)&0xffff)) + qr2SearchID = binary.BigEndian.AppendUint32(qr2SearchID, uint32(sender.SearchID&0xffffffff)) + message = append(message[:0x0A], append(qr2SearchID, message[0x10:0x14]...)...) + } else { + message = message[:0x14] + } + + var matchMessage []byte + matchMessage, ok = common.EncodeMatchCommand(message[8], matchData) + if !ok { + logging.Error(moduleName, "Failed to reencode match command:", aurora.Cyan(message[8])) + return + } + + if len(matchMessage) != 0 { + message = append(message, matchMessage...) + } + } + + 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 + + mutex.Unlock() + + if isNatnegPacket { + cookie := binary.BigEndian.Uint32(message[0x2:0x6]) + logging.Notice(moduleName, "Send NN cookie", aurora.Cyan(cookie), "to", aurora.BrightCyan(destPid)) + } else { + common.LogMatchCommand(moduleName, destPid, message[8], matchData) + } + + payload := createResponseHeader(ClientMessageRequest, destSessionID) + + payload = append(payload, []byte{0, 0, 0, 0}...) + 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()) + } +} diff --git a/qr2/session.go b/qr2/session.go index fbedc88..05a4662 100644 --- a/qr2/session.go +++ b/qr2/session.go @@ -1,8 +1,6 @@ package qr2 import ( - "bytes" - "encoding/binary" "github.com/logrusorgru/aurora/v3" "math/rand" "net" @@ -34,12 +32,29 @@ type Session struct { var ( // I would use a sync.Map instead of the map mutex combo, but this performs better. - sessions = map[uint32]*Session{} - mutex = sync.RWMutex{} + sessions = map[uint32]*Session{} + sessionByPublicIP = map[uint64]*Session{} + sessionBySearchID = map[uint64]*Session{} + mutex = sync.RWMutex{} ) // Remove a session. func removeSession(sessionId uint32) { + mutex.Lock() + defer mutex.Unlock() + + if sessions[sessionId] == nil { + return + } + + // Delete search ID lookup + delete(sessionBySearchID, sessions[sessionId].SearchID) + + // Delete public IP lookup + ip, port := common.IPFormatToInt(sessions[sessionId].Addr.String()) + lookupIP := (uint64(port) << 32) | uint64(uint32(ip)) + delete(sessionByPublicIP, lookupIP) + delete(sessions, sessionId) } @@ -131,7 +146,6 @@ func setSessionData(sessionId uint32, payload map[string]string, addr net.Addr) logging.Notice(moduleName, "Creating session", aurora.Cyan(sessionId).String()) data := Session{ SessionID: sessionId, - SearchID: uint64(rand.Int63n(0x400 << 32)), Addr: addr, Challenge: "", Authenticated: false, @@ -140,7 +154,22 @@ func setSessionData(sessionId uint32, payload map[string]string, addr net.Addr) Data: payload, PacketCount: 0, } - data.Data["+searchid"] = strconv.FormatUint(data.SearchID, 10) + + // Set search ID + for { + searchID := uint64(rand.Int63n(0x400 << 32)) + if _, exists := sessionBySearchID[searchID]; !exists { + data.SearchID = searchID + data.Data["+searchid"] = strconv.FormatUint(searchID, 10) + sessionBySearchID[searchID] = &data + break + } + } + + // Set public IP lookup + ip, port := common.IPFormatToInt(addr.String()) + lookupIP := (uint64(port) << 32) | uint64(uint32(ip)) + sessionByPublicIP[lookupIP] = &data sessions[sessionId] = &data return data, true @@ -185,253 +214,3 @@ func GetSessionServers() []map[string]string { return servers } - -func getSessionByPublicIP(publicIP uint32, publicPort uint16) *Session { - ipStr := strconv.FormatInt(int64(int32(publicIP)), 10) - portStr := strconv.FormatUint(uint64(publicPort), 10) - - currentTime := time.Now().Unix() - mutex.Lock() - defer mutex.Unlock() - - // Find the session with the IP - for _, session := range sessions { - if !session.Authenticated { - continue - } - - // If the last keep alive was over a minute ago then consider the server unreachable - if session.LastKeepAlive < currentTime-60 { - continue - } - - if session.Data["publicip"] == ipStr && session.Data["publicport"] == portStr { - return session - } - } - - return nil -} - -func getSessionBySearchID(searchID uint64) *Session { - currentTime := time.Now().Unix() - mutex.Lock() - defer mutex.Unlock() - - // Find the session with the ID - for _, session := range sessions { - if !session.Authenticated { - continue - } - - // If the last keep alive was over a minute ago then consider the server unreachable - if session.LastKeepAlive < currentTime-60 { - continue - } - - if session.SearchID == searchID { - return session - } - } - - return nil -} - -func SendClientMessage(senderIP string, destSearchID uint64, message []byte) { - moduleName := "QR2/MSG" - - var matchData common.MatchCommandData - var sender *Session - var receiver *Session - senderIPInt, _ := common.IPFormatToInt(senderIP) - - useSearchID := destSearchID < (0x400 << 32) - if useSearchID { - receiver = getSessionBySearchID(destSearchID) - } else { - // It's an IP address, used in some circumstances - receiverIPInt := uint32(destSearchID & 0xffffffff) - receiverPort := uint16(destSearchID >> 32) - receiver = getSessionByPublicIP(receiverIPInt, receiverPort) - } - - if receiver == nil { - logging.Error(moduleName, "Destination", aurora.Cyan(destSearchID), "does not exist") - return - } - - // Decode and validate the message - isNatnegPacket := false - if bytes.Equal(message[:2], []byte{0xfd, 0xfc}) { - // Sending natneg cookie - isNatnegPacket = true - if len(message) != 0xA { - logging.Error(moduleName, "Received invalid length NATNEG packet") - return - } - - senderSessionID := binary.LittleEndian.Uint32(message[0x6:0xA]) - moduleName = "QR2/MSG:s" + strconv.FormatUint(uint64(senderSessionID), 10) - } else if bytes.Equal(message[:4], []byte{0xbb, 0x49, 0xcc, 0x4d}) { - // DWC match command - if len(message) < 0x14 { - logging.Error(moduleName, "Received invalid length match command packet") - return - } - - senderProfileID := binary.LittleEndian.Uint32(message[0x10:0x14]) - moduleName = "QR2/MSG:p" + strconv.FormatUint(uint64(senderProfileID), 10) - - if (int(message[9]) + 0x14) != len(message) { - logging.Error(moduleName, "Received invalid match command packet header") - return - } - - qr2IP := binary.BigEndian.Uint32(message[0x0C:0x10]) - qr2Port := binary.LittleEndian.Uint16(message[0x0A:0x0C]) - - if senderIPInt != int32(qr2IP) { - logging.Error(moduleName, "Wrong QR2 IP in match command packet header") - } - - sender = getSessionByPublicIP(qr2IP, qr2Port) - if sender == nil { - logging.Error(moduleName, "Session does not exist with QR2 IP and port") - return - } - - var ok bool - matchData, ok = common.DecodeMatchCommand(message[8], message[0x14:]) - if !ok { - logging.Error(moduleName, "Received invalid match command:", aurora.Cyan(message[8])) - return - } - - if message[8] == common.MatchReservation { - if matchData.Reservation.MatchType == 3 { - // TODO: Check that this is correct - logging.Error(moduleName, "RESERVATION: Attempt to join a private room over ServerBrowser") - return - } - - if qr2IP != matchData.Reservation.PublicIP { - logging.Error(moduleName, "RESERVATION: Public IP mismatch in header and command") - return - } - - if matchData.Reservation.PublicPort < 1024 { - logging.Error(moduleName, "RESERVATION: Public port is reserved") - return - } - - if matchData.Reservation.LocalPort < 1024 { - logging.Error(moduleName, "RESERVATION: Local port is reserved") - return - } - - if useSearchID { - matchData.Reservation.PublicIP = uint32(sender.SearchID & 0xffffffff) - matchData.Reservation.PublicPort = uint16((sender.SearchID >> 32) & 0xffff) - matchData.Reservation.LocalIP = 0 - matchData.Reservation.LocalPort = 0 - } - } - - if message[8] == common.MatchResvOK { - if qr2IP != matchData.ResvOK.PublicIP { - logging.Error(moduleName, "RESERVATION: Public IP mismatch in header and command") - return - } - - if matchData.ResvOK.PublicPort < 1024 { - logging.Error(moduleName, "RESV_OK: Public port is reserved") - return - } - - if matchData.ResvOK.LocalPort < 1024 { - logging.Error(moduleName, "RESV_OK: Local port is reserved") - return - } - - if matchData.ResvOK.ProfileID != senderProfileID { - logging.Error(moduleName, "RESV_OK: Profile ID mismatch in header") - return - } - - if useSearchID { - matchData.ResvOK.PublicIP = uint32(sender.SearchID & 0xffffffff) - matchData.ResvOK.PublicPort = uint16((sender.SearchID >> 32) & 0xffff) - matchData.ResvOK.LocalIP = 0 - matchData.ResvOK.LocalPort = 0 - } - } - - if message[8] == common.MatchTellAddr { - mutex.Lock() - if sender.Data["publicip"] != receiver.Data["publicip"] { - mutex.Unlock() - logging.Error(moduleName, "TELL_ADDR: Public IP does not match receiver") - return - } - mutex.Unlock() - - if matchData.TellAddr.LocalPort < 1024 { - logging.Error(moduleName, "TELL_ADDR: Local port is reserved") - return - } - } - - if useSearchID { - // Convert public IP to search ID - qr2SearchID := binary.LittleEndian.AppendUint16([]byte{}, uint16((sender.SearchID>>32)&0xffff)) - qr2SearchID = binary.BigEndian.AppendUint32(qr2SearchID, uint32(sender.SearchID&0xffffffff)) - message = append(message[:0x0A], append(qr2SearchID, message[0x10:0x14]...)...) - } else { - message = message[:0x14] - } - - var matchMessage []byte - matchMessage, ok = common.EncodeMatchCommand(message[8], matchData) - if !ok { - logging.Error(moduleName, "Failed to reencode match command:", aurora.Cyan(message[8])) - return - } - - if len(matchMessage) != 0 { - message = append(message, matchMessage...) - } - } - - 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 - - mutex.Unlock() - - if isNatnegPacket { - cookie := binary.BigEndian.Uint32(message[0x2:0x6]) - logging.Notice(moduleName, "Send NN cookie", aurora.Cyan(cookie), "to", aurora.BrightCyan(destPid)) - } else { - common.LogMatchCommand(moduleName, destPid, message[8], matchData) - } - - payload := createResponseHeader(ClientMessageRequest, destSessionID) - - payload = append(payload, []byte{0, 0, 0, 0}...) - binary.BigEndian.PutUint32(payload[len(payload)-4:], packetCount) - payload = append(payload, message...) - - // TODO: Send again if no CLIENT_MESSAGE_ACK is received after - _, err := masterConn.WriteTo(payload, destAddr) - if err != nil { - logging.Error(moduleName, "Error sending message:", err.Error()) - } -}