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/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)
-		})
-}