[management] batch disconnect and lock free reconnect (#8373)

management/server/account.go

@@ -2033,13 +2033,13 @@ func domainIsUpToDate(domain string, domainCategory string, userAuth auth.UserAu
 // derived from syncTime (the moment the gRPC stream opened). Any
 // concurrent stream that started earlier loses the optimistic-lock race
 // in MarkPeerConnected and bails without writing.
-func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
-	peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP}, accountID)
+func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+	peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP, Settings: settings}, accountID)
 	if err != nil {
 		return nil, nil, nil, 0, fmt.Errorf("error syncing peer: %w", err)
 	}
 
-	if err := am.MarkPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap); err != nil {
+	if err := am.markPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap, settings); err != nil {
 		log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peerPubKey, err)
 	}
 
@@ -2059,6 +2059,17 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
 	return nil
 }
 
+// MarkPeersDisconnected marks the peers of the given sessions disconnected in batched writes,
+// skipping every peer whose stored session differs from the given one.
+func (am *DefaultAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+	marked, err := am.Store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+	if err != nil {
+		return fmt.Errorf("mark peers disconnected: %w", err)
+	}
+	log.WithContext(ctx).Infof("marked %d of %d peers as disconnected", marked, len(sessions))
+	return nil
+}
+
 func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error {
 	accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peerPubKey)
 	if err != nil {

management/server/account/manager.go

@@ -124,8 +124,10 @@ type Manager interface {
 	UpdateIntegratedValidator(ctx context.Context, accountID, userID, validator string, groups []string) error
 	GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error)
 	GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, map[string]string, error)
-	SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
+	SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
 	OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error
+	CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
+	MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error
 	SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error
 	FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 	GetAccountIDForPeerKey(ctx context.Context, peerKey string) (string, error)

management/server/account/manager_mock.go

@@ -146,6 +146,20 @@ func (mr *MockManagerMockRecorder) BuildUserInfosForAccount(ctx, accountID, init
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BuildUserInfosForAccount", reflect.TypeOf((*MockManager)(nil).BuildUserInfosForAccount), ctx, accountID, initiatorUserID, accountUsers)
 }
 
+// CheckPeerMaySync mocks base method.
+func (m *MockManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CheckPeerMaySync", ctx, accountID, peerPubKey, settings)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CheckPeerMaySync indicates an expected call of CheckPeerMaySync.
+func (mr *MockManagerMockRecorder) CheckPeerMaySync(ctx, accountID, peerPubKey, settings any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckPeerMaySync", reflect.TypeOf((*MockManager)(nil).CheckPeerMaySync), ctx, accountID, peerPubKey, settings)
+}
+
 // CreateGroup mocks base method.
 func (m *MockManager) CreateGroup(ctx context.Context, accountID, userID string, group *types.Group) error {
 	m.ctrl.T.Helper()
@@ -1369,6 +1383,20 @@ func (mr *MockManagerMockRecorder) MarkPeerDisconnected(ctx, peerKey, accountID,
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeerDisconnected), ctx, peerKey, accountID, sessionStartedAt)
 }
 
+// MarkPeersDisconnected mocks base method.
+func (m *MockManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "MarkPeersDisconnected", ctx, sessions)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// MarkPeersDisconnected indicates an expected call of MarkPeersDisconnected.
+func (mr *MockManagerMockRecorder) MarkPeersDisconnected(ctx, sessions any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeersDisconnected), ctx, sessions)
+}
+
 // OnPeerDisconnected mocks base method.
 func (m *MockManager) OnPeerDisconnected(ctx context.Context, accountID, peerPubKey string, streamStartTime time.Time) error {
 	m.ctrl.T.Helper()
@@ -1569,9 +1597,9 @@ func (mr *MockManagerMockRecorder) StoreEvent(ctx, initiatorID, targetID, accoun
 }
 
 // SyncAndMarkPeer mocks base method.
-func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime)
+	ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
 	ret0, _ := ret[0].(*peer.Peer)
 	ret1, _ := ret[1].(*types.NetworkMap)
 	ret2, _ := ret[2].([]*nmdata.PostureChecks)
@@ -1581,9 +1609,9 @@ func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey
 }
 
 // SyncAndMarkPeer indicates an expected call of SyncAndMarkPeer.
-func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime any) *gomock.Call {
+func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
 }
 
 // SyncPeer mocks base method.

management/server/account_test.go

@@ -3977,7 +3977,7 @@ func BenchmarkSyncAndMarkPeer(b *testing.B) {
 			b.ResetTimer()
 			start := time.Now()
 			for i := 0; i < b.N; i++ {
-				_, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC())
+				_, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC(), nil)
 				assert.NoError(b, err)
 			}
 

management/server/management_proto_test.go

@@ -17,8 +17,10 @@ import (
 	"go.uber.org/mock/gomock"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/status"
 
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/formatter/hook"
@@ -824,3 +826,130 @@ func testLoginPerformance(t *testing.T, loginCalls []func() error) {
 	tAvg := tSum / time.Duration(len(durations))
 	t.Logf("Min: %v, Max: %v, Avg: %v", tMin, tMax, tAvg)
 }
+
+type expireOnConnectStore struct {
+	store.Store
+	once sync.Once
+	err  error
+}
+
+func (s *expireOnConnectStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error) {
+	updated, err := s.Store.MarkPeerConnectedIfNewerSession(ctx, accountID, peerID, newSessionStartedAt)
+	if err != nil {
+		return updated, err
+	}
+
+	s.once.Do(func() {
+		peer, getErr := s.Store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
+		if getErr != nil {
+			s.err = getErr
+			return
+		}
+		peer.Status.LoginExpired = true
+		s.err = s.Store.SavePeerStatus(ctx, accountID, peerID, *peer.Status)
+	})
+	return updated, err
+}
+
+func Test_SyncRejectsPeerExpiredWhileConnecting(t *testing.T) {
+	mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
+		Datadir: t.TempDir(),
+	})
+	defer cleanup()
+	require.NoError(t, err)
+	defer mgmtServer.GracefulStop()
+
+	client, clientConn, err := createRawClient(mgmtAddr)
+	require.NoError(t, err)
+	defer clientConn.Close()
+
+	key, err := wgtypes.GeneratePrivateKey()
+	require.NoError(t, err)
+	_, err = loginPeerWithValidSetupKey(key, client)
+	require.NoError(t, err)
+
+	serverKey, err := getServerKey(client)
+	require.NoError(t, err)
+
+	decorated := &expireOnConnectStore{Store: am.Store}
+	am.Store = decorated
+
+	message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     message,
+	})
+	require.NoError(t, err)
+
+	resp := &mgmtProto.EncryptedMessage{}
+	for err == nil {
+		err = sync.RecvMsg(resp)
+	}
+	require.Equal(t, codes.PermissionDenied, status.Code(err), "sync stream ended with: %v", err)
+	require.NoError(t, decorated.err)
+
+	peer, err := am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+	require.NoError(t, err)
+	require.True(t, peer.Status.LoginExpired)
+	require.Eventually(t, func() bool {
+		peer, err = am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+		return err == nil && !peer.Status.Connected
+	}, 5*time.Second, 50*time.Millisecond)
+}
+
+type settingsCountingStore struct {
+	store.Store
+	reads atomic.Int32
+}
+
+func (s *settingsCountingStore) GetAccountSettings(ctx context.Context, lockStrength store.LockingStrength, accountID string) (*types.Settings, error) {
+	s.reads.Add(1)
+	return s.Store.GetAccountSettings(ctx, lockStrength, accountID)
+}
+
+func Test_LoginAndSyncReuseRequestSettings(t *testing.T) {
+	mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
+		Datadir: t.TempDir(),
+	})
+	defer cleanup()
+	require.NoError(t, err)
+	defer mgmtServer.GracefulStop()
+
+	client, clientConn, err := createRawClient(mgmtAddr)
+	require.NoError(t, err)
+	defer clientConn.Close()
+
+	key, err := wgtypes.GeneratePrivateKey()
+	require.NoError(t, err)
+	_, err = loginPeerWithValidSetupKey(key, client)
+	require.NoError(t, err)
+
+	counting := &settingsCountingStore{Store: am.Store}
+	am.Store = counting
+
+	_, err = loginPeerWithValidSetupKey(key, client)
+	require.NoError(t, err)
+	require.Zero(t, counting.reads.Load(), "login must not read settings through the account manager")
+
+	serverKey, err := getServerKey(client)
+	require.NoError(t, err)
+	message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{WgPubKey: key.PublicKey().String(), Body: message})
+	require.NoError(t, err)
+	require.NoError(t, sync.RecvMsg(&mgmtProto.EncryptedMessage{}))
+
+	require.Eventually(t, func() bool {
+		peer, err := counting.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+		return err == nil && peer.Status.Connected
+	}, 5*time.Second, 50*time.Millisecond)
+	require.Zero(t, counting.reads.Load(), "sync must not read settings through the account manager")
+}

management/server/mock_server/account_mock.go

@@ -42,7 +42,9 @@ type MockAccountManager struct {
 	GetPeersFunc                          func(ctx context.Context, accountID, userID, nameFilter, ipFilter, macFilter string) ([]*nbpeer.Peer, int, error)
 	MarkPeerConnectedFunc                 func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
 	MarkPeerDisconnectedFunc              func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64) error
-	SyncAndMarkPeerFunc                   func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
+	MarkPeersDisconnectedFunc             func(ctx context.Context, sessions []types.PeerSession) error
+	CheckPeerMaySyncFunc                  func(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
+	SyncAndMarkPeerFunc                   func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
 	DeletePeerFunc                        func(ctx context.Context, accountID, peerKey, userID string) error
 	GetNetworkMapFunc                     func(ctx context.Context, peerKey string) (*types.NetworkMap, error)
 	GetPeerNetworkFunc                    func(ctx context.Context, peerKey string) (*types.Network, error)
@@ -242,13 +244,27 @@ func (am *MockAccountManager) DeleteSetupKey(ctx context.Context, accountID, use
 	return status.Errorf(codes.Unimplemented, "method DeleteSetupKey is not implemented")
 }
 
-func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
 	if am.SyncAndMarkPeerFunc != nil {
-		return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime)
+		return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
 	}
 	return nil, nil, nil, 0, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
 }
 
+func (am *MockAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+	if am.MarkPeersDisconnectedFunc != nil {
+		return am.MarkPeersDisconnectedFunc(ctx, sessions)
+	}
+	return nil
+}
+
+func (am *MockAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+	if am.CheckPeerMaySyncFunc != nil {
+		return am.CheckPeerMaySyncFunc(ctx, accountID, peerPubKey, settings)
+	}
+	return status.Errorf(codes.Unimplemented, "method CheckPeerMaySync is not implemented")
+}
+
 func (am *MockAccountManager) OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error {
 	// Mirror DefaultAccountManager.OnPeerDisconnected: drive the fencing
 	// hook so tests that inject MarkPeerDisconnectedFunc actually observe

management/server/peer.go

@@ -88,6 +88,12 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
 // Disconnects use MarkPeerDisconnected and require the session to match
 // exactly; see PeerStatus.SessionStartedAt for the protocol.
 func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
+	return am.markPeerConnected(ctx, peerPubKey, accountID, sessionStartedAt, nmap, nil)
+}
+
+// markPeerConnected is MarkPeerConnected with the account settings the caller already
+// read; nil settings are loaded when needed.
+func (am *DefaultAccountManager) markPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap, settings *types.Settings) error {
 	start := time.Now()
 	defer func() {
 		am.metrics.AccountManagerMetrics().RecordPeerStatusUpdateDuration(telemetry.PeerStatusConnect, time.Since(start))
@@ -119,19 +125,13 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
 
 	log.WithContext(ctx).Debugf("mark peer %s connected", peer.ID)
 
-	if err = am.schedulePeerExpirations(ctx, accountID, peer); err != nil {
+	if err = am.schedulePeerExpirations(ctx, accountID, peer, settings); err != nil {
 		return err
 	}
 
-	// A login-expired peer reconnecting, or an embedded proxy peer flipping to
-	// connected (which triggers SynthesizePrivateServiceZones), must refresh the
-	// peers reachable from it. The embedded-proxy fan-out tolerates a dispatch error.
-	if peer.Status != nil && peer.Status.LoginExpired {
-		affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
-		if err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
-			return fmt.Errorf("notify network map controller of peer update: %w", err)
-		}
-	}
+	// An embedded proxy peer flipping to connected (which triggers
+	// SynthesizePrivateServiceZones) must refresh the peers reachable from it.
+	// The fan-out tolerates a dispatch error.
 	if peer.ProxyMeta.Embedded {
 		affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
 		if err := am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
@@ -166,11 +166,11 @@ func (am *DefaultAccountManager) writePeerStatus(ctx context.Context, accountID
 
 // schedulePeerExpirations arms the account's login/inactivity expiration
 // timers for an SSO peer that just connected.
-func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
+func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer, settings *types.Settings) error {
 	if !peer.AddedWithSSOLogin() {
 		return nil
 	}
-	settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	settings, err := am.accountSettings(ctx, accountID, settings)
 	if err != nil {
 		return err
 	}
@@ -1087,37 +1087,22 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
 	var metaDiff nbpeer.MetaDiff
 	var err error
 
-	settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	settings, err := am.accountSettings(ctx, accountID, sync.Settings)
 	if err != nil {
 		return nil, nil, nil, 0, err
 	}
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
-		peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, sync.WireGuardPubKey)
-		if err != nil {
-			return status.NewPeerNotRegisteredError()
-		}
+	peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, am.Store, store.LockingStrengthNone, sync, settings)
+	if err != nil {
+		return nil, nil, nil, 0, err
+	}
 
-		if peer.UserID != "" {
-			user, err := transaction.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
-			if err != nil {
+	if metaDiff.Updated() {
+		err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
+			peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, transaction, store.LockingStrengthUpdate, sync, settings)
+			if err != nil || !metaDiff.Updated() {
 				return err
 			}
-
-			if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
-				return err
-			}
-		}
-
-		if peerLoginExpired(ctx, peer, settings) {
-			return status.NewPeerLoginExpiredError()
-		}
-
-		oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
-		newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
-		metaDiff = peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
-		ipv6CapabilityChanged = oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
-		if metaDiff.Updated() {
 			am.metrics.AccountManagerMetrics().CountPeerMetUpdate()
 			log.WithContext(ctx).Tracef("peer %s metadata updated", peer.ID)
 			if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
@@ -1126,11 +1111,11 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
 			if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
 				return fmt.Errorf("increment network serial: %w", err)
 			}
+			return nil
+		})
+		if err != nil {
+			return nil, nil, nil, 0, err
 		}
-		return nil
-	})
-	if err != nil {
-		return nil, nil, nil, 0, err
 	}
 
 	peerGroupIDs, err := getPeerGroupIDs(ctx, am.Store, accountID, peer.ID)
@@ -1160,6 +1145,82 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
 	return peer, nmap, resPostureChecks, dnsFwdPort, nil
 }
 
+// prepareSyncPeer loads the syncing peer with lockStrength, checks that it may sync and
+// applies the reported meta to it in memory. The returned diff tells whether the peer
+// has to be saved.
+func (am *DefaultAccountManager) prepareSyncPeer(ctx context.Context, st store.Store, lockStrength store.LockingStrength, sync types.PeerSync, settings *types.Settings) (*nbpeer.Peer, nbpeer.MetaDiff, bool, error) {
+	peer, err := st.GetPeerByPeerPubKey(ctx, lockStrength, sync.WireGuardPubKey)
+	if err != nil {
+		return nil, nbpeer.MetaDiff{}, false, peerLookupError(err)
+	}
+
+	if err = checkPeerMaySync(ctx, st, peer, settings); err != nil {
+		return nil, nbpeer.MetaDiff{}, false, err
+	}
+
+	oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
+	newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
+	metaDiff := peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
+	ipv6CapabilityChanged := oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
+
+	return peer, metaDiff, ipv6CapabilityChanged, nil
+}
+
+// CheckPeerMaySync reads the peer again and rejects it like SyncPeer does when its owner
+// is blocked or its login expired. Nil settings are loaded from the store. The gRPC
+// server calls it once the sync stream can receive updates, so an expiry committed
+// while the stream was being set up is not missed.
+func (am *DefaultAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+	settings, err := am.accountSettings(ctx, accountID, settings)
+	if err != nil {
+		return err
+	}
+
+	peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, peerPubKey)
+	if err != nil {
+		return peerLookupError(err)
+	}
+
+	return checkPeerMaySync(ctx, am.Store, peer, settings)
+}
+
+// peerLookupError maps a missing peer to the not-registered error that makes the client
+// register again, and keeps any other store error so a transient failure is retried.
+func peerLookupError(err error) error {
+	if s, ok := status.FromError(err); ok && s.Type() == status.NotFound {
+		return status.NewPeerNotRegisteredError()
+	}
+	return err
+}
+
+// accountSettings returns settings when the caller already read them for this request
+// and loads them from the store otherwise.
+func (am *DefaultAccountManager) accountSettings(ctx context.Context, accountID string, settings *types.Settings) (*types.Settings, error) {
+	if settings != nil {
+		return settings, nil
+	}
+	return am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+}
+
+// checkPeerMaySync rejects a peer whose owner is blocked or whose login expired.
+func checkPeerMaySync(ctx context.Context, st store.Store, peer *nbpeer.Peer, settings *types.Settings) error {
+	if peer.UserID != "" {
+		user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
+		if err != nil {
+			return err
+		}
+
+		if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
+			return err
+		}
+	}
+
+	if peerLoginExpired(ctx, peer, settings) {
+		return status.NewPeerLoginExpiredError()
+	}
+	return nil
+}
+
 // metaDiffAffectsPosture reports whether the meta change flips the verdict of any of
 // the peer's posture checks, replaying them against the old and new state.
 func metaDiffAffectsPosture(diff *nbpeer.MetaDiff, checks []*nmdata.PostureChecks) bool {
@@ -1206,8 +1267,8 @@ func (am *DefaultAccountManager) syncPeerAffectedPeers(ctx context.Context, acco
 	return affectedPeerIDsFromNetworkMap(nmap, peerID)
 }
 
-// markConnectedAffectedPeers resolves the peers affected when a peer connects
-// (login-expiry reconnect or embedded-proxy connect). The connecting peer's
+// markConnectedAffectedPeers resolves the peers affected when an embedded proxy
+// peer connects. The connecting peer's
 // network map already lists them bidirectionally — the synthesized
 // private-service policy puts proxy access-group members in the proxy peer's own
 // map, and these edges carry no source-posture gate. An invalid peer has an
@@ -1249,8 +1310,13 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
 	// when the client sends a login request with a JWT which is used to get the user ID,
 	// it means that the client has already checked if it needs login and had been through the SSO flow
 	// so, we can skip this check and directly proceed with the login
+	settings, err := am.accountSettings(ctx, accountID, login.Settings)
+	if err != nil {
+		return nil, nil, nil, false, err
+	}
+
 	if login.UserID == "" {
-		err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
+		err = am.checkIFPeerNeedsLoginWithoutLock(ctx, login, settings)
 		if err != nil {
 			return nil, nil, nil, false, err
 		}
@@ -1260,56 +1326,55 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
 	var shouldStorePeer, shouldUpdatePeers bool
 	var peerGroupIDs []string
 
-	settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	peer, err = am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
 	if err != nil {
 		return nil, nil, nil, false, err
 	}
+	if err = am.checkPeerLogin(ctx, am.Store, peer, login); err != nil {
+		return nil, nil, nil, false, err
+	}
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
-		peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
-		if err != nil {
-			return err
-		}
-
-		if login.UserID != "" {
-			if peer.UserID != login.UserID {
-				log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
-				return status.NewPeerLoginMismatchError()
-			}
-
-			changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
+	if peerLoginChangesPeer(ctx, peer, login, settings) {
+		err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
+			peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
 			if err != nil {
 				return err
 			}
-
-			if changed {
-				shouldStorePeer = true
-				shouldUpdatePeers = true
-			}
-		}
-
-		if peer.SSHKey != login.SSHKey {
-			peer.SSHKey = login.SSHKey
-			shouldStorePeer = true
-		}
-
-		if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
-			return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
-		}
-
-		if shouldStorePeer {
-			if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
+			if err = am.checkPeerLogin(ctx, transaction, peer, login); err != nil {
 				return err
 			}
-			if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
-				return fmt.Errorf("increment network serial: %w", err)
-			}
-		}
 
-		return nil
-	})
-	if err != nil {
-		return nil, nil, nil, false, err
+			if login.UserID != "" {
+				changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
+				if err != nil {
+					return err
+				}
+
+				if changed {
+					shouldStorePeer = true
+					shouldUpdatePeers = true
+				}
+			}
+
+			if peer.SSHKey != login.SSHKey {
+				peer.SSHKey = login.SSHKey
+				shouldStorePeer = true
+			}
+
+			if shouldStorePeer {
+				if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
+					return err
+				}
+				if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
+					return fmt.Errorf("increment network serial: %w", err)
+				}
+			}
+
+			return nil
+		})
+		if err != nil {
+			return nil, nil, nil, false, err
+		}
 	}
 
 	// This is needed to keep in memory for the peer config. Otherwise browser client will end in a retry loop
@@ -1341,6 +1406,40 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
 	return peer, network, postureChecks, enableSSH, nil
 }
 
+// checkPeerLogin rejects a login whose user does not own the peer, whose owner is
+// blocked, or that requests extra DNS labels the peer is not allowed to use.
+func (am *DefaultAccountManager) checkPeerLogin(ctx context.Context, st store.Store, peer *nbpeer.Peer, login types.PeerLogin) error {
+	if login.UserID != "" {
+		if peer.UserID != login.UserID {
+			log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
+			return status.NewPeerLoginMismatchError()
+		}
+
+		user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
+		if err != nil {
+			return err
+		}
+		if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
+			return err
+		}
+	}
+
+	if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
+		return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
+	}
+
+	return nil
+}
+
+// peerLoginChangesPeer reports whether logging in writes the peer: an SSO
+// re-login of an expired peer or a new SSH key.
+func peerLoginChangesPeer(ctx context.Context, peer *nbpeer.Peer, login types.PeerLogin, settings *types.Settings) bool {
+	if peer.SSHKey != login.SSHKey {
+		return true
+	}
+	return login.UserID != "" && peerLoginExpired(ctx, peer, settings)
+}
+
 // ExtendPeerSession refreshes the peer's SSO session deadline by updating
 // LastLogin after a successful JWT validation. The tunnel is untouched: no
 // network map sync, no peer reconnect.
@@ -1507,7 +1606,7 @@ func processPeerPostureChecks(policy *types.Policy, peerID string, peerGroupIDs
 // The NetBird client doesn't have a way to check if the peer needs login besides sending a login request
 // with no JWT token and usually no setup-key. As the client can send up to two login request to check if it is expired
 // and before starting the engine, we do the checks without an account lock to avoid piling up requests.
-func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, accountID string, login types.PeerLogin) error {
+func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, login types.PeerLogin, settings *types.Settings) error {
 	peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
 	if err != nil {
 		return err
@@ -1519,11 +1618,6 @@ func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Co
 		return nil
 	}
 
-	settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
-	if err != nil {
-		return err
-	}
-
 	if peerLoginExpired(ctx, peer, settings) {
 		return status.NewPeerLoginExpiredError()
 	}

management/server/peer_test.go

@@ -17,6 +17,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -3323,3 +3324,136 @@ func TestPeerWrites_BumpNetworkSerial(t *testing.T) {
 		assert.Greater(t, networkSerial(t, manager, account.Id), afterDisconnect, "embedded proxy connect must bump")
 	})
 }
+
+type txCountingStore struct {
+	store.Store
+	txs atomic.Int32
+}
+
+func (s *txCountingStore) ExecuteInTransaction(ctx context.Context, f func(store.Store) error) error {
+	s.txs.Add(1)
+	return s.Store.ExecuteInTransaction(ctx, f)
+}
+
+func TestPeerSyncAndLogin_LockOnlyOnChange(t *testing.T) {
+	manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+	counting := &txCountingStore{Store: manager.Store}
+	manager.Store = counting
+
+	current, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+	require.NoError(t, err)
+
+	t.Run("SyncPeer without meta change opens no transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: current.Meta}, account.Id)
+		require.NoError(t, err)
+		assert.Zero(t, counting.txs.Load())
+	})
+
+	t.Run("SyncPeer with meta change saves in one transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		meta := current.Meta
+		meta.Hostname = "lock-on-change-host"
+		_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: meta}, account.Id)
+		require.NoError(t, err)
+		assert.Equal(t, int32(1), counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "lock-on-change-host", saved.Meta.Hostname)
+	})
+
+	t.Run("LoginPeer without change opens no transaction", func(t *testing.T) {
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+
+		counting.txs.Store(0)
+		_, _, _, _, err = manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: saved.SSHKey, Meta: saved.Meta})
+		require.NoError(t, err)
+		assert.Zero(t, counting.txs.Load())
+	})
+
+	t.Run("LoginPeer with new ssh key saves in one transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlockonchange", Meta: current.Meta})
+		require.NoError(t, err)
+		assert.Equal(t, int32(1), counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey)
+	})
+
+	t.Run("LoginPeer rejects a user mismatch without a transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, UserID: "someone-else", SSHKey: "ssh-ed25519 AAAAother", Meta: current.Meta})
+		assert.Equal(t, status.NewPeerLoginMismatchError(), err)
+		assert.Zero(t, counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey, "a rejected login must not store its ssh key")
+	})
+
+	t.Run("LoginPeer rejects disallowed extra DNS labels without a transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlabels", ExtraDNSLabels: []string{"extra"}, Meta: current.Meta})
+		errStatus, ok := status.FromError(err)
+		require.True(t, ok, "expected a status error, got %v", err)
+		assert.Equal(t, status.PreconditionFailed, errStatus.Type())
+		assert.Zero(t, counting.txs.Load())
+	})
+}
+
+func TestCheckPeerMaySync(t *testing.T) {
+	manager, _, account, peer1, peer2, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+
+	require.NoError(t, manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+	expired, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+	require.NoError(t, err)
+	expired.Status.LoginExpired = true
+	require.NoError(t, manager.Store.SavePeer(ctx, account.Id, expired))
+	assert.Equal(t, status.NewPeerLoginExpiredError(), manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+	owner := types.NewRegularUser("blocked-owner", "", "")
+	owner.AccountID = account.Id
+	owner.Blocked = true
+	require.NoError(t, manager.Store.SaveUser(ctx, owner))
+	blocked, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer2.ID)
+	require.NoError(t, err)
+	blocked.UserID = owner.Id
+	require.NoError(t, manager.Store.SavePeer(ctx, account.Id, blocked))
+	err = manager.CheckPeerMaySync(ctx, account.Id, peer2.Key, nil)
+	errStatus, ok := status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.PermissionDenied, errStatus.Type())
+
+	assert.Equal(t, status.NewPeerNotRegisteredError(), manager.CheckPeerMaySync(ctx, account.Id, "unknown-key", nil))
+}
+
+type failingPeerLookupStore struct {
+	store.Store
+}
+
+func (s *failingPeerLookupStore) GetPeerByPeerPubKey(context.Context, store.LockingStrength, string) (*nbpeer.Peer, error) {
+	return nil, status.Errorf(status.Internal, "issue getting peer from store: connection reset")
+}
+
+func TestPeerLookupErrorKeepsTransientStoreErrors(t *testing.T) {
+	manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+	manager.Store = &failingPeerLookupStore{Store: manager.Store}
+
+	err := manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil)
+	errStatus, ok := status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+
+	_, _, _, _, err = manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key}, account.Id)
+	errStatus, ok = status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+}

management/server/store/sql_store_peer.go

@@ -106,7 +106,6 @@ func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountI
 			"peer_status_connected":          true,
 			"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
 			"peer_status_session_started_at": newSessionStartedAt,
-			"peer_status_login_expired":      false,
 		})
 	if result.Error != nil {
 		return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
@@ -142,6 +141,51 @@ func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accoun
 	return result.RowsAffected > 0, nil
 }
 
+// markPeersDisconnectedChunkSize bounds the sessions per statement so the
+// bound parameter count stays below the database limits.
+const markPeersDisconnectedChunkSize = 1000
+
+// MarkPeersDisconnectedIfSameSession is the batched form of
+// MarkPeerDisconnectedIfSameSession: one UPDATE per chunk of sessions, each
+// peer guarded by its own (account, peer, session) tuple. A failed chunk does
+// not stop the others, only a done ctx does; the count then covers the chunks
+// that succeeded and the error joins the failures.
+func (s *SqlStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error) {
+	var total int64
+	var errs []error
+	for start := 0; start < len(sessions); start += markPeersDisconnectedChunkSize {
+		if err := ctx.Err(); err != nil {
+			errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", err))
+			break
+		}
+		end := min(start+markPeersDisconnectedChunkSize, len(sessions))
+		tuples := make([][]any, 0, end-start)
+		for _, session := range sessions[start:end] {
+			if session.SessionStartedAt == 0 {
+				continue
+			}
+			tuples = append(tuples, []any{session.AccountID, session.PeerID, session.SessionStartedAt})
+		}
+		if len(tuples) == 0 {
+			continue
+		}
+		result := s.db.WithContext(ctx).
+			Model(&nbpeer.Peer{}).
+			Where("(account_id, id, peer_status_session_started_at) IN ?", tuples).
+			Updates(map[string]any{
+				"peer_status_connected":          false,
+				"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
+				"peer_status_session_started_at": int64(0),
+			})
+		if result.Error != nil {
+			errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", result.Error))
+			continue
+		}
+		total += result.RowsAffected
+	}
+	return total, errors.Join(errs...)
+}
+
 // ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
 func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
 	result := s.db.Model(&nbpeer.Peer{}).

management/server/store/sql_store_peer_test.go

@@ -3,16 +3,20 @@ package store
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
 	"reflect"
 	"sort"
+	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
 
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -941,3 +945,146 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
 		})
 	})
 }
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSession(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	sqlStore := store.(*SqlStore)
+
+	sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103, "peer-d": 104}
+	for id, session := range sessionOf {
+		require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+			ID:        id,
+			AccountID: accountID,
+			Key:       "key-" + id,
+			DNSLabel:  id,
+			IP:        netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+			IPv6:      netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+		}))
+		require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+			"peer_status_connected":          true,
+			"peer_status_session_started_at": session,
+		}).Error)
+	}
+
+	sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+	for i := 0; i < markPeersDisconnectedChunkSize+5; i++ {
+		sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+	}
+	sessions = append(sessions,
+		types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 999},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 0},
+	)
+
+	marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+	require.NoError(t, err)
+	assert.Equal(t, int64(3), marked)
+
+	for id, wantConnected := range map[string]bool{"peer-a": false, "peer-b": false, "peer-c": false, "peer-d": true} {
+		peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		assert.Equal(t, wantConnected, peer.Status.Connected, id)
+		wantSession := int64(0)
+		if wantConnected {
+			wantSession = sessionOf[id]
+		}
+		assert.Equal(t, wantSession, peer.Status.SessionStartedAt, id)
+	}
+}
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSessionContinuesAfterFailedChunk(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	sqlStore := store.(*SqlStore)
+
+	sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103}
+	for id, session := range sessionOf {
+		require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+			ID:        id,
+			AccountID: accountID,
+			Key:       "key-" + id,
+			DNSLabel:  id,
+			IP:        netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+			IPv6:      netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+		}))
+		require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+			"peer_status_connected":          true,
+			"peer_status_session_started_at": session,
+		}).Error)
+	}
+
+	sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+	for i := 0; i < markPeersDisconnectedChunkSize-1; i++ {
+		sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+	}
+	sessions = append(sessions,
+		types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+	)
+
+	chunkErr := errors.New("deadlock detected")
+	var updates atomic.Int32
+	require.NoError(t, sqlStore.db.Callback().Update().Before("gorm:update").Register("test:fail_first_chunk", func(db *gorm.DB) {
+		if updates.Add(1) == 1 {
+			_ = db.AddError(chunkErr)
+		}
+	}))
+
+	marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+	require.Error(t, err)
+	assert.ErrorContains(t, err, chunkErr.Error())
+	assert.Equal(t, int64(2), marked, "only the second chunk is counted")
+	assert.Equal(t, int32(2), updates.Load(), "the second chunk must still run")
+
+	for id, wantConnected := range map[string]bool{"peer-a": true, "peer-b": false, "peer-c": false} {
+		peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		assert.Equal(t, wantConnected, peer.Status.Connected, id)
+	}
+
+	cancelled, cancel := context.WithCancel(ctx)
+	cancel()
+	updates.Store(1)
+	marked, err = store.MarkPeersDisconnectedIfSameSession(cancelled, sessions)
+	assert.Zero(t, marked)
+	assert.ErrorContains(t, err, context.Canceled.Error())
+	assert.Equal(t, int32(1), updates.Load(), "no chunk runs once the context is done")
+	assert.Equal(t, 1, strings.Count(err.Error(), context.Canceled.Error()), "a done context is reported once")
+}
+
+func TestSqlStore_MarkPeerConnectedKeepsLoginExpired(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+		ID:        "expired-peer",
+		AccountID: accountID,
+		Key:       "key-expired-peer",
+		DNSLabel:  "expired-peer",
+		IP:        netip.AddrFrom4([4]byte{10, 9, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::9:101"),
+		Status:    &nbpeer.PeerStatus{LoginExpired: true},
+	}))
+
+	updated, err := store.MarkPeerConnectedIfNewerSession(ctx, accountID, "expired-peer", 1)
+	require.NoError(t, err)
+	require.True(t, updated)
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, "expired-peer")
+	require.NoError(t, err)
+	assert.True(t, peer.Status.Connected)
+	assert.True(t, peer.Status.LoginExpired, "connecting must not clear a login expiry")
+}

management/server/store/store.go

@@ -196,6 +196,11 @@ type Store interface {
 	// recorded by the database. Returns true when the update happened,
 	// false when a newer session has taken over.
 	MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
+	// MarkPeersDisconnectedIfSameSession applies the semantics of
+	// MarkPeerDisconnectedIfSameSession to many peers in batched statements.
+	// Sessions with a zero SessionStartedAt are skipped. Returns the number
+	// of peers that were marked disconnected.
+	MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error)
 	ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
 	DeletePeer(ctx context.Context, accountID string, peerID string) error
 

management/server/store/store_mock.go

@@ -3419,6 +3419,21 @@ func (mr *MockStoreMockRecorder) MarkPeerDisconnectedIfSameSession(ctx, accountI
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeerDisconnectedIfSameSession), ctx, accountID, peerID, sessionStartedAt)
 }
 
+// MarkPeersDisconnectedIfSameSession mocks base method.
+func (m *MockStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types3.PeerSession) (int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "MarkPeersDisconnectedIfSameSession", ctx, sessions)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// MarkPeersDisconnectedIfSameSession indicates an expected call of MarkPeersDisconnectedIfSameSession.
+func (mr *MockStoreMockRecorder) MarkPeersDisconnectedIfSameSession(ctx, sessions any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeersDisconnectedIfSameSession), ctx, sessions)
+}
+
 // MarkPendingJobsAsFailed mocks base method.
 func (m *MockStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
 	m.ctrl.T.Helper()

management/server/types/peer.go

@@ -18,6 +18,9 @@ type PeerSync struct {
 	// UpdateAccountPeers indicate updating account peers,
 	// which occurs when the peer's metadata is updated
 	UpdateAccountPeers bool
+	// Settings are the account settings the caller already read for this request.
+	// When nil they are loaded from the store.
+	Settings *Settings
 }
 
 // PeerLogin used as a data object between the gRPC API and Manager on Login request.
@@ -37,4 +40,14 @@ type PeerLogin struct {
 
 	// ExtraDNSLabels is a list of extra DNS labels that the peer wants to use
 	ExtraDNSLabels []string
+	// Settings are the account settings the caller already read for this request.
+	// When nil they are loaded from the store.
+	Settings *Settings
+}
+
+// PeerSession identifies the sync session a peer is currently owned by.
+type PeerSession struct {
+	AccountID        string
+	PeerID           string
+	SessionStartedAt int64
 }