Merge remote-tracking branch 'origin/main' into fix_debug_upload_url_from_mgmt

# Conflicts:
#	client/proto/daemon.pb.go
#	management/internals/controllers/network_map/controller/controller.go
#	management/internals/controllers/network_map/nmaptest/runner.go
#	management/internals/shared/grpc/server.go
#	management/server/account.go
#	shared/management/proto/management.pb.go
#	upload-server/server/server.go

management/internals/controllers/network_map/controller/controller.go

@@ -24,7 +24,6 @@ import (
 	"github.com/netbirdio/netbird/management/internals/shared/requestbuffer"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
-	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/settings"
@@ -60,8 +59,6 @@ type Controller struct {
 
 	requestBuffer account.RequestBuffer
 
-	proxyController port_forwarding.Controller
-
 	integratedPeerValidator integrated_validator.IntegratedValidator
 
 	serverSupportedSyncMessageVersion sharedgrpc.SyncMessageVersion
@@ -87,7 +84,7 @@ type bufferAffectedUpdate struct {
 
 var _ network_map.Controller = (*Controller)(nil)
 
-func NewController(ctx context.Context, store store.Store, metrics telemetry.AppMetrics, peersUpdateManager network_map.PeersUpdateManager, requestBuffer account.RequestBuffer, integratedPeerValidator integrated_validator.IntegratedValidator, settingsManager settings.Manager, dnsDomain string, proxyController port_forwarding.Controller, ephemeralPeersManager ephemeral.Manager, config *config.Config, nmdataStore *networkmapdb.NetworkMapDBStoreImpl) *Controller {
+func NewController(ctx context.Context, store store.Store, metrics telemetry.AppMetrics, peersUpdateManager network_map.PeersUpdateManager, requestBuffer account.RequestBuffer, integratedPeerValidator integrated_validator.IntegratedValidator, settingsManager settings.Manager, dnsDomain string, ephemeralPeersManager ephemeral.Manager, config *config.Config, nmdataStore *networkmapdb.NetworkMapDBStoreImpl) *Controller {
 	nMetrics, err := newMetrics(metrics.UpdateChannelMetrics())
 	if err != nil {
 		log.Fatal(fmt.Errorf("error creating metrics: %w", err))
@@ -104,11 +101,10 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 		dnsDomain:               dnsDomain,
 		config:                  config,
 
-		proxyController:                              proxyController,
 		EphemeralPeersManager:                        ephemeralPeersManager,
 		serverSupportedSyncMessageVersion:            sharedgrpc.SyncMessageVersionFromConfig(config.HighestSupportedSyncMessageVersion),
 		perAccountServerSupportedSyncMessageVersions: sharedgrpc.SyncMessageVersionsFromMap(config.PerAccountHighestSupportedSyncMessageVersion),
-		nmdataStore:                                  nmdataStore,
+		nmdataStore: nmdataStore,
 	}
 
 	if nmdataStore != nil {
@@ -131,14 +127,20 @@ func (c *Controller) OnPeerConnected(ctx context.Context, accountID string, peer
 	return c.peersUpdateManager.CreateChannel(ctx, peerID), nil
 }
 
-func (c *Controller) OnPeerDisconnected(ctx context.Context, accountID string, peerID string) {
-	c.peersUpdateManager.CloseChannel(ctx, peerID)
+// OnPeerDisconnected closes the session's updates channel and schedules an ephemeral peer for
+// cleanup. It returns false without touching anything when a newer session owns the peer. A nil
+// session closes any registered channel.
+func (c *Controller) OnPeerDisconnected(ctx context.Context, accountID string, peerID string, session chan *network_map.UpdateMessage) bool {
+	if !c.peersUpdateManager.CloseSessionChannel(ctx, peerID, session) {
+		return false
+	}
 	peer, err := c.repo.GetPeerByID(ctx, accountID, peerID)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to get peer %s: %v", peerID, err)
-		return
+		return true
 	}
 	c.EphemeralPeersManager.OnPeerDisconnected(ctx, peer)
+	return true
 }
 
 // injectAllProxyPolicies prepares an account for the per-peer network-map
@@ -226,12 +228,6 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 	routers := account.GetResourceRoutersMap()
 	groupIDToUserIDs := account.GetActiveGroupUsers()
 
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMapsAll(ctx, accountID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return fmt.Errorf("failed to get proxy network maps: %v", err)
-	}
-
 	extraSetting, err := c.settingsManager.GetExtraSettings(ctx, accountID)
 	if err != nil {
 		return fmt.Errorf("failed to get flow enabled status: %v", err)
@@ -273,7 +269,6 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 			start = time.Now()
 
 			peerGroups := account.GetPeerGroups(p.ID)
-			proxyNetworkMap := proxyNetworkMaps[p.ID]
 			var update *proto.SyncResponse
 
 			commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
@@ -294,10 +289,7 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 				c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
 				start = time.Now()
-				// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
-				// the client merges it into Calculate()'s output the same
-				// way the legacy server did via NetworkMap.Merge.
-				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
+				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, components, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 				c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
 
 				c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -313,10 +305,6 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 
 			c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
-			if proxyNetworkMap != nil {
-				nmap.Merge(proxyNetworkMap)
-			}
-
 			start = time.Now()
 			update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, nmap, dnsDomain, postureChecks, dnsCache, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 			c.metrics.CountToSyncResponseDuration(time.Since(start))
@@ -451,7 +439,7 @@ func (c *Controller) sendUpdatesFromData(ctx context.Context, accountID string,
 				c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
 				start = time.Now()
-				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, nil, dnsDomain, postureChecks, nmData.AccountSettings, extraSettings, peerGroups, dnsFwdPort, c.config.DebugUpload.URL)
+				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, dnsDomain, postureChecks, nmData.AccountSettings, extraSettings, peerGroups, dnsFwdPort, c.config.DebugUpload.URL)
 				c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
 
 				c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -684,12 +672,6 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 	routers := account.GetResourceRoutersMap()
 	groupIDToUserIDs := account.GetActiveGroupUsers()
 
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMapsAll(ctx, accountID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return fmt.Errorf("failed to get proxy network maps: %v", err)
-	}
-
 	extraSetting, err := c.settingsManager.GetExtraSettings(ctx, accountID)
 	if err != nil {
 		return fmt.Errorf("failed to get flow enabled status: %v", err)
@@ -722,7 +704,6 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 			start = time.Now()
 
 			peerGroups := account.GetPeerGroups(p.ID)
-			proxyNetworkMap := proxyNetworkMaps[p.ID]
 			var update *proto.SyncResponse
 
 			commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
@@ -743,10 +724,7 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 				c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
 				start = time.Now()
-				// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
-				// the client merges it into Calculate()'s output the same
-				// way the legacy server did via NetworkMap.Merge.
-				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
+				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, components, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 				c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
 
 				c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -762,10 +740,6 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 
 			c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
-			if proxyNetworkMap != nil {
-				nmap.Merge(proxyNetworkMap)
-			}
-
 			start = time.Now()
 			update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(p), nil, nil, nmap, dnsDomain, postureChecks, dnsCache, types.TwinAccountSettings(account.Settings), extraSetting, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 			c.metrics.CountToSyncResponseDuration(time.Since(start))
@@ -843,19 +817,12 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
 		return fmt.Errorf("failed to get posture checks for peer %s: %v", peerId, err)
 	}
 
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peer.ID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return err
-	}
-
 	accountZones, err := c.repo.GetAccountZones(ctx, account.Id)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to get account zones: %v", err)
 		return err
 	}
 
-	proxyNetworkMap := proxyNetworkMaps[peer.ID]
 	extraSettings, err := c.settingsManager.GetExtraSettings(ctx, peer.AccountID)
 	if err != nil {
 		return fmt.Errorf("failed to get extra settings: %v", err)
@@ -881,10 +848,7 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
 		components := account.GetPeerNetworkMapComponents(
 			ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
 
-		// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
-		// the client merges it into Calculate()'s output the same
-		// way the legacy server did via NetworkMap.Merge.
-		update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(peer), nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSettings, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
+		update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(peer), nil, nil, components, dnsDomain, postureChecks, types.TwinAccountSettings(account.Settings), extraSettings, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 
 		c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
 			Update:      update,
@@ -897,10 +861,6 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
 	nmap := account.GetPeerNetworkMapFromComponents(
 		ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
 
-	if proxyNetworkMap != nil {
-		nmap.Merge(proxyNetworkMap)
-	}
-
 	update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, types.TwinPeer(peer), nil, nil, nmap, dnsDomain, postureChecks, dnsCache, types.TwinAccountSettings(account.Settings), extraSettings, maps.Keys(peerGroups), dnsFwdPort, c.config.DebugUpload.URL)
 
 	c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
@@ -951,17 +911,16 @@ func (c *Controller) BufferUpdateAccountPeers(ctx context.Context, accountID str
 
 // GetValidatedPeerWithComponents is the components-format counterpart of
 // GetValidatedPeerWithMap. It returns raw NetworkMapComponents for capable
-// peers along with the proxy NetworkMap fragment (BYOP / port-forwarding
-// data the legacy server folds in via NetworkMap.Merge). The gRPC layer
-// encodes both into the wire envelope. Callers must gate on capability
-// themselves before dispatching here — this method does NOT branch on it.
-func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+// peers, which the gRPC layer encodes into the wire envelope. Callers must
+// gate on capability themselves before dispatching here — this method does
+// NOT branch on it.
+func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, []*nmdata.PostureChecks, int64, error) {
 	if isRequiresApproval {
 		network, err := c.repo.GetAccountNetwork(ctx, accountID)
 		if err != nil {
-			return nil, nil, nil, nil, 0, err
+			return nil, nil, nil, 0, err
 		}
-		return peer, &types.NetworkMapComponents{Network: types.TwinNetwork(network)}, nil, nil, 0, nil
+		return peer, &types.NetworkMapComponents{Network: types.TwinNetwork(network)}, nil, 0, nil
 	}
 
 	if nmData := c.getNetworkMapData(ctx, accountID); nmData != nil {
@@ -970,39 +929,29 @@ func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequi
 
 	account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 	if err != nil {
-		return nil, nil, nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	// it's possible that the peer gets deleted between the call to "sendInitialSync()" and here, bail out in this case
 	if _, ok := account.Peers[peer.ID]; !ok {
-		return nil, nil, nil, nil, 0, fmt.Errorf("peer '%s' no longer exists", peer.ID)
+		return nil, nil, nil, 0, fmt.Errorf("peer '%s' no longer exists", peer.ID)
 	}
 
 	c.injectAllProxyPolicies(ctx, account)
 
 	approvedPeersMap, err := c.integratedPeerValidator.GetValidatedPeers(ctx, account.Id, types.TwinGroups(maps.Values(account.Groups)), types.TwinPeers(maps.Values(account.Peers)), account.Settings.Extra)
 	if err != nil {
-		return nil, nil, nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	postureChecks, err := c.getPeerPostureChecks(account, peer.ID)
 	if err != nil {
-		return nil, nil, nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	accountZones, err := c.repo.GetAccountZones(ctx, account.Id)
 	if err != nil {
-		return nil, nil, nil, nil, 0, err
-	}
-
-	// Fetch the proxy network map fragment for this peer alongside the
-	// components — same single-account-load path the streaming controller
-	// uses, so initial-sync delivers BYOP/forwarding patches synchronously
-	// instead of waiting for the next streaming push.
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peer.ID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return nil, nil, nil, nil, 0, err
+		return nil, nil, nil, 0, err
 	}
 
 	dnsDomain := c.GetDNSDomain(account.Settings)
@@ -1014,13 +963,12 @@ func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequi
 	components := account.GetPeerNetworkMapComponents(ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
 	dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
 
-	return peer, components, proxyNetworkMaps[peer.ID], postureChecks, dnsFwdPort, nil
+	return peer, components, postureChecks, dnsFwdPort, nil
 }
 
 // getValidatedPeerWithComponentsFromData is the account-free variant of
-// GetValidatedPeerWithComponents. The proxy network map fragment is omitted
-// like on the other nmdata paths.
-func (c *Controller) getValidatedPeerWithComponentsFromData(ctx context.Context, accountID string, peer *nbpeer.Peer, nmData *networkmap.NetworkMapData) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+// GetValidatedPeerWithComponents.
+func (c *Controller) getValidatedPeerWithComponentsFromData(ctx context.Context, accountID string, peer *nbpeer.Peer, nmData *networkmap.NetworkMapData) (*nbpeer.Peer, *types.NetworkMapComponents, []*nmdata.PostureChecks, int64, error) {
 	postureChecks := peerPostureChecksFromData(nmData, peer.ID)
 
 	dnsDomain := c.getDNSDomainFromData(nmData.AccountSettings)
@@ -1029,7 +977,7 @@ func (c *Controller) getValidatedPeerWithComponentsFromData(ctx context.Context,
 	components := nmData.GetPeerNetworkMapComponents(peer.ID, peersCustomZone)
 	dnsFwdPort := ComputeForwarderPortFromData(nmData.Peers, network_map.DnsForwarderPortMinVersion)
 
-	return peer, components, nil, postureChecks, dnsFwdPort, nil
+	return peer, components, postureChecks, dnsFwdPort, nil
 }
 
 // BufferUpdateAffectedPeers accumulates peer IDs and flushes them after the buffer interval.
@@ -1173,22 +1121,11 @@ func (c *Controller) GetValidatedPeerWithMap(ctx context.Context, isRequiresAppr
 	dnsDomain := c.GetDNSDomain(account.Settings)
 	peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
 
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peerID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return nil, nil, 0, err
-	}
-
 	resourcePolicies := account.GetResourcePoliciesMap()
 	routers := account.GetResourceRoutersMap()
 	groupIDToUserIDs := account.GetActiveGroupUsers()
 	networkMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
 
-	proxyNetworkMap, ok := proxyNetworkMaps[peerID]
-	if ok {
-		networkMap.Merge(proxyNetworkMap)
-	}
-
 	dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
 
 	return networkMap, postureChecks, dnsFwdPort, nil
@@ -1457,23 +1394,12 @@ func (c *Controller) GetNetworkMap(ctx context.Context, peerID string) (*types.N
 	dnsDomain := c.GetDNSDomain(account.Settings)
 	peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
 
-	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peerID, account.Peers)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
-		return nil, err
-	}
-
 	c.injectAllProxyPolicies(ctx, account)
 	resourcePolicies := account.GetResourcePoliciesMap()
 	routers := account.GetResourceRoutersMap()
 	groupIDToUserIDs := account.GetActiveGroupUsers()
 	networkMap := account.GetPeerNetworkMapFromComponents(ctx, peer.ID, peersCustomZone, accountZones, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
 
-	proxyNetworkMap, ok := proxyNetworkMaps[peer.ID]
-	if ok {
-		networkMap.Merge(proxyNetworkMap)
-	}
-
 	return networkMap, nil
 }
 

management/internals/controllers/network_map/controller/controller_disconnect_test.go

@@ -0,0 +1,78 @@
+package controller
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+)
+
+type recordingEphemeralManager struct {
+	disconnected []string
+}
+
+func (r *recordingEphemeralManager) LoadInitialPeers(context.Context)              {}
+func (r *recordingEphemeralManager) Stop()                                         {}
+func (r *recordingEphemeralManager) OnPeerConnected(context.Context, *nbpeer.Peer) {}
+func (r *recordingEphemeralManager) OnPeerDisconnected(_ context.Context, peer *nbpeer.Peer) {
+	r.disconnected = append(r.disconnected, peer.ID)
+}
+
+func TestOnPeerDisconnected_OwnSession(t *testing.T) {
+	ctx := context.Background()
+	repo := NewMockRepository(gomock.NewController(t))
+	ephemeralManager := &recordingEphemeralManager{}
+	updateManager := update_channel.NewPeersUpdateManager(nil)
+	c := Controller{repo: repo, peersUpdateManager: updateManager, EphemeralPeersManager: ephemeralManager}
+
+	session := updateManager.CreateChannel(ctx, "peer-1")
+	repo.EXPECT().GetPeerByID(gomock.Any(), "account-1", "peer-1").Return(&nbpeer.Peer{ID: "peer-1", Ephemeral: true}, nil)
+
+	require.True(t, c.OnPeerDisconnected(ctx, "account-1", "peer-1", session))
+	assert.False(t, updateManager.HasChannel("peer-1"))
+	assert.Equal(t, []string{"peer-1"}, ephemeralManager.disconnected)
+}
+
+func TestOnPeerDisconnected_NilSessionClosesOlderChannel(t *testing.T) {
+	ctx := context.Background()
+	repo := NewMockRepository(gomock.NewController(t))
+	ephemeralManager := &recordingEphemeralManager{}
+	updateManager := update_channel.NewPeersUpdateManager(nil)
+	c := Controller{repo: repo, peersUpdateManager: updateManager, EphemeralPeersManager: ephemeralManager}
+
+	older := updateManager.CreateChannel(ctx, "peer-1")
+	repo.EXPECT().GetPeerByID(gomock.Any(), "account-1", "peer-1").Return(&nbpeer.Peer{ID: "peer-1", Ephemeral: true}, nil)
+
+	require.True(t, c.OnPeerDisconnected(ctx, "account-1", "peer-1", nil))
+	assert.False(t, updateManager.HasChannel("peer-1"))
+	_, open := <-older
+	assert.False(t, open, "older channel must be closed")
+	assert.Equal(t, []string{"peer-1"}, ephemeralManager.disconnected)
+}
+
+func TestOnPeerDisconnected_NewerSessionOwnsPeer(t *testing.T) {
+	ctx := context.Background()
+	ephemeralManager := &recordingEphemeralManager{}
+	updateManager := update_channel.NewPeersUpdateManager(nil)
+	c := Controller{repo: NewMockRepository(gomock.NewController(t)), peersUpdateManager: updateManager, EphemeralPeersManager: ephemeralManager}
+
+	stale := updateManager.CreateChannel(ctx, "peer-1")
+	current := updateManager.CreateChannel(ctx, "peer-1")
+
+	require.False(t, c.OnPeerDisconnected(ctx, "account-1", "peer-1", stale))
+	require.True(t, updateManager.HasChannel("peer-1"))
+	updateManager.SendUpdate(ctx, "peer-1", &network_map.UpdateMessage{})
+	select {
+	case _, open := <-current:
+		assert.True(t, open, "newer session channel must stay open")
+	default:
+		t.Fatal("newer session channel did not receive the update")
+	}
+	assert.Empty(t, ephemeralManager.disconnected, "a live peer must not be scheduled for ephemeral cleanup")
+}

management/internals/controllers/network_map/controller/controller_test.go

@@ -122,11 +122,10 @@ func TestGetValidatedPeerWithComponents_DeletedPeer(t *testing.T) {
 	}
 
 	mockrequestBuffer.EXPECT().GetAccountWithBackpressure(gomock.Any(), gomock.Any()).Return(&types.Account{}, nil)
-	peer, components, netmap, posturechecks, dnsforwardPort, err := c.GetValidatedPeerWithComponents(context.TODO(), false, "test-account-id", &nbpeer.Peer{ID: "test-peer-id"})
+	peer, components, posturechecks, dnsforwardPort, err := c.GetValidatedPeerWithComponents(context.TODO(), false, "test-account-id", &nbpeer.Peer{ID: "test-peer-id"})
 
 	assert.Nil(t, peer)
 	assert.Nil(t, components)
-	assert.Nil(t, netmap)
 	assert.Nil(t, posturechecks)
 	assert.Equal(t, int64(0), dnsforwardPort)
 	assert.NotNil(t, err)

management/internals/controllers/network_map/controller/repository.go

@@ -44,7 +44,7 @@ func (r *repository) GetAccountNetwork(ctx context.Context, accountID string) (*
 }
 
 func (r *repository) GetAccountPeers(ctx context.Context, accountID string) ([]*peer.Peer, error) {
-	return r.store.GetAccountPeers(ctx, store.LockingStrengthNone, accountID, "", "")
+	return r.store.GetAccountPeers(ctx, store.LockingStrengthNone, accountID, "", "", "")
 }
 
 func (r *repository) GetAccountByPeerID(ctx context.Context, peerID string) (*types.Account, error) {

management/internals/controllers/network_map/controller/repository_mock.go

@@ -89,6 +89,21 @@ func (mr *MockRepositoryMockRecorder) GetAccountPeers(ctx, accountID any) *gomoc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountPeers", reflect.TypeOf((*MockRepository)(nil).GetAccountPeers), ctx, accountID)
 }
 
+// GetAccountServices mocks base method.
+func (m *MockRepository) GetAccountServices(ctx context.Context, accountID string) ([]*service.Service, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetAccountServices", ctx, accountID)
+	ret0, _ := ret[0].([]*service.Service)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetAccountServices indicates an expected call of GetAccountServices.
+func (mr *MockRepositoryMockRecorder) GetAccountServices(ctx, accountID any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountServices", reflect.TypeOf((*MockRepository)(nil).GetAccountServices), ctx, accountID)
+}
+
 // GetAccountZones mocks base method.
 func (m *MockRepository) GetAccountZones(ctx context.Context, accountID string) ([]*zones.Zone, error) {
 	m.ctrl.T.Helper()

management/internals/controllers/network_map/interface.go

@@ -24,7 +24,7 @@ type Controller interface {
 	UpdateAccountPeer(ctx context.Context, accountId string, peerId string) error
 	BufferUpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error
 	GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*nmdata.PostureChecks, int64, error)
-	GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
+	GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, []*nmdata.PostureChecks, int64, error)
 	GetDNSDomain(settings *types.Settings) string
 	StartWarmup(context.Context)
 	GetNetworkMap(ctx context.Context, peerID string) (*types.NetworkMap, error)
@@ -35,7 +35,10 @@ type Controller interface {
 	OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error
 	DisconnectPeers(ctx context.Context, accountId string, peerIDs []string)
 	OnPeerConnected(ctx context.Context, accountID string, peerID string) (chan *UpdateMessage, error)
-	OnPeerDisconnected(ctx context.Context, accountID string, peerID string)
+	// OnPeerDisconnected tears down the stream state of the peer's session, identified by its
+	// updates channel. It returns false and leaves everything untouched when a newer session
+	// owns the peer. A nil session is the newest session and tears down any registered channel.
+	OnPeerDisconnected(ctx context.Context, accountID string, peerID string, session chan *UpdateMessage) bool
 
 	TrackEphemeralPeer(ctx context.Context, peer *nbpeer.Peer)
 }

management/internals/controllers/network_map/interface_mock.go

@@ -127,16 +127,15 @@ func (mr *MockControllerMockRecorder) GetNetworkMap(ctx, peerID any) *gomock.Cal
 }
 
 // GetValidatedPeerWithComponents mocks base method.
-func (m *MockController) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *peer.Peer) (*peer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+func (m *MockController) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *peer.Peer) (*peer.Peer, *types.NetworkMapComponents, []*nmdata.PostureChecks, int64, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetValidatedPeerWithComponents", ctx, isRequiresApproval, accountID, p)
 	ret0, _ := ret[0].(*peer.Peer)
 	ret1, _ := ret[1].(*types.NetworkMapComponents)
-	ret2, _ := ret[2].(*types.NetworkMap)
-	ret3, _ := ret[3].([]*nmdata.PostureChecks)
-	ret4, _ := ret[4].(int64)
-	ret5, _ := ret[5].(error)
-	return ret0, ret1, ret2, ret3, ret4, ret5
+	ret2, _ := ret[2].([]*nmdata.PostureChecks)
+	ret3, _ := ret[3].(int64)
+	ret4, _ := ret[4].(error)
+	return ret0, ret1, ret2, ret3, ret4
 }
 
 // GetValidatedPeerWithComponents indicates an expected call of GetValidatedPeerWithComponents.
@@ -178,15 +177,17 @@ func (mr *MockControllerMockRecorder) OnPeerConnected(ctx, accountID, peerID any
 }
 
 // OnPeerDisconnected mocks base method.
-func (m *MockController) OnPeerDisconnected(ctx context.Context, accountID, peerID string) {
+func (m *MockController) OnPeerDisconnected(ctx context.Context, accountID, peerID string, session chan *UpdateMessage) bool {
 	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "OnPeerDisconnected", ctx, accountID, peerID)
+	ret := m.ctrl.Call(m, "OnPeerDisconnected", ctx, accountID, peerID, session)
+	ret0, _ := ret[0].(bool)
+	return ret0
 }
 
 // OnPeerDisconnected indicates an expected call of OnPeerDisconnected.
-func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID any) *gomock.Call {
+func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID, session any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDisconnected", reflect.TypeOf((*MockController)(nil).OnPeerDisconnected), ctx, accountID, peerID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDisconnected", reflect.TypeOf((*MockController)(nil).OnPeerDisconnected), ctx, accountID, peerID, session)
 }
 
 // OnPeersAdded mocks base method.

management/internals/controllers/network_map/nmaptest/canonicalize.go

@@ -123,7 +123,6 @@ func canonicalize(nm *proto.NetworkMap) {
 	slices.SortFunc(nm.Routes, cmpRoute)
 	slices.SortFunc(nm.FirewallRules, cmpFirewallRule)
 	slices.SortFunc(nm.RoutesFirewallRules, cmpRouteFirewallRule)
-	slices.SortFunc(nm.ForwardingRules, cmpForwardingRule)
 
 	for _, r := range nm.FirewallRules {
 		slices.SortFunc(r.SourcePrefixes, bytes.Compare)
@@ -353,16 +352,6 @@ func cmpRouteFirewallRule(a, b *proto.RouteFirewallRule) int {
 	return boolCmp(a.IsDynamic, b.IsDynamic)
 }
 
-func cmpForwardingRule(a, b *proto.ForwardingRule) int {
-	if a == nil || b == nil {
-		return boolCmp(a == nil, b == nil)
-	}
-	if c := cmp.Compare(int32(a.Protocol), int32(b.Protocol)); c != 0 {
-		return c
-	}
-	return bytes.Compare(a.TranslatedAddress, b.TranslatedAddress)
-}
-
 func portInfoKey(pi *proto.PortInfo) string {
 	if pi == nil {
 		return ""

management/internals/controllers/network_map/nmaptest/runner.go

@@ -243,7 +243,7 @@ func computeMode(t *testing.T, ctx context.Context, mode Mode, nmData *networkma
 	case ModeEnvelope:
 		components := nmData.GetPeerNetworkMapComponents(peerID, zone)
 		peerGroups := maps.Keys(nmData.GetPeerGroups(peerID))
-		resp := mgmtgrpc.ToComponentSyncResponse(ctx, nil, nil, nil, peer, nil, nil, components, nil,
+		resp := mgmtgrpc.ToComponentSyncResponse(ctx, nil, nil, nil, peer, nil, nil, components,
 			dnsDomain, nil, nmData.AccountSettings, nil, peerGroups, dnsFwdPort, "")
 		res, err := networkmap.EnvelopeToNetworkMap(ctx, resp.NetworkMapEnvelope, peer.Key, dnsDomain, false)
 		require.NoError(t, err, "expand envelope")

management/internals/controllers/network_map/update_channel.go

@@ -6,6 +6,10 @@ type PeersUpdateManager interface {
 	SendUpdate(ctx context.Context, peerID string, update *UpdateMessage)
 	CreateChannel(ctx context.Context, peerID string) chan *UpdateMessage
 	CloseChannel(ctx context.Context, peerID string)
+	// CloseSessionChannel closes the peer's channel and returns true, unless a channel of
+	// another session is registered, in which case it returns false and closes nothing.
+	// A nil session is the newest session and closes any registered channel.
+	CloseSessionChannel(ctx context.Context, peerID string, session chan *UpdateMessage) bool
 	CountStreams() int
 	HasChannel(peerID string) bool
 	CloseChannels(ctx context.Context, peerIDs []string)

management/internals/controllers/network_map/update_channel/updatechannel.go

@@ -133,6 +133,30 @@ func (p *PeersUpdateManager) CloseChannel(ctx context.Context, peerID string) {
 	p.closeChannel(ctx, peerID)
 }
 
+// CloseSessionChannel closes the peer's updates channel only while it is still session's channel.
+// It returns false when a newer stream has registered a different channel, which the stale
+// session must leave to its new owner. A nil session belongs to a stream that failed before
+// registering its own channel; it is the newest session and closes any channel still registered.
+func (p *PeersUpdateManager) CloseSessionChannel(ctx context.Context, peerID string, session chan *network_map.UpdateMessage) bool {
+	start := time.Now()
+
+	p.channelsMux.Lock()
+	defer func() {
+		p.channelsMux.Unlock()
+		if p.metrics != nil {
+			p.metrics.UpdateChannelMetrics().CountCloseChannelDuration(time.Since(start))
+		}
+	}()
+
+	if channel, ok := p.peerChannels[peerID]; ok && session != nil && channel != session {
+		log.WithContext(ctx).Debugf("skipped closing updates channel: peer %s is owned by a newer session", peerID)
+		return false
+	}
+
+	p.closeChannel(ctx, peerID)
+	return true
+}
+
 // GetAllConnectedPeers returns a copy of the connected peers map
 func (p *PeersUpdateManager) GetAllConnectedPeers() map[string]struct{} {
 	start := time.Now()

management/internals/controllers/network_map/update_channel/updatechannel_test.go

@@ -5,6 +5,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
@@ -84,3 +87,55 @@ func TestCloseChannel(t *testing.T) {
 		t.Error("Error closing the channel")
 	}
 }
+
+func TestCloseSessionChannel(t *testing.T) {
+	ctx := context.Background()
+	const peer = "test-close-session"
+
+	t.Run("own channel is closed", func(t *testing.T) {
+		peersUpdater := NewPeersUpdateManager(nil)
+		session := peersUpdater.CreateChannel(ctx, peer)
+
+		require.True(t, peersUpdater.CloseSessionChannel(ctx, peer, session))
+		assert.False(t, peersUpdater.HasChannel(peer))
+		_, open := <-session
+		assert.False(t, open, "own channel must be closed")
+	})
+
+	t.Run("newer session channel is kept", func(t *testing.T) {
+		peersUpdater := NewPeersUpdateManager(nil)
+		stale := peersUpdater.CreateChannel(ctx, peer)
+		current := peersUpdater.CreateChannel(ctx, peer)
+
+		require.False(t, peersUpdater.CloseSessionChannel(ctx, peer, stale))
+		require.True(t, peersUpdater.HasChannel(peer))
+		assert.Equal(t, current, peersUpdater.peerChannels[peer])
+
+		peersUpdater.SendUpdate(ctx, peer, &network_map.UpdateMessage{})
+		select {
+		case _, open := <-current:
+			assert.True(t, open, "newer session channel must stay open")
+		default:
+			t.Fatal("newer session channel did not receive the update")
+		}
+	})
+
+	t.Run("no registered channel", func(t *testing.T) {
+		peersUpdater := NewPeersUpdateManager(nil)
+		session := peersUpdater.CreateChannel(ctx, peer)
+		peersUpdater.CloseChannel(ctx, peer)
+
+		assert.True(t, peersUpdater.CloseSessionChannel(ctx, peer, session))
+		assert.True(t, peersUpdater.CloseSessionChannel(ctx, peer, nil))
+	})
+
+	t.Run("nil session closes the registered channel", func(t *testing.T) {
+		peersUpdater := NewPeersUpdateManager(nil)
+		current := peersUpdater.CreateChannel(ctx, peer)
+
+		require.True(t, peersUpdater.CloseSessionChannel(ctx, peer, nil))
+		assert.False(t, peersUpdater.HasChannel(peer))
+		_, open := <-current
+		assert.False(t, open, "registered channel must be closed")
+	})
+}