From 73829c8d968bc2c0a8dd423449638ecc57134adf Mon Sep 17 00:00:00 2001 From: Mohammadreza Razavian <62871170+smrrazavian@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:16:33 +0330 Subject: [PATCH] [management] Give the /api/setup owner account the single account mode domain (#8147) --- management/server/account_test.go | 44 ++++++++++ .../instance/instance_handler_test.go | 18 ++--- management/server/instance/setup_service.go | 29 ++++++- .../server/instance/setup_service_test.go | 81 +++++++++++++++---- 4 files changed, 145 insertions(+), 27 deletions(-) diff --git a/management/server/account_test.go b/management/server/account_test.go index 8f7b1f18b..983ee8680 100644 --- a/management/server/account_test.go +++ b/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) diff --git a/management/server/http/handlers/instance/instance_handler_test.go b/management/server/http/handlers/instance/instance_handler_test.go index ba59497fa..8c72c112f 100644 --- a/management/server/http/handlers/instance/instance_handler_test.go +++ b/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) { - return "", errors.New("db down") + GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) { + 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) { - return "acc-1", nil + GetAccountIDFromUserAuthFunc: func(_ context.Context, _ auth.UserAuth) (string, string, error) { + return "acc-1", "u1", nil }, CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) { return nil, status.Errorf(status.Internal, "token store unavailable") diff --git a/management/server/instance/setup_service.go b/management/server/instance/setup_service.go index 92a4923be..1ef3e9736 100644 --- a/management/server/instance/setup_service.go +++ b/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. -func (m *SetupService) rollbackSetupAccount(ctx context.Context, accountID string) error { +// owned by instanceManager.RollbackSetup. It keeps an account another user +// 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 +} diff --git a/management/server/instance/setup_service_test.go b/management/server/instance/setup_service_test.go index af3a91b75..4fc210f5b 100644 --- a/management/server/instance/setup_service_test.go +++ b/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) { - return "acc-1", nil + 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, 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) { - return "acc-1", nil + 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, errors.New("token failure")