Skip to content

Commit

Permalink
Update account peers on login on meta change (#2991)
Browse files Browse the repository at this point in the history
* Update account peers on login on meta change

* Factor out LoginPeer peer not found handling
  • Loading branch information
mohamed-essam authored Dec 5, 2024
1 parent e67fe89 commit 713e320
Showing 1 changed file with 25 additions and 16 deletions.
41 changes: 25 additions & 16 deletions management/server/peer.go
Original file line number Diff line number Diff line change
Expand Up @@ -710,26 +710,30 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, validPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
}

func (am *DefaultAccountManager) handlePeerLoginNotFound(ctx context.Context, login PeerLogin, err error) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
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.
newPeer := &nbpeer.Peer{
Key: login.WireGuardPubKey,
Meta: login.Meta,
SSHKey: login.SSHKey,
Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
}

return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
}

log.WithContext(ctx).Errorf("failed while logging in peer %s: %v", login.WireGuardPubKey, err)
return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
}

// 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(ctx context.Context, login PeerLogin) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, 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.
newPeer := &nbpeer.Peer{
Key: login.WireGuardPubKey,
Meta: login.Meta,
SSHKey: login.SSHKey,
Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
}

return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
}

log.WithContext(ctx).Errorf("failed while logging in peer %s: %v", login.WireGuardPubKey, err)
return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
return am.handlePeerLoginNotFound(ctx, login, err)
}

// when the client sends a login request with a JWT which is used to get the user ID,
Expand Down Expand Up @@ -828,7 +832,12 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
return nil, nil, nil, err
}

if updateRemotePeers || isStatusChanged {
postureChecks, err := am.getPeerPostureChecks(account, peer.ID)
if err != nil {
return nil, nil, nil, err
}

if updateRemotePeers || isStatusChanged || (updated && len(postureChecks) > 0) {
am.updateAccountPeers(ctx, accountID)
}

Expand Down

0 comments on commit 713e320

Please sign in to comment.