perf: optimize reverse NAT lookup with O(1) map instead of O(n) iteration

Replace O(n) linear search through NAT table with O(1) reverse lookup map
for reply packet NAT translation.

Changes:
- Add reverseConnKey type for reverse NAT lookups
- Add reverseNatTable map to ProxyHandler for O(1) lookups
- Populate both forward and reverse maps when creating NAT entries
- Replace iteration-based reverse lookup with direct map access

Performance:
- O(n) → O(1) complexity for reverse NAT lookups
- Eliminates lock-held iteration on every reply packet
- Removes string comparisons from hot path
- Expected 10-50x improvement for reverse NAT lookups

This addresses Critical #1 from performance analysis where reply path
was walking the entire NAT table to find original mapping.
This commit is contained in:
Laurence
2025-12-16 08:16:37 +00:00
parent d5e0771094
commit 1cf75b00ff

View File

@@ -136,7 +136,7 @@ func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16) *SubnetRule
return nil return nil
} }
// connKey uniquely identifies a connection for NAT tracking // connKey uniquely identifies a connection for NAT tracking (forward direction)
type connKey struct { type connKey struct {
srcIP string srcIP string
srcPort uint16 srcPort uint16
@@ -145,6 +145,17 @@ type connKey struct {
proto uint8 proto uint8
} }
// reverseConnKey uniquely identifies a connection for reverse NAT lookup (reply direction)
// Key structure: (rewrittenTo, originalSrcIP, originalSrcPort, originalDstPort, proto)
// This allows O(1) lookup of NAT entries for reply packets
type reverseConnKey struct {
rewrittenTo string // The address we rewrote to (becomes src in replies)
originalSrcIP string // Original source IP (becomes dst in replies)
originalSrcPort uint16 // Original source port (becomes dst port in replies)
originalDstPort uint16 // Original destination port (becomes src port in replies)
proto uint8
}
// destKey identifies a destination for handler lookups (without source port since it may change) // destKey identifies a destination for handler lookups (without source port since it may change)
type destKey struct { type destKey struct {
srcIP string srcIP string
@@ -168,7 +179,8 @@ type ProxyHandler struct {
udpHandler *UDPHandler udpHandler *UDPHandler
subnetLookup *SubnetLookup subnetLookup *SubnetLookup
natTable map[connKey]*natState natTable map[connKey]*natState
destRewriteTable map[destKey]netip.Addr // Maps original dest to rewritten dest for handler lookups reverseNatTable map[reverseConnKey]*natState // Reverse lookup map for O(1) reply packet NAT
destRewriteTable map[destKey]netip.Addr // Maps original dest to rewritten dest for handler lookups
natMu sync.RWMutex natMu sync.RWMutex
enabled bool enabled bool
} }
@@ -190,6 +202,7 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
enabled: true, enabled: true,
subnetLookup: NewSubnetLookup(), subnetLookup: NewSubnetLookup(),
natTable: make(map[connKey]*natState), natTable: make(map[connKey]*natState),
reverseNatTable: make(map[reverseConnKey]*natState),
destRewriteTable: make(map[destKey]netip.Addr), destRewriteTable: make(map[destKey]netip.Addr),
proxyEp: channel.New(1024, uint32(options.MTU), ""), proxyEp: channel.New(1024, uint32(options.MTU), ""),
proxyStack: stack.New(stack.Options{ proxyStack: stack.New(stack.Options{
@@ -472,10 +485,23 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
// Store NAT state for this connection // Store NAT state for this connection
p.natMu.Lock() p.natMu.Lock()
p.natTable[key] = &natState{ natEntry := &natState{
originalDst: dstAddr, originalDst: dstAddr,
rewrittenTo: newDst, rewrittenTo: newDst,
} }
p.natTable[key] = natEntry
// Create reverse lookup key for O(1) reply packet lookups
// Key: (rewrittenTo, originalSrcIP, originalSrcPort, originalDstPort, proto)
reverseKey := reverseConnKey{
rewrittenTo: newDst.String(),
originalSrcIP: srcAddr.String(),
originalSrcPort: srcPort,
originalDstPort: dstPort,
proto: uint8(protocol),
}
p.reverseNatTable[reverseKey] = natEntry
// Store destination rewrite for handler lookups // Store destination rewrite for handler lookups
p.destRewriteTable[dKey] = newDst p.destRewriteTable[dKey] = newDst
p.natMu.Unlock() p.natMu.Unlock()
@@ -657,20 +683,22 @@ func (p *ProxyHandler) ReadOutgoingPacket() *buffer.View {
} }
} }
// Look up NAT state for reverse translation // Look up NAT state for reverse translation using O(1) reverse lookup map
// The key uses the original dst (before rewrite), so for replies we need to // Key: (rewrittenTo, originalSrcIP, originalSrcPort, originalDstPort, proto)
// find the entry where the rewritten address matches the current source // For reply packets:
// - reply's srcIP = rewrittenTo (the address we rewrote to)
// - reply's dstIP = originalSrcIP (original source IP)
// - reply's srcPort = originalDstPort (original destination port)
// - reply's dstPort = originalSrcPort (original source port)
p.natMu.RLock() p.natMu.RLock()
var natEntry *natState reverseKey := reverseConnKey{
for k, entry := range p.natTable { rewrittenTo: srcIP.String(), // Reply's source is the rewritten address
// Match: reply's dst should be original src, reply's src should be rewritten dst originalSrcIP: dstIP.String(), // Reply's destination is the original source
if k.srcIP == dstIP.String() && k.srcPort == dstPort && originalSrcPort: dstPort, // Reply's destination port is the original source port
entry.rewrittenTo.String() == srcIP.String() && k.dstPort == srcPort && originalDstPort: srcPort, // Reply's source port is the original destination port
k.proto == uint8(protocol) { proto: uint8(protocol),
natEntry = entry
break
}
} }
natEntry := p.reverseNatTable[reverseKey]
p.natMu.RUnlock() p.natMu.RUnlock()
if natEntry != nil { if natEntry != nil {