From 2aa89f74e64d2adc2e1a61c07001632db5fdd957 Mon Sep 17 00:00:00 2001 From: mkwcat Date: Tue, 12 Dec 2023 06:45:01 -0500 Subject: [PATCH] QR2: Always lookup session by IP address --- qr2/challenge.go | 4 +-- qr2/group.go | 6 ++--- qr2/heartbeat.go | 14 +++++------ qr2/main.go | 21 ++++++++++------ qr2/message.go | 4 +-- qr2/session.go | 64 ++++++++++++++++++++++-------------------------- 6 files changed, 57 insertions(+), 56 deletions(-) diff --git a/qr2/challenge.go b/qr2/challenge.go index 230dbd2..20d83ea 100644 --- a/qr2/challenge.go +++ b/qr2/challenge.go @@ -8,7 +8,7 @@ import ( "wwfc/common" ) -func sendChallenge(conn net.PacketConn, addr net.Addr, session Session) { +func sendChallenge(conn net.PacketConn, addr net.Addr, session Session, lookupAddr uint64) { challenge := session.Challenge if challenge == "" { // Generate challenge @@ -32,7 +32,7 @@ func sendChallenge(conn net.PacketConn, addr net.Addr, session Session) { challenge = common.RandomString(6) + "00" + hexIP + hexPort mutex.Lock() - sessions[session.SessionID].Challenge = challenge + sessions[lookupAddr].Challenge = challenge mutex.Unlock() } diff --git a/qr2/group.go b/qr2/group.go index 48b4876..2baba55 100644 --- a/qr2/group.go +++ b/qr2/group.go @@ -88,13 +88,13 @@ func ProcessGPResvOK(cmd common.MatchCommandDataResvOK, senderIP uint64, senderP mutex.Lock() defer mutex.Unlock() - from := sessionByPublicIP[senderIP] + from := sessions[senderIP] if from == nil { logging.Error(moduleName, "Sender IP does not exist:", aurora.Cyan(fmt.Sprintf("%012x", senderIP))) return false } - to := sessionByPublicIP[destIP] + to := sessions[destIP] if to == nil { logging.Error(moduleName, "Destination IP does not exist:", aurora.Cyan(fmt.Sprintf("%012x", destIP))) return false @@ -117,7 +117,7 @@ func ProcessGPStatusUpdate(senderIP uint64, status string) { mutex.Lock() defer mutex.Unlock() - session := sessionByPublicIP[senderIP] + session := sessions[senderIP] if session == nil || session.GroupPointer == nil { return } diff --git a/qr2/heartbeat.go b/qr2/heartbeat.go index 1fa2d93..afe7bc5 100644 --- a/qr2/heartbeat.go +++ b/qr2/heartbeat.go @@ -4,17 +4,15 @@ import ( "encoding/binary" "github.com/logrusorgru/aurora/v3" "net" - "strconv" "strings" "wwfc/common" "wwfc/logging" ) -func heartbeat(conn net.PacketConn, addr net.Addr, buffer []byte) { +func heartbeat(moduleName string, conn net.PacketConn, addr net.Addr, buffer []byte) { sessionId := binary.BigEndian.Uint32(buffer[1:5]) - moduleName := "QR2:" + strconv.FormatInt(int64(sessionId), 10) + logging.Info(moduleName, "Received heartbeat; session ID:", aurora.BrightCyan(sessionId)) - logging.Info(moduleName, "Received heartbeat from", aurora.BrightCyan(addr)) values := strings.Split(string(buffer[5:]), "\u0000") payload := map[string]string{} @@ -41,6 +39,8 @@ func heartbeat(conn net.PacketConn, addr net.Addr, buffer []byte) { return } + lookupAddr := makeLookupAddr(addr.String()) + if statechanged, ok := payload["statechanged"]; ok { if statechanged == "1" { // TODO: This would be a good place to run the server->client message exploit @@ -52,19 +52,19 @@ func heartbeat(conn net.PacketConn, addr net.Addr, buffer []byte) { if statechanged == "2" { logging.Notice(moduleName, "Client session shutdown") - removeSession(sessionId) + removeSession(lookupAddr) return } } - session, ok := setSessionData(sessionId, payload, addr) + session, ok := setSessionData(moduleName, addr, sessionId, payload) if !ok { return } if !session.Authenticated { logging.Notice(moduleName, "Sending challenge") - sendChallenge(conn, addr, session) + sendChallenge(conn, addr, session, lookupAddr) return } } diff --git a/qr2/main.go b/qr2/main.go index 31867cb..3358cf6 100644 --- a/qr2/main.go +++ b/qr2/main.go @@ -4,7 +4,6 @@ import ( "encoding/binary" "github.com/logrusorgru/aurora/v3" "net" - "strconv" "time" "wwfc/common" "wwfc/logging" @@ -55,15 +54,23 @@ func StartServer() { func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) { packetType := buffer[0] - sessionId := binary.BigEndian.Uint32(buffer[1:5]) - session, ok := sessions[sessionId] - moduleName := "QR2:" + strconv.FormatInt(int64(sessionId), 10) + moduleName := "QR2:" + addr.String() + var session *Session if packetType != HeartbeatRequest && packetType != AvailableRequest { + mutex.Lock() + + var ok bool + session, ok = sessions[makeLookupAddr(addr.String())] if !ok { - logging.Error(moduleName, "Invalid session") + mutex.Unlock() + logging.Error(moduleName, "Cannot find session for this IP address") return } + + session.SessionID = binary.BigEndian.Uint32(buffer[1:5]) + + mutex.Unlock() } switch packetType { @@ -80,7 +87,7 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) { session.Authenticated = true mutex.Unlock() - conn.WriteTo(createResponseHeader(ClientRegisteredReply, sessionId), addr) + conn.WriteTo(createResponseHeader(ClientRegisteredReply, session.SessionID), addr) } else { mutex.Unlock() } @@ -92,7 +99,7 @@ func handleConnection(conn net.PacketConn, addr net.Addr, buffer []byte) { case HeartbeatRequest: logging.Notice(moduleName, "Command:", aurora.Yellow("HEARTBEAT")) - heartbeat(conn, addr, buffer) + heartbeat(moduleName, conn, addr, buffer) break case AddErrorRequest: diff --git a/qr2/message.go b/qr2/message.go index 21a51f3..81e4357 100644 --- a/qr2/message.go +++ b/qr2/message.go @@ -35,7 +35,7 @@ func SendClientMessage(senderIP string, destSearchID uint64, message []byte) { receiver = sessionBySearchID[destSearchID] } else { // It's an IP address, used in some circumstances - receiver = sessionByPublicIP[destSearchID] + receiver = sessions[destSearchID] } if receiver == nil || !receiver.Authenticated { @@ -89,7 +89,7 @@ func SendClientMessage(senderIP string, destSearchID uint64, message []byte) { return } - sender = sessionByPublicIP[(uint64(qr2Port)<<32)|uint64(qr2IP)] + sender = sessions[(uint64(qr2Port)<<32)|uint64(qr2IP)] if sender == nil || !sender.Authenticated { logging.Error(moduleName, "Session does not exist with QR2 IP and port") return diff --git a/qr2/session.go b/qr2/session.go index 28ea3a9..9eb464e 100644 --- a/qr2/session.go +++ b/qr2/session.go @@ -33,19 +33,17 @@ type Session struct { } var ( - // I would use a sync.Map instead of the map mutex combo, but this performs better. - sessions = map[uint32]*Session{} - sessionByPublicIP = map[uint64]*Session{} + sessions = map[uint64]*Session{} sessionBySearchID = map[uint64]*Session{} mutex = sync.RWMutex{} ) // Remove a session. -func removeSession(sessionId uint32) { +func removeSession(addr uint64) { mutex.Lock() defer mutex.Unlock() - session := sessions[sessionId] + session := sessions[addr] if session == nil { return } @@ -64,27 +62,22 @@ func removeSession(sessionId uint32) { } // Delete search ID lookup - delete(sessionBySearchID, sessions[sessionId].SearchID) + delete(sessionBySearchID, sessions[addr].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) + delete(sessions, addr) } // Update session data, creating the session if it doesn't exist. Returns a copy of the session data. -func setSessionData(sessionId uint32, payload map[string]string, addr net.Addr) (Session, bool) { - moduleName := "QR2:" + strconv.FormatInt(int64(sessionId), 10) - +func setSessionData(moduleName string, addr net.Addr, sessionId uint32, payload map[string]string) (Session, bool) { newPID, newPIDValid := payload["dwc_pid"] delete(payload, "dwc_pid") + lookupAddr := makeLookupAddr(addr.String()) + // Moving into performing operations on the session data, so lock the mutex mutex.Lock() defer mutex.Unlock() - session, sessionExists := sessions[sessionId] + session, sessionExists := sessions[lookupAddr] if sessionExists && session.Addr.String() != addr.String() { logging.Error(moduleName, "Session IP mismatch") @@ -123,12 +116,7 @@ func setSessionData(sessionId uint32, payload map[string]string, addr net.Addr) } } - // Set public IP lookup - ip, port := common.IPFormatToInt(addr.String()) - lookupIP := (uint64(port) << 32) | uint64(uint32(ip)) - sessionByPublicIP[lookupIP] = session - - sessions[sessionId] = session + sessions[lookupAddr] = session return *session, true } @@ -141,6 +129,7 @@ func setSessionData(sessionId uint32, payload map[string]string, addr net.Addr) session.Data = payload session.LastKeepAlive = time.Now().Unix() + session.SessionID = sessionId return *session, true } @@ -179,9 +168,9 @@ func (session *Session) setProfileID(moduleName string, newPID string) bool { } // Constraint: only one session can exist with a profile ID - var outdated []uint32 - for sessionID, otherSession := range sessions { - if sessionID == session.SessionID { + var outdated []uint64 + for sessionAddr, otherSession := range sessions { + if otherSession == session { continue } @@ -190,12 +179,12 @@ func (session *Session) setProfileID(moduleName string, newPID string) bool { } // Remove old sessions with the PID - outdated = append(outdated, sessionID) + outdated = append(outdated, sessionAddr) } - for _, sessionID := range outdated { - logging.Notice(moduleName, "Removing outdated session", aurora.BrightCyan(sessionID), "with PID", aurora.Cyan(newPID)) - removeSession(sessionID) + for _, sessionAddr := range outdated { + logging.Notice(moduleName, "Removing outdated session", aurora.BrightCyan(sessions[sessionAddr].Addr.String()), "with PID", aurora.Cyan(newPID)) + removeSession(sessionAddr) } session.Data["dwc_pid"] = newPID @@ -204,20 +193,25 @@ func (session *Session) setProfileID(moduleName string, newPID string) bool { return true } +func makeLookupAddr(addr string) uint64 { + ip, port := common.IPFormatToInt(addr) + return (uint64(port) << 32) | uint64(uint32(ip)) +} + // Get a copy of the list of servers func GetSessionServers() []map[string]string { var servers []map[string]string - var unreachable []uint32 + var unreachable []uint64 currentTime := time.Now().Unix() mutex.Lock() defer mutex.Unlock() - for _, session := range sessions { + for sessionAddr, session := range sessions { // If the last keep alive was over a minute ago then consider the server unreachable if session.LastKeepAlive < currentTime-60 { // If the last keep alive was over an hour ago then remove the server if session.LastKeepAlive < currentTime-((60*60)*1) { - unreachable = append(unreachable, session.SessionID) + unreachable = append(unreachable, sessionAddr) } continue } @@ -230,9 +224,9 @@ func GetSessionServers() []map[string]string { } // Remove unreachable sessions - for _, sessionID := range unreachable { - logging.Notice("QR2", "Removing unreachable session", aurora.BrightCyan(sessionID)) - removeSession(sessionID) + for _, sessionAddr := range unreachable { + logging.Notice("QR2", "Removing unreachable session", aurora.BrightCyan(sessions[sessionAddr].Addr.String())) + removeSession(sessionAddr) } return servers