Keep a modified peer's activation state across remove and re-add

client/internal/engine.go

@@ -863,19 +863,35 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 		}
 	}
 
-	// second, close all modified connections and remove them from the state map
+	// second, close all modified connections and remove them from the state map,
+	// remembering which of them were active
+	active := make(map[string]bool, len(modified))
 	for _, p := range modified {
-		err := e.removePeer(p.GetWgPubKey())
-		if err != nil {
+		peerPubKey := p.GetWgPubKey()
+		if state, err := e.statusRecorder.GetPeer(peerPubKey); err == nil {
+			active[peerPubKey] = state.ConnStatus != peer.StatusIdle
+		}
+		if err := e.removePeer(peerPubKey); err != nil {
 			return err
 		}
 	}
-	// third, add the peer connections again
+	// third, add the peer connections again, restoring each peer's activation
+	// state: under lazy connections a re-added peer starts idle, but the remote
+	// side of an established connection keeps its state and sends no further
+	// offers, so a previously active peer left idle cannot reconnect until the
+	// remote's connection expires.
 	for _, p := range modified {
-		err := e.addNewPeer(p)
-		if err != nil {
+		if err := e.addNewPeer(p); err != nil {
 			return err
 		}
+		if !active[p.GetWgPubKey()] {
+			continue
+		}
+		conn, ok := e.peerStore.PeerConn(p.GetWgPubKey())
+		if !ok {
+			continue
+		}
+		e.connMgr.ActivatePeer(e.ctx, conn)
 	}
 	return nil
 }

client/internal/engine_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/iface/wgproxy"
 	"github.com/netbirdio/netbird/client/internal/dns"
+	"github.com/netbirdio/netbird/client/internal/lazyconn"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
@@ -466,6 +467,128 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 	}
 }
 
+// TestEngine_ModifiedPeerKeepsActivationState verifies that a peer re-added by
+// modifyPeers keeps its previous activation state under lazy connections. A
+// modified peer is removed and re-added, and a re-add defaults to idle; the
+// remote side of an established connection keeps its state and sends no further
+// offers, so a previously active peer parked idle leaves the pair unable to
+// reconnect until the remote's connection expires.
+func TestEngine_ModifiedPeerKeepsActivationState(t *testing.T) {
+	key, err := wgtypes.GeneratePrivateKey()
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(CtxInitState(context.Background()))
+	defer cancel()
+
+	relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU)
+	engine := NewEngine(ctx, cancel, &EngineConfig{
+		WgIfaceName:    "utun103",
+		WgAddr:         wgaddr.MustParseWGAddress("100.64.0.1/24"),
+		WgPrivateKey:   key,
+		WgPort:         33101,
+		MTU:            iface.DefaultMTU,
+		LazyConnection: lazyconn.StateOn,
+	}, EngineServices{
+		SignalClient:   &signal.MockClient{},
+		MgmClient:      &mgmt.MockClient{},
+		RelayManager:   relayMgr,
+		StatusRecorder: peer.NewRecorder("https://mgm"),
+	}, MobileDependency{})
+
+	wgIface := &MockWGIface{
+		NameFunc: func() string { return "utun103" },
+		IsUserspaceBindFunc: func() bool {
+			return false
+		},
+		RemovePeerFunc: func(peerKey string) error {
+			return nil
+		},
+		AddressFunc: func() wgaddr.Address {
+			return wgaddr.Address{
+				IP:      netip.MustParseAddr("10.20.0.1"),
+				Network: netip.MustParsePrefix("10.20.0.0/24"),
+			}
+		},
+		UpdatePeerFunc: func(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error {
+			return nil
+		},
+	}
+	engine.wgInterface = wgIface
+	engine.routeManager = routemanager.NewManager(routemanager.ManagerConfig{
+		Context:          ctx,
+		PublicKey:        key.PublicKey().String(),
+		DNSRouteInterval: time.Minute,
+		WGInterface:      engine.wgInterface,
+		StatusRecorder:   engine.statusRecorder,
+		RelayManager:     relayMgr,
+	})
+	require.NoError(t, engine.routeManager.Init())
+	engine.dnsServer = &dns.MockServer{
+		UpdateDNSServerFunc: func(serial uint64, update nbdns.Config) error { return nil },
+	}
+	udpConn, err := net.ListenUDP("udp4", nil)
+	require.NoError(t, err)
+	engine.udpMux = udpmux.NewUniversalUDPMuxDefault(udpmux.UniversalUDPMuxParams{UDPConn: udpConn, MTU: 1280})
+	engine.ctx = ctx
+	engine.srWatcher = guard.NewSRWatcher(nil, nil, nil, icemaker.Config{})
+	engine.connMgr = NewConnMgr(engine.config, engine.statusRecorder, engine.peerStore, wgIface)
+	engine.connMgr.Start(ctx)
+
+	// No agent version: not lazy-capable, so the connection opens permanently.
+	activePeer := &mgmtProto.RemotePeerConfig{
+		WgPubKey:   "RRHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
+		AllowedIps: []string{"100.64.0.10/24"},
+	}
+	// Lazy-capable, never activated: managed as idle.
+	idlePeer := &mgmtProto.RemotePeerConfig{
+		WgPubKey:     "LLHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
+		AllowedIps:   []string{"100.64.0.11/24"},
+		AgentVersion: "development",
+	}
+
+	err = engine.updateNetworkMap(&mgmtProto.NetworkMap{
+		Serial:      1,
+		RemotePeers: []*mgmtProto.RemotePeerConfig{activePeer, idlePeer},
+	})
+	require.NoError(t, err)
+
+	state, err := engine.statusRecorder.GetPeer(activePeer.WgPubKey)
+	require.NoError(t, err)
+	require.Equal(t, peer.StatusConnecting, state.ConnStatus, "peer without lazy support should open a permanent connection")
+
+	state, err = engine.statusRecorder.GetPeer(idlePeer.WgPubKey)
+	require.NoError(t, err)
+	require.Equal(t, peer.StatusIdle, state.ConnStatus, "lazy-capable peer should be managed as idle")
+
+	// The active peer's agent version changes, as when a peer registered over the
+	// API logs in and fills in its meta; the idle peer's allowed IPs change. Both
+	// count as modified and are removed and re-added.
+	err = engine.updateNetworkMap(&mgmtProto.NetworkMap{
+		Serial: 2,
+		RemotePeers: []*mgmtProto.RemotePeerConfig{
+			{
+				WgPubKey:     activePeer.WgPubKey,
+				AllowedIps:   activePeer.AllowedIps,
+				AgentVersion: "development",
+			},
+			{
+				WgPubKey:     idlePeer.WgPubKey,
+				AllowedIps:   []string{"100.64.0.21/24"},
+				AgentVersion: "development",
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	state, err = engine.statusRecorder.GetPeer(activePeer.WgPubKey)
+	require.NoError(t, err)
+	assert.NotEqual(t, peer.StatusIdle, state.ConnStatus, "previously active peer should stay active after a modify")
+
+	state, err = engine.statusRecorder.GetPeer(idlePeer.WgPubKey)
+	require.NoError(t, err)
+	assert.Equal(t, peer.StatusIdle, state.ConnStatus, "previously idle peer should stay idle after a modify")
+}
+
 func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {
 	testCases := []struct {
 		name                 string