From 340dc93cd266d6df03b9d0c7485e3548db3877b2 Mon Sep 17 00:00:00 2001
From: Barak Amar <barak.amar@treeverse.io>
Date: Wed, 29 Jan 2025 11:41:17 +0200
Subject: [PATCH] Add admin/superuser error handling and cleanup (#8559)

* Add admin/superuser error handling and cleanup

- Added a defer function to clean up the admin user if an error occurs during creation.
- Added tests for creating an admin user, verifying successful creation, attempting to create another admin user, deleting the admin user, creating an admin user without a secret access key, and re-adding the admin user.
- Included a new test `TestAddAdminUser` to verify the behavior of the `AddAdminUser` function in various scenarios.

* Added a logger to log the error message and warn about deleting the user
---
 pkg/auth/setup/setup.go      | 10 +++++
 pkg/auth/setup/setup_test.go | 80 ++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 pkg/auth/setup/setup_test.go

diff --git a/pkg/auth/setup/setup.go b/pkg/auth/setup/setup.go
index fd86505a4d0..19ba28dcf4b 100644
--- a/pkg/auth/setup/setup.go
+++ b/pkg/auth/setup/setup.go
@@ -174,6 +174,16 @@ func AddAdminUser(ctx context.Context, authService auth.Service, user *model.Sup
 	if err != nil {
 		return nil, fmt.Errorf("create user - %w", err)
 	}
+	defer func() {
+		// delete admin on any error to avoid partial setup
+		if err != nil {
+			logger := logging.ContextUnavailable()
+			logger.WithError(err).Warn("Failed to create admin user, deleting user")
+			if delUserErr := authService.DeleteUser(ctx, user.Username); delUserErr != nil {
+				logger.WithError(delUserErr).Error("Failed to delete user")
+			}
+		}
+	}()
 
 	if addToAdmins {
 		// verify the admin group exists
diff --git a/pkg/auth/setup/setup_test.go b/pkg/auth/setup/setup_test.go
new file mode 100644
index 00000000000..ee1f4b02cec
--- /dev/null
+++ b/pkg/auth/setup/setup_test.go
@@ -0,0 +1,80 @@
+package setup_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/treeverse/lakefs/pkg/auth"
+	"github.com/treeverse/lakefs/pkg/auth/crypt"
+	"github.com/treeverse/lakefs/pkg/auth/model"
+	authparams "github.com/treeverse/lakefs/pkg/auth/params"
+	"github.com/treeverse/lakefs/pkg/auth/setup"
+	"github.com/treeverse/lakefs/pkg/kv"
+	"github.com/treeverse/lakefs/pkg/kv/kvtest"
+	"github.com/treeverse/lakefs/pkg/logging"
+)
+
+func SetupService(t *testing.T, secret string) (*auth.BasicAuthService, kv.Store) {
+	t.Helper()
+	kvStore := kvtest.GetStore(context.Background(), t)
+	return auth.NewBasicAuthService(kvStore, crypt.NewSecretStore([]byte(secret)), authparams.ServiceCache{
+		Enabled: false,
+	}, logging.ContextUnavailable()), kvStore
+}
+
+func TestAddAdminUser(t *testing.T) {
+	authService, _ := SetupService(t, "secret")
+	ctx := context.Background()
+
+	// setup admin user
+	superuser := &model.SuperuserConfiguration{
+		User: model.User{
+			CreatedAt: time.Now(),
+			Username:  "testuser",
+		},
+		AccessKeyID:     "key",
+		SecretAccessKey: "secret",
+	}
+	_, err := setup.AddAdminUser(ctx, authService, superuser, false)
+	if err != nil {
+		t.Fatal("failed to add admin user:", err)
+	}
+
+	// try to add another admin user - should fail
+	superuser2 := &model.SuperuserConfiguration{
+		User: model.User{
+			CreatedAt: time.Now(),
+			Username:  "testuser2",
+		},
+		AccessKeyID:     "key2",
+		SecretAccessKey: "secret2",
+	}
+	_, err = setup.AddAdminUser(ctx, authService, superuser2, false)
+	expectedErr := auth.ErrAlreadyExists
+	if !errors.Is(err, expectedErr) {
+		t.Fatalf("adding another admin. err:%v, expected:%v", err, expectedErr)
+	}
+
+	// delete admin user should work
+	err = authService.DeleteUser(ctx, superuser.User.Username)
+	if err != nil {
+		t.Fatal("failed to delete admin user:", err)
+	}
+
+	// setup admin user again without secret access key should fail
+	superuser.SecretAccessKey = ""
+	_, err = setup.AddAdminUser(ctx, authService, superuser, false)
+	expectedErr = auth.ErrNotFound
+	if !errors.Is(err, expectedErr) {
+		t.Fatalf("adding admin without secret access. key err:%v, expected:%v", err, expectedErr)
+	}
+
+	// we should be able to re-add the admin user
+	superuser.SecretAccessKey = "new-secret"
+	_, err = setup.AddAdminUser(ctx, authService, superuser, false)
+	if err != nil {
+		t.Fatal("failed to re-add admin user:", err)
+	}
+}