refactor AddPeer

management/server/peer.go

@@ -561,6 +561,99 @@ func (am *DefaultAccountManager) GetPeerNetwork(ctx context.Context, peerID stri
 	return account.Network.Copy(), err
 }
 
+type peerAddAuthConfig struct {
+	AccountID           string
+	SetupKeyID          string
+	SetupKeyName        string
+	GroupsToAdd         []string
+	AllowExtraDNSLabels bool
+	Ephemeral           bool
+}
+
+func (am *DefaultAccountManager) processPeerAddAuth(ctx context.Context, accountID, userID, encodedHashedKey string, peer *nbpeer.Peer, temporary, addedByUser, addedBySetupKey bool, opEvent *activity.Event) (*peerAddAuthConfig, error) {
+	config := &peerAddAuthConfig{
+		AccountID: accountID,
+		Ephemeral: peer.Ephemeral,
+	}
+
+	switch {
+	case addedByUser:
+		if err := am.handleUserAddedPeer(ctx, accountID, userID, temporary, opEvent, config); err != nil {
+			return nil, err
+		}
+	case addedBySetupKey:
+		if err := am.handleSetupKeyAddedPeer(ctx, encodedHashedKey, peer, opEvent, config); err != nil {
+			return nil, err
+		}
+	default:
+		if peer.ProxyMeta.Embedded {
+			log.WithContext(ctx).Debugf("adding peer for proxy embedded, accountID: %s", accountID)
+		} else {
+			log.WithContext(ctx).Warnf("adding peer without setup key or userID, accountID: %s", accountID)
+		}
+	}
+
+	opEvent.AccountID = config.AccountID
+	if temporary {
+		config.Ephemeral = true
+	}
+
+	return config, nil
+}
+
+func (am *DefaultAccountManager) handleUserAddedPeer(ctx context.Context, accountID, userID string, temporary bool, opEvent *activity.Event, config *peerAddAuthConfig) error {
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthNone, userID)
+	if err != nil {
+		return status.Errorf(status.NotFound, "failed adding new peer: user not found")
+	}
+	if user.PendingApproval {
+		return status.Errorf(status.PermissionDenied, "user pending approval cannot add peers")
+	}
+
+	if temporary {
+		allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Create)
+		if err != nil {
+			return status.NewPermissionValidationError(err)
+		}
+		if !allowed {
+			return status.NewPermissionDeniedError()
+		}
+	} else {
+		config.AccountID = user.AccountID
+		config.GroupsToAdd = user.AutoGroups
+	}
+
+	opEvent.InitiatorID = userID
+	opEvent.Activity = activity.PeerAddedByUser
+	return nil
+}
+
+func (am *DefaultAccountManager) handleSetupKeyAddedPeer(ctx context.Context, encodedHashedKey string, peer *nbpeer.Peer, opEvent *activity.Event, config *peerAddAuthConfig) error {
+	sk, err := am.Store.GetSetupKeyBySecret(ctx, store.LockingStrengthNone, encodedHashedKey)
+	if err != nil {
+		return status.Errorf(status.NotFound, "couldn't add peer: setup key is invalid")
+	}
+
+	if !sk.IsValid() {
+		return status.Errorf(status.NotFound, "couldn't add peer: setup key is invalid")
+	}
+
+	if !sk.AllowExtraDNSLabels && len(peer.ExtraDNSLabels) > 0 {
+		return status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key doesn't allow extra DNS labels")
+	}
+
+	opEvent.InitiatorID = sk.Id
+	opEvent.Activity = activity.PeerAddedWithSetupKey
+	config.GroupsToAdd = sk.AutoGroups
+	config.Ephemeral = sk.Ephemeral
+	config.SetupKeyID = sk.Id
+	config.SetupKeyName = sk.Name
+	config.AllowExtraDNSLabels = sk.AllowExtraDNSLabels
+	config.AccountID = sk.AccountID
+
+	return nil
+}
+
 // AddPeer adds a new peer to the Store.
 // 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.
@@ -596,70 +689,12 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 
 	var newPeer *nbpeer.Peer
 
-	var setupKeyID string
+	peerAddConfig, err := am.processPeerAddAuth(ctx, accountID, userID, encodedHashedKey, peer, temporary, addedByUser, addedBySetupKey, opEvent)
-	var setupKeyName string
+	if err != nil {
-	var groupsToAdd []string
+		return nil, nil, nil, err
-	var allowExtraDNSLabels bool
-	ephemeral := peer.Ephemeral
-	switch {
-	case addedByUser:
-		user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthNone, userID)
-		if err != nil {
-			return nil, nil, nil, status.Errorf(status.NotFound, "failed adding new peer: user not found")
-		}
-		if user.PendingApproval {
-			return nil, nil, nil, status.Errorf(status.PermissionDenied, "user pending approval cannot add peers")
-		}
-		if temporary {
-			allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Create)
-			if err != nil {
-				return nil, nil, nil, status.NewPermissionValidationError(err)
-			}
-
-			if !allowed {
-				return nil, nil, nil, status.NewPermissionDeniedError()
-			}
-		} else {
-			accountID = user.AccountID
-			groupsToAdd = user.AutoGroups
-		}
-		opEvent.InitiatorID = userID
-		opEvent.Activity = activity.PeerAddedByUser
-	case addedBySetupKey:
-		// Validate the setup key
-		sk, err := am.Store.GetSetupKeyBySecret(ctx, store.LockingStrengthNone, encodedHashedKey)
-		if err != nil {
-			return nil, nil, nil, status.Errorf(status.NotFound, "couldn't add peer: setup key is invalid")
-		}
-
-		// we will check key twice for early return
-		if !sk.IsValid() {
-			return nil, nil, nil, status.Errorf(status.NotFound, "couldn't add peer: setup key is invalid")
-		}
-
-		opEvent.InitiatorID = sk.Id
-		opEvent.Activity = activity.PeerAddedWithSetupKey
-		groupsToAdd = sk.AutoGroups
-		ephemeral = sk.Ephemeral
-		setupKeyID = sk.Id
-		setupKeyName = sk.Name
-		allowExtraDNSLabels = sk.AllowExtraDNSLabels
-		accountID = sk.AccountID
-		if !sk.AllowExtraDNSLabels && len(peer.ExtraDNSLabels) > 0 {
-			return nil, nil, nil, status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key doesn't allow extra DNS labels")
-		}
-	default:
-		if peer.ProxyMeta.Embedded {
-			log.WithContext(ctx).Debugf("adding peer for proxy embedded, accountID: %s", accountID)
-		} else {
-			log.WithContext(ctx).Warnf("adding peer without setup key or userID, accountID: %s", accountID)
-		}
-	}
-	opEvent.AccountID = accountID
-
-	if temporary {
-		ephemeral = true
 	}
+	accountID = peerAddConfig.AccountID
+	ephemeral := peerAddConfig.Ephemeral
 
 	if (strings.ToLower(peer.Meta.Hostname) == "iphone" || strings.ToLower(peer.Meta.Hostname) == "ipad") && userID != "" {
 		if am.idpManager != nil {
@@ -693,7 +728,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 		Location:                    peer.Location,
 		InactivityExpirationEnabled: addedByUser && !temporary,
 		ExtraDNSLabels:              peer.ExtraDNSLabels,
-		AllowExtraDNSLabels:         allowExtraDNSLabels,
+		AllowExtraDNSLabels:         peerAddConfig.AllowExtraDNSLabels,
 	}
 	settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
 	if err != nil {
@@ -711,7 +746,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 		}
 	}
 
-	newPeer = am.integratedPeerValidator.PreparePeer(ctx, accountID, newPeer, groupsToAdd, settings.Extra, temporary)
+	newPeer = am.integratedPeerValidator.PreparePeer(ctx, accountID, newPeer, peerAddConfig.GroupsToAdd, settings.Extra, temporary)
 
 	network, err := am.Store.GetAccountNetwork(ctx, store.LockingStrengthNone, accountID)
 	if err != nil {
@@ -747,8 +782,8 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 				return err
 			}
 
-			if len(groupsToAdd) > 0 {
+			if len(peerAddConfig.GroupsToAdd) > 0 {
-				for _, g := range groupsToAdd {
+				for _, g := range peerAddConfig.GroupsToAdd {
 					err = transaction.AddPeerToGroup(ctx, newPeer.AccountID, newPeer.ID, g)
 					if err != nil {
 						return err
@@ -780,7 +815,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 					return status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key is invalid")
 				}
 
-				err = transaction.IncrementSetupKeyUsage(ctx, setupKeyID)
+				err = transaction.IncrementSetupKeyUsage(ctx, peerAddConfig.SetupKeyID)
 				if err != nil {
 					return fmt.Errorf("failed to increment setup key usage: %w", err)
 				}
@@ -821,7 +856,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 	opEvent.TargetID = newPeer.ID
 	opEvent.Meta = newPeer.EventMeta(am.networkMapController.GetDNSDomain(settings))
 	if !addedByUser {
-		opEvent.Meta["setup_key_name"] = setupKeyName
+		opEvent.Meta["setup_key_name"] = peerAddConfig.SetupKeyName
 	}
 
 	am.StoreEvent(ctx, opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)

management/server/peer_test.go

@@ -2489,3 +2489,252 @@ func TestLoginPeer_ApprovedUserCanLogin(t *testing.T) {
 	_, _, _, err = manager.LoginPeer(context.Background(), login)
 	require.NoError(t, err, "Regular user should be able to login peers")
 }
+
+func TestHandleUserAddedPeer(t *testing.T) {
+	manager, _, err := createManager(t)
+	require.NoError(t, err)
+
+	account := newAccountWithId(context.Background(), "test-account", "owner", "", "", "", false)
+	err = manager.Store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	t.Run("regular user can add peer", func(t *testing.T) {
+		regularUser := types.NewRegularUser("regular-user-1", "", "")
+		regularUser.AccountID = account.Id
+		regularUser.AutoGroups = []string{"group1", "group2"}
+		err = manager.Store.SaveUser(context.Background(), regularUser)
+		require.NoError(t, err)
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+
+		err = manager.handleUserAddedPeer(context.Background(), account.Id, regularUser.Id, false, opEvent, config)
+		require.NoError(t, err)
+		assert.Equal(t, account.Id, config.AccountID)
+		assert.Equal(t, regularUser.AutoGroups, config.GroupsToAdd)
+		assert.Equal(t, regularUser.Id, opEvent.InitiatorID)
+		assert.Equal(t, activity.PeerAddedByUser, opEvent.Activity)
+	})
+
+	t.Run("pending approval user cannot add peer", func(t *testing.T) {
+		pendingUser := types.NewRegularUser("pending-user", "", "")
+		pendingUser.AccountID = account.Id
+		pendingUser.PendingApproval = true
+		err = manager.Store.SaveUser(context.Background(), pendingUser)
+		require.NoError(t, err)
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+
+		err = manager.handleUserAddedPeer(context.Background(), account.Id, pendingUser.Id, false, opEvent, config)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "user pending approval cannot add peers")
+	})
+
+	t.Run("user not found", func(t *testing.T) {
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+
+		err = manager.handleUserAddedPeer(context.Background(), account.Id, "non-existent-user", false, opEvent, config)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "user not found")
+	})
+
+	t.Run("temporary peer requires permissions", func(t *testing.T) {
+		regularUser := types.NewRegularUser("regular-user-2", "", "")
+		regularUser.AccountID = account.Id
+		err = manager.Store.SaveUser(context.Background(), regularUser)
+		require.NoError(t, err)
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+
+		// Should fail because user doesn't have permissions for temporary peers
+		err = manager.handleUserAddedPeer(context.Background(), account.Id, regularUser.Id, true, opEvent, config)
+		require.Error(t, err)
+	})
+}
+
+func TestHandleSetupKeyAddedPeer(t *testing.T) {
+	manager, _, err := createManager(t)
+	require.NoError(t, err)
+
+	account := newAccountWithId(context.Background(), "test-account", "owner", "", "", "", false)
+	err = manager.Store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	// Create admin user for setup key creation
+	adminUser := types.NewAdminUser("admin-user")
+	adminUser.AccountID = account.Id
+	err = manager.Store.SaveUser(context.Background(), adminUser)
+	require.NoError(t, err)
+
+	t.Run("valid setup key", func(t *testing.T) {
+		setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", types.SetupKeyReusable, time.Hour, []string{}, 0, adminUser.Id, false, false)
+		require.NoError(t, err)
+
+		upperKey := strings.ToUpper(setupKey.Key)
+		hashedKey := sha256.Sum256([]byte(upperKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+		peer := &nbpeer.Peer{ExtraDNSLabels: []string{}}
+
+		err = manager.handleSetupKeyAddedPeer(context.Background(), encodedHashedKey, peer, opEvent, config)
+		require.NoError(t, err)
+		assert.Equal(t, setupKey.Id, config.SetupKeyID)
+		assert.Equal(t, setupKey.Name, config.SetupKeyName)
+		assert.Equal(t, setupKey.AutoGroups, config.GroupsToAdd)
+		assert.Equal(t, setupKey.Ephemeral, config.Ephemeral)
+		assert.Equal(t, setupKey.Id, opEvent.InitiatorID)
+		assert.Equal(t, activity.PeerAddedWithSetupKey, opEvent.Activity)
+	})
+
+	t.Run("invalid setup key", func(t *testing.T) {
+		invalidKey := "invalid-key"
+		hashedKey := sha256.Sum256([]byte(invalidKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+		peer := &nbpeer.Peer{}
+
+		err = manager.handleSetupKeyAddedPeer(context.Background(), encodedHashedKey, peer, opEvent, config)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "setup key is invalid")
+	})
+
+	t.Run("expired setup key", func(t *testing.T) {
+		setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "expired-key", types.SetupKeyReusable, time.Millisecond, []string{}, 0, adminUser.Id, false, false)
+		require.NoError(t, err)
+
+		// Wait for key to expire
+		time.Sleep(10 * time.Millisecond)
+
+		upperKey := strings.ToUpper(setupKey.Key)
+		hashedKey := sha256.Sum256([]byte(upperKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+		peer := &nbpeer.Peer{}
+
+		err = manager.handleSetupKeyAddedPeer(context.Background(), encodedHashedKey, peer, opEvent, config)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "setup key is invalid")
+	})
+
+	t.Run("extra DNS labels not allowed", func(t *testing.T) {
+		setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "no-dns-key", types.SetupKeyReusable, time.Hour, []string{}, 0, adminUser.Id, false, false)
+		require.NoError(t, err)
+
+		upperKey := strings.ToUpper(setupKey.Key)
+		hashedKey := sha256.Sum256([]byte(upperKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+		peer := &nbpeer.Peer{ExtraDNSLabels: []string{"custom.label"}}
+
+		err = manager.handleSetupKeyAddedPeer(context.Background(), encodedHashedKey, peer, opEvent, config)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "doesn't allow extra DNS labels")
+	})
+
+	t.Run("extra DNS labels allowed", func(t *testing.T) {
+		setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "dns-key", types.SetupKeyReusable, time.Hour, []string{}, 0, adminUser.Id, false, true)
+		require.NoError(t, err)
+
+		upperKey := strings.ToUpper(setupKey.Key)
+		hashedKey := sha256.Sum256([]byte(upperKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{}
+		config := &peerAddAuthConfig{}
+		peer := &nbpeer.Peer{ExtraDNSLabels: []string{"custom.label"}}
+
+		err = manager.handleSetupKeyAddedPeer(context.Background(), encodedHashedKey, peer, opEvent, config)
+		require.NoError(t, err)
+		assert.True(t, config.AllowExtraDNSLabels)
+	})
+}
+
+func TestProcessPeerAddAuth(t *testing.T) {
+	manager, _, err := createManager(t)
+	require.NoError(t, err)
+
+	account := newAccountWithId(context.Background(), "test-account", "owner", "", "", "", false)
+	err = manager.Store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	adminUser := types.NewAdminUser("admin")
+	adminUser.AccountID = account.Id
+	err = manager.Store.SaveUser(context.Background(), adminUser)
+	require.NoError(t, err)
+
+	t.Run("user authentication flow", func(t *testing.T) {
+		regularUser := types.NewRegularUser("user-auth-test", "", "")
+		regularUser.AccountID = account.Id
+		regularUser.AutoGroups = []string{"group1"}
+		err = manager.Store.SaveUser(context.Background(), regularUser)
+		require.NoError(t, err)
+
+		opEvent := &activity.Event{Timestamp: time.Now()}
+		peer := &nbpeer.Peer{Ephemeral: false}
+
+		config, err := manager.processPeerAddAuth(context.Background(), account.Id, regularUser.Id, "", peer, false, true, false, opEvent)
+		require.NoError(t, err)
+		assert.Equal(t, account.Id, config.AccountID)
+		assert.False(t, config.Ephemeral)
+		assert.Equal(t, regularUser.AutoGroups, config.GroupsToAdd)
+		assert.Equal(t, account.Id, opEvent.AccountID)
+	})
+
+	t.Run("setup key authentication flow", func(t *testing.T) {
+		setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "auth-test-key", types.SetupKeyReusable, time.Hour, []string{}, 0, adminUser.Id, true, false)
+		require.NoError(t, err)
+
+		upperKey := strings.ToUpper(setupKey.Key)
+		hashedKey := sha256.Sum256([]byte(upperKey))
+		encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+		opEvent := &activity.Event{Timestamp: time.Now()}
+		peer := &nbpeer.Peer{Ephemeral: false}
+
+		config, err := manager.processPeerAddAuth(context.Background(), account.Id, "", encodedHashedKey, peer, false, false, true, opEvent)
+		require.NoError(t, err)
+		assert.Equal(t, account.Id, config.AccountID)
+		assert.True(t, config.Ephemeral) // setupKey.Ephemeral is true
+		assert.Equal(t, setupKey.AutoGroups, config.GroupsToAdd)
+		assert.Equal(t, account.Id, opEvent.AccountID)
+	})
+
+	t.Run("temporary flag overrides ephemeral", func(t *testing.T) {
+		regularUser := types.NewRegularUser("temp-user", "", "")
+		regularUser.AccountID = account.Id
+		err = manager.Store.SaveUser(context.Background(), regularUser)
+		require.NoError(t, err)
+
+		opEvent := &activity.Event{Timestamp: time.Now()}
+		peer := &nbpeer.Peer{Ephemeral: false}
+
+		config, err := manager.processPeerAddAuth(context.Background(), account.Id, regularUser.Id, "", peer, true, true, false, opEvent)
+		require.Error(t, err) // Will fail permission check but that's expected
+		_ = config            // avoid unused warning
+	})
+
+	t.Run("proxy embedded peer (no auth)", func(t *testing.T) {
+		opEvent := &activity.Event{Timestamp: time.Now()}
+		peer := &nbpeer.Peer{
+			Ephemeral: false,
+			ProxyMeta: nbpeer.ProxyMeta{Embedded: true},
+		}
+
+		config, err := manager.processPeerAddAuth(context.Background(), account.Id, "", "", peer, false, false, false, opEvent)
+		require.NoError(t, err)
+		assert.Equal(t, account.Id, config.AccountID)
+		assert.False(t, config.Ephemeral)
+		assert.Empty(t, config.GroupsToAdd)
+	})
+}