Move Login business logic from gRPC API to Accountmanager (#713)

The Management gRPC API has too much business logic happening while it has to be in the Account manager. This also needs to make more requests to the store through the account manager.
2026-04-21 17:56:39 +00:00 · 2023-03-03 18:35:38 +01:00
parent f64e0754ee
commit e914adb5cd
10 changed files with 385 additions and 451 deletions
--- a/management/server/peer.go
+++ b/management/server/peer.go
@@ -2,7 +2,6 @@ package server

 import (
 	"fmt"
-	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/rs/xid"
@@ -36,6 +35,26 @@ type PeerStatus struct {
 	LoginExpired bool
 }

+// PeerSync used as a data object between the gRPC API and AccountManager on Sync request.
+type PeerSync struct {
+	// WireGuardPubKey is a peers WireGuard public key
+	WireGuardPubKey string
+}
+
+// PeerLogin used as a data object between the gRPC API and AccountManager on Login request.
+type PeerLogin struct {
+	// WireGuardPubKey is a peers WireGuard public key
+	WireGuardPubKey string
+	// SSHKey is a peer's ssh key. Can be empty (e.g., old version do not provide it, or this feature is disabled)
+	SSHKey string
+	// Meta is the system information passed by peer, must be always present.
+	Meta PeerSystemMeta
+	// UserID indicates that JWT was used to log in, and it was valid. Can be empty when SetupKey is used or auth is not required.
+	UserID string
+	// SetupKey references to a server.SetupKey to log in. Can be empty when UserID is used or auth is not required.
+	SetupKey string
+}
+
 // Peer represents a machine connected to the network.
 // The Peer is a WireGuard peer identified by a public key
 type Peer struct {
@@ -93,6 +112,15 @@ func (p *Peer) Copy() *Peer {
 	}
 }

+// UpdateMeta updates peer's system meta data
+func (p *Peer) UpdateMeta(meta PeerSystemMeta) {
+	// Avoid overwriting UIVersion if the update was triggered sole by the CLI client
+	if meta.UIVersion == "" {
+		meta.UIVersion = p.Meta.UIVersion
+	}
+	p.Meta = meta
+}
+
 // MarkLoginExpired marks peer's status expired or not
 func (p *Peer) MarkLoginExpired(expired bool) {
 	newStatus := p.Status.Copy()
@@ -194,6 +222,18 @@ func (am *DefaultAccountManager) GetPeers(accountID, userID string) ([]*Peer, er
 	return peers, nil
 }

+func (am *DefaultAccountManager) markPeerLoginExpired(peer *Peer, account *Account, expired bool) (*Peer, error) {
+	peer.MarkLoginExpired(expired)
+	account.UpdatePeer(peer)
+
+	err := am.Store.SavePeerStatus(account.Id, peer.ID, *peer.Status)
+	if err != nil {
+		return nil, err
+	}
+
+	return peer, nil
+}
+
 // MarkPeerLoginExpired when peer login has expired
 func (am *DefaultAccountManager) MarkPeerLoginExpired(peerPubKey string, loginExpired bool) error {
 	account, err := am.Store.GetAccountByPeerPubKey(peerPubKey)
@@ -423,6 +463,10 @@ func (am *DefaultAccountManager) GetPeerByIP(accountID string, peerIP string) (*
 	return nil, status.Errorf(status.NotFound, "peer with IP %s not found", peerIP)
 }

+func (am *DefaultAccountManager) getNetworkMap(peer *Peer, account *Account) *NetworkMap {
+	return account.GetPeerNetworkMap(peer.ID, am.dnsDomain)
+}
+
 // GetNetworkMap returns Network map for a given peer (omits original peer from the Peers result)
 func (am *DefaultAccountManager) GetNetworkMap(peerID string) (*NetworkMap, error) {

@@ -436,40 +480,7 @@ func (am *DefaultAccountManager) GetNetworkMap(peerID string) (*NetworkMap, erro
 		return nil, status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
 	}

-	aclPeers := account.getPeersByACL(peerID)
-	// exclude expired peers
-	var peersToConnect []*Peer
-	for _, p := range aclPeers {
-		expired, _ := peer.LoginExpired(account.Settings.PeerLoginExpiration)
-		if expired {
-			continue
-		}
-		peersToConnect = append(peersToConnect, p)
-	}
-	// Please mind, that the returned route.Route objects will contain Peer.Key instead of Peer.ID.
-	routesUpdate := account.getRoutesToSync(peerID, peersToConnect)
-
-	dnsManagementStatus := account.getPeerDNSManagementStatus(peerID)
-	dnsUpdate := nbdns.Config{
-		ServiceEnable: dnsManagementStatus,
-	}
-
-	if dnsManagementStatus {
-		var zones []nbdns.CustomZone
-		peersCustomZone := getPeersCustomZone(account, am.dnsDomain)
-		if peersCustomZone.Domain != "" {
-			zones = append(zones, peersCustomZone)
-		}
-		dnsUpdate.CustomZones = zones
-		dnsUpdate.NameServerGroups = getPeerNSGroups(account, peerID)
-	}
-
-	return &NetworkMap{
-		Peers:     peersToConnect,
-		Network:   account.Network.Copy(),
-		Routes:    routesUpdate,
-		DNSConfig: dnsUpdate,
-	}, err
+	return am.getNetworkMap(peer, account), nil
 }

 // GetPeerNetwork returns the Network for a given peer
@@ -484,13 +495,17 @@ func (am *DefaultAccountManager) GetPeerNetwork(peerID string) (*Network, error)
 }

 // AddPeer adds a new peer to the Store.
-// Each Account has a list of pre-authorised SetupKey and if no Account has a given key err with a code codes.Unauthenticated
-// will be returned, meaning the key is invalid
+// Each Account has a list of pre-authorized SetupKey and if no Account has a given key err with a code status.PermissionDenied
+// will be returned, meaning the setup key is invalid or not found.
 // If a User ID is provided, it means that we passed the authentication using JWT, then we look for account by User ID and register the peer
-// to it. We also add the User ID to the peer metadata to identify registrant.
+// to it. We also add the User ID to the peer metadata to identify registrant. If no userID provided, then fail with status.PermissionDenied
 // Each new Peer will be assigned a new next net.IP from the Account.Network and Account.Network.LastIP will be updated (IP's are not reused).
 // The peer property is just a placeholder for the Peer properties to pass further
-func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*Peer, error) {
+func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*Peer, *NetworkMap, error) {
+	if setupKey == "" && userID == "" {
+		// no auth method provided => reject access
+		return nil, nil, status.Errorf(status.Unauthenticated, "no peer auth method provided, please use a setup key or interactive SSO login")
+	}

 	upperKey := strings.ToUpper(setupKey)
 	var account *Account
@@ -503,7 +518,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 		account, err = am.Store.GetAccountBySetupKey(setupKey)
 	}
 	if err != nil {
-		return nil, status.Errorf(status.NotFound, "failed adding new peer: account not found")
+		return nil, nil, status.Errorf(status.NotFound, "failed adding new peer: account not found")
 	}

 	unlock := am.Store.AcquireAccountLock(account.Id)
@@ -512,7 +527,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	// ensure that we consider modification happened meanwhile (because we were outside the account lock when we fetched the account)
 	account, err = am.Store.GetAccount(account.Id)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}

 	opEvent := &activity.Event{
@@ -524,11 +539,11 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 		// validate the setup key if adding with a key
 		sk, err := account.FindSetupKey(upperKey)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}

 		if !sk.IsValid() {
-			return nil, status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key is invalid")
+			return nil, nil, status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key is invalid")
 		}

 		account.SetupKeys[sk.Key] = sk.IncrementUsage()
@@ -542,16 +557,16 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	takenIps := account.getTakenIPs()
 	existingLabels := account.getPeerDNSLabels()

-	newLabel, err := getPeerHostLabel(peer.Name, existingLabels)
+	newLabel, err := getPeerHostLabel(peer.Meta.Hostname, existingLabels)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}

 	peer.DNSLabel = newLabel
 	network := account.Network
 	nextIp, err := AllocatePeerIP(network.Net, takenIps)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}

 	newPeer := &Peer{
@@ -560,7 +575,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 		SetupKey:               upperKey,
 		IP:                     nextIp,
 		Meta:                   peer.Meta,
-		Name:                   peer.Name,
+		Name:                   peer.Meta.Hostname,
 		DNSLabel:               newLabel,
 		UserID:                 userID,
 		Status:                 &PeerStatus{Connected: false, LastSeen: time.Now()},
@@ -573,7 +588,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	// add peer to 'All' group
 	group, err := account.GetGroupAll()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	group.Peers = append(group.Peers, newPeer.ID)

@@ -581,12 +596,12 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	if addedByUser {
 		groupsToAdd, err = account.getUserGroups(userID)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 	} else {
 		groupsToAdd, err = account.getSetupKeyGroups(upperKey)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 	}

@@ -602,50 +617,173 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	account.Network.IncSerial()
 	err = am.Store.SaveAccount(account)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}

 	opEvent.TargetID = newPeer.ID
 	opEvent.Meta = newPeer.EventMeta(am.GetDNSDomain())
 	am.storeEvent(opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)

-	return newPeer, nil
-}
-
-// UpdatePeerLastLogin sets Peer.LastLogin to the current timestamp.
-func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
-	account, err := am.Store.GetAccountByPeerID(peerID)
+	err = am.updateAccountPeers(account)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}

+	networkMap := account.GetPeerNetworkMap(peer.ID, am.dnsDomain)
+	return newPeer, networkMap, nil
+}
+
+func (am *DefaultAccountManager) checkPeerLoginExpiration(loginUserID string, peer *Peer, account *Account) error {
+	if peer.AddedWithSSOLogin() {
+		expired, expiresIn := peer.LoginExpired(account.Settings.PeerLoginExpiration)
+		expired = account.Settings.PeerLoginExpirationEnabled && expired
+		if expired || peer.Status.LoginExpired {
+			log.Debugf("peer %s login expired", peer.ID)
+			if loginUserID == "" {
+				// absence of a user ID indicates that JWT wasn't provided.
+				_, err := am.markPeerLoginExpired(peer, account, true)
+				if err != nil {
+					return err
+				}
+				return status.Errorf(status.PermissionDenied,
+					"peer login has expired %v ago. Please log in once more", expiresIn)
+			}
+			// user ID is there meaning that JWT validation passed successfully in the API layer.
+			if peer.UserID != loginUserID {
+				log.Warnf("user mismatch when loggin in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, loginUserID)
+				return status.Errorf(status.Unauthenticated, "can't login")
+			}
+			_ = am.updatePeerLastLogin(peer, account)
+		}
+	}
+
+	return nil
+}
+
+// SyncPeer checks whether peer is eligible for receiving NetworkMap (authenticated) and returns its NetworkMap if eligible
+func (am *DefaultAccountManager) SyncPeer(sync PeerSync) (*Peer, *NetworkMap, error) {
+	account, err := am.Store.GetAccountByPeerPubKey(sync.WireGuardPubKey)
+	if err != nil {
+		if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
+			return nil, nil, status.Errorf(status.Unauthenticated, "peer is not registered")
+		}
+		return nil, nil, err
+	}
+
+	// we found the peer, and we follow a normal login flow
 	unlock := am.Store.AcquireAccountLock(account.Id)
 	defer unlock()

-	// ensure that we consider modification happened meanwhile (because we were outside the account lock when we fetched the account)
+	// fetch the account from the store once more after acquiring lock to avoid concurrent updates inconsistencies
 	account, err = am.Store.GetAccount(account.Id)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}

-	peer := account.GetPeer(peerID)
-	if peer == nil {
-		return status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
+	peer, err := account.FindPeerByPubKey(sync.WireGuardPubKey)
+	if err != nil {
+		return nil, nil, status.Errorf(status.Unauthenticated, "peer is not registered")
 	}

+	err = am.checkPeerLoginExpiration("", peer, account)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return peer, am.getNetworkMap(peer, account), nil
+
+}
+
+// LoginPeer logs in or registers a peer.
+// If peer doesn't exist the function checks whether a setup key or a user is present and registers a new peer if so.
+func (am *DefaultAccountManager) LoginPeer(login PeerLogin) (*Peer, *NetworkMap, error) {
+
+	account, err := am.Store.GetAccountByPeerPubKey(login.WireGuardPubKey)
+	if err != nil {
+		if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
+			// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
+			// Try registering it.
+			return am.AddPeer(login.SetupKey, login.UserID, &Peer{
+				Key:    login.WireGuardPubKey,
+				Meta:   login.Meta,
+				SSHKey: login.SSHKey,
+			})
+		}
+		log.Errorf("failed while logging in peer %s: %v", login.WireGuardPubKey, err)
+		return nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
+	}
+
+	// we found the peer, and we follow a normal login flow
+	unlock := am.Store.AcquireAccountLock(account.Id)
+	defer unlock()
+
+	// fetch the account from the store once more after acquiring lock to avoid concurrent updates inconsistencies
+	account, err = am.Store.GetAccount(account.Id)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
+	if err != nil {
+		return nil, nil, status.Errorf(status.Unauthenticated, "peer is not registered")
+	}
+
+	err = am.checkPeerLoginExpiration(login.UserID, peer, account)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	peer = updatePeerMeta(peer, login.Meta, account)
+
+	peer, err = am.checkAndUpdatePeerSSHKey(peer, account, login.SSHKey)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.Store.SaveAccount(account)
+	if err != nil {
+		return nil, nil, err
+	}
+	networkMap := account.GetPeerNetworkMap(peer.ID, am.dnsDomain)
+	return peer, networkMap, nil
+
+}
+
+func (am *DefaultAccountManager) updatePeerLastLogin(peer *Peer, account *Account) *Peer {
 	peer.LastLogin = time.Now()
 	newStatus := peer.Status.Copy()
 	newStatus.LoginExpired = false
 	peer.Status = newStatus
-
 	account.UpdatePeer(peer)
+	return peer
+}

-	err = am.Store.SaveAccount(account)
-	if err != nil {
-		return err
+func (am *DefaultAccountManager) checkAndUpdatePeerSSHKey(peer *Peer, account *Account, newSSHKey string) (*Peer, error) {
+	if len(newSSHKey) == 0 {
+		log.Debugf("no new SSH key provided for peer %s, skipping update", peer.ID)
+		return peer, nil
 	}

-	return nil
+	if peer.SSHKey == newSSHKey {
+		log.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(account)
+	if err != nil {
+		return nil, err
+	}
+
+	// trigger network map update
+	err = am.updateAccountPeers(account)
+	if err != nil {
+		return nil, err
+	}
+
+	return peer, nil
 }

 // UpdatePeerSSHKey updates peer's public SSH key
@@ -737,35 +875,10 @@ func (am *DefaultAccountManager) GetPeer(accountID, peerID, userID string) (*Pee
 	return nil, status.Errorf(status.Internal, "user %s has no access to peer %s under account %s", userID, peerID, accountID)
 }

-// UpdatePeerMeta updates peer's system metadata
-func (am *DefaultAccountManager) UpdatePeerMeta(peerID string, meta PeerSystemMeta) error {
-
-	account, err := am.Store.GetAccountByPeerID(peerID)
-	if err != nil {
-		return err
-	}
-
-	unlock := am.Store.AcquireAccountLock(account.Id)
-	defer unlock()
-
-	peer := account.GetPeer(peerID)
-	if peer == nil {
-		return status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
-	}
-
-	// Avoid overwriting UIVersion if the update was triggered sole by the CLI client
-	if meta.UIVersion == "" {
-		meta.UIVersion = peer.Meta.UIVersion
-	}
-
-	peer.Meta = meta
+func updatePeerMeta(peer *Peer, meta PeerSystemMeta, account *Account) *Peer {
+	peer.UpdateMeta(meta)
 	account.UpdatePeer(peer)
-
-	err = am.Store.SaveAccount(account)
-	if err != nil {
-		return err
-	}
-	return nil
+	return peer
 }

 // getPeersByACL returns all peers that given peer has access to.