From abc744c647e7267f9bb161ffb1bf1a410a8ea3c5 Mon Sep 17 00:00:00 2001
From: Laurence <laurence.jones@live.co.uk>
Date: Thu, 12 Mar 2026 12:55:49 +0000
Subject: [PATCH] perf(relay): index WireGuard sessions by receiver index

---
 relay/relay.go | 37 ++++++++++++++++++++-----------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/relay/relay.go b/relay/relay.go
index 22aff76..402c696 100644
--- a/relay/relay.go
+++ b/relay/relay.go
@@ -150,6 +150,9 @@ type UDPProxyServer struct {
 	// Session tracking for WireGuard peers
 	// Key format: "senderIndex:receiverIndex"
 	wgSessions sync.Map
+	// Session index for O(1) lookup by receiver index
+	// Key: receiverIndex (uint32), Value: *WireGuardSession
+	sessionsByReceiverIndex sync.Map
 	// Communication pattern tracking for rebuilding sessions
 	// Key format: "clientIP:clientPort-destIP:destPort"
 	commPatterns sync.Map
@@ -477,12 +480,15 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		sessionKey := fmt.Sprintf("%d:%d", receiverIndex, senderIndex)
 
 		// Store the session information
-		s.wgSessions.Store(sessionKey, &WireGuardSession{
+		session := &WireGuardSession{
 			ReceiverIndex: receiverIndex,
 			SenderIndex:   senderIndex,
 			DestAddr:      remoteAddr,
 			LastSeen:      time.Now(),
-		})
+		}
+		s.wgSessions.Store(sessionKey, session)
+		// Also index by sender index for O(1) lookup in transport data path
+		s.sessionsByReceiverIndex.Store(senderIndex, session)
 
 		// Forward the response to the original sender
 		for _, dest := range proxyMapping.Destinations {
@@ -508,21 +514,15 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		// Data packet: forward only to the established session peer
 		// logger.Debug("Received transport data with receiver index %d from %s", receiverIndex, remoteAddr)
 
-		// Look up the session based on the receiver index
+		// Look up the session based on the receiver index - O(1) lookup instead of O(n) Range
 		var destAddr *net.UDPAddr
 
-		// First check for existing sessions to see if we know where to send this packet
-		s.wgSessions.Range(func(k, v interface{}) bool {
-			session := v.(*WireGuardSession)
-			// Check if session matches (read lock for check)
-			if session.GetSenderIndex() == receiverIndex {
-				// Found matching session - get dest addr and update last seen
-				destAddr = session.GetDestAddr()
-				session.UpdateLastSeen()
-				return false // stop iteration
-			}
-			return true // continue iteration
-		})
+		// Fast path: direct index lookup by receiver index
+		if sessionObj, ok := s.sessionsByReceiverIndex.Load(receiverIndex); ok {
+			session := sessionObj.(*WireGuardSession)
+			destAddr = session.GetDestAddr()
+			session.UpdateLastSeen()
+		}
 
 		if destAddr != nil {
 			// We found a specific peer to forward to
@@ -634,12 +634,15 @@ func (s *UDPProxyServer) handleResponses(conn *net.UDPConn, destAddr *net.UDPAdd
 			if ok && buffer[0] == WireGuardMessageTypeHandshakeResponse {
 				// Store the session mapping for the handshake response
 				sessionKey := fmt.Sprintf("%d:%d", senderIndex, receiverIndex)
-				s.wgSessions.Store(sessionKey, &WireGuardSession{
+				session := &WireGuardSession{
 					ReceiverIndex: receiverIndex,
 					SenderIndex:   senderIndex,
 					DestAddr:      destAddr,
 					LastSeen:      time.Now(),
-				})
+				}
+				s.wgSessions.Store(sessionKey, session)
+				// Also index by sender index for O(1) lookup
+				s.sessionsByReceiverIndex.Store(senderIndex, session)
 				logger.Debug("Stored session mapping: %s -> %s", sessionKey, destAddr.String())
 			} else if ok && buffer[0] == WireGuardMessageTypeTransportData {
 				// Track communication pattern for session rebuilding (reverse direction)