mirror of
https://github.com/fosrl/newt.git
synced 2026-03-05 18:26:42 +00:00
Merge branch 'optimize-reverse-nat-lookup' of github.com:LaurenceJJones/newt into LaurenceJJones-optimize-reverse-nat-lookup
This commit is contained in:
@@ -57,6 +57,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
|
||||||
@@ -81,7 +92,8 @@ type ProxyHandler struct {
|
|||||||
icmpHandler *ICMPHandler
|
icmpHandler *ICMPHandler
|
||||||
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
|
||||||
icmpReplies chan []byte // Channel for ICMP reply packets to be sent back through the tunnel
|
icmpReplies chan []byte // Channel for ICMP reply packets to be sent back through the tunnel
|
||||||
@@ -106,6 +118,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),
|
||||||
icmpReplies: make(chan []byte, 256), // Buffer for ICMP reply packets
|
icmpReplies: make(chan []byte, 256), // Buffer for ICMP reply packets
|
||||||
proxyEp: channel.New(1024, uint32(options.MTU), ""),
|
proxyEp: channel.New(1024, uint32(options.MTU), ""),
|
||||||
@@ -408,10 +421,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()
|
||||||
@@ -610,20 +636,22 @@ func (p *ProxyHandler) ReadOutgoingPacket() *buffer.View {
|
|||||||
return view
|
return 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 {
|
||||||
|
|||||||
Reference in New Issue
Block a user