diff --git a/stack_system_nat.go b/stack_system_nat.go index 636b561..cc46017 100644 --- a/stack_system_nat.go +++ b/stack_system_nat.go @@ -15,10 +15,15 @@ type TCPNat struct { portIndex uint16 portAccess sync.RWMutex addrAccess sync.RWMutex - addrMap map[netip.AddrPort]uint16 + addrMap map[tcpNatKey]uint16 portMap map[uint16]*TCPSession } +type tcpNatKey struct { + Source netip.AddrPort + Destination netip.AddrPort +} + type TCPSession struct { sync.Mutex Source netip.AddrPort @@ -30,7 +35,7 @@ func NewNat(ctx context.Context, timeout time.Duration) *TCPNat { natMap := &TCPNat{ timeout: timeout, portIndex: 10000, - addrMap: make(map[netip.AddrPort]uint16), + addrMap: make(map[tcpNatKey]uint16), portMap: make(map[uint16]*TCPSession), } go natMap.loopCheckTimeout(ctx) @@ -59,7 +64,7 @@ func (n *TCPNat) checkTimeout() { for natPort, session := range n.portMap { session.Lock() if now.Sub(session.LastActive) > n.timeout { - delete(n.addrMap, session.Source) + delete(n.addrMap, tcpNatKey{Source: session.Source, Destination: session.Destination}) delete(n.portMap, natPort) } session.Unlock() @@ -81,8 +86,9 @@ func (n *TCPNat) LookupBack(port uint16) *TCPSession { } func (n *TCPNat) Lookup(source netip.AddrPort, destination netip.AddrPort, handler Handler) (uint16, error) { + key := tcpNatKey{Source: source, Destination: destination} n.addrAccess.RLock() - port, loaded := n.addrMap[source] + port, loaded := n.addrMap[key] n.addrAccess.RUnlock() if loaded { return port, nil @@ -99,7 +105,7 @@ func (n *TCPNat) Lookup(source netip.AddrPort, destination netip.AddrPort, handl } else { n.portIndex++ } - n.addrMap[source] = nextPort + n.addrMap[key] = nextPort n.addrAccess.Unlock() n.portAccess.Lock() n.portMap[nextPort] = &TCPSession{