Refactor Sync

When we delete a peer from an account, we save the account in the file store.
The file store maintains peerID -> accountID and peerKey -> accountID indices.
Those can't be updated when we delete a peer because the store saves the whole account
without a peer already and has no access to the removed peer.
In this PR, we dynamically check if there are stale indices when GetAccountByPeerPubKey
and GetAccountByPeerID.

management/server/account.go

@@ -63,7 +63,6 @@ type AccountManager interface {
 	GetNetworkMap(peerID string) (*NetworkMap, error)
 	GetPeerNetwork(peerID string) (*Network, error)
 	AddPeer(setupKey, userID string, peer *Peer) (*Peer, 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 +95,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, error)
+	SyncPeer(sync PeerSync) (*Peer, *NetworkMap, error)
 }
 
 type DefaultAccountManager struct {

management/server/account_test.go

@@ -544,8 +544,7 @@ func TestAccountManager_AddPeer(t *testing.T) {
 
 	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  expectedPeerKey,
-		Meta: PeerSystemMeta{},
-		Name: expectedPeerKey,
+		Meta: PeerSystemMeta{Hostname: expectedPeerKey},
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v", err)
@@ -613,8 +612,7 @@ func TestAccountManager_AddPeerWithUserID(t *testing.T) {
 
 	peer, err := manager.AddPeer("", userID, &Peer{
 		Key:  expectedPeerKey,
-		Meta: PeerSystemMeta{},
-		Name: expectedPeerKey,
+		Meta: PeerSystemMeta{Hostname: expectedPeerKey},
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v, account users: %v", err, account.CreatedBy)
@@ -696,8 +694,7 @@ func TestAccountManager_NetworkUpdates(t *testing.T) {
 
 		peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
 			Key:  expectedPeerKey,
-			Meta: PeerSystemMeta{},
-			Name: expectedPeerKey,
+			Meta: PeerSystemMeta{Hostname: expectedPeerKey},
 		})
 		if err != nil {
 			t.Fatalf("expecting peer1 to be added, got failure %v", err)
@@ -866,8 +863,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 
 	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey,
-		Meta: PeerSystemMeta{},
-		Name: peerKey,
+		Meta: PeerSystemMeta{Hostname: peerKey},
 	})
 	if err != nil {
 		t.Errorf("expecting peer to be added, got failure %v", err)
@@ -951,7 +947,7 @@ func TestGetUsersFromAccount(t *testing.T) {
 	}
 }
 
-func TestAccountManager_UpdatePeerMeta(t *testing.T) {
+/*func TestAccountManager_UpdatePeerMeta(t *testing.T) {
 	manager, err := createManager(t)
 	if err != nil {
 		t.Fatal(err)
@@ -1018,7 +1014,7 @@ func TestAccountManager_UpdatePeerMeta(t *testing.T) {
 	}
 
 	assert.Equal(t, newMeta, p.Meta)
-}
+}*/
 
 func TestAccount_GetPeerRules(t *testing.T) {
 
@@ -1304,8 +1300,7 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
 	require.NoError(t, err, "unable to generate WireGuard key")
 	peer, err := manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
-		Name:                   "test-peer",
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")
@@ -1353,8 +1348,7 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
 	require.NoError(t, err, "unable to generate WireGuard key")
 	_, err = manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
-		Name:                   "test-peer",
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")
@@ -1395,8 +1389,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 	require.NoError(t, err, "unable to generate WireGuard key")
 	_, err = manager.AddPeer("", userID, &Peer{
 		Key:                    key.PublicKey().String(),
-		Meta:                   PeerSystemMeta{},
-		Name:                   "test-peer",
+		Meta:                   PeerSystemMeta{Hostname: "test-peer"},
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")

management/server/file_store.go

@@ -225,7 +225,6 @@ func (s *FileStore) SaveAccount(account *Account) error {
 
 	accountCopy := account.Copy()
 
-	// todo will override, handle existing keys
 	s.Accounts[accountCopy.Id] = accountCopy
 
 	// todo check that account.Id and keyId are not exist already
@@ -354,6 +353,14 @@ func (s *FileStore) GetAccountByPeerID(peerID string) (*Account, error) {
 		return nil, err
 	}
 
+	// this protection is needed because when we delete a peer, we don't really remove index peerID -> accountID.
+	// check Account.Peers for a match
+	if _, ok := account.Peers[peerID]; !ok {
+		delete(s.PeerID2AccountID, peerID)
+		log.Warnf("removed stale peerID %s to accountID %s index", peerID, accountID)
+		return nil, status.Errorf(status.NotFound, "provided peer doesn't exists %s", peerID)
+	}
+
 	return account.Copy(), nil
 }
 
@@ -372,6 +379,21 @@ func (s *FileStore) GetAccountByPeerPubKey(peerKey string) (*Account, error) {
 		return nil, err
 	}
 
+	// this protection is needed because when we delete a peer, we don't really remove index peerKey -> accountID.
+	// check Account.Peers for a match
+	stale := true
+	for _, peer := range account.Peers {
+		if peer.Key == peerKey {
+			stale = false
+			break
+		}
+	}
+	if stale {
+		delete(s.PeerKeyID2AccountID, peerKey)
+		log.Warnf("removed stale peerKey %s to accountID %s index", peerKey, accountID)
+		return nil, status.Errorf(status.NotFound, "provided peer doesn't exists %s", peerKey)
+	}
+
 	return account.Copy(), nil
 }
 

management/server/file_store_test.go

@@ -14,6 +14,44 @@ type accounts struct {
 	Accounts map[string]*Account
 }
 
+func TestStalePeerIndices(t *testing.T) {
+	storeDir := t.TempDir()
+
+	err := util.CopyFileContents("testdata/store.json", filepath.Join(storeDir, "store.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	store, err := NewFileStore(storeDir)
+	if err != nil {
+		return
+	}
+
+	account, err := store.GetAccount("bf1c8084-ba50-4ce7-9439-34653001fc3b")
+	require.NoError(t, err)
+
+	peerID := "some_peer"
+	peerKey := "some_peer_key"
+	account.Peers[peerID] = &Peer{
+		ID:  peerID,
+		Key: peerKey,
+	}
+
+	err = store.SaveAccount(account)
+	require.NoError(t, err)
+
+	account.DeletePeer(peerID)
+
+	err = store.SaveAccount(account)
+	require.NoError(t, err)
+
+	_, err = store.GetAccountByPeerID(peerID)
+	require.Error(t, err, "expecting to get an error when found stale index")
+
+	_, err = store.GetAccountByPeerPubKey(peerKey)
+	require.Error(t, err, "expecting to get an error when found stale index")
+}
+
 func TestNewStore(t *testing.T) {
 	store := newStore(t)
 

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())
-		msg := status.Errorf(codes.InvalidArgument, "invalid request message from %s,remote addr is %s", peerKey.String(), p.Addr.String())
-		log.Debug(msg)
-		return msg
+		return err
 	}
 
-	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) {
-	var (
-		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
+// maps internal internalStatus.Error to gRPC status.Error
+func mapError(err error) error {
+	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 status.Errorf(codes.FailedPrecondition, e.Message)
+		case internalStatus.NotFound:
+			return status.Errorf(codes.NotFound, e.Message)
+		default:
 		}
-	} else {
-		log.Debugln("using setup key to register peer")
-		reqSetupKey = req.GetSetupKey()
-		userID = ""
 	}
+	return status.Errorf(codes.Internal, "failed handling request")
+}
 
-	meta := req.GetMeta()
-	if meta == nil {
-		return nil, status.Errorf(codes.InvalidArgument, "peer meta data was not provided")
+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(),
 	}
+}
 
-	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(),
-		},
-	})
+func (s *GRPCServer) parseRequest(req *proto.EncryptedMessage, parsed pb.Message) (wgtypes.Key, error) {
+	peerKey, err := wgtypes.ParseKey(req.GetWgPubKey())
 	if err != nil {
-		if e, ok := internalStatus.FromError(err); ok {
-			switch e.Type() {
-			case internalStatus.PreconditionFailed:
-				return nil, status.Errorf(codes.FailedPrecondition, e.Message)
-			case internalStatus.NotFound:
-				return nil, status.Errorf(codes.NotFound, e.Message)
-			default:
-			}
-		}
-		return nil, status.Errorf(codes.Internal, "failed registering new peer")
+		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
-	networkMap, err := s.accountManager.GetNetworkMap(peer.ID)
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, parsed)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "unable to fetch network map after registering peer, error: %v", err)
-	}
-	// 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())
-		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 wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "invalid request message")
 	}
 
-	return peer, nil
+	return peerKey, nil
 }
 
 // Login endpoint first checks whether peer is registered under any account
@@ -323,103 +266,51 @@ 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 err != nil {
-		if errStatus, ok := internalStatus.FromError(err); ok && errStatus.Type() == internalStatus.NotFound {
-			// peer doesn't exist -> check if setup key was provided
-			if loginReq.GetJwtToken() == "" && loginReq.GetSetupKey() == "" {
-				// 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
-			}
+	if loginReq.GetMeta() == nil {
+		msg := status.Errorf(codes.FailedPrecondition,
+			"peer system meta has to be provided to log in. Peer %s, remote addr %s", peerKey.String(),
+			p.Addr.String())
+		log.Warn(msg)
+		return nil, msg
+	}
 
-			// setup key or jwt is present -> try normal registration flow
-			peer, err = s.registerPeer(peerKey, loginReq)
-			if err != nil {
-				return nil, 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(),
-		},
-		)
+	userID := ""
+	// 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() != "" {
+		// todo what about the case when JWT provided is expired?
+		userID, err = s.validateToken(loginReq.GetJwtToken())
 		if err != nil {
-			log.Errorf("failed updating peer system meta data %s", peerKey.String())
-			return nil, status.Error(codes.Internal, "internal server error")
+			log.Warnf("failed validating JWT token sent from peer %s", peerKey)
+			return nil, mapError(err)
 		}
 	}
-
-	// 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 {
-		err = s.accountManager.UpdatePeerSSHKey(peer.ID, string(sshKey))
-		if err != nil {
-			return nil, err
-		}
+	peer, err := s.accountManager.LoginPeer(PeerLogin{
+		WireGuardPubKey: peerKey.String(),
+		SSHKey:          string(sshKey),
+		Meta:            extractPeerMeta(loginReq),
+		UserID:          userID,
+		SetupKey:        loginReq.GetSetupKey(),
+	})
+	if err != nil {
+		log.Warnf("failed logging in peer %s", peerKey)
+		return nil, mapError(err)
 	}
 
 	network, err := s.accountManager.GetPeerNetwork(peer.ID)
 	if err != nil {
+		log.Warnf("failed getting peer %s network on login", peer.ID)
 		return nil, status.Errorf(codes.Internal, "failed getting peer network on login")
 	}
 
@@ -430,6 +321,7 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 	}
 	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 +447,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 {
-	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
-	}
-
+func (s *GRPCServer) sendInitialSync(peerKey wgtypes.Key, peer *Peer, networkMap *NetworkMap, srv proto.ManagementService_SyncServer) error {
 	// make secret time based TURN credentials optional
 	var turnCredentials *TURNCredentials
 	if s.config.TURNConfig.TimeBasedCredentials {

management/server/http/handler.go

@@ -19,6 +19,12 @@ type AuthCfg struct {
 	KeysLocation string
 }
 
+type apiHandler struct {
+	Router         *mux.Router
+	AccountManager s.AccountManager
+	AuthCfg        AuthCfg
+}
+
 // APIHandler creates the Management service HTTP API handler registering all the available endpoints.
 func APIHandler(accountManager s.AccountManager, appMetrics telemetry.AppMetrics, authCfg AuthCfg) (http.Handler, error) {
 	jwtMiddleware, err := middleware.NewJwtMiddleware(
@@ -39,69 +45,27 @@ func APIHandler(accountManager s.AccountManager, appMetrics telemetry.AppMetrics
 	rootRouter := mux.NewRouter()
 	metricsMiddleware := appMetrics.HTTPMiddleware()
 
-	apiHandler := rootRouter.PathPrefix("/api").Subrouter()
-	apiHandler.Use(metricsMiddleware.Handler, corsMiddleware.Handler, jwtMiddleware.Handler, acMiddleware.Handler)
+	router := rootRouter.PathPrefix("/api").Subrouter()
+	router.Use(metricsMiddleware.Handler, corsMiddleware.Handler, jwtMiddleware.Handler, acMiddleware.Handler)
 
-	groupsHandler := NewGroupsHandler(accountManager, authCfg)
-	rulesHandler := NewRulesHandler(accountManager, authCfg)
-	peersHandler := NewPeersHandler(accountManager, authCfg)
-	keysHandler := NewSetupKeysHandler(accountManager, authCfg)
-	userHandler := NewUsersHandler(accountManager, authCfg)
-	routesHandler := NewRoutesHandler(accountManager, authCfg)
-	nameserversHandler := NewNameserversHandler(accountManager, authCfg)
-	eventsHandler := NewEventsHandler(accountManager, authCfg)
-	dnsSettingsHandler := NewDNSSettingsHandler(accountManager, authCfg)
-	accountsHandler := NewAccountsHandler(accountManager, authCfg)
+	api := apiHandler{
+		Router:         router,
+		AccountManager: accountManager,
+		AuthCfg:        authCfg,
+	}
 
-	apiHandler.HandleFunc("/accounts/{id}", accountsHandler.UpdateAccount).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/accounts", accountsHandler.GetAllAccounts).Methods("GET", "OPTIONS")
+	api.addAccountsEndpoint()
+	api.addPeersEndpoint()
+	api.addUsersEndpoint()
+	api.addSetupKeysEndpoint()
+	api.addRulesEndpoint()
+	api.addGroupsEndpoint()
+	api.addRoutesEndpoint()
+	api.addDNSNameserversEndpoint()
+	api.addDNSSettingEndpoint()
+	api.addEventsEndpoint()
 
-	apiHandler.HandleFunc("/peers", peersHandler.GetAllPeers).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/peers/{id}", peersHandler.HandlePeer).
-		Methods("GET", "PUT", "DELETE", "OPTIONS")
-	apiHandler.HandleFunc("/users", userHandler.GetAllUsers).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/users/{id}", userHandler.UpdateUser).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/users", userHandler.CreateUser).Methods("POST", "OPTIONS")
-
-	apiHandler.HandleFunc("/setup-keys", keysHandler.GetAllSetupKeys).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/setup-keys", keysHandler.CreateSetupKey).Methods("POST", "OPTIONS")
-	apiHandler.HandleFunc("/setup-keys/{id}", keysHandler.GetSetupKey).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/setup-keys/{id}", keysHandler.UpdateSetupKey).Methods("PUT", "OPTIONS")
-
-	apiHandler.HandleFunc("/rules", rulesHandler.GetAllRules).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/rules", rulesHandler.CreateRule).Methods("POST", "OPTIONS")
-	apiHandler.HandleFunc("/rules/{id}", rulesHandler.UpdateRule).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/rules/{id}", rulesHandler.PatchRule).Methods("PATCH", "OPTIONS")
-	apiHandler.HandleFunc("/rules/{id}", rulesHandler.GetRule).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/rules/{id}", rulesHandler.DeleteRule).Methods("DELETE", "OPTIONS")
-
-	apiHandler.HandleFunc("/groups", groupsHandler.GetAllGroups).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/groups", groupsHandler.CreateGroup).Methods("POST", "OPTIONS")
-	apiHandler.HandleFunc("/groups/{id}", groupsHandler.UpdateGroup).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/groups/{id}", groupsHandler.PatchGroup).Methods("PATCH", "OPTIONS")
-	apiHandler.HandleFunc("/groups/{id}", groupsHandler.GetGroup).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/groups/{id}", groupsHandler.DeleteGroup).Methods("DELETE", "OPTIONS")
-
-	apiHandler.HandleFunc("/routes", routesHandler.GetAllRoutes).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/routes", routesHandler.CreateRoute).Methods("POST", "OPTIONS")
-	apiHandler.HandleFunc("/routes/{id}", routesHandler.UpdateRoute).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/routes/{id}", routesHandler.PatchRoute).Methods("PATCH", "OPTIONS")
-	apiHandler.HandleFunc("/routes/{id}", routesHandler.GetRoute).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/routes/{id}", routesHandler.DeleteRoute).Methods("DELETE", "OPTIONS")
-
-	apiHandler.HandleFunc("/dns/nameservers", nameserversHandler.GetAllNameservers).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/dns/nameservers", nameserversHandler.CreateNameserverGroup).Methods("POST", "OPTIONS")
-	apiHandler.HandleFunc("/dns/nameservers/{id}", nameserversHandler.UpdateNameserverGroup).Methods("PUT", "OPTIONS")
-	apiHandler.HandleFunc("/dns/nameservers/{id}", nameserversHandler.PatchNameserverGroup).Methods("PATCH", "OPTIONS")
-	apiHandler.HandleFunc("/dns/nameservers/{id}", nameserversHandler.GetNameserverGroup).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/dns/nameservers/{id}", nameserversHandler.DeleteNameserverGroup).Methods("DELETE", "OPTIONS")
-
-	apiHandler.HandleFunc("/events", eventsHandler.GetAllEvents).Methods("GET", "OPTIONS")
-
-	apiHandler.HandleFunc("/dns/settings", dnsSettingsHandler.GetDNSSettings).Methods("GET", "OPTIONS")
-	apiHandler.HandleFunc("/dns/settings", dnsSettingsHandler.UpdateDNSSettings).Methods("PUT", "OPTIONS")
-
-	err = apiHandler.Walk(func(route *mux.Route, _ *mux.Router, _ []*mux.Route) error {
+	err = api.Router.Walk(func(route *mux.Route, _ *mux.Router, _ []*mux.Route) error {
 		methods, err := route.GetMethods()
 		if err != nil {
 			return err
@@ -124,3 +88,82 @@ func APIHandler(accountManager s.AccountManager, appMetrics telemetry.AppMetrics
 
 	return rootRouter, nil
 }
+
+func (apiHandler *apiHandler) addAccountsEndpoint() {
+	accountsHandler := NewAccountsHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/accounts/{id}", accountsHandler.UpdateAccount).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/accounts", accountsHandler.GetAllAccounts).Methods("GET", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addPeersEndpoint() {
+	peersHandler := NewPeersHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/peers", peersHandler.GetAllPeers).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/peers/{id}", peersHandler.HandlePeer).
+		Methods("GET", "PUT", "DELETE", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addUsersEndpoint() {
+	userHandler := NewUsersHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/users", userHandler.GetAllUsers).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/users/{id}", userHandler.UpdateUser).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/users", userHandler.CreateUser).Methods("POST", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addSetupKeysEndpoint() {
+	keysHandler := NewSetupKeysHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/setup-keys", keysHandler.GetAllSetupKeys).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/setup-keys", keysHandler.CreateSetupKey).Methods("POST", "OPTIONS")
+	apiHandler.Router.HandleFunc("/setup-keys/{id}", keysHandler.GetSetupKey).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/setup-keys/{id}", keysHandler.UpdateSetupKey).Methods("PUT", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addRulesEndpoint() {
+	rulesHandler := NewRulesHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/rules", rulesHandler.GetAllRules).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/rules", rulesHandler.CreateRule).Methods("POST", "OPTIONS")
+	apiHandler.Router.HandleFunc("/rules/{id}", rulesHandler.UpdateRule).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/rules/{id}", rulesHandler.PatchRule).Methods("PATCH", "OPTIONS")
+	apiHandler.Router.HandleFunc("/rules/{id}", rulesHandler.GetRule).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/rules/{id}", rulesHandler.DeleteRule).Methods("DELETE", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addGroupsEndpoint() {
+	groupsHandler := NewGroupsHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/groups", groupsHandler.GetAllGroups).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/groups", groupsHandler.CreateGroup).Methods("POST", "OPTIONS")
+	apiHandler.Router.HandleFunc("/groups/{id}", groupsHandler.UpdateGroup).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/groups/{id}", groupsHandler.PatchGroup).Methods("PATCH", "OPTIONS")
+	apiHandler.Router.HandleFunc("/groups/{id}", groupsHandler.GetGroup).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/groups/{id}", groupsHandler.DeleteGroup).Methods("DELETE", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addRoutesEndpoint() {
+	routesHandler := NewRoutesHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/routes", routesHandler.GetAllRoutes).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/routes", routesHandler.CreateRoute).Methods("POST", "OPTIONS")
+	apiHandler.Router.HandleFunc("/routes/{id}", routesHandler.UpdateRoute).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/routes/{id}", routesHandler.PatchRoute).Methods("PATCH", "OPTIONS")
+	apiHandler.Router.HandleFunc("/routes/{id}", routesHandler.GetRoute).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/routes/{id}", routesHandler.DeleteRoute).Methods("DELETE", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addDNSNameserversEndpoint() {
+	nameserversHandler := NewNameserversHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/dns/nameservers", nameserversHandler.GetAllNameservers).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/nameservers", nameserversHandler.CreateNameserverGroup).Methods("POST", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/nameservers/{id}", nameserversHandler.UpdateNameserverGroup).Methods("PUT", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/nameservers/{id}", nameserversHandler.PatchNameserverGroup).Methods("PATCH", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/nameservers/{id}", nameserversHandler.GetNameserverGroup).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/nameservers/{id}", nameserversHandler.DeleteNameserverGroup).Methods("DELETE", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addDNSSettingEndpoint() {
+	dnsSettingsHandler := NewDNSSettingsHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/dns/settings", dnsSettingsHandler.GetDNSSettings).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/dns/settings", dnsSettingsHandler.UpdateDNSSettings).Methods("PUT", "OPTIONS")
+}
+
+func (apiHandler *apiHandler) addEventsEndpoint() {
+	eventsHandler := NewEventsHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/events", eventsHandler.GetAllEvents).Methods("GET", "OPTIONS")
+}

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

@@ -71,8 +71,9 @@ type MockAccountManager struct {
 	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, error)
+	SyncPeerFunc                    func(sync server.PeerSync) (*server.Peer, *server.NetworkMap, error)
 }
 
 // GetUsersFromAccount mock implementation of GetUsersFromAccount from server.AccountManager interface
@@ -549,14 +550,6 @@ func (am *MockAccountManager) GetAccountByPeerID(peerID string) (*server.Account
 	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 +557,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, error) {
+	if am.LoginPeerFunc != nil {
+		return am.LoginPeerFunc(login)
+	}
+	return 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/peer.go

@@ -36,6 +36,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 +113,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()
@@ -190,6 +219,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)
@@ -419,6 +460,34 @@ 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 {
+	aclPeers := account.getPeersByACL(peer.ID)
+	// Please mind, that the returned route.Route objects will contain Peer.Key instead of Peer.ID.
+	routesUpdate := account.getRoutesToSync(peer.ID, aclPeers)
+
+	dnsManagementStatus := account.getPeerDNSManagementStatus(peer.ID)
+	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, peer.ID)
+	}
+
+	return &NetworkMap{
+		Peers:     aclPeers,
+		Network:   account.Network.Copy(),
+		Routes:    routesUpdate,
+		DNSConfig: dnsUpdate,
+	}
+}
+
 // GetNetworkMap returns Network map for a given peer (omits original peer from the Peers result)
 func (am *DefaultAccountManager) GetNetworkMap(peerID string) (*NetworkMap, error) {
 
@@ -432,31 +501,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)
-	// Please mind, that the returned route.Route objects will contain Peer.Key instead of Peer.ID.
-	routesUpdate := account.getRoutesToSync(peerID, aclPeers)
-
-	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:     aclPeers,
-		Network:   account.Network.Copy(),
-		Routes:    routesUpdate,
-		DNSConfig: dnsUpdate,
-	}, err
+	return am.getNetworkMap(peer, account), nil
 }
 
 // GetPeerNetwork returns the Network for a given peer
@@ -471,13 +516,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) {
+	if setupKey == "" && userID == "" {
+		// no auth method provided => reject access
+		return 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
@@ -529,7 +578,7 @@ 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
 	}
@@ -547,7 +596,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()},
@@ -596,43 +645,163 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *Peer) (*
 	opEvent.Meta = newPeer.EventMeta(am.GetDNSDomain())
 	am.storeEvent(opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)
 
+	err = am.updateAccountPeers(account)
+	if err != nil {
+		return nil, err
+	}
+
 	return newPeer, nil
 }
 
-// UpdatePeerLastLogin sets Peer.LastLogin to the current timestamp.
-func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
-	account, err := am.Store.GetAccountByPeerID(peerID)
-	if err != nil {
-		return err
+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)
+			} else {
+				// 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 {
+		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, err
 	}
 
+	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, 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, 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, err
+	}
+
+	peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
+	if err != nil {
+		return nil, err
+	}
+
+	err = am.checkPeerLoginExpiration(login.UserID, peer, account)
+	if err != nil {
+		return nil, err
+	}
+
+	peer = am.updatePeerMeta(peer, peer.Meta, account)
+
+	peer, err = am.checkAndUpdatePeerSSHKey(peer, account, login.SSHKey)
+	if err != nil {
+		return nil, err
+	}
+
+	err = am.Store.SaveAccount(account)
+	if err != nil {
+		return nil, err
+	}
+
+	return peer, 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
@@ -724,35 +893,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 (am *DefaultAccountManager) 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.

management/server/peer_test.go

@@ -92,8 +92,7 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
 
 	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
 	})
 
 	if err != nil {
@@ -108,8 +107,7 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
 	}
 	_, err = manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
 	})
 
 	if err != nil {
@@ -165,8 +163,7 @@ func TestAccountManager_GetNetworkMapWithRule(t *testing.T) {
 
 	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
 	})
 
 	if err != nil {
@@ -181,8 +178,7 @@ func TestAccountManager_GetNetworkMapWithRule(t *testing.T) {
 	}
 	peer2, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
 	})
 
 	if err != nil {
@@ -334,8 +330,7 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
 
 	peer1, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-1"},
 	})
 
 	if err != nil {
@@ -350,8 +345,7 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
 	}
 	_, err = manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
 	})
 
 	if err != nil {
@@ -403,8 +397,7 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 
 	peer1, err := manager.AddPeer("", someUser, &Peer{
 		Key:  peerKey1.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "test-peer-2"},
 	})
 
 	if err != nil {
@@ -421,8 +414,7 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 	// the second peer added with a setup key
 	peer2, err := manager.AddPeer(setupKey.Key, "", &Peer{
 		Key:  peerKey2.PublicKey().String(),
-		Meta: PeerSystemMeta{},
-		Name: "test-peer-2",
+		Meta: PeerSystemMeta{Hostname: "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