[management] use peer.Status.LastSeen field to detect stale peer.Status.Connected values (#7268)

management/server/store/sql_store_activity_test.go

@@ -109,6 +109,88 @@ func TestRefreshPeerLastSeenLeavesSessionStateAlone(t *testing.T) {
 	assert.Equal(t, int64(1234567890), peer.Status.SessionStartedAt, "session token must survive an activity write")
 }
 
+// TestPeerReadsAsDisconnectedWhenStale covers a management instance that stopped
+// without marking its peers disconnected: the stored flag stays set, so reads must
+// fall back to LastSeen, both when loading the peer and when filtering on it.
+func TestPeerReadsAsDisconnectedWhenStale(t *testing.T) {
+	ctx := context.Background()
+
+	tests := []struct {
+		name          string
+		lastSeen      time.Time
+		wantConnected bool
+	}{
+		{name: "seen within StaleAfter", lastSeen: time.Now().UTC().Add(-time.Minute), wantConnected: true},
+		{name: "not seen for longer than StaleAfter", lastSeen: time.Now().UTC().Add(-nbpeer.StaleAfter - time.Minute)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := newActivityTestStore(t)
+			stored := activityPeer(tt.lastSeen)
+			stored.Status.Connected = true
+			require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+			peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantConnected, peer.Status.Connected)
+
+			for _, connected := range []bool{true, false} {
+				peers, _, err := store.GetAccountPeersPaginated(ctx, LockingStrengthNone, activityAccountID, PaginationState{}, PeerFilters{Connected: &connected}, PeerSorting{})
+				require.NoError(t, err)
+				if connected == tt.wantConnected {
+					assert.Len(t, peers, 1, "filter connected=%t", connected)
+				} else {
+					assert.Empty(t, peers, "filter connected=%t", connected)
+				}
+			}
+		})
+	}
+}
+
+func TestHeartbeatPeerSession(t *testing.T) {
+	ctx := context.Background()
+	const session = int64(1234567890)
+	stale := time.Now().UTC().Add(-time.Hour)
+
+	tests := []struct {
+		name          string
+		storedSession int64
+		loginExpired  bool
+		session       int64
+		wantUpdated   bool
+	}{
+		{name: "stream that owns the session", storedSession: session, session: session, wantUpdated: true},
+		{name: "login expired while the stream stayed open", storedSession: session, loginExpired: true, session: session},
+		{name: "session replaced by a newer stream", storedSession: session + 1, session: session},
+		{name: "session already ended", session: session},
+		{name: "zero token never matches", session: 0},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := newActivityTestStore(t)
+			stored := activityPeer(stale)
+			stored.Status.SessionStartedAt = tt.storedSession
+			stored.Status.LoginExpired = tt.loginExpired
+			require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+			updated, err := store.HeartbeatPeerSession(ctx, activityAccountID, "activityPeer", tt.session)
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantUpdated, updated)
+
+			peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+			require.NoError(t, err)
+			assert.Equal(t, tt.storedSession, peer.Status.SessionStartedAt, "the session token must never change")
+			if tt.wantUpdated {
+				assert.True(t, peer.Status.Connected, "a heartbeat must undo a connected flag that was saved back as false")
+				assert.WithinDuration(t, time.Now().UTC(), peer.Status.LastSeen.UTC(), time.Minute)
+				return
+			}
+			assert.False(t, peer.Status.Connected)
+			assert.WithinDuration(t, stale, peer.Status.LastSeen.UTC(), time.Second)
+		})
+	}
+}
+
 func activityPeer(lastSeen time.Time) *nbpeer.Peer {
 	return &nbpeer.Peer{
 		ID:        "activityPeer",

management/server/store/sql_store_get_account_test.go

@@ -233,7 +233,7 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 			SystemManufacturer: "Dell Inc.",
 		},
 		Status: &nbpeer.PeerStatus{
-			LastSeen:         now.Add(-5 * time.Minute),
+			LastSeen:         now.Add(-1 * time.Minute),
 			Connected:        true,
 			LoginExpired:     false,
 			RequiresApproval: false,

management/server/store/sql_store_peer.go

@@ -226,6 +226,30 @@ func (s *SqlStore) RefreshPeerLastSeen(ctx context.Context, accountID, peerID st
 	return result.RowsAffected > 0, nil
 }
 
+// HeartbeatPeerSession stamps LastSeen with the database clock and marks the peer
+// connected, but only while sessionStartedAt still owns the peer and its login has not
+// expired. Setting Connected again undoes a save of a peer that was read as disconnected
+// while its heartbeats were delayed; the session match keeps an ended or replaced
+// session from coming back, and the login check keeps an expired peer disconnected
+// until it logs in again.
+func (s *SqlStore) HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
+	if sessionStartedAt == 0 {
+		return false, nil
+	}
+	result := s.db.WithContext(ctx).
+		Model(&nbpeer.Peer{}).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Where("peer_status_session_started_at = ? AND peer_status_login_expired = ?", sessionStartedAt, false).
+		Updates(map[string]any{
+			"peer_status_connected": true,
+			"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
+		})
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "heartbeat peer session: %v", result.Error)
+	}
+	return result.RowsAffected > 0, nil
+}
+
 func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Peer, error) {
 	const query = `SELECT id, account_id, key, ip, name, dns_label, user_id, ssh_key, ssh_enabled, login_expiration_enabled,
 	inactivity_expiration_enabled, last_login, created_at, ephemeral, extra_dns_labels, allow_extra_dns_labels, meta_hostname,
@@ -296,7 +320,7 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
 				p.Status.SessionStartedAt = peerStatusSessionStartedAt.Int64
 			}
 			if peerStatusConnected.Valid {
-				p.Status.Connected = peerStatusConnected.Bool
+				p.Status.Connected = nbpeer.IsConnected(peerStatusConnected.Bool, p.Status.LastSeen)
 			}
 			if peerStatusLoginExpired.Valid {
 				p.Status.LoginExpired = peerStatusLoginExpired.Bool
@@ -643,7 +667,12 @@ func (f *PeerFilters) AsConditions(db *gorm.DB) *gorm.DB {
 		db = db.Where("user_id=?", f.UserId)
 	}
 	if f.Connected != nil {
-		db = db.Where("peer_status_connected=?", *f.Connected)
+		lastSeenAfter := time.Now().UTC().Add(-nbpeer.StaleAfter)
+		if *f.Connected {
+			db = db.Where("peer_status_connected = ? AND peer_status_last_seen >= ?", true, lastSeenAfter)
+		} else {
+			db = db.Where("(peer_status_connected = ? OR peer_status_last_seen IS NULL OR peer_status_last_seen < ?)", false, lastSeenAfter)
+		}
 	}
 	if f.ApprovalRequired != nil {
 		db = db.Where("peer_status_requires_approval=?", *f.ApprovalRequired)

management/server/store/sql_store_peer_test.go

@@ -968,6 +968,7 @@ func TestSqlStore_MarkPeersDisconnectedIfSameSession(t *testing.T) {
 		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,
+			"peer_status_last_seen":          time.Now().UTC(),
 		}).Error)
 	}
 
@@ -1020,6 +1021,7 @@ func TestSqlStore_MarkPeersDisconnectedIfSameSessionContinuesAfterFailedChunk(t
 		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,
+			"peer_status_last_seen":          time.Now().UTC(),
 		}).Error)
 	}
 

management/server/store/sql_store_peers_filter_test.go

@@ -3,12 +3,21 @@ package store
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestSqlStore_FilterPeers(t *testing.T) {
 	runTestForAllEnginesNoSqliteSeed(t, "../testdata/peers_pagination.sql", func(t *testing.T, engine string, store Store) {
+		for _, peerID := range []string{"peer-id-one", "peer-id-two"} {
+			peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, "account-1", peerID)
+			require.NoError(t, err)
+			peer.Status.Connected = true
+			peer.Status.LastSeen = time.Now().UTC()
+			require.NoError(t, store.SavePeerStatus(context.Background(), "account-1", peerID, *peer.Status))
+		}
+
 		var tests = []struct {
 			description    string
 			filters        PeerFilters

management/server/store/sql_store_setup_key_test.go

@@ -22,7 +22,7 @@ func TestSqlite_GetSetupKeyBySecret(t *testing.T) {
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 
-	plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
+	plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB" // gitleaks:allow
 	hashedKey := sha256.Sum256([]byte(plainKey))
 	encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
 
@@ -47,7 +47,7 @@ func TestSqlite_incrementSetupKeyUsage(t *testing.T) {
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 
-	plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
+	plainKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB" // gitleaks:allow
 	hashedKey := sha256.Sum256([]byte(plainKey))
 	encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
 

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)
+	// HeartbeatPeerSession stamps LastSeen with the database clock and marks the
+	// peer connected, but only while the stored SessionStartedAt equals
+	// sessionStartedAt and the peer's login has not expired. Returns true when
+	// the update happened.
+	HeartbeatPeerSession(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

management/server/store/store_mock.go

@@ -3185,6 +3185,21 @@ func (mr *MockStoreMockRecorder) HasGatewayEndpointByOtherAccount(ctx, host, acc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayEndpointByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayEndpointByOtherAccount), ctx, host, accountID)
 }
 
+// HeartbeatPeerSession mocks base method.
+func (m *MockStore) HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "HeartbeatPeerSession", ctx, accountID, peerID, sessionStartedAt)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// HeartbeatPeerSession indicates an expected call of HeartbeatPeerSession.
+func (mr *MockStoreMockRecorder) HeartbeatPeerSession(ctx, accountID, peerID, sessionStartedAt any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeartbeatPeerSession", reflect.TypeOf((*MockStore)(nil).HeartbeatPeerSession), ctx, accountID, peerID, sessionStartedAt)
+}
+
 // IncrementAgentNetworkConsumption mocks base method.
 func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
 	m.ctrl.T.Helper()