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