Minor fixes

stack_system_nat.go

@@ -54,18 +54,36 @@ func (n *TCPNat) loopCheckTimeout(ctx context.Context) {
 
 func (n *TCPNat) checkTimeout() {
 	now := time.Now()
-	n.portAccess.Lock()
-	defer n.portAccess.Unlock()
-	n.addrAccess.Lock()
-	defer n.addrAccess.Unlock()
+	type expiredSession struct {
+		port    uint16
+		session *TCPSession
+	}
+	var expired []expiredSession
+	n.portAccess.RLock()
 	for natPort, session := range n.portMap {
 		session.Lock()
-		if now.Sub(session.LastActive) > n.timeout {
-			delete(n.addrMap, tcpNatKey{Source: session.Source, Destination: session.Destination})
-			delete(n.portMap, natPort)
-		}
+		timedOut := now.Sub(session.LastActive) > n.timeout
 		session.Unlock()
+		if timedOut {
+			expired = append(expired, expiredSession{port: natPort, session: session})
+		}
 	}
+	n.portAccess.RUnlock()
+	if len(expired) == 0 {
+		return
+	}
+	n.addrAccess.Lock()
+	n.portAccess.Lock()
+	for _, e := range expired {
+		e.session.Lock()
+		if now.Sub(e.session.LastActive) > n.timeout {
+			delete(n.addrMap, tcpNatKey{Source: e.session.Source, Destination: e.session.Destination})
+			delete(n.portMap, e.port)
+		}
+		e.session.Unlock()
+	}
+	n.portAccess.Unlock()
+	n.addrAccess.Unlock()
 }
 
 func (n *TCPNat) LookupBack(port uint16) *TCPSession {
@@ -91,21 +109,37 @@ func (n *TCPNat) Lookup(source netip.AddrPort, destination netip.AddrPort) uint1
 		return port
 	}
 	n.addrAccess.Lock()
-	nextPort := n.portIndex
-	if nextPort == 0 {
-		nextPort = 10000
-		n.portIndex = 10001
-	} else {
-		n.portIndex++
+	defer n.addrAccess.Unlock()
+	if port, loaded = n.addrMap[key]; loaded {
+		return port
 	}
-	n.addrMap[key] = nextPort
-	n.addrAccess.Unlock()
 	n.portAccess.Lock()
+	defer n.portAccess.Unlock()
+	nextPort, ok := n.allocatePortLocked()
+	if !ok {
+		return 0
+	}
 	n.portMap[nextPort] = &TCPSession{
 		Source:      source,
 		Destination: destination,
 		LastActive:  time.Now(),
 	}
-	n.portAccess.Unlock()
+	n.addrMap[key] = nextPort
 	return nextPort
 }
+
+func (n *TCPNat) allocatePortLocked() (uint16, bool) {
+	for range 65535 - 10000 + 1 {
+		nextPort := n.portIndex
+		if nextPort == 0 {
+			nextPort = 10000
+			n.portIndex = 10001
+		} else {
+			n.portIndex++
+		}
+		if _, occupied := n.portMap[nextPort]; !occupied {
+			return nextPort, true
+		}
+	}
+	return 0, false
+}