Skip to content

Commit

Permalink
Refactor login with store.SavePeer (#2334)
Browse files Browse the repository at this point in the history 
This pull request refactors the login functionality by integrating store.SavePeer. The changes aim to improve the handling of peer login processes, particularly focusing on synchronization and error handling.

Changes:
- Refactored login logic to use store.SavePeer.
- Added checks for login without lock for login necessary checks from the client and utilized write lock for full login flow.
- Updated error handling with status.NewPeerLoginExpiredError().
- Moved geoIP check logic to a more appropriate place.
- Removed redundant calls and improved documentation.
- Moved the code to smaller methods to improve readability.
  • Loading branch information
@mlsmaycon
mlsmaycon authored Jul 29, 2024
1 parent 7321046 commit da39c8b
Show file tree
Hide file tree
Showing 2 changed files with 100 additions and 104 deletions.
199 changes: 95 additions & 104 deletions management/server/peer.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -453,6 +453,17 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
Location: peer.Location,
}

if am.geo != nil && newPeer.Location.ConnectionIP != nil {
location, err := am.geo.Lookup(newPeer.Location.ConnectionIP)
if err != nil {
log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", newPeer.Location.ConnectionIP.String(), err)
} else {
newPeer.Location.CountryCode = location.Country.ISOCode
newPeer.Location.CityName = location.City.Names.En
newPeer.Location.GeoNameID = location.City.GeonameID
}
}

// add peer to 'All' group
group, err := account.GetGroupAll()
if err != nil {
Expand Down Expand Up @@ -535,7 +546,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
}

if peerLoginExpired(ctx, peer, account.Settings) {
return nil, nil, nil, status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
return nil, nil, nil, status.NewPeerLoginExpiredError()
}

peer, updated := updatePeerMeta(peer, sync.Meta, account)
Expand Down Expand Up @@ -586,21 +597,10 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
// Try registering it.
newPeer := &nbpeer.Peer{
Key: login.WireGuardPubKey,
Meta: login.Meta,
SSHKey: login.SSHKey,
}
if am.geo != nil && login.ConnectionIP != nil {
location, err := am.geo.Lookup(login.ConnectionIP)
if err != nil {
log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", login.ConnectionIP.String(), err)
} else {
newPeer.Location.ConnectionIP = login.ConnectionIP
newPeer.Location.CountryCode = location.Country.ISOCode
newPeer.Location.CityName = location.City.Names.En
newPeer.Location.GeoNameID = location.City.GeonameID

}
Key: login.WireGuardPubKey,
Meta: login.Meta,
SSHKey: login.SSHKey,
Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
}

return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
Expand All @@ -610,44 +610,17 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
}

peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
if err != nil {
return nil, nil, nil, status.NewPeerNotRegisteredError()
}

accSettings, err := am.Store.GetAccountSettings(ctx, accountID)
if err != nil {
return nil, nil, nil, status.Errorf(status.Internal, "failed to get account settings: %s", err)
}

var isWriteLock bool

// duplicated logic from after the lock to have an early exit
expired := peerLoginExpired(ctx, peer, accSettings)
switch {
case expired:
if err := checkAuth(ctx, login.UserID, peer); err != nil {
// when the client sends a login request with a JWT which is used to get the user ID,
// it means that the client has already checked if it needs login and had been through the SSO flow
// so, we can skip this check and directly proceed with the login
if login.UserID == "" {
err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
if err != nil {
return nil, nil, nil, err
}
isWriteLock = true
log.WithContext(ctx).Debugf("peer login expired, acquiring write lock")

case peer.UpdateMetaIfNew(login.Meta):
isWriteLock = true
log.WithContext(ctx).Debugf("peer changed meta, acquiring write lock")

default:
isWriteLock = false
log.WithContext(ctx).Debugf("peer meta is the same, acquiring read lock")
}

var unlock func()

if isWriteLock {
unlock = am.Store.AcquireAccountWriteLock(ctx, accountID)
} else {
unlock = am.Store.AcquireAccountReadLock(ctx, accountID)
}
unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
defer func() {
if unlock != nil {
unlock()
Expand All @@ -660,7 +633,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
return nil, nil, nil, err
}

peer, err = account.FindPeerByPubKey(login.WireGuardPubKey)
peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
if err != nil {
return nil, nil, nil, status.NewPeerNotRegisteredError()
}
Expand All @@ -671,67 +644,86 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
}

// this flag prevents unnecessary calls to the persistent store.
shouldStoreAccount := false
shouldStorePeer := false
updateRemotePeers := false
if peerLoginExpired(ctx, peer, account.Settings) {
err = checkAuth(ctx, login.UserID, peer)
err = am.handleExpiredPeer(ctx, login, account, peer)
if err != nil {
return nil, nil, nil, err
}
// If peer was expired before and if it reached this point, it is re-authenticated.
// UserID is present, meaning that JWT validation passed successfully in the API layer.
updatePeerLastLogin(peer, account)
updateRemotePeers = true
shouldStoreAccount = true

// sync user last login with peer last login
user, err := account.FindUser(login.UserID)
if err != nil {
return nil, nil, nil, status.Errorf(status.Internal, "couldn't find user")
}
user.updateLastLogin(peer.LastLogin)

am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
shouldStorePeer = true
}

isRequiresApproval, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
if err != nil {
return nil, nil, nil, err
}

peer, updated := updatePeerMeta(peer, login.Meta, account)
if updated {
shouldStoreAccount = true
shouldStorePeer = true
}

peer, err = am.checkAndUpdatePeerSSHKey(ctx, peer, account, login.SSHKey)
if err != nil {
return nil, nil, nil, err
if peer.SSHKey != login.SSHKey {
peer.SSHKey = login.SSHKey
shouldStorePeer = true
}

if shouldStoreAccount {
if !isWriteLock {
log.WithContext(ctx).Errorf("account %s should be stored but is not write locked", accountID)
return nil, nil, nil, status.Errorf(status.Internal, "account should be stored but is not write locked")
}
err = am.Store.SaveAccount(ctx, account)
if shouldStorePeer {
err = am.Store.SavePeer(ctx, accountID, peer)
if err != nil {
return nil, nil, nil, err
}
}

unlock()
unlock = nil

if updateRemotePeers || isStatusChanged {
am.updateAccountPeers(ctx, account)
}

return am.getValidatedPeerWithMap(ctx, isRequiresApproval, account, peer)
}

// checkIFPeerNeedsLoginWithoutLock checks if the peer needs login without acquiring the account lock. The check validate if the peer was not added via SSO
// and if the peer login is expired.
// The NetBird client doesn't have a way to check if the peer needs login besides sending a login request
// with no JWT token and usually no setup-key. As the client can send up to two login request to check if it is expired
// and before starting the engine, we do the checks without an account lock to avoid piling up requests.
func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, accountID string, login PeerLogin) error {
peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
if err != nil {
return err
}

// if the peer was not added with SSO login we can exit early because peers activated with setup-key
// doesn't expire, and we avoid extra databases calls.
if !peer.AddedWithSSOLogin() {
return nil
}

settings, err := am.Store.GetAccountSettings(ctx, accountID)
if err != nil {
return err
}

if peerLoginExpired(ctx, peer, settings) {
return status.NewPeerLoginExpiredError()
}

return nil
}

func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, account *Account, peer *nbpeer.Peer) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
var postureChecks []*posture.Checks

if isRequiresApproval {
emptyMap := &NetworkMap{
Network: account.Network.Copy(),
}
return peer, emptyMap, postureChecks, nil
return peer, emptyMap, nil, nil
}

approvedPeersMap, err := am.GetValidatedPeers(account)
Expand All @@ -743,6 +735,30 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
return peer, account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, approvedPeersMap), postureChecks, nil
}

func (am *DefaultAccountManager) handleExpiredPeer(ctx context.Context, login PeerLogin, account *Account, peer *nbpeer.Peer) error {
err := checkAuth(ctx, login.UserID, peer)
if err != nil {
return err
}
// If peer was expired before and if it reached this point, it is re-authenticated.
// UserID is present, meaning that JWT validation passed successfully in the API layer.
updatePeerLastLogin(peer, account)

// sync user last login with peer last login
user, err := account.FindUser(login.UserID)
if err != nil {
return status.Errorf(status.Internal, "couldn't find user")
}

err = am.Store.SaveUserLastLogin(account.Id, user.Id, peer.LastLogin)
if err != nil {
return err
}

am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
return nil
}

func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
if peer.AddedWithSSOLogin() {
user, err := account.FindUser(peer.UserID)
Expand All @@ -759,11 +775,11 @@ func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
func checkAuth(ctx context.Context, loginUserID string, peer *nbpeer.Peer) error {
if loginUserID == "" {
// absence of a user ID indicates that JWT wasn't provided.
return status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
return status.NewPeerLoginExpiredError()
}
if peer.UserID != loginUserID {
log.WithContext(ctx).Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, loginUserID)
return status.Errorf(status.Unauthenticated, "can't login")
return status.Errorf(status.Unauthenticated, "can't login with this credentials")
}
return nil
}
Expand All @@ -783,31 +799,6 @@ func updatePeerLastLogin(peer *nbpeer.Peer, account *Account) {
account.UpdatePeer(peer)
}

func (am *DefaultAccountManager) checkAndUpdatePeerSSHKey(ctx context.Context, peer *nbpeer.Peer, account *Account, newSSHKey string) (*nbpeer.Peer, error) {
if len(newSSHKey) == 0 {
log.WithContext(ctx).Debugf("no new SSH key provided for peer %s, skipping update", peer.ID)
return peer, nil
}

if peer.SSHKey == newSSHKey {
log.WithContext(ctx).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(ctx, account)
if err != nil {
return nil, err
}

// trigger network map update
am.updateAccountPeers(ctx, account)

return peer, nil
}

// UpdatePeerSSHKey updates peer's public SSH key
func (am *DefaultAccountManager) UpdatePeerSSHKey(ctx context.Context, peerID string, sshKey string) error {
if sshKey == "" {
Expand Down
5 changes: 5 additions & 0 deletions management/server/status/error.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,3 +95,8 @@ func NewUserNotFoundError(userKey string) error {
func NewPeerNotRegisteredError() error {
return Errorf(Unauthenticated, "peer is not registered")
}

// NewPeerLoginExpiredError creates a new Error with PermissionDenied type for an expired peer
func NewPeerLoginExpiredError() error {
return Errorf(PermissionDenied, "peer login has expired, please log in once more")
}

0 comments on commit da39c8b

Please sign in to comment.