[management] Give the /api/setup owner account the single account mode domain (#8147)

management/server/account_test.go

@@ -50,6 +50,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/cache"
 	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools"
 	"github.com/netbirdio/netbird/management/server/idp"
+	"github.com/netbirdio/netbird/management/server/instance"
 	"github.com/netbirdio/netbird/management/server/job"
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
@@ -844,6 +845,49 @@ func TestAccountManager_GetAccountByUserID(t *testing.T) {
 	}
 }
 
+type setupOwnerInstanceManager struct {
+	instance.Manager
+	user *idp.UserData
+}
+
+func (m setupOwnerInstanceManager) CreateOwnerUser(context.Context, string, string, string) (*idp.UserData, error) {
+	return m.user, nil
+}
+
+// The /api/setup owner account must get the single account mode domain, so later
+// users join that account instead of forking a new one.
+func TestSetupOwner_SingleAccountMode_LaterUsersJoinSetupAccount(t *testing.T) {
+	t.Setenv(instance.SetupPATEnabledEnvKey, "true")
+
+	manager, _, err := createManager(t)
+	require.NoError(t, err)
+	manager.singleAccountMode = true
+	manager.singleAccountModeDomain = "netbird.selfhosted"
+
+	ctx := context.Background()
+	owner := &idp.UserData{ID: "owner", Email: "owner@example.com", Name: "Owner"}
+	setup := instance.NewSetupService(setupOwnerInstanceManager{user: owner}, manager)
+	_, err = setup.SetupOwner(ctx, owner.Email, "password", owner.Name, instance.SetupOptions{CreatePAT: true})
+	require.NoError(t, err)
+
+	ownerAccountID, err := manager.Store.GetAccountIDByUserID(ctx, store.LockingStrengthNone, owner.ID)
+	require.NoError(t, err)
+	acc, err := manager.Store.GetAccount(ctx, ownerAccountID)
+	require.NoError(t, err)
+	assert.Equal(t, "netbird.selfhosted", acc.Domain)
+	assert.Equal(t, types.PrivateCategory, acc.DomainCategory)
+	assert.True(t, acc.IsDomainPrimaryAccount)
+
+	jitAccountID, _, err := manager.GetAccountIDFromUserAuth(ctx, auth.UserAuth{
+		UserId:         "jit-user",
+		Email:          "jit@gmail.com",
+		Domain:         "gmail.com",
+		DomainCategory: types.PublicCategory,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, ownerAccountID, jitAccountID)
+}
+
 func createAccount(am *DefaultAccountManager, accountID, userID, domain string) (*types.Account, error) {
 	account := newAccountWithId(context.Background(), accountID, userID, domain, "", "", false)
 	err := am.Store.SaveAccount(context.Background(), account)

management/server/http/handlers/instance/instance_handler_test.go

@@ -363,9 +363,9 @@ func TestSetup_PAT_MissingExpireIn_DefaultsToOneDay(t *testing.T) {
 		},
 	}
 	accountMgr := &mock_server.MockAccountManager{
-		GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+		GetAccountIDFromUserAuthFunc: func(_ context.Context, userAuth auth.UserAuth) (string, string, error) {
 			assert.Equal(t, "u1", userAuth.UserId)
-			return "acc-1", nil
+			return "acc-1", "u1", nil
 		},
 		CreatePATFunc: func(_ context.Context, accountID, initiator, target, name string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
 			assert.Equal(t, "acc-1", accountID)
@@ -425,10 +425,10 @@ func TestSetup_PAT_Success(t *testing.T) {
 		email  string
 	}{}
 	accountMgr := &mock_server.MockAccountManager{
-		GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+		GetAccountIDFromUserAuthFunc: func(_ context.Context, userAuth auth.UserAuth) (string, string, error) {
 			gotAccountArgs.userID = userAuth.UserId
 			gotAccountArgs.email = userAuth.Email
-			return "acc-1", nil
+			return "acc-1", "u1", nil
 		},
 		CreatePATFunc: func(_ context.Context, accountID, initiator, target, name string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
 			assert.Equal(t, "acc-1", accountID)
@@ -479,8 +479,8 @@ func TestSetup_PAT_AccountCreationFails_Rollback(t *testing.T) {
 		},
 	}
 	accountMgr := &mock_server.MockAccountManager{
-		GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+		GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) {
-			return "", errors.New("db down")
+			return "", "", errors.New("db down")
 		},
 		GetStoreFunc: func() nbstore.Store {
 			return accountStore
@@ -505,7 +505,7 @@ func TestSetup_PAT_CreatePATFails_Rollback(t *testing.T) {
 
 	ctrl := gomock.NewController(t)
 	accountStore := nbstore.NewMockStore(ctrl)
-	account := &types.Account{Id: "acc-1"}
+	account := &types.Account{Id: "acc-1", CreatedBy: "owner-id"}
 	accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
 	accountStore.EXPECT().DeleteAccount(gomock.Any(), account).Return(nil)
 
@@ -521,8 +521,8 @@ func TestSetup_PAT_CreatePATFails_Rollback(t *testing.T) {
 		},
 	}
 	accountMgr := &mock_server.MockAccountManager{
-		GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+		GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) {
-			return "acc-1", nil
+			return "acc-1", "u1", nil
 		},
 		CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) {
 			return nil, status.Errorf(status.Internal, "token store unavailable")

management/server/instance/setup_service.go

@@ -10,6 +10,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/shared/auth"
 	"github.com/netbirdio/netbird/shared/management/status"
 )
@@ -111,7 +112,9 @@ func (m *SetupService) SetupOwner(ctx context.Context, email, password, name str
 		Name:   userData.Name,
 	}
 
-	accountID, err := m.accountManager.GetAccountIDByUserID(ctx, userAuth)
+	// Resolve through the same path as a first login so single account mode
+	// stamps the configured domain, private category and primary flag.
+	accountID, _, err := m.accountManager.GetAccountIDFromUserAuth(ctx, userAuth)
 	if err != nil {
 		err = fmt.Errorf("create account for setup user: %w", err)
 		if rollbackErr := m.rollbackSetup(ctx, userData.ID, "account provisioning failed", err, ""); rollbackErr != nil {
@@ -145,7 +148,7 @@ func (m *SetupService) rollbackSetup(ctx context.Context, userID, reason string,
 	}
 
 	if accountID != "" {
-		if err := m.rollbackSetupAccount(ctx, accountID); err != nil {
+		if err := m.rollbackSetupAccount(ctx, accountID, userID); err != nil {
 			rollbackErr := fmt.Errorf("roll back setup account %s: %w", accountID, err)
 			log.WithContext(ctx).Errorf("failed to roll back setup account %s for user %s after %s: original error: %v, rollback error: %v", accountID, userID, reason, origErr, rollbackErr)
 			return rollbackErr
@@ -186,8 +189,9 @@ func (m *SetupService) lookupSetupAccountIDForRollback(ctx context.Context, user
 // rollbackSetupAccount removes only the setup-created account data from the
 // store. It intentionally avoids accountManager.DeleteAccount because the normal
 // account deletion path also deletes users from the IdP; embedded IdP cleanup is
-// owned by instanceManager.RollbackSetup.
+// owned by instanceManager.RollbackSetup. It keeps an account another user
-func (m *SetupService) rollbackSetupAccount(ctx context.Context, accountID string) error {
+// created or joined, which a concurrent first login can produce.
+func (m *SetupService) rollbackSetupAccount(ctx context.Context, accountID, userID string) error {
 	if m.accountManager == nil {
 		return fmt.Errorf("account manager is required to roll back setup account")
 	}
@@ -205,6 +209,11 @@ func (m *SetupService) rollbackSetupAccount(ctx context.Context, accountID strin
 		return fmt.Errorf("get setup account for rollback: %w", err)
 	}
 
+	if !isSetupOnlyAccount(account, userID) {
+		log.WithContext(ctx).Warnf("keeping account %s on setup rollback: shared with another user", accountID)
+		return nil
+	}
+
 	if err := accountStore.DeleteAccount(ctx, account); err != nil {
 		if isNotFoundError(err) {
 			return nil
@@ -214,3 +223,15 @@ func (m *SetupService) rollbackSetupAccount(ctx context.Context, accountID strin
 
 	return nil
 }
+
+func isSetupOnlyAccount(account *types.Account, userID string) bool {
+	if account.CreatedBy != userID {
+		return false
+	}
+	for id := range account.Users {
+		if id != userID {
+			return false
+		}
+	}
+	return true
+}

management/server/instance/setup_service_test.go

@@ -5,9 +5,9 @@ import (
 	"errors"
 	"testing"
 
-	"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/server/idp"
 	"github.com/netbirdio/netbird/management/server/mock_server"
@@ -87,9 +87,9 @@ func TestSetupOwner_PATFeatureEnabled_MissingExpireDefaultsToOneDay(t *testing.T
 			},
 		},
 		&mock_server.MockAccountManager{
-			GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+			GetAccountIDFromUserAuthFunc: func(_ context.Context, userAuth auth.UserAuth) (string, string, error) {
 				assert.Equal(t, "owner-id", userAuth.UserId)
-				return "acc-1", nil
+				return "acc-1", "owner-id", nil
 			},
 			CreatePATFunc: func(_ context.Context, accountID, initiatorUserID, targetUserID, tokenName string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
 				assert.Equal(t, "acc-1", accountID)
@@ -147,7 +147,7 @@ func TestSetupOwner_AccountProvisioningFails_RollsBackSideEffectAccountAndUser(t
 
 	ctrl := gomock.NewController(t)
 	accountStore := nbstore.NewMockStore(ctrl)
-	account := &types.Account{Id: "acc-1"}
+	account := &types.Account{Id: "acc-1", CreatedBy: "owner-id"}
 	accountStore.EXPECT().GetAccountIDByUserID(gomock.Any(), nbstore.LockingStrengthNone, "owner-id").Return("acc-1", nil)
 	accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
 	accountStore.EXPECT().DeleteAccount(gomock.Any(), account).Return(nil)
@@ -163,9 +163,9 @@ func TestSetupOwner_AccountProvisioningFails_RollsBackSideEffectAccountAndUser(t
 			},
 		},
 		&mock_server.MockAccountManager{
-			GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+			GetAccountIDFromUserAuthFunc: func(_ context.Context, userAuth auth.UserAuth) (string, string, error) {
 				assert.Equal(t, "owner-id", userAuth.UserId)
-				return "", errors.New("metadata update failed")
+				return "", "", errors.New("metadata update failed")
 			},
 			GetStoreFunc: func() nbstore.Store {
 				return accountStore
@@ -190,7 +190,7 @@ func TestSetupOwner_CreatePATFails_RollsBackSetupAccountAndUser(t *testing.T) {
 
 	ctrl := gomock.NewController(t)
 	accountStore := nbstore.NewMockStore(ctrl)
-	account := &types.Account{Id: "acc-1"}
+	account := &types.Account{Id: "acc-1", CreatedBy: "owner-id"}
 	accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
 	accountStore.EXPECT().DeleteAccount(gomock.Any(), account).Return(nil)
 
@@ -204,9 +204,9 @@ func TestSetupOwner_CreatePATFails_RollsBackSetupAccountAndUser(t *testing.T) {
 			},
 		},
 		&mock_server.MockAccountManager{
-			GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+			GetAccountIDFromUserAuthFunc: func(_ context.Context, userAuth auth.UserAuth) (string, string, error) {
 				assert.Equal(t, "owner-id", userAuth.UserId)
-				return "acc-1", nil
+				return "acc-1", "owner-id", nil
 			},
 			CreatePATFunc: func(_ context.Context, accountID, initiatorUserID, targetUserID, tokenName string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
 				assert.Equal(t, "acc-1", accountID)
@@ -233,6 +233,59 @@ func TestSetupOwner_CreatePATFails_RollsBackSetupAccountAndUser(t *testing.T) {
 	assert.Equal(t, 1, rollbackCalls)
 }
 
+// A concurrent first login can make the setup account shared: either a JIT user
+// created the primary account first and the owner joined it, or a JIT user joined
+// the setup account. A PAT failure must not delete an account other users rely on.
+// DeleteAccount is not expected, so the mock store fails the test if it is called.
+func TestSetupOwner_CreatePATFails_KeepsSharedAccount(t *testing.T) {
+	tests := map[string]*types.Account{
+		"created by another user": {Id: "acc-1", CreatedBy: "jit-user"},
+		"another user joined": {Id: "acc-1", CreatedBy: "owner-id", Users: map[string]*types.User{
+			"owner-id": {Id: "owner-id"},
+			"jit-user": {Id: "jit-user"},
+		}},
+	}
+
+	for name, account := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Setenv(SetupPATEnabledEnvKey, "true")
+
+			ctrl := gomock.NewController(t)
+			accountStore := nbstore.NewMockStore(ctrl)
+			accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
+
+			rollbackCalls := 0
+			setupManager := NewSetupService(
+				&setupInstanceManagerMock{
+					rollbackSetupFn: func(_ context.Context, _ string) error {
+						rollbackCalls++
+						return nil
+					},
+				},
+				&mock_server.MockAccountManager{
+					GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) {
+						return "acc-1", "owner-id", nil
+					},
+					CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) {
+						return nil, status.Errorf(status.PermissionDenied, "user is blocked")
+					},
+					GetStoreFunc: func() nbstore.Store {
+						return accountStore
+					},
+				},
+			)
+
+			_, err := setupManager.SetupOwner(context.Background(), "admin@example.com", "securepassword123", "Admin", SetupOptions{
+				CreatePAT:       true,
+				PATExpireInDays: intPtr(30),
+			})
+
+			require.Error(t, err)
+			assert.Equal(t, 1, rollbackCalls, "setup user should still be rolled back")
+		})
+	}
+}
+
 func TestSetupOwner_CreatePATFails_AccountAlreadyGoneStillRollsBackUser(t *testing.T) {
 	t.Setenv(SetupPATEnabledEnvKey, "true")
 
@@ -251,8 +304,8 @@ func TestSetupOwner_CreatePATFails_AccountAlreadyGoneStillRollsBackUser(t *testi
 			},
 		},
 		&mock_server.MockAccountManager{
-			GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+			GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) {
-				return "acc-1", nil
+				return "acc-1", "owner-id", nil
 			},
 			CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) {
 				return nil, errors.New("token failure")
@@ -280,7 +333,7 @@ func TestSetupOwner_CreatePATFails_AccountRollbackFailureStopsBeforeUserRollback
 
 	ctrl := gomock.NewController(t)
 	accountStore := nbstore.NewMockStore(ctrl)
-	account := &types.Account{Id: "acc-1"}
+	account := &types.Account{Id: "acc-1", CreatedBy: "owner-id"}
 	accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
 	accountStore.EXPECT().DeleteAccount(gomock.Any(), account).Return(errors.New("delete failed"))
 
@@ -293,8 +346,8 @@ func TestSetupOwner_CreatePATFails_AccountRollbackFailureStopsBeforeUserRollback
 			},
 		},
 		&mock_server.MockAccountManager{
-			GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+			GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) {
-				return "acc-1", nil
+				return "acc-1", "owner-id", nil
 			},
 			CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) {
 				return nil, errors.New("token failure")