[management] expire and disconnect peers while including offline peers (#7467)

This commit is contained in:
Pascal Fischer
2026-09-09 13:44:19 +02:00
committed by GitHub
parent d101f6cc46
commit 269cbadfeb
7 changed files with 350 additions and 32 deletions
+5 -2
View File
@@ -719,8 +719,10 @@ func (am *DefaultAccountManager) schedulePeerLoginExpiration(ctx context.Context
log.WithContext(ctx).Tracef("peer login expiration job for account %s is already scheduled", accountID)
return
}
// The job outlives the request that arms it, so it must not inherit the request's cancellation.
jobCtx := context.WithoutCancel(ctx)
if nextRun, ok := am.getNextPeerExpiration(ctx, accountID); ok {
go am.peerLoginExpiry.Schedule(ctx, nextRun, accountID, am.peerLoginExpirationJob(ctx, accountID))
go am.peerLoginExpiry.Schedule(jobCtx, nextRun, accountID, am.peerLoginExpirationJob(jobCtx, accountID))
}
}
@@ -752,8 +754,9 @@ func (am *DefaultAccountManager) peerInactivityExpirationJob(ctx context.Context
// checkAndSchedulePeerInactivityExpiration periodically checks for inactive peers to end their sessions
func (am *DefaultAccountManager) checkAndSchedulePeerInactivityExpiration(ctx context.Context, accountID string) {
am.peerInactivityExpiry.Cancel(ctx, []string{accountID})
jobCtx := context.WithoutCancel(ctx)
if nextRun, ok := am.getNextInactivePeerExpiration(ctx, accountID); ok {
go am.peerInactivityExpiry.Schedule(ctx, nextRun, accountID, am.peerInactivityExpirationJob(ctx, accountID))
go am.peerInactivityExpiry.Schedule(jobCtx, nextRun, accountID, am.peerInactivityExpirationJob(jobCtx, accountID))
}
}
+176 -3
View File
@@ -1920,6 +1920,154 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
}
}
func TestDefaultAccountManager_SchedulePeerLoginExpiration_IncludesOfflinePeers(t *testing.T) {
manager, updateManager, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to create an account")
connectedKey, offlineKey := addExpiringPeers(t, manager)
_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &types.Settings{
PeerLoginExpiration: time.Hour,
PeerLoginExpirationEnabled: true,
Extra: &types.ExtraSettings{},
})
require.NoError(t, err, "expecting to update account settings successfully but got error")
manager.peerLoginExpiry.CancelAll(context.Background())
// The connected peer logged in just now, so a job computed from connected peers alone
// would be armed for an hour. The offline peer's login expires in two seconds; a
// reconnect of that peer must not have to wait for the connected peer's tick.
now := time.Now().UTC()
setPeerLogin(t, manager, accountID, connectedKey, true, now)
setPeerLogin(t, manager, accountID, offlineKey, false, now.Add(-time.Hour+2*time.Second))
offlinePeer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, offlineKey)
require.NoError(t, err)
updateManager.CreateChannel(context.Background(), offlinePeer.ID)
manager.peerLoginExpiry = NewDefaultScheduler()
t.Cleanup(func() { manager.peerLoginExpiry.CancelAll(context.Background()) })
manager.schedulePeerLoginExpiration(context.Background(), accountID)
// The flag is committed per peer before the disconnect fans out, so wait for both.
require.Eventually(t, func() bool {
peer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, offlineKey)
return err == nil && peer.Status.LoginExpired && !updateManager.HasChannel(offlinePeer.ID)
}, 10*time.Second, 100*time.Millisecond, "offline peer should be expired and disconnected at its own deadline")
connectedPeer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, connectedKey)
require.NoError(t, err)
assert.False(t, connectedPeer.Status.LoginExpired, "connected peer with a fresh login must not expire")
}
func TestDefaultAccountManager_SchedulePeerLoginExpiration_DetachesRequestContext(t *testing.T) {
manager, _, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to create an account")
connectedKey, _ := addExpiringPeers(t, manager)
setPeerLogin(t, manager, accountID, connectedKey, true, time.Now().UTC())
scheduled := make(chan context.Context, 1)
manager.peerLoginExpiry = &MockScheduler{
IsSchedulerRunningFunc: func(string) bool { return false },
ScheduleFunc: func(ctx context.Context, _ time.Duration, _ string, _ func() (time.Duration, bool)) {
scheduled <- ctx
},
}
requestCtx, cancel := context.WithCancel(context.Background())
manager.schedulePeerLoginExpiration(requestCtx, accountID)
cancel()
select {
case jobCtx := <-scheduled:
assert.NoError(t, jobCtx.Err(), "the expiration job must outlive the request that armed it")
case <-time.After(time.Second):
t.Fatal("timeout while waiting for the job to be scheduled")
}
}
func TestDefaultAccountManager_ExpireAndUpdatePeers_SkipsPeerThatLoggedInAgain(t *testing.T) {
manager, updateManager, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to create an account")
reloggedKey, staleKey := addExpiringPeers(t, manager)
_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &types.Settings{
PeerLoginExpiration: time.Hour,
PeerLoginExpirationEnabled: true,
Extra: &types.ExtraSettings{},
})
require.NoError(t, err, "expecting to update account settings successfully but got error")
manager.peerLoginExpiry.CancelAll(context.Background())
expiredLogin := time.Now().UTC().Add(-2 * time.Hour)
setPeerLogin(t, manager, accountID, reloggedKey, true, expiredLogin)
setPeerLogin(t, manager, accountID, staleKey, true, expiredLogin)
expiredPeers, err := manager.getExpiredPeers(context.Background(), accountID)
require.NoError(t, err)
require.Len(t, expiredPeers, 2, "both peers should be due for expiration")
// The job holds the candidate list while one peer completes a fresh login, which
// moves its deadline into the future and must win over the stale candidate entry.
setPeerLogin(t, manager, accountID, reloggedKey, true, time.Now().UTC())
reloggedPeer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, reloggedKey)
require.NoError(t, err)
stalePeer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, staleKey)
require.NoError(t, err)
updateManager.CreateChannel(context.Background(), reloggedPeer.ID)
updateManager.CreateChannel(context.Background(), stalePeer.ID)
err = manager.expireAndUpdatePeers(context.Background(), accountID, expiredPeers, peerExpirationSessionExpired)
require.NoError(t, err)
reloggedPeer, err = manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, reloggedKey)
require.NoError(t, err)
assert.False(t, reloggedPeer.Status.LoginExpired, "a peer that logged in again must not be flagged from the stale candidate list")
assert.True(t, reloggedPeer.Status.Connected, "the re-logged peer must keep its connected status")
assert.True(t, updateManager.HasChannel(reloggedPeer.ID), "the re-logged peer's update channel must stay open")
stalePeer, err = manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, staleKey)
require.NoError(t, err)
assert.True(t, stalePeer.Status.LoginExpired, "a peer that is still due must be flagged")
assert.False(t, updateManager.HasChannel(stalePeer.ID), "the expired peer's update channel must be closed")
}
// addExpiringPeers registers two SSO peers with login expiration enabled and returns their public keys.
func addExpiringPeers(t *testing.T, manager *DefaultAccountManager) (string, string) {
t.Helper()
keys := make([]string, 0, 2)
for _, hostname := range []string{"connected-peer", "offline-peer"} {
key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
_, _, _, _, err = manager.AddPeer(context.Background(), "", "", userID, &nbpeer.Peer{
Key: key.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: hostname},
LoginExpirationEnabled: true,
}, false)
require.NoError(t, err, "unable to add peer")
keys = append(keys, key.PublicKey().String())
}
return keys[0], keys[1]
}
func setPeerLogin(t *testing.T, manager *DefaultAccountManager, accountID, peerKey string, connected bool, lastLogin time.Time) {
t.Helper()
peer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, peerKey)
require.NoError(t, err)
peer.Status.Connected = connected
peer.LastLogin = &lastLogin
require.NoError(t, manager.Store.SavePeer(context.Background(), accountID, peer))
}
func TestDefaultAccountManager_MarkPeerDisconnected_SchedulesInactivityExpiration(t *testing.T) {
manager, _, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
@@ -2702,7 +2850,7 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
expectedNextExpiration: time.Duration(0),
},
{
name: "No connected peers, no expiration",
name: "Offline peer with expiration, return expiration",
peers: map[string]*nbpeer.Peer{
"peer-1": {
Status: &nbpeer.PeerStatus{
@@ -2721,8 +2869,33 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
},
expiration: time.Second,
expirationEnabled: false,
expectedNextRun: false,
expectedNextExpiration: time.Duration(0),
expectedNextRun: true,
expectedNextExpiration: time.Second,
},
{
name: "Offline peer with the earliest deadline defines the next run",
peers: map[string]*nbpeer.Peer{
"peer-1": {
Status: &nbpeer.PeerStatus{
Connected: true,
},
LoginExpirationEnabled: true,
LastLogin: util.ToPtr(time.Now().UTC()),
UserID: userID,
},
"peer-2": {
Status: &nbpeer.PeerStatus{
Connected: false,
},
LoginExpirationEnabled: true,
LastLogin: util.ToPtr(time.Now().UTC().Add(-50 * time.Minute)),
UserID: userID,
},
},
expiration: time.Hour,
expirationEnabled: true,
expectedNextRun: true,
expectedNextExpiration: 10 * time.Minute,
},
{
name: "Connected peers with disabled expiration, no expiration",
+10 -6
View File
@@ -1494,9 +1494,12 @@ func checkAuth(ctx context.Context, loginUserID string, peer *nbpeer.Peer) error
func peerLoginExpired(ctx context.Context, peer *nbpeer.Peer, settings *types.Settings) bool {
expired, expiresIn := peer.LoginExpired(settings.PeerLoginExpiration)
expired = settings.PeerLoginExpirationEnabled && expired
if expired || peer.Status.LoginExpired {
log.WithContext(ctx).Debugf("peer's %s login expired %v ago", peer.ID, expiresIn)
if settings.PeerLoginExpirationEnabled && expired {
log.WithContext(ctx).Debugf("peer's %s login expired %v ago", peer.ID, -expiresIn)
return true
}
if peer.Status.LoginExpired {
log.WithContext(ctx).Debugf("peer's %s login is marked as expired", peer.ID)
return true
}
return false
@@ -1643,7 +1646,9 @@ func (am *DefaultAccountManager) UpdateAccountPeer(ctx context.Context, accountI
// getNextPeerExpiration returns the minimum duration in which the next peer of the account will expire if it was found.
// If there is no peer that expires this function returns false and a duration of 0.
// This function only considers peers that haven't been expired yet and that are connected.
// This function only considers peers that haven't been expired yet. Offline peers count too:
// a running job is never re-armed on connect, so a peer that reconnects with an old login
// must already be part of the scheduled run.
func (am *DefaultAccountManager) getNextPeerExpiration(ctx context.Context, accountID string) (time.Duration, bool) {
peersWithExpiry, err := am.Store.GetAccountPeersWithExpiration(ctx, store.LockingStrengthNone, accountID)
if err != nil {
@@ -1663,8 +1668,7 @@ func (am *DefaultAccountManager) getNextPeerExpiration(ctx context.Context, acco
var nextExpiry *time.Duration
for _, peer := range peersWithExpiry {
// consider only connected peers because others will require login on connecting to the management server
if peer.Status.LoginExpired || !peer.Status.Connected {
if peer.Status.LoginExpired {
continue
}
_, duration := peer.LoginExpired(settings.PeerLoginExpiration)
+7 -2
View File
@@ -117,6 +117,7 @@ func (wm *DefaultScheduler) Schedule(ctx context.Context, in time.Duration, ID s
}
ticker := time.NewTicker(in)
period := in
wm.jobs[ID] = cancel
log.WithContext(ctx).Debugf("scheduled a job %s to run in %s. There are %d total jobs scheduled.", ID, in.String(), len(wm.jobs))
@@ -136,14 +137,18 @@ func (wm *DefaultScheduler) Schedule(ctx context.Context, in time.Duration, ID s
if !reschedule {
wm.mu.Lock()
defer wm.mu.Unlock()
delete(wm.jobs, ID)
// A Cancel during job() may have registered a replacement under this ID.
if current, ok := wm.jobs[ID]; ok && current == cancel {
delete(wm.jobs, ID)
}
log.WithContext(ctx).Debugf("job %s is not scheduled to run again", ID)
ticker.Stop()
return
}
// we need this comparison to avoid resetting the ticker with the same duration and missing the current elapsesed time
if runIn != in {
if runIn != period {
ticker.Reset(runIn)
period = runIn
}
case <-cancel:
log.WithContext(ctx).Debugf("job %s was canceled, stopping timer", ID)
+89
View File
@@ -6,10 +6,12 @@ import (
"math/rand"
"runtime"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestScheduler_Performance(t *testing.T) {
@@ -150,3 +152,90 @@ func TestScheduler_Schedule(t *testing.T) {
scheduler.cancel(context.Background(), jobID)
}
func TestScheduler_Schedule_ResetsTickerAfterReturningInitialInterval(t *testing.T) {
jobID := "test-scheduler-job-2"
scheduler := NewDefaultScheduler()
defer scheduler.Cancel(context.Background(), []string{jobID})
initial := 30 * time.Millisecond
stretched := 400 * time.Millisecond
runs := make(chan time.Time, 3)
count := 0
// The first run stretches the period; the second returns the initial interval again,
// which must shrink the period back instead of keeping the stretched one.
job := func() (nextRunIn time.Duration, reschedule bool) {
count++
runs <- time.Now()
switch count {
case 1:
return stretched, true
case 2:
return initial, true
default:
return 0, false
}
}
scheduler.Schedule(context.Background(), initial, jobID, job)
var stamps []time.Time
for len(stamps) < 3 {
select {
case ts := <-runs:
stamps = append(stamps, ts)
case <-time.After(2 * time.Second):
t.Fatalf("timed out after %d runs", len(stamps))
}
}
assert.Less(t, stamps[2].Sub(stamps[1]), stretched/2, "returning the initial interval must reset the stretched ticker")
}
func TestScheduler_Schedule_StaleCompletionKeepsReplacement(t *testing.T) {
jobID := "test-scheduler-job-3"
scheduler := NewDefaultScheduler()
defer scheduler.Cancel(context.Background(), []string{jobID})
started := make(chan struct{})
release := make(chan struct{})
staleJob := func() (nextRunIn time.Duration, reschedule bool) {
close(started)
<-release
return 0, false
}
scheduler.Schedule(context.Background(), 10*time.Millisecond, jobID, staleJob)
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("timed out waiting for the first job to start")
}
// Cancel the job while it is still executing and register a replacement under the
// same ID, as the expiration paths do on a settings change.
scheduler.Cancel(context.Background(), []string{jobID})
var replacementRuns atomic.Int32
scheduler.Schedule(context.Background(), 20*time.Millisecond, jobID, func() (nextRunIn time.Duration, reschedule bool) {
replacementRuns.Add(1)
return 20 * time.Millisecond, true
})
require.True(t, scheduler.IsSchedulerRunning(jobID), "replacement must be registered")
// The stale job now completes without rescheduling; its cleanup must leave the
// replacement's entry in place.
close(release)
assert.Never(t, func() bool { return !scheduler.IsSchedulerRunning(jobID) }, 200*time.Millisecond, 10*time.Millisecond,
"stale completion must not drop the replacement job")
var duplicateRuns atomic.Int32
scheduler.Schedule(context.Background(), 10*time.Millisecond, jobID, func() (nextRunIn time.Duration, reschedule bool) {
duplicateRuns.Add(1)
return 10 * time.Millisecond, true
})
assert.Never(t, func() bool { return duplicateRuns.Load() > 0 }, 100*time.Millisecond, 10*time.Millisecond,
"a duplicate schedule must be refused while the replacement is registered")
scheduler.Cancel(context.Background(), []string{jobID})
assert.False(t, scheduler.IsSchedulerRunning(jobID), "cancel must find and remove the replacement")
runsAfterCancel := replacementRuns.Load()
assert.Never(t, func() bool { return replacementRuns.Load() > runsAfterCancel+1 }, 150*time.Millisecond, 10*time.Millisecond,
"the replacement must stop after cancel")
}
+2 -3
View File
@@ -404,7 +404,7 @@ func (a *Account) GetExpiredPeers() []*nbpeer.Peer {
// GetNextPeerExpiration returns the minimum duration in which the next peer of the account will expire if it was found.
// If there is no peer that expires this function returns false and a duration of 0.
// This function only considers peers that haven't been expired yet and that are connected.
// This function only considers peers that haven't been expired yet, whether connected or not.
func (a *Account) GetNextPeerExpiration() (time.Duration, bool) {
peersWithExpiry := a.GetPeersWithExpiration()
if len(peersWithExpiry) == 0 {
@@ -412,8 +412,7 @@ func (a *Account) GetNextPeerExpiration() (time.Duration, bool) {
}
var nextExpiry *time.Duration
for _, peer := range peersWithExpiry {
// consider only connected peers because others will require login on connecting to the management server
if peer.Status.LoginExpired || !peer.Status.Connected {
if peer.Status.LoginExpired {
continue
}
_, duration := peer.LoginExpired(a.Settings.PeerLoginExpiration)
+61 -16
View File
@@ -1177,28 +1177,35 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
dnsDomain := am.networkMapController.GetDNSDomain(settings)
var peerIDs []string
for _, peer := range peers {
defer func() {
if len(peerIDs) == 0 {
return
}
// this will trigger peer disconnect from the management service
log.Debugf("Expiring %d peers for account %s", len(peerIDs), accountID)
am.networkMapController.DisconnectPeers(ctx, accountID, peerIDs)
}()
for _, candidate := range peers {
// nolint:staticcheck
ctx = context.WithValue(ctx, nbcontext.PeerIDKey, peer.Key)
peerCtx := context.WithValue(ctx, nbcontext.PeerIDKey, candidate.Key)
if peer.UserID == "" {
if candidate.UserID == "" {
// we do not want to expire peers that are added via setup key
continue
}
if peer.Status.LoginExpired {
peer, err := am.expirePeerIfStillDue(peerCtx, accountID, candidate.ID, settings, reason)
if err != nil {
return err
}
if peer == nil {
continue
}
peerIDs = append(peerIDs, peer.ID)
peer.MarkLoginExpired(true)
if err := am.Store.SavePeerStatus(ctx, accountID, peer.ID, *peer.Status); err != nil {
return err
}
meta := peer.EventMeta(dnsDomain)
meta["reason"] = string(reason)
am.StoreEvent(
ctx,
peerCtx,
peer.UserID, peer.ID, accountID,
activity.PeerLoginExpired, meta,
)
@@ -1215,15 +1222,53 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
if err != nil {
return fmt.Errorf("notify network map controller of peer update: %w", err)
}
if len(peerIDs) != 0 {
// this will trigger peer disconnect from the management service
log.Debugf("Expiring %d peers for account %s", len(peerIDs), accountID)
am.networkMapController.DisconnectPeers(ctx, accountID, peerIDs)
}
return nil
}
// expirePeerIfStillDue flags the peer as login-expired and returns its fresh copy, or nil
// when it no longer qualifies. The candidate list is read without a lock, so a login that
// landed in between would otherwise be overwritten with a stale expired status.
func (am *DefaultAccountManager) expirePeerIfStillDue(ctx context.Context, accountID, peerID string, settings *types.Settings, reason peerExpirationReason) (*nbpeer.Peer, error) {
var expired *nbpeer.Peer
err := am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
peer, err := transaction.GetPeerByID(ctx, store.LockingStrengthUpdate, accountID, peerID)
if err != nil {
if s, ok := status.FromError(err); ok && s.Type() == status.NotFound {
return nil
}
return err
}
if peer.Status.LoginExpired || !peerExpirationDue(peer, settings, reason) {
return nil
}
peer.MarkLoginExpired(true)
if err := transaction.SavePeerStatus(ctx, accountID, peer.ID, *peer.Status); err != nil {
return err
}
expired = peer
return nil
})
if err != nil {
return nil, err
}
return expired, nil
}
// peerExpirationDue re-evaluates a time-based expiry against the peer's current state.
// Administrative reasons expire the peer unconditionally.
func peerExpirationDue(peer *nbpeer.Peer, settings *types.Settings, reason peerExpirationReason) bool {
switch reason {
case peerExpirationSessionExpired:
expired, _ := peer.LoginExpired(settings.PeerLoginExpiration)
return settings.PeerLoginExpirationEnabled && expired
case peerExpirationInactivity:
expired, _ := peer.SessionExpired(settings.PeerInactivityExpiration)
return settings.PeerInactivityExpirationEnabled && expired
default:
return true
}
}
func (am *DefaultAccountManager) deleteUserFromIDP(ctx context.Context, targetUserID, accountID string) error {
if am.userDeleteFromIDPEnabled {
log.WithContext(ctx).Debugf("user %s deleted from IdP", targetUserID)