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

.github/workflows/golang-test-darwin.yml

@@ -15,13 +15,14 @@ jobs:
     name: "Client / Unit"
     runs-on: macos-latest
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
-      - name: Checkout code
-        uses: actions/checkout@v4
 
       - name: Cache Go modules
         uses: actions/cache@v4

.github/workflows/golang-test-freebsd.yml

@@ -25,7 +25,7 @@ jobs:
           release: "14.2"
           prepare: |
             pkg install -y curl pkgconf xorg
-            GO_TARBALL="go1.23.12.freebsd-amd64.tar.gz"
+            GO_TARBALL="go1.24.10.freebsd-amd64.tar.gz"
             GO_URL="https://go.dev/dl/$GO_TARBALL"
             curl -vLO "$GO_URL"
             tar -C /usr/local -vxzf "$GO_TARBALL"            

.github/workflows/golang-test-linux.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
       - name: Get Go environment
@@ -106,15 +106,15 @@ jobs:
         arch: [ '386','amd64' ]
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -151,15 +151,15 @@ jobs:
     needs: [ build-cache ]
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Get Go environment
         id: go-env
         run: |
@@ -200,7 +200,7 @@ jobs:
             -e GOCACHE=${CONTAINER_GOCACHE} \
             -e GOMODCACHE=${CONTAINER_GOMODCACHE} \
             -e CONTAINER=${CONTAINER} \
-            golang:1.23-alpine \
+            golang:1.24-alpine \
             sh -c ' \
               apk update; apk add --no-cache \
                 ca-certificates iptables ip6tables dbus dbus-dev libpcap-dev build-base; \
@@ -220,15 +220,15 @@ jobs:
             raceFlag: "-race"
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Install dependencies
         if: steps.cache.outputs.cache-hit != 'true'
         run: sudo apt update && sudo apt install -y gcc-multilib g++-multilib libc6-dev-i386
@@ -270,15 +270,15 @@ jobs:
         arch: [ '386','amd64' ]
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Install dependencies
         if: steps.cache.outputs.cache-hit != 'true'
         run: sudo apt update && sudo apt install -y gcc-multilib g++-multilib libc6-dev-i386
@@ -321,15 +321,15 @@ jobs:
         store: [ 'sqlite', 'postgres', 'mysql' ]
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -408,15 +408,16 @@ jobs:
             -v $PWD/prometheus.yml:/etc/prometheus/prometheus.yml \
             -p 9090:9090 \
             prom/prometheus
-      - name: Install Go
-        uses: actions/setup-go@v5
-        with:
-          go-version: "1.23.x"
-          cache: false
 
       - name: Checkout code
         uses: actions/checkout@v4
 
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version-file: "go.mod"
+          cache: false
+
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -497,15 +498,15 @@ jobs:
             -p 9090:9090 \
             prom/prometheus
 
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -561,15 +562,15 @@ jobs:
         store: [ 'sqlite', 'postgres']
     runs-on: ubuntu-22.04
     steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV

.github/workflows/golang-test-windows.yml

@@ -24,7 +24,7 @@ jobs:
         uses: actions/setup-go@v5
         id: go
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
 
       - name: Get Go environment

.github/workflows/golangci-lint.yml

@@ -46,7 +46,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
           cache: false
       - name: Install dependencies
         if: matrix.os == 'ubuntu-latest'

.github/workflows/mobile-build-validation.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
       - name: Setup Android SDK
         uses: android-actions/setup-android@v3
         with:
@@ -39,7 +39,7 @@ jobs:
       - name: Setup NDK
         run: /usr/local/lib/android/sdk/cmdline-tools/7.0/bin/sdkmanager --install "ndk;23.1.7779620"
       - name: install gomobile
-        run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20240404231514-09dbf07665ed
+        run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20251113184115-a159579294ab
       - name: gomobile init
         run: gomobile init
       - name: build android netbird lib
@@ -56,9 +56,9 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
       - name: install gomobile
-        run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20240404231514-09dbf07665ed
+        run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20251113184115-a159579294ab
       - name: gomobile init
         run: gomobile init
       - name: build iOS netbird lib

.github/workflows/release.yml

@@ -20,7 +20,7 @@ concurrency:
 
 jobs:
   release:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest-m
     env:
       flags: ""
     steps:
@@ -40,7 +40,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23"
+          go-version-file: "go.mod"
           cache: false
       - name: Cache Go modules
         uses: actions/cache@v4
@@ -136,7 +136,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23"
+          go-version-file: "go.mod"
           cache: false
       - name: Cache Go modules
         uses: actions/cache@v4
@@ -200,7 +200,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23"
+          go-version-file: "go.mod"
           cache: false
       - name: Cache Go modules
         uses: actions/cache@v4

.github/workflows/test-infrastructure-files.yml

@@ -67,10 +67,13 @@ jobs:
       - name: Install curl
         run: sudo apt-get install -y curl
 
+      - name: Checkout code
+        uses: actions/checkout@v4
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
 
       - name: Cache Go modules
         uses: actions/cache@v4
@@ -80,9 +83,6 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-go-
 
-      - name: Checkout code
-        uses: actions/checkout@v4
-
       - name: Setup MySQL privileges
         if: matrix.store == 'mysql'
         run: |

.github/workflows/wasm-build-validation.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
       - name: Install dependencies
         run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev libpcap-dev
       - name: Install golangci-lint
@@ -45,7 +45,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: "1.23.x"
+          go-version-file: "go.mod"
       - name: Build Wasm client
         run: GOOS=js GOARCH=wasm go build -o netbird.wasm ./client/wasm/cmd
         env:
@@ -60,8 +60,8 @@ jobs:
 
           echo "Size: ${SIZE} bytes (${SIZE_MB} MB)"
 
-          if [ ${SIZE} -gt 52428800 ]; then
-            echo "Wasm binary size (${SIZE_MB}MB) exceeds 50MB limit!"
+          if [ ${SIZE} -gt 57671680 ]; then
+            echo "Wasm binary size (${SIZE_MB}MB) exceeds 55MB limit!"
             exit 1
           fi
 

client/android/client.go

@@ -92,7 +92,7 @@ func NewClient(platformFiles PlatformFiles, androidSDKVersion int, deviceName st
 }
 
 // Run start the internal client. It is a blocker function
-func (c *Client) Run(urlOpener URLOpener, dns *DNSList, dnsReadyListener DnsReadyListener, envList *EnvList) error {
+func (c *Client) Run(urlOpener URLOpener, isAndroidTV bool, dns *DNSList, dnsReadyListener DnsReadyListener, envList *EnvList) error {
 	exportEnvList(envList)
 	cfg, err := profilemanager.UpdateOrCreateConfig(profilemanager.ConfigInput{
 		ConfigPath: c.cfgFile,
@@ -115,7 +115,7 @@ func (c *Client) Run(urlOpener URLOpener, dns *DNSList, dnsReadyListener DnsRead
 	c.ctxCancelLock.Unlock()
 
 	auth := NewAuthWithConfig(ctx, cfg)
-	err = auth.login(urlOpener)
+	err = auth.login(urlOpener, isAndroidTV)
 	if err != nil {
 		return err
 	}

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"
@@ -32,7 +24,7 @@ type ErrListener interface {
 // URLOpener it is a callback interface. The Open function will be triggered if
 // the backend want to show an url for the user
 type URLOpener interface {
-	Open(string)
+	Open(url string, userCode string)
 	OnLoginSuccess()
 }
 
@@ -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,26 +110,24 @@ 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)
 }
 
 // Login try register the client on the server
-func (a *Auth) Login(resultListener ErrListener, urlOpener URLOpener) {
+func (a *Auth) Login(resultListener ErrListener, urlOpener URLOpener, isAndroidTV bool) {
 	go func() {
-		err := a.login(urlOpener)
+		err := a.login(urlOpener, isAndroidTV)
 		if err != nil {
 			resultListener.OnError(err)
 		} else {
@@ -159,48 +136,40 @@ func (a *Auth) Login(resultListener ErrListener, urlOpener URLOpener) {
 	}()
 }
 
-func (a *Auth) login(urlOpener URLOpener) error {
-	var needsLogin bool
+func (a *Auth) login(urlOpener URLOpener, isAndroidTV bool) error {
+	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 := ""
 	if needsLogin {
-		tokenInfo, err := a.foregroundGetTokenInfo(urlOpener)
+		tokenInfo, err := a.foregroundGetTokenInfo(urlOpener, isAndroidTV)
 		if err != nil {
 			return fmt.Errorf("interactive sso login failed: %v", err)
 		}
 		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
 }
 
-func (a *Auth) foregroundGetTokenInfo(urlOpener URLOpener) (*auth.TokenInfo, error) {
-	oAuthFlow, err := auth.NewOAuthFlow(a.ctx, a.config, false, "")
+func (a *Auth) foregroundGetTokenInfo(urlOpener URLOpener, isAndroidTV bool) (*auth.TokenInfo, error) {
+	oAuthFlow, err := auth.NewOAuthFlow(a.ctx, a.config, false, isAndroidTV, "")
 	if err != nil {
 		return nil, err
 	}
@@ -210,24 +179,12 @@ func (a *Auth) foregroundGetTokenInfo(urlOpener URLOpener) (*auth.TokenInfo, err
 		return nil, fmt.Errorf("getting a request OAuth flow info failed: %v", err)
 	}
 
-	go urlOpener.Open(flowInfo.VerificationURIComplete)
+	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,13 +2,13 @@ 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"
@@ -21,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"
 )
 
@@ -277,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 := ""
@@ -300,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
@@ -332,7 +320,7 @@ func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *pro
 		hint = profileState.Email
 	}
 
-	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isUnixRunningDesktop(), hint)
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isUnixRunningDesktop(), false, hint)
 	if err != nil {
 		return nil, err
 	}
@@ -344,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)
 	}

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/firewall/nftables/router_linux.go

@@ -91,11 +91,7 @@ func newRouter(workTable *nftables.Table, wgIface iFaceMapper, mtu uint16) (*rou
 	var err error
 	r.filterTable, err = r.loadFilterTable()
 	if err != nil {
-		if errors.Is(err, errFilterTableNotFound) {
-			log.Warnf("table 'filter' not found for forward rules")
-		} else {
-			return nil, fmt.Errorf("load filter table: %w", err)
-		}
+		log.Warnf("failed to load filter table, skipping accept rules: %v", err)
 	}
 
 	return r, nil
@@ -175,7 +171,7 @@ func (r *router) removeNatPreroutingRules() error {
 func (r *router) loadFilterTable() (*nftables.Table, error) {
 	tables, err := r.conn.ListTablesOfFamily(nftables.TableFamilyIPv4)
 	if err != nil {
-		return nil, fmt.Errorf("unable to list tables: %v", err)
+		return nil, fmt.Errorf("list tables: %w", err)
 	}
 
 	for _, table := range tables {
@@ -193,8 +189,6 @@ func (r *router) createContainers() error {
 		Table: r.workTable,
 	})
 
-	insertReturnTrafficRule(r.conn, r.workTable, r.chains[chainNameRoutingFw])
-
 	prio := *nftables.ChainPriorityNATSource - 1
 	r.chains[chainNameRoutingNat] = r.conn.AddChain(&nftables.Chain{
 		Name:     chainNameRoutingNat,
@@ -236,9 +230,12 @@ func (r *router) createContainers() error {
 		Type:     nftables.ChainTypeFilter,
 	})
 
-	// Add the single NAT rule that matches on mark
-	if err := r.addPostroutingRules(); err != nil {
-		return fmt.Errorf("add single nat rule: %v", err)
+	insertReturnTrafficRule(r.conn, r.workTable, r.chains[chainNameRoutingFw])
+
+	r.addPostroutingRules()
+
+	if err := r.conn.Flush(); err != nil {
+		return fmt.Errorf("initialize tables: %v", err)
 	}
 
 	if err := r.addMSSClampingRules(); err != nil {
@@ -250,11 +247,7 @@ func (r *router) createContainers() error {
 	}
 
 	if err := r.refreshRulesMap(); err != nil {
-		log.Errorf("failed to clean up rules from FORWARD chain: %s", err)
-	}
-
-	if err := r.conn.Flush(); err != nil {
-		return fmt.Errorf("initialize tables: %v", err)
+		log.Errorf("failed to refresh rules: %s", err)
 	}
 
 	return nil
@@ -695,7 +688,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 }
 
 // addPostroutingRules adds the masquerade rules
-func (r *router) addPostroutingRules() error {
+func (r *router) addPostroutingRules() {
 	// First masquerade rule for traffic coming in from WireGuard interface
 	exprs := []expr.Any{
 		// Match on the first fwmark
@@ -761,8 +754,6 @@ func (r *router) addPostroutingRules() error {
 		Chain: r.chains[chainNameRoutingNat],
 		Exprs: exprs2,
 	})
-
-	return nil
 }
 
 // addMSSClampingRules adds MSS clamping rules to prevent fragmentation for forwarded traffic.
@@ -839,7 +830,7 @@ func (r *router) addMSSClampingRules() error {
 		Exprs: exprsOut,
 	})
 
-	return nil
+	return r.conn.Flush()
 }
 
 // addLegacyRouteRule adds a legacy routing rule for mgmt servers pre route acls
@@ -1068,7 +1059,7 @@ func (r *router) acceptFilterRulesNftables() error {
 	}
 	r.conn.InsertRule(inputRule)
 
-	return nil
+	return r.conn.Flush()
 }
 
 func (r *router) removeAcceptFilterRules() error {
@@ -1196,7 +1187,7 @@ func (r *router) refreshRulesMap() error {
 	for _, chain := range r.chains {
 		rules, err := r.conn.GetRules(chain.Table, chain)
 		if err != nil {
-			return fmt.Errorf(" unable to list rules: %v", err)
+			return fmt.Errorf("list rules: %w", err)
 		}
 		for _, rule := range rules {
 			if len(rule.UserData) > 0 {

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"
 )
 
@@ -60,14 +59,19 @@ func (t TokenInfo) GetTokenToUse() string {
 	return t.AccessToken
 }
 
+func shouldUseDeviceFlow(force bool, isUnixDesktopClient bool) bool {
+	return force || (runtime.GOOS == "linux" || runtime.GOOS == "freebsd") && !isUnixDesktopClient
+}
+
 // NewOAuthFlow initializes and returns the appropriate OAuth flow based on the management configuration
 //
 // It starts by initializing the PKCE.If this process fails, it resorts to the Device Code Flow,
 // and if that also fails, the authentication process is deemed unsuccessful
 //
 // On Linux distros without desktop environment support, it only tries to initialize the Device Code Flow
-func NewOAuthFlow(ctx context.Context, config *profilemanager.Config, isUnixDesktopClient bool, hint string) (OAuthFlow, error) {
-	if (runtime.GOOS == "linux" || runtime.GOOS == "freebsd") && !isUnixDesktopClient {
+// forceDeviceCodeFlow can be used to skip PKCE and go directly to Device Code Flow (e.g., for Android TV)
+func NewOAuthFlow(ctx context.Context, config *profilemanager.Config, isUnixDesktopClient bool, forceDeviceCodeFlow bool, hint string) (OAuthFlow, error) {
+	if shouldUseDeviceFlow(forceDeviceCodeFlow, isUnixDesktopClient) {
 		return authenticateWithDeviceCodeFlow(ctx, config, hint)
 	}
 
@@ -82,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:
@@ -109,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

@@ -19,8 +19,8 @@ import (
 	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"
 )
 
 var _ OAuthFlow = &PKCEAuthorizationFlow{}
@@ -33,17 +33,67 @@ 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
 
 	// find the first available redirect URL
@@ -121,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)
 
@@ -135,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 {
@@ -146,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:

client/internal/auth/pkce_flow_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/netbirdio/netbird/client/internal"
 	mgm "github.com/netbirdio/netbird/shared/management/client/common"
 )
 
@@ -41,7 +40,7 @@ func TestPromptLogin(t *testing.T) {
 
 	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",

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/dns.go

@@ -76,6 +76,9 @@ func collectPTRRecords(config *nbdns.Config, prefix netip.Prefix) []nbdns.Simple
 	var records []nbdns.SimpleRecord
 
 	for _, zone := range config.CustomZones {
+		if zone.SkipPTRProcess {
+			continue
+		}
 		for _, record := range zone.Records {
 			if record.Type != int(dns.TypeA) {
 				continue
@@ -106,8 +109,9 @@ func addReverseZone(config *nbdns.Config, network netip.Prefix) {
 	records := collectPTRRecords(config, network)
 
 	reverseZone := nbdns.CustomZone{
-		Domain:  zoneName,
-		Records: records,
+		Domain:               zoneName,
+		Records:              records,
+		SearchDomainDisabled: true,
 	}
 
 	config.CustomZones = append(config.CustomZones, reverseZone)

client/internal/dns/host.go

@@ -11,11 +11,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 )
 
-const (
-	ipv4ReverseZone = ".in-addr.arpa."
-	ipv6ReverseZone = ".ip6.arpa."
-)
-
 type hostManager interface {
 	applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error
 	restoreHostDNS() error
@@ -110,10 +105,9 @@ func dnsConfigToHostDNSConfig(dnsConfig nbdns.Config, ip netip.Addr, port int) H
 	}
 
 	for _, customZone := range dnsConfig.CustomZones {
-		matchOnly := strings.HasSuffix(customZone.Domain, ipv4ReverseZone) || strings.HasSuffix(customZone.Domain, ipv6ReverseZone)
 		config.Domains = append(config.Domains, DomainConfig{
 			Domain:    strings.ToLower(dns.Fqdn(customZone.Domain)),
-			MatchOnly: matchOnly,
+			MatchOnly: customZone.SearchDomainDisabled,
 		})
 	}
 

client/internal/dns/upstream.go

@@ -197,7 +197,7 @@ func (u *upstreamResolverBase) handleUpstreamError(err error, upstream netip.Add
 		timeoutMsg += " " + peerInfo
 	}
 	timeoutMsg += fmt.Sprintf(" - error: %v", err)
-	logger.Warnf(timeoutMsg)
+	logger.Warn(timeoutMsg)
 }
 
 func (u *upstreamResolverBase) writeSuccessResponse(w dns.ResponseWriter, rm *dns.Msg, upstream netip.AddrPort, domain string, t time.Duration, logger *log.Entry) bool {

client/internal/engine.go

@@ -891,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
 	}
@@ -1207,7 +1207,9 @@ func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig, network netip.Prefix) nbdns
 
 	for _, zone := range protoDNSConfig.GetCustomZones() {
 		dnsZone := nbdns.CustomZone{
-			Domain: zone.GetDomain(),
+			Domain:               zone.GetDomain(),
+			SearchDomainDisabled: zone.GetSearchDomainDisabled(),
+			SkipPTRProcess:       zone.GetSkipPTRProcess(),
 		}
 		for _, record := range zone.Records {
 			dnsRecord := nbdns.SimpleRecord{
@@ -1515,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
 	}
@@ -1664,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

@@ -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,

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/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
 }
 
@@ -228,7 +234,7 @@ func (c *Client) LoginForMobile() string {
 		ConfigPath: c.cfgFile,
 	})
 
-	oAuthFlow, err := auth.NewOAuthFlow(ctx, cfg, false, "")
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, cfg, false, false, "")
 	if err != nil {
 		return err.Error()
 	}
@@ -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/server/server.go

@@ -278,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
@@ -504,7 +512,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 		if msg.Hint != nil {
 			hint = *msg.Hint
 		}
-		oAuthFlow, err := auth.NewOAuthFlow(ctx, config, msg.IsUnixDesktopClient, hint)
+		oAuthFlow, err := auth.NewOAuthFlow(ctx, config, msg.IsUnixDesktopClient, false, hint)
 		if err != nil {
 			state.Set(internal.StatusLoginFailed)
 			return nil, err
@@ -606,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()
@@ -1002,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() {
@@ -1012,7 +1019,7 @@ func (s *Server) sendLogoutRequestWithConfig(ctx context.Context, config *profil
 		}
 	}()
 
-	return mgmClient.Logout()
+	return mgmClient.Logout(ctx)
 }
 
 // Status returns the daemon status
@@ -1235,7 +1242,7 @@ func (s *Server) RequestJWTAuth(
 	}
 
 	isDesktop := isUnixRunningDesktop()
-	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isDesktop, hint)
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isDesktop, false, hint)
 	if err != nil {
 		return nil, gstatus.Errorf(codes.Internal, "failed to create OAuth flow: %v", err)
 	}

client/ssh/testutil/user_helpers.go

@@ -72,7 +72,8 @@ func IsSystemAccount(username string) bool {
 			return true
 		}
 	}
-	return false
+
+	return strings.HasSuffix(username, "$")
 }
 
 // RegisterTestUserCleanup registers a test user for cleanup

client/ssh/testutil/user_helpers_test.go

@@ -0,0 +1,115 @@
+package testutil
+
+import (
+	"os/user"
+	"runtime"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestUserCurrentBehavior validates user.Current() behavior on Windows.
+// When running as SYSTEM on a domain-joined machine, user.Current() returns:
+// - Username: Computer account name (e.g., "DOMAIN\MACHINE$")
+// - SID: SYSTEM SID (S-1-5-18)
+func TestUserCurrentBehavior(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip("Windows-specific test")
+	}
+
+	currentUser, err := user.Current()
+	require.NoError(t, err, "Should be able to get current user")
+
+	t.Logf("Current user - Username: %s, SID: %s", currentUser.Username, currentUser.Uid)
+
+	// When running as SYSTEM, validate expected behavior
+	if currentUser.Uid == "S-1-5-18" {
+		t.Run("SYSTEM_account_behavior", func(t *testing.T) {
+			// SID must be S-1-5-18 for SYSTEM
+			require.Equal(t, "S-1-5-18", currentUser.Uid,
+				"SYSTEM account must have SID S-1-5-18")
+
+			// Username can be either "NT AUTHORITY\SYSTEM" (standalone)
+			// or "DOMAIN\MACHINE$" (domain-joined)
+			username := currentUser.Username
+			isNTAuthority := strings.Contains(strings.ToUpper(username), "NT AUTHORITY")
+			isComputerAccount := strings.HasSuffix(username, "$")
+
+			assert.True(t, isNTAuthority || isComputerAccount,
+				"Username should be either 'NT AUTHORITY\\SYSTEM' or computer account (ending with $), got: %s",
+				username)
+
+			if isComputerAccount {
+				t.Logf("SYSTEM as computer account: %s", username)
+			} else if isNTAuthority {
+				t.Logf("SYSTEM as NT AUTHORITY\\SYSTEM")
+			}
+		})
+	}
+
+	// Validate that IsSystemAccount correctly identifies system accounts
+	t.Run("IsSystemAccount_validation", func(t *testing.T) {
+		// Test with current user if it's a system account
+		if currentUser.Uid == "S-1-5-18" || // SYSTEM
+			currentUser.Uid == "S-1-5-19" || // LOCAL SERVICE
+			currentUser.Uid == "S-1-5-20" { // NETWORK SERVICE
+
+			result := IsSystemAccount(currentUser.Username)
+			assert.True(t, result,
+				"IsSystemAccount should recognize system account: %s (SID: %s)",
+				currentUser.Username, currentUser.Uid)
+		}
+
+		// Test explicit cases
+		testCases := []struct {
+			username string
+			expected bool
+			reason   string
+		}{
+			{"NT AUTHORITY\\SYSTEM", true, "NT AUTHORITY\\SYSTEM"},
+			{"system", true, "system"},
+			{"SYSTEM", true, "SYSTEM (case insensitive)"},
+			{"NT AUTHORITY\\LOCAL SERVICE", true, "LOCAL SERVICE"},
+			{"NT AUTHORITY\\NETWORK SERVICE", true, "NETWORK SERVICE"},
+			{"DOMAIN\\MACHINE$", true, "computer account (ends with $)"},
+			{"WORKGROUP\\WIN2K19-C2$", true, "computer account (ends with $)"},
+			{"Administrator", false, "Administrator is not a system account"},
+			{"alice", false, "regular user"},
+			{"DOMAIN\\alice", false, "domain user"},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.username, func(t *testing.T) {
+				result := IsSystemAccount(tc.username)
+				assert.Equal(t, tc.expected, result,
+					"IsSystemAccount(%q) should be %v because: %s",
+					tc.username, tc.expected, tc.reason)
+			})
+		}
+	})
+}
+
+// TestComputerAccountDetection validates computer account detection.
+func TestComputerAccountDetection(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip("Windows-specific test")
+	}
+
+	computerAccounts := []string{
+		"MACHINE$",
+		"WIN2K19-C2$",
+		"DOMAIN\\MACHINE$",
+		"WORKGROUP\\SERVER$",
+		"server.domain.com$",
+	}
+
+	for _, account := range computerAccounts {
+		t.Run(account, func(t *testing.T) {
+			result := IsSystemAccount(account)
+			assert.True(t, result,
+				"Computer account %q should be recognized as system account", account)
+		})
+	}
+}

dns/dns.go

@@ -45,6 +45,10 @@ type CustomZone struct {
 	Domain string
 	// Records custom zone records
 	Records []SimpleRecord
+	// SearchDomainDisabled indicates whether to add match domains to a search domains list or not
+	SearchDomainDisabled bool
+	// SkipPTRProcess indicates whether a client should process PTR records from custom zones
+	SkipPTRProcess bool
 }
 
 // SimpleRecord provides a simple DNS record specification for CNAME, A and AAAA records

go.mod

@@ -1,6 +1,6 @@
 module github.com/netbirdio/netbird
 
-go 1.23.1
+go 1.24.10
 
 require (
 	cunicu.li/go-rosenpass v0.4.0
@@ -17,8 +17,8 @@ require (
 	github.com/spf13/cobra v1.7.0
 	github.com/spf13/pflag v1.0.5
 	github.com/vishvananda/netlink v1.3.1
-	golang.org/x/crypto v0.41.0
-	golang.org/x/sys v0.35.0
+	golang.org/x/crypto v0.45.0
+	golang.org/x/sys v0.38.0
 	golang.zx2c4.com/wireguard v0.0.0-20230704135630-469159ecf7d1
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	golang.zx2c4.com/wireguard/windows v0.5.3
@@ -105,12 +105,12 @@ require (
 	go.uber.org/zap v1.27.0
 	goauthentik.io/api/v3 v3.2023051.3
 	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
-	golang.org/x/mobile v0.0.0-20231127183840-76ac6878050a
-	golang.org/x/mod v0.26.0
-	golang.org/x/net v0.42.0
+	golang.org/x/mobile v0.0.0-20251113184115-a159579294ab
+	golang.org/x/mod v0.30.0
+	golang.org/x/net v0.47.0
 	golang.org/x/oauth2 v0.30.0
-	golang.org/x/sync v0.16.0
-	golang.org/x/term v0.34.0
+	golang.org/x/sync v0.18.0
+	golang.org/x/term v0.37.0
 	golang.org/x/time v0.12.0
 	google.golang.org/api v0.177.0
 	gopkg.in/yaml.v3 v3.0.1
@@ -251,9 +251,9 @@ require (
 	go.opentelemetry.io/otel/sdk v1.35.0 // indirect
 	go.opentelemetry.io/otel/trace v1.35.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/image v0.24.0 // indirect
-	golang.org/x/text v0.28.0 // indirect
-	golang.org/x/tools v0.35.0 // indirect
+	golang.org/x/image v0.33.0 // indirect
+	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/tools v0.39.0 // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20250707201910-8d1bb00bc6a7 // indirect
 	gopkg.in/square/go-jose.v2 v2.6.0 // indirect

go.sum

@@ -600,19 +600,19 @@ golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1m
 golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
 golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
 golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
-golang.org/x/crypto v0.41.0 h1:WKYxWedPGCTVVl5+WHSSrOBT0O8lx32+zxmHxijgXp4=
-golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sUc=
+golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q=
+golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM=
 golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc=
-golang.org/x/image v0.24.0 h1:AN7zRgVsbvmTfNyqIbbOraYL8mSwcKncEj8ofjgzcMQ=
-golang.org/x/image v0.24.0/go.mod h1:4b/ITuLfqYq1hqZcjofwctIhi7sZh2WaCjvsBNjjya8=
+golang.org/x/image v0.33.0 h1:LXRZRnv1+zGd5XBUVRFmYEphyyKJjQjCRiOuAP3sZfQ=
+golang.org/x/image v0.33.0/go.mod h1:DD3OsTYT9chzuzTQt+zMcOlBHgfoKQb1gry8p76Y1sc=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
-golang.org/x/mobile v0.0.0-20231127183840-76ac6878050a h1:sYbmY3FwUWCBTodZL1S3JUuOvaW6kM2o+clDzzDNBWg=
-golang.org/x/mobile v0.0.0-20231127183840-76ac6878050a/go.mod h1:Ede7gF0KGoHlj822RtphAHK1jLdrcuRBZg0sF1Q+SPc=
+golang.org/x/mobile v0.0.0-20251113184115-a159579294ab h1:Iqyc+2zr7aGyLuEadIm0KRJP0Wwt+fhlXLa51Fxf1+Q=
+golang.org/x/mobile v0.0.0-20251113184115-a159579294ab/go.mod h1:Eq3Nh/5pFSWug2ohiudJ1iyU59SO78QFuh4qTTN++I0=
 golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
@@ -622,8 +622,8 @@ golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
 golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
-golang.org/x/mod v0.26.0 h1:EGMPT//Ezu+ylkCijjPc+f4Aih7sZvaAr+O3EHBxvZg=
-golang.org/x/mod v0.26.0/go.mod h1:/j6NAhSk8iQ723BGAUyoAcn7SlD7s15Dp9Nd/SfeaFQ=
+golang.org/x/mod v0.30.0 h1:fDEXFVZ/fmCKProc/yAXXUijritrDzahmwwefnjoPFk=
+golang.org/x/mod v0.30.0/go.mod h1:lAsf5O2EvJeSFMiBxXDki7sCgAxEUcZHXoXMKT4GJKc=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -647,8 +647,8 @@ golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
 golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
 golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
 golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
-golang.org/x/net v0.42.0 h1:jzkYrhi3YQWD6MLBJcsklgQsoAcw89EcZbJw8Z614hs=
-golang.org/x/net v0.42.0/go.mod h1:FF1RA5d3u7nAYA4z2TkclSCKh68eSXtiFwcWQpPXdt8=
+golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
+golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.8.0/go.mod h1:yr7u4HXZRm1R1kBWqr/xKNqewf0plRYoB7sla+BCIXE=
 golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
@@ -665,8 +665,8 @@ golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
 golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
-golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
-golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
+golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
+golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -703,8 +703,8 @@ golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.19.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
-golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
+golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
 golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
@@ -717,8 +717,8 @@ golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY=
 golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
 golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
 golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
-golang.org/x/term v0.34.0 h1:O/2T7POpk0ZZ7MAzMeWFSg6S5IpWd/RXDlM9hgM3DR4=
-golang.org/x/term v0.34.0/go.mod h1:5jC53AEywhIVebHgPVeg0mj8OD3VO9OzclacVrqpaAw=
+golang.org/x/term v0.37.0 h1:8EGAD0qCmHYZg6J17DvsMy9/wJ7/D/4pV/wfnld5lTU=
+golang.org/x/term v0.37.0/go.mod h1:5pB4lxRNYYVZuTLmy8oR2BH8dflOR+IbTYFD8fi3254=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
@@ -730,8 +730,8 @@ golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
-golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng=
-golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
+golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
+golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
 golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
 golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
@@ -749,8 +749,8 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
-golang.org/x/tools v0.35.0 h1:mBffYraMEf7aa0sB+NuKnuCy8qI/9Bughn8dC2Gu5r0=
-golang.org/x/tools v0.35.0/go.mod h1:NKdj5HkL/73byiZSJjqJgKn3ep7KjFkBOkR/Hps3VPw=
+golang.org/x/tools v0.39.0 h1:ik4ho21kwuQln40uelmciQPp9SipgNDdrafrYA4TmQQ=
+golang.org/x/tools v0.39.0/go.mod h1:JnefbkDPyD8UU2kI5fuf8ZX4/yUeh9W877ZeBONxUqQ=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

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

@@ -1,16 +1,9 @@
 package controller
 
 import (
-	"context"
-	"sync"
-	"sync/atomic"
 	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
-	"github.com/netbirdio/netbird/management/server/mock_server"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 )
 
@@ -114,131 +107,3 @@ func TestComputeForwarderPort(t *testing.T) {
 		t.Errorf("Expected %d for peers with unknown version, got %d", network_map.OldForwarderPort, result)
 	}
 }
-
-func TestBufferUpdateAccountPeers(t *testing.T) {
-	const (
-		peersCount            = 1000
-		updateAccountInterval = 50 * time.Millisecond
-	)
-
-	var (
-		deletedPeers, updatePeersDeleted, updatePeersRuns atomic.Int32
-		uapLastRun, dpLastRun                             atomic.Int64
-
-		totalNewRuns, totalOldRuns int
-	)
-
-	uap := func(ctx context.Context, accountID string) {
-		updatePeersDeleted.Store(deletedPeers.Load())
-		updatePeersRuns.Add(1)
-		uapLastRun.Store(time.Now().UnixMilli())
-		time.Sleep(100 * time.Millisecond)
-	}
-
-	t.Run("new approach", func(t *testing.T) {
-		updatePeersRuns.Store(0)
-		updatePeersDeleted.Store(0)
-		deletedPeers.Store(0)
-
-		var mustore sync.Map
-		bufupd := func(ctx context.Context, accountID string) {
-			mu, _ := mustore.LoadOrStore(accountID, &bufferUpdate{})
-			b := mu.(*bufferUpdate)
-
-			if !b.mu.TryLock() {
-				b.update.Store(true)
-				return
-			}
-
-			if b.next != nil {
-				b.next.Stop()
-			}
-
-			go func() {
-				defer b.mu.Unlock()
-				uap(ctx, accountID)
-				if !b.update.Load() {
-					return
-				}
-				b.update.Store(false)
-				b.next = time.AfterFunc(updateAccountInterval, func() {
-					uap(ctx, accountID)
-				})
-			}()
-		}
-		dp := func(ctx context.Context, accountID, peerID, userID string) error {
-			deletedPeers.Add(1)
-			dpLastRun.Store(time.Now().UnixMilli())
-			time.Sleep(10 * time.Millisecond)
-			bufupd(ctx, accountID)
-			return nil
-		}
-
-		am := mock_server.MockAccountManager{
-			UpdateAccountPeersFunc:       uap,
-			BufferUpdateAccountPeersFunc: bufupd,
-			DeletePeerFunc:               dp,
-		}
-		empty := ""
-		for range peersCount {
-			//nolint
-			am.DeletePeer(context.Background(), empty, empty, empty)
-		}
-		time.Sleep(100 * time.Millisecond)
-
-		assert.Equal(t, peersCount, int(deletedPeers.Load()), "Expected all peers to be deleted")
-		assert.Equal(t, peersCount, int(updatePeersDeleted.Load()), "Expected all peers to be updated in the buffer")
-		assert.GreaterOrEqual(t, uapLastRun.Load(), dpLastRun.Load(), "Expected update account peers to run after delete peer")
-
-		totalNewRuns = int(updatePeersRuns.Load())
-	})
-
-	t.Run("old approach", func(t *testing.T) {
-		updatePeersRuns.Store(0)
-		updatePeersDeleted.Store(0)
-		deletedPeers.Store(0)
-
-		var mustore sync.Map
-		bufupd := func(ctx context.Context, accountID string) {
-			mu, _ := mustore.LoadOrStore(accountID, &sync.Mutex{})
-			b := mu.(*sync.Mutex)
-
-			if !b.TryLock() {
-				return
-			}
-
-			go func() {
-				time.Sleep(updateAccountInterval)
-				b.Unlock()
-				uap(ctx, accountID)
-			}()
-		}
-		dp := func(ctx context.Context, accountID, peerID, userID string) error {
-			deletedPeers.Add(1)
-			dpLastRun.Store(time.Now().UnixMilli())
-			time.Sleep(10 * time.Millisecond)
-			bufupd(ctx, accountID)
-			return nil
-		}
-
-		am := mock_server.MockAccountManager{
-			UpdateAccountPeersFunc:       uap,
-			BufferUpdateAccountPeersFunc: bufupd,
-			DeletePeerFunc:               dp,
-		}
-		empty := ""
-		for range peersCount {
-			//nolint
-			am.DeletePeer(context.Background(), empty, empty, empty)
-		}
-		time.Sleep(100 * time.Millisecond)
-
-		assert.Equal(t, peersCount, int(deletedPeers.Load()), "Expected all peers to be deleted")
-		assert.Equal(t, peersCount, int(updatePeersDeleted.Load()), "Expected all peers to be updated in the buffer")
-		assert.GreaterOrEqual(t, uapLastRun.Load(), dpLastRun.Load(), "Expected update account peers to run after delete peer")
-
-		totalOldRuns = int(updatePeersRuns.Load())
-	})
-	assert.Less(t, totalNewRuns, totalOldRuns, "Expected new approach to run less than old approach. New runs: %d, Old runs: %d", totalNewRuns, totalOldRuns)
-	t.Logf("New runs: %d, Old runs: %d", totalNewRuns, totalOldRuns)
-}

management/internals/shared/grpc/conversion.go

@@ -369,7 +369,7 @@ func buildJWTConfig(config *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfi
 	}
 
 	issuer := strings.TrimSpace(config.AuthIssuer)
-	if issuer == "" || deviceFlowConfig != nil {
+	if issuer == "" && deviceFlowConfig != nil {
 		if d := deriveIssuerFromTokenEndpoint(deviceFlowConfig.ProviderConfig.TokenEndpoint); d != "" {
 			issuer = d
 		}

management/server/http/handlers/policies/posture_checks_handler_test.go

@@ -46,7 +46,7 @@ func initPostureChecksTestData(postureChecks ...*posture.Checks) *postureChecksH
 				testPostureChecks[postureChecks.ID] = postureChecks
 
 				if err := postureChecks.Validate(); err != nil {
-					return nil, status.Errorf(status.InvalidArgument, err.Error()) //nolint
+					return nil, status.Errorf(status.InvalidArgument, "%s", err.Error()) //nolint
 				}
 
 				return postureChecks, nil

management/server/posture_checks.go

@@ -158,7 +158,7 @@ func arePostureCheckChangesAffectPeers(ctx context.Context, transaction store.St
 // validatePostureChecks validates the posture checks.
 func validatePostureChecks(ctx context.Context, transaction store.Store, accountID string, postureChecks *posture.Checks) error {
 	if err := postureChecks.Validate(); err != nil {
-		return status.Errorf(status.InvalidArgument, err.Error()) //nolint
+		return status.Errorf(status.InvalidArgument, "%s", err.Error()) //nolint
 	}
 
 	// If the posture check already has an ID, verify its existence in the store.

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

@@ -193,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")
 	}
@@ -216,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")
 	}
@@ -243,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) {
@@ -272,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)
 	}
@@ -293,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)
 	}
@@ -354,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
@@ -400,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)
 	}
@@ -489,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{
@@ -512,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")
 	}
@@ -533,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{
@@ -558,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/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/proto/management.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.26.0
-// 	protoc        v6.33.0
+// 	protoc        v6.32.1
 // source: management.proto
 
 package proto
@@ -2682,8 +2682,10 @@ type CustomZone struct {
 	sizeCache     protoimpl.SizeCache
 	unknownFields protoimpl.UnknownFields
 
-	Domain  string          `protobuf:"bytes,1,opt,name=Domain,proto3" json:"Domain,omitempty"`
-	Records []*SimpleRecord `protobuf:"bytes,2,rep,name=Records,proto3" json:"Records,omitempty"`
+	Domain               string          `protobuf:"bytes,1,opt,name=Domain,proto3" json:"Domain,omitempty"`
+	Records              []*SimpleRecord `protobuf:"bytes,2,rep,name=Records,proto3" json:"Records,omitempty"`
+	SearchDomainDisabled bool            `protobuf:"varint,3,opt,name=SearchDomainDisabled,proto3" json:"SearchDomainDisabled,omitempty"`
+	SkipPTRProcess       bool            `protobuf:"varint,4,opt,name=SkipPTRProcess,proto3" json:"SkipPTRProcess,omitempty"`
 }
 
 func (x *CustomZone) Reset() {
@@ -2732,6 +2734,20 @@ func (x *CustomZone) GetRecords() []*SimpleRecord {
 	return nil
 }
 
+func (x *CustomZone) GetSearchDomainDisabled() bool {
+	if x != nil {
+		return x.SearchDomainDisabled
+	}
+	return false
+}
+
+func (x *CustomZone) GetSkipPTRProcess() bool {
+	if x != nil {
+		return x.SkipPTRProcess
+	}
+	return false
+}
+
 // SimpleRecord represents a dns.SimpleRecord
 type SimpleRecord struct {
 	state         protoimpl.MessageState
@@ -3893,157 +3909,163 @@ var file_management_proto_rawDesc = []byte{
 	0x73, 0x74, 0x6f, 0x6d, 0x5a, 0x6f, 0x6e, 0x65, 0x73, 0x12, 0x28, 0x0a, 0x0d, 0x46, 0x6f, 0x72,
 	0x77, 0x61, 0x72, 0x64, 0x65, 0x72, 0x50, 0x6f, 0x72, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x03,
 	0x42, 0x02, 0x18, 0x01, 0x52, 0x0d, 0x46, 0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x65, 0x72, 0x50,
-	0x6f, 0x72, 0x74, 0x22, 0x58, 0x0a, 0x0a, 0x43, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x5a, 0x6f, 0x6e,
-	0x65, 0x12, 0x16, 0x0a, 0x06, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x06, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x12, 0x32, 0x0a, 0x07, 0x52, 0x65, 0x63,
-	0x6f, 0x72, 0x64, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x18, 0x2e, 0x6d, 0x61, 0x6e,
-	0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x53, 0x69, 0x6d, 0x70, 0x6c, 0x65, 0x52, 0x65,
-	0x63, 0x6f, 0x72, 0x64, 0x52, 0x07, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x22, 0x74, 0x0a,
-	0x0c, 0x53, 0x69, 0x6d, 0x70, 0x6c, 0x65, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, 0x12, 0x0a,
-	0x04, 0x4e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x4e, 0x61, 0x6d,
-	0x65, 0x12, 0x12, 0x0a, 0x04, 0x54, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52,
-	0x04, 0x54, 0x79, 0x70, 0x65, 0x12, 0x14, 0x0a, 0x05, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x18, 0x03,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x12, 0x10, 0x0a, 0x03, 0x54,
-	0x54, 0x4c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x03, 0x52, 0x03, 0x54, 0x54, 0x4c, 0x12, 0x14, 0x0a,
-	0x05, 0x52, 0x44, 0x61, 0x74, 0x61, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x52, 0x44,
-	0x61, 0x74, 0x61, 0x22, 0xb3, 0x01, 0x0a, 0x0f, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76,
-	0x65, 0x72, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x12, 0x38, 0x0a, 0x0b, 0x4e, 0x61, 0x6d, 0x65, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x6d,
-	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65,
-	0x72, 0x76, 0x65, 0x72, 0x52, 0x0b, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72,
-	0x73, 0x12, 0x18, 0x0a, 0x07, 0x50, 0x72, 0x69, 0x6d, 0x61, 0x72, 0x79, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x08, 0x52, 0x07, 0x50, 0x72, 0x69, 0x6d, 0x61, 0x72, 0x79, 0x12, 0x18, 0x0a, 0x07, 0x44,
-	0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, 0x09, 0x52, 0x07, 0x44, 0x6f,
-	0x6d, 0x61, 0x69, 0x6e, 0x73, 0x12, 0x32, 0x0a, 0x14, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x44,
-	0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x04, 0x20,
-	0x01, 0x28, 0x08, 0x52, 0x14, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x44, 0x6f, 0x6d, 0x61, 0x69,
-	0x6e, 0x73, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x22, 0x48, 0x0a, 0x0a, 0x4e, 0x61, 0x6d,
-	0x65, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x49, 0x50, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x50, 0x12, 0x16, 0x0a, 0x06, 0x4e, 0x53, 0x54, 0x79, 0x70,
-	0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x06, 0x4e, 0x53, 0x54, 0x79, 0x70, 0x65, 0x12,
-	0x12, 0x0a, 0x04, 0x50, 0x6f, 0x72, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x04, 0x50,
-	0x6f, 0x72, 0x74, 0x22, 0xa7, 0x02, 0x0a, 0x0c, 0x46, 0x69, 0x72, 0x65, 0x77, 0x61, 0x6c, 0x6c,
-	0x52, 0x75, 0x6c, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x50, 0x65, 0x65, 0x72, 0x49, 0x50, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x50, 0x65, 0x65, 0x72, 0x49, 0x50, 0x12, 0x37, 0x0a, 0x09,
-	0x44, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0e, 0x32,
-	0x19, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c,
-	0x65, 0x44, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x09, 0x44, 0x69, 0x72, 0x65,
-	0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x2e, 0x0a, 0x06, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18,
-	0x03, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x16, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65,
-	0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x06, 0x41,
-	0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f,
-	0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x18, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
-	0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f,
-	0x6c, 0x52, 0x08, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x50,
-	0x6f, 0x72, 0x74, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x50, 0x6f, 0x72, 0x74, 0x12,
-	0x30, 0x0a, 0x08, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x18, 0x06, 0x20, 0x01, 0x28,
+	0x6f, 0x72, 0x74, 0x22, 0xb4, 0x01, 0x0a, 0x0a, 0x43, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x5a, 0x6f,
+	0x6e, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x18, 0x01, 0x20, 0x01,
+	0x28, 0x09, 0x52, 0x06, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x12, 0x32, 0x0a, 0x07, 0x52, 0x65,
+	0x63, 0x6f, 0x72, 0x64, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x18, 0x2e, 0x6d, 0x61,
+	0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x53, 0x69, 0x6d, 0x70, 0x6c, 0x65, 0x52,
+	0x65, 0x63, 0x6f, 0x72, 0x64, 0x52, 0x07, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x73, 0x12, 0x32,
+	0x0a, 0x14, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x44, 0x69,
+	0x73, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x14, 0x53, 0x65,
+	0x61, 0x72, 0x63, 0x68, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x44, 0x69, 0x73, 0x61, 0x62, 0x6c,
+	0x65, 0x64, 0x12, 0x26, 0x0a, 0x0e, 0x53, 0x6b, 0x69, 0x70, 0x50, 0x54, 0x52, 0x50, 0x72, 0x6f,
+	0x63, 0x65, 0x73, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0e, 0x53, 0x6b, 0x69, 0x70,
+	0x50, 0x54, 0x52, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x22, 0x74, 0x0a, 0x0c, 0x53, 0x69,
+	0x6d, 0x70, 0x6c, 0x65, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x4e, 0x61,
+	0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x4e, 0x61, 0x6d, 0x65, 0x12, 0x12,
+	0x0a, 0x04, 0x54, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x04, 0x54, 0x79,
+	0x70, 0x65, 0x12, 0x14, 0x0a, 0x05, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28,
+	0x09, 0x52, 0x05, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x12, 0x10, 0x0a, 0x03, 0x54, 0x54, 0x4c, 0x18,
+	0x04, 0x20, 0x01, 0x28, 0x03, 0x52, 0x03, 0x54, 0x54, 0x4c, 0x12, 0x14, 0x0a, 0x05, 0x52, 0x44,
+	0x61, 0x74, 0x61, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x52, 0x44, 0x61, 0x74, 0x61,
+	0x22, 0xb3, 0x01, 0x0a, 0x0f, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x47,
+	0x72, 0x6f, 0x75, 0x70, 0x12, 0x38, 0x0a, 0x0b, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76,
+	0x65, 0x72, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
+	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76, 0x65,
+	0x72, 0x52, 0x0b, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x12, 0x18,
+	0x0a, 0x07, 0x50, 0x72, 0x69, 0x6d, 0x61, 0x72, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52,
+	0x07, 0x50, 0x72, 0x69, 0x6d, 0x61, 0x72, 0x79, 0x12, 0x18, 0x0a, 0x07, 0x44, 0x6f, 0x6d, 0x61,
+	0x69, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, 0x09, 0x52, 0x07, 0x44, 0x6f, 0x6d, 0x61, 0x69,
+	0x6e, 0x73, 0x12, 0x32, 0x0a, 0x14, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x44, 0x6f, 0x6d, 0x61,
+	0x69, 0x6e, 0x73, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08,
+	0x52, 0x14, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x44, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x45,
+	0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x22, 0x48, 0x0a, 0x0a, 0x4e, 0x61, 0x6d, 0x65, 0x53, 0x65,
+	0x72, 0x76, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x49, 0x50, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09,
+	0x52, 0x02, 0x49, 0x50, 0x12, 0x16, 0x0a, 0x06, 0x4e, 0x53, 0x54, 0x79, 0x70, 0x65, 0x18, 0x02,
+	0x20, 0x01, 0x28, 0x03, 0x52, 0x06, 0x4e, 0x53, 0x54, 0x79, 0x70, 0x65, 0x12, 0x12, 0x0a, 0x04,
+	0x50, 0x6f, 0x72, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x04, 0x50, 0x6f, 0x72, 0x74,
+	0x22, 0xa7, 0x02, 0x0a, 0x0c, 0x46, 0x69, 0x72, 0x65, 0x77, 0x61, 0x6c, 0x6c, 0x52, 0x75, 0x6c,
+	0x65, 0x12, 0x16, 0x0a, 0x06, 0x50, 0x65, 0x65, 0x72, 0x49, 0x50, 0x18, 0x01, 0x20, 0x01, 0x28,
+	0x09, 0x52, 0x06, 0x50, 0x65, 0x65, 0x72, 0x49, 0x50, 0x12, 0x37, 0x0a, 0x09, 0x44, 0x69, 0x72,
+	0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x19, 0x2e, 0x6d,
+	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x44, 0x69,
+	0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x09, 0x44, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69,
+	0x6f, 0x6e, 0x12, 0x2e, 0x0a, 0x06, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01,
+	0x28, 0x0e, 0x32, 0x16, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e,
+	0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x06, 0x41, 0x63, 0x74, 0x69,
+	0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x04,
+	0x20, 0x01, 0x28, 0x0e, 0x32, 0x18, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e,
+	0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x52, 0x08,
+	0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x50, 0x6f, 0x72, 0x74,
+	0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x50, 0x6f, 0x72, 0x74, 0x12, 0x30, 0x0a, 0x08,
+	0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14,
+	0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74,
+	0x49, 0x6e, 0x66, 0x6f, 0x52, 0x08, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x1a,
+	0x0a, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x18, 0x07, 0x20, 0x01, 0x28, 0x0c,
+	0x52, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x22, 0x38, 0x0a, 0x0e, 0x4e, 0x65,
+	0x74, 0x77, 0x6f, 0x72, 0x6b, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x12, 0x14, 0x0a, 0x05,
+	0x6e, 0x65, 0x74, 0x49, 0x50, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6e, 0x65, 0x74,
+	0x49, 0x50, 0x12, 0x10, 0x0a, 0x03, 0x6d, 0x61, 0x63, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52,
+	0x03, 0x6d, 0x61, 0x63, 0x22, 0x1e, 0x0a, 0x06, 0x43, 0x68, 0x65, 0x63, 0x6b, 0x73, 0x12, 0x14,
+	0x0a, 0x05, 0x46, 0x69, 0x6c, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x05, 0x46,
+	0x69, 0x6c, 0x65, 0x73, 0x22, 0x96, 0x01, 0x0a, 0x08, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66,
+	0x6f, 0x12, 0x14, 0x0a, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0d, 0x48,
+	0x00, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x12, 0x32, 0x0a, 0x05, 0x72, 0x61, 0x6e, 0x67, 0x65,
+	0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d,
+	0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x2e, 0x52, 0x61, 0x6e,
+	0x67, 0x65, 0x48, 0x00, 0x52, 0x05, 0x72, 0x61, 0x6e, 0x67, 0x65, 0x1a, 0x2f, 0x0a, 0x05, 0x52,
+	0x61, 0x6e, 0x67, 0x65, 0x12, 0x14, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x72, 0x74, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x0d, 0x52, 0x05, 0x73, 0x74, 0x61, 0x72, 0x74, 0x12, 0x10, 0x0a, 0x03, 0x65, 0x6e,
+	0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x03, 0x65, 0x6e, 0x64, 0x42, 0x0f, 0x0a, 0x0d,
+	0x70, 0x6f, 0x72, 0x74, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x22, 0x87, 0x03,
+	0x0a, 0x11, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x46, 0x69, 0x72, 0x65, 0x77, 0x61, 0x6c, 0x6c, 0x52,
+	0x75, 0x6c, 0x65, 0x12, 0x22, 0x0a, 0x0c, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x52, 0x61, 0x6e,
+	0x67, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x0c, 0x73, 0x6f, 0x75, 0x72, 0x63,
+	0x65, 0x52, 0x61, 0x6e, 0x67, 0x65, 0x73, 0x12, 0x2e, 0x0a, 0x06, 0x61, 0x63, 0x74, 0x69, 0x6f,
+	0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x16, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
+	0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52,
+	0x06, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x0a, 0x0b, 0x64, 0x65, 0x73, 0x74, 0x69,
+	0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, 0x64, 0x65,
+	0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08, 0x70, 0x72, 0x6f,
+	0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x18, 0x2e, 0x6d, 0x61,
+	0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x50, 0x72, 0x6f,
+	0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12,
+	0x30, 0x0a, 0x08, 0x70, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x18, 0x05, 0x20, 0x01, 0x28,
 	0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50,
-	0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x08, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66,
-	0x6f, 0x12, 0x1a, 0x0a, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x18, 0x07, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x22, 0x38, 0x0a,
-	0x0e, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x12,
-	0x14, 0x0a, 0x05, 0x6e, 0x65, 0x74, 0x49, 0x50, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05,
-	0x6e, 0x65, 0x74, 0x49, 0x50, 0x12, 0x10, 0x0a, 0x03, 0x6d, 0x61, 0x63, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x03, 0x6d, 0x61, 0x63, 0x22, 0x1e, 0x0a, 0x06, 0x43, 0x68, 0x65, 0x63, 0x6b,
-	0x73, 0x12, 0x14, 0x0a, 0x05, 0x46, 0x69, 0x6c, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09,
-	0x52, 0x05, 0x46, 0x69, 0x6c, 0x65, 0x73, 0x22, 0x96, 0x01, 0x0a, 0x08, 0x50, 0x6f, 0x72, 0x74,
-	0x49, 0x6e, 0x66, 0x6f, 0x12, 0x14, 0x0a, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0d, 0x48, 0x00, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x12, 0x32, 0x0a, 0x05, 0x72, 0x61,
-	0x6e, 0x67, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
-	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x2e,
-	0x52, 0x61, 0x6e, 0x67, 0x65, 0x48, 0x00, 0x52, 0x05, 0x72, 0x61, 0x6e, 0x67, 0x65, 0x1a, 0x2f,
-	0x0a, 0x05, 0x52, 0x61, 0x6e, 0x67, 0x65, 0x12, 0x14, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x72, 0x74,
-	0x18, 0x01, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x05, 0x73, 0x74, 0x61, 0x72, 0x74, 0x12, 0x10, 0x0a,
-	0x03, 0x65, 0x6e, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x03, 0x65, 0x6e, 0x64, 0x42,
-	0x0f, 0x0a, 0x0d, 0x70, 0x6f, 0x72, 0x74, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e,
-	0x22, 0x87, 0x03, 0x0a, 0x11, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x46, 0x69, 0x72, 0x65, 0x77, 0x61,
-	0x6c, 0x6c, 0x52, 0x75, 0x6c, 0x65, 0x12, 0x22, 0x0a, 0x0c, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65,
-	0x52, 0x61, 0x6e, 0x67, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, 0x0c, 0x73, 0x6f,
-	0x75, 0x72, 0x63, 0x65, 0x52, 0x61, 0x6e, 0x67, 0x65, 0x73, 0x12, 0x2e, 0x0a, 0x06, 0x61, 0x63,
-	0x74, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x16, 0x2e, 0x6d, 0x61, 0x6e,
-	0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69,
-	0x6f, 0x6e, 0x52, 0x06, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x0a, 0x0b, 0x64, 0x65,
-	0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x0b, 0x64, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x0a, 0x08,
-	0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x18,
-	0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65,
-	0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63,
-	0x6f, 0x6c, 0x12, 0x30, 0x0a, 0x08, 0x70, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x18, 0x05,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e,
-	0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x08, 0x70, 0x6f, 0x72, 0x74,
-	0x49, 0x6e, 0x66, 0x6f, 0x12, 0x1c, 0x0a, 0x09, 0x69, 0x73, 0x44, 0x79, 0x6e, 0x61, 0x6d, 0x69,
-	0x63, 0x18, 0x06, 0x20, 0x01, 0x28, 0x08, 0x52, 0x09, 0x69, 0x73, 0x44, 0x79, 0x6e, 0x61, 0x6d,
-	0x69, 0x63, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x18, 0x07, 0x20,
-	0x03, 0x28, 0x09, 0x52, 0x07, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x12, 0x26, 0x0a, 0x0e,
-	0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x08,
-	0x20, 0x01, 0x28, 0x0d, 0x52, 0x0e, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x50, 0x72, 0x6f, 0x74,
-	0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x1a, 0x0a, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44,
-	0x18, 0x09, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44,
-	0x12, 0x18, 0x0a, 0x07, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x49, 0x44, 0x18, 0x0a, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x07, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x49, 0x44, 0x22, 0xf2, 0x01, 0x0a, 0x0e, 0x46,
-	0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x75, 0x6c, 0x65, 0x12, 0x34, 0x0a,
-	0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0e, 0x32,
-	0x18, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c,
-	0x65, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f,
-	0x63, 0x6f, 0x6c, 0x12, 0x3e, 0x0a, 0x0f, 0x64, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x50, 0x6f, 0x72, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d,
-	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e,
-	0x66, 0x6f, 0x52, 0x0f, 0x64, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50,
-	0x6f, 0x72, 0x74, 0x12, 0x2c, 0x0a, 0x11, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65,
-	0x64, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x11,
-	0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73,
-	0x73, 0x12, 0x3c, 0x0a, 0x0e, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x50,
-	0x6f, 0x72, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
+	0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x08, 0x70, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66,
+	0x6f, 0x12, 0x1c, 0x0a, 0x09, 0x69, 0x73, 0x44, 0x79, 0x6e, 0x61, 0x6d, 0x69, 0x63, 0x18, 0x06,
+	0x20, 0x01, 0x28, 0x08, 0x52, 0x09, 0x69, 0x73, 0x44, 0x79, 0x6e, 0x61, 0x6d, 0x69, 0x63, 0x12,
+	0x18, 0x0a, 0x07, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x18, 0x07, 0x20, 0x03, 0x28, 0x09,
+	0x52, 0x07, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x73, 0x12, 0x26, 0x0a, 0x0e, 0x63, 0x75, 0x73,
+	0x74, 0x6f, 0x6d, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x08, 0x20, 0x01, 0x28,
+	0x0d, 0x52, 0x0e, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f,
+	0x6c, 0x12, 0x1a, 0x0a, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x18, 0x09, 0x20,
+	0x01, 0x28, 0x0c, 0x52, 0x08, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x49, 0x44, 0x12, 0x18, 0x0a,
+	0x07, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x49, 0x44, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07,
+	0x52, 0x6f, 0x75, 0x74, 0x65, 0x49, 0x44, 0x22, 0xf2, 0x01, 0x0a, 0x0e, 0x46, 0x6f, 0x72, 0x77,
+	0x61, 0x72, 0x64, 0x69, 0x6e, 0x67, 0x52, 0x75, 0x6c, 0x65, 0x12, 0x34, 0x0a, 0x08, 0x70, 0x72,
+	0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x18, 0x2e, 0x6d,
+	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x75, 0x6c, 0x65, 0x50, 0x72,
+	0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c,
+	0x12, 0x3e, 0x0a, 0x0f, 0x64, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50,
+	0x6f, 0x72, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
 	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52,
-	0x0e, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x50, 0x6f, 0x72, 0x74, 0x2a,
-	0x4c, 0x0a, 0x0c, 0x52, 0x75, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12,
-	0x0b, 0x0a, 0x07, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x07, 0x0a, 0x03,
-	0x41, 0x4c, 0x4c, 0x10, 0x01, 0x12, 0x07, 0x0a, 0x03, 0x54, 0x43, 0x50, 0x10, 0x02, 0x12, 0x07,
-	0x0a, 0x03, 0x55, 0x44, 0x50, 0x10, 0x03, 0x12, 0x08, 0x0a, 0x04, 0x49, 0x43, 0x4d, 0x50, 0x10,
-	0x04, 0x12, 0x0a, 0x0a, 0x06, 0x43, 0x55, 0x53, 0x54, 0x4f, 0x4d, 0x10, 0x05, 0x2a, 0x20, 0x0a,
-	0x0d, 0x52, 0x75, 0x6c, 0x65, 0x44, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x06,
-	0x0a, 0x02, 0x49, 0x4e, 0x10, 0x00, 0x12, 0x07, 0x0a, 0x03, 0x4f, 0x55, 0x54, 0x10, 0x01, 0x2a,
-	0x22, 0x0a, 0x0a, 0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x0a, 0x0a,
-	0x06, 0x41, 0x43, 0x43, 0x45, 0x50, 0x54, 0x10, 0x00, 0x12, 0x08, 0x0a, 0x04, 0x44, 0x52, 0x4f,
-	0x50, 0x10, 0x01, 0x32, 0xcd, 0x04, 0x0a, 0x11, 0x4d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65,
-	0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x45, 0x0a, 0x05, 0x4c, 0x6f, 0x67,
-	0x69, 0x6e, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e,
-	0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65,
-	0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e,
-	0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00,
-	0x12, 0x46, 0x0a, 0x04, 0x53, 0x79, 0x6e, 0x63, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67,
-	0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d,
-	0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d,
-	0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73,
-	0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x42, 0x0a, 0x0c, 0x47, 0x65, 0x74, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x12, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67,
-	0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x1a, 0x1d, 0x2e, 0x6d, 0x61,
-	0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x4b,
-	0x65, 0x79, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x33, 0x0a, 0x09,
-	0x69, 0x73, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x79, 0x12, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
-	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x1a, 0x11, 0x2e, 0x6d,
-	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x22,
-	0x00, 0x12, 0x5a, 0x0a, 0x1a, 0x47, 0x65, 0x74, 0x44, 0x65, 0x76, 0x69, 0x63, 0x65, 0x41, 0x75,
-	0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x46, 0x6c, 0x6f, 0x77, 0x12,
+	0x0f, 0x64, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x72, 0x74,
+	0x12, 0x2c, 0x0a, 0x11, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x41, 0x64,
+	0x64, 0x72, 0x65, 0x73, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x11, 0x74, 0x72, 0x61,
+	0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x12, 0x3c,
+	0x0a, 0x0e, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x50, 0x6f, 0x72, 0x74,
+	0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d,
+	0x65, 0x6e, 0x74, 0x2e, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x0e, 0x74, 0x72,
+	0x61, 0x6e, 0x73, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x50, 0x6f, 0x72, 0x74, 0x2a, 0x4c, 0x0a, 0x0c,
+	0x52, 0x75, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x0b, 0x0a, 0x07,
+	0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x07, 0x0a, 0x03, 0x41, 0x4c, 0x4c,
+	0x10, 0x01, 0x12, 0x07, 0x0a, 0x03, 0x54, 0x43, 0x50, 0x10, 0x02, 0x12, 0x07, 0x0a, 0x03, 0x55,
+	0x44, 0x50, 0x10, 0x03, 0x12, 0x08, 0x0a, 0x04, 0x49, 0x43, 0x4d, 0x50, 0x10, 0x04, 0x12, 0x0a,
+	0x0a, 0x06, 0x43, 0x55, 0x53, 0x54, 0x4f, 0x4d, 0x10, 0x05, 0x2a, 0x20, 0x0a, 0x0d, 0x52, 0x75,
+	0x6c, 0x65, 0x44, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x06, 0x0a, 0x02, 0x49,
+	0x4e, 0x10, 0x00, 0x12, 0x07, 0x0a, 0x03, 0x4f, 0x55, 0x54, 0x10, 0x01, 0x2a, 0x22, 0x0a, 0x0a,
+	0x52, 0x75, 0x6c, 0x65, 0x41, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x0a, 0x0a, 0x06, 0x41, 0x43,
+	0x43, 0x45, 0x50, 0x54, 0x10, 0x00, 0x12, 0x08, 0x0a, 0x04, 0x44, 0x52, 0x4f, 0x50, 0x10, 0x01,
+	0x32, 0xcd, 0x04, 0x0a, 0x11, 0x4d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x53,
+	0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x45, 0x0a, 0x05, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x12,
 	0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63,
 	0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e,
 	0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79,
-	0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x58, 0x0a,
-	0x18, 0x47, 0x65, 0x74, 0x50, 0x4b, 0x43, 0x45, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x46, 0x6c, 0x6f, 0x77, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
-	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64,
-	0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
-	0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65,
-	0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x3d, 0x0a, 0x08, 0x53, 0x79, 0x6e, 0x63, 0x4d,
-	0x65, 0x74, 0x61, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74,
+	0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x46, 0x0a,
+	0x04, 0x53, 0x79, 0x6e, 0x63, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65,
+	0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73,
+	0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74,
 	0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67,
-	0x65, 0x1a, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45,
-	0x6d, 0x70, 0x74, 0x79, 0x22, 0x00, 0x12, 0x3b, 0x0a, 0x06, 0x4c, 0x6f, 0x67, 0x6f, 0x75, 0x74,
+	0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x42, 0x0a, 0x0c, 0x47, 0x65, 0x74, 0x53, 0x65, 0x72, 0x76,
+	0x65, 0x72, 0x4b, 0x65, 0x79, 0x12, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65,
+	0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x1a, 0x1d, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67,
+	0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x52,
+	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x33, 0x0a, 0x09, 0x69, 0x73, 0x48,
+	0x65, 0x61, 0x6c, 0x74, 0x68, 0x79, 0x12, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d,
+	0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x1a, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
+	0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x22, 0x00, 0x12, 0x5a,
+	0x0a, 0x1a, 0x47, 0x65, 0x74, 0x44, 0x65, 0x76, 0x69, 0x63, 0x65, 0x41, 0x75, 0x74, 0x68, 0x6f,
+	0x72, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x46, 0x6c, 0x6f, 0x77, 0x12, 0x1c, 0x2e, 0x6d,
+	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70,
+	0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e,
+	0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65,
+	0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x58, 0x0a, 0x18, 0x47, 0x65,
+	0x74, 0x50, 0x4b, 0x43, 0x45, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x61, 0x74, 0x69,
+	0x6f, 0x6e, 0x46, 0x6c, 0x6f, 0x77, 0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d,
+	0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73,
+	0x73, 0x61, 0x67, 0x65, 0x1a, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e,
+	0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61,
+	0x67, 0x65, 0x22, 0x00, 0x12, 0x3d, 0x0a, 0x08, 0x53, 0x79, 0x6e, 0x63, 0x4d, 0x65, 0x74, 0x61,
 	0x12, 0x1c, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e,
 	0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x11,
 	0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74,
-	0x79, 0x22, 0x00, 0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70,
-	0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x79, 0x22, 0x00, 0x12, 0x3b, 0x0a, 0x06, 0x4c, 0x6f, 0x67, 0x6f, 0x75, 0x74, 0x12, 0x1c, 0x2e,
+	0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79,
+	0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x11, 0x2e, 0x6d, 0x61,
+	0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x22, 0x00,
+	0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
+	0x6f, 0x33,
 }
 
 var (

shared/management/proto/management.proto

@@ -433,6 +433,8 @@ message DNSConfig {
 message CustomZone {
   string Domain = 1;
   repeated SimpleRecord Records = 2;
+  bool SearchDomainDisabled = 3;
+  bool SkipPTRProcess = 4;
 }
 
 // SimpleRecord represents a dns.SimpleRecord