[management] Keep embedded IdP deployments on a single account (#7380)

combined/cmd/root.go

@@ -365,7 +365,6 @@ func setupServerHooks(servers *serverInstances, cfg *CombinedConfig) {
 			})
 		}
 	}
-
 }
 
 func startServers(wg *sync.WaitGroup, srv *relayServer.Server, httpHealthcheck *healthcheck.Server, stunServer *stun.Server, metricsServer *sharedMetrics.Metrics) {
@@ -539,7 +538,7 @@ func createManagementServer(cfg *CombinedConfig, mgmtConfig *nbconfig.Config) (m
 		&mgmtServer.Config{
 			NbConfig:                mgmtConfig,
 			DNSDomain:               "",
-			MgmtSingleAccModeDomain: "",
+			MgmtSingleAccModeDomain: mgmtServer.DefaultSelfHostedDomain,
 			AutoResolveDomains:      true,
 			MgmtPort:                mgmtPort,
 			MgmtMetricsPort:         cfg.Server.MetricsPort,
@@ -554,7 +553,7 @@ func createManagementServer(cfg *CombinedConfig, mgmtConfig *nbconfig.Config) (m
 }
 
 // createCombinedHandler creates an HTTP handler that multiplexes Management, Signal (via wsproxy), and Relay WebSocket traffic
-func createCombinedHandler(grpcServer *grpc.Server, httpHandler http.Handler, idpHandler http.Handler, relaySrv *relayServer.Server, meter metric.Meter, cfg *CombinedConfig) http.Handler {
+func createCombinedHandler(grpcServer *grpc.Server, httpHandler, idpHandler http.Handler, relaySrv *relayServer.Server, meter metric.Meter, cfg *CombinedConfig) http.Handler {
 	wsProxy := wsproxyserver.New(grpcServer, wsproxyserver.WithOTelMeter(meter))
 
 	var relayAcceptFn func(conn listener.Conn)

management/cmd/management.go

@@ -236,6 +236,9 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config) error {
 		// Embedded IdP requires single account mode - multiple account mode is not supported
 		return fmt.Errorf("embedded IdP requires single account mode; multiple account mode is not supported with embedded IdP. Please remove --disable-single-account-mode flag")
 	}
+	if mgmtSingleAccModeDomain == "" {
+		return fmt.Errorf("embedded IdP requires single account mode; --single-account-mode-domain must not be empty")
+	}
 	// Enable user deletion from IDP by default if EmbeddedIdP is enabled
 	userDeleteFromIDPEnabled = true
 

management/cmd/management_test.go

@@ -5,8 +5,12 @@ import (
 	"os"
 	"testing"
 
-	"github.com/netbirdio/netbird/shared/management/grpc"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/idp"
+	"github.com/netbirdio/netbird/shared/management/grpc"
 )
 
 const (
@@ -60,6 +64,22 @@ func Test_LoadMgmtConfig_Empty(t *testing.T) {
 	assert.Nil(t, cfg.PerAccountHighestSupportedSyncMessageVersion)
 }
 
+func TestApplyEmbeddedIdPConfigRequiresSingleAccountDomain(t *testing.T) {
+	previousDomain := mgmtSingleAccModeDomain
+	previousDisabled := disableSingleAccMode
+	t.Cleanup(func() {
+		mgmtSingleAccModeDomain = previousDomain
+		disableSingleAccMode = previousDisabled
+	})
+
+	mgmtSingleAccModeDomain = ""
+	disableSingleAccMode = false
+	cfg := &nbconfig.Config{
+		EmbeddedIdP: &idp.EmbeddedIdPConfig{Enabled: true},
+	}
+	require.ErrorContains(t, ApplyEmbeddedIdPConfig(context.Background(), cfg), "embedded IdP requires single account mode")
+}
+
 func createConfig(config string) (string, error) {
 	tmpfile, err := os.CreateTemp("", "config.json")
 	if err != nil {

management/server/account.go

@@ -14,10 +14,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
-	"github.com/netbirdio/netbird/management/server/job"
-	"github.com/netbirdio/netbird/shared/auth"
-
 	cacheStore "github.com/eko/gocache/lib/v4/store"
 	"github.com/eko/gocache/store/redis/v4"
 	"github.com/rs/xid"
@@ -29,6 +25,7 @@ import (
 	"github.com/netbirdio/netbird/formatter/hook"
 	"github.com/netbirdio/netbird/idp/dex"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
@@ -39,6 +36,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/job"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
@@ -50,6 +48,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/management/server/util"
 	"github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/shared/auth"
 	nbdomain "github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/netbirdio/netbird/shared/management/status"
@@ -238,6 +237,10 @@ func BuildManager(
 		log.WithContext(ctx).Error(err)
 	}
 
+	if IsEmbeddedIdp(idpManager) && accountsCounter > 1 {
+		log.WithContext(ctx).Warnf("embedded IdP requires a single account, found %d", accountsCounter)
+	}
+
 	// enable single account mode only if configured by user and number of existing accounts is not grater than 1
 	am.singleAccountMode = singleAccountModeDomain != "" && accountsCounter <= 1
 	if am.singleAccountMode {
@@ -1592,7 +1595,10 @@ func (am *DefaultAccountManager) updateUserAuthWithSingleMode(ctx context.Contex
 	if err != nil {
 		return err
 	}
-	userAuth.Domain = domain
+	// Keep the configured single account domain when the existing account has none
+	if domain != "" {
+		userAuth.Domain = domain
+	}
 
 	log.WithContext(ctx).Debugf("overriding JWT Domain and DomainCategory claims since single account mode is enabled")
 	return nil
@@ -1837,6 +1843,7 @@ func (am *DefaultAccountManager) getAccountIDWithAuthorizationClaims(ctx context
 
 	return am.addNewPrivateAccount(ctx, domainAccountID, userAuth)
 }
+
 func (am *DefaultAccountManager) getPrivateDomainWithGlobalLock(ctx context.Context, domain string) (string, context.CancelFunc, error) {
 	domainAccountID, err := am.Store.GetAccountIDByPrivateDomain(ctx, store.LockingStrengthNone, domain)
 	if handleNotFound(err) != nil {

management/server/identity_provider_test.go

@@ -10,9 +10,9 @@ import (
 	"testing"
 	"time"
 
-	"go.uber.org/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
@@ -34,6 +34,20 @@ import (
 
 func createManagerWithEmbeddedIdP(t testing.TB) (*DefaultAccountManager, *update_channel.PeersUpdateManager, error) {
 	t.Helper()
+	return createManagerWithEmbeddedIdPMode(t, "netbird.selfhosted")
+}
+
+func createManagerWithEmbeddedIdPMode(t testing.TB, singleAccountModeDomain string) (*DefaultAccountManager, *update_channel.PeersUpdateManager, error) {
+	t.Helper()
+	return createManagerWithEmbeddedIdPModeAndSetup(t, singleAccountModeDomain, nil)
+}
+
+func createManagerWithEmbeddedIdPModeAndSetup(
+	t testing.TB,
+	singleAccountModeDomain string,
+	setupStore func(context.Context, store.Store) error,
+) (*DefaultAccountManager, *update_channel.PeersUpdateManager, error) {
+	t.Helper()
 
 	ctx := context.Background()
 
@@ -43,6 +57,11 @@ func createManagerWithEmbeddedIdP(t testing.TB) (*DefaultAccountManager, *update
 		return nil, nil, err
 	}
 	t.Cleanup(cleanUp)
+	if setupStore != nil {
+		if err := setupStore(ctx, testStore); err != nil {
+			return nil, nil, err
+		}
+	}
 
 	// Create embedded IdP manager
 	embeddedConfig := &idp.EmbeddedIdPConfig{
@@ -93,7 +112,7 @@ func createManagerWithEmbeddedIdP(t testing.TB) (*DefaultAccountManager, *update
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, testStore)
 	networkMapController := controller.NewController(ctx, testStore, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(testStore, peersManager), &config.Config{}, nil)
-	manager, err := BuildManager(ctx, &config.Config{}, testStore, networkMapController, job.NewJobManager(nil, testStore, peersManager), idpManager, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false, cacheStore)
+	manager, err := BuildManager(ctx, &config.Config{}, testStore, networkMapController, job.NewJobManager(nil, testStore, peersManager), idpManager, singleAccountModeDomain, eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false, cacheStore)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -196,6 +215,23 @@ func TestDefaultAccountManager_GetIdentityProvider_NotFound(t *testing.T) {
 	assert.Contains(t, err.Error(), "not found")
 }
 
+func TestUpdateUserAuthWithSingleModeKeepsConfiguredDomain(t *testing.T) {
+	ctx := context.Background()
+	manager, _, err := createManagerWithEmbeddedIdPModeAndSetup(t, "netbird.selfhosted", func(ctx context.Context, testStore store.Store) error {
+		// An account with no domain, as left behind by an IdP that emitted no domain claims.
+		return testStore.SaveAccount(ctx, newAccountWithId(ctx, "account-1", "user-1", "", "", "", false))
+	})
+	require.NoError(t, err)
+	require.True(t, manager.singleAccountMode)
+
+	userAuth := auth.UserAuth{UserId: "user-2"}
+	require.NoError(t, manager.updateUserAuthWithSingleMode(ctx, &userAuth))
+
+	assert.Equal(t, "netbird.selfhosted", userAuth.Domain,
+		"An empty account domain must not clear the configured single account domain")
+	assert.Equal(t, types.PrivateCategory, userAuth.DomainCategory)
+}
+
 func TestDefaultAccountManager_UpdateIdentityProvider_Validation(t *testing.T) {
 	manager, _, err := createManager(t)
 	require.NoError(t, err)

management/server/idp/migration/migration.go

@@ -10,6 +10,8 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"regexp"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 
@@ -25,8 +27,10 @@ type Server interface {
 	EventStore() EventStore // may return nil
 }
 
-const idpSeedInfoKey = "IDP_SEED_INFO"
-const dryRunEnvKey = "NB_IDP_MIGRATION_DRY_RUN"
+const (
+	idpSeedInfoKey = "IDP_SEED_INFO"
+	dryRunEnvKey   = "NB_IDP_MIGRATION_DRY_RUN"
+)
 
 func isDryRun() bool {
 	return os.Getenv(dryRunEnvKey) == "true"
@@ -233,3 +237,163 @@ func PopulateUserInfo(s Server, idpManager idp.Manager, dryRun bool) error {
 
 	return nil
 }
+
+const DefaultSingleAccountDomain = "netbird.selfhosted"
+
+var (
+	ErrMultipleAccounts = errors.New("the embedded IdP supports a single account only")
+	ErrUnusableDomain   = errors.New("domain cannot be resolved in single account mode")
+	ErrDomainConflict   = errors.New("requested domain conflicts with the account domain")
+)
+
+var resolvableDomainRegexp = regexp.MustCompile(`^([a-z0-9]+(-[a-z0-9]+)*\.)+[a-z]{2,}$`)
+
+// RequireSingleAccount refuses to migrate an instance that holds more than one account.
+func RequireSingleAccount(s Server) error {
+	accountsCounter, err := s.Store().GetAccountsCounter(context.Background())
+	if err != nil {
+		return fmt.Errorf("failed to count accounts: %w", err)
+	}
+
+	if accountsCounter > 1 {
+		return errMultipleAccounts(accountsCounter)
+	}
+
+	return nil
+}
+
+func errMultipleAccounts(accountsCounter int64) error {
+	return fmt.Errorf("%w: this instance has %d accounts. Identity provider connectors are stored without "+
+		"an account scope, so every account would share and be able to manage the same connectors. "+
+		"Consolidate this instance to a single account, or keep using an external IdP, before migrating",
+		ErrMultipleAccounts, accountsCounter)
+}
+
+func NormalizeSingleAccountDomain(singleAccountDomain string) (string, error) {
+	if singleAccountDomain == "" {
+		singleAccountDomain = DefaultSingleAccountDomain
+	}
+
+	singleAccountDomain = strings.ToLower(singleAccountDomain)
+	if !resolvableDomainRegexp.MatchString(singleAccountDomain) {
+		return "", fmt.Errorf("%w: %q must contain at least one dot and only lowercase letters, digits and "+
+			"hyphens, otherwise users cannot join the existing account", ErrUnusableDomain, singleAccountDomain)
+	}
+
+	return singleAccountDomain, nil
+}
+
+// resolveAccountDomain picks the domain the account should end up with. The account keeps a usable
+// domain of its own, the configured one only fills a blank. Anything else is a conflict to report.
+func resolveAccountDomain(accountID, accountDomain, singleAccountDomain string, requested bool) (string, error) {
+	accountDomain = strings.ToLower(accountDomain)
+
+	if accountDomain == "" {
+		return singleAccountDomain, nil
+	}
+
+	if !resolvableDomainRegexp.MatchString(accountDomain) {
+		return "", fmt.Errorf("%w: account %s has domain %q, which must contain at least one dot and only "+
+			"lowercase letters, digits and hyphens. Correct the account domain before migrating",
+			ErrUnusableDomain, accountID, accountDomain)
+	}
+
+	if requested && accountDomain != singleAccountDomain {
+		return "", fmt.Errorf("%w: account %s already uses domain %q but %q was requested. Re-run without "+
+			"--single-account-mode-domain to keep %q, or correct the account domain first",
+			ErrDomainConflict, accountID, accountDomain, singleAccountDomain, accountDomain)
+	}
+
+	return accountDomain, nil
+}
+
+// EnsureSingleAccountDomain gives the remaining account the domain attributes single account mode
+// resolves against, so users can still join it after the migration.
+func EnsureSingleAccountDomain(s Server, singleAccountDomain string) error {
+	plan, err := planSingleAccountDomain(s, singleAccountDomain)
+	if err != nil {
+		return err
+	}
+	if plan.skip {
+		return nil
+	}
+
+	if isDryRun() {
+		log.Infof("[DRY RUN] would set account %s domain to %q, category to %q and mark it as the primary domain account "+
+			"(currently domain=%q primary=%v)", plan.accountID, plan.domain, types.PrivateCategory,
+			plan.currentDomain, plan.isPrimary)
+		return nil
+	}
+
+	if err := s.Store().UpdateAccountDomainAttributes(context.Background(), plan.accountID, plan.domain,
+		types.PrivateCategory, true); err != nil {
+		return fmt.Errorf("failed to update domain attributes of account %s: %w", plan.accountID, err)
+	}
+
+	log.Infof("account %s now resolves in single account mode with domain %q", plan.accountID, plan.domain)
+	return nil
+}
+
+// CheckSingleAccountDomain reports whether EnsureSingleAccountDomain would succeed, without writing.
+func CheckSingleAccountDomain(s Server, singleAccountDomain string) error {
+	_, err := planSingleAccountDomain(s, singleAccountDomain)
+	return err
+}
+
+type singleAccountDomainPlan struct {
+	accountID     string
+	domain        string
+	currentDomain string
+	isPrimary     bool
+	skip          bool
+}
+
+// planSingleAccountDomain decides what the account's domain attributes should become. It reads
+// only, so it can run both as a preflight and as the first half of the update.
+func planSingleAccountDomain(s Server, singleAccountDomain string) (singleAccountDomainPlan, error) {
+	ctx := context.Background()
+
+	// An empty value means the operator did not pick a domain, so the default is only a fallback.
+	requested := singleAccountDomain != ""
+
+	singleAccountDomain, err := NormalizeSingleAccountDomain(singleAccountDomain)
+	if err != nil {
+		return singleAccountDomainPlan{}, err
+	}
+
+	accountsCounter, err := s.Store().GetAccountsCounter(ctx)
+	if err != nil {
+		return singleAccountDomainPlan{}, fmt.Errorf("failed to count accounts: %w", err)
+	}
+	// The count is checked again here: it is read long after RequireSingleAccount, and marking an
+	// arbitrary account as the primary one for the domain would be wrong.
+	switch {
+	case accountsCounter == 0:
+		log.Info("no accounts yet, nothing to prepare for single account mode")
+		return singleAccountDomainPlan{skip: true}, nil
+	case accountsCounter > 1:
+		return singleAccountDomainPlan{}, errMultipleAccounts(accountsCounter)
+	}
+
+	accountID, err := s.Store().GetAnyAccountID(ctx)
+	if err != nil {
+		return singleAccountDomainPlan{}, fmt.Errorf("failed to get the existing account: %w", err)
+	}
+
+	isPrimary, accountDomain, err := s.Store().IsPrimaryAccount(ctx, accountID)
+	if err != nil {
+		return singleAccountDomainPlan{}, fmt.Errorf("failed to read domain attributes of account %s: %w", accountID, err)
+	}
+
+	domain, err := resolveAccountDomain(accountID, accountDomain, singleAccountDomain, requested)
+	if err != nil {
+		return singleAccountDomainPlan{}, err
+	}
+
+	return singleAccountDomainPlan{
+		accountID:     accountID,
+		domain:        domain,
+		currentDomain: accountDomain,
+		isPrimary:     isPrimary,
+	}, nil
+}

management/server/idp/migration/migration_test.go

@@ -24,6 +24,17 @@ type testStore struct {
 	checkSchemaFunc    func(checks []SchemaCheck) []SchemaError
 	updateCalls        []updateUserIDCall
 	updateInfoCalls    []updateUserInfoCall
+
+	accountsCounter int64
+	accounts        map[string]*types.Account
+	domainAttrCalls []domainAttrCall
+}
+
+type domainAttrCall struct {
+	AccountID string
+	Domain    string
+	Category  string
+	IsPrimary bool
 }
 
 type updateUserIDCall struct {
@@ -38,6 +49,35 @@ type updateUserInfoCall struct {
 	Name   string
 }
 
+func (s *testStore) GetAccountsCounter(context.Context) (int64, error) {
+	return s.accountsCounter, nil
+}
+
+func (s *testStore) GetAnyAccountID(context.Context) (string, error) {
+	for id := range s.accounts {
+		return id, nil
+	}
+	return "", fmt.Errorf("no accounts")
+}
+
+func (s *testStore) IsPrimaryAccount(_ context.Context, accountID string) (bool, string, error) {
+	account, ok := s.accounts[accountID]
+	if !ok {
+		return false, "", fmt.Errorf("account %s not found", accountID)
+	}
+	return account.IsDomainPrimaryAccount, account.Domain, nil
+}
+
+func (s *testStore) UpdateAccountDomainAttributes(_ context.Context, accountID, domain, category string, isPrimaryDomain bool) error {
+	s.domainAttrCalls = append(s.domainAttrCalls, domainAttrCall{accountID, domain, category, isPrimaryDomain})
+	if account, ok := s.accounts[accountID]; ok {
+		account.Domain = domain
+		account.DomainCategory = category
+		account.IsDomainPrimaryAccount = isPrimaryDomain
+	}
+	return nil
+}
+
 func (s *testStore) ListUsers(ctx context.Context) ([]*types.User, error) {
 	return s.listUsersFunc(ctx)
 }
@@ -826,3 +866,212 @@ func TestCheckSchema_MockStore(t *testing.T) {
 		assert.Equal(t, "email", errs[0].Column)
 	})
 }
+
+func TestRequireSingleAccount(t *testing.T) {
+	tests := []struct {
+		name      string
+		accounts  int64
+		expectErr bool
+	}{
+		{name: "fresh install", accounts: 0},
+		{name: "single account", accounts: 1},
+		{name: "multiple accounts", accounts: 3, expectErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			srv := &testServer{store: &testStore{accountsCounter: tt.accounts}}
+
+			err := RequireSingleAccount(srv)
+			if !tt.expectErr {
+				require.NoError(t, err)
+				return
+			}
+
+			require.Error(t, err)
+			assert.ErrorIs(t, err, ErrMultipleAccounts)
+		})
+	}
+}
+
+func TestEnsureSingleAccountDomain(t *testing.T) {
+	tests := []struct {
+		name            string
+		account         *types.Account
+		requestedDomain string
+		expectedDomain  string
+	}{
+		{
+			name:           "account migrated from an IdP without domain claims",
+			account:        &types.Account{Id: "account-1"},
+			expectedDomain: DefaultSingleAccountDomain,
+		},
+		{
+			name:            "requested domain is applied to an account without one",
+			account:         &types.Account{Id: "account-1"},
+			requestedDomain: "corp.example.com",
+			expectedDomain:  "corp.example.com",
+		},
+		{
+			name:           "account keeps its own domain",
+			account:        &types.Account{Id: "account-1", Domain: "acme.com"},
+			expectedDomain: "acme.com",
+		},
+		{
+			name:            "requesting the domain the account already has is not a conflict",
+			account:         &types.Account{Id: "account-1", Domain: "acme.com"},
+			requestedDomain: "acme.com",
+			expectedDomain:  "acme.com",
+		},
+		{
+			name: "already resolvable account is rewritten with the same values",
+			account: &types.Account{
+				Id:                     "account-1",
+				Domain:                 "acme.com",
+				DomainCategory:         types.PrivateCategory,
+				IsDomainPrimaryAccount: true,
+			},
+			expectedDomain: "acme.com",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := &testStore{
+				accountsCounter: 1,
+				accounts:        map[string]*types.Account{tt.account.Id: tt.account},
+			}
+
+			require.NoError(t, EnsureSingleAccountDomain(&testServer{store: store}, tt.requestedDomain))
+
+			require.Len(t, store.domainAttrCalls, 1)
+			assert.Equal(t, domainAttrCall{
+				AccountID: tt.account.Id,
+				Domain:    tt.expectedDomain,
+				Category:  types.PrivateCategory,
+				IsPrimary: true,
+			}, store.domainAttrCalls[0])
+		})
+	}
+}
+
+func TestEnsureSingleAccountDomainDryRun(t *testing.T) {
+	t.Setenv(dryRunEnvKey, "true")
+
+	store := &testStore{
+		accountsCounter: 1,
+		accounts:        map[string]*types.Account{"account-1": {Id: "account-1"}},
+	}
+
+	require.NoError(t, EnsureSingleAccountDomain(&testServer{store: store}, ""))
+	assert.Empty(t, store.domainAttrCalls, "Dry run must not write anything")
+}
+
+func TestEnsureSingleAccountDomainRejectsUnresolvableDomains(t *testing.T) {
+	t.Run("account domain that cannot resolve is reported", func(t *testing.T) {
+		account := &types.Account{Id: "account-1", Domain: "corp"}
+		store := &testStore{
+			accountsCounter: 1,
+			accounts:        map[string]*types.Account{account.Id: account},
+		}
+
+		err := EnsureSingleAccountDomain(&testServer{store: store}, "")
+		require.Error(t, err)
+		assert.ErrorIs(t, err, ErrUnusableDomain)
+		assert.Empty(t, store.domainAttrCalls, "A broken account domain must not be replaced silently")
+	})
+
+	t.Run("requested domain conflicting with the account domain is reported", func(t *testing.T) {
+		account := &types.Account{Id: "account-1", Domain: "acme.com"}
+		store := &testStore{
+			accountsCounter: 1,
+			accounts:        map[string]*types.Account{account.Id: account},
+		}
+
+		err := EnsureSingleAccountDomain(&testServer{store: store}, "corp.example.com")
+		require.Error(t, err)
+		assert.ErrorIs(t, err, ErrDomainConflict)
+		assert.Empty(t, store.domainAttrCalls, "A conflict must not overwrite the account domain")
+	})
+
+	t.Run("configured domain that cannot resolve is rejected", func(t *testing.T) {
+		store := &testStore{
+			accountsCounter: 1,
+			accounts:        map[string]*types.Account{"account-1": {Id: "account-1"}},
+		}
+
+		err := EnsureSingleAccountDomain(&testServer{store: store}, "corp")
+		require.Error(t, err)
+		assert.ErrorIs(t, err, ErrUnusableDomain)
+		assert.Empty(t, store.domainAttrCalls)
+	})
+
+	t.Run("account appearing after the preflight is rejected", func(t *testing.T) {
+		store := &testStore{
+			accountsCounter: 2,
+			accounts: map[string]*types.Account{
+				"account-1": {Id: "account-1"},
+				"account-2": {Id: "account-2"},
+			},
+		}
+
+		err := EnsureSingleAccountDomain(&testServer{store: store}, "")
+		require.Error(t, err)
+		assert.ErrorIs(t, err, ErrMultipleAccounts)
+		assert.Empty(t, store.domainAttrCalls, "No account may be marked primary when several exist")
+	})
+
+	t.Run("fresh install with no accounts is a no-op", func(t *testing.T) {
+		store := &testStore{accountsCounter: 0, accounts: map[string]*types.Account{}}
+
+		require.NoError(t, EnsureSingleAccountDomain(&testServer{store: store}, ""))
+		assert.Empty(t, store.domainAttrCalls)
+	})
+}
+
+func TestCheckSingleAccountDomain(t *testing.T) {
+	tests := []struct {
+		name      string
+		account   *types.Account
+		requested string
+		expectErr error
+	}{
+		{
+			name:    "usable account domain passes",
+			account: &types.Account{Id: "account-1", Domain: "acme.com"},
+		},
+		{
+			name:    "empty account domain passes",
+			account: &types.Account{Id: "account-1"},
+		},
+		{
+			name:      "unresolvable account domain fails",
+			account:   &types.Account{Id: "account-1", Domain: "corp"},
+			expectErr: ErrUnusableDomain,
+		},
+		{
+			name:      "conflicting request fails",
+			account:   &types.Account{Id: "account-1", Domain: "acme.com"},
+			requested: "corp.example.com",
+			expectErr: ErrDomainConflict,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := &testStore{
+				accountsCounter: 1,
+				accounts:        map[string]*types.Account{tt.account.Id: tt.account},
+			}
+
+			err := CheckSingleAccountDomain(&testServer{store: store}, tt.requested)
+			if tt.expectErr == nil {
+				require.NoError(t, err)
+			} else {
+				require.ErrorIs(t, err, tt.expectErr)
+			}
+
+			assert.Empty(t, store.domainAttrCalls, "The preflight must not write anything")
+		})
+	}
+}

management/server/idp/migration/store.go

@@ -60,6 +60,20 @@ type Store interface {
 	// CheckSchema verifies that all tables and columns required by the migration
 	// exist in the database. Returns a list of problems; an empty slice means OK.
 	CheckSchema(checks []SchemaCheck) []SchemaError
+
+	// GetAccountsCounter returns the total number of accounts in the store.
+	GetAccountsCounter(ctx context.Context) (int64, error)
+
+	// GetAnyAccountID returns the ID of one of the existing accounts.
+	GetAnyAccountID(ctx context.Context) (string, error)
+
+	// IsPrimaryAccount returns whether the account is the primary account for its domain,
+	// along with that domain.
+	IsPrimaryAccount(ctx context.Context, accountID string) (bool, string, error)
+
+	// UpdateAccountDomainAttributes sets the domain, domain category and primary
+	// domain flag of an account.
+	UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error
 }
 
 // RequiredEventSchema lists all tables and columns that the migration tool needs

tools/idp-migrate/DEVELOPMENT.md

@@ -50,6 +50,7 @@ The build requires `CGO_ENABLED=1` because it links the SQLite driver used by `S
 | `--domain` | string | `""` | Sets both dashboard and API domain (convenience shorthand) |
 | `--dashboard-domain` | string | *(required)* | Dashboard domain (for redirect URIs) |
 | `--api-domain` | string | *(required)* | API domain (for Dex issuer and callback URLs) |
+| `--single-account-mode-domain` | string | `netbird.selfhosted` | Domain single account mode groups users under. Used only when the account has no domain of its own; passing one that conflicts with the account's existing domain is an error |
 | `--dry-run` | bool | `false` | Preview changes without writing |
 | `--force` | bool | `false` | Skip interactive confirmation prompt |
 | `--skip-config` | bool | `false` | Skip config generation (DB-only migration) |
@@ -68,6 +69,7 @@ All flags can be overridden via environment variables. Env vars take precedence
 | `NETBIRD_CONFIG_PATH` | `--config` |
 | `NETBIRD_DATA_DIR` | `--datadir` |
 | `NETBIRD_IDP_SEED_INFO` | `--idp-seed-info` |
+| `NETBIRD_SINGLE_ACCOUNT_MODE_DOMAIN` | `--single-account-mode-domain` |
 | `NETBIRD_DRY_RUN` | `--dry-run` (set to `"true"`) |
 | `NETBIRD_FORCE` | `--force` (set to `"true"`) |
 | `NETBIRD_SKIP_CONFIG` | `--skip-config` (set to `"true"`) |

tools/idp-migrate/config.go

@@ -6,16 +6,18 @@ import (
 	"os"
 	"strconv"
 
+	"github.com/netbirdio/netbird/management/server/idp/migration"
 	"github.com/netbirdio/netbird/util"
 )
 
 type migrationConfig struct {
 	// Data
-	dashboardURL string
-	apiURL       string
-	configPath   string
-	dataDir      string
-	idpSeedInfo  string
+	dashboardURL        string
+	apiURL              string
+	configPath          string
+	dataDir             string
+	idpSeedInfo         string
+	singleAccountDomain string
 
 	// Options
 	dryRun               bool
@@ -51,6 +53,7 @@ func configFromArgs(args []string) (*migrationConfig, error) {
 	fs.StringVar(&cfg.configPath, "config", "", "path to management.json (required)")
 	fs.StringVar(&cfg.dataDir, "datadir", "", "override data directory from config")
 	fs.StringVar(&cfg.idpSeedInfo, "idp-seed-info", "", "base64-encoded connector JSON (overrides auto-detection)")
+	fs.StringVar(&cfg.singleAccountDomain, "single-account-mode-domain", "", "domain single account mode groups users under, used only when the account has no domain of its own (default "+migration.DefaultSingleAccountDomain+")")
 	fs.BoolVar(&cfg.dryRun, "dry-run", false, "preview changes without writing")
 	fs.BoolVar(&cfg.force, "force", false, "skip confirmation prompt")
 	fs.BoolVar(&cfg.skipConfig, "skip-config", false, "skip config generation (DB migration only)")
@@ -118,6 +121,10 @@ func applyOverrides(cfg *migrationConfig, domain string) {
 		cfg.idpSeedInfo = val
 	}
 
+	if val, ok := os.LookupEnv("NETBIRD_SINGLE_ACCOUNT_MODE_DOMAIN"); ok {
+		cfg.singleAccountDomain = val
+	}
+
 	// Enforce dry run if any value is provided
 	if sval, ok := os.LookupEnv("NETBIRD_DRY_RUN"); ok {
 		if val, err := strconv.ParseBool(sval); err == nil {
@@ -170,5 +177,9 @@ func validateConfig(cfg *migrationConfig) error {
 		return fmt.Errorf("--dashboard-domain is required")
 	}
 
+	if _, err := migration.NormalizeSingleAccountDomain(cfg.singleAccountDomain); err != nil {
+		return err
+	}
+
 	return nil
 }

tools/idp-migrate/main.go

@@ -71,6 +71,10 @@ func run(cfg *migrationConfig) error {
 		return err
 	}
 
+	if err := preflightAccounts(cfg, mgmtConfig); err != nil {
+		return err
+	}
+
 	if !cfg.skipPopulateUserInfo {
 		err := populateUserInfoFromIDP(cfg, mgmtConfig)
 		if err != nil {
@@ -102,6 +106,22 @@ func run(cfg *migrationConfig) error {
 	return generateConfig(cfg, connectorConfig)
 }
 
+func preflightAccounts(cfg *migrationConfig, mgmtConfig *nbconfig.Config) error {
+	ctx := context.Background()
+	migStore, migEventStore, cleanup, err := openStores(ctx, mgmtConfig, cfg.dataDir)
+	if err != nil {
+		return err
+	}
+	defer cleanup()
+
+	srv := &migrationServer{store: migStore, eventStore: migEventStore}
+	if err := migration.RequireSingleAccount(srv); err != nil {
+		return err
+	}
+
+	return migration.CheckSingleAccountDomain(srv, cfg.singleAccountDomain)
+}
+
 // validateSchema opens the store and checks that all required tables and columns
 // exist. If anything is missing, it returns a descriptive error telling the user
 // to upgrade their management server.
@@ -224,6 +244,8 @@ func migrateDB(cfg *migrationConfig, mgmtConfig *nbconfig.Config, connectorConfi
 	}
 	defer cleanup()
 
+	srv := &migrationServer{store: migStore, eventStore: migEventStore}
+
 	pending, err := previewUsers(ctx, migStore)
 	if err != nil {
 		return err
@@ -243,11 +265,14 @@ func migrateDB(cfg *migrationConfig, mgmtConfig *nbconfig.Config, connectorConfi
 		}
 	}
 
-	srv := &migrationServer{store: migStore, eventStore: migEventStore}
 	if err := migration.MigrateUsersToStaticConnectors(srv, connectorConfig); err != nil {
 		return fmt.Errorf("migrate users: %w", err)
 	}
 
+	if err := migration.EnsureSingleAccountDomain(srv, cfg.singleAccountDomain); err != nil {
+		return fmt.Errorf("prepare single account mode: %w", err)
+	}
+
 	if !cfg.dryRun {
 		log.Info("DB migration completed successfully")
 	}

tools/idp-migrate/main_test.go

@@ -485,3 +485,77 @@ func TestGenerateConfig(t *testing.T) {
 		assert.True(t, os.IsNotExist(err))
 	})
 }
+
+func TestValidateConfigRejectsUnusableSingleAccountDomain(t *testing.T) {
+	base := func() migrationConfig {
+		return migrationConfig{
+			configPath:          "/tmp/management.json",
+			dataDir:             "/tmp/datadir",
+			idpSeedInfo:         "seed",
+			apiURL:              "https://api.example.com",
+			dashboardURL:        "https://app.example.com",
+			singleAccountDomain: migration.DefaultSingleAccountDomain,
+		}
+	}
+
+	t.Run("usable domain is accepted", func(t *testing.T) {
+		cfg := base()
+		require.NoError(t, validateConfig(&cfg))
+	})
+
+	t.Run("empty falls back to the default", func(t *testing.T) {
+		cfg := base()
+		cfg.singleAccountDomain = ""
+		require.NoError(t, validateConfig(&cfg))
+	})
+
+	// Rejected up front so the migration cannot fail after it has rewritten user IDs.
+	t.Run("single label domain is rejected", func(t *testing.T) {
+		cfg := base()
+		cfg.singleAccountDomain = "corp"
+		err := validateConfig(&cfg)
+		require.Error(t, err)
+		assert.ErrorIs(t, err, migration.ErrUnusableDomain)
+	})
+}
+
+func TestApplyOverrides_SingleAccountDomainFromEnv(t *testing.T) {
+	t.Run("env var overrides the flag", func(t *testing.T) {
+		t.Setenv("NETBIRD_SINGLE_ACCOUNT_MODE_DOMAIN", "corp.example.com")
+
+		cfg, err := configFromArgs([]string{
+			"--config", "/tmp/management.json",
+			"--datadir", "/tmp/datadir",
+			"--idp-seed-info", "seed",
+			"--domain", "example.com",
+			"--single-account-mode-domain", "flag.example.com",
+		})
+		require.NoError(t, err)
+		assert.Equal(t, "corp.example.com", cfg.singleAccountDomain)
+	})
+
+	t.Run("unset leaves the flag value", func(t *testing.T) {
+		cfg, err := configFromArgs([]string{
+			"--config", "/tmp/management.json",
+			"--datadir", "/tmp/datadir",
+			"--idp-seed-info", "seed",
+			"--domain", "example.com",
+			"--single-account-mode-domain", "flag.example.com",
+		})
+		require.NoError(t, err)
+		assert.Equal(t, "flag.example.com", cfg.singleAccountDomain)
+	})
+
+	t.Run("unusable env value is rejected", func(t *testing.T) {
+		t.Setenv("NETBIRD_SINGLE_ACCOUNT_MODE_DOMAIN", "corp")
+
+		_, err := configFromArgs([]string{
+			"--config", "/tmp/management.json",
+			"--datadir", "/tmp/datadir",
+			"--idp-seed-info", "seed",
+			"--domain", "example.com",
+		})
+		require.Error(t, err)
+		assert.ErrorIs(t, err, migration.ErrUnusableDomain)
+	})
+}