diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2466b7d..82ce101 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -19,7 +19,7 @@ jobs: - name: Set up Go uses: actions/setup-go@4b73464bb391d4059bd26b0524d20df3927bd417 # v6.3.0 with: - go-version: 1.25 + go-version: 1.26 - name: Build go run: go build diff --git a/relay/relay.go b/relay/relay.go index c065c29..dad47c1 100644 --- a/relay/relay.go +++ b/relay/relay.go @@ -161,6 +161,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 @@ -586,6 +589,8 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD } else { metrics.RecordSession(relayIfname, 1) } + // 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 { @@ -615,21 +620,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 @@ -767,6 +766,8 @@ func (s *UDPProxyServer) handleResponses(conn *net.UDPConn, destAddr *net.UDPAdd } else { metrics.RecordSession(relayIfname, 1) } + // 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)