QR2: Always lookup session by IP address

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()
 	}
 

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
 		}

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
 	}
 }

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:

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

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