[management] Enable PAT creation during setup (#6003)

* enable pat creation on setup * remove logic from handler towards setup service * fix lint issue * fix rollback on account id returning empty * fix coderabbit comments * fix setup PAT rollback behavior
2026-05-05 16:46:39 +00:00 · 2026-04-30 17:21:35 +02:00
parent f29f5a0978
commit dcd1db42ef
11 changed files with 997 additions and 19 deletions
--- a/management/server/http/handlers/instance/instance_handler_test.go
+++ b/management/server/http/handlers/instance/instance_handler_test.go
@@ -10,12 +10,18 @@ import (
 	"net/mail"
 	"testing"

+	"github.com/golang/mock/gomock"
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"

+	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/idp"
 	nbinstance "github.com/netbirdio/netbird/management/server/instance"
+	"github.com/netbirdio/netbird/management/server/mock_server"
+	nbstore "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/http/api"
 	"github.com/netbirdio/netbird/shared/management/status"
 )
@@ -25,6 +31,7 @@ type mockInstanceManager struct {
 	isSetupRequired   bool
 	isSetupRequiredFn func(ctx context.Context) (bool, error)
 	createOwnerUserFn func(ctx context.Context, email, password, name string) (*idp.UserData, error)
+	rollbackSetupFn   func(ctx context.Context, userID string) error
 	getVersionInfoFn  func(ctx context.Context) (*nbinstance.VersionInfo, error)
 }

@@ -67,6 +74,13 @@ func (m *mockInstanceManager) CreateOwnerUser(ctx context.Context, email, passwo
 	}, nil
 }

+func (m *mockInstanceManager) RollbackSetup(ctx context.Context, userID string) error {
+	if m.rollbackSetupFn != nil {
+		return m.rollbackSetupFn(ctx, userID)
+	}
+	return nil
+}
+
 func (m *mockInstanceManager) GetVersionInfo(ctx context.Context) (*nbinstance.VersionInfo, error) {
 	if m.getVersionInfoFn != nil {
 		return m.getVersionInfoFn(ctx)
@@ -82,8 +96,12 @@ func (m *mockInstanceManager) GetVersionInfo(ctx context.Context) (*nbinstance.V
 var _ nbinstance.Manager = (*mockInstanceManager)(nil)

 func setupTestRouter(manager nbinstance.Manager) *mux.Router {
+	return setupTestRouterWithPAT(manager, nil)
+}
+
+func setupTestRouterWithPAT(manager nbinstance.Manager, accountManager account.Manager) *mux.Router {
 	router := mux.NewRouter()
-	AddEndpoints(manager, router)
+	AddEndpoints(manager, accountManager, router)
 	return router
 }

@@ -161,6 +179,7 @@ func TestSetup_Success(t *testing.T) {
 	router.ServeHTTP(rec, req)

 	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"))

 	var response api.SetupResponse
 	err := json.NewDecoder(rec.Body).Decode(&response)
@@ -293,6 +312,239 @@ func TestSetup_ManagerError(t *testing.T) {
 	assert.Equal(t, http.StatusInternalServerError, rec.Code)
 }

+func TestSetup_PAT_FeatureDisabled_IgnoresCreatePAT(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "false")
+
+	manager := &mockInstanceManager{isSetupRequired: true}
+	// NB_SETUP_PAT_ENABLED=false: request fields must be silently ignored
+	router := setupTestRouterWithPAT(manager, &mock_server.MockAccountManager{})
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	var response api.SetupResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&response))
+	assert.Nil(t, response.PersonalAccessToken)
+}
+
+func TestSetup_PAT_FlagOmitted_NoPAT(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	manager := &mockInstanceManager{isSetupRequired: true}
+	router := setupTestRouterWithPAT(manager, &mock_server.MockAccountManager{})
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin"}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	var response api.SetupResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&response))
+	assert.Nil(t, response.PersonalAccessToken)
+}
+
+func TestSetup_PAT_MissingExpireIn_DefaultsToOneDay(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	createCalled := false
+	manager := &mockInstanceManager{
+		isSetupRequired: true,
+		createOwnerUserFn: func(ctx context.Context, email, password, name string) (*idp.UserData, error) {
+			createCalled = true
+			return &idp.UserData{ID: "u1", Email: email, Name: name}, nil
+		},
+	}
+	accountMgr := &mock_server.MockAccountManager{
+		GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+			assert.Equal(t, "u1", userAuth.UserId)
+			return "acc-1", nil
+		},
+		CreatePATFunc: func(_ context.Context, accountID, initiator, target, name string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
+			assert.Equal(t, "acc-1", accountID)
+			assert.Equal(t, "u1", initiator)
+			assert.Equal(t, "u1", target)
+			assert.Equal(t, "setup-token", name)
+			assert.Equal(t, 1, expiresIn)
+			return &types.PersonalAccessTokenGenerated{PlainToken: "nbp_plain"}, nil
+		},
+	}
+	router := setupTestRouterWithPAT(manager, accountMgr)
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"))
+	assert.True(t, createCalled)
+	var response api.SetupResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&response))
+	require.NotNil(t, response.PersonalAccessToken)
+	assert.Equal(t, "nbp_plain", *response.PersonalAccessToken)
+}
+
+func TestSetup_PAT_ExpireOutOfRange(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	manager := &mockInstanceManager{isSetupRequired: true}
+	router := setupTestRouterWithPAT(manager, &mock_server.MockAccountManager{})
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true, "pat_expire_in": 0}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusUnprocessableEntity, rec.Code)
+}
+
+func TestSetup_PAT_Success(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	manager := &mockInstanceManager{
+		isSetupRequired: true,
+		createOwnerUserFn: func(ctx context.Context, email, password, name string) (*idp.UserData, error) {
+			return &idp.UserData{ID: "owner-id", Email: email, Name: name}, nil
+		},
+	}
+
+	gotAccountArgs := struct {
+		userID string
+		email  string
+	}{}
+	accountMgr := &mock_server.MockAccountManager{
+		GetAccountIDByUserIdFunc: func(_ context.Context, userAuth auth.UserAuth) (string, error) {
+			gotAccountArgs.userID = userAuth.UserId
+			gotAccountArgs.email = userAuth.Email
+			return "acc-1", nil
+		},
+		CreatePATFunc: func(_ context.Context, accountID, initiator, target, name string, expiresIn int) (*types.PersonalAccessTokenGenerated, error) {
+			assert.Equal(t, "acc-1", accountID)
+			assert.Equal(t, "owner-id", initiator)
+			assert.Equal(t, "owner-id", target)
+			assert.Equal(t, "setup-token", name)
+			assert.Equal(t, 30, expiresIn)
+			return &types.PersonalAccessTokenGenerated{PlainToken: "nbp_plain"}, nil
+		},
+	}
+
+	router := setupTestRouterWithPAT(manager, accountMgr)
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true, "pat_expire_in": 30}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"))
+	var response api.SetupResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&response))
+	assert.Equal(t, "owner-id", response.UserId)
+	require.NotNil(t, response.PersonalAccessToken)
+	assert.Equal(t, "nbp_plain", *response.PersonalAccessToken)
+	assert.Equal(t, "owner-id", gotAccountArgs.userID)
+	assert.Equal(t, "admin@example.com", gotAccountArgs.email)
+}
+
+func TestSetup_PAT_AccountCreationFails_Rollback(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	ctrl := gomock.NewController(t)
+	accountStore := nbstore.NewMockStore(ctrl)
+	accountStore.EXPECT().GetAccountIDByUserID(gomock.Any(), nbstore.LockingStrengthNone, "owner-id").Return("", status.NewAccountNotFoundError("owner-id"))
+
+	rolledBackFor := ""
+	manager := &mockInstanceManager{
+		isSetupRequired: true,
+		createOwnerUserFn: func(ctx context.Context, email, password, name string) (*idp.UserData, error) {
+			return &idp.UserData{ID: "owner-id", Email: email, Name: name}, nil
+		},
+		rollbackSetupFn: func(_ context.Context, userID string) error {
+			rolledBackFor = userID
+			return nil
+		},
+	}
+	accountMgr := &mock_server.MockAccountManager{
+		GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+			return "", errors.New("db down")
+		},
+		GetStoreFunc: func() nbstore.Store {
+			return accountStore
+		},
+	}
+
+	router := setupTestRouterWithPAT(manager, accountMgr)
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true, "pat_expire_in": 30}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusInternalServerError, rec.Code)
+	assert.Equal(t, "owner-id", rolledBackFor, "RollbackSetup must be called with the created user id")
+}
+
+func TestSetup_PAT_CreatePATFails_Rollback(t *testing.T) {
+	t.Setenv(nbinstance.SetupPATEnabledEnvKey, "true")
+
+	ctrl := gomock.NewController(t)
+	accountStore := nbstore.NewMockStore(ctrl)
+	account := &types.Account{Id: "acc-1"}
+	accountStore.EXPECT().GetAccount(gomock.Any(), "acc-1").Return(account, nil)
+	accountStore.EXPECT().DeleteAccount(gomock.Any(), account).Return(nil)
+
+	rolledBackFor := ""
+	manager := &mockInstanceManager{
+		isSetupRequired: true,
+		createOwnerUserFn: func(ctx context.Context, email, password, name string) (*idp.UserData, error) {
+			return &idp.UserData{ID: "owner-id", Email: email, Name: name}, nil
+		},
+		rollbackSetupFn: func(_ context.Context, userID string) error {
+			rolledBackFor = userID
+			return nil
+		},
+	}
+	accountMgr := &mock_server.MockAccountManager{
+		GetAccountIDByUserIdFunc: func(_ context.Context, _ auth.UserAuth) (string, error) {
+			return "acc-1", nil
+		},
+		CreatePATFunc: func(_ context.Context, _, _, _, _ string, _ int) (*types.PersonalAccessTokenGenerated, error) {
+			return nil, status.Errorf(status.Internal, "token store unavailable")
+		},
+		GetStoreFunc: func() nbstore.Store {
+			return accountStore
+		},
+	}
+
+	router := setupTestRouterWithPAT(manager, accountMgr)
+
+	body := `{"email": "admin@example.com", "password": "securepassword123", "name": "Admin", "create_pat": true, "pat_expire_in": 30}`
+	req := httptest.NewRequest(http.MethodPost, "/setup", bytes.NewBufferString(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusInternalServerError, rec.Code)
+	assert.Equal(t, "owner-id", rolledBackFor, "RollbackSetup must be called when CreatePAT fails")
+}
+
 func TestGetVersionInfo_Success(t *testing.T) {
 	manager := &mockInstanceManager{}
 	router := mux.NewRouter()