Refactor gRPC auth and retry logic

- Centralize retry logic in auth layer
- Decouple gRPC connection logic with new Connect method
- Refactor management client to fetch server public key internally
- Add dedicated HealthCheck method for connection verification
- Simplify getServerPublicKey by removing retry logic

client/android/login.go

@@ -3,15 +3,7 @@ package android
 import (
 	"context"
 	"fmt"
-	"time"
 
-	"github.com/cenkalti/backoff/v4"
-	log "github.com/sirupsen/logrus"
-	"google.golang.org/grpc/codes"
-	gstatus "google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/cmd"
-	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/system"
@@ -84,34 +76,21 @@ func (a *Auth) SaveConfigIfSSOSupported(listener SSOListener) {
 }
 
 func (a *Auth) saveConfigIfSSOSupported() (bool, error) {
-	supportsSSO := true
-	err := a.withBackOff(a.ctx, func() (err error) {
-		_, err = internal.GetPKCEAuthorizationFlowInfo(a.ctx, a.config.PrivateKey, a.config.ManagementURL, nil)
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.NotFound || s.Code() == codes.Unimplemented) {
-			_, err = internal.GetDeviceAuthorizationFlowInfo(a.ctx, a.config.PrivateKey, a.config.ManagementURL)
-			s, ok := gstatus.FromError(err)
-			if !ok {
-				return err
-			}
-			if s.Code() == codes.NotFound || s.Code() == codes.Unimplemented {
-				supportsSSO = false
-				err = nil
-			}
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
+	if err != nil {
+		return false, fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
 
-			return err
-		}
-
-		return err
-	})
+	supportsSSO, err := authClient.IsSSOSupported(a.ctx)
+	if err != nil {
+		return false, fmt.Errorf("failed to check SSO support: %v", err)
+	}
 
 	if !supportsSSO {
 		return false, nil
 	}
 
-	if err != nil {
-		return false, fmt.Errorf("backoff cycle failed: %v", err)
-	}
-
 	err = profilemanager.WriteOutConfig(a.cfgPath, a.config)
 	return true, err
 }
@@ -131,17 +110,15 @@ func (a *Auth) LoginWithSetupKeyAndSaveConfig(resultListener ErrListener, setupK
 func (a *Auth) loginWithSetupKeyAndSaveConfig(setupKey string, deviceName string) error {
 	//nolint
 	ctxWithValues := context.WithValue(a.ctx, system.DeviceNameCtxKey, deviceName)
-
-	err := a.withBackOff(a.ctx, func() error {
-		backoffErr := internal.Login(ctxWithValues, a.config, setupKey, "")
-		if s, ok := gstatus.FromError(backoffErr); ok && (s.Code() == codes.PermissionDenied) {
-			// we got an answer from management, exit backoff earlier
-			return backoff.Permanent(backoffErr)
-		}
-		return backoffErr
-	})
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
+
+	err = authClient.Login(ctxWithValues, setupKey, "")
+	if err != nil {
+		return fmt.Errorf("login failed: %v", err)
 	}
 
 	return profilemanager.WriteOutConfig(a.cfgPath, a.config)
@@ -160,15 +137,16 @@ func (a *Auth) Login(resultListener ErrListener, urlOpener URLOpener, isAndroidT
 }
 
 func (a *Auth) login(urlOpener URLOpener, isAndroidTV bool) error {
-	var needsLogin bool
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
+	if err != nil {
+		return fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
 
 	// check if we need to generate JWT token
-	err := a.withBackOff(a.ctx, func() (err error) {
-		needsLogin, err = internal.IsLoginRequired(a.ctx, a.config)
-		return
-	})
+	needsLogin, err := authClient.IsLoginRequired(a.ctx)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("failed to check login requirement: %v", err)
 	}
 
 	jwtToken := ""
@@ -180,22 +158,13 @@ func (a *Auth) login(urlOpener URLOpener, isAndroidTV bool) error {
 		jwtToken = tokenInfo.GetTokenToUse()
 	}
 
-	err = a.withBackOff(a.ctx, func() error {
-		err := internal.Login(a.ctx, a.config, "", jwtToken)
-
-		if err == nil {
-			go urlOpener.OnLoginSuccess()
-		}
-
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
-			return nil
-		}
-		return err
-	})
+	err = authClient.Login(a.ctx, "", jwtToken)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("login failed: %v", err)
 	}
 
+	go urlOpener.OnLoginSuccess()
+
 	return nil
 }
 
@@ -212,22 +181,10 @@ func (a *Auth) foregroundGetTokenInfo(urlOpener URLOpener, isAndroidTV bool) (*a
 
 	go urlOpener.Open(flowInfo.VerificationURIComplete, flowInfo.UserCode)
 
-	waitTimeout := time.Duration(flowInfo.ExpiresIn) * time.Second
-	waitCTX, cancel := context.WithTimeout(a.ctx, waitTimeout)
-	defer cancel()
-	tokenInfo, err := oAuthFlow.WaitToken(waitCTX, flowInfo)
+	tokenInfo, err := oAuthFlow.WaitToken(a.ctx, flowInfo)
 	if err != nil {
 		return nil, fmt.Errorf("waiting for browser login failed: %v", err)
 	}
 
 	return &tokenInfo, nil
 }
-
-func (a *Auth) withBackOff(ctx context.Context, bf func() error) error {
-	return backoff.RetryNotify(
-		bf,
-		backoff.WithContext(cmd.CLIBackOffSettings, ctx),
-		func(err error, duration time.Duration) {
-			log.Warnf("retrying Login to the Management service in %v due to error %v", duration, err)
-		})
-}

client/cmd/login.go

@@ -2,14 +2,16 @@ package cmd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
+	"os/exec"
 	"os/user"
 	"runtime"
 	"strings"
-	"time"
 
 	log "github.com/sirupsen/logrus"
+	"github.com/skratchdot/open-golang/open"
 	"github.com/spf13/cobra"
 	"google.golang.org/grpc/codes"
 	gstatus "google.golang.org/grpc/status"
@@ -19,6 +21,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/system"
+	mgm "github.com/netbirdio/netbird/shared/management/client"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -275,18 +278,19 @@ func handleSSOLogin(ctx context.Context, cmd *cobra.Command, loginResp *proto.Lo
 }
 
 func foregroundLogin(ctx context.Context, cmd *cobra.Command, config *profilemanager.Config, setupKey, profileName string) error {
+	authClient, err := auth.NewAuth(ctx, config.PrivateKey, config.ManagementURL, config)
+	if err != nil {
+		return fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
+
 	needsLogin := false
 
-	err := WithBackOff(func() error {
-		err := internal.Login(ctx, config, "", "")
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
-			needsLogin = true
-			return nil
-		}
-		return err
-	})
-	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+	err = authClient.Login(ctx, "", "")
+	if errors.Is(err, mgm.ErrPermissionDenied) || errors.Is(err, mgm.ErrInvalidArgument) || errors.Is(err, mgm.ErrUnauthenticated) {
+		needsLogin = true
+	} else if err != nil {
+		return fmt.Errorf("login check failed: %v", err)
 	}
 
 	jwtToken := ""
@@ -298,23 +302,9 @@ func foregroundLogin(ctx context.Context, cmd *cobra.Command, config *profileman
 		jwtToken = tokenInfo.GetTokenToUse()
 	}
 
-	var lastError error
-
-	err = WithBackOff(func() error {
-		err := internal.Login(ctx, config, setupKey, jwtToken)
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
-			lastError = err
-			return nil
-		}
-		return err
-	})
-
-	if lastError != nil {
-		return fmt.Errorf("login failed: %v", lastError)
-	}
-
+	err = authClient.Login(ctx, setupKey, jwtToken)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("login failed: %v", err)
 	}
 
 	return nil
@@ -342,11 +332,7 @@ func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *pro
 
 	openURL(cmd, flowInfo.VerificationURIComplete, flowInfo.UserCode, noBrowser)
 
-	waitTimeout := time.Duration(flowInfo.ExpiresIn) * time.Second
-	waitCTX, c := context.WithTimeout(context.TODO(), waitTimeout)
-	defer c()
-
-	tokenInfo, err := oAuthFlow.WaitToken(waitCTX, flowInfo)
+	tokenInfo, err := oAuthFlow.WaitToken(context.TODO(), flowInfo)
 	if err != nil {
 		return nil, fmt.Errorf("waiting for browser login failed: %v", err)
 	}
@@ -371,13 +357,21 @@ func openURL(cmd *cobra.Command, verificationURIComplete, userCode string, noBro
 	cmd.Println("")
 
 	if !noBrowser {
-		if err := util.OpenBrowser(verificationURIComplete); err != nil {
+		if err := openBrowser(verificationURIComplete); err != nil {
 			cmd.Println("\nAlternatively, you may want to use a setup key, see:\n\n" +
 				"https://docs.netbird.io/how-to/register-machines-using-setup-keys")
 		}
 	}
 }
 
+// openBrowser opens the URL in a browser, respecting the BROWSER environment variable.
+func openBrowser(url string) error {
+	if browser := os.Getenv("BROWSER"); browser != "" {
+		return exec.Command(browser, url).Start()
+	}
+	return open.Run(url)
+}
+
 // isUnixRunningDesktop checks if a Linux OS is running desktop environment
 func isUnixRunningDesktop() bool {
 	if runtime.GOOS != "linux" && runtime.GOOS != "freebsd" {

client/cmd/ssh.go

@@ -51,7 +51,6 @@ var (
 	identityFile          string
 	skipCachedToken       bool
 	requestPTY            bool
-	sshNoBrowser          bool
 )
 
 var (
@@ -82,7 +81,6 @@ func init() {
 	sshCmd.PersistentFlags().StringVarP(&identityFile, "identity", "i", "", "Path to SSH private key file (deprecated)")
 	_ = sshCmd.PersistentFlags().MarkDeprecated("identity", "this flag is no longer used")
 	sshCmd.PersistentFlags().BoolVar(&skipCachedToken, "no-cache", false, "Skip cached JWT token and force fresh authentication")
-	sshCmd.PersistentFlags().BoolVar(&sshNoBrowser, noBrowserFlag, false, noBrowserDesc)
 
 	sshCmd.PersistentFlags().StringArrayP("L", "L", []string{}, "Local port forwarding [bind_address:]port:host:hostport")
 	sshCmd.PersistentFlags().StringArrayP("R", "R", []string{}, "Remote port forwarding [bind_address:]port:host:hostport")
@@ -187,21 +185,6 @@ func getEnvOrDefault(flagName, defaultValue string) string {
 	return defaultValue
 }
 
-// getBoolEnvOrDefault checks for boolean environment variables with WT_ and NB_ prefixes
-func getBoolEnvOrDefault(flagName string, defaultValue bool) bool {
-	if envValue := os.Getenv("WT_" + flagName); envValue != "" {
-		if parsed, err := strconv.ParseBool(envValue); err == nil {
-			return parsed
-		}
-	}
-	if envValue := os.Getenv("NB_" + flagName); envValue != "" {
-		if parsed, err := strconv.ParseBool(envValue); err == nil {
-			return parsed
-		}
-	}
-	return defaultValue
-}
-
 // resetSSHGlobals sets SSH globals to their default values
 func resetSSHGlobals() {
 	port = sshserver.DefaultSSHPort
@@ -213,7 +196,6 @@ func resetSSHGlobals() {
 	strictHostKeyChecking = true
 	knownHostsFile = ""
 	identityFile = ""
-	sshNoBrowser = false
 }
 
 // parseCustomSSHFlags extracts -L, -R flags and returns filtered args
@@ -388,7 +370,6 @@ type sshFlags struct {
 	KnownHostsFile        string
 	IdentityFile          string
 	SkipCachedToken       bool
-	NoBrowser             bool
 	ConfigPath            string
 	LogLevel              string
 	LocalForwards         []string
@@ -400,7 +381,6 @@ type sshFlags struct {
 func createSSHFlagSet() (*flag.FlagSet, *sshFlags) {
 	defaultConfigPath := getEnvOrDefault("CONFIG", configPath)
 	defaultLogLevel := getEnvOrDefault("LOG_LEVEL", logLevel)
-	defaultNoBrowser := getBoolEnvOrDefault("NO_BROWSER", false)
 
 	fs := flag.NewFlagSet("ssh-flags", flag.ContinueOnError)
 	fs.SetOutput(nil)
@@ -421,7 +401,6 @@ func createSSHFlagSet() (*flag.FlagSet, *sshFlags) {
 	fs.StringVar(&flags.IdentityFile, "i", "", "Path to SSH private key file")
 	fs.StringVar(&flags.IdentityFile, "identity", "", "Path to SSH private key file")
 	fs.BoolVar(&flags.SkipCachedToken, "no-cache", false, "Skip cached JWT token and force fresh authentication")
-	fs.BoolVar(&flags.NoBrowser, "no-browser", defaultNoBrowser, noBrowserDesc)
 
 	fs.StringVar(&flags.ConfigPath, "c", defaultConfigPath, "Netbird config file location")
 	fs.StringVar(&flags.ConfigPath, "config", defaultConfigPath, "Netbird config file location")
@@ -470,7 +449,6 @@ func validateSSHArgsWithoutFlagParsing(_ *cobra.Command, args []string) error {
 	knownHostsFile = flags.KnownHostsFile
 	identityFile = flags.IdentityFile
 	skipCachedToken = flags.SkipCachedToken
-	sshNoBrowser = flags.NoBrowser
 
 	if flags.ConfigPath != getEnvOrDefault("CONFIG", configPath) {
 		configPath = flags.ConfigPath
@@ -530,7 +508,6 @@ func runSSH(ctx context.Context, addr string, cmd *cobra.Command) error {
 		DaemonAddr:         daemonAddr,
 		SkipCachedToken:    skipCachedToken,
 		InsecureSkipVerify: !strictHostKeyChecking,
-		NoBrowser:          sshNoBrowser,
 	})
 
 	if err != nil {
@@ -786,15 +763,7 @@ func sshProxyFn(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("invalid port: %s", portStr)
 	}
 
-	// Check env var for browser setting since this command is invoked via SSH ProxyCommand
-	// where command-line flags cannot be passed. Default is to open browser.
-	noBrowser := getBoolEnvOrDefault("NO_BROWSER", false)
-	var browserOpener func(string) error
-	if !noBrowser {
-		browserOpener = util.OpenBrowser
-	}
-
-	proxy, err := sshproxy.New(daemonAddr, host, port, cmd.ErrOrStderr(), browserOpener)
+	proxy, err := sshproxy.New(daemonAddr, host, port, cmd.ErrOrStderr())
 	if err != nil {
 		return fmt.Errorf("create SSH proxy: %w", err)
 	}

client/cmd/testutil_test.go

@@ -15,8 +15,6 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 
 	clientProto "github.com/netbirdio/netbird/client/proto"
@@ -26,6 +24,8 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/groups"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/peers"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -116,18 +116,15 @@ func startManagement(t *testing.T, config *config.Config, testFile string) (*grp
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := mgmt.NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, mgmt.MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), manager.NewEphemeralManager(store, peersmanager), config)
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, mgmt.MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), config)
 
 	accountManager, err := mgmt.BuildManager(context.Background(), config, store, networkMapController, nil, "", eventStore, nil, false, iv, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock, false)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		t.Fatal(err)
-	}
-	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, secretsManager, nil, nil, &mgmt.MockIntegratedValidator{}, networkMapController)
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
+	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, updateManager, secretsManager, nil, &manager.EphemeralManager{}, nil, &mgmt.MockIntegratedValidator{}, networkMapController)
 	if err != nil {
 		t.Fatal(err)
 	}

client/embed/embed.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/iface/netstack"
 	"github.com/netbirdio/netbird/client/internal"
+	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	sshcommon "github.com/netbirdio/netbird/client/ssh"
@@ -168,7 +169,13 @@ func (c *Client) Start(startCtx context.Context) error {
 	ctx := internal.CtxInitState(context.Background())
 	// nolint:staticcheck
 	ctx = context.WithValue(ctx, system.DeviceNameCtxKey, c.deviceName)
-	if err := internal.Login(ctx, c.config, c.setupKey, c.jwtToken); err != nil {
+	authClient, err := auth.NewAuth(ctx, c.config.PrivateKey, c.config.ManagementURL, c.config)
+	if err != nil {
+		return fmt.Errorf("create auth client: %w", err)
+	}
+	defer authClient.Close()
+
+	if err := authClient.Login(ctx, c.setupKey, c.jwtToken); err != nil {
 		return fmt.Errorf("login: %w", err)
 	}
 

client/internal/auth/auth.go

@@ -0,0 +1,287 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/url"
+
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/ssh"
+	"github.com/netbirdio/netbird/client/system"
+	mgm "github.com/netbirdio/netbird/shared/management/client"
+	"github.com/netbirdio/netbird/shared/management/client/common"
+)
+
+// Auth manages authentication operations with the management server
+// The underlying management client handles connection retry and reconnection automatically
+type Auth struct {
+	client *mgm.GrpcClient
+	config *profilemanager.Config
+}
+
+// NewAuth creates a new Auth instance that manages authentication flows
+// It establishes a connection to the management server that will be reused for all operations
+// The management client handles connection retry and reconnection automatically
+func NewAuth(ctx context.Context, privateKey string, mgmURL *url.URL, config *profilemanager.Config) (*Auth, error) {
+	// Validate WireGuard private key
+	myPrivateKey, err := wgtypes.ParseKey(privateKey)
+	if err != nil {
+		log.Errorf("failed parsing Wireguard key %s: [%s]", privateKey, err.Error())
+		return nil, err
+	}
+
+	// Determine TLS setting based on URL scheme
+	mgmTLSEnabled := mgmURL.Scheme == "https"
+
+	log.Debugf("connecting to Management Service %s", mgmURL.String())
+	mgmClient := mgm.NewClient(mgmURL.Host, myPrivateKey, mgmTLSEnabled)
+	if err := mgmClient.Connect(ctx); err != nil {
+		log.Errorf("failed connecting to Management Service %s: %v", mgmURL.String(), err)
+		return nil, err
+	}
+
+	log.Debugf("connected to the Management service %s", mgmURL.String())
+
+	return &Auth{
+		client: mgmClient,
+		config: config,
+	}, nil
+}
+
+// Close closes the management client connection
+func (a *Auth) Close() error {
+	if a.client == nil {
+		return nil
+	}
+	return a.client.Close()
+}
+
+// IsSSOSupported checks if the management server supports SSO by attempting to retrieve auth flow configurations.
+// Returns true if either PKCE or Device authorization flow is supported, false otherwise.
+func (a *Auth) IsSSOSupported(ctx context.Context) (bool, error) {
+	// Try PKCE flow first
+	_, err := a.getPKCEFlow(ctx)
+	if err == nil {
+		return true, nil
+	}
+
+	// Check if PKCE is not supported
+	if errors.Is(err, mgm.ErrNotFound) || errors.Is(err, mgm.ErrUnimplemented) {
+		// PKCE not supported, try Device flow
+		_, err = a.getDeviceFlow(ctx)
+		if err == nil {
+			return true, nil
+		}
+
+		// Check if Device flow is also not supported
+		if errors.Is(err, mgm.ErrNotFound) || errors.Is(err, mgm.ErrUnimplemented) {
+			// Neither PKCE nor Device flow is supported
+			return false, nil
+		}
+
+		// Device flow check returned an error other than NotFound/Unimplemented
+		return false, err
+	}
+
+	// PKCE flow check returned an error other than NotFound/Unimplemented
+	return false, err
+}
+
+// IsLoginRequired checks if login is required by attempting to authenticate with the server
+func (a *Auth) IsLoginRequired(ctx context.Context) (bool, error) {
+	pubSSHKey, err := ssh.GeneratePublicKey([]byte(a.config.SSHKey))
+	if err != nil {
+		return false, err
+	}
+
+	err = a.doMgmLogin(ctx, pubSSHKey)
+	if isLoginNeeded(err) {
+		return true, nil
+	}
+
+	return false, err
+}
+
+// Login attempts to log in or register the client with the management server
+// Returns custom errors from mgm package: ErrPermissionDenied, ErrInvalidArgument, ErrUnauthenticated
+func (a *Auth) Login(ctx context.Context, setupKey string, jwtToken string) error {
+	pubSSHKey, err := ssh.GeneratePublicKey([]byte(a.config.SSHKey))
+	if err != nil {
+		return fmt.Errorf("generate SSH public key: %w", err)
+	}
+
+	err = a.doMgmLogin(ctx, pubSSHKey)
+	if isRegistrationNeeded(err) {
+		log.Debugf("peer registration required")
+		return a.registerPeer(ctx, setupKey, jwtToken, pubSSHKey)
+	}
+	return err
+}
+
+// getPKCEFlow retrieves PKCE authorization flow configuration and creates a flow instance
+func (a *Auth) getPKCEFlow(ctx context.Context) (*PKCEAuthorizationFlow, error) {
+	protoFlow, err := a.client.GetPKCEAuthorizationFlow(ctx)
+	if err != nil {
+		if errors.Is(err, mgm.ErrNotFound) {
+			log.Warnf("server couldn't find pkce flow, contact admin: %v", err)
+			return nil, err
+		}
+		log.Errorf("failed to retrieve pkce flow: %v", err)
+		return nil, err
+	}
+
+	protoConfig := protoFlow.GetProviderConfig()
+	config := &PKCEAuthProviderConfig{
+		Audience:              protoConfig.GetAudience(),
+		ClientID:              protoConfig.GetClientID(),
+		ClientSecret:          protoConfig.GetClientSecret(),
+		TokenEndpoint:         protoConfig.GetTokenEndpoint(),
+		AuthorizationEndpoint: protoConfig.GetAuthorizationEndpoint(),
+		Scope:                 protoConfig.GetScope(),
+		RedirectURLs:          protoConfig.GetRedirectURLs(),
+		UseIDToken:            protoConfig.GetUseIDToken(),
+		ClientCertPair:        a.config.ClientCertKeyPair,
+		DisablePromptLogin:    protoConfig.GetDisablePromptLogin(),
+		LoginFlag:             common.LoginFlag(protoConfig.GetLoginFlag()),
+	}
+
+	if err := validatePKCEConfig(config); err != nil {
+		return nil, err
+	}
+
+	flow, err := NewPKCEAuthorizationFlow(*config)
+	if err != nil {
+		return nil, err
+	}
+
+	return flow, nil
+}
+
+// getDeviceFlow retrieves device authorization flow configuration and creates a flow instance
+func (a *Auth) getDeviceFlow(ctx context.Context) (*DeviceAuthorizationFlow, error) {
+	protoFlow, err := a.client.GetDeviceAuthorizationFlow(ctx)
+	if err != nil {
+		if errors.Is(err, mgm.ErrNotFound) {
+			log.Warnf("server couldn't find device flow, contact admin: %v", err)
+			return nil, err
+		}
+		log.Errorf("failed to retrieve device flow: %v", err)
+		return nil, err
+	}
+
+	protoConfig := protoFlow.GetProviderConfig()
+	config := &DeviceAuthProviderConfig{
+		Audience:           protoConfig.GetAudience(),
+		ClientID:           protoConfig.GetClientID(),
+		ClientSecret:       protoConfig.GetClientSecret(),
+		Domain:             protoConfig.Domain,
+		TokenEndpoint:      protoConfig.GetTokenEndpoint(),
+		DeviceAuthEndpoint: protoConfig.GetDeviceAuthEndpoint(),
+		Scope:              protoConfig.GetScope(),
+		UseIDToken:         protoConfig.GetUseIDToken(),
+	}
+
+	// Keep compatibility with older management versions
+	if config.Scope == "" {
+		config.Scope = "openid"
+	}
+
+	if err := validateDeviceAuthConfig(config); err != nil {
+		return nil, err
+	}
+
+	flow, err := NewDeviceAuthorizationFlow(*config)
+	if err != nil {
+		return nil, err
+	}
+
+	return flow, nil
+}
+
+// doMgmLogin performs the actual login operation with the management service
+func (a *Auth) doMgmLogin(ctx context.Context, pubSSHKey []byte) error {
+	sysInfo := system.GetInfo(ctx)
+	sysInfo.SetFlags(
+		a.config.RosenpassEnabled,
+		a.config.RosenpassPermissive,
+		a.config.ServerSSHAllowed,
+		a.config.DisableClientRoutes,
+		a.config.DisableServerRoutes,
+		a.config.DisableDNS,
+		a.config.DisableFirewall,
+		a.config.BlockLANAccess,
+		a.config.BlockInbound,
+		a.config.LazyConnectionEnabled,
+		a.config.EnableSSHRoot,
+		a.config.EnableSSHSFTP,
+		a.config.EnableSSHLocalPortForwarding,
+		a.config.EnableSSHRemotePortForwarding,
+		a.config.DisableSSHAuth,
+	)
+	_, err := a.client.Login(ctx, sysInfo, pubSSHKey, a.config.DNSLabels)
+	return err
+}
+
+// registerPeer checks whether setupKey was provided via cmd line and if not then it prompts user to enter a key.
+// Otherwise tries to register with the provided setupKey via command line.
+func (a *Auth) registerPeer(ctx context.Context, setupKey string, jwtToken string, pubSSHKey []byte) error {
+	validSetupKey, err := uuid.Parse(setupKey)
+	if err != nil && jwtToken == "" {
+		return fmt.Errorf("%w: invalid setup-key or no SSO information provided: %v", mgm.ErrInvalidArgument, err)
+	}
+
+	log.Debugf("sending peer registration request to Management Service")
+	info := system.GetInfo(ctx)
+	info.SetFlags(
+		a.config.RosenpassEnabled,
+		a.config.RosenpassPermissive,
+		a.config.ServerSSHAllowed,
+		a.config.DisableClientRoutes,
+		a.config.DisableServerRoutes,
+		a.config.DisableDNS,
+		a.config.DisableFirewall,
+		a.config.BlockLANAccess,
+		a.config.BlockInbound,
+		a.config.LazyConnectionEnabled,
+		a.config.EnableSSHRoot,
+		a.config.EnableSSHSFTP,
+		a.config.EnableSSHLocalPortForwarding,
+		a.config.EnableSSHRemotePortForwarding,
+		a.config.DisableSSHAuth,
+	)
+
+	// todo: fix error handling of validSetupKey
+	if err := a.client.Register(ctx, validSetupKey.String(), jwtToken, info, pubSSHKey, a.config.DNSLabels); err != nil {
+		log.Errorf("failed registering peer %v", err)
+		return err
+	}
+
+	log.Infof("peer has been successfully registered on Management Service")
+
+	return nil
+}
+
+// isPermissionDenied checks if the error is a PermissionDenied error
+func isPermissionDenied(err error) bool {
+	return errors.Is(err, mgm.ErrPermissionDenied)
+}
+
+// isLoginNeeded checks if the error indicates login is required
+func isLoginNeeded(err error) bool {
+	if err == nil {
+		return false
+	}
+	return errors.Is(err, mgm.ErrInvalidArgument) ||
+		errors.Is(err, mgm.ErrPermissionDenied) ||
+		errors.Is(err, mgm.ErrUnauthenticated)
+}
+
+// isRegistrationNeeded checks if the error indicates peer registration is needed
+func isRegistrationNeeded(err error) bool {
+	return isPermissionDenied(err)
+}

client/internal/auth/device_flow.go

@@ -15,7 +15,6 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/util/embeddedroots"
 )
 
@@ -26,12 +25,56 @@ const (
 
 var _ OAuthFlow = &DeviceAuthorizationFlow{}
 
+// DeviceAuthProviderConfig has all attributes needed to initiate a device authorization flow
+type DeviceAuthProviderConfig struct {
+	// ClientID An IDP application client id
+	ClientID string
+	// ClientSecret An IDP application client secret
+	ClientSecret string
+	// Domain An IDP API domain
+	// Deprecated. Use OIDCConfigEndpoint instead
+	Domain string
+	// Audience An Audience for to authorization validation
+	Audience string
+	// TokenEndpoint is the endpoint of an IDP manager where clients can obtain access token
+	TokenEndpoint string
+	// DeviceAuthEndpoint is the endpoint of an IDP manager where clients can obtain device authorization code
+	DeviceAuthEndpoint string
+	// Scopes provides the scopes to be included in the token request
+	Scope string
+	// UseIDToken indicates if the id token should be used for authentication
+	UseIDToken bool
+	// LoginHint is used to pre-fill the email/username field during authentication
+	LoginHint string
+}
+
+// validateDeviceAuthConfig validates device authorization provider configuration
+func validateDeviceAuthConfig(config *DeviceAuthProviderConfig) error {
+	errorMsgFormat := "invalid provider configuration received from management: %s value is empty. Contact your NetBird administrator"
+
+	if config.Audience == "" {
+		return fmt.Errorf(errorMsgFormat, "Audience")
+	}
+	if config.ClientID == "" {
+		return fmt.Errorf(errorMsgFormat, "Client ID")
+	}
+	if config.TokenEndpoint == "" {
+		return fmt.Errorf(errorMsgFormat, "Token Endpoint")
+	}
+	if config.DeviceAuthEndpoint == "" {
+		return fmt.Errorf(errorMsgFormat, "Device Auth Endpoint")
+	}
+	if config.Scope == "" {
+		return fmt.Errorf(errorMsgFormat, "Device Auth Scopes")
+	}
+	return nil
+}
+
 // DeviceAuthorizationFlow implements the OAuthFlow interface,
 // for the Device Authorization Flow.
 type DeviceAuthorizationFlow struct {
-	providerConfig internal.DeviceAuthProviderConfig
-
-	HTTPClient HTTPClient
+	providerConfig DeviceAuthProviderConfig
+	HTTPClient     HTTPClient
 }
 
 // RequestDeviceCodePayload used for request device code payload for auth0
@@ -57,7 +100,7 @@ type TokenRequestResponse struct {
 }
 
 // NewDeviceAuthorizationFlow returns device authorization flow client
-func NewDeviceAuthorizationFlow(config internal.DeviceAuthProviderConfig) (*DeviceAuthorizationFlow, error) {
+func NewDeviceAuthorizationFlow(config DeviceAuthProviderConfig) (*DeviceAuthorizationFlow, error) {
 	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
 	httpTransport.MaxIdleConns = 5
 
@@ -89,6 +132,11 @@ func (d *DeviceAuthorizationFlow) GetClientID(ctx context.Context) string {
 	return d.providerConfig.ClientID
 }
 
+// SetLoginHint sets the login hint for the device authorization flow
+func (d *DeviceAuthorizationFlow) SetLoginHint(hint string) {
+	d.providerConfig.LoginHint = hint
+}
+
 // RequestAuthInfo requests a device code login flow information from Hosted
 func (d *DeviceAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowInfo, error) {
 	form := url.Values{}
@@ -199,14 +247,22 @@ func (d *DeviceAuthorizationFlow) requestToken(info AuthFlowInfo) (TokenRequestR
 }
 
 // WaitToken waits user's login and authorize the app. Once the user's authorize
-// it retrieves the access token from Hosted's endpoint and validates it before returning
+// it retrieves the access token from Hosted's endpoint and validates it before returning.
+// The method creates a timeout context internally based on info.ExpiresIn.
 func (d *DeviceAuthorizationFlow) WaitToken(ctx context.Context, info AuthFlowInfo) (TokenInfo, error) {
+	// Create timeout context based on flow expiration
+	timeout := time.Duration(info.ExpiresIn) * time.Second
+	waitCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	interval := time.Duration(info.Interval) * time.Second
 	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
 	for {
 		select {
-		case <-ctx.Done():
-			return TokenInfo{}, ctx.Err()
+		case <-waitCtx.Done():
+			return TokenInfo{}, waitCtx.Err()
 		case <-ticker.C:
 
 			tokenResponse, err := d.requestToken(info)

client/internal/auth/device_flow_test.go

@@ -12,8 +12,6 @@ import (
 
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/internal"
 )
 
 type mockHTTPClient struct {
@@ -115,18 +113,19 @@ func TestHosted_RequestDeviceCode(t *testing.T) {
 				err:     testCase.inputReqError,
 			}
 
-			deviceFlow := &DeviceAuthorizationFlow{
-				providerConfig: internal.DeviceAuthProviderConfig{
-					Audience:           expectedAudience,
-					ClientID:           expectedClientID,
-					Scope:              expectedScope,
-					TokenEndpoint:      "test.hosted.com/token",
-					DeviceAuthEndpoint: "test.hosted.com/device/auth",
-					UseIDToken:         false,
-				},
-				HTTPClient: &httpClient,
+			config := DeviceAuthProviderConfig{
+				Audience:           expectedAudience,
+				ClientID:           expectedClientID,
+				Scope:              expectedScope,
+				TokenEndpoint:      "test.hosted.com/token",
+				DeviceAuthEndpoint: "test.hosted.com/device/auth",
+				UseIDToken:         false,
 			}
 
+			deviceFlow, err := NewDeviceAuthorizationFlow(config)
+			require.NoError(t, err, "creating device flow should not fail")
+			deviceFlow.HTTPClient = &httpClient
+
 			authInfo, err := deviceFlow.RequestAuthInfo(context.TODO())
 			testCase.testingErrFunc(t, err, testCase.expectedErrorMSG)
 
@@ -280,18 +279,19 @@ func TestHosted_WaitToken(t *testing.T) {
 				countResBody: testCase.inputCountResBody,
 			}
 
-			deviceFlow := DeviceAuthorizationFlow{
-				providerConfig: internal.DeviceAuthProviderConfig{
-					Audience:           testCase.inputAudience,
-					ClientID:           clientID,
-					TokenEndpoint:      "test.hosted.com/token",
-					DeviceAuthEndpoint: "test.hosted.com/device/auth",
-					Scope:              "openid",
-					UseIDToken:         false,
-				},
-				HTTPClient: &httpClient,
+			config := DeviceAuthProviderConfig{
+				Audience:           testCase.inputAudience,
+				ClientID:           clientID,
+				TokenEndpoint:      "test.hosted.com/token",
+				DeviceAuthEndpoint: "test.hosted.com/device/auth",
+				Scope:              "openid",
+				UseIDToken:         false,
 			}
 
+			deviceFlow, err := NewDeviceAuthorizationFlow(config)
+			require.NoError(t, err, "creating device flow should not fail")
+			deviceFlow.HTTPClient = &httpClient
+
 			ctx, cancel := context.WithTimeout(context.TODO(), testCase.inputTimeout)
 			defer cancel()
 			tokenInfo, err := deviceFlow.WaitToken(ctx, testCase.inputInfo)

client/internal/auth/oauth.go

@@ -10,7 +10,6 @@ import (
 	"google.golang.org/grpc/codes"
 	gstatus "google.golang.org/grpc/status"
 
-	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 )
 
@@ -87,19 +86,33 @@ func NewOAuthFlow(ctx context.Context, config *profilemanager.Config, isUnixDesk
 
 // authenticateWithPKCEFlow initializes the Proof Key for Code Exchange flow auth flow
 func authenticateWithPKCEFlow(ctx context.Context, config *profilemanager.Config, hint string) (OAuthFlow, error) {
-	pkceFlowInfo, err := internal.GetPKCEAuthorizationFlowInfo(ctx, config.PrivateKey, config.ManagementURL, config.ClientCertKeyPair)
+	authClient, err := NewAuth(ctx, config.PrivateKey, config.ManagementURL, config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
+
+	pkceFlowInfo, err := authClient.getPKCEFlow(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("getting pkce authorization flow info failed with error: %v", err)
 	}
 
-	pkceFlowInfo.ProviderConfig.LoginHint = hint
+	if hint != "" {
+		pkceFlowInfo.SetLoginHint(hint)
+	}
 
-	return NewPKCEAuthorizationFlow(pkceFlowInfo.ProviderConfig)
+	return pkceFlowInfo, nil
 }
 
 // authenticateWithDeviceCodeFlow initializes the Device Code auth Flow
 func authenticateWithDeviceCodeFlow(ctx context.Context, config *profilemanager.Config, hint string) (OAuthFlow, error) {
-	deviceFlowInfo, err := internal.GetDeviceAuthorizationFlowInfo(ctx, config.PrivateKey, config.ManagementURL)
+	authClient, err := NewAuth(ctx, config.PrivateKey, config.ManagementURL, config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
+
+	deviceFlowInfo, err := authClient.getDeviceFlow(ctx)
 	if err != nil {
 		switch s, ok := gstatus.FromError(err); {
 		case ok && s.Code() == codes.NotFound:
@@ -114,7 +127,9 @@ func authenticateWithDeviceCodeFlow(ctx context.Context, config *profilemanager.
 		}
 	}
 
-	deviceFlowInfo.ProviderConfig.LoginHint = hint
+	if hint != "" {
+		deviceFlowInfo.SetLoginHint(hint)
+	}
 
-	return NewDeviceAuthorizationFlow(deviceFlowInfo.ProviderConfig)
+	return deviceFlowInfo, nil
 }

client/internal/auth/pkce_flow.go

@@ -13,14 +13,12 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"strconv"
 	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/oauth2"
 
-	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/templates"
 	"github.com/netbirdio/netbird/shared/management/client/common"
 )
@@ -35,23 +33,72 @@ const (
 	defaultPKCETimeoutSeconds = 300
 )
 
+// PKCEAuthProviderConfig has all attributes needed to initiate PKCE authorization flow
+type PKCEAuthProviderConfig struct {
+	// ClientID An IDP application client id
+	ClientID string
+	// ClientSecret An IDP application client secret
+	ClientSecret string
+	// Audience An Audience for to authorization validation
+	Audience string
+	// TokenEndpoint is the endpoint of an IDP manager where clients can obtain access token
+	TokenEndpoint string
+	// AuthorizationEndpoint is the endpoint of an IDP manager where clients can obtain authorization code
+	AuthorizationEndpoint string
+	// Scopes provides the scopes to be included in the token request
+	Scope string
+	// RedirectURL handles authorization code from IDP manager
+	RedirectURLs []string
+	// UseIDToken indicates if the id token should be used for authentication
+	UseIDToken bool
+	// ClientCertPair is used for mTLS authentication to the IDP
+	ClientCertPair *tls.Certificate
+	// DisablePromptLogin makes the PKCE flow to not prompt the user for login
+	DisablePromptLogin bool
+	// LoginFlag is used to configure the PKCE flow login behavior
+	LoginFlag common.LoginFlag
+	// LoginHint is used to pre-fill the email/username field during authentication
+	LoginHint string
+}
+
+// validatePKCEConfig validates PKCE provider configuration
+func validatePKCEConfig(config *PKCEAuthProviderConfig) error {
+	errorMsgFormat := "invalid provider configuration received from management: %s value is empty. Contact your NetBird administrator"
+
+	if config.ClientID == "" {
+		return fmt.Errorf(errorMsgFormat, "Client ID")
+	}
+	if config.TokenEndpoint == "" {
+		return fmt.Errorf(errorMsgFormat, "Token Endpoint")
+	}
+	if config.AuthorizationEndpoint == "" {
+		return fmt.Errorf(errorMsgFormat, "Authorization Auth Endpoint")
+	}
+	if config.Scope == "" {
+		return fmt.Errorf(errorMsgFormat, "PKCE Auth Scopes")
+	}
+	if config.RedirectURLs == nil {
+		return fmt.Errorf(errorMsgFormat, "PKCE Redirect URLs")
+	}
+	return nil
+}
+
 // PKCEAuthorizationFlow implements the OAuthFlow interface for
 // the Authorization Code Flow with PKCE.
 type PKCEAuthorizationFlow struct {
-	providerConfig internal.PKCEAuthProviderConfig
+	providerConfig PKCEAuthProviderConfig
 	state          string
 	codeVerifier   string
 	oAuthConfig    *oauth2.Config
 }
 
 // NewPKCEAuthorizationFlow returns new PKCE authorization code flow.
-func NewPKCEAuthorizationFlow(config internal.PKCEAuthProviderConfig) (*PKCEAuthorizationFlow, error) {
+func NewPKCEAuthorizationFlow(config PKCEAuthProviderConfig) (*PKCEAuthorizationFlow, error) {
 	var availableRedirectURL string
 
-	excludedRanges := getSystemExcludedPortRanges()
-
+	// find the first available redirect URL
 	for _, redirectURL := range config.RedirectURLs {
-		if !isRedirectURLPortUsed(redirectURL, excludedRanges) {
+		if !isRedirectURLPortUsed(redirectURL) {
 			availableRedirectURL = redirectURL
 			break
 		}
@@ -105,10 +152,10 @@ func (p *PKCEAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowIn
 		oauth2.SetAuthURLParam("audience", p.providerConfig.Audience),
 	}
 	if !p.providerConfig.DisablePromptLogin {
-		switch p.providerConfig.LoginFlag {
-		case common.LoginFlagPromptLogin:
+		if p.providerConfig.LoginFlag.IsPromptLogin() {
 			params = append(params, oauth2.SetAuthURLParam("prompt", "login"))
-		case common.LoginFlagMaxAge0:
+		}
+		if p.providerConfig.LoginFlag.IsMaxAge0Login() {
 			params = append(params, oauth2.SetAuthURLParam("max_age", "0"))
 		}
 	}
@@ -124,10 +171,21 @@ func (p *PKCEAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowIn
 	}, nil
 }
 
+// SetLoginHint sets the login hint for the PKCE authorization flow
+func (p *PKCEAuthorizationFlow) SetLoginHint(hint string) {
+	p.providerConfig.LoginHint = hint
+}
+
 // WaitToken waits for the OAuth token in the PKCE Authorization Flow.
 // It starts an HTTP server to receive the OAuth token callback and waits for the token or an error.
 // Once the token is received, it is converted to TokenInfo and validated before returning.
-func (p *PKCEAuthorizationFlow) WaitToken(ctx context.Context, _ AuthFlowInfo) (TokenInfo, error) {
+// The method creates a timeout context internally based on info.ExpiresIn.
+func (p *PKCEAuthorizationFlow) WaitToken(ctx context.Context, info AuthFlowInfo) (TokenInfo, error) {
+	// Create timeout context based on flow expiration
+	timeout := time.Duration(info.ExpiresIn) * time.Second
+	waitCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	tokenChan := make(chan *oauth2.Token, 1)
 	errChan := make(chan error, 1)
 
@@ -138,7 +196,7 @@ func (p *PKCEAuthorizationFlow) WaitToken(ctx context.Context, _ AuthFlowInfo) (
 
 	server := &http.Server{Addr: fmt.Sprintf(":%s", parsedURL.Port())}
 	defer func() {
-		shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
 
 		if err := server.Shutdown(shutdownCtx); err != nil {
@@ -149,8 +207,8 @@ func (p *PKCEAuthorizationFlow) WaitToken(ctx context.Context, _ AuthFlowInfo) (
 	go p.startServer(server, tokenChan, errChan)
 
 	select {
-	case <-ctx.Done():
-		return TokenInfo{}, ctx.Err()
+	case <-waitCtx.Done():
+		return TokenInfo{}, waitCtx.Err()
 	case token := <-tokenChan:
 		return p.parseOAuthToken(token)
 	case err := <-errChan:
@@ -285,22 +343,15 @@ func createCodeChallenge(codeVerifier string) string {
 	return base64.RawURLEncoding.EncodeToString(sha2[:])
 }
 
-// isRedirectURLPortUsed checks if the port used in the redirect URL is in use or excluded on Windows.
-func isRedirectURLPortUsed(redirectURL string, excludedRanges []excludedPortRange) bool {
+// isRedirectURLPortUsed checks if the port used in the redirect URL is in use.
+func isRedirectURLPortUsed(redirectURL string) bool {
 	parsedURL, err := url.Parse(redirectURL)
 	if err != nil {
 		log.Errorf("failed to parse redirect URL: %v", err)
 		return true
 	}
 
-	port := parsedURL.Port()
-
-	if isPortInExcludedRange(port, excludedRanges) {
-		log.Warnf("port %s is in Windows excluded port range, skipping", port)
-		return true
-	}
-
-	addr := fmt.Sprintf(":%s", port)
+	addr := fmt.Sprintf(":%s", parsedURL.Port())
 	conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
 	if err != nil {
 		return false
@@ -314,33 +365,6 @@ func isRedirectURLPortUsed(redirectURL string, excludedRanges []excludedPortRang
 	return true
 }
 
-// excludedPortRange represents a range of excluded ports.
-type excludedPortRange struct {
-	start int
-	end   int
-}
-
-// isPortInExcludedRange checks if the given port is in any of the excluded ranges.
-func isPortInExcludedRange(port string, excludedRanges []excludedPortRange) bool {
-	if len(excludedRanges) == 0 {
-		return false
-	}
-
-	portNum, err := strconv.Atoi(port)
-	if err != nil {
-		log.Debugf("invalid port number %s: %v", port, err)
-		return false
-	}
-
-	for _, r := range excludedRanges {
-		if portNum >= r.start && portNum <= r.end {
-			return true
-		}
-	}
-
-	return false
-}
-
 func renderPKCEFlowTmpl(w http.ResponseWriter, authError error) {
 	tmpl, err := template.New("pkce-auth-flow").Parse(templates.PKCEAuthMsgTmpl)
 	if err != nil {

client/internal/auth/pkce_flow_other.go

@@ -1,8 +0,0 @@
-//go:build !windows
-
-package auth
-
-// getSystemExcludedPortRanges returns nil on non-Windows platforms.
-func getSystemExcludedPortRanges() []excludedPortRange {
-	return nil
-}

client/internal/auth/pkce_flow_test.go

@@ -2,14 +2,10 @@ package auth
 
 import (
 	"context"
-	"fmt"
-	"net"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/netbirdio/netbird/client/internal"
 	mgm "github.com/netbirdio/netbird/shared/management/client/common"
 )
 
@@ -23,34 +19,28 @@ func TestPromptLogin(t *testing.T) {
 		name               string
 		loginFlag          mgm.LoginFlag
 		disablePromptLogin bool
-		expectContains     []string
+		expect             string
 	}{
 		{
-			name:           "Prompt login",
-			loginFlag:      mgm.LoginFlagPromptLogin,
-			expectContains: []string{promptLogin},
+			name:      "Prompt login",
+			loginFlag: mgm.LoginFlagPrompt,
+			expect:    promptLogin,
 		},
 		{
-			name:           "Max age 0",
-			loginFlag:      mgm.LoginFlagMaxAge0,
-			expectContains: []string{maxAge0},
+			name:      "Max age 0 login",
+			loginFlag: mgm.LoginFlagMaxAge0,
+			expect:    maxAge0,
 		},
 		{
 			name:               "Disable prompt login",
-			loginFlag:          mgm.LoginFlagPromptLogin,
+			loginFlag:          mgm.LoginFlagPrompt,
 			disablePromptLogin: true,
-			expectContains:     []string{},
-		},
-		{
-			name:           "None flag should not add parameters",
-			loginFlag:      mgm.LoginFlagNone,
-			expectContains: []string{},
 		},
 	}
 
 	for _, tc := range tt {
 		t.Run(tc.name, func(t *testing.T) {
-			config := internal.PKCEAuthProviderConfig{
+			config := PKCEAuthProviderConfig{
 				ClientID:              "test-client-id",
 				Audience:              "test-audience",
 				TokenEndpoint:         "https://test-token-endpoint.com/token",
@@ -59,7 +49,6 @@ func TestPromptLogin(t *testing.T) {
 				RedirectURLs:          []string{"http://127.0.0.1:33992/"},
 				UseIDToken:            true,
 				LoginFlag:             tc.loginFlag,
-				DisablePromptLogin:    tc.disablePromptLogin,
 			}
 			pkce, err := NewPKCEAuthorizationFlow(config)
 			if err != nil {
@@ -70,153 +59,12 @@ func TestPromptLogin(t *testing.T) {
 				t.Fatalf("Failed to request auth info: %v", err)
 			}
 
-			for _, expected := range tc.expectContains {
-				require.Contains(t, authInfo.VerificationURIComplete, expected)
+			if !tc.disablePromptLogin {
+				require.Contains(t, authInfo.VerificationURIComplete, tc.expect)
+			} else {
+				require.Contains(t, authInfo.VerificationURIComplete, promptLogin)
+				require.NotContains(t, authInfo.VerificationURIComplete, maxAge0)
 			}
 		})
 	}
 }
-
-func TestIsPortInExcludedRange(t *testing.T) {
-	tests := []struct {
-		name            string
-		port            string
-		excludedRanges  []excludedPortRange
-		expectedBlocked bool
-	}{
-		{
-			name:            "Port in excluded range",
-			port:            "8080",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: true,
-		},
-		{
-			name:            "Port at start of range",
-			port:            "8000",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: true,
-		},
-		{
-			name:            "Port at end of range",
-			port:            "8100",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: true,
-		},
-		{
-			name:            "Port before range",
-			port:            "7999",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: false,
-		},
-		{
-			name:            "Port after range",
-			port:            "8101",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: false,
-		},
-		{
-			name:            "Empty excluded ranges",
-			port:            "8080",
-			excludedRanges:  []excludedPortRange{},
-			expectedBlocked: false,
-		},
-		{
-			name:            "Nil excluded ranges",
-			port:            "8080",
-			excludedRanges:  nil,
-			expectedBlocked: false,
-		},
-		{
-			name: "Multiple ranges - port in second range",
-			port: "9050",
-			excludedRanges: []excludedPortRange{
-				{start: 8000, end: 8100},
-				{start: 9000, end: 9100},
-			},
-			expectedBlocked: true,
-		},
-		{
-			name: "Multiple ranges - port not in any range",
-			port: "8500",
-			excludedRanges: []excludedPortRange{
-				{start: 8000, end: 8100},
-				{start: 9000, end: 9100},
-			},
-			expectedBlocked: false,
-		},
-		{
-			name:            "Invalid port string",
-			port:            "invalid",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: false,
-		},
-		{
-			name:            "Empty port string",
-			port:            "",
-			excludedRanges:  []excludedPortRange{{start: 8000, end: 8100}},
-			expectedBlocked: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := isPortInExcludedRange(tt.port, tt.excludedRanges)
-			assert.Equal(t, tt.expectedBlocked, result, "Port exclusion check mismatch")
-		})
-	}
-}
-
-func TestIsRedirectURLPortUsed(t *testing.T) {
-	listener, err := net.Listen("tcp", "127.0.0.1:0")
-	require.NoError(t, err)
-	defer func() {
-		_ = listener.Close()
-	}()
-
-	usedPort := listener.Addr().(*net.TCPAddr).Port
-
-	tests := []struct {
-		name           string
-		redirectURL    string
-		excludedRanges []excludedPortRange
-		expectedUsed   bool
-	}{
-		{
-			name:           "Port in excluded range",
-			redirectURL:    "http://127.0.0.1:8080/",
-			excludedRanges: []excludedPortRange{{start: 8000, end: 8100}},
-			expectedUsed:   true,
-		},
-		{
-			name:           "Port actually in use",
-			redirectURL:    fmt.Sprintf("http://127.0.0.1:%d/", usedPort),
-			excludedRanges: nil,
-			expectedUsed:   true,
-		},
-		{
-			name:           "Port not in use and not excluded",
-			redirectURL:    "http://127.0.0.1:65432/",
-			excludedRanges: nil,
-			expectedUsed:   false,
-		},
-		{
-			name:           "Invalid URL without port",
-			redirectURL:    "not-a-valid-url",
-			excludedRanges: nil,
-			expectedUsed:   false,
-		},
-		{
-			name:           "Port excluded even if not in use",
-			redirectURL:    "http://127.0.0.1:8050/",
-			excludedRanges: []excludedPortRange{{start: 8000, end: 8100}},
-			expectedUsed:   true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := isRedirectURLPortUsed(tt.redirectURL, tt.excludedRanges)
-			assert.Equal(t, tt.expectedUsed, result, "Port usage check mismatch")
-		})
-	}
-}

client/internal/auth/pkce_flow_windows.go

@@ -1,86 +0,0 @@
-//go:build windows
-
-package auth
-
-import (
-	"bufio"
-	"fmt"
-	"os/exec"
-	"strconv"
-	"strings"
-
-	log "github.com/sirupsen/logrus"
-)
-
-// getSystemExcludedPortRanges retrieves the excluded port ranges from Windows using netsh.
-func getSystemExcludedPortRanges() []excludedPortRange {
-	ranges, err := getExcludedPortRangesFromNetsh()
-	if err != nil {
-		log.Debugf("failed to get Windows excluded port ranges: %v", err)
-		return nil
-	}
-
-	return ranges
-}
-
-// getExcludedPortRangesFromNetsh retrieves excluded port ranges using netsh command.
-func getExcludedPortRangesFromNetsh() ([]excludedPortRange, error) {
-	cmd := exec.Command("netsh", "interface", "ipv4", "show", "excludedportrange", "protocol=tcp")
-	output, err := cmd.Output()
-	if err != nil {
-		return nil, fmt.Errorf("netsh command: %w", err)
-	}
-
-	return parseExcludedPortRanges(string(output))
-}
-
-// parseExcludedPortRanges parses the output of the netsh command to extract port ranges.
-func parseExcludedPortRanges(output string) ([]excludedPortRange, error) {
-	var ranges []excludedPortRange
-	scanner := bufio.NewScanner(strings.NewReader(output))
-
-	foundHeader := false
-	for scanner.Scan() {
-		line := strings.TrimSpace(scanner.Text())
-
-		if strings.Contains(line, "Start Port") && strings.Contains(line, "End Port") {
-			foundHeader = true
-			continue
-		}
-
-		if !foundHeader {
-			continue
-		}
-
-		if strings.Contains(line, "----------") {
-			continue
-		}
-
-		if line == "" {
-			continue
-		}
-
-		fields := strings.Fields(line)
-		if len(fields) < 2 {
-			continue
-		}
-
-		startPort, err := strconv.Atoi(fields[0])
-		if err != nil {
-			continue
-		}
-
-		endPort, err := strconv.Atoi(fields[1])
-		if err != nil {
-			continue
-		}
-
-		ranges = append(ranges, excludedPortRange{start: startPort, end: endPort})
-	}
-
-	if err := scanner.Err(); err != nil {
-		return nil, fmt.Errorf("scan output: %w", err)
-	}
-
-	return ranges, nil
-}

client/internal/auth/pkce_flow_windows_test.go

@@ -1,116 +0,0 @@
-//go:build windows
-
-package auth
-
-import (
-	"fmt"
-	"net"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/internal"
-)
-
-func TestParseExcludedPortRanges(t *testing.T) {
-	tests := []struct {
-		name           string
-		netshOutput    string
-		expectedRanges []excludedPortRange
-		expectError    bool
-	}{
-		{
-			name: "Valid netsh output with multiple ranges",
-			netshOutput: `
-Protocol tcp Dynamic Port Range
----------------------------------
-Start Port      : 49152
-Number of Ports : 16384
-
-Protocol tcp Excluded Port Ranges
----------------------------------
-Start Port    End Port
-----------    --------
-     5357        5357      *
-    50000       50059      *
-`,
-			expectedRanges: []excludedPortRange{
-				{start: 5357, end: 5357},
-				{start: 50000, end: 50059},
-			},
-			expectError: false,
-		},
-		{
-			name: "Empty output",
-			netshOutput: `
-Protocol tcp Dynamic Port Range
----------------------------------
-Start Port      : 49152
-Number of Ports : 16384
-`,
-			expectedRanges: nil,
-			expectError:    false,
-		},
-		{
-			name: "Single range",
-			netshOutput: `
-Protocol tcp Excluded Port Ranges
----------------------------------
-Start Port    End Port
-----------    --------
-     8080        8090
-`,
-			expectedRanges: []excludedPortRange{
-				{start: 8080, end: 8090},
-			},
-			expectError: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			ranges, err := parseExcludedPortRanges(tt.netshOutput)
-
-			if tt.expectError {
-				assert.Error(t, err)
-			} else {
-				require.NoError(t, err)
-				assert.Equal(t, tt.expectedRanges, ranges)
-			}
-		})
-	}
-}
-
-func TestNewPKCEAuthorizationFlow_WithActualExcludedPorts(t *testing.T) {
-	ranges := getSystemExcludedPortRanges()
-	t.Logf("Found %d excluded port ranges on this system", len(ranges))
-
-	listener1, err := net.Listen("tcp", "127.0.0.1:0")
-	require.NoError(t, err)
-	defer func() {
-		_ = listener1.Close()
-	}()
-	usedPort1 := listener1.Addr().(*net.TCPAddr).Port
-
-	availablePort := 65432
-
-	config := internal.PKCEAuthProviderConfig{
-		ClientID:              "test-client-id",
-		Audience:              "test-audience",
-		TokenEndpoint:         "https://test-token-endpoint.com/token",
-		Scope:                 "openid email profile",
-		AuthorizationEndpoint: "https://test-auth-endpoint.com/authorize",
-		RedirectURLs: []string{
-			fmt.Sprintf("http://127.0.0.1:%d/", usedPort1),
-			fmt.Sprintf("http://127.0.0.1:%d/", availablePort),
-		},
-		UseIDToken: true,
-	}
-
-	flow, err := NewPKCEAuthorizationFlow(config)
-	require.NoError(t, err)
-	require.NotNil(t, flow)
-	assert.Contains(t, flow.oAuthConfig.RedirectURL, fmt.Sprintf(":%d", availablePort),
-		"Should skip port in use and select available port")
-}

client/internal/connect.go

@@ -162,6 +162,11 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		return err
 	}
 
+	// Create management client once outside retry loop
+	mgmClient := mgm.NewClient(c.config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
+	mgmNotifier := statusRecorderToMgmConnStateNotifier(c.statusRecorder)
+	mgmClient.SetConnStateListener(mgmNotifier)
+
 	defer c.statusRecorder.ClientStop()
 	operation := func() error {
 		// if context cancelled we not start new backoff cycle
@@ -180,12 +185,9 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		}()
 
 		log.Debugf("connecting to the Management service %s", c.config.ManagementURL.Host)
-		mgmClient, err := mgm.NewClient(engineCtx, c.config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
-		if err != nil {
-			return wrapErr(gstatus.Errorf(codes.FailedPrecondition, "failed connecting to Management Service : %s", err))
+		if err := mgmClient.ConnectWithoutRetry(engineCtx); err != nil {
+			return wrapErr(fmt.Errorf("failed connecting to Management Service: %w", err))
 		}
-		mgmNotifier := statusRecorderToMgmConnStateNotifier(c.statusRecorder)
-		mgmClient.SetConnStateListener(mgmNotifier)
 
 		log.Debugf("connected to the Management service %s", c.config.ManagementURL.Host)
 		defer func() {
@@ -198,7 +200,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		loginResp, err := loginToManagement(engineCtx, mgmClient, publicSSHKey, c.config)
 		if err != nil {
 			log.Debug(err)
-			if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.PermissionDenied) {
+			if errors.Is(err, mgm.ErrPermissionDenied) {
 				state.Set(StatusNeedsLogin)
 				_ = c.Stop()
 				return backoff.Permanent(wrapErr(err)) // unrecoverable error
@@ -320,7 +322,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 	err = backoff.Retry(operation, backOff)
 	if err != nil {
 		log.Debugf("exiting client retry loop due to unrecoverable error: %s", err)
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.PermissionDenied) {
+		if errors.Is(err, mgm.ErrPermissionDenied) {
 			state.Set(StatusNeedsLogin)
 			_ = c.Stop()
 		}
@@ -504,12 +506,6 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourP
 
 // loginToManagement creates Management ServiceDependencies client, establishes a connection, logs-in and gets a global Netbird config (signal, turn, stun hosts, etc)
 func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte, config *profilemanager.Config) (*mgmProto.LoginResponse, error) {
-
-	serverPublicKey, err := client.GetServerPublicKey()
-	if err != nil {
-		return nil, gstatus.Errorf(codes.FailedPrecondition, "failed while getting Management Service public key: %s", err)
-	}
-
 	sysInfo := system.GetInfo(ctx)
 	sysInfo.SetFlags(
 		config.RosenpassEnabled,
@@ -528,7 +524,7 @@ func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte,
 		config.EnableSSHRemotePortForwarding,
 		config.DisableSSHAuth,
 	)
-	loginResp, err := client.Login(*serverPublicKey, sysInfo, pubSSHKey, config.DNSLabels)
+	loginResp, err := client.Login(ctx, sysInfo, pubSSHKey, config.DNSLabels)
 	if err != nil {
 		return nil, err
 	}

client/internal/device_auth.go

@@ -1,136 +0,0 @@
-package internal
-
-import (
-	"context"
-	"fmt"
-	"net/url"
-
-	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
-
-	mgm "github.com/netbirdio/netbird/shared/management/client"
-)
-
-// DeviceAuthorizationFlow represents Device Authorization Flow information
-type DeviceAuthorizationFlow struct {
-	Provider       string
-	ProviderConfig DeviceAuthProviderConfig
-}
-
-// DeviceAuthProviderConfig has all attributes needed to initiate a device authorization flow
-type DeviceAuthProviderConfig struct {
-	// ClientID An IDP application client id
-	ClientID string
-	// ClientSecret An IDP application client secret
-	ClientSecret string
-	// Domain An IDP API domain
-	// Deprecated. Use OIDCConfigEndpoint instead
-	Domain string
-	// Audience An Audience for to authorization validation
-	Audience string
-	// TokenEndpoint is the endpoint of an IDP manager where clients can obtain access token
-	TokenEndpoint string
-	// DeviceAuthEndpoint is the endpoint of an IDP manager where clients can obtain device authorization code
-	DeviceAuthEndpoint string
-	// Scopes provides the scopes to be included in the token request
-	Scope string
-	// UseIDToken indicates if the id token should be used for authentication
-	UseIDToken bool
-	// LoginHint is used to pre-fill the email/username field during authentication
-	LoginHint string
-}
-
-// GetDeviceAuthorizationFlowInfo initialize a DeviceAuthorizationFlow instance and return with it
-func GetDeviceAuthorizationFlowInfo(ctx context.Context, privateKey string, mgmURL *url.URL) (DeviceAuthorizationFlow, error) {
-	// validate our peer's Wireguard PRIVATE key
-	myPrivateKey, err := wgtypes.ParseKey(privateKey)
-	if err != nil {
-		log.Errorf("failed parsing Wireguard key %s: [%s]", privateKey, err.Error())
-		return DeviceAuthorizationFlow{}, err
-	}
-
-	var mgmTLSEnabled bool
-	if mgmURL.Scheme == "https" {
-		mgmTLSEnabled = true
-	}
-
-	log.Debugf("connecting to Management Service %s", mgmURL.String())
-	mgmClient, err := mgm.NewClient(ctx, mgmURL.Host, myPrivateKey, mgmTLSEnabled)
-	if err != nil {
-		log.Errorf("failed connecting to Management Service %s %v", mgmURL.String(), err)
-		return DeviceAuthorizationFlow{}, err
-	}
-	log.Debugf("connected to the Management service %s", mgmURL.String())
-
-	defer func() {
-		err = mgmClient.Close()
-		if err != nil {
-			log.Warnf("failed to close the Management service client %v", err)
-		}
-	}()
-
-	serverKey, err := mgmClient.GetServerPublicKey()
-	if err != nil {
-		log.Errorf("failed while getting Management Service public key: %v", err)
-		return DeviceAuthorizationFlow{}, err
-	}
-
-	protoDeviceAuthorizationFlow, err := mgmClient.GetDeviceAuthorizationFlow(*serverKey)
-	if err != nil {
-		if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
-			log.Warnf("server couldn't find device flow, contact admin: %v", err)
-			return DeviceAuthorizationFlow{}, err
-		}
-		log.Errorf("failed to retrieve device flow: %v", err)
-		return DeviceAuthorizationFlow{}, err
-	}
-
-	deviceAuthorizationFlow := DeviceAuthorizationFlow{
-		Provider: protoDeviceAuthorizationFlow.Provider.String(),
-
-		ProviderConfig: DeviceAuthProviderConfig{
-			Audience:           protoDeviceAuthorizationFlow.GetProviderConfig().GetAudience(),
-			ClientID:           protoDeviceAuthorizationFlow.GetProviderConfig().GetClientID(),
-			ClientSecret:       protoDeviceAuthorizationFlow.GetProviderConfig().GetClientSecret(),
-			Domain:             protoDeviceAuthorizationFlow.GetProviderConfig().Domain,
-			TokenEndpoint:      protoDeviceAuthorizationFlow.GetProviderConfig().GetTokenEndpoint(),
-			DeviceAuthEndpoint: protoDeviceAuthorizationFlow.GetProviderConfig().GetDeviceAuthEndpoint(),
-			Scope:              protoDeviceAuthorizationFlow.GetProviderConfig().GetScope(),
-			UseIDToken:         protoDeviceAuthorizationFlow.GetProviderConfig().GetUseIDToken(),
-		},
-	}
-
-	// keep compatibility with older management versions
-	if deviceAuthorizationFlow.ProviderConfig.Scope == "" {
-		deviceAuthorizationFlow.ProviderConfig.Scope = "openid"
-	}
-
-	err = isDeviceAuthProviderConfigValid(deviceAuthorizationFlow.ProviderConfig)
-	if err != nil {
-		return DeviceAuthorizationFlow{}, err
-	}
-
-	return deviceAuthorizationFlow, nil
-}
-
-func isDeviceAuthProviderConfigValid(config DeviceAuthProviderConfig) error {
-	errorMSGFormat := "invalid provider configuration received from management: %s value is empty. Contact your NetBird administrator"
-	if config.Audience == "" {
-		return fmt.Errorf(errorMSGFormat, "Audience")
-	}
-	if config.ClientID == "" {
-		return fmt.Errorf(errorMSGFormat, "Client ID")
-	}
-	if config.TokenEndpoint == "" {
-		return fmt.Errorf(errorMSGFormat, "Token Endpoint")
-	}
-	if config.DeviceAuthEndpoint == "" {
-		return fmt.Errorf(errorMSGFormat, "Device Auth Endpoint")
-	}
-	if config.Scope == "" {
-		return fmt.Errorf(errorMSGFormat, "Device Auth Scopes")
-	}
-	return nil
-}

client/internal/engine.go

@@ -292,12 +292,6 @@ func (e *Engine) Stop() error {
 	}
 	log.Info("Network monitor: stopped")
 
-	if os.Getenv("NB_REMOVE_BEFORE_DNS") == "true" && os.Getenv("NB_REMOVE_BEFORE_ROUTES") != "true" {
-		log.Info("removing peers before dns")
-		if err := e.removeAllPeers(); err != nil {
-			return fmt.Errorf("failed to remove all peers: %s", err)
-		}
-	}
 	if err := e.stopSSHServer(); err != nil {
 		log.Warnf("failed to stop SSH server: %v", err)
 	}
@@ -316,13 +310,6 @@ func (e *Engine) Stop() error {
 
 	e.stopDNSForwarder()
 
-	if os.Getenv("NB_REMOVE_BEFORE_ROUTES") == "true" && os.Getenv("NB_REMOVE_BEFORE_DNS") != "true" {
-		log.Info("removing peers before routes")
-		if err := e.removeAllPeers(); err != nil {
-			return fmt.Errorf("failed to remove all peers: %s", err)
-		}
-	}
-
 	if e.routeManager != nil {
 		e.routeManager.Stop(e.stateManager)
 	}
@@ -330,16 +317,13 @@ func (e *Engine) Stop() error {
 	if e.srWatcher != nil {
 		e.srWatcher.Close()
 	}
-	log.Info("cleaning up status recorder states")
+
 	e.statusRecorder.ReplaceOfflinePeers([]peer.State{})
 	e.statusRecorder.UpdateDNSStates([]peer.NSGroupState{})
 	e.statusRecorder.UpdateRelayStates([]relay.ProbeResult{})
 
-	if os.Getenv("NB_REMOVE_BEFORE_DNS") != "true" && os.Getenv("NB_REMOVE_BEFORE_ROUTES") != "true" {
-		log.Info("removing peers after dns and routes")
-		if err := e.removeAllPeers(); err != nil {
-			return fmt.Errorf("failed to remove all peers: %s", err)
-		}
+	if err := e.removeAllPeers(); err != nil {
+		return fmt.Errorf("failed to remove all peers: %s", err)
 	}
 
 	if e.cancel != nil {
@@ -907,7 +891,7 @@ func (e *Engine) updateChecksIfNew(checks []*mgmProto.Checks) error {
 		e.config.DisableSSHAuth,
 	)
 
-	if err := e.mgmClient.SyncMeta(info); err != nil {
+	if err := e.mgmClient.SyncMeta(e.ctx, info); err != nil {
 		log.Errorf("could not sync meta: error %s", err)
 		return err
 	}
@@ -1533,7 +1517,7 @@ func (e *Engine) readInitialSettings() ([]*route.Route, *nbdns.Config, bool, err
 		e.config.DisableSSHAuth,
 	)
 
-	netMap, err := e.mgmClient.GetNetworkMap(info)
+	netMap, err := e.mgmClient.GetNetworkMap(e.ctx, info)
 	if err != nil {
 		return nil, nil, false, err
 	}
@@ -1682,7 +1666,7 @@ func (e *Engine) RunHealthProbes(waitForResult bool) bool {
 	signalHealthy := e.signal.IsHealthy()
 	log.Debugf("signal health check: healthy=%t", signalHealthy)
 
-	managementHealthy := e.mgmClient.IsHealthy()
+	managementHealthy := e.mgmClient.IsHealthy(e.ctx)
 	log.Debugf("management health check: healthy=%t", managementHealthy)
 
 	stuns := slices.Clone(e.STUNs)

client/internal/engine_test.go

@@ -30,12 +30,11 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/groups"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/configurer"
@@ -55,6 +54,7 @@ import (
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/peers"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -1503,8 +1503,8 @@ func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey strin
 	if err != nil {
 		return nil, err
 	}
-	mgmtClient, err := mgmt.NewClient(ctx, mgmtAddr, key, false)
-	if err != nil {
+	mgmtClient := mgmt.NewClient(mgmtAddr, key, false)
+	if err := mgmtClient.Connect(ctx); err != nil {
 		return nil, err
 	}
 	signalClient, err := signal.NewClient(ctx, signalAddr, key, false)
@@ -1512,13 +1512,8 @@ func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey strin
 		return nil, err
 	}
 
-	publicKey, err := mgmtClient.GetServerPublicKey()
-	if err != nil {
-		return nil, err
-	}
-
 	info := system.GetInfo(ctx)
-	resp, err := mgmtClient.Register(*publicKey, setupKey, "", info, nil, nil)
+	err = mgmtClient.Register(ctx, setupKey, "", info, nil, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -1531,9 +1526,10 @@ func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey strin
 	}
 
 	wgPort := 33100 + i
+	testAddr := fmt.Sprintf("100.64.0.%d/24", i+1)
 	conf := &EngineConfig{
 		WgIfaceName:  ifaceName,
-		WgAddr:       resp.PeerConfig.Address,
+		WgAddr:       testAddr,
 		WgPrivateKey: key,
 		WgPort:       wgPort,
 		MTU:          iface.DefaultMTU,
@@ -1628,17 +1624,14 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := server.NewAccountRequestBuffer(context.Background(), store)
-	networkMapController := controller.NewController(context.Background(), store, metrics, updateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), manager.NewEphemeralManager(store, peersManager), config)
+	networkMapController := controller.NewController(context.Background(), store, metrics, updateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), config)
 	accountManager, err := server.BuildManager(context.Background(), config, store, networkMapController, nil, "", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	if err != nil {
 		return nil, "", err
 	}
 
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		return nil, "", err
-	}
-	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, secretsManager, nil, nil, &server.MockIntegratedValidator{}, networkMapController)
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
+	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, updateManager, secretsManager, nil, &manager.EphemeralManager{}, nil, &server.MockIntegratedValidator{}, networkMapController)
 	if err != nil {
 		return nil, "", err
 	}

client/internal/login.go

@@ -1,201 +0,0 @@
-package internal
-
-import (
-	"context"
-	"net/url"
-
-	"github.com/google/uuid"
-	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/internal/profilemanager"
-	"github.com/netbirdio/netbird/client/ssh"
-	"github.com/netbirdio/netbird/client/system"
-	mgm "github.com/netbirdio/netbird/shared/management/client"
-	mgmProto "github.com/netbirdio/netbird/shared/management/proto"
-)
-
-// IsLoginRequired check that the server is support SSO or not
-func IsLoginRequired(ctx context.Context, config *profilemanager.Config) (bool, error) {
-	mgmURL := config.ManagementURL
-	mgmClient, err := getMgmClient(ctx, config.PrivateKey, mgmURL)
-	if err != nil {
-		return false, err
-	}
-	defer func() {
-		err = mgmClient.Close()
-		if err != nil {
-			cStatus, ok := status.FromError(err)
-			if !ok || ok && cStatus.Code() != codes.Canceled {
-				log.Warnf("failed to close the Management service client, err: %v", err)
-			}
-		}
-	}()
-	log.Debugf("connected to the Management service %s", mgmURL.String())
-
-	pubSSHKey, err := ssh.GeneratePublicKey([]byte(config.SSHKey))
-	if err != nil {
-		return false, err
-	}
-
-	_, _, err = doMgmLogin(ctx, mgmClient, pubSSHKey, config)
-	if isLoginNeeded(err) {
-		return true, nil
-	}
-	return false, err
-}
-
-// Login or register the client
-func Login(ctx context.Context, config *profilemanager.Config, setupKey string, jwtToken string) error {
-	mgmClient, err := getMgmClient(ctx, config.PrivateKey, config.ManagementURL)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		err = mgmClient.Close()
-		if err != nil {
-			cStatus, ok := status.FromError(err)
-			if !ok || ok && cStatus.Code() != codes.Canceled {
-				log.Warnf("failed to close the Management service client, err: %v", err)
-			}
-		}
-	}()
-	log.Debugf("connected to the Management service %s", config.ManagementURL.String())
-
-	pubSSHKey, err := ssh.GeneratePublicKey([]byte(config.SSHKey))
-	if err != nil {
-		return err
-	}
-
-	serverKey, _, err := doMgmLogin(ctx, mgmClient, pubSSHKey, config)
-	if serverKey != nil && isRegistrationNeeded(err) {
-		log.Debugf("peer registration required")
-		_, err = registerPeer(ctx, *serverKey, mgmClient, setupKey, jwtToken, pubSSHKey, config)
-		if err != nil {
-			return err
-		}
-	} else if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-func getMgmClient(ctx context.Context, privateKey string, mgmURL *url.URL) (*mgm.GrpcClient, error) {
-	// validate our peer's Wireguard PRIVATE key
-	myPrivateKey, err := wgtypes.ParseKey(privateKey)
-	if err != nil {
-		log.Errorf("failed parsing Wireguard key %s: [%s]", privateKey, err.Error())
-		return nil, err
-	}
-
-	var mgmTlsEnabled bool
-	if mgmURL.Scheme == "https" {
-		mgmTlsEnabled = true
-	}
-
-	log.Debugf("connecting to the Management service %s", mgmURL.String())
-	mgmClient, err := mgm.NewClient(ctx, mgmURL.Host, myPrivateKey, mgmTlsEnabled)
-	if err != nil {
-		log.Errorf("failed connecting to the Management service %s %v", mgmURL.String(), err)
-		return nil, err
-	}
-	return mgmClient, err
-}
-
-func doMgmLogin(ctx context.Context, mgmClient *mgm.GrpcClient, pubSSHKey []byte, config *profilemanager.Config) (*wgtypes.Key, *mgmProto.LoginResponse, error) {
-	serverKey, err := mgmClient.GetServerPublicKey()
-	if err != nil {
-		log.Errorf("failed while getting Management Service public key: %v", err)
-		return nil, nil, err
-	}
-
-	sysInfo := system.GetInfo(ctx)
-	sysInfo.SetFlags(
-		config.RosenpassEnabled,
-		config.RosenpassPermissive,
-		config.ServerSSHAllowed,
-		config.DisableClientRoutes,
-		config.DisableServerRoutes,
-		config.DisableDNS,
-		config.DisableFirewall,
-		config.BlockLANAccess,
-		config.BlockInbound,
-		config.LazyConnectionEnabled,
-		config.EnableSSHRoot,
-		config.EnableSSHSFTP,
-		config.EnableSSHLocalPortForwarding,
-		config.EnableSSHRemotePortForwarding,
-		config.DisableSSHAuth,
-	)
-	loginResp, err := mgmClient.Login(*serverKey, sysInfo, pubSSHKey, config.DNSLabels)
-	return serverKey, loginResp, err
-}
-
-// registerPeer checks whether setupKey was provided via cmd line and if not then it prompts user to enter a key.
-// Otherwise tries to register with the provided setupKey via command line.
-func registerPeer(ctx context.Context, serverPublicKey wgtypes.Key, client *mgm.GrpcClient, setupKey string, jwtToken string, pubSSHKey []byte, config *profilemanager.Config) (*mgmProto.LoginResponse, error) {
-	validSetupKey, err := uuid.Parse(setupKey)
-	if err != nil && jwtToken == "" {
-		return nil, status.Errorf(codes.InvalidArgument, "invalid setup-key or no sso information provided, err: %v", err)
-	}
-
-	log.Debugf("sending peer registration request to Management Service")
-	info := system.GetInfo(ctx)
-	info.SetFlags(
-		config.RosenpassEnabled,
-		config.RosenpassPermissive,
-		config.ServerSSHAllowed,
-		config.DisableClientRoutes,
-		config.DisableServerRoutes,
-		config.DisableDNS,
-		config.DisableFirewall,
-		config.BlockLANAccess,
-		config.BlockInbound,
-		config.LazyConnectionEnabled,
-		config.EnableSSHRoot,
-		config.EnableSSHSFTP,
-		config.EnableSSHLocalPortForwarding,
-		config.EnableSSHRemotePortForwarding,
-		config.DisableSSHAuth,
-	)
-	loginResp, err := client.Register(serverPublicKey, validSetupKey.String(), jwtToken, info, pubSSHKey, config.DNSLabels)
-	if err != nil {
-		log.Errorf("failed registering peer %v", err)
-		return nil, err
-	}
-
-	log.Infof("peer has been successfully registered on Management Service")
-
-	return loginResp, nil
-}
-
-func isLoginNeeded(err error) bool {
-	if err == nil {
-		return false
-	}
-	s, ok := status.FromError(err)
-	if !ok {
-		return false
-	}
-	if s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied {
-		return true
-	}
-	return false
-}
-
-func isRegistrationNeeded(err error) bool {
-	if err == nil {
-		return false
-	}
-	s, ok := status.FromError(err)
-	if !ok {
-		return false
-	}
-	if s.Code() == codes.PermissionDenied {
-		return true
-	}
-	return false
-}

client/internal/pkce_auth.go

@@ -1,138 +0,0 @@
-package internal
-
-import (
-	"context"
-	"crypto/tls"
-	"fmt"
-	"net/url"
-
-	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
-
-	mgm "github.com/netbirdio/netbird/shared/management/client"
-	"github.com/netbirdio/netbird/shared/management/client/common"
-)
-
-// PKCEAuthorizationFlow represents PKCE Authorization Flow information
-type PKCEAuthorizationFlow struct {
-	ProviderConfig PKCEAuthProviderConfig
-}
-
-// PKCEAuthProviderConfig has all attributes needed to initiate pkce authorization flow
-type PKCEAuthProviderConfig struct {
-	// ClientID An IDP application client id
-	ClientID string
-	// ClientSecret An IDP application client secret
-	ClientSecret string
-	// Audience An Audience for to authorization validation
-	Audience string
-	// TokenEndpoint is the endpoint of an IDP manager where clients can obtain access token
-	TokenEndpoint string
-	// AuthorizationEndpoint is the endpoint of an IDP manager where clients can obtain authorization code
-	AuthorizationEndpoint string
-	// Scopes provides the scopes to be included in the token request
-	Scope string
-	// RedirectURL handles authorization code from IDP manager
-	RedirectURLs []string
-	// UseIDToken indicates if the id token should be used for authentication
-	UseIDToken bool
-	// ClientCertPair is used for mTLS authentication to the IDP
-	ClientCertPair *tls.Certificate
-	// DisablePromptLogin makes the PKCE flow to not prompt the user for login
-	DisablePromptLogin bool
-	// LoginFlag is used to configure the PKCE flow login behavior
-	LoginFlag common.LoginFlag
-	// LoginHint is used to pre-fill the email/username field during authentication
-	LoginHint string
-}
-
-// GetPKCEAuthorizationFlowInfo initialize a PKCEAuthorizationFlow instance and return with it
-func GetPKCEAuthorizationFlowInfo(ctx context.Context, privateKey string, mgmURL *url.URL, clientCert *tls.Certificate) (PKCEAuthorizationFlow, error) {
-	// validate our peer's Wireguard PRIVATE key
-	myPrivateKey, err := wgtypes.ParseKey(privateKey)
-	if err != nil {
-		log.Errorf("failed parsing Wireguard key %s: [%s]", privateKey, err.Error())
-		return PKCEAuthorizationFlow{}, err
-	}
-
-	var mgmTLSEnabled bool
-	if mgmURL.Scheme == "https" {
-		mgmTLSEnabled = true
-	}
-
-	log.Debugf("connecting to Management Service %s", mgmURL.String())
-	mgmClient, err := mgm.NewClient(ctx, mgmURL.Host, myPrivateKey, mgmTLSEnabled)
-	if err != nil {
-		log.Errorf("failed connecting to Management Service %s %v", mgmURL.String(), err)
-		return PKCEAuthorizationFlow{}, err
-	}
-	log.Debugf("connected to the Management service %s", mgmURL.String())
-
-	defer func() {
-		err = mgmClient.Close()
-		if err != nil {
-			log.Warnf("failed to close the Management service client %v", err)
-		}
-	}()
-
-	serverKey, err := mgmClient.GetServerPublicKey()
-	if err != nil {
-		log.Errorf("failed while getting Management Service public key: %v", err)
-		return PKCEAuthorizationFlow{}, err
-	}
-
-	protoPKCEAuthorizationFlow, err := mgmClient.GetPKCEAuthorizationFlow(*serverKey)
-	if err != nil {
-		if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
-			log.Warnf("server couldn't find pkce flow, contact admin: %v", err)
-			return PKCEAuthorizationFlow{}, err
-		}
-		log.Errorf("failed to retrieve pkce flow: %v", err)
-		return PKCEAuthorizationFlow{}, err
-	}
-
-	authFlow := PKCEAuthorizationFlow{
-		ProviderConfig: PKCEAuthProviderConfig{
-			Audience:              protoPKCEAuthorizationFlow.GetProviderConfig().GetAudience(),
-			ClientID:              protoPKCEAuthorizationFlow.GetProviderConfig().GetClientID(),
-			ClientSecret:          protoPKCEAuthorizationFlow.GetProviderConfig().GetClientSecret(),
-			TokenEndpoint:         protoPKCEAuthorizationFlow.GetProviderConfig().GetTokenEndpoint(),
-			AuthorizationEndpoint: protoPKCEAuthorizationFlow.GetProviderConfig().GetAuthorizationEndpoint(),
-			Scope:                 protoPKCEAuthorizationFlow.GetProviderConfig().GetScope(),
-			RedirectURLs:          protoPKCEAuthorizationFlow.GetProviderConfig().GetRedirectURLs(),
-			UseIDToken:            protoPKCEAuthorizationFlow.GetProviderConfig().GetUseIDToken(),
-			ClientCertPair:        clientCert,
-			DisablePromptLogin:    protoPKCEAuthorizationFlow.GetProviderConfig().GetDisablePromptLogin(),
-			LoginFlag:             common.LoginFlag(protoPKCEAuthorizationFlow.GetProviderConfig().GetLoginFlag()),
-		},
-	}
-
-	err = isPKCEProviderConfigValid(authFlow.ProviderConfig)
-	if err != nil {
-		return PKCEAuthorizationFlow{}, err
-	}
-
-	return authFlow, nil
-}
-
-func isPKCEProviderConfigValid(config PKCEAuthProviderConfig) error {
-	errorMSGFormat := "invalid provider configuration received from management: %s value is empty. Contact your NetBird administrator"
-	if config.ClientID == "" {
-		return fmt.Errorf(errorMSGFormat, "Client ID")
-	}
-	if config.TokenEndpoint == "" {
-		return fmt.Errorf(errorMSGFormat, "Token Endpoint")
-	}
-	if config.AuthorizationEndpoint == "" {
-		return fmt.Errorf(errorMSGFormat, "Authorization Auth Endpoint")
-	}
-	if config.Scope == "" {
-		return fmt.Errorf(errorMSGFormat, "PKCE Auth Scopes")
-	}
-	if config.RedirectURLs == nil {
-		return fmt.Errorf(errorMSGFormat, "PKCE Redirect URLs")
-	}
-	return nil
-}

client/internal/profilemanager/config.go

@@ -730,8 +730,8 @@ func UpdateOldManagementURL(ctx context.Context, config *Config, configPath stri
 		return config, err
 	}
 
-	client, err := mgm.NewClient(ctx, newURL.Host, key, mgmTlsEnabled)
-	if err != nil {
+	client := mgm.NewClient(newURL.Host, key, mgmTlsEnabled)
+	if err := client.Connect(ctx); err != nil {
 		log.Infof("couldn't switch to the new Management %s", newURL.String())
 		return config, err
 	}
@@ -743,8 +743,7 @@ func UpdateOldManagementURL(ctx context.Context, config *Config, configPath stri
 	}()
 
 	// gRPC check
-	_, err = client.GetServerPublicKey()
-	if err != nil {
+	if err := client.HealthCheck(ctx); err != nil {
 		log.Infof("couldn't switch to the new Management %s", newURL.String())
 		return nil, err
 	}

client/internal/sleep/detector_darwin.go

@@ -1,218 +0,0 @@
-//go:build darwin && !ios
-
-package sleep
-
-/*
-#cgo LDFLAGS: -framework IOKit -framework CoreFoundation
-#include <IOKit/pwr_mgt/IOPMLib.h>
-#include <IOKit/IOMessage.h>
-#include <CoreFoundation/CoreFoundation.h>
-
-extern void sleepCallbackBridge();
-extern void poweredOnCallbackBridge();
-extern void suspendedCallbackBridge();
-extern void resumedCallbackBridge();
-
-
-// C global variables for IOKit state
-static IONotificationPortRef g_notifyPortRef = NULL;
-static io_object_t g_notifierObject = 0;
-static io_object_t g_generalInterestNotifier = 0;
-static io_connect_t g_rootPort = 0;
-static CFRunLoopRef g_runLoop = NULL;
-
-static void sleepCallback(void* refCon, io_service_t service, natural_t messageType, void* messageArgument) {
-	switch (messageType) {
-		case kIOMessageSystemWillSleep:
-			sleepCallbackBridge();
-			IOAllowPowerChange(g_rootPort, (long)messageArgument);
-			break;
-        case kIOMessageSystemHasPoweredOn:
-          	poweredOnCallbackBridge();
-          	break;
-        case kIOMessageServiceIsSuspended:
-			suspendedCallbackBridge();
-			break;
-		case kIOMessageServiceIsResumed:
-			resumedCallbackBridge();
-			break;
-		default:
-			break;
-	}
-}
-
-static void registerNotifications() {
-	g_rootPort = IORegisterForSystemPower(
-		NULL,
-		&g_notifyPortRef,
-		(IOServiceInterestCallback)sleepCallback,
-		&g_notifierObject
-	);
-
-	if (g_rootPort == 0) {
-		return;
-	}
-
-	CFRunLoopAddSource(CFRunLoopGetCurrent(),
-		IONotificationPortGetRunLoopSource(g_notifyPortRef),
-		kCFRunLoopCommonModes);
-
-	g_runLoop = CFRunLoopGetCurrent();
-	CFRunLoopRun();
-}
-
-static void unregisterNotifications() {
-	CFRunLoopRemoveSource(g_runLoop,
-		IONotificationPortGetRunLoopSource(g_notifyPortRef),
-		kCFRunLoopCommonModes);
-
-	IODeregisterForSystemPower(&g_notifierObject);
-	IOServiceClose(g_rootPort);
-	IONotificationPortDestroy(g_notifyPortRef);
-	CFRunLoopStop(g_runLoop);
-
-	g_notifyPortRef = NULL;
-	g_notifierObject = 0;
-	g_rootPort = 0;
-	g_runLoop = NULL;
-}
-
-*/
-import "C"
-
-import (
-	"context"
-	"fmt"
-	"runtime"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-)
-
-var (
-	serviceRegistry   = make(map[*Detector]struct{})
-	serviceRegistryMu sync.Mutex
-)
-
-//export sleepCallbackBridge
-func sleepCallbackBridge() {
-	log.Info("sleepCallbackBridge event triggered")
-
-	serviceRegistryMu.Lock()
-	defer serviceRegistryMu.Unlock()
-
-	for svc := range serviceRegistry {
-		svc.triggerCallback(EventTypeSleep)
-	}
-}
-
-//export resumedCallbackBridge
-func resumedCallbackBridge() {
-	log.Info("resumedCallbackBridge event triggered")
-}
-
-//export suspendedCallbackBridge
-func suspendedCallbackBridge() {
-	log.Info("suspendedCallbackBridge event triggered")
-}
-
-//export poweredOnCallbackBridge
-func poweredOnCallbackBridge() {
-	log.Info("poweredOnCallbackBridge event triggered")
-	serviceRegistryMu.Lock()
-	defer serviceRegistryMu.Unlock()
-
-	for svc := range serviceRegistry {
-		svc.triggerCallback(EventTypeWakeUp)
-	}
-}
-
-type Detector struct {
-	callback func(event EventType)
-	ctx      context.Context
-	cancel   context.CancelFunc
-}
-
-func NewDetector() (*Detector, error) {
-	return &Detector{}, nil
-}
-
-func (d *Detector) Register(callback func(event EventType)) error {
-	serviceRegistryMu.Lock()
-	defer serviceRegistryMu.Unlock()
-
-	if _, exists := serviceRegistry[d]; exists {
-		return fmt.Errorf("detector service already registered")
-	}
-
-	d.callback = callback
-
-	d.ctx, d.cancel = context.WithCancel(context.Background())
-
-	if len(serviceRegistry) > 0 {
-		serviceRegistry[d] = struct{}{}
-		return nil
-	}
-
-	serviceRegistry[d] = struct{}{}
-
-	// CFRunLoop must run on a single fixed OS thread
-	go func() {
-		runtime.LockOSThread()
-		defer runtime.UnlockOSThread()
-
-		C.registerNotifications()
-	}()
-
-	log.Info("sleep detection service started on macOS")
-	return nil
-}
-
-// Deregister removes the detector. When the last detector is removed, IOKit registration is torn down
-// and the runloop is stopped and cleaned up.
-func (d *Detector) Deregister() error {
-	serviceRegistryMu.Lock()
-	defer serviceRegistryMu.Unlock()
-	_, exists := serviceRegistry[d]
-	if !exists {
-		return nil
-	}
-
-	// cancel and remove this detector
-	d.cancel()
-	delete(serviceRegistry, d)
-
-	// If other Detectors still exist, leave IOKit running
-	if len(serviceRegistry) > 0 {
-		return nil
-	}
-
-	log.Info("sleep detection service stopping (deregister)")
-
-	// Deregister IOKit notifications, stop runloop, and free resources
-	C.unregisterNotifications()
-
-	return nil
-}
-
-func (d *Detector) triggerCallback(event EventType) {
-	doneChan := make(chan struct{})
-
-	timeout := time.NewTimer(500 * time.Millisecond)
-	defer timeout.Stop()
-
-	cb := d.callback
-	go func(callback func(event EventType)) {
-		log.Info("sleep detection event fired")
-		callback(event)
-		close(doneChan)
-	}(cb)
-
-	select {
-	case <-doneChan:
-	case <-d.ctx.Done():
-	case <-timeout.C:
-		log.Warnf("sleep callback timed out")
-	}
-}

client/internal/sleep/detector_notsupported.go

@@ -1,9 +0,0 @@
-//go:build !darwin || ios
-
-package sleep
-
-import "fmt"
-
-func NewDetector() (detector, error) {
-	return nil, fmt.Errorf("sleep not supported on this platform")
-}

client/internal/sleep/service.go

@@ -1,37 +0,0 @@
-package sleep
-
-var (
-	EventTypeUnknown EventType = 0
-	EventTypeSleep   EventType = 1
-	EventTypeWakeUp  EventType = 2
-)
-
-type EventType int
-
-type detector interface {
-	Register(callback func(eventType EventType)) error
-	Deregister() error
-}
-
-type Service struct {
-	detector detector
-}
-
-func New() (*Service, error) {
-	d, err := NewDetector()
-	if err != nil {
-		return nil, err
-	}
-
-	return &Service{
-		detector: d,
-	}, nil
-}
-
-func (s *Service) Register(callback func(eventType EventType)) error {
-	return s.detector.Register(callback)
-}
-
-func (s *Service) Deregister() error {
-	return s.detector.Deregister()
-}

client/ios/NetBirdSDK/client.go

@@ -208,7 +208,13 @@ func (c *Client) IsLoginRequired() bool {
 		ConfigPath: c.cfgFile,
 	})
 
-	needsLogin, _ := internal.IsLoginRequired(ctx, cfg)
+	authClient, err := auth.NewAuth(ctx, cfg.PrivateKey, cfg.ManagementURL, cfg)
+	if err != nil {
+		return true // Assume login is required if we can't create auth client
+	}
+	defer authClient.Close()
+
+	needsLogin, _ := authClient.IsLoginRequired(ctx)
 	return needsLogin
 }
 
@@ -240,15 +246,17 @@ func (c *Client) LoginForMobile() string {
 
 	// This could cause a potential race condition with loading the extension which need to be handled on swift side
 	go func() {
-		waitTimeout := time.Duration(flowInfo.ExpiresIn) * time.Second
-		waitCTX, cancel := context.WithTimeout(ctx, waitTimeout)
-		defer cancel()
-		tokenInfo, err := oAuthFlow.WaitToken(waitCTX, flowInfo)
+		tokenInfo, err := oAuthFlow.WaitToken(ctx, flowInfo)
 		if err != nil {
 			return
 		}
 		jwtToken := tokenInfo.GetTokenToUse()
-		_ = internal.Login(ctx, cfg, "", jwtToken)
+		authClient, err := auth.NewAuth(ctx, cfg.PrivateKey, cfg.ManagementURL, cfg)
+		if err != nil {
+			return
+		}
+		defer authClient.Close()
+		_ = authClient.Login(ctx, "", jwtToken)
 		c.loginComplete = true
 	}()
 

client/ios/NetBirdSDK/login.go

@@ -3,15 +3,8 @@ package NetBirdSDK
 import (
 	"context"
 	"fmt"
-	"time"
 
-	"github.com/cenkalti/backoff/v4"
-	log "github.com/sirupsen/logrus"
-	"google.golang.org/grpc/codes"
-	gstatus "google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/cmd"
-	"github.com/netbirdio/netbird/client/internal"
+	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/system"
 )
@@ -71,30 +64,21 @@ func NewAuthWithConfig(ctx context.Context, config *profilemanager.Config) *Auth
 // If it returns a flow info than save the configuration and return true. If it gets a codes.NotFound, it means that SSO
 // is not supported and returns false without saving the configuration. For other errors return false.
 func (a *Auth) SaveConfigIfSSOSupported() (bool, error) {
-	supportsSSO := true
-	err := a.withBackOff(a.ctx, func() (err error) {
-		_, err = internal.GetDeviceAuthorizationFlowInfo(a.ctx, a.config.PrivateKey, a.config.ManagementURL)
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.NotFound || s.Code() == codes.Unimplemented) {
-			_, err = internal.GetPKCEAuthorizationFlowInfo(a.ctx, a.config.PrivateKey, a.config.ManagementURL, nil)
-			if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.NotFound || s.Code() == codes.Unimplemented) {
-				supportsSSO = false
-				err = nil
-			}
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
+	if err != nil {
+		return false, fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
 
-			return err
-		}
-
-		return err
-	})
+	supportsSSO, err := authClient.IsSSOSupported(a.ctx)
+	if err != nil {
+		return false, fmt.Errorf("failed to check SSO support: %v", err)
+	}
 
 	if !supportsSSO {
 		return false, nil
 	}
 
-	if err != nil {
-		return false, fmt.Errorf("backoff cycle failed: %v", err)
-	}
-
 	err = profilemanager.WriteOutConfig(a.cfgPath, a.config)
 	return true, err
 }
@@ -103,32 +87,31 @@ func (a *Auth) SaveConfigIfSSOSupported() (bool, error) {
 func (a *Auth) LoginWithSetupKeyAndSaveConfig(setupKey string, deviceName string) error {
 	//nolint
 	ctxWithValues := context.WithValue(a.ctx, system.DeviceNameCtxKey, deviceName)
-
-	err := a.withBackOff(a.ctx, func() error {
-		backoffErr := internal.Login(ctxWithValues, a.config, setupKey, "")
-		if s, ok := gstatus.FromError(backoffErr); ok && (s.Code() == codes.PermissionDenied) {
-			// we got an answer from management, exit backoff earlier
-			return backoff.Permanent(backoffErr)
-		}
-		return backoffErr
-	})
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
+
+	err = authClient.Login(ctxWithValues, setupKey, "")
+	if err != nil {
+		return fmt.Errorf("login failed: %v", err)
 	}
 
 	return profilemanager.WriteOutConfig(a.cfgPath, a.config)
 }
 
 func (a *Auth) Login() error {
-	var needsLogin bool
+	authClient, err := auth.NewAuth(a.ctx, a.config.PrivateKey, a.config.ManagementURL, a.config)
+	if err != nil {
+		return fmt.Errorf("failed to create auth client: %v", err)
+	}
+	defer authClient.Close()
 
 	// check if we need to generate JWT token
-	err := a.withBackOff(a.ctx, func() (err error) {
-		needsLogin, err = internal.IsLoginRequired(a.ctx, a.config)
-		return
-	})
+	needsLogin, err := authClient.IsLoginRequired(a.ctx)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("failed to check login requirement: %v", err)
 	}
 
 	jwtToken := ""
@@ -136,25 +119,10 @@ func (a *Auth) Login() error {
 		return fmt.Errorf("Not authenticated")
 	}
 
-	err = a.withBackOff(a.ctx, func() error {
-		err := internal.Login(a.ctx, a.config, "", jwtToken)
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
-			return nil
-		}
-		return err
-	})
+	err = authClient.Login(a.ctx, "", jwtToken)
 	if err != nil {
-		return fmt.Errorf("backoff cycle failed: %v", err)
+		return fmt.Errorf("login failed: %v", err)
 	}
 
 	return nil
 }
-
-func (a *Auth) withBackOff(ctx context.Context, bf func() error) error {
-	return backoff.RetryNotify(
-		bf,
-		backoff.WithContext(cmd.CLIBackOffSettings, ctx),
-		func(err error, duration time.Duration) {
-			log.Warnf("retrying Login to the Management service in %v due to error %v", duration, err)
-		})
-}

client/proto/daemon.proto

@@ -24,7 +24,7 @@ service DaemonService {
   // Status of the service.
   rpc Status(StatusRequest) returns (StatusResponse) {}
 
-  // Down stops engine work in the daemon.
+  // Down engine work in the daemon.
   rpc Down(DownRequest) returns (DownResponse) {}
 
   // GetConfig of the daemon.
@@ -93,26 +93,9 @@ service DaemonService {
 
   // WaitJWTToken waits for JWT authentication completion
   rpc WaitJWTToken(WaitJWTTokenRequest) returns (WaitJWTTokenResponse) {}
-
-  rpc NotifyOSLifecycle(OSLifecycleRequest) returns(OSLifecycleResponse) {}
 }
 
 
-
-message OSLifecycleRequest {
-  // avoid collision with loglevel enum
-  enum CycleType {
-    UNKNOWN = 0;
-    SLEEP = 1;
-    WAKEUP = 2;
-  }
-
-  CycleType type = 1;
-}
-
-message OSLifecycleResponse {}
-
-
 message LoginRequest {
   // setupKey netbird setup key.
   string setupKey = 1;

client/proto/daemon_grpc.pb.go

@@ -27,7 +27,7 @@ type DaemonServiceClient interface {
 	Up(ctx context.Context, in *UpRequest, opts ...grpc.CallOption) (*UpResponse, error)
 	// Status of the service.
 	Status(ctx context.Context, in *StatusRequest, opts ...grpc.CallOption) (*StatusResponse, error)
-	// Down stops engine work in the daemon.
+	// Down engine work in the daemon.
 	Down(ctx context.Context, in *DownRequest, opts ...grpc.CallOption) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(ctx context.Context, in *GetConfigRequest, opts ...grpc.CallOption) (*GetConfigResponse, error)
@@ -70,7 +70,6 @@ type DaemonServiceClient interface {
 	RequestJWTAuth(ctx context.Context, in *RequestJWTAuthRequest, opts ...grpc.CallOption) (*RequestJWTAuthResponse, error)
 	// WaitJWTToken waits for JWT authentication completion
 	WaitJWTToken(ctx context.Context, in *WaitJWTTokenRequest, opts ...grpc.CallOption) (*WaitJWTTokenResponse, error)
-	NotifyOSLifecycle(ctx context.Context, in *OSLifecycleRequest, opts ...grpc.CallOption) (*OSLifecycleResponse, error)
 }
 
 type daemonServiceClient struct {
@@ -383,15 +382,6 @@ func (c *daemonServiceClient) WaitJWTToken(ctx context.Context, in *WaitJWTToken
 	return out, nil
 }
 
-func (c *daemonServiceClient) NotifyOSLifecycle(ctx context.Context, in *OSLifecycleRequest, opts ...grpc.CallOption) (*OSLifecycleResponse, error) {
-	out := new(OSLifecycleResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/NotifyOSLifecycle", in, out, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return out, nil
-}
-
 // DaemonServiceServer is the server API for DaemonService service.
 // All implementations must embed UnimplementedDaemonServiceServer
 // for forward compatibility
@@ -405,7 +395,7 @@ type DaemonServiceServer interface {
 	Up(context.Context, *UpRequest) (*UpResponse, error)
 	// Status of the service.
 	Status(context.Context, *StatusRequest) (*StatusResponse, error)
-	// Down stops engine work in the daemon.
+	// Down engine work in the daemon.
 	Down(context.Context, *DownRequest) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(context.Context, *GetConfigRequest) (*GetConfigResponse, error)
@@ -448,7 +438,6 @@ type DaemonServiceServer interface {
 	RequestJWTAuth(context.Context, *RequestJWTAuthRequest) (*RequestJWTAuthResponse, error)
 	// WaitJWTToken waits for JWT authentication completion
 	WaitJWTToken(context.Context, *WaitJWTTokenRequest) (*WaitJWTTokenResponse, error)
-	NotifyOSLifecycle(context.Context, *OSLifecycleRequest) (*OSLifecycleResponse, error)
 	mustEmbedUnimplementedDaemonServiceServer()
 }
 
@@ -549,9 +538,6 @@ func (UnimplementedDaemonServiceServer) RequestJWTAuth(context.Context, *Request
 func (UnimplementedDaemonServiceServer) WaitJWTToken(context.Context, *WaitJWTTokenRequest) (*WaitJWTTokenResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method WaitJWTToken not implemented")
 }
-func (UnimplementedDaemonServiceServer) NotifyOSLifecycle(context.Context, *OSLifecycleRequest) (*OSLifecycleResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method NotifyOSLifecycle not implemented")
-}
 func (UnimplementedDaemonServiceServer) mustEmbedUnimplementedDaemonServiceServer() {}
 
 // UnsafeDaemonServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -1126,24 +1112,6 @@ func _DaemonService_WaitJWTToken_Handler(srv interface{}, ctx context.Context, d
 	return interceptor(ctx, in, info, handler)
 }
 
-func _DaemonService_NotifyOSLifecycle_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(OSLifecycleRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(DaemonServiceServer).NotifyOSLifecycle(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: "/daemon.DaemonService/NotifyOSLifecycle",
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).NotifyOSLifecycle(ctx, req.(*OSLifecycleRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
 // DaemonService_ServiceDesc is the grpc.ServiceDesc for DaemonService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -1271,10 +1239,6 @@ var DaemonService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "WaitJWTToken",
 			Handler:    _DaemonService_WaitJWTToken_Handler,
 		},
-		{
-			MethodName: "NotifyOSLifecycle",
-			Handler:    _DaemonService_NotifyOSLifecycle_Handler,
-		},
 	},
 	Streams: []grpc.StreamDesc{
 		{

client/server/lifecycle.go

@@ -1,77 +0,0 @@
-package server
-
-import (
-	"context"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal"
-	"github.com/netbirdio/netbird/client/proto"
-)
-
-// NotifyOSLifecycle handles operating system lifecycle events by executing appropriate logic based on the request type.
-func (s *Server) NotifyOSLifecycle(callerCtx context.Context, req *proto.OSLifecycleRequest) (*proto.OSLifecycleResponse, error) {
-	switch req.GetType() {
-	case proto.OSLifecycleRequest_WAKEUP:
-		return s.handleWakeUp(callerCtx)
-	case proto.OSLifecycleRequest_SLEEP:
-		return s.handleSleep(callerCtx)
-	default:
-		log.Errorf("unknown OSLifecycleRequest type: %v", req.GetType())
-	}
-	return &proto.OSLifecycleResponse{}, nil
-}
-
-// handleWakeUp processes a wake-up event by triggering the Up command if the system was previously put to sleep.
-// It resets the sleep state and logs the process. Returns a response or an error if the Up command fails.
-func (s *Server) handleWakeUp(callerCtx context.Context) (*proto.OSLifecycleResponse, error) {
-	if !s.sleepTriggeredDown.Load() {
-		log.Info("skipping up because wasn't sleep down")
-		return &proto.OSLifecycleResponse{}, nil
-	}
-
-	// avoid other wakeup runs if sleep didn't make the computer sleep
-	s.sleepTriggeredDown.Store(false)
-
-	log.Info("running up after wake up")
-	_, err := s.Up(callerCtx, &proto.UpRequest{})
-	if err != nil {
-		log.Errorf("running up failed: %v", err)
-		return &proto.OSLifecycleResponse{}, err
-	}
-
-	log.Info("running up command executed successfully")
-	return &proto.OSLifecycleResponse{}, nil
-}
-
-// handleSleep handles the sleep event by initiating a "down" sequence if the system is in a connected or connecting state.
-func (s *Server) handleSleep(callerCtx context.Context) (*proto.OSLifecycleResponse, error) {
-	s.mutex.Lock()
-
-	state := internal.CtxGetState(s.rootCtx)
-	status, err := state.Status()
-	if err != nil {
-		s.mutex.Unlock()
-		return &proto.OSLifecycleResponse{}, err
-	}
-
-	if status != internal.StatusConnecting && status != internal.StatusConnected {
-		log.Infof("skipping setting the agent down because status is %s", status)
-		s.mutex.Unlock()
-		return &proto.OSLifecycleResponse{}, nil
-	}
-	s.mutex.Unlock()
-
-	log.Info("running down after system started sleeping")
-
-	_, err = s.Down(callerCtx, &proto.DownRequest{})
-	if err != nil {
-		log.Errorf("running down failed: %v", err)
-		return &proto.OSLifecycleResponse{}, err
-	}
-
-	s.sleepTriggeredDown.Store(true)
-
-	log.Info("running down executed successfully")
-	return &proto.OSLifecycleResponse{}, nil
-}

client/server/lifecycle_test.go

@@ -1,219 +0,0 @@
-package server
-
-import (
-	"context"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/internal"
-	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/client/proto"
-)
-
-func newTestServer() *Server {
-	ctx := internal.CtxInitState(context.Background())
-	return &Server{
-		rootCtx:        ctx,
-		statusRecorder: peer.NewRecorder(""),
-	}
-}
-
-func TestNotifyOSLifecycle_WakeUp_SkipsWhenNotSleepTriggered(t *testing.T) {
-	s := newTestServer()
-
-	// sleepTriggeredDown is false by default
-	assert.False(t, s.sleepTriggeredDown.Load())
-
-	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_WAKEUP,
-	})
-
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false")
-}
-
-func TestNotifyOSLifecycle_Sleep_SkipsWhenStatusIdle(t *testing.T) {
-	s := newTestServer()
-
-	state := internal.CtxGetState(s.rootCtx)
-	state.Set(internal.StatusIdle)
-
-	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_SLEEP,
-	})
-
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false when status is Idle")
-}
-
-func TestNotifyOSLifecycle_Sleep_SkipsWhenStatusNeedsLogin(t *testing.T) {
-	s := newTestServer()
-
-	state := internal.CtxGetState(s.rootCtx)
-	state.Set(internal.StatusNeedsLogin)
-
-	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_SLEEP,
-	})
-
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false when status is NeedsLogin")
-}
-
-func TestNotifyOSLifecycle_Sleep_SetsFlag_WhenConnecting(t *testing.T) {
-	s := newTestServer()
-
-	state := internal.CtxGetState(s.rootCtx)
-	state.Set(internal.StatusConnecting)
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	s.actCancel = cancel
-
-	resp, err := s.NotifyOSLifecycle(ctx, &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_SLEEP,
-	})
-
-	require.NoError(t, err)
-	assert.NotNil(t, resp, "handleSleep returns not nil response on success")
-	assert.True(t, s.sleepTriggeredDown.Load(), "flag should be set after sleep when connecting")
-}
-
-func TestNotifyOSLifecycle_Sleep_SetsFlag_WhenConnected(t *testing.T) {
-	s := newTestServer()
-
-	state := internal.CtxGetState(s.rootCtx)
-	state.Set(internal.StatusConnected)
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	s.actCancel = cancel
-
-	resp, err := s.NotifyOSLifecycle(ctx, &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_SLEEP,
-	})
-
-	require.NoError(t, err)
-	assert.NotNil(t, resp, "handleSleep returns not nil response on success")
-	assert.True(t, s.sleepTriggeredDown.Load(), "flag should be set after sleep when connected")
-}
-
-func TestNotifyOSLifecycle_WakeUp_ResetsFlag(t *testing.T) {
-	s := newTestServer()
-
-	// Manually set the flag to simulate prior sleep down
-	s.sleepTriggeredDown.Store(true)
-
-	// WakeUp will try to call Up which fails without proper setup, but flag should reset first
-	_, _ = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_WAKEUP,
-	})
-
-	assert.False(t, s.sleepTriggeredDown.Load(), "flag should be reset after WakeUp attempt")
-}
-
-func TestNotifyOSLifecycle_MultipleWakeUpCalls(t *testing.T) {
-	s := newTestServer()
-
-	// First wakeup without prior sleep - should be no-op
-	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_WAKEUP,
-	})
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	assert.False(t, s.sleepTriggeredDown.Load())
-
-	// Simulate prior sleep
-	s.sleepTriggeredDown.Store(true)
-
-	// First wakeup after sleep - should reset flag
-	_, _ = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_WAKEUP,
-	})
-	assert.False(t, s.sleepTriggeredDown.Load())
-
-	// Second wakeup - should be no-op
-	resp, err = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
-		Type: proto.OSLifecycleRequest_WAKEUP,
-	})
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-	assert.False(t, s.sleepTriggeredDown.Load())
-}
-
-func TestHandleWakeUp_SkipsWhenFlagFalse(t *testing.T) {
-	s := newTestServer()
-
-	resp, err := s.handleWakeUp(context.Background())
-
-	require.NoError(t, err)
-	require.NotNil(t, resp)
-}
-
-func TestHandleWakeUp_ResetsFlagBeforeUp(t *testing.T) {
-	s := newTestServer()
-	s.sleepTriggeredDown.Store(true)
-
-	// Even if Up fails, flag should be reset
-	_, _ = s.handleWakeUp(context.Background())
-
-	assert.False(t, s.sleepTriggeredDown.Load(), "flag must be reset before calling Up")
-}
-
-func TestHandleSleep_SkipsForNonActiveStates(t *testing.T) {
-	tests := []struct {
-		name   string
-		status internal.StatusType
-	}{
-		{"Idle", internal.StatusIdle},
-		{"NeedsLogin", internal.StatusNeedsLogin},
-		{"LoginFailed", internal.StatusLoginFailed},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			s := newTestServer()
-			state := internal.CtxGetState(s.rootCtx)
-			state.Set(tt.status)
-
-			resp, err := s.handleSleep(context.Background())
-
-			require.NoError(t, err)
-			require.NotNil(t, resp)
-			assert.False(t, s.sleepTriggeredDown.Load())
-		})
-	}
-}
-
-func TestHandleSleep_ProceedsForActiveStates(t *testing.T) {
-	tests := []struct {
-		name   string
-		status internal.StatusType
-	}{
-		{"Connecting", internal.StatusConnecting},
-		{"Connected", internal.StatusConnected},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			s := newTestServer()
-			state := internal.CtxGetState(s.rootCtx)
-			state.Set(tt.status)
-
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-			s.actCancel = cancel
-
-			resp, err := s.handleSleep(ctx)
-
-			require.NoError(t, err)
-			assert.NotNil(t, resp)
-			assert.True(t, s.sleepTriggeredDown.Load())
-		})
-	}
-}

client/server/server.go

@@ -85,9 +85,6 @@ type Server struct {
 	profilesDisabled       bool
 	updateSettingsDisabled bool
 
-	// sleepTriggeredDown holds a state indicated if the sleep handler triggered the last client down
-	sleepTriggeredDown atomic.Bool
-
 	jwtCache *jwtCache
 }
 
@@ -281,14 +278,22 @@ func (s *Server) connectWithRetryRuns(ctx context.Context, profileConfig *profil
 
 // loginAttempt attempts to login using the provided information. it returns a status in case something fails
 func (s *Server) loginAttempt(ctx context.Context, setupKey, jwtToken string) (internal.StatusType, error) {
-	var status internal.StatusType
-	err := internal.Login(ctx, s.config, setupKey, jwtToken)
+	authClient, err := auth.NewAuth(ctx, s.config.PrivateKey, s.config.ManagementURL, s.config)
 	if err != nil {
-		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
-			log.Warnf("failed login: %v", err)
+		log.Errorf("failed to create auth client: %v", err)
+		return internal.StatusLoginFailed, err
+	}
+	defer authClient.Close()
+
+	var status internal.StatusType
+	err = authClient.Login(ctx, setupKey, jwtToken)
+	if err != nil {
+		// Check if it's an authentication error (permission denied, invalid credentials, unauthenticated)
+		if errors.Is(err, mgm.ErrPermissionDenied) || errors.Is(err, mgm.ErrInvalidArgument) || errors.Is(err, mgm.ErrUnauthenticated) {
+			log.Warnf("authentication failed: %v", err)
 			status = internal.StatusNeedsLogin
 		} else {
-			log.Errorf("failed login: %v", err)
+			log.Errorf("login failed: %v", err)
 			status = internal.StatusLoginFailed
 		}
 		return status, err
@@ -609,8 +614,7 @@ func (s *Server) WaitSSOLogin(callerCtx context.Context, msg *proto.WaitSSOLogin
 		s.oauthAuthFlow.waitCancel()
 	}
 
-	waitTimeout := time.Until(s.oauthAuthFlow.expiresAt)
-	waitCTX, cancel := context.WithTimeout(ctx, waitTimeout)
+	waitCTX, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	s.mutex.Lock()
@@ -1005,8 +1009,8 @@ func (s *Server) sendLogoutRequestWithConfig(ctx context.Context, config *profil
 	}
 
 	mgmTlsEnabled := config.ManagementURL.Scheme == "https"
-	mgmClient, err := mgm.NewClient(ctx, config.ManagementURL.Host, key, mgmTlsEnabled)
-	if err != nil {
+	mgmClient := mgm.NewClient(config.ManagementURL.Host, key, mgmTlsEnabled)
+	if err := mgmClient.Connect(ctx); err != nil {
 		return fmt.Errorf("connect to management server: %w", err)
 	}
 	defer func() {
@@ -1015,7 +1019,7 @@ func (s *Server) sendLogoutRequestWithConfig(ctx context.Context, config *profil
 		}
 	}()
 
-	return mgmClient.Logout()
+	return mgmClient.Logout(ctx)
 }
 
 // Status returns the daemon status

client/server/server_test.go

@@ -17,12 +17,11 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/groups"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
@@ -36,6 +35,7 @@ import (
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/peers"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -316,17 +316,14 @@ func startManagement(t *testing.T, signalAddr string, counter *int) (*grpc.Serve
 
 	requestBuffer := server.NewAccountRequestBuffer(context.Background(), store)
 	peersUpdateManager := update_channel.NewPeersUpdateManager(metrics)
-	networkMapController := controller.NewController(context.Background(), store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), manager.NewEphemeralManager(store, peersManager), config)
+	networkMapController := controller.NewController(context.Background(), store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), config)
 	accountManager, err := server.BuildManager(context.Background(), config, store, networkMapController, nil, "", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock, false)
 	if err != nil {
 		return nil, "", err
 	}
 
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		return nil, "", err
-	}
-	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, secretsManager, nil, nil, &server.MockIntegratedValidator{}, networkMapController)
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
+	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, &manager.EphemeralManager{}, nil, &server.MockIntegratedValidator{}, networkMapController)
 	if err != nil {
 		return nil, "", err
 	}

client/ssh/client/client.go

@@ -24,7 +24,6 @@ import (
 	"github.com/netbirdio/netbird/client/proto"
 	nbssh "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/ssh/detection"
-	"github.com/netbirdio/netbird/util"
 )
 
 const (
@@ -279,7 +278,6 @@ type DialOptions struct {
 	DaemonAddr         string
 	SkipCachedToken    bool
 	InsecureSkipVerify bool
-	NoBrowser          bool
 }
 
 // Dial connects to the given ssh server with specified options
@@ -309,7 +307,7 @@ func Dial(ctx context.Context, addr, user string, opts DialOptions) (*Client, er
 		config.Auth = append(config.Auth, authMethod)
 	}
 
-	return dialWithJWT(ctx, "tcp", addr, config, daemonAddr, opts.SkipCachedToken, opts.NoBrowser)
+	return dialWithJWT(ctx, "tcp", addr, config, daemonAddr, opts.SkipCachedToken)
 }
 
 // dialSSH establishes an SSH connection without JWT authentication
@@ -335,7 +333,7 @@ func dialSSH(ctx context.Context, network, addr string, config *ssh.ClientConfig
 }
 
 // dialWithJWT establishes an SSH connection with optional JWT authentication based on server detection
-func dialWithJWT(ctx context.Context, network, addr string, config *ssh.ClientConfig, daemonAddr string, skipCache, noBrowser bool) (*Client, error) {
+func dialWithJWT(ctx context.Context, network, addr string, config *ssh.ClientConfig, daemonAddr string, skipCache bool) (*Client, error) {
 	host, portStr, err := net.SplitHostPort(addr)
 	if err != nil {
 		return nil, fmt.Errorf("parse address %s: %w", addr, err)
@@ -361,7 +359,7 @@ func dialWithJWT(ctx context.Context, network, addr string, config *ssh.ClientCo
 	jwtCtx, cancel := context.WithTimeout(ctx, config.Timeout)
 	defer cancel()
 
-	jwtToken, err := requestJWTToken(jwtCtx, daemonAddr, skipCache, noBrowser)
+	jwtToken, err := requestJWTToken(jwtCtx, daemonAddr, skipCache)
 	if err != nil {
 		return nil, fmt.Errorf("request JWT token: %w", err)
 	}
@@ -371,7 +369,7 @@ func dialWithJWT(ctx context.Context, network, addr string, config *ssh.ClientCo
 }
 
 // requestJWTToken requests a JWT token from the NetBird daemon
-func requestJWTToken(ctx context.Context, daemonAddr string, skipCache, noBrowser bool) (string, error) {
+func requestJWTToken(ctx context.Context, daemonAddr string, skipCache bool) (string, error) {
 	hint := profilemanager.GetLoginHint()
 
 	conn, err := connectToDaemon(daemonAddr)
@@ -381,13 +379,7 @@ func requestJWTToken(ctx context.Context, daemonAddr string, skipCache, noBrowse
 	defer conn.Close()
 
 	client := proto.NewDaemonServiceClient(conn)
-
-	var browserOpener func(string) error
-	if !noBrowser {
-		browserOpener = util.OpenBrowser
-	}
-
-	return nbssh.RequestJWTToken(ctx, client, os.Stdout, os.Stderr, !skipCache, hint, browserOpener)
+	return nbssh.RequestJWTToken(ctx, client, os.Stdout, os.Stderr, !skipCache, hint)
 }
 
 // verifyHostKeyViaDaemon verifies SSH host key by querying the NetBird daemon

client/ssh/common.go

@@ -67,31 +67,8 @@ func (d *DaemonHostKeyVerifier) VerifySSHHostKey(peerAddress string, presentedKe
 	return VerifyHostKey(storedKeyData, presentedKey, peerAddress)
 }
 
-// printAuthInstructions prints authentication instructions to stderr
-func printAuthInstructions(stderr io.Writer, authResponse *proto.RequestJWTAuthResponse, browserWillOpen bool) {
-	_, _ = fmt.Fprintln(stderr, "SSH authentication required.")
-
-	if browserWillOpen {
-		_, _ = fmt.Fprintln(stderr, "Please do the SSO login in your browser.")
-		_, _ = fmt.Fprintln(stderr, "If your browser didn't open automatically, use this URL to log in:")
-		_, _ = fmt.Fprintln(stderr)
-	}
-
-	_, _ = fmt.Fprintf(stderr, "%s\n", authResponse.VerificationURIComplete)
-
-	if authResponse.UserCode != "" {
-		_, _ = fmt.Fprintf(stderr, "Or visit: %s and enter code: %s\n", authResponse.VerificationURI, authResponse.UserCode)
-	}
-
-	if browserWillOpen {
-		_, _ = fmt.Fprintln(stderr)
-	}
-
-	_, _ = fmt.Fprintln(stderr, "Waiting for authentication...")
-}
-
 // RequestJWTToken requests or retrieves a JWT token for SSH authentication
-func RequestJWTToken(ctx context.Context, client proto.DaemonServiceClient, stdout, stderr io.Writer, useCache bool, hint string, openBrowser func(string) error) (string, error) {
+func RequestJWTToken(ctx context.Context, client proto.DaemonServiceClient, stdout, stderr io.Writer, useCache bool, hint string) (string, error) {
 	req := &proto.RequestJWTAuthRequest{}
 	if hint != "" {
 		req.Hint = &hint
@@ -107,13 +84,12 @@ func RequestJWTToken(ctx context.Context, client proto.DaemonServiceClient, stdo
 	}
 
 	if stderr != nil {
-		printAuthInstructions(stderr, authResponse, openBrowser != nil)
-	}
-
-	if openBrowser != nil {
-		if err := openBrowser(authResponse.VerificationURIComplete); err != nil {
-			log.Debugf("open browser: %v", err)
+		_, _ = fmt.Fprintln(stderr, "SSH authentication required.")
+		_, _ = fmt.Fprintf(stderr, "Please visit: %s\n", authResponse.VerificationURIComplete)
+		if authResponse.UserCode != "" {
+			_, _ = fmt.Fprintf(stderr, "Or visit: %s and enter code: %s\n", authResponse.VerificationURI, authResponse.UserCode)
 		}
+		_, _ = fmt.Fprintln(stderr, "Waiting for authentication...")
 	}
 
 	tokenResponse, err := client.WaitJWTToken(ctx, &proto.WaitJWTTokenRequest{

client/ssh/proxy/proxy.go

@@ -35,16 +35,15 @@ const (
 )
 
 type SSHProxy struct {
-	daemonAddr    string
-	targetHost    string
-	targetPort    int
-	stderr        io.Writer
-	conn          *grpc.ClientConn
-	daemonClient  proto.DaemonServiceClient
-	browserOpener func(string) error
+	daemonAddr   string
+	targetHost   string
+	targetPort   int
+	stderr       io.Writer
+	conn         *grpc.ClientConn
+	daemonClient proto.DaemonServiceClient
 }
 
-func New(daemonAddr, targetHost string, targetPort int, stderr io.Writer, browserOpener func(string) error) (*SSHProxy, error) {
+func New(daemonAddr, targetHost string, targetPort int, stderr io.Writer) (*SSHProxy, error) {
 	grpcAddr := strings.TrimPrefix(daemonAddr, "tcp://")
 	grpcConn, err := grpc.NewClient(grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
@@ -52,13 +51,12 @@ func New(daemonAddr, targetHost string, targetPort int, stderr io.Writer, browse
 	}
 
 	return &SSHProxy{
-		daemonAddr:    daemonAddr,
-		targetHost:    targetHost,
-		targetPort:    targetPort,
-		stderr:        stderr,
-		conn:          grpcConn,
-		daemonClient:  proto.NewDaemonServiceClient(grpcConn),
-		browserOpener: browserOpener,
+		daemonAddr:   daemonAddr,
+		targetHost:   targetHost,
+		targetPort:   targetPort,
+		stderr:       stderr,
+		conn:         grpcConn,
+		daemonClient: proto.NewDaemonServiceClient(grpcConn),
 	}, nil
 }
 
@@ -72,7 +70,7 @@ func (p *SSHProxy) Close() error {
 func (p *SSHProxy) Connect(ctx context.Context) error {
 	hint := profilemanager.GetLoginHint()
 
-	jwtToken, err := nbssh.RequestJWTToken(ctx, p.daemonClient, nil, p.stderr, true, hint, p.browserOpener)
+	jwtToken, err := nbssh.RequestJWTToken(ctx, p.daemonClient, nil, p.stderr, true, hint)
 	if err != nil {
 		return fmt.Errorf(jwtAuthErrorMsg, err)
 	}

client/ssh/proxy/proxy_test.go

@@ -153,7 +153,7 @@ func TestSSHProxy_Connect(t *testing.T) {
 	validToken := generateValidJWT(t, privateKey, issuer, audience)
 	mockDaemon.setJWTToken(validToken)
 
-	proxyInstance, err := New(mockDaemon.addr, host, port, nil, nil)
+	proxyInstance, err := New(mockDaemon.addr, host, port, nil)
 	require.NoError(t, err)
 
 	clientConn, proxyConn := net.Pipe()

client/ui/client_ui.go

@@ -38,7 +38,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
-	"github.com/netbirdio/netbird/client/internal/sleep"
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/ui/desktop"
 	"github.com/netbirdio/netbird/client/ui/event"
@@ -210,11 +209,10 @@ var iconConnectedDot []byte
 var iconDisconnectedDot []byte
 
 type serviceClient struct {
-	ctx      context.Context
-	cancel   context.CancelFunc
-	addr     string
-	conn     proto.DaemonServiceClient
-	connLock sync.Mutex
+	ctx    context.Context
+	cancel context.CancelFunc
+	addr   string
+	conn   proto.DaemonServiceClient
 
 	eventHandler *eventHandler
 
@@ -1100,9 +1098,6 @@ func (s *serviceClient) onTrayReady() {
 
 	go s.eventManager.Start(s.ctx)
 	go s.eventHandler.listen(s.ctx)
-
-	// Start sleep detection listener
-	go s.startSleepListener()
 }
 
 func (s *serviceClient) attachOutput(cmd *exec.Cmd) *os.File {
@@ -1139,8 +1134,6 @@ func (s *serviceClient) onTrayExit() {
 
 // getSrvClient connection to the service.
 func (s *serviceClient) getSrvClient(timeout time.Duration) (proto.DaemonServiceClient, error) {
-	s.connLock.Lock()
-	defer s.connLock.Unlock()
 	if s.conn != nil {
 		return s.conn, nil
 	}
@@ -1163,62 +1156,6 @@ func (s *serviceClient) getSrvClient(timeout time.Duration) (proto.DaemonService
 	return s.conn, nil
 }
 
-// startSleepListener initializes the sleep detection service and listens for sleep events
-func (s *serviceClient) startSleepListener() {
-	sleepService, err := sleep.New()
-	if err != nil {
-		log.Warnf("%v", err)
-		return
-	}
-
-	if err := sleepService.Register(s.handleSleepEvents); err != nil {
-		log.Errorf("failed to start sleep detection: %v", err)
-		return
-	}
-
-	log.Info("sleep detection service initialized")
-
-	// Cleanup on context cancellation
-	go func() {
-		<-s.ctx.Done()
-		log.Info("stopping sleep event listener")
-		if err := sleepService.Deregister(); err != nil {
-			log.Errorf("failed to deregister sleep detection: %v", err)
-		}
-	}()
-}
-
-// handleSleepEvents sends a sleep notification to the daemon via gRPC
-func (s *serviceClient) handleSleepEvents(event sleep.EventType) {
-	conn, err := s.getSrvClient(0)
-	if err != nil {
-		log.Errorf("failed to get daemon client for sleep notification: %v", err)
-		return
-	}
-
-	req := &proto.OSLifecycleRequest{}
-
-	switch event {
-	case sleep.EventTypeWakeUp:
-		log.Infof("handle wakeup event: %v", event)
-		req.Type = proto.OSLifecycleRequest_WAKEUP
-	case sleep.EventTypeSleep:
-		log.Infof("handle sleep event: %v", event)
-		req.Type = proto.OSLifecycleRequest_SLEEP
-	default:
-		log.Infof("unknown event: %v", event)
-		return
-	}
-
-	_, err = conn.NotifyOSLifecycle(s.ctx, req)
-	if err != nil {
-		log.Errorf("failed to notify daemon about os lifecycle notification: %v", err)
-		return
-	}
-
-	log.Info("successfully notified daemon about os lifecycle")
-}
-
 // setSettingsEnabled enables or disables the settings menu based on the provided state
 func (s *serviceClient) setSettingsEnabled(enabled bool) {
 	if s.mSettings != nil {

client/ui/process/process.go

@@ -28,8 +28,7 @@ func IsAnotherProcessRunning() (int32, bool, error) {
 			continue
 		}
 
-		runningProcessName := strings.ToLower(filepath.Base(runningProcessPath))
-		if runningProcessName == processName && isProcessOwnedByCurrentUser(p) {
+		if strings.Contains(strings.ToLower(runningProcessPath), processName) && isProcessOwnedByCurrentUser(p) {
 			return p.Pid, true, nil
 		}
 	}

go.mod

@@ -64,7 +64,7 @@ require (
 	github.com/mdlayher/socket v0.5.1
 	github.com/miekg/dns v1.1.59
 	github.com/mitchellh/hashstructure/v2 v2.0.2
-	github.com/netbirdio/management-integrations/integrations v0.0.0-20251202114414-534cf891e0ba
+	github.com/netbirdio/management-integrations/integrations v0.0.0-20251027212525-d751b79f5d48
 	github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45
 	github.com/okta/okta-sdk-golang/v2 v2.18.0
 	github.com/oschwald/maxminddb-golang v1.12.0

go.sum

@@ -368,8 +368,8 @@ github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944 h1:TDtJKmM6S
 github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944/go.mod h1:sHA6TRxjQ6RLbnI+3R4DZo2Eseg/iKiPRfNmcuNySVQ=
 github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51 h1:Ov4qdafATOgGMB1wbSuh+0aAHcwz9hdvB6VZjh1mVMI=
 github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51/go.mod h1:ZSIbPdBn5hePO8CpF1PekH2SfpTxg1PDhEwtbqZS7R8=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20251202114414-534cf891e0ba h1:pD6eygRJ5EYAlgzeNskPU3WqszMz6/HhPuc6/Bc/580=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20251202114414-534cf891e0ba/go.mod h1:qzLCKeR253jtsWhfZTt4fyegI5zei32jKZykV+oSQOo=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20251027212525-d751b79f5d48 h1:moJbL1uuaWR35yUgHZ6suijjqqW8/qGCuPPBXu5MeWQ=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20251027212525-d751b79f5d48/go.mod h1:ifKa2jGPsOzZhJFo72v2AE5nMP3GYvlhoZ9JV6lHlJ8=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502 h1:3tHlFmhTdX9axERMVN63dqyFqnvuD+EMJHzM7mNGON8=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45 h1:ujgviVYmx243Ksy7NdSwrdGPSRNE3pb8kEDSpH0QuAQ=

management/internals/controllers/network_map/controller/controller.go

@@ -19,7 +19,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server/account"
@@ -43,7 +42,6 @@ type Controller struct {
 	accountManagerMetrics *telemetry.AccountManagerMetrics
 	peersUpdateManager    network_map.PeersUpdateManager
 	settingsManager       settings.Manager
-	EphemeralPeersManager ephemeral.Manager
 
 	accountUpdateLocks               sync.Map
 	sendAccountUpdateLocks           sync.Map
@@ -72,7 +70,7 @@ type bufferUpdate struct {
 
 var _ network_map.Controller = (*Controller)(nil)
 
-func NewController(ctx context.Context, store store.Store, metrics telemetry.AppMetrics, peersUpdateManager network_map.PeersUpdateManager, requestBuffer account.RequestBuffer, integratedPeerValidator integrated_validator.IntegratedValidator, settingsManager settings.Manager, dnsDomain string, proxyController port_forwarding.Controller, ephemeralPeersManager ephemeral.Manager, config *config.Config) *Controller {
+func NewController(ctx context.Context, store store.Store, metrics telemetry.AppMetrics, peersUpdateManager network_map.PeersUpdateManager, requestBuffer account.RequestBuffer, integratedPeerValidator integrated_validator.IntegratedValidator, settingsManager settings.Manager, dnsDomain string, proxyController port_forwarding.Controller, config *config.Config) *Controller {
 	nMetrics, err := newMetrics(metrics.UpdateChannelMetrics())
 	if err != nil {
 		log.Fatal(fmt.Errorf("error creating metrics: %w", err))
@@ -101,8 +99,7 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 		dnsDomain:               dnsDomain,
 		config:                  config,
 
-		proxyController:       proxyController,
-		EphemeralPeersManager: ephemeralPeersManager,
+		proxyController: proxyController,
 
 		holder:               types.NewHolder(),
 		expNewNetworkMap:     newNetworkMapBuilder,
@@ -110,31 +107,6 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 	}
 }
 
-func (c *Controller) OnPeerConnected(ctx context.Context, accountID string, peerID string) (chan *network_map.UpdateMessage, error) {
-	peer, err := c.repo.GetPeerByID(ctx, accountID, peerID)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get peer %s: %v", peerID, err)
-	}
-
-	c.EphemeralPeersManager.OnPeerConnected(ctx, peer)
-
-	return c.peersUpdateManager.CreateChannel(ctx, peerID), nil
-}
-
-func (c *Controller) OnPeerDisconnected(ctx context.Context, accountID string, peerID string) {
-	c.peersUpdateManager.CloseChannel(ctx, peerID)
-	peer, err := c.repo.GetPeerByID(ctx, accountID, peerID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get peer %s: %v", peerID, err)
-		return
-	}
-	c.EphemeralPeersManager.OnPeerDisconnected(ctx, peer)
-}
-
-func (c *Controller) CountStreams() int {
-	return c.peersUpdateManager.CountStreams()
-}
-
 func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID string) error {
 	log.WithContext(ctx).Tracef("updating peers for account %s from %s", accountID, util.GetCallerName())
 	var (
@@ -394,6 +366,38 @@ func (c *Controller) BufferUpdateAccountPeers(ctx context.Context, accountID str
 	return nil
 }
 
+func (c *Controller) DeletePeer(ctx context.Context, accountId string, peerId string) error {
+	network, err := c.repo.GetAccountNetwork(ctx, accountId)
+	if err != nil {
+		return err
+	}
+
+	peers, err := c.repo.GetAccountPeers(ctx, accountId)
+	if err != nil {
+		return err
+	}
+
+	dnsFwdPort := computeForwarderPort(peers, network_map.DnsForwarderPortMinVersion)
+	c.peersUpdateManager.SendUpdate(ctx, peerId, &network_map.UpdateMessage{
+		Update: &proto.SyncResponse{
+			RemotePeers:        []*proto.RemotePeerConfig{},
+			RemotePeersIsEmpty: true,
+			NetworkMap: &proto.NetworkMap{
+				Serial:               network.CurrentSerial(),
+				RemotePeers:          []*proto.RemotePeerConfig{},
+				RemotePeersIsEmpty:   true,
+				FirewallRules:        []*proto.FirewallRule{},
+				FirewallRulesIsEmpty: true,
+				DNSConfig: &proto.DNSConfig{
+					ForwarderPort: dnsFwdPort,
+				},
+			},
+		},
+	})
+	c.peersUpdateManager.CloseChannel(ctx, peerId)
+	return nil
+}
+
 func (c *Controller) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, int64, error) {
 	if isRequiresApproval {
 		network, err := c.repo.GetAccountNetwork(ctx, accountID)
@@ -694,83 +698,35 @@ func isPeerInPolicySourceGroups(account *types.Account, peerID string, policy *t
 	return false, nil
 }
 
-func (c *Controller) OnPeersUpdated(ctx context.Context, accountID string, peerIDs []string) error {
-	peers, err := c.repo.GetPeersByIDs(ctx, accountID, peerIDs)
-	if err != nil {
-		return fmt.Errorf("failed to get peers by ids: %w", err)
-	}
-
-	for _, peer := range peers {
-		c.UpdatePeerInNetworkMapCache(accountID, peer)
-	}
-
-	err = c.bufferSendUpdateAccountPeers(ctx, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to buffer update account peers for peer update in account %s: %v", accountID, err)
-	}
-
-	return nil
+func (c *Controller) OnPeerUpdated(accountId string, peer *nbpeer.Peer) {
+	c.UpdatePeerInNetworkMapCache(accountId, peer)
+	_ = c.bufferSendUpdateAccountPeers(context.Background(), accountId)
 }
 
-func (c *Controller) OnPeersAdded(ctx context.Context, accountID string, peerIDs []string) error {
-	for _, peerID := range peerIDs {
-		if c.experimentalNetworkMap(accountID) {
-			account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
-			if err != nil {
-				return err
-			}
+func (c *Controller) OnPeerAdded(ctx context.Context, accountID string, peerID string) error {
+	if c.experimentalNetworkMap(accountID) {
+		account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
+		if err != nil {
+			return err
+		}
 
-			err = c.onPeerAddedUpdNetworkMapCache(account, peerID)
-			if err != nil {
-				return err
-			}
+		err = c.onPeerAddedUpdNetworkMapCache(account, peerID)
+		if err != nil {
+			return err
 		}
 	}
 	return c.bufferSendUpdateAccountPeers(ctx, accountID)
 }
 
-func (c *Controller) OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string) error {
-	network, err := c.repo.GetAccountNetwork(ctx, accountID)
-	if err != nil {
-		return err
-	}
-
-	peers, err := c.repo.GetAccountPeers(ctx, accountID)
-	if err != nil {
-		return err
-	}
-
-	dnsFwdPort := computeForwarderPort(peers, network_map.DnsForwarderPortMinVersion)
-	for _, peerID := range peerIDs {
-		c.peersUpdateManager.SendUpdate(ctx, peerID, &network_map.UpdateMessage{
-			Update: &proto.SyncResponse{
-				RemotePeers:        []*proto.RemotePeerConfig{},
-				RemotePeersIsEmpty: true,
-				NetworkMap: &proto.NetworkMap{
-					Serial:               network.CurrentSerial(),
-					RemotePeers:          []*proto.RemotePeerConfig{},
-					RemotePeersIsEmpty:   true,
-					FirewallRules:        []*proto.FirewallRule{},
-					FirewallRulesIsEmpty: true,
-					DNSConfig: &proto.DNSConfig{
-						ForwarderPort: dnsFwdPort,
-					},
-				},
-			},
-		})
-		c.peersUpdateManager.CloseChannel(ctx, peerID)
-
-		if c.experimentalNetworkMap(accountID) {
-			account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
-			if err != nil {
-				log.WithContext(ctx).Errorf("failed to get account %s: %v", accountID, err)
-				continue
-			}
-			err = c.onPeerDeletedUpdNetworkMapCache(account, peerID)
-			if err != nil {
-				log.WithContext(ctx).Errorf("failed to update network map cache for deleted peer %s in account %s: %v", peerID, accountID, err)
-				continue
-			}
+func (c *Controller) OnPeerDeleted(ctx context.Context, accountID string, peerID string) error {
+	if c.experimentalNetworkMap(accountID) {
+		account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
+		if err != nil {
+			return err
+		}
+		err = c.onPeerDeletedUpdNetworkMapCache(account, peerID)
+		if err != nil {
+			return err
 		}
 	}
 
@@ -822,6 +778,10 @@ func (c *Controller) GetNetworkMap(ctx context.Context, peerID string) (*types.N
 	return networkMap, nil
 }
 
-func (c *Controller) DisconnectPeers(ctx context.Context, accountId string, peerIDs []string) {
+func (c *Controller) DisconnectPeers(ctx context.Context, peerIDs []string) {
 	c.peersUpdateManager.CloseChannels(ctx, peerIDs)
 }
+
+func (c *Controller) IsConnected(peerID string) bool {
+	return c.peersUpdateManager.HasChannel(peerID)
+}

management/internals/controllers/network_map/controller/repository.go

@@ -12,8 +12,6 @@ type Repository interface {
 	GetAccountNetwork(ctx context.Context, accountID string) (*types.Network, error)
 	GetAccountPeers(ctx context.Context, accountID string) ([]*peer.Peer, error)
 	GetAccountByPeerID(ctx context.Context, peerID string) (*types.Account, error)
-	GetPeersByIDs(ctx context.Context, accountID string, peerIDs []string) (map[string]*peer.Peer, error)
-	GetPeerByID(ctx context.Context, accountID string, peerID string) (*peer.Peer, error)
 }
 
 type repository struct {
@@ -39,11 +37,3 @@ func (r *repository) GetAccountPeers(ctx context.Context, accountID string) ([]*
 func (r *repository) GetAccountByPeerID(ctx context.Context, peerID string) (*types.Account, error) {
 	return r.store.GetAccountByPeerID(ctx, peerID)
 }
-
-func (r *repository) GetPeersByIDs(ctx context.Context, accountID string, peerIDs []string) (map[string]*peer.Peer, error) {
-	return r.store.GetPeersByIDs(ctx, store.LockingStrengthNone, accountID, peerIDs)
-}
-
-func (r *repository) GetPeerByID(ctx context.Context, accountID string, peerID string) (*peer.Peer, error) {
-	return r.store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
-}

management/internals/controllers/network_map/interface.go

@@ -28,12 +28,12 @@ type Controller interface {
 	GetDNSDomain(settings *types.Settings) string
 	StartWarmup(context.Context)
 	GetNetworkMap(ctx context.Context, peerID string) (*types.NetworkMap, error)
-	CountStreams() int
 
-	OnPeersUpdated(ctx context.Context, accountId string, peerIDs []string) error
-	OnPeersAdded(ctx context.Context, accountID string, peerIDs []string) error
-	OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string) error
-	DisconnectPeers(ctx context.Context, accountId string, peerIDs []string)
-	OnPeerConnected(ctx context.Context, accountID string, peerID string) (chan *UpdateMessage, error)
-	OnPeerDisconnected(ctx context.Context, accountID string, peerID string)
+	DeletePeer(ctx context.Context, accountId string, peerId string) error
+
+	OnPeerUpdated(accountId string, peer *nbpeer.Peer)
+	OnPeerAdded(ctx context.Context, accountID string, peerID string) error
+	OnPeerDeleted(ctx context.Context, accountID string, peerID string) error
+	DisconnectPeers(ctx context.Context, peerIDs []string)
+	IsConnected(peerID string) bool
 }

management/internals/controllers/network_map/interface_mock.go

@@ -57,30 +57,30 @@ func (mr *MockControllerMockRecorder) BufferUpdateAccountPeers(ctx, accountID an
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BufferUpdateAccountPeers", reflect.TypeOf((*MockController)(nil).BufferUpdateAccountPeers), ctx, accountID)
 }
 
-// CountStreams mocks base method.
-func (m *MockController) CountStreams() int {
+// DeletePeer mocks base method.
+func (m *MockController) DeletePeer(ctx context.Context, accountId, peerId string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "CountStreams")
-	ret0, _ := ret[0].(int)
+	ret := m.ctrl.Call(m, "DeletePeer", ctx, accountId, peerId)
+	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-// CountStreams indicates an expected call of CountStreams.
-func (mr *MockControllerMockRecorder) CountStreams() *gomock.Call {
+// DeletePeer indicates an expected call of DeletePeer.
+func (mr *MockControllerMockRecorder) DeletePeer(ctx, accountId, peerId any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountStreams", reflect.TypeOf((*MockController)(nil).CountStreams))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeletePeer", reflect.TypeOf((*MockController)(nil).DeletePeer), ctx, accountId, peerId)
 }
 
 // DisconnectPeers mocks base method.
-func (m *MockController) DisconnectPeers(ctx context.Context, accountId string, peerIDs []string) {
+func (m *MockController) DisconnectPeers(ctx context.Context, peerIDs []string) {
 	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "DisconnectPeers", ctx, accountId, peerIDs)
+	m.ctrl.Call(m, "DisconnectPeers", ctx, peerIDs)
 }
 
 // DisconnectPeers indicates an expected call of DisconnectPeers.
-func (mr *MockControllerMockRecorder) DisconnectPeers(ctx, accountId, peerIDs any) *gomock.Call {
+func (mr *MockControllerMockRecorder) DisconnectPeers(ctx, peerIDs any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DisconnectPeers", reflect.TypeOf((*MockController)(nil).DisconnectPeers), ctx, accountId, peerIDs)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DisconnectPeers", reflect.TypeOf((*MockController)(nil).DisconnectPeers), ctx, peerIDs)
 }
 
 // GetDNSDomain mocks base method.
@@ -130,73 +130,58 @@ func (mr *MockControllerMockRecorder) GetValidatedPeerWithMap(ctx, isRequiresApp
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetValidatedPeerWithMap", reflect.TypeOf((*MockController)(nil).GetValidatedPeerWithMap), ctx, isRequiresApproval, accountID, p)
 }
 
-// OnPeerConnected mocks base method.
-func (m *MockController) OnPeerConnected(ctx context.Context, accountID, peerID string) (chan *UpdateMessage, error) {
+// IsConnected mocks base method.
+func (m *MockController) IsConnected(peerID string) bool {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "OnPeerConnected", ctx, accountID, peerID)
-	ret0, _ := ret[0].(chan *UpdateMessage)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret := m.ctrl.Call(m, "IsConnected", peerID)
+	ret0, _ := ret[0].(bool)
+	return ret0
 }
 
-// OnPeerConnected indicates an expected call of OnPeerConnected.
-func (mr *MockControllerMockRecorder) OnPeerConnected(ctx, accountID, peerID any) *gomock.Call {
+// IsConnected indicates an expected call of IsConnected.
+func (mr *MockControllerMockRecorder) IsConnected(peerID any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerConnected", reflect.TypeOf((*MockController)(nil).OnPeerConnected), ctx, accountID, peerID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsConnected", reflect.TypeOf((*MockController)(nil).IsConnected), peerID)
 }
 
-// OnPeerDisconnected mocks base method.
-func (m *MockController) OnPeerDisconnected(ctx context.Context, accountID, peerID string) {
+// OnPeerAdded mocks base method.
+func (m *MockController) OnPeerAdded(ctx context.Context, accountID, peerID string) error {
 	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "OnPeerDisconnected", ctx, accountID, peerID)
-}
-
-// OnPeerDisconnected indicates an expected call of OnPeerDisconnected.
-func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID any) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDisconnected", reflect.TypeOf((*MockController)(nil).OnPeerDisconnected), ctx, accountID, peerID)
-}
-
-// OnPeersAdded mocks base method.
-func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs []string) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "OnPeersAdded", ctx, accountID, peerIDs)
+	ret := m.ctrl.Call(m, "OnPeerAdded", ctx, accountID, peerID)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-// OnPeersAdded indicates an expected call of OnPeersAdded.
-func (mr *MockControllerMockRecorder) OnPeersAdded(ctx, accountID, peerIDs any) *gomock.Call {
+// OnPeerAdded indicates an expected call of OnPeerAdded.
+func (mr *MockControllerMockRecorder) OnPeerAdded(ctx, accountID, peerID any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeersAdded", reflect.TypeOf((*MockController)(nil).OnPeersAdded), ctx, accountID, peerIDs)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerAdded", reflect.TypeOf((*MockController)(nil).OnPeerAdded), ctx, accountID, peerID)
 }
 
-// OnPeersDeleted mocks base method.
-func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string) error {
+// OnPeerDeleted mocks base method.
+func (m *MockController) OnPeerDeleted(ctx context.Context, accountID, peerID string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "OnPeersDeleted", ctx, accountID, peerIDs)
+	ret := m.ctrl.Call(m, "OnPeerDeleted", ctx, accountID, peerID)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-// OnPeersDeleted indicates an expected call of OnPeersDeleted.
-func (mr *MockControllerMockRecorder) OnPeersDeleted(ctx, accountID, peerIDs any) *gomock.Call {
+// OnPeerDeleted indicates an expected call of OnPeerDeleted.
+func (mr *MockControllerMockRecorder) OnPeerDeleted(ctx, accountID, peerID any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeersDeleted", reflect.TypeOf((*MockController)(nil).OnPeersDeleted), ctx, accountID, peerIDs)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDeleted", reflect.TypeOf((*MockController)(nil).OnPeerDeleted), ctx, accountID, peerID)
 }
 
-// OnPeersUpdated mocks base method.
-func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs []string) error {
+// OnPeerUpdated mocks base method.
+func (m *MockController) OnPeerUpdated(accountId string, peer *peer.Peer) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "OnPeersUpdated", ctx, accountId, peerIDs)
-	ret0, _ := ret[0].(error)
-	return ret0
+	m.ctrl.Call(m, "OnPeerUpdated", accountId, peer)
 }
 
-// OnPeersUpdated indicates an expected call of OnPeersUpdated.
-func (mr *MockControllerMockRecorder) OnPeersUpdated(ctx, accountId, peerIDs any) *gomock.Call {
+// OnPeerUpdated indicates an expected call of OnPeerUpdated.
+func (mr *MockControllerMockRecorder) OnPeerUpdated(accountId, peer any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeersUpdated", reflect.TypeOf((*MockController)(nil).OnPeersUpdated), ctx, accountId, peerIDs)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerUpdated", reflect.TypeOf((*MockController)(nil).OnPeerUpdated), accountId, peer)
 }
 
 // StartWarmup mocks base method.

management/internals/modules/peers/manager.go

@@ -1,162 +0,0 @@
-package peers
-
-//go:generate go run github.com/golang/mock/mockgen -package peers -destination=manager_mock.go -source=./manager.go -build_flags=-mod=mod
-
-import (
-	"context"
-	"fmt"
-	"time"
-
-	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
-	"github.com/netbirdio/netbird/management/server/account"
-	"github.com/netbirdio/netbird/management/server/activity"
-	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
-	"github.com/netbirdio/netbird/management/server/peer"
-	"github.com/netbirdio/netbird/management/server/permissions"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
-	"github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/shared/management/status"
-)
-
-type Manager interface {
-	GetPeer(ctx context.Context, accountID, userID, peerID string) (*peer.Peer, error)
-	GetPeerAccountID(ctx context.Context, peerID string) (string, error)
-	GetAllPeers(ctx context.Context, accountID, userID string) ([]*peer.Peer, error)
-	GetPeersByGroupIDs(ctx context.Context, accountID string, groupsIDs []string) ([]*peer.Peer, error)
-	DeletePeers(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error
-	SetNetworkMapController(networkMapController network_map.Controller)
-	SetIntegratedPeerValidator(integratedPeerValidator integrated_validator.IntegratedValidator)
-	SetAccountManager(accountManager account.Manager)
-}
-
-type managerImpl struct {
-	store                   store.Store
-	permissionsManager      permissions.Manager
-	integratedPeerValidator integrated_validator.IntegratedValidator
-	accountManager          account.Manager
-
-	networkMapController network_map.Controller
-}
-
-func NewManager(store store.Store, permissionsManager permissions.Manager) Manager {
-	return &managerImpl{
-		store:              store,
-		permissionsManager: permissionsManager,
-	}
-}
-
-func (m *managerImpl) SetNetworkMapController(networkMapController network_map.Controller) {
-	m.networkMapController = networkMapController
-}
-
-func (m *managerImpl) SetIntegratedPeerValidator(integratedPeerValidator integrated_validator.IntegratedValidator) {
-	m.integratedPeerValidator = integratedPeerValidator
-}
-
-func (m *managerImpl) SetAccountManager(accountManager account.Manager) {
-	m.accountManager = accountManager
-}
-
-func (m *managerImpl) GetPeer(ctx context.Context, accountID, userID, peerID string) (*peer.Peer, error) {
-	allowed, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Read)
-	if err != nil {
-		return nil, fmt.Errorf("failed to validate user permissions: %w", err)
-	}
-
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
-	}
-
-	return m.store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
-}
-
-func (m *managerImpl) GetAllPeers(ctx context.Context, accountID, userID string) ([]*peer.Peer, error) {
-	allowed, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Read)
-	if err != nil {
-		return nil, fmt.Errorf("failed to validate user permissions: %w", err)
-	}
-
-	if !allowed {
-		return m.store.GetUserPeers(ctx, store.LockingStrengthNone, accountID, userID)
-	}
-
-	return m.store.GetAccountPeers(ctx, store.LockingStrengthNone, accountID, "", "")
-}
-
-func (m *managerImpl) GetPeerAccountID(ctx context.Context, peerID string) (string, error) {
-	return m.store.GetAccountIDByPeerID(ctx, store.LockingStrengthNone, peerID)
-}
-
-func (m *managerImpl) GetPeersByGroupIDs(ctx context.Context, accountID string, groupsIDs []string) ([]*peer.Peer, error) {
-	return m.store.GetPeersByGroupIDs(ctx, accountID, groupsIDs)
-}
-
-func (m *managerImpl) DeletePeers(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-	settings, err := m.store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
-	if err != nil {
-		return err
-	}
-	dnsDomain := m.networkMapController.GetDNSDomain(settings)
-
-	for _, peerID := range peerIDs {
-		var eventsToStore []func()
-		err := m.store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
-			peer, err := transaction.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
-			if err != nil {
-				return err
-			}
-
-			if checkConnected && (peer.Status.Connected || peer.Status.LastSeen.After(time.Now().Add(-(ephemeral.EphemeralLifeTime - 10*time.Second)))) {
-				return nil
-			}
-
-			if err := transaction.RemovePeerFromAllGroups(ctx, peerID); err != nil {
-				return fmt.Errorf("failed to remove peer %s from groups", peerID)
-			}
-
-			if err := m.integratedPeerValidator.PeerDeleted(ctx, accountID, peerID, settings.Extra); err != nil {
-				return err
-			}
-
-			peerPolicyRules, err := transaction.GetPolicyRulesByResourceID(ctx, store.LockingStrengthNone, accountID, peerID)
-			if err != nil {
-				return err
-			}
-			for _, rule := range peerPolicyRules {
-				policy, err := transaction.GetPolicyByID(ctx, store.LockingStrengthNone, accountID, rule.PolicyID)
-				if err != nil {
-					return err
-				}
-
-				err = transaction.DeletePolicy(ctx, accountID, rule.PolicyID)
-				if err != nil {
-					return err
-				}
-
-				eventsToStore = append(eventsToStore, func() {
-					m.accountManager.StoreEvent(ctx, userID, peer.ID, accountID, activity.PolicyRemoved, policy.EventMeta())
-				})
-			}
-
-			if err = transaction.DeletePeer(ctx, accountID, peerID); err != nil {
-				return err
-			}
-
-			eventsToStore = append(eventsToStore, func() {
-				m.accountManager.StoreEvent(ctx, userID, peer.ID, accountID, activity.PeerRemovedByUser, peer.EventMeta(dnsDomain))
-			})
-
-			return nil
-		})
-		if err != nil {
-			return err
-		}
-		for _, event := range eventsToStore {
-			event()
-		}
-	}
-
-	return nil
-}

management/internals/server/boot.go

@@ -57,7 +57,7 @@ func (s *BaseServer) Metrics() telemetry.AppMetrics {
 
 func (s *BaseServer) Store() store.Store {
 	return Create(s, func() store.Store {
-		store, err := store.NewStore(context.Background(), s.Config.StoreConfig.Engine, s.Config.Datadir, s.Metrics(), false)
+		store, err := store.NewStore(context.Background(), s.config.StoreConfig.Engine, s.config.Datadir, s.Metrics(), false)
 		if err != nil {
 			log.Fatalf("failed to create store: %v", err)
 		}
@@ -73,17 +73,17 @@ func (s *BaseServer) EventStore() activity.Store {
 			log.Fatalf("failed to initialize integration metrics: %v", err)
 		}
 
-		eventStore, key, err := integrations.InitEventStore(context.Background(), s.Config.Datadir, s.Config.DataStoreEncryptionKey, integrationMetrics)
+		eventStore, key, err := integrations.InitEventStore(context.Background(), s.config.Datadir, s.config.DataStoreEncryptionKey, integrationMetrics)
 		if err != nil {
 			log.Fatalf("failed to initialize event store: %v", err)
 		}
 
-		if s.Config.DataStoreEncryptionKey != key {
-			log.WithContext(context.Background()).Infof("update Config with activity store key")
-			s.Config.DataStoreEncryptionKey = key
-			err := updateMgmtConfig(context.Background(), nbconfig.MgmtConfigPath, s.Config)
+		if s.config.DataStoreEncryptionKey != key {
+			log.WithContext(context.Background()).Infof("update config with activity store key")
+			s.config.DataStoreEncryptionKey = key
+			err := updateMgmtConfig(context.Background(), nbconfig.MgmtConfigPath, s.config)
 			if err != nil {
-				log.Fatalf("failed to update Config with activity store: %v", err)
+				log.Fatalf("failed to update config with activity store: %v", err)
 			}
 		}
 
@@ -103,14 +103,14 @@ func (s *BaseServer) APIHandler() http.Handler {
 
 func (s *BaseServer) GRPCServer() *grpc.Server {
 	return Create(s, func() *grpc.Server {
-		trustedPeers := s.Config.ReverseProxy.TrustedPeers
+		trustedPeers := s.config.ReverseProxy.TrustedPeers
 		defaultTrustedPeers := []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")}
 		if len(trustedPeers) == 0 || slices.Equal[[]netip.Prefix](trustedPeers, defaultTrustedPeers) {
 			log.WithContext(context.Background()).Warn("TrustedPeers are configured to default value '0.0.0.0/0', '::/0'. This allows connection IP spoofing.")
 			trustedPeers = defaultTrustedPeers
 		}
-		trustedHTTPProxies := s.Config.ReverseProxy.TrustedHTTPProxies
-		trustedProxiesCount := s.Config.ReverseProxy.TrustedHTTPProxiesCount
+		trustedHTTPProxies := s.config.ReverseProxy.TrustedHTTPProxies
+		trustedProxiesCount := s.config.ReverseProxy.TrustedHTTPProxiesCount
 		if len(trustedHTTPProxies) > 0 && trustedProxiesCount > 0 {
 			log.WithContext(context.Background()).Warn("TrustedHTTPProxies and TrustedHTTPProxiesCount both are configured. " +
 				"This is not recommended way to extract X-Forwarded-For. Consider using one of these options.")
@@ -128,15 +128,15 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 			grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor),
 		}
 
-		if s.Config.HttpConfig.LetsEncryptDomain != "" {
-			certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
+		if s.config.HttpConfig.LetsEncryptDomain != "" {
+			certManager, err := encryption.CreateCertManager(s.config.Datadir, s.config.HttpConfig.LetsEncryptDomain)
 			if err != nil {
 				log.Fatalf("failed to create certificate manager: %v", err)
 			}
 			transportCredentials := credentials.NewTLS(certManager.TLSConfig())
 			gRPCOpts = append(gRPCOpts, grpc.Creds(transportCredentials))
-		} else if s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "" {
-			tlsConfig, err := loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
+		} else if s.config.HttpConfig.CertFile != "" && s.config.HttpConfig.CertKey != "" {
+			tlsConfig, err := loadTLSConfig(s.config.HttpConfig.CertFile, s.config.HttpConfig.CertKey)
 			if err != nil {
 				log.Fatalf("cannot load TLS credentials: %v", err)
 			}
@@ -145,7 +145,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		}
 
 		gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
-		srv, err := nbgrpc.NewServer(s.Config, s.AccountManager(), s.SettingsManager(), s.SecretsManager(), s.Metrics(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController())
+		srv, err := nbgrpc.NewServer(s.config, s.AccountManager(), s.SettingsManager(), s.PeersUpdateManager(), s.SecretsManager(), s.Metrics(), s.EphemeralManager(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController())
 		if err != nil {
 			log.Fatalf("failed to create management server: %v", err)
 		}

management/internals/server/controllers.go

@@ -9,17 +9,17 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	nmapcontroller "github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/auth"
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 )
 
 func (s *BaseServer) PeersUpdateManager() network_map.PeersUpdateManager {
-	return Create(s, func() network_map.PeersUpdateManager {
+	return Create(s, func() *update_channel.PeersUpdateManager {
 		return update_channel.NewPeersUpdateManager(s.Metrics())
 	})
 }
@@ -44,37 +44,33 @@ func (s *BaseServer) ProxyController() port_forwarding.Controller {
 	})
 }
 
-func (s *BaseServer) SecretsManager() grpc.SecretsManager {
-	return Create(s, func() grpc.SecretsManager {
-		secretsManager, err := grpc.NewTimeBasedAuthSecretsManager(s.PeersUpdateManager(), s.Config.TURNConfig, s.Config.Relay, s.SettingsManager(), s.GroupsManager())
-		if err != nil {
-			log.Fatalf("failed to create secrets manager: %v", err)
-		}
-		return secretsManager
+func (s *BaseServer) SecretsManager() *grpc.TimeBasedAuthSecretsManager {
+	return Create(s, func() *grpc.TimeBasedAuthSecretsManager {
+		return grpc.NewTimeBasedAuthSecretsManager(s.PeersUpdateManager(), s.config.TURNConfig, s.config.Relay, s.SettingsManager(), s.GroupsManager())
 	})
 }
 
 func (s *BaseServer) AuthManager() auth.Manager {
 	return Create(s, func() auth.Manager {
 		return auth.NewManager(s.Store(),
-			s.Config.HttpConfig.AuthIssuer,
-			s.Config.HttpConfig.AuthAudience,
-			s.Config.HttpConfig.AuthKeysLocation,
-			s.Config.HttpConfig.AuthUserIDClaim,
-			s.Config.GetAuthAudiences(),
-			s.Config.HttpConfig.IdpSignKeyRefreshEnabled)
+			s.config.HttpConfig.AuthIssuer,
+			s.config.HttpConfig.AuthAudience,
+			s.config.HttpConfig.AuthKeysLocation,
+			s.config.HttpConfig.AuthUserIDClaim,
+			s.config.GetAuthAudiences(),
+			s.config.HttpConfig.IdpSignKeyRefreshEnabled)
 	})
 }
 
 func (s *BaseServer) EphemeralManager() ephemeral.Manager {
 	return Create(s, func() ephemeral.Manager {
-		return manager.NewEphemeralManager(s.Store(), s.PeersManager())
+		return manager.NewEphemeralManager(s.Store(), s.AccountManager())
 	})
 }
 
 func (s *BaseServer) NetworkMapController() network_map.Controller {
-	return Create(s, func() network_map.Controller {
-		return nmapcontroller.NewController(context.Background(), s.Store(), s.Metrics(), s.PeersUpdateManager(), s.AccountRequestBuffer(), s.IntegratedValidator(), s.SettingsManager(), s.DNSDomain(), s.ProxyController(), s.EphemeralManager(), s.Config)
+	return Create(s, func() *nmapcontroller.Controller {
+		return nmapcontroller.NewController(context.Background(), s.Store(), s.Metrics(), s.PeersUpdateManager(), s.AccountRequestBuffer(), s.IntegratedValidator(), s.SettingsManager(), s.dnsDomain, s.ProxyController(), s.config)
 	})
 }
 
@@ -83,7 +79,3 @@ func (s *BaseServer) AccountRequestBuffer() *server.AccountRequestBuffer {
 		return server.NewAccountRequestBuffer(context.Background(), s.Store())
 	})
 }
-
-func (s *BaseServer) DNSDomain() string {
-	return s.dnsDomain
-}

management/internals/server/modules.go

@@ -2,12 +2,10 @@ package server
 
 import (
 	"context"
-	"os"
 
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/management-integrations/integrations"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/geolocation"
@@ -16,29 +14,20 @@ import (
 	"github.com/netbirdio/netbird/management/server/networks"
 	"github.com/netbirdio/netbird/management/server/networks/resources"
 	"github.com/netbirdio/netbird/management/server/networks/routers"
-
+	"github.com/netbirdio/netbird/management/server/peers"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/users"
 )
 
-const (
-	geolocationDisabledKey = "NB_DISABLE_GEOLOCATION"
-)
-
 func (s *BaseServer) GeoLocationManager() geolocation.Geolocation {
-	if os.Getenv(geolocationDisabledKey) == "true" {
-		log.Info("geolocation service is disabled, skipping initialization")
-		return nil
-	}
-
 	return Create(s, func() geolocation.Geolocation {
-		geo, err := geolocation.NewGeolocation(context.Background(), s.Config.Datadir, !s.disableGeoliteUpdate)
+		geo, err := geolocation.NewGeolocation(context.Background(), s.config.Datadir, !s.disableGeoliteUpdate)
 		if err != nil {
 			log.Fatalf("could not initialize geolocation service: %v", err)
 		}
 
-		log.Infof("geolocation service has been initialized from %s", s.Config.Datadir)
+		log.Infof("geolocation service has been initialized from %s", s.config.Datadir)
 
 		return geo
 	})
@@ -71,22 +60,20 @@ func (s *BaseServer) SettingsManager() settings.Manager {
 
 func (s *BaseServer) PeersManager() peers.Manager {
 	return Create(s, func() peers.Manager {
-		manager := peers.NewManager(s.Store(), s.PermissionsManager())
-		s.AfterInit(func(s *BaseServer) {
-			manager.SetNetworkMapController(s.NetworkMapController())
-			manager.SetIntegratedPeerValidator(s.IntegratedValidator())
-			manager.SetAccountManager(s.AccountManager())
-		})
-		return manager
+		return peers.NewManager(s.Store(), s.PermissionsManager())
 	})
 }
 
 func (s *BaseServer) AccountManager() account.Manager {
 	return Create(s, func() account.Manager {
-		accountManager, err := server.BuildManager(context.Background(), s.Config, s.Store(), s.NetworkMapController(), s.IdpManager(), s.mgmtSingleAccModeDomain, s.EventStore(), s.GeoLocationManager(), s.userDeleteFromIDPEnabled, s.IntegratedValidator(), s.Metrics(), s.ProxyController(), s.SettingsManager(), s.PermissionsManager(), s.Config.DisableDefaultPolicy)
+		accountManager, err := server.BuildManager(context.Background(), s.config, s.Store(), s.NetworkMapController(), s.IdpManager(), s.mgmtSingleAccModeDomain, s.EventStore(), s.GeoLocationManager(), s.userDeleteFromIDPEnabled, s.IntegratedValidator(), s.Metrics(), s.ProxyController(), s.SettingsManager(), s.PermissionsManager(), s.config.DisableDefaultPolicy)
 		if err != nil {
 			log.Fatalf("failed to create account manager: %v", err)
 		}
+
+		s.AfterInit(func(s *BaseServer) {
+			accountManager.SetEphemeralManager(s.EphemeralManager())
+		})
 		return accountManager
 	})
 }
@@ -95,8 +82,8 @@ func (s *BaseServer) IdpManager() idp.Manager {
 	return Create(s, func() idp.Manager {
 		var idpManager idp.Manager
 		var err error
-		if s.Config.IdpManagerConfig != nil {
-			idpManager, err = idp.NewManager(context.Background(), *s.Config.IdpManagerConfig, s.Metrics())
+		if s.config.IdpManagerConfig != nil {
+			idpManager, err = idp.NewManager(context.Background(), *s.config.IdpManagerConfig, s.Metrics())
 			if err != nil {
 				log.Fatalf("failed to create IDP manager: %v", err)
 			}

management/internals/server/server.go

@@ -41,10 +41,10 @@ type Server interface {
 }
 
 // Server holds the HTTP BaseServer instance.
-// Add any additional fields you need, such as database connections, Config, etc.
+// Add any additional fields you need, such as database connections, config, etc.
 type BaseServer struct {
-	// Config holds the server configuration
-	Config *nbconfig.Config
+	// config holds the server configuration
+	config *nbconfig.Config
 	// container of dependencies, each dependency is identified by a unique string.
 	container map[string]any
 	// AfterInit is a function that will be called after the server is initialized
@@ -70,7 +70,7 @@ type BaseServer struct {
 // NewServer initializes and configures a new Server instance
 func NewServer(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool) *BaseServer {
 	return &BaseServer{
-		Config:                   config,
+		config:                   config,
 		container:                make(map[string]any),
 		dnsDomain:                dnsDomain,
 		mgmtSingleAccModeDomain:  mgmtSingleAccModeDomain,
@@ -103,14 +103,14 @@ func (s *BaseServer) Start(ctx context.Context) error {
 
 	var tlsConfig *tls.Config
 	tlsEnabled := false
-	if s.Config.HttpConfig.LetsEncryptDomain != "" {
-		s.certManager, err = encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
+	if s.config.HttpConfig.LetsEncryptDomain != "" {
+		s.certManager, err = encryption.CreateCertManager(s.config.Datadir, s.config.HttpConfig.LetsEncryptDomain)
 		if err != nil {
 			return fmt.Errorf("failed creating LetsEncrypt cert manager: %v", err)
 		}
 		tlsEnabled = true
-	} else if s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "" {
-		tlsConfig, err = loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
+	} else if s.config.HttpConfig.CertFile != "" && s.config.HttpConfig.CertKey != "" {
+		tlsConfig, err = loadTLSConfig(s.config.HttpConfig.CertFile, s.config.HttpConfig.CertKey)
 		if err != nil {
 			log.WithContext(srvCtx).Errorf("cannot load TLS credentials: %v", err)
 			return err
@@ -126,8 +126,8 @@ func (s *BaseServer) Start(ctx context.Context) error {
 
 	if !s.disableMetrics {
 		idpManager := "disabled"
-		if s.Config.IdpManagerConfig != nil && s.Config.IdpManagerConfig.ManagerType != "" {
-			idpManager = s.Config.IdpManagerConfig.ManagerType
+		if s.config.IdpManagerConfig != nil && s.config.IdpManagerConfig.ManagerType != "" {
+			idpManager = s.config.IdpManagerConfig.ManagerType
 		}
 		metricsWorker := metrics.NewWorker(srvCtx, installationID, s.Store(), s.PeersUpdateManager(), idpManager)
 		go metricsWorker.Run(srvCtx)

management/internals/shared/grpc/server.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral"
 
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -54,12 +55,15 @@ const (
 type Server struct {
 	accountManager  account.Manager
 	settingsManager settings.Manager
+	wgKey           wgtypes.Key
 	proto.UnimplementedManagementServiceServer
-	config         *nbconfig.Config
-	secretsManager SecretsManager
-	appMetrics     telemetry.AppMetrics
-	peerLocks      sync.Map
-	authManager    auth.Manager
+	peersUpdateManager network_map.PeersUpdateManager
+	config             *nbconfig.Config
+	secretsManager     SecretsManager
+	appMetrics         telemetry.AppMetrics
+	ephemeralManager   ephemeral.Manager
+	peerLocks          sync.Map
+	authManager        auth.Manager
 
 	logBlockedPeers          bool
 	blockPeersWithSameConfig bool
@@ -78,16 +82,23 @@ func NewServer(
 	config *nbconfig.Config,
 	accountManager account.Manager,
 	settingsManager settings.Manager,
+	peersUpdateManager network_map.PeersUpdateManager,
 	secretsManager SecretsManager,
 	appMetrics telemetry.AppMetrics,
+	ephemeralManager ephemeral.Manager,
 	authManager auth.Manager,
 	integratedPeerValidator integrated_validator.IntegratedValidator,
 	networkMapController network_map.Controller,
 ) (*Server, error) {
+	key, err := wgtypes.GeneratePrivateKey()
+	if err != nil {
+		return nil, err
+	}
+
 	if appMetrics != nil {
 		// update gauge based on number of connected peers which is equal to open gRPC streams
-		err := appMetrics.GRPCMetrics().RegisterConnectedStreams(func() int64 {
-			return int64(networkMapController.CountStreams())
+		err = appMetrics.GRPCMetrics().RegisterConnectedStreams(func() int64 {
+			return int64(peersUpdateManager.CountStreams())
 		})
 		if err != nil {
 			return nil, err
@@ -109,12 +120,16 @@ func NewServer(
 	}
 
 	return &Server{
+		wgKey: key,
+		// peerKey -> event channel
+		peersUpdateManager:       peersUpdateManager,
 		accountManager:           accountManager,
 		settingsManager:          settingsManager,
 		config:                   config,
 		secretsManager:           secretsManager,
 		authManager:              authManager,
 		appMetrics:               appMetrics,
+		ephemeralManager:         ephemeralManager,
 		logBlockedPeers:          logBlockedPeers,
 		blockPeersWithSameConfig: blockPeersWithSameConfig,
 		integratedPeerValidator:  integratedPeerValidator,
@@ -148,14 +163,8 @@ func (s *Server) GetServerKey(ctx context.Context, req *proto.Empty) (*proto.Ser
 	nanos := int32(now.Nanosecond())
 	expiresAt := &timestamp.Timestamp{Seconds: secs, Nanos: nanos}
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to get wireguard key: %v", err)
-		return nil, errors.New("failed to get wireguard key")
-	}
-
 	return &proto.ServerKeyResponse{
-		Key:       key.PublicKey().String(),
+		Key:       s.wgKey.PublicKey().String(),
 		ExpiresAt: expiresAt,
 	}, nil
 }
@@ -260,13 +269,9 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 		return err
 	}
 
-	updates, err := s.networkMapController.OnPeerConnected(ctx, accountID, peer.ID)
-	if err != nil {
-		log.WithContext(ctx).Debugf("error while notify peer connected for %s: %v", peerKey.String(), err)
-		s.syncSem.Add(-1)
-		s.cancelPeerRoutines(ctx, accountID, peer)
-		return err
-	}
+	updates := s.peersUpdateManager.CreateChannel(ctx, peer.ID)
+
+	s.ephemeralManager.OnPeerConnected(ctx, peer)
 
 	s.secretsManager.SetupRefresh(ctx, accountID, peer.ID)
 
@@ -318,19 +323,13 @@ func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wg
 // sendUpdate encrypts the update message using the peer key and the server's wireguard key,
 // then sends the encrypted message to the connected peer via the sync server.
 func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer) error {
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update)
+	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, update.Update)
 	if err != nil {
 		s.cancelPeerRoutines(ctx, accountID, peer)
 		return status.Errorf(codes.Internal, "failed processing update message")
 	}
 	err = srv.SendMsg(&proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
+		WgPubKey: s.wgKey.PublicKey().String(),
 		Body:     encryptedResp,
 	})
 	if err != nil {
@@ -349,8 +348,9 @@ func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
 	}
-	s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID)
+	s.peersUpdateManager.CloseChannel(ctx, peer.ID)
 	s.secretsManager.CancelRefresh(peer.ID)
+	s.ephemeralManager.OnPeerDisconnected(ctx, peer)
 
 	log.WithContext(ctx).Tracef("peer %s has been disconnected", peer.Key)
 }
@@ -504,12 +504,7 @@ func (s *Server) parseRequest(ctx context.Context, req *proto.EncryptedMessage,
 		return wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", req.WgPubKey)
 	}
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		return wgtypes.Key{}, status.Errorf(codes.Internal, "failed processing request")
-	}
-
-	err = encryption.DecryptMessage(peerKey, key, req.Body, parsed)
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, parsed)
 	if err != nil {
 		return wgtypes.Key{}, status.Errorf(codes.InvalidArgument, "invalid request message")
 	}
@@ -606,6 +601,12 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 
 	log.WithContext(ctx).Debugf("Login: LoginPeer since start %v", time.Since(reqStart))
 
+	// if the login request contains setup key then it is a registration request
+	if loginReq.GetSetupKey() != "" {
+		s.ephemeralManager.OnPeerDisconnected(ctx, peer)
+		log.WithContext(ctx).Debugf("Login: OnPeerDisconnected since start %v", time.Since(reqStart))
+	}
+
 	loginResp, err := s.prepareLoginResponse(ctx, peer, netMap, postureChecks)
 	if err != nil {
 		log.WithContext(ctx).Warnf("failed preparing login response for peer %s: %s", peerKey, err)
@@ -614,20 +615,14 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 
 	log.WithContext(ctx).Debugf("Login: prepareLoginResponse since start %v", time.Since(reqStart))
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		log.WithContext(ctx).Warnf("failed getting server's WireGuard private key: %s", err)
-		return nil, status.Errorf(codes.Internal, "failed logging in peer")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, loginResp)
+	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, loginResp)
 	if err != nil {
 		log.WithContext(ctx).Warnf("failed encrypting peer %s message", peer.ID)
 		return nil, status.Errorf(codes.Internal, "failed logging in peer")
 	}
 
 	return &proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
+		WgPubKey: s.wgKey.PublicKey().String(),
 		Body:     encryptedResp,
 	}, nil
 }
@@ -720,19 +715,14 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
 
 	plainResp := ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, peer, turnToken, relayToken, networkMap, s.networkMapController.GetDNSDomain(settings), postureChecks, nil, settings, settings.Extra, peerGroups, dnsFwdPort)
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		return status.Errorf(codes.Internal, "failed getting server key")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, plainResp)
+	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, plainResp)
 	if err != nil {
 		return status.Errorf(codes.Internal, "error handling request")
 	}
 
 	sendStart := time.Now()
 	err = srv.Send(&proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
+		WgPubKey: s.wgKey.PublicKey().String(),
 		Body:     encryptedResp,
 	})
 	log.WithContext(ctx).Debugf("sendInitialSync: sending response took %s", time.Since(sendStart))
@@ -762,12 +752,7 @@ func (s *Server) GetDeviceAuthorizationFlow(ctx context.Context, req *proto.Encr
 		return nil, status.Error(codes.InvalidArgument, errMSG)
 	}
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to get server key")
-	}
-
-	err = encryption.DecryptMessage(peerKey, key, req.Body, &proto.DeviceAuthorizationFlowRequest{})
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, &proto.DeviceAuthorizationFlowRequest{})
 	if err != nil {
 		errMSG := fmt.Sprintf("error while decrypting peer's message with Wireguard public key %s.", req.WgPubKey)
 		log.WithContext(ctx).Warn(errMSG)
@@ -797,13 +782,13 @@ func (s *Server) GetDeviceAuthorizationFlow(ctx context.Context, req *proto.Encr
 		},
 	}
 
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, flowInfoResp)
+	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, flowInfoResp)
 	if err != nil {
 		return nil, status.Error(codes.Internal, "failed to encrypt no device authorization flow information")
 	}
 
 	return &proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
+		WgPubKey: s.wgKey.PublicKey().String(),
 		Body:     encryptedResp,
 	}, nil
 }
@@ -825,12 +810,7 @@ func (s *Server) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.Encryp
 		return nil, status.Error(codes.InvalidArgument, errMSG)
 	}
 
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to get server key")
-	}
-
-	err = encryption.DecryptMessage(peerKey, key, req.Body, &proto.PKCEAuthorizationFlowRequest{})
+	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, &proto.PKCEAuthorizationFlowRequest{})
 	if err != nil {
 		errMSG := fmt.Sprintf("error while decrypting peer's message with Wireguard public key %s.", req.WgPubKey)
 		log.WithContext(ctx).Warn(errMSG)
@@ -858,13 +838,13 @@ func (s *Server) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.Encryp
 
 	flowInfoResp := s.integratedPeerValidator.ValidateFlowResponse(ctx, peerKey.String(), initInfoFlow)
 
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, flowInfoResp)
+	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, flowInfoResp)
 	if err != nil {
 		return nil, status.Error(codes.Internal, "failed to encrypt no pkce authorization flow information")
 	}
 
 	return &proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
+		WgPubKey: s.wgKey.PublicKey().String(),
 		Body:     encryptedResp,
 	}, nil
 }

management/internals/shared/grpc/server_test.go

@@ -73,17 +73,15 @@ func TestServer_GetDeviceAuthorizationFlow(t *testing.T) {
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
 			mgmtServer := &Server{
-				secretsManager: &TimeBasedAuthSecretsManager{wgKey: testingServerKey},
+				wgKey: testingServerKey,
 				config: &config.Config{
 					DeviceAuthorizationFlow: testCase.inputFlow,
 				},
 			}
 
 			message := &mgmtProto.DeviceAuthorizationFlowRequest{}
-			key, err := mgmtServer.secretsManager.GetWGKey()
-			require.NoError(t, err, "should be able to get server key")
 
-			encryptedMSG, err := encryption.EncryptMessage(testingClientKey.PublicKey(), key, message)
+			encryptedMSG, err := encryption.EncryptMessage(testingClientKey.PublicKey(), mgmtServer.wgKey, message)
 			require.NoError(t, err, "should be able to encrypt message")
 
 			resp, err := mgmtServer.GetDeviceAuthorizationFlow(
@@ -97,7 +95,7 @@ func TestServer_GetDeviceAuthorizationFlow(t *testing.T) {
 			if testCase.expectedComparisonFunc != nil {
 				flowInfoResp := &mgmtProto.DeviceAuthorizationFlow{}
 
-				err = encryption.DecryptMessage(key.PublicKey(), testingClientKey, resp.Body, flowInfoResp)
+				err = encryption.DecryptMessage(mgmtServer.wgKey.PublicKey(), testingClientKey, resp.Body, flowInfoResp)
 				require.NoError(t, err, "should be able to decrypt")
 
 				testCase.expectedComparisonFunc(t, testCase.expectedFlow.Provider, flowInfoResp.Provider, testCase.expectedComparisonMSG)

management/internals/shared/grpc/token_mgr.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
@@ -30,7 +29,6 @@ type SecretsManager interface {
 	GenerateRelayToken() (*Token, error)
 	SetupRefresh(ctx context.Context, accountID, peerKey string)
 	CancelRefresh(peerKey string)
-	GetWGKey() (wgtypes.Key, error)
 }
 
 // TimeBasedAuthSecretsManager generates credentials with TTL and using pre-shared secret known to TURN server
@@ -45,17 +43,11 @@ type TimeBasedAuthSecretsManager struct {
 	groupsManager   groups.Manager
 	turnCancelMap   map[string]chan struct{}
 	relayCancelMap  map[string]chan struct{}
-	wgKey           wgtypes.Key
 }
 
 type Token auth.Token
 
-func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager, turnCfg *nbconfig.TURNConfig, relayCfg *nbconfig.Relay, settingsManager settings.Manager, groupsManager groups.Manager) (*TimeBasedAuthSecretsManager, error) {
-	key, err := wgtypes.GeneratePrivateKey()
-	if err != nil {
-		return nil, err
-	}
-
+func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager, turnCfg *nbconfig.TURNConfig, relayCfg *nbconfig.Relay, settingsManager settings.Manager, groupsManager groups.Manager) *TimeBasedAuthSecretsManager {
 	mgr := &TimeBasedAuthSecretsManager{
 		updateManager:   updateManager,
 		turnCfg:         turnCfg,
@@ -64,7 +56,6 @@ func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager
 		relayCancelMap:  make(map[string]chan struct{}),
 		settingsManager: settingsManager,
 		groupsManager:   groupsManager,
-		wgKey:           key,
 	}
 
 	if turnCfg != nil {
@@ -90,12 +81,7 @@ func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager
 		}
 	}
 
-	return mgr, nil
-}
-
-// GetWGKey returns WireGuard private key used to generate peer keys
-func (m *TimeBasedAuthSecretsManager) GetWGKey() (wgtypes.Key, error) {
-	return m.wgKey, nil
+	return mgr
 }
 
 // GenerateTurnToken generates new time-based secret credentials for TURN

management/internals/shared/grpc/token_mgr_test.go

@@ -46,13 +46,12 @@ func TestTimeBasedAuthSecretsManager_GenerateCredentials(t *testing.T) {
 	settingsMockManager := settings.NewMockManager(ctrl)
 	groupsManager := groups.NewManagerMock()
 
-	tested, err := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
+	tested := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
 		CredentialsTTL:       ttl,
 		Secret:               secret,
 		Turns:                []*config.Host{TurnTestHost},
 		TimeBasedCredentials: true,
 	}, rc, settingsMockManager, groupsManager)
-	require.NoError(t, err)
 
 	turnCredentials, err := tested.GenerateTurnToken()
 	require.NoError(t, err)
@@ -99,13 +98,12 @@ func TestTimeBasedAuthSecretsManager_SetupRefresh(t *testing.T) {
 	settingsMockManager.EXPECT().GetExtraSettings(gomock.Any(), "someAccountID").Return(&types.ExtraSettings{}, nil).AnyTimes()
 	groupsManager := groups.NewManagerMock()
 
-	tested, err := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
+	tested := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
 		CredentialsTTL:       ttl,
 		Secret:               secret,
 		Turns:                []*config.Host{TurnTestHost},
 		TimeBasedCredentials: true,
 	}, rc, settingsMockManager, groupsManager)
-	require.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -203,13 +201,12 @@ func TestTimeBasedAuthSecretsManager_CancelRefresh(t *testing.T) {
 	settingsMockManager := settings.NewMockManager(ctrl)
 	groupsManager := groups.NewManagerMock()
 
-	tested, err := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
+	tested := NewTimeBasedAuthSecretsManager(peersManager, &config.TURNConfig{
 		CredentialsTTL:       ttl,
 		Secret:               secret,
 		Turns:                []*config.Host{TurnTestHost},
 		TimeBasedCredentials: true,
 	}, rc, settingsMockManager, groupsManager)
-	require.NoError(t, err)
 
 	tested.SetupRefresh(context.Background(), "someAccountID", peer)
 	if _, ok := tested.turnCancelMap[peer]; !ok {

management/server/account.go

@@ -37,6 +37,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
 	"github.com/netbirdio/netbird/management/server/permissions/operations"
@@ -76,6 +77,7 @@ type DefaultAccountManager struct {
 	ctx                  context.Context
 	eventStore           activity.Store
 	geo                  geolocation.Geolocation
+	ephemeralManager     ephemeral.Manager
 
 	requestBuffer *AccountRequestBuffer
 
@@ -236,7 +238,7 @@ func BuildManager(
 		log.WithContext(ctx).Infof("single account mode disabled, accounts number %d", accountsCounter)
 	}
 
-	cacheStore, err := nbcache.NewStore(ctx, nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval, nbcache.DefaultIDPCacheOpenConn)
+	cacheStore, err := nbcache.NewStore(ctx, nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
 	if err != nil {
 		return nil, fmt.Errorf("getting cache store: %s", err)
 	}
@@ -261,6 +263,10 @@ func BuildManager(
 	return am, nil
 }
 
+func (am *DefaultAccountManager) SetEphemeralManager(em ephemeral.Manager) {
+	am.ephemeralManager = em
+}
+
 func (am *DefaultAccountManager) GetExternalCacheManager() account.ExternalCacheManager {
 	return am.externalCacheManager
 }
@@ -319,9 +325,6 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
 			}
 		}
 
-		newSettings.Extra.IntegratedValidatorGroups = oldSettings.Extra.IntegratedValidatorGroups
-		newSettings.Extra.IntegratedValidator = oldSettings.Extra.IntegratedValidator
-
 		if err = transaction.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
 			return err
 		}
@@ -2070,10 +2073,7 @@ func (am *DefaultAccountManager) UpdatePeerIP(ctx context.Context, accountID, us
 		if err != nil {
 			return err
 		}
-		err = am.networkMapController.OnPeersUpdated(ctx, peer.AccountID, []string{peerID})
-		if err != nil {
-			return fmt.Errorf("notify network map controller of peer update: %w", err)
-		}
+		am.networkMapController.OnPeerUpdated(peer.AccountID, peer)
 	}
 	return nil
 }

management/server/account/manager.go

@@ -13,6 +13,7 @@ import (
 	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	"github.com/netbirdio/netbird/management/server/idp"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -123,4 +124,5 @@ type Manager interface {
 	UpdateToPrimaryAccount(ctx context.Context, accountId string) error
 	GetOwnerInfo(ctx context.Context, accountId string) (*types.UserInfo, error)
 	GetCurrentUserInfo(ctx context.Context, userAuth auth.UserAuth) (*users.UserInfoWithPermissions, error)
+	SetEphemeralManager(em ephemeral.Manager)
 }

management/server/account_test.go

@@ -25,8 +25,6 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
@@ -2961,8 +2959,8 @@ func createManager(t testing.TB) (*DefaultAccountManager, *update_channel.PeersU
 
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager)), &config.Config{})
-	manager, err := BuildManager(ctx, &config.Config{}, store, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), &config.Config{})
+	manager, err := BuildManager(ctx, nil, store, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -3373,7 +3371,7 @@ func TestDefaultAccountManager_IsCacheCold(t *testing.T) {
 
 	t.Run("memory cache", func(t *testing.T) {
 		t.Run("should always return true", func(t *testing.T) {
-			cacheStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+			cacheStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond)
 			require.NoError(t, err)
 
 			cold, err := manager.isCacheCold(context.Background(), cacheStore)
@@ -3388,7 +3386,7 @@ func TestDefaultAccountManager_IsCacheCold(t *testing.T) {
 		t.Cleanup(cleanup)
 		t.Setenv(cache.RedisStoreEnvVar, redisURL)
 
-		cacheStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+		cacheStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond)
 		require.NoError(t, err)
 
 		t.Run("should return true when no account exists", func(t *testing.T) {

management/server/activity/codes.go

@@ -179,7 +179,6 @@ const (
 	PeerIPUpdated              Activity = 88
 	UserApproved               Activity = 89
 	UserRejected               Activity = 90
-	UserCreated                Activity = 91
 
 	AccountDeleted Activity = 99999
 )
@@ -289,7 +288,6 @@ var activityMap = map[Activity]Code{
 	PeerIPUpdated: {"Peer IP updated", "peer.ip.update"},
 	UserApproved:  {"User approved", "user.approve"},
 	UserRejected:  {"User rejected", "user.reject"},
-	UserCreated:   {"User created", "user.create"},
 }
 
 // StringCode returns a string code of the activity

management/server/cache/idp.go

@@ -18,7 +18,6 @@ const (
 	DefaultIDPCacheExpirationMax   = 7 * 24 * time.Hour // 7 days
 	DefaultIDPCacheExpirationMin   = 3 * 24 * time.Hour // 3 days
 	DefaultIDPCacheCleanupInterval = 30 * time.Minute
-	DefaultIDPCacheOpenConn        = 100
 )
 
 // UserDataCache is an interface that wraps the basic Get, Set and Delete methods for idp.UserData objects.

management/server/cache/idp_test.go

@@ -33,7 +33,7 @@ func TestNewIDPCacheManagers(t *testing.T) {
 				t.Cleanup(cleanup)
 				t.Setenv(cache.RedisStoreEnvVar, redisURL)
 			}
-			cacheStore, err := cache.NewStore(context.Background(), cache.DefaultIDPCacheExpirationMax, cache.DefaultIDPCacheCleanupInterval, cache.DefaultIDPCacheOpenConn)
+			cacheStore, err := cache.NewStore(context.Background(), cache.DefaultIDPCacheExpirationMax, cache.DefaultIDPCacheCleanupInterval)
 			if err != nil {
 				t.Fatalf("couldn't create cache store: %s", err)
 			}

management/server/cache/store.go

@@ -3,7 +3,6 @@ package cache
 import (
 	"context"
 	"fmt"
-	"math"
 	"os"
 	"time"
 
@@ -21,27 +20,24 @@ const RedisStoreEnvVar = "NB_IDP_CACHE_REDIS_ADDRESS"
 
 // NewStore creates a new cache store with the given max timeout and cleanup interval. It checks for the environment Variable RedisStoreEnvVar
 // to determine if a redis store should be used. If the environment variable is set, it will attempt to connect to the redis store.
-func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (store.StoreInterface, error) {
+func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration) (store.StoreInterface, error) {
 	redisAddr := os.Getenv(RedisStoreEnvVar)
 	if redisAddr != "" {
-		return getRedisStore(ctx, redisAddr, maxConn)
+		return getRedisStore(ctx, redisAddr)
 	}
 	goc := gocache.New(maxTimeout, cleanupInterval)
 	return gocache_store.NewGoCache(goc), nil
 }
 
-func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (store.StoreInterface, error) {
+func getRedisStore(ctx context.Context, redisEnvAddr string) (store.StoreInterface, error) {
 	options, err := redis.ParseURL(redisEnvAddr)
 	if err != nil {
 		return nil, fmt.Errorf("parsing redis cache url: %s", err)
 	}
 
-	options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns
-	options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns
-	options.MaxActiveConns = maxConn
-	options.ConnMaxIdleTime = 30 * time.Minute
-	options.ConnMaxLifetime = 0
-	options.PoolTimeout = 10 * time.Second
+	options.MaxIdleConns = 6
+	options.MinIdleConns = 3
+	options.MaxActiveConns = 100
 	redisClient := redis.NewClient(options)
 	subCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
 	defer cancel()

management/server/cache/store_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 func TestMemoryStore(t *testing.T) {
-	memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+	memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond)
 	if err != nil {
 		t.Fatalf("couldn't create memory store: %s", err)
 	}
@@ -42,7 +42,7 @@ func TestMemoryStore(t *testing.T) {
 
 func TestRedisStoreConnectionFailure(t *testing.T) {
 	t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379")
-	_, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond, 100)
+	_, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond)
 	if err == nil {
 		t.Fatal("getting redis cache store should return error")
 	}
@@ -65,7 +65,7 @@ func TestRedisStoreConnectionSuccess(t *testing.T) {
 	}
 
 	t.Setenv(cache.RedisStoreEnvVar, redisURL)
-	redisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+	redisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond)
 	if err != nil {
 		t.Fatalf("couldn't create redis store: %s", err)
 	}

management/server/dns_test.go

@@ -12,8 +12,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	"github.com/netbirdio/netbird/management/server/permissions"
@@ -225,7 +223,7 @@ func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.test", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.test", port_forwarding.NewControllerMock(), &config.Config{})
 
 	return BuildManager(context.Background(), nil, store, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 }

management/server/http/handler.go

@@ -21,7 +21,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	"github.com/netbirdio/netbird/management/server/permissions"
 
-	nbpeers "github.com/netbirdio/netbird/management/internals/modules/peers"
 	"github.com/netbirdio/netbird/management/server/auth"
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	nbgroups "github.com/netbirdio/netbird/management/server/groups"
@@ -40,6 +39,7 @@ import (
 	nbnetworks "github.com/netbirdio/netbird/management/server/networks"
 	"github.com/netbirdio/netbird/management/server/networks/resources"
 	"github.com/netbirdio/netbird/management/server/networks/routers"
+	nbpeers "github.com/netbirdio/netbird/management/server/peers"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 )
 
@@ -105,7 +105,6 @@ func NewAPIHandler(
 		accountManager.SyncUserJWTGroups,
 		accountManager.GetUserFromUserAuth,
 		rateLimitingConfig,
-		appMetrics.GetMeter(),
 	)
 
 	corsMiddleware := cors.AllowAll()

management/server/http/handlers/groups/groups_handler.go

@@ -48,29 +48,6 @@ func (h *handler) getAllGroups(w http.ResponseWriter, r *http.Request) {
 	}
 	accountID, userID := userAuth.AccountId, userAuth.UserId
 
-	// Check if filtering by name
-	groupName := r.URL.Query().Get("name")
-	if groupName != "" {
-		// Get single group by name
-		group, err := h.accountManager.GetGroupByName(r.Context(), groupName, accountID)
-		if err != nil {
-			util.WriteError(r.Context(), err, w)
-			return
-		}
-
-		accountPeers, err := h.accountManager.GetPeers(r.Context(), accountID, userID, "", "")
-		if err != nil {
-			util.WriteError(r.Context(), err, w)
-			return
-		}
-
-		// Return as array with single element to maintain API consistency
-		groupsResponse := []*api.Group{toGroupResponse(accountPeers, group)}
-		util.WriteJSONObject(r.Context(), w, groupsResponse)
-		return
-	}
-
-	// Get all groups
 	groups, err := h.accountManager.GetAllGroups(r.Context(), accountID, userID)
 	if err != nil {
 		util.WriteError(r.Context(), err, w)

management/server/http/handlers/groups/groups_handler_test.go

@@ -60,23 +60,12 @@ func initGroupTestData(initGroups ...*types.Group) *handler {
 
 				return group, nil
 			},
-			GetAllGroupsFunc: func(ctx context.Context, accountID, userID string) ([]*types.Group, error) {
-				groups := []*types.Group{
-					{ID: "id-jwt-group", Name: "From JWT", Issued: types.GroupIssuedJWT},
-					{ID: "id-existed", Name: "Existed", Peers: []string{"A", "B"}, Issued: types.GroupIssuedAPI},
-					{ID: "id-all", Name: "All", Issued: types.GroupIssuedAPI},
-				}
-
-				groups = append(groups, initGroups...)
-
-				return groups, nil
-			},
 			GetGroupByNameFunc: func(ctx context.Context, groupName, _ string) (*types.Group, error) {
 				if groupName == "All" {
 					return &types.Group{ID: "id-all", Name: "All", Issued: types.GroupIssuedAPI}, nil
 				}
 
-				return nil, status.Errorf(status.NotFound, "unknown group name")
+				return nil, fmt.Errorf("unknown group name")
 			},
 			GetPeersFunc: func(ctx context.Context, accountID, userID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error) {
 				return maps.Values(TestPeers), nil
@@ -298,84 +287,6 @@ func TestWriteGroup(t *testing.T) {
 	}
 }
 
-func TestGetAllGroups(t *testing.T) {
-	tt := []struct {
-		name           string
-		expectedStatus int
-		expectedBody   bool
-		requestType    string
-		requestPath    string
-		expectedCount  int
-	}{
-		{
-			name:           "Get All Groups",
-			expectedBody:   true,
-			requestType:    http.MethodGet,
-			requestPath:    "/api/groups",
-			expectedStatus: http.StatusOK,
-			expectedCount:  3, // id-jwt-group, id-existed, id-all
-		},
-		{
-			name:           "Get Group By Name - Existing",
-			expectedBody:   true,
-			requestType:    http.MethodGet,
-			requestPath:    "/api/groups?name=All",
-			expectedStatus: http.StatusOK,
-			expectedCount:  1,
-		},
-		{
-			name:           "Get Group By Name - Not Found",
-			expectedBody:   false,
-			requestType:    http.MethodGet,
-			requestPath:    "/api/groups?name=NonExistent",
-			expectedStatus: http.StatusNotFound,
-		},
-	}
-
-	p := initGroupTestData()
-
-	for _, tc := range tt {
-		t.Run(tc.name, func(t *testing.T) {
-			recorder := httptest.NewRecorder()
-			req := httptest.NewRequest(tc.requestType, tc.requestPath, nil)
-			req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{
-				UserId:    "test_user",
-				Domain:    "hotmail.com",
-				AccountId: "test_id",
-			})
-
-			router := mux.NewRouter()
-			router.HandleFunc("/api/groups", p.getAllGroups).Methods("GET")
-			router.ServeHTTP(recorder, req)
-
-			res := recorder.Result()
-			defer res.Body.Close()
-
-			if status := recorder.Code; status != tc.expectedStatus {
-				t.Errorf("handler returned wrong status code: got %v want %v",
-					status, tc.expectedStatus)
-				return
-			}
-
-			if !tc.expectedBody {
-				return
-			}
-
-			content, err := io.ReadAll(res.Body)
-			if err != nil {
-				t.Fatalf("Failed to read response body: %v", err)
-			}
-
-			var groups []api.Group
-			if err = json.Unmarshal(content, &groups); err != nil {
-				t.Fatalf("Response is not in correct json format; %v", err)
-			}
-
-			assert.Equal(t, tc.expectedCount, len(groups))
-		})
-	}
-}
-
 func TestDeleteGroup(t *testing.T) {
 	tt := []struct {
 		name           string

management/server/http/handlers/peers/peers_handler.go

@@ -45,6 +45,19 @@ func NewHandler(accountManager account.Manager, networkMapController network_map
 	}
 }
 
+func (h *Handler) checkPeerStatus(peer *nbpeer.Peer) (*nbpeer.Peer, error) {
+	peerToReturn := peer.Copy()
+	if peer.Status.Connected {
+		// Although we have online status in store we do not yet have an updated channel so have to show it as disconnected
+		// This may happen after server restart when not all peers are yet connected
+		if !h.networkMapController.IsConnected(peer.ID) {
+			peerToReturn.Status.Connected = false
+		}
+	}
+
+	return peerToReturn, nil
+}
+
 func (h *Handler) getPeer(ctx context.Context, accountID, peerID, userID string, w http.ResponseWriter) {
 	peer, err := h.accountManager.GetPeer(ctx, accountID, peerID, userID)
 	if err != nil {
@@ -52,6 +65,11 @@ func (h *Handler) getPeer(ctx context.Context, accountID, peerID, userID string,
 		return
 	}
 
+	peerToReturn, err := h.checkPeerStatus(peer)
+	if err != nil {
+		util.WriteError(ctx, err, w)
+		return
+	}
 	settings, err := h.accountManager.GetAccountSettings(ctx, accountID, activity.SystemInitiator)
 	if err != nil {
 		util.WriteError(ctx, err, w)
@@ -73,7 +91,7 @@ func (h *Handler) getPeer(ctx context.Context, accountID, peerID, userID string,
 	_, valid := validPeers[peer.ID]
 	reason := invalidPeers[peer.ID]
 
-	util.WriteJSONObject(ctx, w, toSinglePeerResponse(peer, grpsInfoMap[peerID], dnsDomain, valid, reason))
+	util.WriteJSONObject(ctx, w, toSinglePeerResponse(peerToReturn, grpsInfoMap[peerID], dnsDomain, valid, reason))
 }
 
 func (h *Handler) updatePeer(ctx context.Context, accountID, userID, peerID string, w http.ResponseWriter, r *http.Request) {
@@ -219,7 +237,13 @@ func (h *Handler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 	grpsInfoMap := groups.ToGroupsInfoMap(grps, len(peers))
 	respBody := make([]*api.PeerBatch, 0, len(peers))
 	for _, peer := range peers {
-		respBody = append(respBody, toPeerListItemResponse(peer, grpsInfoMap[peer.ID], dnsDomain, 0))
+		peerToReturn, err := h.checkPeerStatus(peer)
+		if err != nil {
+			util.WriteError(r.Context(), err, w)
+			return
+		}
+
+		respBody = append(respBody, toPeerListItemResponse(peerToReturn, grpsInfoMap[peer.ID], dnsDomain, 0))
 	}
 
 	validPeersMap, invalidPeersMap, err := h.accountManager.GetValidatedPeers(r.Context(), accountID)

management/server/http/handlers/peers/peers_handler_test.go

@@ -109,6 +109,14 @@ func initTestMetaData(t *testing.T, peers ...*nbpeer.Peer) *Handler {
 		GetDNSDomain(gomock.Any()).
 		Return("domain").
 		AnyTimes()
+	networkMapController.EXPECT().
+		IsConnected(noUpdateChannelTestPeerID).
+		Return(false).
+		AnyTimes()
+	networkMapController.EXPECT().
+		IsConnected(gomock.Any()).
+		Return(true).
+		AnyTimes()
 
 	return &Handler{
 		accountManager: &mock_server.MockAccountManager{
@@ -261,6 +269,14 @@ func TestGetPeers(t *testing.T) {
 			expectedArray:  false,
 			expectedPeer:   peer,
 		},
+		{
+			name:           "GetPeer with no update channel",
+			requestType:    http.MethodGet,
+			requestPath:    "/api/peers/" + peer1.ID,
+			expectedStatus: http.StatusOK,
+			expectedArray:  false,
+			expectedPeer:   expectedPeer1,
+		},
 		{
 			name:           "PutPeer",
 			requestType:    http.MethodPut,
@@ -320,6 +336,8 @@ func TestGetPeers(t *testing.T) {
 				for _, peer := range respBody {
 					if peer.Id == testPeerID {
 						got = peer
+					} else {
+						assert.Equal(t, peer.Connected, false)
 					}
 				}
 
@@ -333,14 +351,14 @@ func TestGetPeers(t *testing.T) {
 
 			t.Log(got)
 
-			assert.Equal(t, tc.expectedPeer.Name, got.Name)
-			assert.Equal(t, tc.expectedPeer.Meta.WtVersion, got.Version)
-			assert.Equal(t, tc.expectedPeer.IP.String(), got.Ip)
-			assert.Equal(t, "OS core", got.Os)
-			assert.Equal(t, tc.expectedPeer.LoginExpirationEnabled, got.LoginExpirationEnabled)
-			assert.Equal(t, tc.expectedPeer.SSHEnabled, got.SshEnabled)
-			assert.Equal(t, tc.expectedPeer.Status.Connected, got.Connected)
-			assert.Equal(t, tc.expectedPeer.Meta.SystemSerialNumber, got.SerialNumber)
+			assert.Equal(t, got.Name, tc.expectedPeer.Name)
+			assert.Equal(t, got.Version, tc.expectedPeer.Meta.WtVersion)
+			assert.Equal(t, got.Ip, tc.expectedPeer.IP.String())
+			assert.Equal(t, got.Os, "OS core")
+			assert.Equal(t, got.LoginExpirationEnabled, tc.expectedPeer.LoginExpirationEnabled)
+			assert.Equal(t, got.SshEnabled, tc.expectedPeer.SSHEnabled)
+			assert.Equal(t, got.Connected, tc.expectedPeer.Status.Connected)
+			assert.Equal(t, got.SerialNumber, tc.expectedPeer.Meta.SystemSerialNumber)
 		})
 	}
 }

management/server/http/middleware/auth_middleware.go

@@ -9,7 +9,6 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"go.opentelemetry.io/otel/metric"
 
 	serverauth "github.com/netbirdio/netbird/management/server/auth"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
@@ -32,7 +31,6 @@ type AuthMiddleware struct {
 	getUserFromUserAuth GetUserFromUserAuthFunc
 	syncUserJWTGroups   SyncUserJWTGroupsFunc
 	rateLimiter         *APIRateLimiter
-	patUsageTracker     *PATUsageTracker
 }
 
 // NewAuthMiddleware instance constructor
@@ -42,29 +40,18 @@ func NewAuthMiddleware(
 	syncUserJWTGroups SyncUserJWTGroupsFunc,
 	getUserFromUserAuth GetUserFromUserAuthFunc,
 	rateLimiterConfig *RateLimiterConfig,
-	meter metric.Meter,
 ) *AuthMiddleware {
 	var rateLimiter *APIRateLimiter
 	if rateLimiterConfig != nil {
 		rateLimiter = NewAPIRateLimiter(rateLimiterConfig)
 	}
 
-	var patUsageTracker *PATUsageTracker
-	if meter != nil {
-		var err error
-		patUsageTracker, err = NewPATUsageTracker(context.Background(), meter)
-		if err != nil {
-			log.Errorf("Failed to create PAT usage tracker: %s", err)
-		}
-	}
-
 	return &AuthMiddleware{
 		authManager:         authManager,
 		ensureAccount:       ensureAccount,
 		syncUserJWTGroups:   syncUserJWTGroups,
 		getUserFromUserAuth: getUserFromUserAuth,
 		rateLimiter:         rateLimiter,
-		patUsageTracker:     patUsageTracker,
 	}
 }
 
@@ -171,10 +158,6 @@ func (m *AuthMiddleware) checkPATFromRequest(r *http.Request, authHeaderParts []
 		return r, fmt.Errorf("error extracting token: %w", err)
 	}
 
-	if m.patUsageTracker != nil {
-		m.patUsageTracker.IncrementUsage(token)
-	}
-
 	if m.rateLimiter != nil {
 		if !m.rateLimiter.Allow(token) {
 			return r, status.Errorf(status.TooManyRequests, "too many requests")

management/server/http/middleware/auth_middleware_test.go

@@ -208,7 +208,6 @@ func TestAuthMiddleware_Handler(t *testing.T) {
 			return &types.User{}, nil
 		},
 		nil,
-		nil,
 	)
 
 	handlerToTest := authMiddleware.Handler(nextHandler)
@@ -267,7 +266,6 @@ func TestAuthMiddleware_RateLimiting(t *testing.T) {
 				return &types.User{}, nil
 			},
 			rateLimitConfig,
-			nil,
 		)
 
 		handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -319,7 +317,6 @@ func TestAuthMiddleware_RateLimiting(t *testing.T) {
 				return &types.User{}, nil
 			},
 			rateLimitConfig,
-			nil,
 		)
 
 		handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -362,7 +359,6 @@ func TestAuthMiddleware_RateLimiting(t *testing.T) {
 				return &types.User{}, nil
 			},
 			rateLimitConfig,
-			nil,
 		)
 
 		handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -406,7 +402,6 @@ func TestAuthMiddleware_RateLimiting(t *testing.T) {
 				return &types.User{}, nil
 			},
 			rateLimitConfig,
-			nil,
 		)
 
 		handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -470,7 +465,6 @@ func TestAuthMiddleware_RateLimiting(t *testing.T) {
 				return &types.User{}, nil
 			},
 			rateLimitConfig,
-			nil,
 		)
 
 		handler := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -587,7 +581,6 @@ func TestAuthMiddleware_Handler_Child(t *testing.T) {
 			return &types.User{}, nil
 		},
 		nil,
-		nil,
 	)
 
 	for _, tc := range tt {

management/server/http/middleware/pat_usage_tracker.go

@@ -1,87 +0,0 @@
-package middleware
-
-import (
-	"context"
-	"maps"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-	"go.opentelemetry.io/otel/metric"
-)
-
-// PATUsageTracker tracks PAT usage metrics
-type PATUsageTracker struct {
-	usageCounters map[string]int64
-	mu            sync.Mutex
-	stopChan      chan struct{}
-	ctx           context.Context
-	histogram     metric.Int64Histogram
-}
-
-// NewPATUsageTracker creates a new PAT usage tracker with metrics
-func NewPATUsageTracker(ctx context.Context, meter metric.Meter) (*PATUsageTracker, error) {
-	histogram, err := meter.Int64Histogram(
-		"management.pat.usage_distribution",
-		metric.WithUnit("1"),
-		metric.WithDescription("Distribution of PAT token usage counts per minute"),
-	)
-	if err != nil {
-		return nil, err
-	}
-
-	tracker := &PATUsageTracker{
-		usageCounters: make(map[string]int64),
-		stopChan:      make(chan struct{}),
-		ctx:           ctx,
-		histogram:     histogram,
-	}
-
-	go tracker.reportLoop()
-
-	return tracker, nil
-}
-
-// IncrementUsage increments the usage counter for a given token
-func (t *PATUsageTracker) IncrementUsage(token string) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	t.usageCounters[token]++
-}
-
-// reportLoop reports the usage buckets every minute
-func (t *PATUsageTracker) reportLoop() {
-	ticker := time.NewTicker(1 * time.Minute)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ticker.C:
-			t.reportUsageBuckets()
-		case <-t.stopChan:
-			return
-		}
-	}
-}
-
-// reportUsageBuckets reports all token usage counts and resets counters
-func (t *PATUsageTracker) reportUsageBuckets() {
-	t.mu.Lock()
-	snapshot := maps.Clone(t.usageCounters)
-
-	clear(t.usageCounters)
-	t.mu.Unlock()
-
-	totalTokens := len(snapshot)
-	if totalTokens > 0 {
-		for _, count := range snapshot {
-			t.histogram.Record(t.ctx, count)
-		}
-		log.Debugf("PAT usage in last minute: %d unique tokens used", totalTokens)
-	}
-}
-
-// Stop stops the reporting goroutine
-func (t *PATUsageTracker) Stop() {
-	close(t.stopChan)
-}

management/server/http/testing/testing_tools/channel/channel.go

@@ -15,8 +15,6 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 
 	"github.com/netbirdio/netbird/management/server"
@@ -30,6 +28,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/networks"
 	"github.com/netbirdio/netbird/management/server/networks/resources"
 	"github.com/netbirdio/netbird/management/server/networks/routers"
+	"github.com/netbirdio/netbird/management/server/peers"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -73,7 +72,7 @@ func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPee
 
 	ctx := context.Background()
 	requestBuffer := server.NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), &config.Config{})
 	am, err := server.BuildManager(ctx, nil, store, networkMapController, nil, "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager, false)
 	if err != nil {
 		t.Fatalf("Failed to create manager: %v", err)

management/server/management_proto_test.go

@@ -24,14 +24,13 @@ import (
 	"github.com/netbirdio/netbird/formatter/hook"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/groups"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -364,9 +363,7 @@ func startManagementForTest(t *testing.T, testFile string, config *config.Config
 
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, store)
-	ephemeralMgr := manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager))
-
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), ephemeralMgr, config)
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), config)
 	accountManager, err := BuildManager(ctx, nil, store, networkMapController, nil, "",
 		eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 
@@ -375,13 +372,10 @@ func startManagementForTest(t *testing.T, testFile string, config *config.Config
 		return nil, nil, "", cleanup, err
 	}
 
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		cleanup()
-		return nil, nil, "", cleanup, err
-	}
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
 
-	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, secretsManager, nil, nil, MockIntegratedValidator{}, networkMapController)
+	ephemeralMgr := manager.NewEphemeralManager(store, accountManager)
+	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, updateManager, secretsManager, nil, ephemeralMgr, nil, MockIntegratedValidator{}, networkMapController)
 	if err != nil {
 		return nil, nil, "", cleanup, err
 	}

management/server/management_test.go

@@ -22,14 +22,13 @@ import (
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/groups"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -206,7 +205,7 @@ func startServer(
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := server.NewAccountRequestBuffer(ctx, str)
-	networkMapController := controller.NewController(ctx, str, metrics, updateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(str, peers.NewManager(str, permissionsManager)), config)
+	networkMapController := controller.NewController(ctx, str, metrics, updateManager, requestBuffer, server.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), config)
 
 	accountManager, err := server.BuildManager(
 		context.Background(),
@@ -229,16 +228,15 @@ func startServer(
 	}
 
 	groupsManager := groups.NewManager(str, permissionsManager, accountManager)
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		t.Fatalf("failed creating secrets manager: %v", err)
-	}
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
 	mgmtServer, err := nbgrpc.NewServer(
 		config,
 		accountManager,
 		settingsMockManager,
+		updateManager,
 		secretsManager,
 		nil,
+		&manager.EphemeralManager{},
 		nil,
 		server.MockIntegratedValidator{},
 		networkMapController,

management/server/mock_server/account_mock.go

@@ -15,6 +15,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/idp"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -975,6 +976,11 @@ func (am *MockAccountManager) GetCurrentUserInfo(ctx context.Context, userAuth a
 	return nil, status.Errorf(codes.Unimplemented, "method GetCurrentUserInfo is not implemented")
 }
 
+// SetEphemeralManager mocks SetEphemeralManager of the AccountManager interface
+func (am *MockAccountManager) SetEphemeralManager(em ephemeral.Manager) {
+	// Mock implementation - does nothing
+}
+
 func (am *MockAccountManager) AllowSync(key string, hash uint64) bool {
 	if am.AllowSyncFunc != nil {
 		return am.AllowSyncFunc(key, hash)

management/server/nameserver_test.go

@@ -13,8 +13,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
@@ -794,7 +792,7 @@ func createNSManager(t *testing.T) (*DefaultAccountManager, error) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), &config.Config{})
 
 	return BuildManager(context.Background(), nil, store, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 }

management/server/peer.go

@@ -136,10 +136,7 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
 	}
 
 	if expired {
-		err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID})
-		if err != nil {
-			return fmt.Errorf("notify network map controller of peer update: %w", err)
-		}
+		am.networkMapController.OnPeerUpdated(accountID, peer)
 	}
 
 	return nil
@@ -312,10 +309,7 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		}
 	}
 
-	err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID})
-	if err != nil {
-		return nil, fmt.Errorf("notify network map controller of peer update: %w", err)
-	}
+	am.networkMapController.OnPeerUpdated(accountID, peer)
 
 	return peer, nil
 }
@@ -371,8 +365,13 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 		storeEvent()
 	}
 
-	if err := am.networkMapController.OnPeersDeleted(ctx, accountID, []string{peerID}); err != nil {
-		log.WithContext(ctx).Errorf("failed to delete peer %s from network map: %v", peerID, err)
+	err = am.networkMapController.DeletePeer(ctx, accountID, peer.ID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to delete peer %s from network map: %v", peer.ID, err)
+	}
+
+	if err := am.networkMapController.OnPeerDeleted(ctx, accountID, peerID); err != nil {
+		log.WithContext(ctx).Errorf("failed to update network map cache for peer %s: %v", peerID, err)
 	}
 
 	return nil
@@ -584,6 +583,11 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 				return fmt.Errorf("failed adding peer to All group: %w", err)
 			}
 
+			if temporary {
+				// we are running the on disconnect handler so that it is considered not connected as we are adding the peer manually
+				am.ephemeralManager.OnPeerDisconnected(ctx, newPeer)
+			}
+
 			if addedByUser {
 				err := transaction.SaveUserLastLogin(ctx, accountID, userID, newPeer.GetLastLogin())
 				if err != nil {
@@ -641,7 +645,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKe
 
 	am.StoreEvent(ctx, opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)
 
-	if err := am.networkMapController.OnPeersAdded(ctx, accountID, []string{newPeer.ID}); err != nil {
+	if err := am.networkMapController.OnPeerAdded(ctx, accountID, newPeer.ID); err != nil {
 		log.WithContext(ctx).Errorf("failed to update network map cache for peer %s: %v", newPeer.ID, err)
 	}
 
@@ -725,10 +729,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
 	}
 
 	if isStatusChanged || sync.UpdateAccountPeers || (updated && (len(postureChecks) > 0 || versionChanged)) {
-		err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID})
-		if err != nil {
-			return nil, nil, nil, 0, fmt.Errorf("notify network map controller of peer update: %w", err)
-		}
+		am.networkMapController.OnPeerUpdated(accountID, peer)
 	}
 
 	return am.networkMapController.GetValidatedPeerWithMap(ctx, peerNotValid, accountID, peer)
@@ -856,10 +857,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
 	log.WithContext(ctx).Debugf("LoginPeer: transaction took %v", time.Since(startTransaction))
 
 	if updateRemotePeers || isStatusChanged || (isPeerUpdated && len(postureChecks) > 0) {
-		err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID})
-		if err != nil {
-			return nil, nil, nil, fmt.Errorf("notify network map controller of peer update: %w", err)
-		}
+		am.networkMapController.OnPeerUpdated(accountID, peer)
 	}
 
 	p, nmap, pc, _, err := am.networkMapController.GetValidatedPeerWithMap(ctx, isRequiresApproval, accountID, peer)

management/server/peer_test.go

@@ -28,8 +28,6 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools"
@@ -1060,7 +1058,6 @@ func testUpdateAccountPeers(t *testing.T) {
 
 			for _, channel := range peerChannels {
 				update := <-channel
-				assert.Nil(t, update.Update.NetbirdConfig)
 				assert.Equal(t, tc.peers, len(update.Update.NetworkMap.RemotePeers))
 				assert.Equal(t, tc.peers*2, len(update.Update.NetworkMap.FirewallRules))
 			}
@@ -1293,7 +1290,7 @@ func Test_RegisterPeerByUser(t *testing.T) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, s)
-	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(s, peers.NewManager(s, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), &config.Config{})
 
 	am, err := BuildManager(context.Background(), nil, s, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
@@ -1378,7 +1375,7 @@ func Test_RegisterPeerBySetupKey(t *testing.T) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, s)
-	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(s, peers.NewManager(s, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), &config.Config{})
 
 	am, err := BuildManager(context.Background(), nil, s, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
@@ -1531,7 +1528,7 @@ func Test_RegisterPeerRollbackOnFailure(t *testing.T) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, s)
-	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(s, peers.NewManager(s, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), &config.Config{})
 
 	am, err := BuildManager(context.Background(), nil, s, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
@@ -1611,7 +1608,7 @@ func Test_LoginPeer(t *testing.T) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, s)
-	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(s, peers.NewManager(s, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, s, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.cloud", port_forwarding.NewControllerMock(), &config.Config{})
 
 	am, err := BuildManager(context.Background(), nil, s, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)

management/server/peers/ephemeral/interface.go

@@ -2,15 +2,10 @@ package ephemeral
 
 import (
 	"context"
-	"time"
 
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 )
 
-const (
-	EphemeralLifeTime = 10 * time.Minute
-)
-
 type Manager interface {
 	LoadInitialPeers(ctx context.Context)
 	Stop()

management/server/peers/ephemeral/manager/ephemeral.go

@@ -7,15 +7,14 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
+	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
-
 	"github.com/netbirdio/netbird/management/server/store"
 )
 
 const (
+	ephemeralLifeTime = 10 * time.Minute
 	// cleanupWindow is the time window to wait after nearest peer deadline to start the cleanup procedure.
 	cleanupWindow = 1 * time.Minute
 )
@@ -34,11 +33,11 @@ type ephemeralPeer struct {
 // todo: consider to remove peer from ephemeral list when the peer has been deleted via API. If we do not do it
 // in worst case we will get invalid error message in this manager.
 
-// EphemeralManager keep a list of ephemeral peers. After EphemeralLifeTime inactivity the peer will be deleted
+// EphemeralManager keep a list of ephemeral peers. After ephemeralLifeTime inactivity the peer will be deleted
 // automatically. Inactivity means the peer disconnected from the Management server.
 type EphemeralManager struct {
-	store        store.Store
-	peersManager peers.Manager
+	store          store.Store
+	accountManager nbAccount.Manager
 
 	headPeer  *ephemeralPeer
 	tailPeer  *ephemeralPeer
@@ -50,12 +49,12 @@ type EphemeralManager struct {
 }
 
 // NewEphemeralManager instantiate new EphemeralManager
-func NewEphemeralManager(store store.Store, peersManager peers.Manager) *EphemeralManager {
+func NewEphemeralManager(store store.Store, accountManager nbAccount.Manager) *EphemeralManager {
 	return &EphemeralManager{
-		store:        store,
-		peersManager: peersManager,
+		store:          store,
+		accountManager: accountManager,
 
-		lifeTime:      ephemeral.EphemeralLifeTime,
+		lifeTime:      ephemeralLifeTime,
 		cleanupWindow: cleanupWindow,
 	}
 }
@@ -107,7 +106,7 @@ func (e *EphemeralManager) OnPeerConnected(ctx context.Context, peer *nbpeer.Pee
 }
 
 // OnPeerDisconnected add the peer to the linked list of ephemeral peers. Because of the peer
-// is inactive it will be deleted after the EphemeralLifeTime period.
+// is inactive it will be deleted after the ephemeralLifeTime period.
 func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.Peer) {
 	if !peer.Ephemeral {
 		return
@@ -181,18 +180,20 @@ func (e *EphemeralManager) cleanup(ctx context.Context) {
 
 	e.peersLock.Unlock()
 
-	peerIDsPerAccount := make(map[string][]string)
-	for id, p := range deletePeers {
-		peerIDsPerAccount[p.accountID] = append(peerIDsPerAccount[p.accountID], id)
-	}
+	bufferAccountCall := make(map[string]struct{})
 
-	for accountID, peerIDs := range peerIDsPerAccount {
-		log.WithContext(ctx).Debugf("delete ephemeral peers for account: %s", accountID)
-		err := e.peersManager.DeletePeers(ctx, accountID, peerIDs, activity.SystemInitiator, true)
+	for id, p := range deletePeers {
+		log.WithContext(ctx).Debugf("delete ephemeral peer: %s", id)
+		err := e.accountManager.DeletePeer(ctx, p.accountID, id, activity.SystemInitiator)
 		if err != nil {
 			log.WithContext(ctx).Errorf("failed to delete ephemeral peer: %s", err)
+		} else {
+			bufferAccountCall[p.accountID] = struct{}{}
 		}
 	}
+	for accountID := range bufferAccountCall {
+		e.accountManager.BufferUpdateAccountPeers(ctx, accountID)
+	}
 }
 
 func (e *EphemeralManager) addPeer(accountID string, peerID string, deadline time.Time) {

management/server/peers/ephemeral/manager/ephemeral_test.go

@@ -7,13 +7,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/mock/gomock"
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 
 	nbdns "github.com/netbirdio/netbird/dns"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
 	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -94,27 +91,17 @@ func TestNewManager(t *testing.T) {
 	}
 
 	store := &MockStore{}
-	ctrl := gomock.NewController(t)
-	peersManager := peers.NewMockManager(ctrl)
+	am := MockAccountManager{
+		store: store,
+	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	// Expect DeletePeers to be called for ephemeral peers
-	peersManager.EXPECT().
-		DeletePeers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), true).
-		DoAndReturn(func(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-			for _, peerID := range peerIDs {
-				delete(store.account.Peers, peerID)
-			}
-			return nil
-		}).
-		AnyTimes()
-
-	mgr := NewEphemeralManager(store, peersManager)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
-	startTime = startTime.Add(ephemeral.EphemeralLifeTime + 1)
+	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
 	if len(store.account.Peers) != numberOfPeers {
@@ -132,29 +119,19 @@ func TestNewManagerPeerConnected(t *testing.T) {
 	}
 
 	store := &MockStore{}
-	ctrl := gomock.NewController(t)
-	peersManager := peers.NewMockManager(ctrl)
+	am := MockAccountManager{
+		store: store,
+	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	// Expect DeletePeers to be called for ephemeral peers (except the connected one)
-	peersManager.EXPECT().
-		DeletePeers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), true).
-		DoAndReturn(func(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-			for _, peerID := range peerIDs {
-				delete(store.account.Peers, peerID)
-			}
-			return nil
-		}).
-		AnyTimes()
-
-	mgr := NewEphemeralManager(store, peersManager)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
 	mgr.OnPeerConnected(context.Background(), store.account.Peers["ephemeral_peer_0"])
 
-	startTime = startTime.Add(ephemeral.EphemeralLifeTime + 1)
+	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
 	expected := numberOfPeers + 1
@@ -173,25 +150,15 @@ func TestNewManagerPeerDisconnected(t *testing.T) {
 	}
 
 	store := &MockStore{}
-	ctrl := gomock.NewController(t)
-	peersManager := peers.NewMockManager(ctrl)
+	am := MockAccountManager{
+		store: store,
+	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	// Expect DeletePeers to be called for the one disconnected peer
-	peersManager.EXPECT().
-		DeletePeers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), true).
-		DoAndReturn(func(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-			for _, peerID := range peerIDs {
-				delete(store.account.Peers, peerID)
-			}
-			return nil
-		}).
-		AnyTimes()
-
-	mgr := NewEphemeralManager(store, peersManager)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
 	for _, v := range store.account.Peers {
 		mgr.OnPeerConnected(context.Background(), v)
@@ -199,7 +166,7 @@ func TestNewManagerPeerDisconnected(t *testing.T) {
 	}
 	mgr.OnPeerDisconnected(context.Background(), store.account.Peers["ephemeral_peer_0"])
 
-	startTime = startTime.Add(ephemeral.EphemeralLifeTime + 1)
+	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
 	expected := numberOfPeers + numberOfEphemeralPeers - 1
@@ -214,63 +181,25 @@ func TestCleanupSchedulingBehaviorIsBatched(t *testing.T) {
 		testLifeTime      = 1 * time.Second
 		testCleanupWindow = 100 * time.Millisecond
 	)
-
-	t.Cleanup(func() {
-		timeNow = time.Now
-	})
-	startTime := time.Now()
-	timeNow = func() time.Time {
-		return startTime
-	}
-
 	mockStore := &MockStore{}
-	account := newAccountWithId(context.Background(), "account", "", "", false)
-	mockStore.account = account
-
-	wg := &sync.WaitGroup{}
-	wg.Add(ephemeralPeers)
 	mockAM := &MockAccountManager{
 		store: mockStore,
-		wg:    wg,
 	}
-
-	ctrl := gomock.NewController(t)
-	peersManager := peers.NewMockManager(ctrl)
-
-	// Set up expectation that DeletePeers will be called once with all peer IDs
-	peersManager.EXPECT().
-		DeletePeers(gomock.Any(), account.Id, gomock.Any(), gomock.Any(), true).
-		DoAndReturn(func(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-			// Simulate the actual deletion behavior
-			for _, peerID := range peerIDs {
-				err := mockAM.DeletePeer(ctx, accountID, peerID, userID)
-				if err != nil {
-					return err
-				}
-			}
-			mockAM.BufferUpdateAccountPeers(ctx, accountID)
-			return nil
-		}).
-		Times(1)
-
-	mgr := NewEphemeralManager(mockStore, peersManager)
+	mockAM.wg = &sync.WaitGroup{}
+	mockAM.wg.Add(ephemeralPeers)
+	mgr := NewEphemeralManager(mockStore, mockAM)
 	mgr.lifeTime = testLifeTime
 	mgr.cleanupWindow = testCleanupWindow
 
-	// Add peers and disconnect them at slightly different times (within cleanup window)
+	account := newAccountWithId(context.Background(), "account", "", "", false)
+	mockStore.account = account
 	for i := range ephemeralPeers {
 		p := &nbpeer.Peer{ID: fmt.Sprintf("peer-%d", i), AccountID: account.Id, Ephemeral: true}
 		mockStore.account.Peers[p.ID] = p
+		time.Sleep(testCleanupWindow / ephemeralPeers)
 		mgr.OnPeerDisconnected(context.Background(), p)
-		startTime = startTime.Add(testCleanupWindow / (ephemeralPeers * 2))
 	}
-
-	// Advance time past the lifetime to trigger cleanup
-	startTime = startTime.Add(testLifeTime + testCleanupWindow)
-
-	// Wait for all deletions to complete
-	wg.Wait()
-
+	mockAM.wg.Wait()
 	assert.Len(t, mockStore.account.Peers, 0, "all ephemeral peers should be cleaned up after the lifetime")
 	assert.Equal(t, 1, mockAM.GetBufferUpdateCalls(account.Id), "buffer update should be called once")
 	assert.Equal(t, ephemeralPeers, mockAM.GetDeletePeerCalls(), "should have deleted all peers")

management/server/peers/manager.go

@@ -0,0 +1,68 @@
+package peers
+
+//go:generate go run github.com/golang/mock/mockgen -package peers -destination=manager_mock.go -source=./manager.go -build_flags=-mod=mod
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/permissions"
+	"github.com/netbirdio/netbird/management/server/permissions/modules"
+	"github.com/netbirdio/netbird/management/server/permissions/operations"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+type Manager interface {
+	GetPeer(ctx context.Context, accountID, userID, peerID string) (*peer.Peer, error)
+	GetPeerAccountID(ctx context.Context, peerID string) (string, error)
+	GetAllPeers(ctx context.Context, accountID, userID string) ([]*peer.Peer, error)
+	GetPeersByGroupIDs(ctx context.Context, accountID string, groupsIDs []string) ([]*peer.Peer, error)
+}
+
+type managerImpl struct {
+	store              store.Store
+	permissionsManager permissions.Manager
+}
+
+func NewManager(store store.Store, permissionsManager permissions.Manager) Manager {
+	return &managerImpl{
+		store:              store,
+		permissionsManager: permissionsManager,
+	}
+}
+
+func (m *managerImpl) GetPeer(ctx context.Context, accountID, userID, peerID string) (*peer.Peer, error) {
+	allowed, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Read)
+	if err != nil {
+		return nil, fmt.Errorf("failed to validate user permissions: %w", err)
+	}
+
+	if !allowed {
+		return nil, status.NewPermissionDeniedError()
+	}
+
+	return m.store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
+}
+
+func (m *managerImpl) GetAllPeers(ctx context.Context, accountID, userID string) ([]*peer.Peer, error) {
+	allowed, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Peers, operations.Read)
+	if err != nil {
+		return nil, fmt.Errorf("failed to validate user permissions: %w", err)
+	}
+
+	if !allowed {
+		return m.store.GetUserPeers(ctx, store.LockingStrengthNone, accountID, userID)
+	}
+
+	return m.store.GetAccountPeers(ctx, store.LockingStrengthNone, accountID, "", "")
+}
+
+func (m *managerImpl) GetPeerAccountID(ctx context.Context, peerID string) (string, error) {
+	return m.store.GetAccountIDByPeerID(ctx, store.LockingStrengthNone, peerID)
+}
+
+func (m *managerImpl) GetPeersByGroupIDs(ctx context.Context, accountID string, groupsIDs []string) ([]*peer.Peer, error) {
+	return m.store.GetPeersByGroupIDs(ctx, accountID, groupsIDs)
+}

management/server/peers/manager_mock.go

@@ -9,9 +9,6 @@ import (
 	reflect "reflect"
 
 	gomock "github.com/golang/mock/gomock"
-	network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map"
-	account "github.com/netbirdio/netbird/management/server/account"
-	integrated_validator "github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	peer "github.com/netbirdio/netbird/management/server/peer"
 )
 
@@ -38,20 +35,6 @@ func (m *MockManager) EXPECT() *MockManagerMockRecorder {
 	return m.recorder
 }
 
-// DeletePeers mocks base method.
-func (m *MockManager) DeletePeers(ctx context.Context, accountID string, peerIDs []string, userID string, checkConnected bool) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DeletePeers", ctx, accountID, peerIDs, userID, checkConnected)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// DeletePeers indicates an expected call of DeletePeers.
-func (mr *MockManagerMockRecorder) DeletePeers(ctx, accountID, peerIDs, userID, checkConnected interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeletePeers", reflect.TypeOf((*MockManager)(nil).DeletePeers), ctx, accountID, peerIDs, userID, checkConnected)
-}
-
 // GetAllPeers mocks base method.
 func (m *MockManager) GetAllPeers(ctx context.Context, accountID, userID string) ([]*peer.Peer, error) {
 	m.ctrl.T.Helper()
@@ -111,39 +94,3 @@ func (mr *MockManagerMockRecorder) GetPeersByGroupIDs(ctx, accountID, groupsIDs
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPeersByGroupIDs", reflect.TypeOf((*MockManager)(nil).GetPeersByGroupIDs), ctx, accountID, groupsIDs)
 }
-
-// SetAccountManager mocks base method.
-func (m *MockManager) SetAccountManager(accountManager account.Manager) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "SetAccountManager", accountManager)
-}
-
-// SetAccountManager indicates an expected call of SetAccountManager.
-func (mr *MockManagerMockRecorder) SetAccountManager(accountManager interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetAccountManager", reflect.TypeOf((*MockManager)(nil).SetAccountManager), accountManager)
-}
-
-// SetIntegratedPeerValidator mocks base method.
-func (m *MockManager) SetIntegratedPeerValidator(integratedPeerValidator integrated_validator.IntegratedValidator) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "SetIntegratedPeerValidator", integratedPeerValidator)
-}
-
-// SetIntegratedPeerValidator indicates an expected call of SetIntegratedPeerValidator.
-func (mr *MockManagerMockRecorder) SetIntegratedPeerValidator(integratedPeerValidator interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetIntegratedPeerValidator", reflect.TypeOf((*MockManager)(nil).SetIntegratedPeerValidator), integratedPeerValidator)
-}
-
-// SetNetworkMapController mocks base method.
-func (m *MockManager) SetNetworkMapController(networkMapController network_map.Controller) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "SetNetworkMapController", networkMapController)
-}
-
-// SetNetworkMapController indicates an expected call of SetNetworkMapController.
-func (mr *MockManagerMockRecorder) SetNetworkMapController(networkMapController interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetNetworkMapController", reflect.TypeOf((*MockManager)(nil).SetNetworkMapController), networkMapController)
-}

management/server/route_test.go

@@ -16,8 +16,6 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
@@ -1293,7 +1291,7 @@ func createRouterManager(t *testing.T) (*DefaultAccountManager, *update_channel.
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peers.NewManager(store, permissionsManager)), &config.Config{})
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), &config.Config{})
 
 	am, err := BuildManager(context.Background(), nil, store, networkMapController, nil, "", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	if err != nil {

management/server/user.go

@@ -263,11 +263,15 @@ func (am *DefaultAccountManager) DeleteUser(ctx context.Context, accountID, init
 		return err
 	}
 
-	_, err = am.deleteRegularUser(ctx, accountID, initiatorUserID, userInfo)
+	updateAccountPeers, err := am.deleteRegularUser(ctx, accountID, initiatorUserID, userInfo)
 	if err != nil {
 		return err
 	}
 
+	if updateAccountPeers {
+		am.UpdateAccountPeers(ctx, accountID)
+	}
+
 	return nil
 }
 
@@ -592,15 +596,9 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 }
 
 // prepareUserUpdateEvents prepares a list user update events based on the changes between the old and new user data.
-func (am *DefaultAccountManager) prepareUserUpdateEvents(ctx context.Context, accountID string, initiatorUserID string, oldUser, newUser *types.User, transferredOwnerRole bool, isNewUser bool, removedGroupIDs, addedGroupIDs []string, tx store.Store) []func() {
+func (am *DefaultAccountManager) prepareUserUpdateEvents(ctx context.Context, accountID string, initiatorUserID string, oldUser, newUser *types.User, transferredOwnerRole bool, removedGroupIDs, addedGroupIDs []string, tx store.Store) []func() {
 	var eventsToStore []func()
 
-	if isNewUser {
-		eventsToStore = append(eventsToStore, func() {
-			am.StoreEvent(ctx, initiatorUserID, newUser.Id, accountID, activity.UserCreated, nil)
-		})
-	}
-
 	if oldUser.IsBlocked() != newUser.IsBlocked() {
 		if newUser.IsBlocked() {
 			eventsToStore = append(eventsToStore, func() {
@@ -663,7 +661,7 @@ func (am *DefaultAccountManager) processUserUpdate(ctx context.Context, transact
 		return false, nil, nil, nil, status.Errorf(status.InvalidArgument, "provided user update is nil")
 	}
 
-	oldUser, isNewUser, err := getUserOrCreateIfNotExists(ctx, transaction, accountID, update, addIfNotExists)
+	oldUser, err := getUserOrCreateIfNotExists(ctx, transaction, accountID, update, addIfNotExists)
 	if err != nil {
 		return false, nil, nil, nil, err
 	}
@@ -718,30 +716,30 @@ func (am *DefaultAccountManager) processUserUpdate(ctx context.Context, transact
 	}
 
 	updateAccountPeers := len(userPeers) > 0
-	userEventsToAdd := am.prepareUserUpdateEvents(ctx, updatedUser.AccountID, initiatorUserId, oldUser, updatedUser, transferredOwnerRole, isNewUser, removedGroups, addedGroups, transaction)
+	userEventsToAdd := am.prepareUserUpdateEvents(ctx, updatedUser.AccountID, initiatorUserId, oldUser, updatedUser, transferredOwnerRole, removedGroups, addedGroups, transaction)
 
 	return updateAccountPeers, updatedUser, peersToExpire, userEventsToAdd, nil
 }
 
 // getUserOrCreateIfNotExists retrieves the existing user or creates a new one if it doesn't exist.
-func getUserOrCreateIfNotExists(ctx context.Context, transaction store.Store, accountID string, update *types.User, addIfNotExists bool) (*types.User, bool, error) {
+func getUserOrCreateIfNotExists(ctx context.Context, transaction store.Store, accountID string, update *types.User, addIfNotExists bool) (*types.User, error) {
 	existingUser, err := transaction.GetUserByUserID(ctx, store.LockingStrengthNone, update.Id)
 	if err != nil {
 		if sErr, ok := status.FromError(err); ok && sErr.Type() == status.NotFound {
 			if !addIfNotExists {
-				return nil, false, status.Errorf(status.NotFound, "user to update doesn't exist: %s", update.Id)
+				return nil, status.Errorf(status.NotFound, "user to update doesn't exist: %s", update.Id)
 			}
 			update.AccountID = accountID
-			return update, true, nil // use all fields from update if addIfNotExists is true
+			return update, nil // use all fields from update if addIfNotExists is true
 		}
-		return nil, false, err
+		return nil, err
 	}
 
 	if existingUser.AccountID != accountID {
-		return nil, false, status.Errorf(status.InvalidArgument, "user account ID mismatch")
+		return nil, status.Errorf(status.InvalidArgument, "user account ID mismatch")
 	}
 
-	return existingUser, false, nil
+	return existingUser, nil
 }
 
 func handleOwnerRoleTransfer(ctx context.Context, transaction store.Store, initiatorUser, update *types.User) (bool, error) {
@@ -994,17 +992,14 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
 			peer.UserID, peer.ID, accountID,
 			activity.PeerLoginExpired, peer.EventMeta(dnsDomain),
 		)
-	}
 
-	err = am.networkMapController.OnPeersUpdated(ctx, accountID, peerIDs)
-	if err != nil {
-		return fmt.Errorf("notify network map controller of peer update: %w", err)
+		am.networkMapController.OnPeerUpdated(accountID, peer)
 	}
 
 	if len(peerIDs) != 0 {
 		// this will trigger peer disconnect from the management service
 		log.Debugf("Expiring %d peers for account %s", len(peerIDs), accountID)
-		am.networkMapController.DisconnectPeers(ctx, accountID, peerIDs)
+		am.networkMapController.DisconnectPeers(ctx, peerIDs)
 	}
 	return nil
 }
@@ -1050,6 +1045,7 @@ func (am *DefaultAccountManager) DeleteRegularUsers(ctx context.Context, account
 	}
 
 	var allErrors error
+	var updateAccountPeers bool
 
 	for _, targetUserID := range targetUserIDs {
 		if initiatorUserID == targetUserID {
@@ -1080,11 +1076,19 @@ func (am *DefaultAccountManager) DeleteRegularUsers(ctx context.Context, account
 			continue
 		}
 
-		_, err = am.deleteRegularUser(ctx, accountID, initiatorUserID, userInfo)
+		userHadPeers, err := am.deleteRegularUser(ctx, accountID, initiatorUserID, userInfo)
 		if err != nil {
 			allErrors = errors.Join(allErrors, err)
 			continue
 		}
+
+		if userHadPeers {
+			updateAccountPeers = true
+		}
+	}
+
+	if updateAccountPeers {
+		am.UpdateAccountPeers(ctx, accountID)
 	}
 
 	return allErrors
@@ -1142,12 +1146,15 @@ func (am *DefaultAccountManager) deleteRegularUser(ctx context.Context, accountI
 		return false, err
 	}
 
-	var peerIDs []string
 	for _, peer := range userPeers {
-		peerIDs = append(peerIDs, peer.ID)
-	}
-	if err := am.networkMapController.OnPeersDeleted(ctx, accountID, peerIDs); err != nil {
-		log.WithContext(ctx).Errorf("failed to delete peers %s from network map: %v", peerIDs, err)
+		err = am.networkMapController.DeletePeer(ctx, accountID, peer.ID)
+		if err != nil {
+			log.WithContext(ctx).Errorf("failed to delete peer %s from network map: %v", peer.ID, err)
+		}
+
+		if err := am.networkMapController.OnPeerDeleted(ctx, accountID, peer.ID); err != nil {
+			log.WithContext(ctx).Errorf("failed to update network map cache for peer %s: %v", peer.ID, err)
+		}
 	}
 
 	for _, addPeerRemovedEvent := range addPeerRemovedEvents {

management/server/user_test.go

@@ -8,10 +8,8 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
-	"go.uber.org/mock/gomock"
 	"golang.org/x/exp/maps"
 
-	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
@@ -549,7 +547,7 @@ func TestUser_InviteNewUser(t *testing.T) {
 		permissionsManager: permissionsManager,
 	}
 
-	cs, err := nbcache.NewStore(context.Background(), nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval, nbcache.DefaultIDPCacheOpenConn)
+	cs, err := nbcache.NewStore(context.Background(), nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
 	require.NoError(t, err)
 
 	am.cacheManager = nbcache.NewAccountUserDataCache(am.loadAccount, cs)
@@ -741,18 +739,11 @@ func TestUser_DeleteUser_regularUser(t *testing.T) {
 		t.Fatalf("Error when saving account: %s", err)
 	}
 
-	ctrl := gomock.NewController(t)
-	networkMapControllerMock := network_map.NewMockController(ctrl)
-	networkMapControllerMock.EXPECT().
-		OnPeersDeleted(gomock.Any(), gomock.Any(), gomock.Any()).
-		Return(nil)
-
 	permissionsManager := permissions.NewManager(store)
 	am := DefaultAccountManager{
-		Store:                store,
-		eventStore:           &activity.InMemoryEventStore{},
-		permissionsManager:   permissionsManager,
-		networkMapController: networkMapControllerMock,
+		Store:              store,
+		eventStore:         &activity.InMemoryEventStore{},
+		permissionsManager: permissionsManager,
 	}
 
 	testCases := []struct {
@@ -857,20 +848,12 @@ func TestUser_DeleteUser_RegularUsers(t *testing.T) {
 		t.Fatalf("Error when saving account: %s", err)
 	}
 
-	ctrl := gomock.NewController(t)
-	networkMapControllerMock := network_map.NewMockController(ctrl)
-	networkMapControllerMock.EXPECT().
-		OnPeersDeleted(gomock.Any(), gomock.Any(), gomock.Any()).
-		Return(nil).
-		AnyTimes()
-
 	permissionsManager := permissions.NewManager(store)
 	am := DefaultAccountManager{
 		Store:                   store,
 		eventStore:              &activity.InMemoryEventStore{},
 		integratedPeerValidator: MockIntegratedValidator{},
 		permissionsManager:      permissionsManager,
-		networkMapController:    networkMapControllerMock,
 	}
 
 	testCases := []struct {
@@ -1073,7 +1056,7 @@ func TestDefaultAccountManager_ExternalCache(t *testing.T) {
 		permissionsManager: permissionsManager,
 	}
 
-	cacheStore, err := nbcache.NewStore(context.Background(), nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval, nbcache.DefaultIDPCacheOpenConn)
+	cacheStore, err := nbcache.NewStore(context.Background(), nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
 	assert.NoError(t, err)
 	am.externalCacheManager = nbcache.NewUserDataCache(cacheStore)
 	am.cacheManager = nbcache.NewAccountUserDataCache(am.loadAccount, cacheStore)
@@ -1429,7 +1412,7 @@ func TestUserAccountPeersUpdate(t *testing.T) {
 	t.Run("deleting user with no linked peers", func(t *testing.T) {
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
+			peerShouldNotReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 

relay/cmd/root.go

@@ -160,8 +160,7 @@ func execute(cmd *cobra.Command, args []string) error {
 		log.Debugf("failed to create relay server: %v", err)
 		return fmt.Errorf("failed to create relay server: %v", err)
 	}
-	instanceURL := srv.InstanceURL()
-	log.Infof("server will be available on: %s", instanceURL.String())
+	log.Infof("server will be available on: %s", srv.InstanceURL())
 	wg.Add(1)
 	go func() {
 		defer wg.Done()

relay/healthcheck/healthcheck.go

@@ -6,14 +6,14 @@ import (
 	"errors"
 	"net"
 	"net/http"
-	"net/url"
 	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/relay/protocol"
-	"github.com/netbirdio/netbird/relay/server"
+	"github.com/netbirdio/netbird/relay/server/listener/quic"
+	"github.com/netbirdio/netbird/relay/server/listener/ws"
 )
 
 const (
@@ -27,7 +27,7 @@ const (
 
 type ServiceChecker interface {
 	ListenerProtocols() []protocol.Protocol
-	InstanceURL() url.URL
+	ListenAddress() string
 }
 
 type HealthStatus struct {
@@ -135,11 +135,7 @@ func (s *Server) getHealthStatus(ctx context.Context) (*HealthStatus, bool) {
 	}
 	status.Listeners = listeners
 
-	if s.config.ServiceChecker.InstanceURL().Scheme != server.SchemeRELS {
-		status.CertificateValid = false
-	}
-
-	if ok := s.validateConnection(ctx); !ok {
+	if ok := s.validateCertificate(ctx); !ok {
 		status.Status = statusUnhealthy
 		status.CertificateValid = false
 		healthy = false
@@ -156,13 +152,32 @@ func (s *Server) validateListeners() ([]protocol.Protocol, bool) {
 	return listeners, true
 }
 
-func (s *Server) validateConnection(ctx context.Context) bool {
-	addr := s.config.ServiceChecker.InstanceURL()
-	if err := dialWS(ctx, addr); err != nil {
-		log.Errorf("failed to dial WebSocket listener at %s: %v", addr.String(), err)
+func (s *Server) validateCertificate(ctx context.Context) bool {
+	listenAddress := s.config.ServiceChecker.ListenAddress()
+	if listenAddress == "" {
+		log.Warn("listen address is empty")
 		return false
 	}
 
+	dAddr := dialAddress(listenAddress)
+
+	for _, proto := range s.config.ServiceChecker.ListenerProtocols() {
+		switch proto {
+		case ws.Proto:
+			if err := dialWS(ctx, dAddr); err != nil {
+				log.Errorf("failed to dial WebSocket listener: %v", err)
+				return false
+			}
+		case quic.Proto:
+			if err := dialQUIC(ctx, dAddr); err != nil {
+				log.Errorf("failed to dial QUIC listener: %v", err)
+				return false
+			}
+		default:
+			log.Warnf("unknown protocol for healthcheck: %s", proto)
+			return false
+		}
+	}
 	return true
 }
 
@@ -172,9 +187,8 @@ func dialAddress(listenAddress string) string {
 		return listenAddress // fallback, might be invalid for dialing
 	}
 
-	// When listening on all interfaces, show localhost for better readability
 	if host == "" || host == "::" || host == "0.0.0.0" {
-		host = "localhost"
+		host = "0.0.0.0"
 	}
 
 	return net.JoinHostPort(host, port)

relay/healthcheck/quic.go

@@ -0,0 +1,31 @@
+package healthcheck
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"time"
+
+	"github.com/quic-go/quic-go"
+
+	tlsnb "github.com/netbirdio/netbird/shared/relay/tls"
+)
+
+func dialQUIC(ctx context.Context, address string) error {
+	tlsConfig := &tls.Config{
+		InsecureSkipVerify: false, // Keep certificate validation enabled
+		NextProtos:         []string{tlsnb.NBalpn},
+	}
+
+	conn, err := quic.DialAddr(ctx, address, tlsConfig, &quic.Config{
+		MaxIdleTimeout:  30 * time.Second,
+		KeepAlivePeriod: 10 * time.Second,
+		EnableDatagrams: true,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to connect to QUIC server: %w", err)
+	}
+
+	_ = conn.CloseWithError(0, "availability check complete")
+	return nil
+}

relay/healthcheck/ws.go

@@ -3,27 +3,19 @@ package healthcheck
 import (
 	"context"
 	"fmt"
-	"net/url"
 
 	"github.com/coder/websocket"
 
-	"github.com/netbirdio/netbird/relay/server"
 	"github.com/netbirdio/netbird/shared/relay"
 )
 
-func dialWS(ctx context.Context, address url.URL) error {
-	scheme := "ws"
-	if address.Scheme == server.SchemeRELS {
-		scheme = "wss"
-	}
-	wsURL := fmt.Sprintf("%s://%s%s", scheme, address.Host, relay.WebSocketURLPath)
+func dialWS(ctx context.Context, address string) error {
+	url := fmt.Sprintf("wss://%s%s", address, relay.WebSocketURLPath)
 
-	conn, resp, err := websocket.Dial(ctx, wsURL, nil)
+	conn, resp, err := websocket.Dial(ctx, url, nil)
 	if resp != nil {
 		defer func() {
-			if resp.Body != nil {
-				_ = resp.Body.Close()
-			}
+			_ = resp.Body.Close()
 		}()
 
 	}

relay/server/relay.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"net/url"
 	"sync"
 	"time"
 
@@ -23,7 +22,7 @@ type Config struct {
 	TLSSupport     bool
 	AuthValidator  Validator
 
-	instanceURL url.URL
+	instanceURL string
 }
 
 func (c *Config) validate() error {
@@ -38,7 +37,7 @@ func (c *Config) validate() error {
 	if err != nil {
 		return fmt.Errorf("invalid url: %v", err)
 	}
-	c.instanceURL = *instanceURL
+	c.instanceURL = instanceURL
 
 	if c.AuthValidator == nil {
 		return fmt.Errorf("auth validator is required")
@@ -52,11 +51,10 @@ type Relay struct {
 	metricsCancel context.CancelFunc
 	validator     Validator
 
-	store          *store.Store
-	notifier       *store.PeerNotifier
-	instanceURL    url.URL
-	exposedAddress string
-	preparedMsg    *preparedMsg
+	store       *store.Store
+	notifier    *store.PeerNotifier
+	instanceURL string
+	preparedMsg *preparedMsg
 
 	closed  bool
 	closeMu sync.RWMutex
@@ -89,16 +87,15 @@ func NewRelay(config Config) (*Relay, error) {
 	}
 
 	r := &Relay{
-		metrics:        m,
-		metricsCancel:  metricsCancel,
-		validator:      config.AuthValidator,
-		instanceURL:    config.instanceURL,
-		exposedAddress: config.ExposedAddress,
-		store:          store.NewStore(),
-		notifier:       store.NewPeerNotifier(),
+		metrics:       m,
+		metricsCancel: metricsCancel,
+		validator:     config.AuthValidator,
+		instanceURL:   config.instanceURL,
+		store:         store.NewStore(),
+		notifier:      store.NewPeerNotifier(),
 	}
 
-	r.preparedMsg, err = newPreparedMsg(r.instanceURL.String())
+	r.preparedMsg, err = newPreparedMsg(r.instanceURL)
 	if err != nil {
 		metricsCancel()
 		return nil, fmt.Errorf("prepare message: %v", err)
@@ -178,6 +175,6 @@ func (r *Relay) Shutdown(ctx context.Context) {
 }
 
 // InstanceURL returns the instance URL of the relay server
-func (r *Relay) InstanceURL() url.URL {
+func (r *Relay) InstanceURL() string {
 	return r.instanceURL
 }

relay/server/relay_test.go

@@ -13,7 +13,7 @@ func TestGetInstanceURL(t *testing.T) {
 		{"Valid address with TLS", "example.com", true, "rels://example.com", false},
 		{"Valid address without TLS", "example.com", false, "rel://example.com", false},
 		{"Valid address with scheme", "rel://example.com", false, "rel://example.com", false},
-		{"Invalid address with non TLS scheme and TLS true", "rel://example.com", true, "", true},
+		{"Valid address with non TLS scheme and TLS true", "rel://example.com", true, "rel://example.com", false},
 		{"Valid address with TLS scheme", "rels://example.com", true, "rels://example.com", false},
 		{"Valid address with TLS scheme and TLS false", "rels://example.com", false, "rels://example.com", false},
 		{"Valid address with TLS scheme and custom port", "rels://example.com:9300", true, "rels://example.com:9300", false},
@@ -28,11 +28,8 @@ func TestGetInstanceURL(t *testing.T) {
 			if (err != nil) != tt.expectError {
 				t.Errorf("expected error: %v, got: %v", tt.expectError, err)
 			}
-			if !tt.expectError && url != nil && url.String() != tt.expectedURL {
-				t.Errorf("expected URL: %s, got: %s", tt.expectedURL, url.String())
-			}
-			if tt.expectError && url != nil {
-				t.Errorf("expected nil URL on error, got: %s", url.String())
+			if url != tt.expectedURL {
+				t.Errorf("expected URL: %s, got: %s", tt.expectedURL, url)
 			}
 		})
 	}

relay/server/server.go

@@ -3,7 +3,6 @@ package server
 import (
 	"context"
 	"crypto/tls"
-	"net/url"
 	"sync"
 
 	"github.com/hashicorp/go-multierror"
@@ -29,6 +28,8 @@ type ListenerConfig struct {
 // It is the gate between the WebSocket listener and the Relay server logic.
 // In a new HTTP connection, the server will accept the connection and pass it to the Relay server via the Accept method.
 type Server struct {
+	listenAddr string
+
 	relay       *Relay
 	listeners   []listener.Listener
 	listenerMux sync.Mutex
@@ -40,7 +41,7 @@ type Server struct {
 //
 //	config: A Config struct containing the necessary configuration:
 //	  - Meter: An OpenTelemetry metric.Meter used for recording metrics. If nil, a default no-op meter is used.
-//	  - InstanceURL: The public address (in domain:port format) used as the server's instance URL. Required.
+//	  - ExposedAddress: The public address (in domain:port format) used as the server's instance URL. Required.
 //	  - TLSSupport: A boolean indicating whether TLS is enabled for the server.
 //	  - AuthValidator: A Validator used to authenticate peers. Required.
 //
@@ -61,6 +62,8 @@ func NewServer(config Config) (*Server, error) {
 
 // Listen starts the relay server.
 func (r *Server) Listen(cfg ListenerConfig) error {
+	r.listenAddr = cfg.Address
+
 	wSListener := &ws.Listener{
 		Address:   cfg.Address,
 		TLSConfig: cfg.TLSConfig,
@@ -120,6 +123,11 @@ func (r *Server) Shutdown(ctx context.Context) error {
 	return nberrors.FormatErrorOrNil(multiErr)
 }
 
+// InstanceURL returns the instance URL of the relay server.
+func (r *Server) InstanceURL() string {
+	return r.relay.instanceURL
+}
+
 func (r *Server) ListenerProtocols() []protocol.Protocol {
 	result := make([]protocol.Protocol, 0)
 
@@ -131,6 +139,6 @@ func (r *Server) ListenerProtocols() []protocol.Protocol {
 	return result
 }
 
-func (r *Server) InstanceURL() url.URL {
-	return r.relay.InstanceURL()
+func (r *Server) ListenAddress() string {
+	return r.listenAddr
 }

relay/server/url.go

@@ -6,14 +6,9 @@ import (
 	"strings"
 )
 
-const (
-	SchemeREL  = "rel"
-	SchemeRELS = "rels"
-)
-
 // getInstanceURL checks if user supplied a URL scheme otherwise adds to the
 // provided address according to TLS definition and parses the address before returning it
-func getInstanceURL(exposedAddress string, tlsSupported bool) (*url.URL, error) {
+func getInstanceURL(exposedAddress string, tlsSupported bool) (string, error) {
 	addr := exposedAddress
 	split := strings.Split(exposedAddress, "://")
 	switch {
@@ -22,22 +17,17 @@ func getInstanceURL(exposedAddress string, tlsSupported bool) (*url.URL, error)
 	case len(split) == 1 && !tlsSupported:
 		addr = "rel://" + exposedAddress
 	case len(split) > 2:
-		return nil, fmt.Errorf("invalid exposed address: %s", exposedAddress)
+		return "", fmt.Errorf("invalid exposed address: %s", exposedAddress)
 	}
 
 	parsedURL, err := url.ParseRequestURI(addr)
 	if err != nil {
-		return nil, fmt.Errorf("invalid exposed address: %v", err)
+		return "", fmt.Errorf("invalid exposed address: %v", err)
 	}
 
-	if parsedURL.Scheme != SchemeREL && parsedURL.Scheme != SchemeRELS {
-		return nil, fmt.Errorf("invalid scheme: %s", parsedURL.Scheme)
+	if parsedURL.Scheme != "rel" && parsedURL.Scheme != "rels" {
+		return "", fmt.Errorf("invalid scheme: %s", parsedURL.Scheme)
 	}
 
-	// Validate scheme matches TLS configuration
-	if tlsSupported && parsedURL.Scheme == SchemeREL {
-		return nil, fmt.Errorf("non-TLS scheme '%s' provided but TLS is supported", SchemeREL)
-	}
-
-	return parsedURL, nil
+	return parsedURL.String(), nil
 }

shared/management/client/client.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"io"
 
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/proto"
@@ -14,13 +12,13 @@ import (
 type Client interface {
 	io.Closer
 	Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
-	GetServerPublicKey() (*wgtypes.Key, error)
-	Register(serverKey wgtypes.Key, setupKey string, jwtToken string, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
-	Login(serverKey wgtypes.Key, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
-	GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error)
-	GetPKCEAuthorizationFlow(serverKey wgtypes.Key) (*proto.PKCEAuthorizationFlow, error)
-	GetNetworkMap(sysInfo *system.Info) (*proto.NetworkMap, error)
-	IsHealthy() bool
-	SyncMeta(sysInfo *system.Info) error
-	Logout() error
+	Register(ctx context.Context, setupKey string, jwtToken string, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) error
+	Login(ctx context.Context, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
+	GetDeviceAuthorizationFlow(ctx context.Context) (*proto.DeviceAuthorizationFlow, error)
+	GetPKCEAuthorizationFlow(ctx context.Context) (*proto.PKCEAuthorizationFlow, error)
+	GetNetworkMap(ctx context.Context, sysInfo *system.Info) (*proto.NetworkMap, error)
+	IsHealthy(ctx context.Context) bool
+	HealthCheck(ctx context.Context) error
+	SyncMeta(ctx context.Context, sysInfo *system.Info) error
+	Logout(ctx context.Context) error
 }

shared/management/client/client_test.go

@@ -21,8 +21,6 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
-	"github.com/netbirdio/netbird/management/internals/modules/peers"
-	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 
 	"github.com/netbirdio/netbird/client/system"
@@ -33,6 +31,8 @@ import (
 	"github.com/netbirdio/netbird/management/server/groups"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	"github.com/netbirdio/netbird/management/server/mock_server"
+	"github.com/netbirdio/netbird/management/server/peers"
+	"github.com/netbirdio/netbird/management/server/peers/ephemeral/manager"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -117,7 +117,7 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 	ctx := context.Background()
 	updateManager := update_channel.NewPeersUpdateManager(metrics)
 	requestBuffer := mgmt.NewAccountRequestBuffer(ctx, store)
-	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, mgmt.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), manager.NewEphemeralManager(store, peersManger), config)
+	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, mgmt.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", port_forwarding.NewControllerMock(), config)
 	accountManager, err := mgmt.BuildManager(context.Background(), config, store, networkMapController, nil, "", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock, false)
 	if err != nil {
 		t.Fatal(err)
@@ -125,11 +125,8 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 
 	groupsManager := groups.NewManagerMock()
 
-	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
-	if err != nil {
-		t.Fatal(err)
-	}
-	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, secretsManager, nil, nil, mgmt.MockIntegratedValidator{}, networkMapController)
+	secretsManager := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
+	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, updateManager, secretsManager, nil, &manager.EphemeralManager{}, nil, mgmt.MockIntegratedValidator{}, networkMapController)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -196,12 +193,12 @@ func TestClient_GetServerPublicKey(t *testing.T) {
 	s, listener := startManagement(t)
 	defer closeManagementSilently(s, listener)
 
-	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
-	if err != nil {
+	client := NewClient(listener.Addr().String(), testKey, false)
+	if err := client.Connect(ctx); err != nil {
 		t.Fatal(err)
 	}
 
-	key, err := client.GetServerPublicKey()
+	key, err := client.getServerPublicKey(ctx)
 	if err != nil {
 		t.Error("couldn't retrieve management public key")
 	}
@@ -219,16 +216,12 @@ func TestClient_LoginUnregistered_ShouldThrow_401(t *testing.T) {
 	s, listener := startManagement(t)
 	defer closeManagementSilently(s, listener)
 
-	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
-	if err != nil {
-		t.Fatal(err)
-	}
-	key, err := client.GetServerPublicKey()
-	if err != nil {
+	client := NewClient(listener.Addr().String(), testKey, false)
+	if err := client.Connect(ctx); err != nil {
 		t.Fatal(err)
 	}
 	sysInfo := system.GetInfo(context.TODO())
-	_, err = client.Login(*key, sysInfo, nil, nil)
+	_, err = client.Login(ctx, sysInfo, nil, nil)
 	if err == nil {
 		t.Error("expecting err on unregistered login, got nil")
 	}
@@ -246,24 +239,16 @@ func TestClient_LoginRegistered(t *testing.T) {
 	s, listener := startManagement(t)
 	defer closeManagementSilently(s, listener)
 
-	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
-	if err != nil {
+	client := NewClient(listener.Addr().String(), testKey, false)
+	if err := client.Connect(ctx); err != nil {
 		t.Fatal(err)
 	}
 
-	key, err := client.GetServerPublicKey()
-	if err != nil {
-		t.Error(err)
-	}
 	info := system.GetInfo(context.TODO())
-	resp, err := client.Register(*key, ValidKey, "", info, nil, nil)
+	err = client.Register(ctx, ValidKey, "", info, nil, nil)
 	if err != nil {
 		t.Error(err)
 	}
-
-	if resp == nil {
-		t.Error("expecting non nil response, got nil")
-	}
 }
 
 func TestClient_Sync(t *testing.T) {
@@ -275,18 +260,13 @@ func TestClient_Sync(t *testing.T) {
 	s, listener := startManagement(t)
 	defer closeManagementSilently(s, listener)
 
-	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
-	if err != nil {
+	client := NewClient(listener.Addr().String(), testKey, false)
+	if err := client.Connect(ctx); err != nil {
 		t.Fatal(err)
 	}
 
-	serverKey, err := client.GetServerPublicKey()
-	if err != nil {
-		t.Error(err)
-	}
-
 	info := system.GetInfo(context.TODO())
-	_, err = client.Register(*serverKey, ValidKey, "", info, nil, nil)
+	err = client.Register(ctx, ValidKey, "", info, nil, nil)
 	if err != nil {
 		t.Error(err)
 	}
@@ -296,13 +276,14 @@ func TestClient_Sync(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	remoteClient, err := NewClient(context.TODO(), listener.Addr().String(), remoteKey, false)
-	if err != nil {
+	remoteClient := NewClient(listener.Addr().String(), remoteKey, false)
+	remoteCtx := context.TODO()
+	if err := remoteClient.Connect(remoteCtx); err != nil {
 		t.Fatal(err)
 	}
 
 	info = system.GetInfo(context.TODO())
-	_, err = remoteClient.Register(*serverKey, ValidKey, "", info, nil, nil)
+	err = remoteClient.Register(remoteCtx, ValidKey, "", info, nil, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -357,14 +338,9 @@ func Test_SystemMetaDataFromClient(t *testing.T) {
 	serverAddr := lis.Addr().String()
 	ctx := context.Background()
 
-	testClient, err := NewClient(ctx, serverAddr, testKey, false)
-	if err != nil {
-		t.Fatalf("error while creating testClient: %v", err)
-	}
-
-	key, err := testClient.GetServerPublicKey()
-	if err != nil {
-		t.Fatalf("error while getting server public key from testclient, %v", err)
+	testClient := NewClient(serverAddr, testKey, false)
+	if err := testClient.Connect(ctx); err != nil {
+		t.Fatalf("error while connecting testClient: %v", err)
 	}
 
 	var actualMeta *mgmtProto.PeerSystemMeta
@@ -403,7 +379,7 @@ func Test_SystemMetaDataFromClient(t *testing.T) {
 	}
 
 	info := system.GetInfo(context.TODO())
-	_, err = testClient.Register(*key, ValidKey, "", info, nil, nil)
+	err = testClient.Register(ctx, ValidKey, "", info, nil, nil)
 	if err != nil {
 		t.Errorf("error while trying to register client: %v", err)
 	}
@@ -492,9 +468,9 @@ func Test_GetDeviceAuthorizationFlow(t *testing.T) {
 	serverAddr := lis.Addr().String()
 	ctx := context.Background()
 
-	client, err := NewClient(ctx, serverAddr, testKey, false)
-	if err != nil {
-		t.Fatalf("error while creating testClient: %v", err)
+	client := NewClient(serverAddr, testKey, false)
+	if err := client.Connect(ctx); err != nil {
+		t.Fatalf("error while connecting testClient: %v", err)
 	}
 
 	expectedFlowInfo := &mgmtProto.DeviceAuthorizationFlow{
@@ -515,7 +491,7 @@ func Test_GetDeviceAuthorizationFlow(t *testing.T) {
 		}, nil
 	}
 
-	flowInfo, err := client.GetDeviceAuthorizationFlow(serverKey)
+	flowInfo, err := client.GetDeviceAuthorizationFlow(ctx)
 	if err != nil {
 		t.Error("error while retrieving device auth flow information")
 	}
@@ -536,9 +512,9 @@ func Test_GetPKCEAuthorizationFlow(t *testing.T) {
 	serverAddr := lis.Addr().String()
 	ctx := context.Background()
 
-	client, err := NewClient(ctx, serverAddr, testKey, false)
-	if err != nil {
-		t.Fatalf("error while creating testClient: %v", err)
+	client := NewClient(serverAddr, testKey, false)
+	if err := client.Connect(ctx); err != nil {
+		t.Fatalf("error while connecting testClient: %v", err)
 	}
 
 	expectedFlowInfo := &mgmtProto.PKCEAuthorizationFlow{
@@ -561,7 +537,7 @@ func Test_GetPKCEAuthorizationFlow(t *testing.T) {
 		}, nil
 	}
 
-	flowInfo, err := client.GetPKCEAuthorizationFlow(serverKey)
+	flowInfo, err := client.GetPKCEAuthorizationFlow(ctx)
 	if err != nil {
 		t.Error("error while retrieving pkce auth flow information")
 	}

shared/management/client/common/types.go

@@ -1,20 +1,19 @@
 package common
 
-// LoginFlag introduces additional login flags to the PKCE authorization request.
-//
-// # Config Values
-//
-//	| Value | Flag                 | OAuth Parameters                        |
-//	|-------|----------------------|-----------------------------------------|
-//	| 0     | LoginFlagPromptLogin | prompt=login                            |
-//	| 1     | LoginFlagMaxAge0     | max_age=0                               |
+// LoginFlag introduces additional login flags to the PKCE authorization request
 type LoginFlag uint8
 
 const (
-	// LoginFlagPromptLogin adds prompt=login to the authorization request
-	LoginFlagPromptLogin LoginFlag = iota
+	// LoginFlagPrompt adds prompt=login to the authorization request
+	LoginFlagPrompt LoginFlag = iota
 	// LoginFlagMaxAge0 adds max_age=0 to the authorization request
 	LoginFlagMaxAge0
-	// LoginFlagNone disables all login flags
-	LoginFlagNone
 )
+
+func (l LoginFlag) IsPromptLogin() bool {
+	return l == LoginFlagPrompt
+}
+
+func (l LoginFlag) IsMaxAge0Login() bool {
+	return l == LoginFlagMaxAge0
+}

shared/management/client/grpc.go

@@ -25,6 +25,24 @@ import (
 	"github.com/netbirdio/netbird/util/wsproxy"
 )
 
+// Custom management client errors that abstract away gRPC error codes
+var (
+	// ErrPermissionDenied is returned when the server denies access to a resource
+	ErrPermissionDenied = errors.New("permission denied")
+
+	// ErrInvalidArgument is returned when the request contains invalid arguments
+	ErrInvalidArgument = errors.New("invalid argument")
+
+	// ErrUnauthenticated is returned when authentication is required
+	ErrUnauthenticated = errors.New("unauthenticated")
+
+	// ErrNotFound is returned when the requested resource is not found
+	ErrNotFound = errors.New("not found")
+
+	// ErrUnimplemented is returned when the operation is not implemented
+	ErrUnimplemented = errors.New("not implemented")
+)
+
 const ConnectTimeout = 10 * time.Second
 
 const (
@@ -41,40 +59,67 @@ type ConnStateNotifier interface {
 type GrpcClient struct {
 	key                   wgtypes.Key
 	realClient            proto.ManagementServiceClient
-	ctx                   context.Context
 	conn                  *grpc.ClientConn
 	connStateCallback     ConnStateNotifier
 	connStateCallbackLock sync.RWMutex
+	addr                  string
+	tlsEnabled            bool
+	reconnectMutex        sync.Mutex
 }
 
 // NewClient creates a new client to Management service
-func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) (*GrpcClient, error) {
+// The client is not connected after creation - call Connect to establish the connection
+func NewClient(addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) *GrpcClient {
+	return &GrpcClient{
+		key:                   ourPrivateKey,
+		addr:                  addr,
+		tlsEnabled:            tlsEnabled,
+		connStateCallbackLock: sync.RWMutex{},
+		reconnectMutex:        sync.Mutex{},
+	}
+}
+
+// Connect establishes a connection to the Management Service with retry logic
+// Retries connection attempts with exponential backoff on failure
+func (c *GrpcClient) Connect(ctx context.Context) error {
 	var conn *grpc.ClientConn
 
 	operation := func() error {
 		var err error
-		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.ManagementComponent)
+		conn, err = nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent)
 		if err != nil {
-			return fmt.Errorf("create connection: %w", err)
+			log.Warnf("failed to connect to Management Service: %v", err)
+			return err
 		}
 		return nil
 	}
 
-	err := backoff.Retry(operation, nbgrpc.Backoff(ctx))
-	if err != nil {
-		log.Errorf("failed creating connection to Management Service: %v", err)
-		return nil, err
+	if err := backoff.Retry(operation, defaultBackoff(ctx)); err != nil {
+		log.Errorf("failed creating connection to Management Service after retries: %v", err)
+		return fmt.Errorf("create connection: %w", err)
 	}
 
-	realClient := proto.NewManagementServiceClient(conn)
+	c.conn = conn
+	c.realClient = proto.NewManagementServiceClient(conn)
 
-	return &GrpcClient{
-		key:                   ourPrivateKey,
-		realClient:            realClient,
-		ctx:                   ctx,
-		conn:                  conn,
-		connStateCallbackLock: sync.RWMutex{},
-	}, nil
+	log.Infof("connected to the Management Service at %s", c.addr)
+	return nil
+}
+
+// ConnectWithoutRetry establishes a connection to the Management Service without retry logic
+// Performs a single connection attempt - callers should implement their own retry logic if needed
+func (c *GrpcClient) ConnectWithoutRetry(ctx context.Context) error {
+	conn, err := nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent)
+	if err != nil {
+		log.Warnf("failed to connect to Management Service: %v", err)
+		return fmt.Errorf("create connection: %w", err)
+	}
+
+	c.conn = conn
+	c.realClient = proto.NewManagementServiceClient(conn)
+
+	log.Debugf("connected to the Management Service at %s", c.addr)
+	return nil
 }
 
 // Close closes connection to the Management Service
@@ -89,19 +134,6 @@ func (c *GrpcClient) SetConnStateListener(notifier ConnStateNotifier) {
 	c.connStateCallback = notifier
 }
 
-// defaultBackoff is a basic backoff mechanism for general issues
-func defaultBackoff(ctx context.Context) backoff.BackOff {
-	return backoff.WithContext(&backoff.ExponentialBackOff{
-		InitialInterval:     800 * time.Millisecond,
-		RandomizationFactor: 1,
-		Multiplier:          1.7,
-		MaxInterval:         10 * time.Second,
-		MaxElapsedTime:      3 * 30 * 24 * time.Hour, // 3 months
-		Stop:                backoff.Stop,
-		Clock:               backoff.SystemClock,
-	}, ctx)
-}
-
 // ready indicates whether the client is okay and ready to be used
 // for now it just checks whether gRPC connection to the service is ready
 func (c *GrpcClient) ready() bool {
@@ -122,7 +154,7 @@ func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler
 			return fmt.Errorf("connection to management is not ready and in %s state", connState)
 		}
 
-		serverPubKey, err := c.GetServerPublicKey()
+		serverPubKey, err := c.getServerPublicKey(ctx)
 		if err != nil {
 			log.Debugf(errMsgMgmtPublicKey, err)
 			return err
@@ -177,15 +209,13 @@ func (c *GrpcClient) handleStream(ctx context.Context, serverPubKey wgtypes.Key,
 }
 
 // GetNetworkMap return with the network map
-func (c *GrpcClient) GetNetworkMap(sysInfo *system.Info) (*proto.NetworkMap, error) {
-	serverPubKey, err := c.GetServerPublicKey()
+func (c *GrpcClient) GetNetworkMap(ctx context.Context, sysInfo *system.Info) (*proto.NetworkMap, error) {
+	serverPubKey, err := c.getServerPublicKey(ctx)
 	if err != nil {
 		log.Debugf("failed getting Management Service public key: %s", err)
 		return nil, err
 	}
 
-	ctx, cancelStream := context.WithCancel(c.ctx)
-	defer cancelStream()
 	stream, err := c.connectToStream(ctx, *serverPubKey, sysInfo)
 	if err != nil {
 		log.Debugf("failed to open Management Service stream: %s", err)
@@ -264,30 +294,32 @@ func (c *GrpcClient) receiveEvents(stream proto.ManagementService_SyncClient, se
 	}
 }
 
-// GetServerPublicKey returns server's WireGuard public key (used later for encrypting messages sent to the server)
-func (c *GrpcClient) GetServerPublicKey() (*wgtypes.Key, error) {
+// getServerPublicKey returns server's WireGuard public key (used later for encrypting messages sent to the server)
+// This is a simple operation without retry logic - callers should handle retries at the operation level
+func (c *GrpcClient) getServerPublicKey(ctx context.Context) (*wgtypes.Key, error) {
 	if !c.ready() {
 		return nil, errors.New(errMsgNoMgmtConnection)
 	}
 
-	mgmCtx, cancel := context.WithTimeout(c.ctx, 5*time.Second)
+	mgmCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
+
 	resp, err := c.realClient.GetServerKey(mgmCtx, &proto.Empty{})
 	if err != nil {
 		log.Errorf("failed while getting Management Service public key: %v", err)
 		return nil, fmt.Errorf("failed while getting Management Service public key")
 	}
 
-	serverKey, err := wgtypes.ParseKey(resp.Key)
+	key, err := wgtypes.ParseKey(resp.Key)
 	if err != nil {
 		return nil, err
 	}
 
-	return &serverKey, nil
+	return &key, nil
 }
 
 // IsHealthy probes the gRPC connection and returns false on errors
-func (c *GrpcClient) IsHealthy() bool {
+func (c *GrpcClient) IsHealthy(ctx context.Context) bool {
 	switch c.conn.GetState() {
 	case connectivity.TransientFailure:
 		return false
@@ -299,10 +331,10 @@ func (c *GrpcClient) IsHealthy() bool {
 	case connectivity.Ready:
 	}
 
-	ctx, cancel := context.WithTimeout(c.ctx, 1*time.Second)
+	healthCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
 	defer cancel()
 
-	_, err := c.realClient.GetServerKey(ctx, &proto.Empty{})
+	_, err := c.realClient.GetServerKey(healthCtx, &proto.Empty{})
 	if err != nil {
 		c.notifyDisconnected(err)
 		log.Warnf("health check returned: %s", err)
@@ -312,12 +344,26 @@ func (c *GrpcClient) IsHealthy() bool {
 	return true
 }
 
-func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*proto.LoginResponse, error) {
+// HealthCheck verifies connectivity to the management server
+// Returns an error if the server is not reachable
+// Internally uses getServerPublicKey to verify the connection
+func (c *GrpcClient) HealthCheck(ctx context.Context) error {
+	_, err := c.getServerPublicKey(ctx)
+	return err
+}
+
+func (c *GrpcClient) login(ctx context.Context, req *proto.LoginRequest) (*proto.LoginResponse, error) {
 	if !c.ready() {
 		return nil, errors.New(errMsgNoMgmtConnection)
 	}
 
-	loginReq, err := encryption.EncryptMessage(serverKey, c.key, req)
+	serverKey, err := c.getServerPublicKey(ctx)
+	if err != nil {
+		log.Debugf(errMsgMgmtPublicKey, err)
+		return nil, err
+	}
+
+	loginReq, err := encryption.EncryptMessage(*serverKey, c.key, req)
 	if err != nil {
 		log.Errorf("failed to encrypt message: %s", err)
 		return nil, err
@@ -325,7 +371,7 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 
 	var resp *proto.EncryptedMessage
 	operation := func() error {
-		mgmCtx, cancel := context.WithTimeout(context.Background(), ConnectTimeout)
+		mgmCtx, cancel := context.WithTimeout(ctx, ConnectTimeout)
 		defer cancel()
 
 		var err error
@@ -344,14 +390,14 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 		return nil
 	}
 
-	err = backoff.Retry(operation, nbgrpc.Backoff(c.ctx))
+	err = backoff.Retry(operation, nbgrpc.Backoff(ctx))
 	if err != nil {
 		log.Errorf("failed to login to Management Service: %v", err)
 		return nil, err
 	}
 
 	loginResp := &proto.LoginResponse{}
-	err = encryption.DecryptMessage(serverKey, c.key, resp.Body, loginResp)
+	err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, loginResp)
 	if err != nil {
 		log.Errorf("failed to decrypt login response: %s", err)
 		return nil, err
@@ -363,99 +409,135 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 // Register registers peer on Management Server. It actually calls a Login endpoint with a provided setup key
 // Takes care of encrypting and decrypting messages.
 // This method will also collect system info and send it with the request (e.g. hostname, os, etc)
-func (c *GrpcClient) Register(serverKey wgtypes.Key, setupKey string, jwtToken string, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
+// Returns custom errors: ErrPermissionDenied, ErrInvalidArgument, ErrUnauthenticated
+func (c *GrpcClient) Register(ctx context.Context, setupKey string, jwtToken string, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) error {
 	keys := &proto.PeerKeys{
 		SshPubKey: pubSSHKey,
 		WgPubKey:  []byte(c.key.PublicKey().String()),
 	}
-	return c.login(serverKey, &proto.LoginRequest{SetupKey: setupKey, Meta: infoToMetaData(sysInfo), JwtToken: jwtToken, PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()})
+	_, err := c.login(ctx, &proto.LoginRequest{SetupKey: setupKey, Meta: infoToMetaData(sysInfo), JwtToken: jwtToken, PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()})
+	return wrapGRPCError(err)
 }
 
 // Login attempts login to Management Server. Takes care of encrypting and decrypting messages.
-func (c *GrpcClient) Login(serverKey wgtypes.Key, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
+// Returns custom errors: ErrPermissionDenied, ErrInvalidArgument, ErrUnauthenticated
+func (c *GrpcClient) Login(ctx context.Context, sysInfo *system.Info, pubSSHKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
 	keys := &proto.PeerKeys{
 		SshPubKey: pubSSHKey,
 		WgPubKey:  []byte(c.key.PublicKey().String()),
 	}
-	return c.login(serverKey, &proto.LoginRequest{Meta: infoToMetaData(sysInfo), PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()})
+	resp, err := c.login(ctx, &proto.LoginRequest{Meta: infoToMetaData(sysInfo), PeerKeys: keys, DnsLabels: dnsLabels.ToPunycodeList()})
+	return resp, wrapGRPCError(err)
 }
 
 // GetDeviceAuthorizationFlow returns a device authorization flow information.
 // It also takes care of encrypting and decrypting messages.
-func (c *GrpcClient) GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error) {
-	if !c.ready() {
-		return nil, fmt.Errorf("no connection to management in order to get device authorization flow")
-	}
-	mgmCtx, cancel := context.WithTimeout(c.ctx, time.Second*2)
-	defer cancel()
+// It automatically retries with backoff and reconnection on connection errors.
+// Returns custom errors: ErrNotFound, ErrUnimplemented
+func (c *GrpcClient) GetDeviceAuthorizationFlow(ctx context.Context) (*proto.DeviceAuthorizationFlow, error) {
+	var flowInfoResp *proto.DeviceAuthorizationFlow
 
-	message := &proto.DeviceAuthorizationFlowRequest{}
-	encryptedMSG, err := encryption.EncryptMessage(serverKey, c.key, message)
-	if err != nil {
-		return nil, err
-	}
+	err := c.withRetry(ctx, func() error {
+		if !c.ready() {
+			return fmt.Errorf("no connection to management in order to get device authorization flow")
+		}
 
-	resp, err := c.realClient.GetDeviceAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{
-		WgPubKey: c.key.PublicKey().String(),
-		Body:     encryptedMSG},
-	)
-	if err != nil {
-		return nil, err
-	}
+		serverKey, err := c.getServerPublicKey(ctx)
+		if err != nil {
+			log.Debugf(errMsgMgmtPublicKey, err)
+			return err
+		}
 
-	flowInfoResp := &proto.DeviceAuthorizationFlow{}
-	err = encryption.DecryptMessage(serverKey, c.key, resp.Body, flowInfoResp)
-	if err != nil {
-		errWithMSG := fmt.Errorf("failed to decrypt device authorization flow message: %s", err)
-		log.Error(errWithMSG)
-		return nil, errWithMSG
-	}
+		mgmCtx, cancel := context.WithTimeout(ctx, time.Second*2)
+		defer cancel()
 
-	return flowInfoResp, nil
+		message := &proto.DeviceAuthorizationFlowRequest{}
+		encryptedMSG, err := encryption.EncryptMessage(*serverKey, c.key, message)
+		if err != nil {
+			return err
+		}
+
+		resp, err := c.realClient.GetDeviceAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{
+			WgPubKey: c.key.PublicKey().String(),
+			Body:     encryptedMSG},
+		)
+		if err != nil {
+			return err
+		}
+
+		flowInfo := &proto.DeviceAuthorizationFlow{}
+		err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, flowInfo)
+		if err != nil {
+			errWithMSG := fmt.Errorf("failed to decrypt device authorization flow message: %s", err)
+			log.Error(errWithMSG)
+			return errWithMSG
+		}
+
+		flowInfoResp = flowInfo
+		return nil
+	})
+
+	return flowInfoResp, wrapGRPCError(err)
 }
 
 // GetPKCEAuthorizationFlow returns a pkce authorization flow information.
 // It also takes care of encrypting and decrypting messages.
-func (c *GrpcClient) GetPKCEAuthorizationFlow(serverKey wgtypes.Key) (*proto.PKCEAuthorizationFlow, error) {
-	if !c.ready() {
-		return nil, fmt.Errorf("no connection to management in order to get pkce authorization flow")
-	}
-	mgmCtx, cancel := context.WithTimeout(c.ctx, time.Second*2)
-	defer cancel()
+// It automatically retries with backoff and reconnection on connection errors.
+// Returns custom errors: ErrNotFound, ErrUnimplemented
+func (c *GrpcClient) GetPKCEAuthorizationFlow(ctx context.Context) (*proto.PKCEAuthorizationFlow, error) {
+	var flowInfoResp *proto.PKCEAuthorizationFlow
 
-	message := &proto.PKCEAuthorizationFlowRequest{}
-	encryptedMSG, err := encryption.EncryptMessage(serverKey, c.key, message)
-	if err != nil {
-		return nil, err
-	}
+	err := c.withRetry(ctx, func() error {
+		if !c.ready() {
+			return fmt.Errorf("no connection to management in order to get pkce authorization flow")
+		}
 
-	resp, err := c.realClient.GetPKCEAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{
-		WgPubKey: c.key.PublicKey().String(),
-		Body:     encryptedMSG,
+		serverKey, err := c.getServerPublicKey(ctx)
+		if err != nil {
+			log.Debugf(errMsgMgmtPublicKey, err)
+			return err
+		}
+
+		mgmCtx, cancel := context.WithTimeout(ctx, time.Second*2)
+		defer cancel()
+
+		message := &proto.PKCEAuthorizationFlowRequest{}
+		encryptedMSG, err := encryption.EncryptMessage(*serverKey, c.key, message)
+		if err != nil {
+			return err
+		}
+
+		resp, err := c.realClient.GetPKCEAuthorizationFlow(mgmCtx, &proto.EncryptedMessage{
+			WgPubKey: c.key.PublicKey().String(),
+			Body:     encryptedMSG,
+		})
+		if err != nil {
+			return err
+		}
+
+		flowInfo := &proto.PKCEAuthorizationFlow{}
+		err = encryption.DecryptMessage(*serverKey, c.key, resp.Body, flowInfo)
+		if err != nil {
+			errWithMSG := fmt.Errorf("failed to decrypt pkce authorization flow message: %s", err)
+			log.Error(errWithMSG)
+			return errWithMSG
+		}
+
+		flowInfoResp = flowInfo
+		return nil
 	})
-	if err != nil {
-		return nil, err
-	}
 
-	flowInfoResp := &proto.PKCEAuthorizationFlow{}
-	err = encryption.DecryptMessage(serverKey, c.key, resp.Body, flowInfoResp)
-	if err != nil {
-		errWithMSG := fmt.Errorf("failed to decrypt pkce authorization flow message: %s", err)
-		log.Error(errWithMSG)
-		return nil, errWithMSG
-	}
-
-	return flowInfoResp, nil
+	return flowInfoResp, wrapGRPCError(err)
 }
 
 // SyncMeta sends updated system metadata to the Management Service.
 // It should be used if there is changes on peer posture check after initial sync.
-func (c *GrpcClient) SyncMeta(sysInfo *system.Info) error {
+func (c *GrpcClient) SyncMeta(ctx context.Context, sysInfo *system.Info) error {
 	if !c.ready() {
 		return errors.New(errMsgNoMgmtConnection)
 	}
 
-	serverPubKey, err := c.GetServerPublicKey()
+	serverPubKey, err := c.getServerPublicKey(ctx)
 	if err != nil {
 		log.Debugf(errMsgMgmtPublicKey, err)
 		return err
@@ -467,7 +549,7 @@ func (c *GrpcClient) SyncMeta(sysInfo *system.Info) error {
 		return err
 	}
 
-	mgmCtx, cancel := context.WithTimeout(c.ctx, ConnectTimeout)
+	mgmCtx, cancel := context.WithTimeout(ctx, ConnectTimeout)
 	defer cancel()
 
 	_, err = c.realClient.SyncMeta(mgmCtx, &proto.EncryptedMessage{
@@ -497,13 +579,13 @@ func (c *GrpcClient) notifyConnected() {
 	c.connStateCallback.MarkManagementConnected()
 }
 
-func (c *GrpcClient) Logout() error {
-	serverKey, err := c.GetServerPublicKey()
+func (c *GrpcClient) Logout(ctx context.Context) error {
+	serverKey, err := c.getServerPublicKey(ctx)
 	if err != nil {
 		return fmt.Errorf("get server public key: %w", err)
 	}
 
-	mgmCtx, cancel := context.WithTimeout(c.ctx, time.Second*15)
+	mgmCtx, cancel := context.WithTimeout(ctx, time.Second*15)
 	defer cancel()
 
 	message := &proto.Empty{}
@@ -523,6 +605,156 @@ func (c *GrpcClient) Logout() error {
 	return nil
 }
 
+// reconnect closes the current connection and creates a new one
+func (c *GrpcClient) reconnect(ctx context.Context) error {
+	c.reconnectMutex.Lock()
+	defer c.reconnectMutex.Unlock()
+
+	// Close existing connection
+	if c.conn != nil {
+		if err := c.conn.Close(); err != nil {
+			log.Debugf("error closing old connection: %v", err)
+		}
+	}
+
+	// Create new connection
+	log.Debugf("reconnecting to Management Service %s", c.addr)
+	conn, err := nbgrpc.CreateConnection(ctx, c.addr, c.tlsEnabled, wsproxy.ManagementComponent)
+	if err != nil {
+		log.Errorf("failed reconnecting to Management Service %s: %v", c.addr, err)
+		return fmt.Errorf("reconnect: create connection: %w", err)
+	}
+
+	c.conn = conn
+	c.realClient = proto.NewManagementServiceClient(conn)
+	log.Debugf("successfully reconnected to Management service %s", c.addr)
+	return nil
+}
+
+// withRetry wraps an operation with exponential backoff retry logic
+// It automatically reconnects on connection errors
+func (c *GrpcClient) withRetry(ctx context.Context, operation func() error) error {
+	backoffSettings := &backoff.ExponentialBackOff{
+		InitialInterval:     500 * time.Millisecond,
+		RandomizationFactor: 0.5,
+		Multiplier:          1.5,
+		MaxInterval:         10 * time.Second,
+		MaxElapsedTime:      2 * time.Minute,
+		Stop:                backoff.Stop,
+		Clock:               backoff.SystemClock,
+	}
+	backoffSettings.Reset()
+
+	return backoff.RetryNotify(
+		func() error {
+			err := operation()
+			if err == nil {
+				return nil
+			}
+
+			// If it's a connection error, attempt reconnection
+			if isConnectionError(err) {
+				log.Warnf("connection error detected, attempting reconnection: %v", err)
+				if reconnectErr := c.reconnect(ctx); reconnectErr != nil {
+					log.Errorf("reconnection failed: %v", reconnectErr)
+					return reconnectErr
+				}
+				// Return the original error to trigger retry with the new connection
+				return err
+			}
+
+			// For authentication errors (InvalidArgument, PermissionDenied), don't retry
+			if isAuthenticationError(err) {
+				return backoff.Permanent(err)
+			}
+
+			return err
+		},
+		backoff.WithContext(backoffSettings, ctx),
+		func(err error, duration time.Duration) {
+			log.Warnf("operation failed, retrying in %v: %v", duration, err)
+		},
+	)
+}
+
+// defaultBackoff is a basic backoff mechanism for general issues
+func defaultBackoff(ctx context.Context) backoff.BackOff {
+	return backoff.WithContext(&backoff.ExponentialBackOff{
+		InitialInterval:     800 * time.Millisecond,
+		RandomizationFactor: 1,
+		Multiplier:          1.7,
+		MaxInterval:         10 * time.Second,
+		MaxElapsedTime:      3 * 30 * 24 * time.Hour, // 3 months
+		Stop:                backoff.Stop,
+		Clock:               backoff.SystemClock,
+	}, ctx)
+}
+
+// isConnectionError checks if the error is a connection-related error that should trigger reconnection
+func isConnectionError(err error) bool {
+	if err == nil {
+		return false
+	}
+	s, ok := gstatus.FromError(err)
+	if !ok {
+		return false
+	}
+	// These error codes indicate connection issues
+	return s.Code() == codes.Unavailable ||
+		s.Code() == codes.DeadlineExceeded ||
+		s.Code() == codes.Canceled ||
+		s.Code() == codes.Internal
+}
+
+// isAuthenticationError checks if the error is an authentication-related error that should not be retried
+func isAuthenticationError(err error) bool {
+	if err == nil {
+		return false
+	}
+	s, ok := gstatus.FromError(err)
+	if !ok {
+		return false
+	}
+	return s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied
+}
+
+// wrapGRPCError converts gRPC errors to custom management client errors
+func wrapGRPCError(err error) error {
+	if err == nil {
+		return nil
+	}
+
+	// Check if it's already a custom error
+	if errors.Is(err, ErrPermissionDenied) ||
+		errors.Is(err, ErrInvalidArgument) ||
+		errors.Is(err, ErrUnauthenticated) ||
+		errors.Is(err, ErrNotFound) ||
+		errors.Is(err, ErrUnimplemented) {
+		return err
+	}
+
+	// Convert gRPC status errors to custom errors
+	s, ok := gstatus.FromError(err)
+	if !ok {
+		return err
+	}
+
+	switch s.Code() {
+	case codes.PermissionDenied:
+		return fmt.Errorf("%w: %s", ErrPermissionDenied, s.Message())
+	case codes.InvalidArgument:
+		return fmt.Errorf("%w: %s", ErrInvalidArgument, s.Message())
+	case codes.Unauthenticated:
+		return fmt.Errorf("%w: %s", ErrUnauthenticated, s.Message())
+	case codes.NotFound:
+		return fmt.Errorf("%w: %s", ErrNotFound, s.Message())
+	case codes.Unimplemented:
+		return fmt.Errorf("%w: %s", ErrUnimplemented, s.Message())
+	default:
+		return err
+	}
+}
+
 func infoToMetaData(info *system.Info) *proto.PeerSystemMeta {
 	if info == nil {
 		return nil

shared/management/client/mock.go

@@ -3,8 +3,6 @@ package client
 import (
 	"context"
 
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/proto"
@@ -13,17 +11,21 @@ import (
 type MockClient struct {
 	CloseFunc                      func() error
 	SyncFunc                       func(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
-	GetServerPublicKeyFunc         func() (*wgtypes.Key, error)
-	RegisterFunc                   func(serverKey wgtypes.Key, setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
-	LoginFunc                      func(serverKey wgtypes.Key, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
-	GetDeviceAuthorizationFlowFunc func(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error)
-	GetPKCEAuthorizationFlowFunc   func(serverKey wgtypes.Key) (*proto.PKCEAuthorizationFlow, error)
-	SyncMetaFunc                   func(sysInfo *system.Info) error
-	LogoutFunc                     func() error
+	RegisterFunc                   func(ctx context.Context, setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) error
+	LoginFunc                      func(ctx context.Context, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
+	GetDeviceAuthorizationFlowFunc func(ctx context.Context) (*proto.DeviceAuthorizationFlow, error)
+	GetPKCEAuthorizationFlowFunc   func(ctx context.Context) (*proto.PKCEAuthorizationFlow, error)
+	SyncMetaFunc                   func(ctx context.Context, sysInfo *system.Info) error
+	HealthCheckFunc                func(ctx context.Context) error
+	LogoutFunc                     func(ctx context.Context) error
+	IsHealthyFunc                  func(ctx context.Context) bool
 }
 
-func (m *MockClient) IsHealthy() bool {
-	return true
+func (m *MockClient) IsHealthy(ctx context.Context) bool {
+	if m.IsHealthyFunc == nil {
+		return true
+	}
+	return m.IsHealthyFunc(ctx)
 }
 
 func (m *MockClient) Close() error {
@@ -40,56 +42,56 @@ func (m *MockClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler
 	return m.SyncFunc(ctx, sysInfo, msgHandler)
 }
 
-func (m *MockClient) GetServerPublicKey() (*wgtypes.Key, error) {
-	if m.GetServerPublicKeyFunc == nil {
-		return nil, nil
-	}
-	return m.GetServerPublicKeyFunc()
-}
-
-func (m *MockClient) Register(serverKey wgtypes.Key, setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
+func (m *MockClient) Register(ctx context.Context, setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) error {
 	if m.RegisterFunc == nil {
-		return nil, nil
+		return nil
 	}
-	return m.RegisterFunc(serverKey, setupKey, jwtToken, info, sshKey, dnsLabels)
+	return m.RegisterFunc(ctx, setupKey, jwtToken, info, sshKey, dnsLabels)
 }
 
-func (m *MockClient) Login(serverKey wgtypes.Key, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
+func (m *MockClient) Login(ctx context.Context, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) {
 	if m.LoginFunc == nil {
 		return nil, nil
 	}
-	return m.LoginFunc(serverKey, info, sshKey, dnsLabels)
+	return m.LoginFunc(ctx, info, sshKey, dnsLabels)
 }
 
-func (m *MockClient) GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error) {
+func (m *MockClient) GetDeviceAuthorizationFlow(ctx context.Context) (*proto.DeviceAuthorizationFlow, error) {
 	if m.GetDeviceAuthorizationFlowFunc == nil {
 		return nil, nil
 	}
-	return m.GetDeviceAuthorizationFlowFunc(serverKey)
+	return m.GetDeviceAuthorizationFlowFunc(ctx)
 }
 
-func (m *MockClient) GetPKCEAuthorizationFlow(serverKey wgtypes.Key) (*proto.PKCEAuthorizationFlow, error) {
+func (m *MockClient) GetPKCEAuthorizationFlow(ctx context.Context) (*proto.PKCEAuthorizationFlow, error) {
 	if m.GetPKCEAuthorizationFlowFunc == nil {
 		return nil, nil
 	}
-	return m.GetPKCEAuthorizationFlow(serverKey)
+	return m.GetPKCEAuthorizationFlowFunc(ctx)
 }
 
 // GetNetworkMap mock implementation of GetNetworkMap from mgm.Client interface
-func (m *MockClient) GetNetworkMap(_ *system.Info) (*proto.NetworkMap, error) {
+func (m *MockClient) GetNetworkMap(ctx context.Context, _ *system.Info) (*proto.NetworkMap, error) {
 	return nil, nil
 }
 
-func (m *MockClient) SyncMeta(sysInfo *system.Info) error {
+func (m *MockClient) SyncMeta(ctx context.Context, sysInfo *system.Info) error {
 	if m.SyncMetaFunc == nil {
 		return nil
 	}
-	return m.SyncMetaFunc(sysInfo)
+	return m.SyncMetaFunc(ctx, sysInfo)
 }
 
-func (m *MockClient) Logout() error {
+func (m *MockClient) HealthCheck(ctx context.Context) error {
+	if m.HealthCheckFunc == nil {
+		return nil
+	}
+	return m.HealthCheckFunc(ctx)
+}
+
+func (m *MockClient) Logout(ctx context.Context) error {
 	if m.LogoutFunc == nil {
 		return nil
 	}
-	return m.LogoutFunc()
+	return m.LogoutFunc(ctx)
 }

shared/management/client/rest/groups.go

@@ -4,14 +4,10 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
-	"errors"
 
 	"github.com/netbirdio/netbird/shared/management/http/api"
 )
 
-// ErrGroupNotFound is returned when a group is not found
-var ErrGroupNotFound = errors.New("group not found")
-
 // GroupsAPI APIs for Groups, do not use directly
 type GroupsAPI struct {
 	c *Client
@@ -31,27 +27,6 @@ func (a *GroupsAPI) List(ctx context.Context) ([]api.Group, error) {
 	return ret, err
 }
 
-// GetByName get group by name
-// See more: https://docs.netbird.io/api/resources/groups#list-all-groups
-func (a *GroupsAPI) GetByName(ctx context.Context, groupName string) (*api.Group, error) {
-	params := map[string]string{"name": groupName}
-	resp, err := a.c.NewRequest(ctx, "GET", "/api/groups", nil, params)
-	if err != nil {
-		return nil, err
-	}
-	if resp.Body != nil {
-		defer resp.Body.Close()
-	}
-	ret, err := parseResponse[[]api.Group](resp)
-	if err != nil {
-		return nil, err
-	}
-	if len(ret) == 0 {
-		return nil, ErrGroupNotFound
-	}
-	return &ret[0], nil
-}
-
 // Get get group info
 // See more: https://docs.netbird.io/api/resources/groups#retrieve-a-group
 func (a *GroupsAPI) Get(ctx context.Context, groupID string) (*api.Group, error) {