Refactor AddPeer to ensure consistency (#557)

management/server/account.go

@@ -15,6 +15,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"math/rand"
+	"net"
 	"net/netip"
 	"reflect"
 	"regexp"
@@ -62,7 +63,7 @@ type AccountManager interface {
 	UpdatePeer(accountID string, peer *Peer) (*Peer, error)
 	GetNetworkMap(peerKey string) (*NetworkMap, error)
 	GetPeerNetwork(peerKey string) (*Network, error)
-	AddPeer(setupKey string, userId string, peer *Peer) (*Peer, error)
+	AddPeer(setupKey, userID string, peer *Peer) (*Peer, error)
 	UpdatePeerMeta(peerKey string, meta PeerSystemMeta) error
 	UpdatePeerSSHKey(peerKey string, sshKey string) error
 	GetUsersFromAccount(accountID, userID string) ([]*UserInfo, error)
@@ -277,7 +278,41 @@ func (a *Account) FindUser(userID string) (*User, error) {
 	return user, nil
 }
 
-// getPeerDNSLabels return the account's peers' dns labels
+// FindSetupKey looks for a given SetupKey in the Account or returns error if it wasn't found.
+func (a *Account) FindSetupKey(setupKey string) (*SetupKey, error) {
+	key := a.SetupKeys[setupKey]
+	if key == nil {
+		return nil, Errorf(SetupKeyNotFound, "setup key not found")
+	}
+
+	return key, nil
+}
+
+func (a *Account) getUserGroups(userID string) ([]string, error) {
+	user, err := a.FindUser(userID)
+	if err != nil {
+		return nil, err
+	}
+	return user.AutoGroups, nil
+}
+
+func (a *Account) getSetupKeyGroups(setupKey string) ([]string, error) {
+	key, err := a.FindSetupKey(setupKey)
+	if err != nil {
+		return nil, err
+	}
+	return key.AutoGroups, nil
+}
+
+func (a *Account) getTakenIPs() []net.IP {
+	var takenIps []net.IP
+	for _, existingPeer := range a.Peers {
+		takenIps = append(takenIps, existingPeer.IP)
+	}
+
+	return takenIps
+}
+
 func (a *Account) getPeerDNSLabels() lookupMap {
 	existingLabels := make(lookupMap)
 	for _, peer := range a.Peers {
@@ -925,15 +960,6 @@ func newAccountWithId(accountId, userId, domain string) *Account {
 	return acc
 }
 
-func getAccountSetupKeyByKey(acc *Account, key string) *SetupKey {
-	for _, k := range acc.SetupKeys {
-		if key == k.Key {
-			return k
-		}
-	}
-	return nil
-}
-
 func removeFromList(inputList []string, toRemove []string) []string {
 	toRemoveMap := make(map[string]struct{})
 	for _, item := range toRemove {

management/server/dns.go

@@ -130,7 +130,7 @@ func getPeerHostLabel(name string, peerLabels lookupMap) (string, error) {
 
 	uniqueLabel := getUniqueHostLabel(label, peerLabels)
 	if uniqueLabel == "" {
-		return "", fmt.Errorf("couldn't find a unique valid lavel for %s, parsed label %s", name, label)
+		return "", fmt.Errorf("couldn't find a unique valid label for %s, parsed label %s", name, label)
 	}
 	return uniqueLabel, nil
 }

management/server/error.go

@@ -6,17 +6,20 @@ import (
 
 const (
 	// UserAlreadyExists indicates that user already exists
-	UserAlreadyExists ErrorType = 1
+	UserAlreadyExists ErrorType = iota
 	// AccountNotFound indicates that specified account hasn't been found
-	AccountNotFound ErrorType = iota
+	AccountNotFound
 	// PreconditionFailed indicates that some pre-condition for the operation hasn't been fulfilled
-	PreconditionFailed ErrorType = iota
+	PreconditionFailed
 
 	// UserNotFound indicates that user wasn't found in the system (or under a given Account)
-	UserNotFound ErrorType = iota
+	UserNotFound
 
 	// PermissionDenied indicates that user has no permissions to view data
-	PermissionDenied ErrorType = iota
+	PermissionDenied
+
+	// SetupKeyNotFound indicates that the setup key wasn't found in the system (or under a given Account)
+	SetupKeyNotFound
 )
 
 // ErrorType is a type of the Error

management/server/grpcserver.go

@@ -185,7 +185,7 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 func (s *GRPCServer) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest) (*Peer, error) {
 	var (
 		reqSetupKey string
-		userId      string
+		userID      string
 	)
 
 	if req.GetJwtToken() != "" {
@@ -204,12 +204,11 @@ func (s *GRPCServer) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest)
 		if err != nil {
 			return nil, status.Errorf(codes.Internal, "unable to fetch account with claims, err: %v", err)
 		}
-		userId = claims.UserId
+		userID = claims.UserId
 	} else {
 		log.Debugln("using setup key to register peer")
-
 		reqSetupKey = req.GetSetupKey()
-		userId = ""
+		userID = ""
 	}
 
 	meta := req.GetMeta()
@@ -222,7 +221,7 @@ func (s *GRPCServer) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest)
 		sshKey = req.GetPeerKeys().GetSshPubKey()
 	}
 
-	peer, err := s.accountManager.AddPeer(reqSetupKey, userId, &Peer{
+	peer, err := s.accountManager.AddPeer(reqSetupKey, userID, &Peer{
 		Key:    peerKey.String(),
 		Name:   meta.GetHostname(),
 		SSHKey: string(sshKey),
@@ -238,13 +237,18 @@ func (s *GRPCServer) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest)
 		},
 	})
 	if err != nil {
-		s, ok := status.FromError(err)
+		if e, ok := FromError(err); ok {
-		if ok {
+			switch e.Type() {
-			if s.Code() == codes.FailedPrecondition || s.Code() == codes.OutOfRange {
+			case PreconditionFailed:
-				return nil, err
+				return nil, status.Errorf(codes.FailedPrecondition, e.message)
+			case AccountNotFound:
+			case SetupKeyNotFound:
+			case UserNotFound:
+				return nil, status.Errorf(codes.NotFound, e.message)
+			default:
 			}
 		}
-		return nil, status.Errorf(codes.NotFound, "provided setup key doesn't exists")
+		return nil, status.Errorf(codes.Internal, "failed registering new peer")
 	}
 
 	// todo move to DefaultAccountManager the code below

management/server/network.go

@@ -5,8 +5,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
 	"github.com/rs/xid"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 	"math/rand"
 	"net"
 	"sync"
@@ -95,7 +93,7 @@ func AllocatePeerIP(ipNet net.IPNet, takenIps []net.IP) (net.IP, error) {
 	ips, _ := generateIPs(&ipNet, takenIPMap)
 
 	if len(ips) == 0 {
-		return nil, status.Errorf(codes.OutOfRange, "failed allocating new IP for the ipNet %s - network is out of IPs", ipNet.String())
+		return nil, Errorf(PreconditionFailed, "failed allocating new IP for the ipNet %s - network is out of IPs", ipNet.String())
 	}
 
 	// pick a random IP

management/server/peer.go

@@ -319,59 +319,20 @@ func (am *DefaultAccountManager) GetPeerNetwork(peerPubKey string) (*Network, er
 // to it. We also add the User ID to the peer metadata to identify registrant.
 // 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 string, userID string, peer *Peer) (*Peer, error) {
+func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*Peer, error) {
 
 	upperKey := strings.ToUpper(setupKey)
-
 	var account *Account
 	var err error
-	var sk *SetupKey
+	addedByUser := false
-	// auto-assign groups that are coming with a SetupKey or a User
+	if len(userID) > 0 {
-	var groupsToAdd []string
+		addedByUser = true
-	if len(upperKey) != 0 {
-		account, err = am.Store.GetAccountBySetupKey(upperKey)
-		if err != nil {
-			return nil, status.Errorf(
-				codes.NotFound,
-				"unable to register peer, unable to find account with setupKey %s",
-				upperKey,
-			)
-		}
-
-		sk = getAccountSetupKeyByKey(account, upperKey)
-		if sk == nil {
-			// shouldn't happen actually
-			return nil, status.Errorf(
-				codes.NotFound,
-				"unable to register peer, unknown setupKey %s",
-				upperKey,
-			)
-		}
-
-		if !sk.IsValid() {
-			return nil, status.Errorf(
-				codes.FailedPrecondition,
-				"unable to register peer, its setup key is invalid (expired, overused or revoked)",
-			)
-		}
-
-		groupsToAdd = sk.AutoGroups
-
-	} else if len(userID) != 0 {
 		account, err = am.Store.GetAccountByUser(userID)
-		if err != nil {
-			return nil, status.Errorf(codes.NotFound, "unable to register peer, unknown user with ID: %s", userID)
-		}
-		user, ok := account.Users[userID]
-		if !ok {
-			return nil, status.Errorf(codes.NotFound, "unable to register peer, unknown user with ID: %s", userID)
-		}
-
-		groupsToAdd = user.AutoGroups
-
 	} else {
-		// Empty setup key and jwt fail
+		account, err = am.Store.GetAccountBySetupKey(setupKey)
-		return nil, status.Errorf(codes.InvalidArgument, "no setup key or user id provided")
+	}
+	if err != nil {
+		return nil, Errorf(AccountNotFound, "failed adding new peer: account not found")
 	}
 
 	unlock := am.Store.AcquireAccountLock(account.Id)
@@ -383,22 +344,29 @@ func (am *DefaultAccountManager) AddPeer(setupKey string, userID string, peer *P
 		return nil, err
 	}
 
-	var takenIps []net.IP
+	if !addedByUser {
-	existingLabels := make(lookupMap)
+		// validate the setup key if adding with a key
-	for _, existingPeer := range account.Peers {
+		sk, err := account.FindSetupKey(upperKey)
-		takenIps = append(takenIps, existingPeer.IP)
+		if err != nil {
-		if existingPeer.DNSLabel != "" {
+			return nil, err
-			existingLabels[existingPeer.DNSLabel] = struct{}{}
 		}
+
+		if !sk.IsValid() {
+			return nil, Errorf(PreconditionFailed, "couldn't add peer: setup key is invalid")
 		}
 
+		account.SetupKeys[sk.Key] = sk.IncrementUsage()
+	}
+
+	takenIps := account.getTakenIPs()
+	existingLabels := account.getPeerDNSLabels()
+
 	newLabel, err := getPeerHostLabel(peer.Name, existingLabels)
 	if err != nil {
 		return nil, err
 	}
 
 	peer.DNSLabel = newLabel
-
 	network := account.Network
 	nextIp, err := AllocatePeerIP(network.Net, takenIps)
 	if err != nil {
@@ -425,6 +393,19 @@ func (am *DefaultAccountManager) AddPeer(setupKey string, userID string, peer *P
 	}
 	group.Peers = append(group.Peers, newPeer.Key)
 
+	var groupsToAdd []string
+	if addedByUser {
+		groupsToAdd, err = account.getUserGroups(userID)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		groupsToAdd, err = account.getSetupKeyGroups(upperKey)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if len(groupsToAdd) > 0 {
 		for _, s := range groupsToAdd {
 			if g, ok := account.Groups[s]; ok && g.Name != "All" {
@@ -434,11 +415,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey string, userID string, peer *P
 	}
 
 	account.Peers[newPeer.Key] = newPeer
-	if len(upperKey) != 0 {
-		account.SetupKeys[sk.Key] = sk.IncrementUsage()
-	}
 	account.Network.IncSerial()
-
 	err = am.Store.SaveAccount(account)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed adding peer")