Move peer login to account manager

management/server/account.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"github.com/netbirdio/netbird/management/server/http/middleware"
 	"math/rand"
 	"net"
 	"net/netip"
@@ -63,7 +64,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)
@@ -98,6 +98,7 @@ type AccountManager interface {
 	GetPeer(accountID, peerID, userID string) (*Peer, error)
 	UpdatePeerLastLogin(peerID string) error
 	UpdateAccountSettings(accountID, userID string, newSettings *Settings) (*Account, error)
+	LoginPeer(login PeerLogin) (*Peer, error)
 }
 
 type DefaultAccountManager struct {
@@ -119,8 +120,10 @@ type DefaultAccountManager struct {
 	// singleAccountModeDomain is a domain to use in singleAccountMode setup
 	singleAccountModeDomain string
 	// dnsDomain is used for peer resolution. This is appended to the peer's name
-	dnsDomain       string
-	peerLoginExpiry Scheduler
+	dnsDomain          string
+	peerLoginExpiry    Scheduler
+	jwtMiddleware      *middleware.JWTMiddleware
+	jwtClaimsExtractor *jwtclaims.ClaimsExtractor
 }
 
 // Settings represents Account settings structure that can be modified via API and Dashboard

management/server/account_test.go

@@ -951,7 +951,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 +1018,7 @@ func TestAccountManager_UpdatePeerMeta(t *testing.T) {
 	}
 
 	assert.Equal(t, newMeta, p.Meta)
-}
+}*/
 
 func TestAccount_GetPeerRules(t *testing.T) {
 

management/server/grpcserver.go

@@ -230,84 +230,54 @@ 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) parseLoginRequest(req *proto.EncryptedMessage) (*proto.LoginRequest, 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 nil, 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)
+	loginReq := &proto.LoginRequest{}
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, loginReq)
 	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 nil, wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "invalid request message")
 	}
 
-	return peer, nil
+	return loginReq, peerKey, nil
+
 }
 
 // Login endpoint first checks whether peer is registered under any account
@@ -323,99 +293,43 @@ 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())
+	loginReq, peerKey, err := s.parseLoginRequest(req)
 	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)
+		return nil, err
 	}
 
-	loginReq := &proto.LoginRequest{}
-	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, loginReq)
-	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "invalid request message")
+	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
 	}
 
-	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
-			}
-
-			// 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 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")
+			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 {
+		return nil, mapError(err)
 	}
 
 	network, err := s.accountManager.GetPeerNetwork(peer.ID)

management/server/peer.go

@@ -36,6 +36,20 @@ type PeerStatus struct {
 	LoginExpired bool
 }
 
+// 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 +107,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 +213,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)
@@ -471,13 +506,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
@@ -547,7 +586,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,9 +635,94 @@ 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
 }
 
+// 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
+	}
+
+	expired, expiresIn := peer.LoginExpired(account.Settings.PeerLoginExpiration)
+	expired = account.Settings.PeerLoginExpirationEnabled && expired
+	if peer.UserID != "" && (expired || peer.Status.LoginExpired) {
+		if login.UserID == "" {
+			// absence of a user ID indicates that JWT wasn't provided.
+			peer, err = am.markPeerLoginExpired(peer, account, true)
+			if err != nil {
+				return nil, err
+			}
+			return nil, 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 != login.UserID {
+				log.Warnf("user mismatch when loggin in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
+				return nil, status.Errorf(status.Unauthenticated, "can't login")
+			}
+			peer = am.updatePeerLastLogin(peer, account)
+		}
+	}
+
+	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
+}
+
 // UpdatePeerLastLogin sets Peer.LastLogin to the current timestamp.
 func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
 	account, err := am.Store.GetAccountByPeerID(peerID)
@@ -624,7 +748,6 @@ func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
 	newStatus := peer.Status.Copy()
 	newStatus.LoginExpired = false
 	peer.Status = newStatus
-
 	account.UpdatePeer(peer)
 
 	err = am.Store.SaveAccount(account)
@@ -635,6 +758,34 @@ func (am *DefaultAccountManager) UpdatePeerLastLogin(peerID string) error {
 	return nil
 }
 
+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
+	}
+
+	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
 func (am *DefaultAccountManager) UpdatePeerSSHKey(peerID string, sshKey string) error {
 
@@ -724,35 +875,10 @@ func (am *DefaultAccountManager) GetPeer(accountID, peerID, userID string) (*Pee
 	return nil, status.Errorf(status.Internal, "user %s has no access to peer %s under account %s", userID, peerID, accountID)
 }
 
-// UpdatePeerMeta updates peer's system metadata
-func (am *DefaultAccountManager) UpdatePeerMeta(peerID string, meta PeerSystemMeta) error {
-
-	account, err := am.Store.GetAccountByPeerID(peerID)
-	if err != nil {
-		return err
-	}
-
-	unlock := am.Store.AcquireAccountLock(account.Id)
-	defer unlock()
-
-	peer := account.GetPeer(peerID)
-	if peer == nil {
-		return status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
-	}
-
-	// Avoid overwriting UIVersion if the update was triggered sole by the CLI client
-	if meta.UIVersion == "" {
-		meta.UIVersion = peer.Meta.UIVersion
-	}
-
-	peer.Meta = meta
+func (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/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