Refactor login with store.SavePeer (#2334)

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.
This commit is contained in:
Maycon Santos 2024-07-29 13:30:27 +02:00 committed by GitHub
parent 7321046cd6
commit da39c8bbca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 100 additions and 104 deletions

View File

@ -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 {
@ -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)
@ -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)
@ -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()
@ -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()
}
@ -671,53 +644,39 @@ 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
@ -725,13 +684,46 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
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)
@ -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)
@ -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
}
@ -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 == "" {

View File

@ -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")
}