pangolin_mirror/gerbil
1
0
Fork 0
mirror of https://github.com/fosrl/gerbil.git synced 2026-05-12 03:09:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(relay): index WireGuard sessions by receiver index

Browse Source
This commit is contained in:
Laurence
2026-03-12 12:55:49 +00:00
parent b9261b8fea
commit abc744c647
1 changed files with 20 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
relay/relay.go
View File

@@ -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)