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.

management/server/account.go

@@ -49,21 +49,18 @@ type AccountManager interface {
 	SaveUser(accountID, userID string, update *User) (*UserInfo, error)
 	GetSetupKey(accountID, userID, keyID string) (*SetupKey, error)
 	GetAccountByUserOrAccountID(userID, accountID, domain string) (*Account, error)
-	GetAccountByPeerID(peerID string) (*Account, error)
 	GetAccountFromToken(claims jwtclaims.AuthorizationClaims) (*Account, *User, error)
 	IsUserAdmin(claims jwtclaims.AuthorizationClaims) (bool, error)
 	AccountExists(accountId string) (*bool, error)
 	GetPeerByKey(peerKey string) (*Peer, error)
 	GetPeers(accountID, userID string) ([]*Peer, error)
 	MarkPeerConnected(peerKey string, connected bool) error
-	MarkPeerLoginExpired(peerPubKey string, loginExpired bool) error
 	DeletePeer(accountID, peerID, userID string) (*Peer, error)
 	GetPeerByIP(accountId string, peerIP string) (*Peer, error)
 	UpdatePeer(accountID, userID string, peer *Peer) (*Peer, error)
 	GetNetworkMap(peerID string) (*NetworkMap, error)
 	GetPeerNetwork(peerID string) (*Network, error)
-	AddPeer(setupKey, userID string, peer *Peer) (*Peer, error)
+	AddPeer(setupKey, userID string, peer *Peer) (*Peer, *NetworkMap, error)
-	UpdatePeerMeta(peerID string, meta PeerSystemMeta) error
 	UpdatePeerSSHKey(peerID string, sshKey string) error
 	GetUsersFromAccount(accountID, userID string) ([]*UserInfo, error)
 	GetGroup(accountId, groupID string) (*Group, error)
@@ -96,8 +93,9 @@ type AccountManager interface {
 	GetDNSSettings(accountID string, userID string) (*DNSSettings, error)
 	SaveDNSSettings(accountID string, userID string, dnsSettingsToSave *DNSSettings) error
 	GetPeer(accountID, peerID, userID string) (*Peer, error)
-	UpdatePeerLastLogin(peerID string) error
 	UpdateAccountSettings(accountID, userID string, newSettings *Settings) (*Account, error)
+	LoginPeer(login PeerLogin) (*Peer, *NetworkMap, error) //used by peer gRPC API
+	SyncPeer(sync PeerSync) (*Peer, *NetworkMap, error)    //used by peer gRPC API
 }
 
 type DefaultAccountManager struct {
@@ -308,6 +306,44 @@ func (a *Account) GetGroup(groupID string) *Group {
 	return a.Groups[groupID]
 }
 
+// GetPeerNetworkMap returns a group by ID if exists, nil otherwise
+func (a *Account) GetPeerNetworkMap(peerID, dnsDomain string) *NetworkMap {
+	aclPeers := a.getPeersByACL(peerID)
+	// exclude expired peers
+	var peersToConnect []*Peer
+	for _, p := range aclPeers {
+		expired, _ := p.LoginExpired(a.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 := a.getRoutesToSync(peerID, peersToConnect)
+
+	dnsManagementStatus := a.getPeerDNSManagementStatus(peerID)
+	dnsUpdate := nbdns.Config{
+		ServiceEnable: dnsManagementStatus,
+	}
+
+	if dnsManagementStatus {
+		var zones []nbdns.CustomZone
+		peersCustomZone := getPeersCustomZone(a, dnsDomain)
+		if peersCustomZone.Domain != "" {
+			zones = append(zones, peersCustomZone)
+		}
+		dnsUpdate.CustomZones = zones
+		dnsUpdate.NameServerGroups = getPeerNSGroups(a, peerID)
+	}
+
+	return &NetworkMap{
+		Peers:     peersToConnect,
+		Network:   a.Network.Copy(),
+		Routes:    routesUpdate,
+		DNSConfig: dnsUpdate,
+	}
+}
+
 // GetExpiredPeers returns peers that have been expired
 func (a *Account) GetExpiredPeers() []*Peer {
 	var peers []*Peer
@@ -803,11 +839,6 @@ func (am *DefaultAccountManager) warmupIDPCache() error {
 	return nil
 }
 
-// GetAccountByPeerID returns account from the store by a provided peer ID
-func (am *DefaultAccountManager) GetAccountByPeerID(peerID string) (*Account, error) {
-	return am.Store.GetAccountByPeerID(peerID)
-}
-
 // GetAccountByUserOrAccountID looks for an account by user or accountID, if no account is provided and
 // userID doesn't have an account associated with it, one account is created
 func (am *DefaultAccountManager) GetAccountByUserOrAccountID(userID, accountID, domain string) (*Account, error) {

management/server/account_test.go

@@ -38,7 +38,7 @@ func verifyCanAddPeerToAccount(t *testing.T, manager AccountManager, account *Ac
 		setupKey = key.Key
 	}
 
-	_, err := manager.AddPeer(setupKey, userID, peer)
+	_, _, err := manager.AddPeer(setupKey, userID, peer)
 	if err != nil {
 		t.Error("expected to add new peer successfully after creating new account, but failed", err)
 	}
@@ -542,10 +542,9 @@ func TestAccountManager_AddPeer(t *testing.T) {
 	expectedPeerKey := key.PublicKey().String()
 	expectedSetupKey := setupKey.Key
 
-	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  expectedPeerKey,
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: expectedPeerKey},
-		Name: expectedPeerKey,
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v", err)
@@ -611,10 +610,9 @@ func TestAccountManager_AddPeerWithUserID(t *testing.T) {
 	expectedPeerKey := key.PublicKey().String()
 	expectedUserID := userID
 
-	peer, err := manager.AddPeer("", userID, &Peer{
+	peer, _, err := manager.AddPeer("", userID, &Peer{
 		Key:  expectedPeerKey,
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: expectedPeerKey},
-		Name: expectedPeerKey,
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v, account users: %v", err, account.CreatedBy)
@@ -694,10 +692,9 @@ func TestAccountManager_NetworkUpdates(t *testing.T) {
 		}
 		expectedPeerKey := key.PublicKey().String()
 
-		peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
+		peer, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 			Key:  expectedPeerKey,
-			Meta: PeerSystemMeta{},
+			Meta: PeerSystemMeta{Hostname: expectedPeerKey},
-			Name: expectedPeerKey,
 		})
 		if err != nil {
 			t.Fatalf("expecting peer1 to be added, got failure %v", err)
@@ -864,10 +861,9 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 
 	peerKey := key.PublicKey().String()
 
-	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey,
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: peerKey},
-		Name: peerKey,
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v", err)
@@ -951,75 +947,6 @@ func TestGetUsersFromAccount(t *testing.T) {
 	}
 }
 
-func TestAccountManager_UpdatePeerMeta(t *testing.T) {
-	manager, err := createManager(t)
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
-
-	account, err := createAccount(manager, "test_account", "account_creator", "")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	var setupKey *SetupKey
-	for _, key := range account.SetupKeys {
-		setupKey = key
-	}
-
-	key, err := wgtypes.GeneratePrivateKey()
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
-
-	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
-		Key: key.PublicKey().String(),
-		Meta: PeerSystemMeta{
-			Hostname:  "Hostname",
-			GoOS:      "GoOS",
-			Kernel:    "Kernel",
-			Core:      "Core",
-			Platform:  "Platform",
-			OS:        "OS",
-			WtVersion: "WtVersion",
-		},
-		Name: key.PublicKey().String(),
-	})
-	if err != nil {
-		t.Errorf("expecting peer to be added, got failure %v", err)
-		return
-	}
-
-	newMeta := PeerSystemMeta{
-		Hostname:  "new-Hostname",
-		GoOS:      "new-GoOS",
-		Kernel:    "new-Kernel",
-		Core:      "new-Core",
-		Platform:  "new-Platform",
-		OS:        "new-OS",
-		WtVersion: "new-WtVersion",
-	}
-	err = manager.UpdatePeerMeta(peer.ID, newMeta)
-	if err != nil {
-		t.Error(err)
-		return
-	}
-
-	p, err := manager.GetPeerByKey(peer.Key)
-	if err != nil {
-		return
-	}
-
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
-
-	assert.Equal(t, newMeta, p.Meta)
-}
-
 func TestAccount_GetPeerRules(t *testing.T) {
 
 	groups := map[string]*Group{
@@ -1302,10 +1229,9 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
 
 	key, err := wgtypes.GenerateKey()
 	require.NoError(t, err, "unable to generate WireGuard key")
-	peer, err := manager.AddPeer("", userID, &Peer{
+	peer, _, err := manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
-		Name:                   "test-peer",
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")
@@ -1351,10 +1277,9 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
 
 	key, err := wgtypes.GenerateKey()
 	require.NoError(t, err, "unable to generate WireGuard key")
-	_, err = manager.AddPeer("", userID, &Peer{
+	_, _, err = manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
-		Name:                   "test-peer",
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")
@@ -1393,10 +1318,9 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 
 	key, err := wgtypes.GenerateKey()
 	require.NoError(t, err, "unable to generate WireGuard key")
-	_, err = manager.AddPeer("", userID, &Peer{
+	_, _, err = manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
-		Name:                   "test-peer",
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")

management/server/dns_test.go

@@ -244,11 +244,11 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 		return nil, err
 	}
 
-	_, err = am.AddPeer("", dnsAdminUserID, peer1)
+	_, _, err = am.AddPeer("", dnsAdminUserID, peer1)
 	if err != nil {
 		return nil, err
 	}
-	_, err = am.AddPeer("", dnsAdminUserID, peer2)
+	_, _, err = am.AddPeer("", dnsAdminUserID, peer2)
 	if err != nil {
 		return nil, err
 	}

management/server/grpcserver.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"fmt"
+	pb "github.com/golang/protobuf/proto" //nolint
 	"strings"
 	"time"
 
@@ -118,44 +119,18 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 		log.Debugf("Sync request from peer [%s] [%s]", req.WgPubKey, p.Addr.String())
 	}
 
-	peerKey, err := wgtypes.ParseKey(req.GetWgPubKey())
-	if err != nil {
-		log.Warnf("error while parsing peer's Wireguard public key %s on Sync request.", peerKey.String())
-		return status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", peerKey.String())
-	}
-
-	peer, err := s.accountManager.GetPeerByKey(peerKey.String())
-	if err != nil {
-		p, _ := gRPCPeer.FromContext(srv.Context())
-		msg := status.Errorf(codes.PermissionDenied, "provided peer with the key wgPubKey %s is not registered, remote addr is %s", peerKey.String(), p.Addr.String())
-		log.Debug(msg)
-		return msg
-	}
-
-	account, err := s.accountManager.GetAccountByPeerID(peer.ID)
-	if err != nil {
-		return status.Error(codes.Internal, "internal server error")
-	}
-	expired, left := peer.LoginExpired(account.Settings.PeerLoginExpiration)
-	expired = account.Settings.PeerLoginExpirationEnabled && expired
-	if peer.UserID != "" && (expired || peer.Status.LoginExpired) {
-		err = s.accountManager.MarkPeerLoginExpired(peerKey.String(), true)
-		if err != nil {
-			log.Warnf("failed marking peer login expired %s %v", peerKey, err)
-		}
-		return status.Errorf(codes.PermissionDenied, "peer login has expired %v ago. Please log in once more", left)
-	}
-
 	syncReq := &proto.SyncRequest{}
-	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, syncReq)
+	peerKey, err := s.parseRequest(req, syncReq)
 	if err != nil {
-		p, _ := gRPCPeer.FromContext(srv.Context())
+		return err
-		msg := status.Errorf(codes.InvalidArgument, "invalid request message from %s,remote addr is %s", peerKey.String(), p.Addr.String())
-		log.Debug(msg)
-		return msg
 	}
 
-	err = s.sendInitialSync(peerKey, peer, srv)
+	peer, netMap, err := s.accountManager.SyncPeer(PeerSync{WireGuardPubKey: peerKey.String()})
+	if err != nil {
+		return mapError(err)
+	}
+
+	err = s.sendInitialSync(peerKey, peer, netMap, srv)
 	if err != nil {
 		log.Debugf("error while sending initial sync for %s: %v", peerKey.String(), err)
 		return err
@@ -218,7 +193,7 @@ func (s *GRPCServer) validateToken(jwtToken string) (string, error) {
 
 	token, err := s.jwtMiddleware.ValidateAndParse(jwtToken)
 	if err != nil {
-		return "", status.Errorf(codes.Internal, "invalid jwt token, err: %v", err)
+		return "", status.Errorf(codes.InvalidArgument, "invalid jwt token, err: %v", err)
 	}
 	claims := s.jwtClaimsExtractor.FromToken(token)
 	// we need to call this method because if user is new, we will automatically add it to existing or create a new account
@@ -230,84 +205,52 @@ func (s *GRPCServer) validateToken(jwtToken string) (string, error) {
 	return claims.UserId, nil
 }
 
-func (s *GRPCServer) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest) (*Peer, error) {
+// maps internal internalStatus.Error to gRPC status.Error
-	var (
+func mapError(err error) error {
-		reqSetupKey string
-		userID      string
-		err         error
-	)
-
-	if req.GetJwtToken() != "" {
-		log.Debugln("using jwt token to register peer")
-		userID, err = s.validateToken(req.JwtToken)
-		if err != nil {
-			return nil, err
-		}
-	} else {
-		log.Debugln("using setup key to register peer")
-		reqSetupKey = req.GetSetupKey()
-		userID = ""
-	}
-
-	meta := req.GetMeta()
-	if meta == nil {
-		return nil, status.Errorf(codes.InvalidArgument, "peer meta data was not provided")
-	}
-
-	var sshKey []byte
-	if req.GetPeerKeys() != nil {
-		sshKey = req.GetPeerKeys().GetSshPubKey()
-	}
-
-	peer, err := s.accountManager.AddPeer(reqSetupKey, userID, &Peer{
-		Key:    peerKey.String(),
-		Name:   meta.GetHostname(),
-		SSHKey: string(sshKey),
-		Meta: PeerSystemMeta{
-			Hostname:  meta.GetHostname(),
-			GoOS:      meta.GetGoOS(),
-			Kernel:    meta.GetKernel(),
-			Core:      meta.GetCore(),
-			Platform:  meta.GetPlatform(),
-			OS:        meta.GetOS(),
-			WtVersion: meta.GetWiretrusteeVersion(),
-			UIVersion: meta.GetUiVersion(),
-		},
-	})
-	if err != nil {
 	if e, ok := internalStatus.FromError(err); ok {
 		switch e.Type() {
+		case internalStatus.PermissionDenied:
+			return status.Errorf(codes.PermissionDenied, e.Message)
+		case internalStatus.Unauthorized:
+			return status.Errorf(codes.PermissionDenied, e.Message)
+		case internalStatus.Unauthenticated:
+			return status.Errorf(codes.PermissionDenied, e.Message)
 		case internalStatus.PreconditionFailed:
-				return nil, status.Errorf(codes.FailedPrecondition, e.Message)
+			return status.Errorf(codes.FailedPrecondition, e.Message)
 		case internalStatus.NotFound:
-				return nil, status.Errorf(codes.NotFound, e.Message)
+			return status.Errorf(codes.NotFound, e.Message)
 		default:
 		}
 	}
-		return nil, status.Errorf(codes.Internal, "failed registering new peer")
+	return status.Errorf(codes.Internal, "failed handling request")
+}
+
+func extractPeerMeta(loginReq *proto.LoginRequest) PeerSystemMeta {
+	return PeerSystemMeta{
+		Hostname:  loginReq.GetMeta().GetHostname(),
+		GoOS:      loginReq.GetMeta().GetGoOS(),
+		Kernel:    loginReq.GetMeta().GetKernel(),
+		Core:      loginReq.GetMeta().GetCore(),
+		Platform:  loginReq.GetMeta().GetPlatform(),
+		OS:        loginReq.GetMeta().GetOS(),
+		WtVersion: loginReq.GetMeta().GetWiretrusteeVersion(),
+		UIVersion: loginReq.GetMeta().GetUiVersion(),
+	}
+}
+
+func (s *GRPCServer) parseRequest(req *proto.EncryptedMessage, parsed pb.Message) (wgtypes.Key, error) {
+	peerKey, err := wgtypes.ParseKey(req.GetWgPubKey())
+	if err != nil {
+		log.Warnf("error while parsing peer's WireGuard public key %s.", req.WgPubKey)
+		return wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", req.WgPubKey)
 	}
 
-	// todo move to DefaultAccountManager the code below
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, parsed)
-	networkMap, err := s.accountManager.GetNetworkMap(peer.ID)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "unable to fetch network map after registering peer, error: %v", err)
+		return wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "invalid request message")
-	}
-	// notify other peers of our registration
-	for _, remotePeer := range networkMap.Peers {
-		remotePeerNetworkMap, err := s.accountManager.GetNetworkMap(remotePeer.ID)
-		if err != nil {
-			return nil, status.Errorf(codes.Internal, "unable to fetch network map after registering peer, error: %v", err)
 	}
 
-		update := toSyncResponse(s.config, remotePeer, nil, remotePeerNetworkMap, s.accountManager.GetDNSDomain())
+	return peerKey, nil
-		err = s.peersUpdateManager.SendUpdate(remotePeer.ID, &UpdateMessage{Update: update})
-		if err != nil {
-			// todo rethink if we should keep this return
-			return nil, status.Errorf(codes.Internal, "unable to send update after registering peer, error: %v", err)
-		}
-	}
-
-	return peer, nil
 }
 
 // Login endpoint first checks whether peer is registered under any account
@@ -323,113 +266,55 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 		log.Debugf("Login request from peer [%s] [%s]", req.WgPubKey, p.Addr.String())
 	}
 
-	peerKey, err := wgtypes.ParseKey(req.GetWgPubKey())
-	if err != nil {
-		log.Warnf("error while parsing peer's Wireguard public key %s on Sync request.", req.WgPubKey)
-		return nil, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", req.WgPubKey)
-	}
-
 	loginReq := &proto.LoginRequest{}
-	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, loginReq)
+	peerKey, err := s.parseRequest(req, loginReq)
 	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "invalid request message")
+		return nil, err
 	}
 
-	peer, err := s.accountManager.GetPeerByKey(peerKey.String())
+	if loginReq.GetMeta() == nil {
-	if err != nil {
+		msg := status.Errorf(codes.FailedPrecondition,
-		if errStatus, ok := internalStatus.FromError(err); ok && errStatus.Type() == internalStatus.NotFound {
+			"peer system meta has to be provided to log in. Peer %s, remote addr %s", peerKey.String(),
-			// peer doesn't exist -> check if setup key was provided
+			p.Addr.String())
-			if loginReq.GetJwtToken() == "" && loginReq.GetSetupKey() == "" {
+		log.Warn(msg)
-				// absent setup key or jwt -> permission denied
-				p, _ := gRPCPeer.FromContext(ctx)
-				msg := status.Errorf(codes.PermissionDenied,
-					"provided peer with the key wgPubKey %s is not registered and no setup key or jwt was provided,"+
-						" remote addr is %s", peerKey.String(), p.Addr.String())
-				log.Debug(msg)
 		return nil, msg
 	}
 
-			// setup key or jwt is present -> try normal registration flow
+	userID := ""
-			peer, err = s.registerPeer(peerKey, loginReq)
+	// JWT token is not always provided, it is fine for userID to be empty cuz it might be that peer is already registered,
+	// or it uses a setup key to register.
+	if loginReq.GetJwtToken() != "" {
+		userID, err = s.validateToken(loginReq.GetJwtToken())
 		if err != nil {
-				return nil, err
+			log.Warnf("failed validating JWT token sent from peer %s", peerKey)
-			}
+			return nil, mapError(err)
-
-		} else {
-			return nil, status.Error(codes.Internal, "internal server error")
-		}
-	} else if loginReq.GetMeta() != nil {
-		// update peer's system meta data on Login
-		err = s.accountManager.UpdatePeerMeta(peer.ID, PeerSystemMeta{
-			Hostname:  loginReq.GetMeta().GetHostname(),
-			GoOS:      loginReq.GetMeta().GetGoOS(),
-			Kernel:    loginReq.GetMeta().GetKernel(),
-			Core:      loginReq.GetMeta().GetCore(),
-			Platform:  loginReq.GetMeta().GetPlatform(),
-			OS:        loginReq.GetMeta().GetOS(),
-			WtVersion: loginReq.GetMeta().GetWiretrusteeVersion(),
-			UIVersion: loginReq.GetMeta().GetUiVersion(),
-		},
-		)
-		if err != nil {
-			log.Errorf("failed updating peer system meta data %s", peerKey.String())
-			return nil, status.Error(codes.Internal, "internal server error")
 		}
 	}
-
-	// check if peer login has expired
-	account, err := s.accountManager.GetAccountByPeerID(peer.ID)
-	if err != nil {
-		return nil, status.Error(codes.Internal, "internal server error")
-	}
-
-	expired, left := peer.LoginExpired(account.Settings.PeerLoginExpiration)
-	expired = account.Settings.PeerLoginExpirationEnabled && expired
-	if peer.UserID != "" && (expired || peer.Status.LoginExpired) {
-		// it might be that peer expired but user has logged in already, check token then
-		if loginReq.GetJwtToken() == "" {
-			err = s.accountManager.MarkPeerLoginExpired(peerKey.String(), true)
-			if err != nil {
-				log.Warnf("failed marking peer login expired %s %v", peerKey, err)
-			}
-			return nil, status.Errorf(codes.PermissionDenied,
-				"peer login has expired %v ago. Please log in once more", left)
-		}
-		_, err = s.validateToken(loginReq.GetJwtToken())
-		if err != nil {
-			return nil, err
-		}
-
-		err = s.accountManager.UpdatePeerLastLogin(peer.ID)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	var sshKey []byte
 	if loginReq.GetPeerKeys() != nil {
 		sshKey = loginReq.GetPeerKeys().GetSshPubKey()
 	}
 
-	if len(sshKey) > 0 {
+	peer, netMap, err := s.accountManager.LoginPeer(PeerLogin{
-		err = s.accountManager.UpdatePeerSSHKey(peer.ID, string(sshKey))
+		WireGuardPubKey: peerKey.String(),
+		SSHKey:          string(sshKey),
+		Meta:            extractPeerMeta(loginReq),
+		UserID:          userID,
+		SetupKey:        loginReq.GetSetupKey(),
+	})
 	if err != nil {
-			return nil, err
+		log.Warnf("failed logging in peer %s", peerKey)
-		}
+		return nil, mapError(err)
-	}
-
-	network, err := s.accountManager.GetPeerNetwork(peer.ID)
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed getting peer network on login")
 	}
 
 	// if peer has reached this point then it has logged in
 	loginResp := &proto.LoginResponse{
 		WiretrusteeConfig: toWiretrusteeConfig(s.config, nil),
-		PeerConfig:        toPeerConfig(peer, network, s.accountManager.GetDNSDomain()),
+		PeerConfig:        toPeerConfig(peer, netMap.Network, s.accountManager.GetDNSDomain()),
 	}
 	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, loginResp)
 	if err != nil {
+		log.Warnf("failed encrypting peer %s message", peer.ID)
 		return nil, status.Errorf(codes.Internal, "failed logging in peer")
 	}
 
@@ -555,13 +440,7 @@ func (s *GRPCServer) IsHealthy(ctx context.Context, req *proto.Empty) (*proto.Em
 }
 
 // sendInitialSync sends initial proto.SyncResponse to the peer requesting synchronization
-func (s *GRPCServer) sendInitialSync(peerKey wgtypes.Key, peer *Peer, srv proto.ManagementService_SyncServer) error {
+func (s *GRPCServer) sendInitialSync(peerKey wgtypes.Key, peer *Peer, networkMap *NetworkMap, srv proto.ManagementService_SyncServer) error {
-	networkMap, err := s.accountManager.GetNetworkMap(peer.ID)
-	if err != nil {
-		log.Warnf("error getting a list of peers for a peer %s", peer.ID)
-		return err
-	}
-
 	// make secret time based TURN credentials optional
 	var turnCredentials *TURNCredentials
 	if s.config.TURNConfig.TimeBasedCredentials {

management/server/management_test.go

@@ -240,7 +240,8 @@ var _ = Describe("Management service", func() {
 		Context("with an invalid setup key", func() {
 			Specify("an error is returned", func() {
 				key, _ := wgtypes.GenerateKey()
-				message, err := encryption.EncryptMessage(serverPubKey, key, &mgmtProto.LoginRequest{SetupKey: "invalid setup key"})
+				message, err := encryption.EncryptMessage(serverPubKey, key, &mgmtProto.LoginRequest{SetupKey: "invalid setup key",
+					Meta: &mgmtProto.PeerSystemMeta{}})
 				Expect(err).NotTo(HaveOccurred())
 
 				resp, err := client.Login(context.TODO(), &mgmtProto.EncryptedMessage{
@@ -269,7 +270,7 @@ var _ = Describe("Management service", func() {
 				Expect(regResp).NotTo(BeNil())
 
 				// just login without registration
-				message, err := encryption.EncryptMessage(serverPubKey, key, &mgmtProto.LoginRequest{})
+				message, err := encryption.EncryptMessage(serverPubKey, key, &mgmtProto.LoginRequest{Meta: &mgmtProto.PeerSystemMeta{}})
 				Expect(err).NotTo(HaveOccurred())
 				loginResp, err := client.Login(context.TODO(), &mgmtProto.EncryptedMessage{
 					WgPubKey: key.PublicKey().String(),

management/server/mock_server/account_mock.go

@@ -25,12 +25,11 @@ type MockAccountManager struct {
 	GetPeerByKeyFunc                func(peerKey string) (*server.Peer, error)
 	GetPeersFunc                    func(accountID, userID string) ([]*server.Peer, error)
 	MarkPeerConnectedFunc           func(peerKey string, connected bool) error
-	MarkPeerLoginExpiredFunc        func(peerPubKey string, loginExpired bool) error
 	DeletePeerFunc                  func(accountID, peerKey, userID string) (*server.Peer, error)
 	GetPeerByIPFunc                 func(accountId string, peerIP string) (*server.Peer, error)
 	GetNetworkMapFunc               func(peerKey string) (*server.NetworkMap, error)
 	GetPeerNetworkFunc              func(peerKey string) (*server.Network, error)
-	AddPeerFunc                     func(setupKey string, userId string, peer *server.Peer) (*server.Peer, error)
+	AddPeerFunc                     func(setupKey string, userId string, peer *server.Peer) (*server.Peer, *server.NetworkMap, error)
 	GetGroupFunc                    func(accountID, groupID string) (*server.Group, error)
 	SaveGroupFunc                   func(accountID, userID string, group *server.Group) error
 	UpdateGroupFunc                 func(accountID string, groupID string, operations []server.GroupUpdateOperation) (*server.Group, error)
@@ -70,9 +69,9 @@ type MockAccountManager struct {
 	GetDNSSettingsFunc              func(accountID, userID string) (*server.DNSSettings, error)
 	SaveDNSSettingsFunc             func(accountID, userID string, dnsSettingsToSave *server.DNSSettings) error
 	GetPeerFunc                     func(accountID, peerID, userID string) (*server.Peer, error)
-	GetAccountByPeerIDFunc          func(peerID string) (*server.Account, error)
-	UpdatePeerLastLoginFunc         func(peerID string) error
 	UpdateAccountSettingsFunc       func(accountID, userID string, newSettings *server.Settings) (*server.Account, error)
+	LoginPeerFunc                   func(login server.PeerLogin) (*server.Peer, *server.NetworkMap, error)
+	SyncPeerFunc                    func(sync server.PeerSync) (*server.Peer, *server.NetworkMap, error)
 }
 
 // GetUsersFromAccount mock implementation of GetUsersFromAccount from server.AccountManager interface
@@ -165,14 +164,6 @@ func (am *MockAccountManager) MarkPeerConnected(peerKey string, connected bool)
 	return status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
 }
 
-// MarkPeerLoginExpired mock implementation of MarkPeerLoginExpired from server.AccountManager interface
-func (am *MockAccountManager) MarkPeerLoginExpired(peerPubKey string, loginExpired bool) error {
-	if am.MarkPeerLoginExpiredFunc != nil {
-		return am.MarkPeerLoginExpiredFunc(peerPubKey, loginExpired)
-	}
-	return status.Errorf(codes.Unimplemented, "method MarkPeerLoginExpired is not implemented")
-}
-
 // GetPeerByIP mock implementation of GetPeerByIP from server.AccountManager interface
 func (am *MockAccountManager) GetPeerByIP(accountId string, peerIP string) (*server.Peer, error) {
 	if am.GetPeerByIPFunc != nil {
@@ -202,11 +193,11 @@ func (am *MockAccountManager) AddPeer(
 	setupKey string,
 	userId string,
 	peer *server.Peer,
-) (*server.Peer, error) {
+) (*server.Peer, *server.NetworkMap, error) {
 	if am.AddPeerFunc != nil {
 		return am.AddPeerFunc(setupKey, userId, peer)
 	}
-	return nil, status.Errorf(codes.Unimplemented, "method AddPeer is not implemented")
+	return nil, nil, status.Errorf(codes.Unimplemented, "method AddPeer is not implemented")
 }
 
 // GetGroup mock implementation of GetGroup from server.AccountManager interface
@@ -541,22 +532,6 @@ func (am *MockAccountManager) GetPeer(accountID, peerID, userID string) (*server
 	return nil, status.Errorf(codes.Unimplemented, "method GetPeer is not implemented")
 }
 
-// GetAccountByPeerID mocks GetAccountByPeerID of the AccountManager interface
-func (am *MockAccountManager) GetAccountByPeerID(peerID string) (*server.Account, error) {
-	if am.GetAccountByPeerIDFunc != nil {
-		return am.GetAccountByPeerIDFunc(peerID)
-	}
-	return nil, status.Errorf(codes.Unimplemented, "method GetAccountByPeerID is not implemented")
-}
-
-// UpdatePeerLastLogin mocks UpdatePeerLastLogin of the AccountManager interface
-func (am *MockAccountManager) UpdatePeerLastLogin(peerID string) error {
-	if am.UpdatePeerLastLoginFunc != nil {
-		return am.UpdatePeerLastLoginFunc(peerID)
-	}
-	return status.Errorf(codes.Unimplemented, "method UpdatePeerLastLogin is not implemented")
-}
-
 // UpdateAccountSettings mocks UpdateAccountSettings of the AccountManager interface
 func (am *MockAccountManager) UpdateAccountSettings(accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
 	if am.UpdateAccountSettingsFunc != nil {
@@ -564,3 +539,19 @@ func (am *MockAccountManager) UpdateAccountSettings(accountID, userID string, ne
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method UpdateAccountSettings is not implemented")
 }
+
+// LoginPeer mocks LoginPeer of the AccountManager interface
+func (am *MockAccountManager) LoginPeer(login server.PeerLogin) (*server.Peer, *server.NetworkMap, error) {
+	if am.LoginPeerFunc != nil {
+		return am.LoginPeerFunc(login)
+	}
+	return nil, nil, status.Errorf(codes.Unimplemented, "method LoginPeer is not implemented")
+}
+
+// SyncPeer mocks SyncPeer of the AccountManager interface
+func (am *MockAccountManager) SyncPeer(sync server.PeerSync) (*server.Peer, *server.NetworkMap, error) {
+	if am.SyncPeerFunc != nil {
+		return am.SyncPeerFunc(sync)
+	}
+	return nil, nil, status.Errorf(codes.Unimplemented, "method SyncPeer is not implemented")
+}

management/server/nameserver_test.go

@@ -1147,11 +1147,11 @@ func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error
 		return nil, err
 	}
 
-	_, err = am.AddPeer("", userID, peer1)
+	_, _, err = am.AddPeer("", userID, peer1)
 	if err != nil {
 		return nil, err
 	}
-	_, err = am.AddPeer("", userID, peer2)
+	_, _, err = am.AddPeer("", userID, peer2)
 	if err != nil {
 		return nil, err
 	}

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)
+	return am.getNetworkMap(peer, account), nil
-	// 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
 }
 
 // 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
+// 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 key is invalid
+// 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
+	err = am.updateAccountPeers(account)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	networkMap := account.GetPeerNetworkMap(peer.ID, am.dnsDomain)
+	return newPeer, networkMap, nil
 }
 
-// UpdatePeerLastLogin sets Peer.LastLogin to the current timestamp.
+func (am *DefaultAccountManager) checkPeerLoginExpiration(loginUserID string, peer *Peer, account *Account) error {
-func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
+	if peer.AddedWithSSOLogin() {
-	account, err := am.Store.GetAccountByPeerID(peerID)
+		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)
+	peer, err := account.FindPeerByPubKey(sync.WireGuardPubKey)
-	if peer == nil {
+	if err != nil {
-		return status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
+		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)
+func (am *DefaultAccountManager) checkAndUpdatePeerSSHKey(peer *Peer, account *Account, newSSHKey string) (*Peer, error) {
-	if err != nil {
+	if len(newSSHKey) == 0 {
-		return err
+		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 updatePeerMeta(peer *Peer, meta PeerSystemMeta, account *Account) *Peer {
-func (am *DefaultAccountManager) UpdatePeerMeta(peerID string, meta PeerSystemMeta) error {
+	peer.UpdateMeta(meta)
-
-	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
 	account.UpdatePeer(peer)
-
+	return peer
-	err = am.Store.SaveAccount(account)
-	if err != nil {
-		return err
-	}
-	return nil
 }
 
 // getPeersByACL returns all peers that given peer has access to.

management/server/peer_test.go

@@ -91,10 +91,9 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
 		return
 	}
 
-	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer1, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -107,10 +106,9 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
 		t.Fatal(err)
 		return
 	}
-	_, err = manager.AddPeer(setupKey.Key, "", &Peer{
+	_, _, err = manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -164,10 +162,9 @@ func TestAccountManager_GetNetworkMapWithRule(t *testing.T) {
 		return
 	}
 
-	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer1, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -180,10 +177,9 @@ func TestAccountManager_GetNetworkMapWithRule(t *testing.T) {
 		t.Fatal(err)
 		return
 	}
-	peer2, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer2, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -333,10 +329,9 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
 		return
 	}
 
-	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer1, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -349,10 +344,9 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
 		t.Fatal(err)
 		return
 	}
-	_, err = manager.AddPeer(setupKey.Key, "", &Peer{
+	_, _, err = manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -402,10 +396,9 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 		return
 	}
 
-	peer1, err := manager.AddPeer("", someUser, &Peer{
+	peer1, _, err := manager.AddPeer("", someUser, &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {
@@ -420,10 +413,9 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 	}
 
 	// the second peer added with a setup key
-	peer2, err := manager.AddPeer(setupKey.Key, "", &Peer{
+	peer2, _, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
-		Name: "test-peer-2",
 	})
 
 	if err != nil {

management/server/status/error.go

@@ -31,6 +31,9 @@ const (
 
 	// BadRequest indicates that user is not authorized
 	BadRequest Type = 9
+
+	// Unauthenticated indicates that user is not authenticated due to absence of valid credentials
+	Unauthenticated Type = 10
 )
 
 // Type is a type of the Error