From 28dbd5898115c997e4a614b0674f21be30933834 Mon Sep 17 00:00:00 2001 From: pascal Date: Tue, 27 Oct 2026 13:36:40 +0100 Subject: [PATCH] review comments: skip update on nmap, fix other readers and add timeout to avoid staleness --- .../internals/network_map_db/pgsql/peer.go | 2 +- .../internals/network_map_db/shared_types.go | 6 +- .../network_map_db/shared_types_test.go | 32 ++++++++ .../internals/network_map_db/sqlite/peer.go | 2 +- .../shared/grpc/peer_update_handler.go | 14 ++++ .../shared/grpc/peer_update_handler_test.go | 62 +++++++++++++- management/internals/shared/grpc/server.go | 19 +++-- management/server/peer/peer.go | 27 ++++++ management/server/peer/peer_test.go | 20 +++++ .../server/store/sql_store_activity_test.go | 82 +++++++++++++++++++ .../store/sql_store_get_account_test.go | 2 +- management/server/store/sql_store_peer.go | 33 +++++++- .../store/sql_store_peers_filter_test.go | 9 ++ management/server/store/store.go | 5 ++ management/server/store/store_mock.go | 15 ++++ 15 files changed, 313 insertions(+), 17 deletions(-) diff --git a/management/internals/network_map_db/pgsql/peer.go b/management/internals/network_map_db/pgsql/peer.go index 108ad2a5e..452a8c825 100644 --- a/management/internals/network_map_db/pgsql/peer.go +++ b/management/internals/network_map_db/pgsql/peer.go @@ -11,7 +11,7 @@ import ( const ( GetPeersQuery = ` select id, key, ssh_key, dns_label, extra_dns_labels, user_id, ssh_enabled, login_expiration_enabled, last_login, ip, ipv6, - peer_status_requires_approval, peer_status_connected, proxy_meta_embedded, proxy_meta_cluster, + peer_status_requires_approval, peer_status_connected, peer_status_last_seen, proxy_meta_embedded, proxy_meta_cluster, meta_wt_version, meta_go_os, meta_os_version, meta_kernel_version, meta_network_addresses, meta_files, meta_certificates, meta_capabilities, meta_flags, meta_sync_message_version, location_country_code, location_city_name, location_connection_ip from peers diff --git a/management/internals/network_map_db/shared_types.go b/management/internals/network_map_db/shared_types.go index ca92d4eea..9ea3eac54 100644 --- a/management/internals/network_map_db/shared_types.go +++ b/management/internals/network_map_db/shared_types.go @@ -11,6 +11,7 @@ import ( "github.com/miekg/dns" "github.com/netbirdio/netbird/management/server/integrations/integrated_validator" + nbpeer "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/settings" "github.com/netbirdio/netbird/shared/management/networkmap" "github.com/netbirdio/netbird/shared/management/networkmap/nmdata" @@ -186,6 +187,7 @@ type Peer struct { IPv6 []byte `nmap:"json"` PeerStatusRequiresApproval sql.NullBool `nmap:"map_to:RequiresApproval"` PeerStatusConnected sql.NullBool `nmap:"skip"` + PeerStatusLastSeen sql.NullTime `nmap:"skip"` ProxyMetaEmbedded sql.NullBool `nmap:"skip"` ProxyMetaCluster sql.NullString `nmap:"skip"` MetaWtVersion sql.NullString `nmap:"skip"` @@ -318,8 +320,8 @@ func ConvertToNmdataPeers(peers []Peer) ([]nmdata.Peer, map[string][]*nmdata.Pee } dp.ProxyMeta.Cluster = p.ProxyMetaCluster.String // This is only used to build private service candidates, not connected peers are skipped - dp.Connected = p.PeerStatusConnected.Bool - if dp.ProxyMeta.Embedded && p.PeerStatusConnected.Bool { + dp.Connected = nbpeer.IsConnected(p.PeerStatusConnected.Bool, p.PeerStatusLastSeen.Time) + if dp.ProxyMeta.Embedded && dp.Connected { clusterToPeerIdx[p.ProxyMetaCluster.String] = append(clusterToPeerIdx[p.ProxyMetaCluster.String], &dp) } if p.MetaWtVersion.Valid { diff --git a/management/internals/network_map_db/shared_types_test.go b/management/internals/network_map_db/shared_types_test.go index 34b69d265..fb73b6712 100644 --- a/management/internals/network_map_db/shared_types_test.go +++ b/management/internals/network_map_db/shared_types_test.go @@ -1,10 +1,14 @@ package networkmapdb import ( + "database/sql" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + nbpeer "github.com/netbirdio/netbird/management/server/peer" ) func TestRecordTypeAndRdata(t *testing.T) { @@ -58,3 +62,31 @@ func TestConvertToNmdataPeers_RejectsMalformedCertificates(t *testing.T) { assert.Error(t, err, "a stored value that is not a list of chains must not load as no certificates") } + +// TestConvertToNmdataPeers_StaleEmbeddedProxyIsNotACandidate: a management instance that +// stops without marking its peers disconnected leaves the connected flag set. An embedded +// proxy peer in that state must not become a private service DNS target, or clients +// resolve the service to a tunnel IP that no longer answers. +func TestConvertToNmdataPeers_StaleEmbeddedProxyIsNotACandidate(t *testing.T) { + embedded := func(id string, lastSeen time.Time) Peer { + return Peer{ + ID: id, + PeerStatusConnected: sql.NullBool{Bool: true, Valid: true}, + PeerStatusLastSeen: sql.NullTime{Time: lastSeen, Valid: true}, + ProxyMetaEmbedded: sql.NullBool{Bool: true, Valid: true}, + ProxyMetaCluster: sql.NullString{String: "cluster", Valid: true}, + } + } + + peers, clusterToPeers, err := ConvertToNmdataPeers([]Peer{ + embedded("live", time.Now().Add(-time.Minute)), + embedded("stale", time.Now().Add(-nbpeer.StaleAfter-time.Minute)), + }) + + require.NoError(t, err) + require.Len(t, peers, 2) + assert.True(t, peers[0].Connected) + assert.False(t, peers[1].Connected, "a peer not seen for longer than StaleAfter must read as disconnected") + require.Len(t, clusterToPeers["cluster"], 1) + assert.Equal(t, "live", clusterToPeers["cluster"][0].ID) +} diff --git a/management/internals/network_map_db/sqlite/peer.go b/management/internals/network_map_db/sqlite/peer.go index c4570658d..57f6ada12 100644 --- a/management/internals/network_map_db/sqlite/peer.go +++ b/management/internals/network_map_db/sqlite/peer.go @@ -10,7 +10,7 @@ import ( const ( GetPeersQuery = ` select id, key, ssh_key, dns_label, extra_dns_labels, user_id, ssh_enabled, login_expiration_enabled, last_login, ip, ipv6, - peer_status_requires_approval, peer_status_connected, proxy_meta_embedded, proxy_meta_cluster, + peer_status_requires_approval, peer_status_connected, peer_status_last_seen, proxy_meta_embedded, proxy_meta_cluster, meta_wt_version, meta_go_os, meta_os_version, meta_kernel_version, meta_network_addresses, meta_files, meta_certificates, meta_capabilities, meta_flags, meta_sync_message_version, location_country_code, location_city_name, location_connection_ip from peers diff --git a/management/internals/shared/grpc/peer_update_handler.go b/management/internals/shared/grpc/peer_update_handler.go index 5037083c1..53734993f 100644 --- a/management/internals/shared/grpc/peer_update_handler.go +++ b/management/internals/shared/grpc/peer_update_handler.go @@ -6,6 +6,7 @@ import ( "github.com/netbirdio/netbird/encryption" "github.com/netbirdio/netbird/management/internals/controllers/network_map" + nbpeer "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/telemetry" "github.com/netbirdio/netbird/shared/management/certposture" "github.com/netbirdio/netbird/shared/management/proto" @@ -21,6 +22,7 @@ func PeerUpdateHandlerFactory( secretsManager SecretsManager, challenger *certposture.Challenger, onChallengeStamped func(), + heartbeat func(), srv proto.ManagementService_SyncServer, cleanupfunc func()) *PeerUpdateHandler { return &PeerUpdateHandler{ @@ -29,6 +31,7 @@ func PeerUpdateHandlerFactory( secretsManager: secretsManager, challenger: challenger, onChallengeStamped: onChallengeStamped, + heartbeat: heartbeat, srv: srv, encrypter: encryption.DefaultEncrypter{}, debouncer: NewUpdateDebouncer(1000 * time.Millisecond), @@ -54,6 +57,7 @@ type PeerUpdateHandler struct { encrypter encryption.Encrypter debouncer Debouncer cleanupFunc func() + heartbeat func() } func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler { @@ -72,6 +76,13 @@ func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error { defer pu.debouncer.Stop() + var heartbeat <-chan time.Time + if pu.heartbeat != nil { + ticker := time.NewTicker(nbpeer.HeartbeatInterval) + defer ticker.Stop() + heartbeat = ticker.C + } + for { select { // condition when there are some updates @@ -110,6 +121,9 @@ func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error { } } + case <-heartbeat: + pu.heartbeat() + // condition when client <-> server connection has been terminated case <-pu.srv.Context().Done(): // happens when connection drops, e.g. client disconnects diff --git a/management/internals/shared/grpc/peer_update_handler_test.go b/management/internals/shared/grpc/peer_update_handler_test.go index 02de49c47..548fc94d6 100644 --- a/management/internals/shared/grpc/peer_update_handler_test.go +++ b/management/internals/shared/grpc/peer_update_handler_test.go @@ -5,14 +5,17 @@ import ( "fmt" "sync" "testing" + "testing/synctest" "time" pb "github.com/golang/protobuf/proto" //nolint - "github.com/netbirdio/netbird/management/internals/controllers/network_map" - "github.com/netbirdio/netbird/shared/management/proto" "github.com/stretchr/testify/assert" "go.uber.org/mock/gomock" "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + + "github.com/netbirdio/netbird/management/internals/controllers/network_map" + nbpeer "github.com/netbirdio/netbird/management/server/peer" + "github.com/netbirdio/netbird/shared/management/proto" ) func TestSendPeerUpdates_FirstUpdate(t *testing.T) { @@ -118,6 +121,61 @@ func TestSendPeerUpdates_ServerContextDone(t *testing.T) { wg.Wait() } +func TestSendPeerUpdates_HeartbeatOnlyOnInterval(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctrl := gomock.NewController(t) + secretsManager := NewMockSecretsManager(ctrl) + updateDebouncer := NewMockDebouncer(ctrl) + syncSender := NewMocksyncSender(ctrl) + + var heartbeats int + pu := PeerUpdateHandler{ + peerKey: mustGenerateKey(t), + updates: make(chan *network_map.UpdateMessage), + secretsManager: secretsManager, + encrypter: testEncrypter{}, + debouncer: updateDebouncer, + srv: syncSender, + cleanupFunc: func() {}, + heartbeat: func() { heartbeats++ }, + } + + msg := network_map.UpdateMessage{ + Update: &proto.SyncResponse{Version: 1}, + } + + srvCtx, cancel := context.WithCancel(context.Background()) + srvKey := mustGenerateKey(t) + updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true) + updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(make(chan time.Time)) + syncSender.EXPECT().Context().AnyTimes().Return(srvCtx) + secretsManager.EXPECT().GetWGKey().Return(srvKey, nil) + syncSender.EXPECT().Send(gomock.Any()) + updateDebouncer.EXPECT().Stop() + + done := make(chan struct{}) + go func() { + defer close(done) + pu.HandleUpdates(context.Background()) //nolint:errcheck + }() + + pu.updates <- &msg + synctest.Wait() + assert.Zero(t, heartbeats, "an update must not write a heartbeat") + + time.Sleep(nbpeer.HeartbeatInterval) + synctest.Wait() + assert.Equal(t, 1, heartbeats, "one heartbeat per interval") + + time.Sleep(nbpeer.HeartbeatInterval) + synctest.Wait() + assert.Equal(t, 2, heartbeats, "one heartbeat per interval") + + cancel() + <-done + }) +} + func mustGenerateKey(t *testing.T) wgtypes.Key { t.Helper() k, err := wgtypes.GenerateKey() diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go index c2aa2fc22..ff466c5a6 100644 --- a/management/internals/shared/grpc/server.go +++ b/management/internals/shared/grpc/server.go @@ -57,6 +57,8 @@ const ( envConcurrentSyncs = "NB_MAX_CONCURRENT_SYNCS" defaultSyncLim = 1000 + + peerHeartbeatTimeout = 5 * time.Second ) // Server an instance of a Management gRPC API server @@ -348,7 +350,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S s.syncSem.Add(-1) return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, s.challenger, - trackChallenge, func() { s.refreshPeerLastSeen(ctx, accountID, peer.ID) }, + trackChallenge, func() { s.heartbeatPeerSession(ctx, accountID, peer.ID, syncStart) }, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart, updates) }). WithMetrics(s.appMetrics).HandleUpdates(ctx) } @@ -442,13 +444,14 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even return nil } -// refreshPeerLastSeen records that the peer of a live sync stream is still connected, -// so the stale peer sweep leaves it alone. A failed refresh keeps the stream open: the -// sweep only acts after several missed heartbeats. -func (s *Server) refreshPeerLastSeen(ctx context.Context, accountID, peerID string) { - staleBefore := time.Now().UTC().Add(-nbpeer.HeartbeatInterval / 2) - if _, err := s.accountManager.GetStore().RefreshPeerLastSeen(ctx, accountID, peerID, staleBefore); err != nil && ctx.Err() == nil { - log.WithContext(ctx).Warnf("refresh last seen of peer %s: %v", peerID, err) +// heartbeatPeerSession records that the peer of a live sync stream is still connected. +// A failed heartbeat keeps the stream open: the peer only reads as disconnected after +// nbpeer.StaleAfter without one, and reads as connected again once a heartbeat lands. +func (s *Server) heartbeatPeerSession(ctx context.Context, accountID, peerID string, streamStartTime time.Time) { + ctx, cancel := context.WithTimeout(ctx, peerHeartbeatTimeout) + defer cancel() + if _, err := s.accountManager.GetStore().HeartbeatPeerSession(ctx, accountID, peerID, streamStartTime.UnixNano()); err != nil && !errors.Is(ctx.Err(), context.Canceled) { + log.WithContext(ctx).Warnf("heartbeat of peer %s: %v", peerID, err) } } diff --git a/management/server/peer/peer.go b/management/server/peer/peer.go index 5ed19f239..4bc126d6c 100644 --- a/management/server/peer/peer.go +++ b/management/server/peer/peer.go @@ -10,6 +10,7 @@ import ( "time" log "github.com/sirupsen/logrus" + "gorm.io/gorm" "github.com/netbirdio/netbird/management/server/util" "github.com/netbirdio/netbird/shared/management/http/api" @@ -101,6 +102,32 @@ type PeerStatus struct { //nolint:revive RequiresApproval bool } +const ( + // HeartbeatInterval is how often the management instance serving a peer's sync + // stream refreshes the peer's LastSeen. + HeartbeatInterval = time.Minute + // StaleAfter is how long a connected peer may go without a LastSeen refresh before + // it reads as disconnected. It spans several heartbeats so one missed refresh does + // not disconnect a live peer. + StaleAfter = 5 * time.Minute +) + +// IsConnected reports whether a peer stored with the given connected flag and LastSeen +// counts as connected. A management instance that stops without marking its peers +// disconnected leaves the flag set, so a peer whose LastSeen is older than StaleAfter +// reads as disconnected. +func IsConnected(connected bool, lastSeen time.Time) bool { + return connected && time.Since(lastSeen) < StaleAfter +} + +// AfterFind applies IsConnected to every peer loaded through gorm. +func (p *Peer) AfterFind(*gorm.DB) error { + if p.Status != nil { + p.Status.Connected = IsConnected(p.Status.Connected, p.Status.LastSeen) + } + return nil +} + // Location is a geo location information of a Peer based on public connection IP type Location struct { ConnectionIP net.IP `gorm:"serializer:json"` // from grpc peer or reverse proxy headers depends on setup diff --git a/management/server/peer/peer_test.go b/management/server/peer/peer_test.go index c5b512069..99c0adc79 100644 --- a/management/server/peer/peer_test.go +++ b/management/server/peer/peer_test.go @@ -4,6 +4,7 @@ import ( "fmt" "net/netip" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -164,3 +165,22 @@ func TestPeerCapabilities(t *testing.T) { }) } } + +func TestIsConnected(t *testing.T) { + tests := []struct { + name string + connected bool + lastSeen time.Time + want bool + }{ + {name: "connected and recently seen", connected: true, lastSeen: time.Now().Add(-time.Minute), want: true}, + {name: "connected but not seen for longer than StaleAfter", connected: true, lastSeen: time.Now().Add(-StaleAfter - time.Minute)}, + {name: "connected but never seen", connected: true}, + {name: "disconnected and recently seen", lastSeen: time.Now()}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, IsConnected(tt.connected, tt.lastSeen)) + }) + } +} diff --git a/management/server/store/sql_store_activity_test.go b/management/server/store/sql_store_activity_test.go index b7aef47b1..411b7c184 100644 --- a/management/server/store/sql_store_activity_test.go +++ b/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", diff --git a/management/server/store/sql_store_get_account_test.go b/management/server/store/sql_store_get_account_test.go index aca8d8a12..a3da90d1d 100644 --- a/management/server/store/sql_store_get_account_test.go +++ b/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, diff --git a/management/server/store/sql_store_peer.go b/management/server/store/sql_store_peer.go index 77642eea7..aab9d8fa1 100644 --- a/management/server/store/sql_store_peer.go +++ b/management/server/store/sql_store_peer.go @@ -182,6 +182,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, @@ -252,7 +276,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 @@ -599,7 +623,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) diff --git a/management/server/store/sql_store_peers_filter_test.go b/management/server/store/sql_store_peers_filter_test.go index dabfcf2ae..b9ef706f9 100644 --- a/management/server/store/sql_store_peers_filter_test.go +++ b/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 diff --git a/management/server/store/store.go b/management/server/store/store.go index 670bd3956..c35e84a6e 100644 --- a/management/server/store/store.go +++ b/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) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) DeletePeer(ctx context.Context, accountID string, peerID string) error diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go index 8b89219e5..bed6708d4 100644 --- a/management/server/store/store_mock.go +++ b/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()