Fix clients targets sync and remove nat entries

netstack2/handlers.go

@@ -166,6 +166,13 @@ func (h *TCPHandler) handleTCPConn(netstackConn *gonet.TCPConn, id stack.Transpo
 
 	defer netstackConn.Close()
 
+	// Release this connection's NAT state once it fully closes, so a rule
+	// change (e.g. RewriteTo) takes effect for the next connection on this
+	// tuple instead of being masked by a stale cached resolution forever.
+	if h.proxyHandler != nil {
+		defer h.proxyHandler.releaseConnectionState(srcIP, srcPort, dstIP, dstPort, uint8(tcp.ProtocolNumber))
+	}
+
 	logger.Info("TCP Forwarder: Handling connection %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
 
 	// Check if there's a destination rewrite for this connection (e.g., localhost targets)
@@ -315,6 +322,13 @@ func (h *UDPHandler) handleUDPConn(netstackConn *gonet.UDPConn, id stack.Transpo
 	dstIP := id.LocalAddress.String()
 	dstPort := id.LocalPort
 
+	// Release this session's NAT state once it fully closes (session end or
+	// idle timeout), so a rule change takes effect for the next session on
+	// this tuple instead of being masked by a stale cached resolution.
+	if h.proxyHandler != nil {
+		defer h.proxyHandler.releaseConnectionState(srcIP, srcPort, dstIP, dstPort, uint8(udp.ProtocolNumber))
+	}
+
 	logger.Info("UDP Forwarder: Handling connection %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
 
 	// Drop connection if blocking is enabled

netstack2/proxy.go

@@ -272,6 +272,18 @@ func (p *ProxyHandler) RemoveSubnetRule(sourcePrefix, destPrefix netip.Prefix) {
 	p.subnetLookup.RemoveSubnet(sourcePrefix, destPrefix)
 }
 
+// ReplaceAllSubnetRules atomically replaces the full set of subnet rules.
+// Intended for full-state syncs where the desired rule set is authoritative,
+// so any stale rule is guaranteed to be cleared even if its key
+// (SourcePrefix, DestPrefix) matches a still-desired rule but its contents
+// (e.g. RewriteTo) have changed.
+func (p *ProxyHandler) ReplaceAllSubnetRules(rules []SubnetRule) {
+	if p == nil || !p.enabled {
+		return
+	}
+	p.subnetLookup.ReplaceAll(rules)
+}
+
 // GetAllRules returns all subnet rules from the proxy handler
 func (p *ProxyHandler) GetAllRules() []SubnetRule {
 	if p == nil || !p.enabled {
@@ -335,6 +347,51 @@ func (p *ProxyHandler) SetHTTPRequestLogSender(fn SendFunc) {
 	p.httpRequestLogger.SetSendFunc(fn)
 }
 
+// releaseConnectionState removes the per-connection NAT state for a single
+// (srcIP, srcPort, dstIP, dstPort, proto) tuple. Callers must only invoke
+// this once that exact connection has fully closed (both directions torn
+// down), since a new connection can never be accepted on the same 5-tuple
+// before then.
+//
+// This intentionally does NOT touch destRewriteTable/resourceTable: those
+// are keyed without srcPort (destKey), so they are shared across every
+// concurrent connection from the same source to the same destination
+// service. They don't need connection-scoped cleanup - each new connection's
+// first packet already refreshes them via HandleIncomingPacket - and
+// deleting them here on a single connection's close could break other
+// connections still in flight to the same destination.
+func (p *ProxyHandler) releaseConnectionState(srcIP string, srcPort uint16, dstIP string, dstPort uint16, proto uint8) {
+	if p == nil || !p.enabled {
+		return
+	}
+
+	key := connKey{
+		srcIP:   srcIP,
+		srcPort: srcPort,
+		dstIP:   dstIP,
+		dstPort: dstPort,
+		proto:   proto,
+	}
+
+	p.natMu.Lock()
+	defer p.natMu.Unlock()
+
+	entry, ok := p.natTable[key]
+	if !ok {
+		return
+	}
+	delete(p.natTable, key)
+
+	reverseKey := reverseConnKey{
+		rewrittenTo:     entry.rewrittenTo.String(),
+		originalSrcIP:   srcIP,
+		originalSrcPort: srcPort,
+		originalDstPort: dstPort,
+		proto:           proto,
+	}
+	delete(p.reverseNatTable, reverseKey)
+}
+
 // LookupDestinationRewrite looks up the rewritten destination for a connection
 // This is used by TCP/UDP handlers to find the actual target address
 func (p *ProxyHandler) LookupDestinationRewrite(srcIP, dstIP string, dstPort uint16, proto uint8) (netip.Addr, bool) {

netstack2/subnet_lookup.go

@@ -52,6 +52,13 @@ func (sl *SubnetLookup) AddSubnet(rule SubnetRule) {
 	sl.mu.Lock()
 	defer sl.mu.Unlock()
 
+	sl.addSubnetLocked(rule)
+}
+
+// addSubnetLocked is the lock-free body of AddSubnet, factored out so
+// ReplaceAll can insert many rules under a single lock acquisition.
+// Callers must hold sl.mu for writing.
+func (sl *SubnetLookup) addSubnetLocked(rule SubnetRule) {
 	rulePtr := &rule
 
 	// Canonicalize source prefix to handle host bits correctly
@@ -89,6 +96,21 @@ func (sl *SubnetLookup) AddSubnet(rule SubnetRule) {
 	destTriePtr.rules = newRules
 }
 
+// ReplaceAll atomically replaces the entire rule set with the given rules.
+// This guarantees no stale rule can survive a sync, even when a rule's key
+// (SourcePrefix, DestPrefix) is unchanged but other fields (e.g. RewriteTo)
+// differ - a case that an add/remove diff keyed only on prefixes would miss.
+func (sl *SubnetLookup) ReplaceAll(rules []SubnetRule) {
+	sl.mu.Lock()
+	defer sl.mu.Unlock()
+
+	sl.sourceTrie = &bart.Table[*destTrie]{}
+
+	for _, rule := range rules {
+		sl.addSubnetLocked(rule)
+	}
+}
+
 // RemoveSubnet removes a subnet rule from the lookup table
 func (sl *SubnetLookup) RemoveSubnet(sourcePrefix, destPrefix netip.Prefix) {
 	sl.mu.Lock()

netstack2/tun.go

@@ -364,6 +364,15 @@ func (net *Net) RemoveProxySubnetRule(sourcePrefix, destPrefix netip.Prefix) {
 	}
 }
 
+// ReplaceProxySubnetRules atomically replaces the full set of subnet rules
+// on the proxy handler with the given rules.
+func (net *Net) ReplaceProxySubnetRules(rules []SubnetRule) {
+	tun := (*netTun)(net)
+	if tun.proxyHandler != nil {
+		tun.proxyHandler.ReplaceAllSubnetRules(rules)
+	}
+}
+
 // GetProxySubnetRules returns all subnet rules from the proxy handler
 func (net *Net) GetProxySubnetRules() []SubnetRule {
 	tun := (*netTun)(net)