From e6292e3124f85fffce6042547ffaea0d3e45d7c0 Mon Sep 17 00:00:00 2001
From: Misha Bragin <bangvalo@gmail.com>
Date: Thu, 23 Mar 2023 17:47:53 +0100
Subject: [PATCH] Disable peer expiration of peers added with setup keys (#758)

---
 management/server/account.go      |  4 ++--
 management/server/account_test.go | 37 +++++++++++++++++++++++++++++++
 management/server/peer.go         |  2 +-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index 1d4c1072131..01cae2e6430 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -358,11 +358,11 @@ func (a *Account) GetNextPeerExpiration() (time.Duration, bool) {
 	return *nextExpiry, true
 }
 
-// GetPeersWithExpiration returns a list of peers that have Peer.LoginExpirationEnabled set to true
+// GetPeersWithExpiration returns a list of peers that have Peer.LoginExpirationEnabled set to true and that were added by a user
 func (a *Account) GetPeersWithExpiration() []*Peer {
 	peers := make([]*Peer, 0)
 	for _, peer := range a.Peers {
-		if peer.LoginExpirationEnabled {
+		if peer.LoginExpirationEnabled && peer.AddedWithSSOLogin() {
 			peers = append(peers, peer)
 		}
 	}
diff --git a/management/server/account_test.go b/management/server/account_test.go
index 5b4b1cc174c..af894817b50 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -1605,9 +1605,11 @@ func TestAccount_GetPeersWithExpiration(t *testing.T) {
 			peers: map[string]*Peer{
 				"peer-1": {
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 				"peer-2": {
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 			},
 			expectedPeers: map[string]struct{}{},
@@ -1618,9 +1620,11 @@ func TestAccount_GetPeersWithExpiration(t *testing.T) {
 				"peer-1": {
 					ID:                     "peer-1",
 					LoginExpirationEnabled: true,
+					UserID:                 userID,
 				},
 				"peer-2": {
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 			},
 			expectedPeers: map[string]struct{}{
@@ -1680,12 +1684,14 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 						Connected: false,
 					},
 					LoginExpirationEnabled: true,
+					UserID:                 userID,
 				},
 				"peer-2": {
 					Status: &PeerStatus{
 						Connected: true,
 					},
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 			},
 			expiration:             time.Second,
@@ -1701,12 +1707,14 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 						Connected: true,
 					},
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 				"peer-2": {
 					Status: &PeerStatus{
 						Connected: true,
 					},
 					LoginExpirationEnabled: false,
+					UserID:                 userID,
 				},
 			},
 			expiration:             time.Second,
@@ -1723,6 +1731,7 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 						LoginExpired: true,
 					},
 					LoginExpirationEnabled: true,
+					UserID:                 userID,
 				},
 				"peer-2": {
 					Status: &PeerStatus{
@@ -1730,6 +1739,7 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 						LoginExpired: true,
 					},
 					LoginExpirationEnabled: true,
+					UserID:                 userID,
 				},
 			},
 			expiration:             time.Second,
@@ -1747,6 +1757,7 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 					},
 					LoginExpirationEnabled: true,
 					LastLogin:              time.Now(),
+					UserID:                 userID,
 				},
 				"peer-2": {
 					Status: &PeerStatus{
@@ -1754,6 +1765,7 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 						LoginExpired: true,
 					},
 					LoginExpirationEnabled: true,
+					UserID:                 userID,
 				},
 			},
 			expiration:             time.Minute,
@@ -1761,6 +1773,31 @@ func TestAccount_GetNextPeerExpiration(t *testing.T) {
 			expectedNextRun:        true,
 			expectedNextExpiration: expectedNextExpiration,
 		},
+		{
+			name: "Peers added with setup keys, no expiration",
+			peers: map[string]*Peer{
+				"peer-1": {
+					Status: &PeerStatus{
+						Connected:    true,
+						LoginExpired: false,
+					},
+					LoginExpirationEnabled: true,
+					SetupKey:               "key",
+				},
+				"peer-2": {
+					Status: &PeerStatus{
+						Connected:    true,
+						LoginExpired: false,
+					},
+					LoginExpirationEnabled: true,
+					SetupKey:               "key",
+				},
+			},
+			expiration:             time.Second,
+			expirationEnabled:      false,
+			expectedNextRun:        false,
+			expectedNextExpiration: time.Duration(0),
+		},
 	}
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
diff --git a/management/server/peer.go b/management/server/peer.go
index b5505f91243..7b5ca539ff9 100644
--- a/management/server/peer.go
+++ b/management/server/peer.go
@@ -528,7 +528,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 		SSHEnabled:             false,
 		SSHKey:                 peer.SSHKey,
 		LastLogin:              time.Now(),
-		LoginExpirationEnabled: true,
+		LoginExpirationEnabled: addedByUser,
 	}
 
 	// add peer to 'All' group