Refactor login with store.SavePeer (#2334)

This pull request refactors the login functionality by integrating store.SavePeer. The changes aim to improve the handling of peer login processes, particularly focusing on synchronization and error handling.

Changes:
- Refactored login logic to use store.SavePeer.
- Added checks for login without lock for login necessary checks from the client and utilized write lock for full login flow.
- Updated error handling with status.NewPeerLoginExpiredError().
- Moved geoIP check logic to a more appropriate place.
- Removed redundant calls and improved documentation.
- Moved the code to smaller methods to improve readability.

management/server/peer.go

@@ -453,6 +453,17 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		Location:               peer.Location,
 	}
 
+	if am.geo != nil && newPeer.Location.ConnectionIP != nil {
+		location, err := am.geo.Lookup(newPeer.Location.ConnectionIP)
+		if err != nil {
+			log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", newPeer.Location.ConnectionIP.String(), err)
+		} else {
+			newPeer.Location.CountryCode = location.Country.ISOCode
+			newPeer.Location.CityName = location.City.Names.En
+			newPeer.Location.GeoNameID = location.City.GeonameID
+		}
+	}
+
 	// add peer to 'All' group
 	group, err := account.GetGroupAll()
 	if err != nil {
@@ -535,7 +546,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	}
 
 	if peerLoginExpired(ctx, peer, account.Settings) {
-		return nil, nil, nil, status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
+		return nil, nil, nil, status.NewPeerLoginExpiredError()
 	}
 
 	peer, updated := updatePeerMeta(peer, sync.Meta, account)
@@ -586,21 +597,10 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 			// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
 			// Try registering it.
 			newPeer := &nbpeer.Peer{
-				Key:    login.WireGuardPubKey,
+				Key:      login.WireGuardPubKey,
-				Meta:   login.Meta,
+				Meta:     login.Meta,
-				SSHKey: login.SSHKey,
+				SSHKey:   login.SSHKey,
-			}
+				Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
-			if am.geo != nil && login.ConnectionIP != nil {
-				location, err := am.geo.Lookup(login.ConnectionIP)
-				if err != nil {
-					log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", login.ConnectionIP.String(), err)
-				} else {
-					newPeer.Location.ConnectionIP = login.ConnectionIP
-					newPeer.Location.CountryCode = location.Country.ISOCode
-					newPeer.Location.CityName = location.City.Names.En
-					newPeer.Location.GeoNameID = location.City.GeonameID
-
-				}
 			}
 
 			return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
@@ -610,44 +610,17 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
 	}
 
-	peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
+	// when the client sends a login request with a JWT which is used to get the user ID,
-	if err != nil {
+	// it means that the client has already checked if it needs login and had been through the SSO flow
-		return nil, nil, nil, status.NewPeerNotRegisteredError()
+	// so, we can skip this check and directly proceed with the login
-	}
+	if login.UserID == "" {
-
+		err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
-	accSettings, err := am.Store.GetAccountSettings(ctx, accountID)
+		if err != nil {
-	if err != nil {
-		return nil, nil, nil, status.Errorf(status.Internal, "failed to get account settings: %s", err)
-	}
-
-	var isWriteLock bool
-
-	// duplicated logic from after the lock to have an early exit
-	expired := peerLoginExpired(ctx, peer, accSettings)
-	switch {
-	case expired:
-		if err := checkAuth(ctx, login.UserID, peer); err != nil {
 			return nil, nil, nil, err
 		}
-		isWriteLock = true
-		log.WithContext(ctx).Debugf("peer login expired, acquiring write lock")
-
-	case peer.UpdateMetaIfNew(login.Meta):
-		isWriteLock = true
-		log.WithContext(ctx).Debugf("peer changed meta, acquiring write lock")
-
-	default:
-		isWriteLock = false
-		log.WithContext(ctx).Debugf("peer meta is the same, acquiring read lock")
 	}
 
-	var unlock func()
+	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
-
-	if isWriteLock {
-		unlock = am.Store.AcquireAccountWriteLock(ctx, accountID)
-	} else {
-		unlock = am.Store.AcquireAccountReadLock(ctx, accountID)
-	}
 	defer func() {
 		if unlock != nil {
 			unlock()
@@ -660,7 +633,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		return nil, nil, nil, err
 	}
 
-	peer, err = account.FindPeerByPubKey(login.WireGuardPubKey)
+	peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
 	if err != nil {
 		return nil, nil, nil, status.NewPeerNotRegisteredError()
 	}
@@ -671,53 +644,39 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 	}
 
 	// this flag prevents unnecessary calls to the persistent store.
-	shouldStoreAccount := false
+	shouldStorePeer := false
 	updateRemotePeers := false
 	if peerLoginExpired(ctx, peer, account.Settings) {
-		err = checkAuth(ctx, login.UserID, peer)
+		err = am.handleExpiredPeer(ctx, login, account, peer)
 		if err != nil {
 			return nil, nil, nil, err
 		}
-		// If peer was expired before and if it reached this point, it is re-authenticated.
-		// UserID is present, meaning that JWT validation passed successfully in the API layer.
-		updatePeerLastLogin(peer, account)
 		updateRemotePeers = true
-		shouldStoreAccount = true
+		shouldStorePeer = true
-
-		// sync user last login with peer last login
-		user, err := account.FindUser(login.UserID)
-		if err != nil {
-			return nil, nil, nil, status.Errorf(status.Internal, "couldn't find user")
-		}
-		user.updateLastLogin(peer.LastLogin)
-
-		am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
 	}
 
 	isRequiresApproval, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
 		return nil, nil, nil, err
 	}
+
 	peer, updated := updatePeerMeta(peer, login.Meta, account)
 	if updated {
-		shouldStoreAccount = true
+		shouldStorePeer = true
 	}
 
-	peer, err = am.checkAndUpdatePeerSSHKey(ctx, peer, account, login.SSHKey)
+	if peer.SSHKey != login.SSHKey {
-	if err != nil {
+		peer.SSHKey = login.SSHKey
-		return nil, nil, nil, err
+		shouldStorePeer = true
 	}
 
-	if shouldStoreAccount {
+	if shouldStorePeer {
-		if !isWriteLock {
+		err = am.Store.SavePeer(ctx, accountID, peer)
-			log.WithContext(ctx).Errorf("account %s should be stored but is not write locked", accountID)
-			return nil, nil, nil, status.Errorf(status.Internal, "account should be stored but is not write locked")
-		}
-		err = am.Store.SaveAccount(ctx, account)
 		if err != nil {
 			return nil, nil, nil, err
 		}
 	}
+
 	unlock()
 	unlock = nil
 
@@ -725,13 +684,46 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		am.updateAccountPeers(ctx, account)
 	}
 
+	return am.getValidatedPeerWithMap(ctx, isRequiresApproval, account, peer)
+}
+
+// checkIFPeerNeedsLoginWithoutLock checks if the peer needs login without acquiring the account lock. The check validate if the peer was not added via SSO
+// and if the peer login is expired.
+// 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 PeerLogin) error {
+	peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
+	if err != nil {
+		return err
+	}
+
+	// if the peer was not added with SSO login we can exit early because peers activated with setup-key
+	// doesn't expire, and we avoid extra databases calls.
+	if !peer.AddedWithSSOLogin() {
+		return nil
+	}
+
+	settings, err := am.Store.GetAccountSettings(ctx, accountID)
+	if err != nil {
+		return err
+	}
+
+	if peerLoginExpired(ctx, peer, settings) {
+		return status.NewPeerLoginExpiredError()
+	}
+
+	return nil
+}
+
+func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, account *Account, peer *nbpeer.Peer) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
 	var postureChecks []*posture.Checks
 
 	if isRequiresApproval {
 		emptyMap := &NetworkMap{
 			Network: account.Network.Copy(),
 		}
-		return peer, emptyMap, postureChecks, nil
+		return peer, emptyMap, nil, nil
 	}
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
@@ -743,6 +735,30 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 	return peer, account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, approvedPeersMap), postureChecks, nil
 }
 
+func (am *DefaultAccountManager) handleExpiredPeer(ctx context.Context, login PeerLogin, account *Account, peer *nbpeer.Peer) error {
+	err := checkAuth(ctx, login.UserID, peer)
+	if err != nil {
+		return err
+	}
+	// If peer was expired before and if it reached this point, it is re-authenticated.
+	// UserID is present, meaning that JWT validation passed successfully in the API layer.
+	updatePeerLastLogin(peer, account)
+
+	// sync user last login with peer last login
+	user, err := account.FindUser(login.UserID)
+	if err != nil {
+		return status.Errorf(status.Internal, "couldn't find user")
+	}
+
+	err = am.Store.SaveUserLastLogin(account.Id, user.Id, peer.LastLogin)
+	if err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
+	return nil
+}
+
 func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
 	if peer.AddedWithSSOLogin() {
 		user, err := account.FindUser(peer.UserID)
@@ -759,11 +775,11 @@ func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
 func checkAuth(ctx context.Context, loginUserID string, peer *nbpeer.Peer) error {
 	if loginUserID == "" {
 		// absence of a user ID indicates that JWT wasn't provided.
-		return status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
+		return status.NewPeerLoginExpiredError()
 	}
 	if peer.UserID != loginUserID {
 		log.WithContext(ctx).Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, loginUserID)
-		return status.Errorf(status.Unauthenticated, "can't login")
+		return status.Errorf(status.Unauthenticated, "can't login with this credentials")
 	}
 	return nil
 }
@@ -783,31 +799,6 @@ func updatePeerLastLogin(peer *nbpeer.Peer, account *Account) {
 	account.UpdatePeer(peer)
 }
 
-func (am *DefaultAccountManager) checkAndUpdatePeerSSHKey(ctx context.Context, peer *nbpeer.Peer, account *Account, newSSHKey string) (*nbpeer.Peer, error) {
-	if len(newSSHKey) == 0 {
-		log.WithContext(ctx).Debugf("no new SSH key provided for peer %s, skipping update", peer.ID)
-		return peer, nil
-	}
-
-	if peer.SSHKey == newSSHKey {
-		log.WithContext(ctx).Debugf("same SSH key provided for peer %s, skipping update", peer.ID)
-		return peer, nil
-	}
-
-	peer.SSHKey = newSSHKey
-	account.UpdatePeer(peer)
-
-	err := am.Store.SaveAccount(ctx, account)
-	if err != nil {
-		return nil, err
-	}
-
-	// trigger network map update
-	am.updateAccountPeers(ctx, account)
-
-	return peer, nil
-}
-
 // UpdatePeerSSHKey updates peer's public SSH key
 func (am *DefaultAccountManager) UpdatePeerSSHKey(ctx context.Context, peerID string, sshKey string) error {
 	if sshKey == "" {

management/server/status/error.go

@@ -95,3 +95,8 @@ func NewUserNotFoundError(userKey string) error {
 func NewPeerNotRegisteredError() error {
 	return Errorf(Unauthenticated, "peer is not registered")
 }
+
+// NewPeerLoginExpiredError creates a new Error with PermissionDenied type for an expired peer
+func NewPeerLoginExpiredError() error {
+	return Errorf(PermissionDenied, "peer login has expired, please log in once more")
+}