reliability fixes and disconnect-all command

management/cmd/admin/admin.go

@@ -11,34 +11,49 @@ import (
 	"log/slog"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/dexidp/dex/storage"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.org/x/crypto/bcrypt"
 
+	"github.com/netbirdio/netbird/formatter/hook"
 	nbdex "github.com/netbirdio/netbird/idp/dex"
+	"github.com/netbirdio/netbird/management/cmd/proxy"
+	"github.com/netbirdio/netbird/management/cmd/token"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
 )
 
-const (
-	localConnectorID           = "local"
-	dashboardClientID          = "netbird-dashboard"
-	cliClientID                = "netbird-cli"
-	defaultTOTPAuthenticatorID = "default-totp"
-)
-
 // Resources contains the storages required by the admin commands.
 type Resources struct {
-	Store      store.Store
-	IDPStorage storage.Storage
+	Store          store.Store
+	IDPStorage     storage.Storage
+	IDPStorageFile string
+	EventStore     activity.Store
 }
 
 // Opener initializes command resources from the command context and calls fn.
 type Opener func(cmd *cobra.Command, fn func(ctx context.Context, resources Resources) error) error
 
+// StoreOpener initializes only the management store from the command context and calls fn.
+type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
+
+// IDPOpener initializes only the embedded IdP storage from the command context and calls fn.
+type IDPOpener func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error
+
+// Openers contains the resource openers needed by the admin command tree.
+type Openers struct {
+	Resources Opener
+	Store     StoreOpener
+	IDP       IDPOpener
+}
+
 type userSelector struct {
 	email  string
 	userID string
@@ -59,8 +74,8 @@ func (s userSelector) validate() error {
 	return nil
 }
 
-// NewCommands creates the admin command tree with the given resource opener.
-func NewCommands(opener Opener) *cobra.Command {
+// NewCommands creates the admin command tree with the given resource openers.
+func NewCommands(openers Openers) *cobra.Command {
 	adminCmd := &cobra.Command{
 		Use:   "admin",
 		Short: "Self-hosted administrator helpers",
@@ -81,12 +96,15 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short:   "Change a local user's password",
 		Args:    cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
+			if err := passwordSelector.validate(); err != nil {
+				return err
+			}
 			newPassword, err := resolvePasswordInput(cmd, password, passwordFile)
 			if err != nil {
 				return err
 			}
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
-				return runChangePassword(ctx, resources.IDPStorage, cmd.OutOrStdout(), passwordSelector, newPassword)
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runChangePassword(ctx, idpStorage, cmd.OutOrStdout(), passwordSelector, newPassword, storageFile)
 			})
 		},
 	}
@@ -100,8 +118,11 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Reset a local user's MFA enrollment",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
-				return runResetMFA(ctx, resources.IDPStorage, cmd.OutOrStdout(), resetSelector)
+			if err := resetSelector.validate(); err != nil {
+				return err
+			}
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runResetMFA(ctx, idpStorage, cmd.OutOrStdout(), resetSelector, storageFile)
 			})
 		},
 	}
@@ -119,7 +140,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Enable MFA for local embedded IdP users",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), true)
 			})
 		},
@@ -130,7 +151,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Disable MFA for local embedded IdP users",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), false)
 			})
 		},
@@ -141,7 +162,7 @@ func NewCommands(opener Opener) *cobra.Command {
 		Short: "Show local MFA status",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			return opener(cmd, func(ctx context.Context, resources Resources) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
 				return runMFAStatus(ctx, resources, cmd.OutOrStdout())
 			})
 		},
@@ -149,6 +170,10 @@ func NewCommands(opener Opener) *cobra.Command {
 
 	mfaCmd.AddCommand(enableCmd, disableCmd, statusCmd)
 	adminCmd.AddCommand(userCmd, mfaCmd)
+	if openers.Store != nil {
+		adminCmd.AddCommand(tokencmd.NewCommands(tokencmd.StoreOpener(openers.Store)))
+		adminCmd.AddCommand(proxycmd.NewCommands(proxycmd.StoreOpener(openers.Store)))
+	}
 	return adminCmd
 }
 
@@ -171,6 +196,45 @@ func OpenEmbeddedIDPStorage(cfg *idp.EmbeddedIdPConfig) (storage.Storage, error)
 	return st, nil
 }
 
+// CloseStore closes the management store and logs cleanup errors at debug level.
+func CloseStore(ctx context.Context, s store.Store) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(ctx); err != nil {
+		log.Debugf("close store: %v", err)
+	}
+}
+
+// OpenIDPStorage opens embedded IdP storage and returns its sqlite file path when applicable.
+func OpenIDPStorage(config *nbconfig.Config) (storage.Storage, string, error) {
+	if config == nil {
+		return nil, "", fmt.Errorf("management config is required")
+	}
+	idpStorage, err := OpenEmbeddedIDPStorage(config.EmbeddedIdP)
+	if err != nil {
+		return nil, "", err
+	}
+	return idpStorage, embeddedIDPStorageFile(config), nil
+}
+
+func embeddedIDPStorageFile(config *nbconfig.Config) string {
+	if config.EmbeddedIdP == nil || config.EmbeddedIdP.Storage.Type != "sqlite3" {
+		return ""
+	}
+	return config.EmbeddedIdP.Storage.Config.File
+}
+
+// CloseIDPStorage closes embedded IdP storage and logs cleanup errors at debug level.
+func CloseIDPStorage(s storage.Storage) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(); err != nil {
+		log.Debugf("close embedded IdP storage: %v", err)
+	}
+}
+
 func addUserSelectorFlags(cmd *cobra.Command, selector *userSelector) {
 	cmd.Flags().StringVar(&selector.email, "email", "", "User email")
 	cmd.Flags().StringVar(&selector.userID, "user-id", "", "User ID")
@@ -197,7 +261,7 @@ func resolvePasswordInput(cmd *cobra.Command, password, passwordFile string) (st
 	return strings.TrimRight(string(data), "\r\n"), nil
 }
 
-func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string) error {
+func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string, idpStorageFile string) error {
 	if idpStorage == nil {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
@@ -212,7 +276,7 @@ func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Wri
 		return fmt.Errorf("invalid password: %w", err)
 	}
 
-	user, err := findLocalUser(ctx, idpStorage, selector)
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
 	if err != nil {
 		return err
 	}
@@ -237,7 +301,7 @@ func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Wri
 	return nil
 }
 
-func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector) error {
+func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, idpStorageFile string) error {
 	if idpStorage == nil {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
@@ -246,13 +310,13 @@ func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, s
 		return err
 	}
 
-	user, err := findLocalUser(ctx, idpStorage, selector)
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
 	if err != nil {
 		return err
 	}
 
 	reset := false
-	err = idpStorage.UpdateUserIdentity(ctx, user.UserID, localConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
+	err = idpStorage.UpdateUserIdentity(ctx, user.UserID, idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
 		reset = reset || len(old.MFASecrets) > 0 || len(old.WebAuthnCredentials) > 0
 		old.MFASecrets = map[string]*storage.MFASecret{}
 		old.WebAuthnCredentials = map[string][]storage.WebAuthnCredential{}
@@ -289,24 +353,30 @@ func runSetMFAEnabled(ctx context.Context, resources Resources, w io.Writer, ena
 		return fmt.Errorf("embedded IdP storage is required")
 	}
 
-	accounts := resources.Store.GetAllAccounts(ctx)
-	if len(accounts) != 1 {
-		return fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", len(accounts))
+	accountID, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
 	}
 
-	settings := &types.Settings{}
-	if accounts[0].Settings != nil {
-		settings = accounts[0].Settings.Copy()
-	}
-	settings.LocalMfaEnabled = enabled
-	if err := resources.Store.SaveAccountSettings(ctx, accounts[0].Id, settings); err != nil {
-		return fmt.Errorf("save local MFA account setting: %w", err)
-	}
+	oldEnabled := settings.LocalMfaEnabled
+	newSettings := settings.Copy()
+	newSettings.LocalMfaEnabled = enabled
 
 	if err := setIDPClientsMFA(ctx, resources.IDPStorage, enabled); err != nil {
 		return err
 	}
 
+	if err := resources.Store.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
+		if rollbackErr := setIDPClientsMFA(ctx, resources.IDPStorage, oldEnabled); rollbackErr != nil {
+			return fmt.Errorf("save local MFA account setting: %w (also failed to roll back embedded IdP MFA state: %v)", err, rollbackErr)
+		}
+		return fmt.Errorf("save local MFA account setting: %w", err)
+	}
+
+	if err := storeMFAActivity(ctx, resources.EventStore, accountID, enabled); err != nil {
+		_, _ = fmt.Fprintf(w, "Warning: failed to record audit event: %v\n", err)
+	}
+
 	state := "disabled"
 	if enabled {
 		state = "enabled"
@@ -323,13 +393,13 @@ func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
 		return fmt.Errorf("embedded IdP storage is required")
 	}
 
-	accounts := resources.Store.GetAllAccounts(ctx)
-	accountStatus := "unknown"
-	if len(accounts) == 1 && accounts[0].Settings != nil {
-		accountStatus = "disabled"
-		if accounts[0].Settings.LocalMfaEnabled {
-			accountStatus = "enabled"
-		}
+	_, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
+	}
+	accountStatus := "disabled"
+	if settings.LocalMfaEnabled {
+		accountStatus = "enabled"
 	}
 
 	clientStatus, err := idpClientsMFAStatus(ctx, resources.IDPStorage)
@@ -342,7 +412,52 @@ func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
 	return nil
 }
 
-func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector) (storage.Password, error) {
+func getSingleAccountSettings(ctx context.Context, s store.Store) (string, *types.Settings, error) {
+	count, err := s.GetAccountsCounter(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("count accounts: %w", err)
+	}
+	if count != 1 {
+		return "", nil, fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", count)
+	}
+
+	accountID, err := s.GetAnyAccountID(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account ID: %w", err)
+	}
+
+	settings, err := s.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account settings: %w", err)
+	}
+	if settings == nil {
+		settings = &types.Settings{}
+	}
+	return accountID, settings, nil
+}
+
+func storeMFAActivity(ctx context.Context, eventStore activity.Store, accountID string, enabled bool) error {
+	if eventStore == nil {
+		return nil
+	}
+	event := activity.AccountLocalMfaDisabled
+	if enabled {
+		event = activity.AccountLocalMfaEnabled
+	}
+	_, err := eventStore.Save(ctx, &activity.Event{
+		Timestamp:   time.Now().UTC(),
+		Activity:    event,
+		InitiatorID: string(hook.SystemSource),
+		TargetID:    accountID,
+		AccountID:   accountID,
+	})
+	if err != nil {
+		return fmt.Errorf("save local MFA audit event: %w", err)
+	}
+	return nil
+}
+
+func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector, idpStorageFile string) (storage.Password, error) {
 	selector = selector.normalized()
 	if err := selector.validate(); err != nil {
 		return storage.Password{}, err
@@ -351,6 +466,11 @@ func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector use
 	if selector.email != "" {
 		user, err := idpStorage.GetPassword(ctx, selector.email)
 		if errors.Is(err, storage.ErrNotFound) {
+			if empty, listErr := localUsersEmpty(ctx, idpStorage); listErr != nil {
+				return storage.Password{}, listErr
+			} else if empty {
+				return storage.Password{}, noLocalUsersError(idpStorageFile)
+			}
 			return storage.Password{}, fmt.Errorf("local user with email %q not found", selector.email)
 		}
 		if err != nil {
@@ -374,11 +494,31 @@ func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector use
 		}
 	}
 
+	if len(users) == 0 {
+		return storage.Password{}, noLocalUsersError(idpStorageFile)
+	}
+
 	return storage.Password{}, fmt.Errorf("local user with ID %q not found", selector.userID)
 }
 
+func localUsersEmpty(ctx context.Context, idpStorage storage.Storage) (bool, error) {
+	users, err := idpStorage.ListPasswords(ctx)
+	if err != nil {
+		return false, fmt.Errorf("list local users: %w", err)
+	}
+	return len(users) == 0, nil
+}
+
+func noLocalUsersError(idpStorageFile string) error {
+	location := ""
+	if idpStorageFile != "" {
+		location = fmt.Sprintf(" (%s)", idpStorageFile)
+	}
+	return fmt.Errorf("no local users exist in the embedded IdP storage%s; the management server may never have started with this config, or --datadir points at the wrong location", location)
+}
+
 func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, userID string) error {
-	err := idpStorage.DeleteAuthSession(ctx, userID, localConnectorID)
+	err := idpStorage.DeleteAuthSession(ctx, userID, idp.LocalConnectorID)
 	if err == nil || errors.Is(err, storage.ErrNotFound) {
 		return nil
 	}
@@ -388,25 +528,21 @@ func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, use
 func setIDPClientsMFA(ctx context.Context, idpStorage storage.Storage, enabled bool) error {
 	var mfaChain []string
 	if enabled {
-		mfaChain = []string{defaultTOTPAuthenticatorID}
+		mfaChain = []string{idp.DefaultTOTPAuthenticatorID}
 	}
 
-	for _, clientID := range []string{cliClientID, dashboardClientID} {
-		if err := idpStorage.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
-			old.MFAChain = mfaChain
-			return old, nil
-		}); err != nil {
-			if errors.Is(err, storage.ErrNotFound) {
-				return fmt.Errorf("embedded IdP client %q not found; start the management server once before toggling MFA", clientID)
-			}
-			return fmt.Errorf("update MFA chain on embedded IdP client %q: %w", clientID, err)
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
+	if err := nbdex.SetClientsMFAChain(ctx, idpStorage, clientIDs, mfaChain); err != nil {
+		if errors.Is(err, storage.ErrNotFound) {
+			return fmt.Errorf("embedded IdP client not found; start the management server once before toggling MFA: %w", err)
 		}
+		return fmt.Errorf("update MFA chain on embedded IdP clients: %w", err)
 	}
 	return nil
 }
 
 func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (string, error) {
-	clientIDs := []string{cliClientID, dashboardClientID}
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
 	enabledCount := 0
 	for _, clientID := range clientIDs {
 		client, err := idpStorage.GetClient(ctx, clientID)
@@ -416,7 +552,7 @@ func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (strin
 		if err != nil {
 			return "unknown", fmt.Errorf("get embedded IdP client %q: %w", clientID, err)
 		}
-		if hasAuthenticator(client.MFAChain, defaultTOTPAuthenticatorID) {
+		if hasAuthenticator(client.MFAChain, idp.DefaultTOTPAuthenticatorID) {
 			enabledCount++
 		}
 	}