Repository: netbirdio/netbird

Issue:
Filter-by-status for relayed & P2P state
I'd like to filter the detailed status for all peers which are connected via relay or P2P
Example: `netbird status -d --filter-by-status P2P`

Discussion:
@Silex: In the meantime you can use this:

``` shell
netbird status --json | jq '.peers.details[] | select(.connectionType == "P2P")'
```@florian-obradovic: @Silex awesome - thank you very much!
`netbird status --json | jq -r '.peers.details[] | select(.connectionType == "P2P") | [.fqdn, .netbirdIp, .status, .connectionType] | @tsv' | column -t`

or 

`netbird status --json | jq -r '.peers.details[] | select(.connectionType == "Relayed") | [.fqdn, .netbirdIp, .status, .connectionType] | @tsv' | column -t
`

![Image](https://github.com/user-attachments/assets/92b11e5e-5b65-4fe1-b9bd-e8cbcf1a5307)

Solution patch:
diff --git a/.github/workflows/git-town.yml b/.github/workflows/git-town.yml
index c54fcb44..699ed7d9 100644
--- a/.github/workflows/git-town.yml
+++ b/.github/workflows/git-town.yml
@@ -16,6 +16,6 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: git-town/action@v1
+      - uses: git-town/action@v1.2.1
         with:
-          skip-single-stacks: true
\ No newline at end of file
+          skip-single-stacks: true
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 00898ab2..44e02f45 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -9,7 +9,7 @@ on:
   pull_request:
 
 env:
-  SIGN_PIPE_VER: "v0.0.20"
+  SIGN_PIPE_VER: "v0.0.21"
   GORELEASER_VER: "v2.3.2"
   PRODUCT_NAME: "NetBird"
   COPYRIGHT: "NetBird GmbH"
@@ -231,3 +231,17 @@ jobs:
           ref: ${{ env.SIGN_PIPE_VER }}
           token: ${{ secrets.SIGN_GITHUB_TOKEN }}
           inputs: '{ "tag": "${{ github.ref }}", "skipRelease": false }'
+
+  post_on_forum:
+    runs-on: ubuntu-latest
+    continue-on-error: true
+    needs: [trigger_signer]
+    steps:
+      - uses: Codixer/discourse-topic-github-release-action@v2.0.1
+        with:
+          discourse-api-key: ${{ secrets.DISCOURSE_RELEASES_API_KEY }}
+          discourse-base-url: https://forum.netbird.io
+          discourse-author-username: NetBird
+          discourse-category: 17
+          discourse-tags:
+            releases          
diff --git a/client/cmd/debug.go b/client/cmd/debug.go
index 4036bb8f..3f13a0c3 100644
--- a/client/cmd/debug.go
+++ b/client/cmd/debug.go
@@ -307,7 +307,7 @@ func getStatusOutput(cmd *cobra.Command, anon bool) string {
 		cmd.PrintErrf("Failed to get status: %v\n", err)
 	} else {
 		statusOutputString = nbstatus.ParseToFullDetailSummary(
-			nbstatus.ConvertToStatusOutputOverview(statusResp, anon, "", nil, nil, nil),
+			nbstatus.ConvertToStatusOutputOverview(statusResp, anon, "", nil, nil, nil, ""),
 		)
 	}
 	return statusOutputString
diff --git a/client/cmd/status.go b/client/cmd/status.go
index b108ca57..2d6e41bc 100644
--- a/client/cmd/status.go
+++ b/client/cmd/status.go
@@ -26,6 +26,7 @@ var (
 	statusFilter         string
 	ipsFilterMap         map[string]struct{}
 	prefixNamesFilterMap map[string]struct{}
+	connectionTypeFilter string
 )
 
 var statusCmd = &cobra.Command{
@@ -45,6 +46,7 @@ func init() {
 	statusCmd.PersistentFlags().StringSliceVar(&ipsFilter, "filter-by-ips", []string{}, "filters the detailed output by a list of one or more IPs, e.g., --filter-by-ips 100.64.0.100,100.64.0.200")
 	statusCmd.PersistentFlags().StringSliceVar(&prefixNamesFilter, "filter-by-names", []string{}, "filters the detailed output by a list of one or more peer FQDN or hostnames, e.g., --filter-by-names peer-a,peer-b.netbird.cloud")
 	statusCmd.PersistentFlags().StringVar(&statusFilter, "filter-by-status", "", "filters the detailed output by connection status(idle|connecting|connected), e.g., --filter-by-status connected")
+	statusCmd.PersistentFlags().StringVar(&connectionTypeFilter, "filter-by-connection-type", "", "filters the detailed output by connection type (P2P|Relayed), e.g., --filter-by-connection-type P2P")
 }
 
 func statusFunc(cmd *cobra.Command, args []string) error {
@@ -89,7 +91,7 @@ func statusFunc(cmd *cobra.Command, args []string) error {
 		return nil
 	}
 
-	var outputInformationHolder = nbstatus.ConvertToStatusOutputOverview(resp, anonymizeFlag, statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilterMap)
+	var outputInformationHolder = nbstatus.ConvertToStatusOutputOverview(resp, anonymizeFlag, statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilterMap, connectionTypeFilter)
 	var statusOutputString string
 	switch {
 	case detailFlag:
@@ -156,6 +158,15 @@ func parseFilters() error {
 		enableDetailFlagWhenFilterFlag()
 	}
 
+	switch strings.ToLower(connectionTypeFilter) {
+	case "", "p2p", "relayed":
+		if strings.ToLower(connectionTypeFilter) != "" {
+			enableDetailFlagWhenFilterFlag()
+		}
+	default:
+		return fmt.Errorf("wrong connection-type filter, should be one of P2P|Relayed, got: %s", connectionTypeFilter)
+	}
+
 	return nil
 }
 
diff --git a/client/proto/daemon.pb.go b/client/proto/daemon.pb.go
index 26e58d18..753aa62d 100644
--- a/client/proto/daemon.pb.go
+++ b/client/proto/daemon.pb.go
@@ -1330,6 +1330,13 @@ func (x *PeerState) GetRelayAddress() string {
 	return ""
 }
 
+func (x *PeerState) GetConnectionType() string {
+	if x.Relayed {
+		return "Relayed"
+	}
+	return "P2P"
+}
+
 // LocalPeerState contains the latest state of the local peer
 type LocalPeerState struct {
 	state               protoimpl.MessageState `protogen:"open.v1"`
diff --git a/client/status/status.go b/client/status/status.go
index 18056e36..507c7ea8 100644
--- a/client/status/status.go
+++ b/client/status/status.go
@@ -100,7 +100,7 @@ type OutputOverview struct {
 	LazyConnectionEnabled   bool                       `json:"lazyConnectionEnabled" yaml:"lazyConnectionEnabled"`
 }
 
-func ConvertToStatusOutputOverview(resp *proto.StatusResponse, anon bool, statusFilter string, prefixNamesFilter []string, prefixNamesFilterMap map[string]struct{}, ipsFilter map[string]struct{}) OutputOverview {
+func ConvertToStatusOutputOverview(resp *proto.StatusResponse, anon bool, statusFilter string, prefixNamesFilter []string, prefixNamesFilterMap map[string]struct{}, ipsFilter map[string]struct{}, connectionTypeFilter string) OutputOverview {
 	pbFullStatus := resp.GetFullStatus()
 
 	managementState := pbFullStatus.GetManagementState()
@@ -118,7 +118,7 @@ func ConvertToStatusOutputOverview(resp *proto.StatusResponse, anon bool, status
 	}
 
 	relayOverview := mapRelays(pbFullStatus.GetRelays())
-	peersOverview := mapPeers(resp.GetFullStatus().GetPeers(), statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilter)
+	peersOverview := mapPeers(resp.GetFullStatus().GetPeers(), statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilter, connectionTypeFilter)
 
 	overview := OutputOverview{
 		Peers:                   peersOverview,
@@ -193,6 +193,7 @@ func mapPeers(
 	prefixNamesFilter []string,
 	prefixNamesFilterMap map[string]struct{},
 	ipsFilter map[string]struct{},
+	connectionTypeFilter string,
 ) PeersStateOutput {
 	var peersStateDetail []PeerStateDetailOutput
 	peersConnected := 0
@@ -208,7 +209,7 @@ func mapPeers(
 		transferSent := int64(0)
 
 		isPeerConnected := pbPeerState.ConnStatus == peer.StatusConnected.String()
-		if skipDetailByFilters(pbPeerState, pbPeerState.ConnStatus, statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilter) {
+		if skipDetailByFilters(pbPeerState, pbPeerState.ConnStatus, statusFilter, prefixNamesFilter, prefixNamesFilterMap, ipsFilter, connectionTypeFilter) {
 			continue
 		}
 		if isPeerConnected {
@@ -218,10 +219,7 @@ func mapPeers(
 			remoteICE = pbPeerState.GetRemoteIceCandidateType()
 			localICEEndpoint = pbPeerState.GetLocalIceCandidateEndpoint()
 			remoteICEEndpoint = pbPeerState.GetRemoteIceCandidateEndpoint()
-			connType = "P2P"
-			if pbPeerState.Relayed {
-				connType = "Relayed"
-			}
+			connType = pbPeerState.GetConnectionType()
 			relayServerAddress = pbPeerState.GetRelayAddress()
 			lastHandshake = pbPeerState.GetLastWireguardHandshake().AsTime().Local()
 			transferReceived = pbPeerState.GetBytesRx()
@@ -542,10 +540,11 @@ func parsePeers(peers PeersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 	return peersString
 }
 
-func skipDetailByFilters(peerState *proto.PeerState, peerStatus string, statusFilter string, prefixNamesFilter []string, prefixNamesFilterMap map[string]struct{}, ipsFilter map[string]struct{}) bool {
+func skipDetailByFilters(peerState *proto.PeerState, peerStatus string, statusFilter string, prefixNamesFilter []string, prefixNamesFilterMap map[string]struct{}, ipsFilter map[string]struct{}, connectionTypeFilter string) bool {
 	statusEval := false
 	ipEval := false
 	nameEval := true
+	connectionTypeEval := false
 
 	if statusFilter != "" {
 		if !strings.EqualFold(peerStatus, statusFilter) {
@@ -570,8 +569,11 @@ func skipDetailByFilters(peerState *proto.PeerState, peerStatus string, statusFi
 	} else {
 		nameEval = false
 	}
+	if connectionTypeFilter != "" && !strings.EqualFold(peerState.GetConnectionType(), connectionTypeFilter) {
+		connectionTypeEval = true
+	}
 
-	return statusEval || ipEval || nameEval
+	return statusEval || ipEval || nameEval || connectionTypeEval
 }
 
 func toIEC(b int64) string {
diff --git a/client/ui/debug.go b/client/ui/debug.go
index ab7dba37..55829de1 100644
--- a/client/ui/debug.go
+++ b/client/ui/debug.go
@@ -433,7 +433,7 @@ func (s *serviceClient) collectDebugData(
 
 	var postUpStatusOutput string
 	if postUpStatus != nil {
-		overview := nbstatus.ConvertToStatusOutputOverview(postUpStatus, params.anonymize, "", nil, nil, nil)
+		overview := nbstatus.ConvertToStatusOutputOverview(postUpStatus, params.anonymize, "", nil, nil, nil, "")
 		postUpStatusOutput = nbstatus.ParseToFullDetailSummary(overview)
 	}
 	headerPostUp := fmt.Sprintf("----- NetBird post-up - Timestamp: %s", time.Now().Format(time.RFC3339))
@@ -450,7 +450,7 @@ func (s *serviceClient) collectDebugData(
 
 	var preDownStatusOutput string
 	if preDownStatus != nil {
-		overview := nbstatus.ConvertToStatusOutputOverview(preDownStatus, params.anonymize, "", nil, nil, nil)
+		overview := nbstatus.ConvertToStatusOutputOverview(preDownStatus, params.anonymize, "", nil, nil, nil, "")
 		preDownStatusOutput = nbstatus.ParseToFullDetailSummary(overview)
 	}
 	headerPreDown := fmt.Sprintf("----- NetBird pre-down - Timestamp: %s - Duration: %s",
@@ -581,7 +581,7 @@ func (s *serviceClient) createDebugBundle(anonymize bool, systemInfo bool, uploa
 
 	var statusOutput string
 	if statusResp != nil {
-		overview := nbstatus.ConvertToStatusOutputOverview(statusResp, anonymize, "", nil, nil, nil)
+		overview := nbstatus.ConvertToStatusOutputOverview(statusResp, anonymize, "", nil, nil, nil, "")
 		statusOutput = nbstatus.ParseToFullDetailSummary(overview)
 	}
 
diff --git a/go.mod b/go.mod
index 4a972737..cf2a2375 100644
--- a/go.mod
+++ b/go.mod
@@ -63,7 +63,7 @@ require (
 	github.com/miekg/dns v1.1.59
 	github.com/mitchellh/hashstructure/v2 v2.0.2
 	github.com/nadoo/ipset v0.5.0
-	github.com/netbirdio/management-integrations/integrations v0.0.0-20250612164546-6bd7e2338d65
+	github.com/netbirdio/management-integrations/integrations v0.0.0-20250718071730-f4d133556ff5
 	github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250514131221-a464fd5f30cb
 	github.com/okta/okta-sdk-golang/v2 v2.18.0
 	github.com/oschwald/maxminddb-golang v1.12.0
diff --git a/go.sum b/go.sum
index a622f203..699a832d 100644
--- a/go.sum
+++ b/go.sum
@@ -503,8 +503,8 @@ github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944 h1:TDtJKmM6S
 github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944/go.mod h1:sHA6TRxjQ6RLbnI+3R4DZo2Eseg/iKiPRfNmcuNySVQ=
 github.com/netbirdio/ice/v3 v3.0.0-20240315174635-e72a50fcb64e h1:PURA50S8u4mF6RrkYYCAvvPCixhqqEiEy3Ej6avh04c=
 github.com/netbirdio/ice/v3 v3.0.0-20240315174635-e72a50fcb64e/go.mod h1:YMLU7qbKfVjmEv7EoZPIVEI+kNYxWCdPK3VS0BU+U4Q=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20250612164546-6bd7e2338d65 h1:5OfYiLjpr4dbQYJI5ouZaylkVdi2KlErLFOwBeBo5Hw=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20250612164546-6bd7e2338d65/go.mod h1:Gi9raplYzCCyh07Olw/DVfCJTFgpr1WCXJ/Q+8TSA9Q=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20250718071730-f4d133556ff5 h1:Zfn8d83OVyELCdxgprcyXR3D8uqoxHtXE9PUxVXDx/w=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20250718071730-f4d133556ff5/go.mod h1:Gi9raplYzCCyh07Olw/DVfCJTFgpr1WCXJ/Q+8TSA9Q=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502 h1:3tHlFmhTdX9axERMVN63dqyFqnvuD+EMJHzM7mNGON8=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250514131221-a464fd5f30cb h1:Cr6age+ePALqlSvtp7wc6lYY97XN7rkD1K4XEDmY+TU=
diff --git a/management/cmd/management.go b/management/cmd/management.go
index 878e4c39..24c260e9 100644
--- a/management/cmd/management.go
+++ b/management/cmd/management.go
@@ -292,7 +292,7 @@ var (
 			ephemeralManager.LoadInitialPeers(ctx)
 
 			gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
-			srv, err := server.NewServer(ctx, config, accountManager, settingsManager, peersUpdateManager, secretsManager, appMetrics, ephemeralManager, authManager)
+			srv, err := server.NewServer(ctx, config, accountManager, settingsManager, peersUpdateManager, secretsManager, appMetrics, ephemeralManager, authManager, integratedPeerValidator)
 			if err != nil {
 				return fmt.Errorf("failed creating gRPC API handler: %v", err)
 			}
diff --git a/management/server/grpcserver.go b/management/server/grpcserver.go
index 2b27f9e0..2f1bc367 100644
--- a/management/server/grpcserver.go
+++ b/management/server/grpcserver.go
@@ -19,6 +19,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
+	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/management/proto"
@@ -40,13 +41,14 @@ type GRPCServer struct {
 	settingsManager settings.Manager
 	wgKey           wgtypes.Key
 	proto.UnimplementedManagementServiceServer
-	peersUpdateManager *PeersUpdateManager
-	config             *types.Config
-	secretsManager     SecretsManager
-	appMetrics         telemetry.AppMetrics
-	ephemeralManager   *EphemeralManager
-	peerLocks          sync.Map
-	authManager        auth.Manager
+	peersUpdateManager      *PeersUpdateManager
+	config                  *types.Config
+	secretsManager          SecretsManager
+	appMetrics              telemetry.AppMetrics
+	ephemeralManager        *EphemeralManager
+	peerLocks               sync.Map
+	authManager             auth.Manager
+	integratedPeerValidator integrated_validator.IntegratedValidator
 }
 
 // NewServer creates a new Management server
@@ -60,6 +62,7 @@ func NewServer(
 	appMetrics telemetry.AppMetrics,
 	ephemeralManager *EphemeralManager,
 	authManager auth.Manager,
+	integratedPeerValidator integrated_validator.IntegratedValidator,
 ) (*GRPCServer, error) {
 	key, err := wgtypes.GeneratePrivateKey()
 	if err != nil {
@@ -79,14 +82,15 @@ func NewServer(
 	return &GRPCServer{
 		wgKey: key,
 		// peerKey -> event channel
-		peersUpdateManager: peersUpdateManager,
-		accountManager:     accountManager,
-		settingsManager:    settingsManager,
-		config:             config,
-		secretsManager:     secretsManager,
-		authManager:        authManager,
-		appMetrics:         appMetrics,
-		ephemeralManager:   ephemeralManager,
+		peersUpdateManager:      peersUpdateManager,
+		accountManager:          accountManager,
+		settingsManager:         settingsManager,
+		config:                  config,
+		secretsManager:          secretsManager,
+		authManager:             authManager,
+		appMetrics:              appMetrics,
+		ephemeralManager:        ephemeralManager,
+		integratedPeerValidator: integratedPeerValidator,
 	}, nil
 }
 
@@ -850,7 +854,7 @@ func (s *GRPCServer) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.En
 		return nil, status.Error(codes.NotFound, "no pkce authorization flow information available")
 	}
 
-	flowInfoResp := &proto.PKCEAuthorizationFlow{
+	initInfoFlow := &proto.PKCEAuthorizationFlow{
 		ProviderConfig: &proto.ProviderConfig{
 			Audience:              s.config.PKCEAuthorizationFlow.ProviderConfig.Audience,
 			ClientID:              s.config.PKCEAuthorizationFlow.ProviderConfig.ClientID,
@@ -865,6 +869,8 @@ func (s *GRPCServer) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.En
 		},
 	}
 
+	flowInfoResp := s.integratedPeerValidator.ValidateFlowResponse(ctx, peerKey.String(), initInfoFlow)
+
 	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, flowInfoResp)
 	if err != nil {
 		return nil, status.Error(codes.Internal, "failed to encrypt no pkce authorization flow information")
diff --git a/management/server/integrated_validator.go b/management/server/integrated_validator.go
index cfde7c61..e3e47441 100644
--- a/management/server/integrated_validator.go
+++ b/management/server/integrated_validator.go
@@ -6,6 +6,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -101,22 +102,23 @@ func (am *DefaultAccountManager) GetValidatedPeers(ctx context.Context, accountI
 	return am.integratedPeerValidator.GetValidatedPeers(accountID, groups, peers, settings.Extra)
 }
 
-type MocIntegratedValidator struct {
+type MockIntegratedValidator struct {
+	integrated_validator.IntegratedValidator
 	ValidatePeerFunc func(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *types.ExtraSettings) (*nbpeer.Peer, bool, error)
 }
 
-func (a MocIntegratedValidator) ValidateExtraSettings(_ context.Context, newExtraSettings *types.ExtraSettings, oldExtraSettings *types.ExtraSettings, peers map[string]*nbpeer.Peer, userID string, accountID string) error {
+func (a MockIntegratedValidator) ValidateExtraSettings(_ context.Context, newExtraSettings *types.ExtraSettings, oldExtraSettings *types.ExtraSettings, peers map[string]*nbpeer.Peer, userID string, accountID string) error {
 	return nil
 }
 
-func (a MocIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *types.ExtraSettings) (*nbpeer.Peer, bool, error) {
+func (a MockIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *types.ExtraSettings) (*nbpeer.Peer, bool, error) {
 	if a.ValidatePeerFunc != nil {
 		return a.ValidatePeerFunc(context.Background(), update, peer, userID, accountID, dnsDomain, peersGroup, extraSettings)
 	}
 	return update, false, nil
 }
 
-func (a MocIntegratedValidator) GetValidatedPeers(accountID string, groups []*types.Group, peers []*nbpeer.Peer, extraSettings *types.ExtraSettings) (map[string]struct{}, error) {
+func (a MockIntegratedValidator) GetValidatedPeers(accountID string, groups []*types.Group, peers []*nbpeer.Peer, extraSettings *types.ExtraSettings) (map[string]struct{}, error) {
 	validatedPeers := make(map[string]struct{})
 	for _, peer := range peers {
 		validatedPeers[peer.ID] = struct{}{}
@@ -124,22 +126,22 @@ func (a MocIntegratedValidator) GetValidatedPeers(accountID string, groups []*ty
 	return validatedPeers, nil
 }
 
-func (MocIntegratedValidator) PreparePeer(_ context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *types.ExtraSettings) *nbpeer.Peer {
+func (MockIntegratedValidator) PreparePeer(_ context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *types.ExtraSettings) *nbpeer.Peer {
 	return peer
 }
 
-func (MocIntegratedValidator) IsNotValidPeer(_ context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *types.ExtraSettings) (bool, bool, error) {
+func (MockIntegratedValidator) IsNotValidPeer(_ context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *types.ExtraSettings) (bool, bool, error) {
 	return false, false, nil
 }
 
-func (MocIntegratedValidator) PeerDeleted(_ context.Context, _, _ string) error {
+func (MockIntegratedValidator) PeerDeleted(_ context.Context, _, _ string) error {
 	return nil
 }
 
-func (MocIntegratedValidator) SetPeerInvalidationListener(func(accountID string)) {
+func (MockIntegratedValidator) SetPeerInvalidationListener(func(accountID string)) {
 	// just a dummy
 }
 
-func (MocIntegratedValidator) Stop(_ context.Context) {
+func (MockIntegratedValidator) Stop(_ context.Context) {
 	// just a dummy
 }
diff --git a/management/server/integrations/integrated_validator/interface.go b/management/server/integrations/integrated_validator/interface.go
index 083baa65..245c0168 100644
--- a/management/server/integrations/integrated_validator/interface.go
+++ b/management/server/integrations/integrated_validator/interface.go
@@ -3,6 +3,7 @@ package integrated_validator
 import (
 	"context"
 
+	"github.com/netbirdio/netbird/management/proto"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/types"
 )
@@ -17,4 +18,5 @@ type IntegratedValidator interface {
 	PeerDeleted(ctx context.Context, accountID, peerID string) error
 	SetPeerInvalidationListener(fn func(accountID string))
 	Stop(ctx context.Context)
+	ValidateFlowResponse(ctx context.Context, peerKey string, flowResponse *proto.PKCEAuthorizationFlow) *proto.PKCEAuthorizationFlow
 }
diff --git a/management/server/migration/migration.go b/management/server/migration/migration.go
index ab11be73..c2f1a5ab 100644
--- a/management/server/migration/migration.go
+++ b/management/server/migration/migration.go
@@ -283,7 +283,7 @@ func MigrateSetupKeyToHashedSetupKey[T any](ctx context.Context, db *gorm.DB) er
 			}
 		}
 
-		if err := tx.Exec(fmt.Sprintf("ALTER TABLE %s DROP COLUMN %s", "peers", "setup_key")).Error; err != nil {
+		if err := tx.Exec(fmt.Sprintf("ALTER TABLE %s DROP COLUMN IF EXISTS %s", "peers", "setup_key")).Error; err != nil {
 			log.WithContext(ctx).Errorf("Failed to drop column %s: %v", "setup_key", err)
 		}
 
@@ -377,6 +377,11 @@ func DropIndex[T any](ctx context.Context, db *gorm.DB, indexName string) error
 func CreateIndexIfNotExists[T any](ctx context.Context, db *gorm.DB, indexName string, columns ...string) error {
 	var model T
 
+	if !db.Migrator().HasTable(&model) {
+		log.WithContext(ctx).Debugf("table for %T does not exist, no migration needed", model)
+		return nil
+	}
+
 	stmt := &gorm.Statement{DB: db}
 	if err := stmt.Parse(&model); err != nil {
 		return fmt.Errorf("failed to parse model schema: %w", err)
@@ -384,6 +389,11 @@ func CreateIndexIfNotExists[T any](ctx context.Context, db *gorm.DB, indexName s
 	tableName := stmt.Schema.Table
 	dialect := db.Dialector.Name()
 
+	if db.Migrator().HasIndex(&model, indexName) {
+		log.WithContext(ctx).Infof("index %s already exists on table %s", indexName, tableName)
+		return nil
+	}
+
 	var columnClause string
 	if dialect == "mysql" {
 		var withLength []string
diff --git a/relay/client/client.go b/relay/client/client.go
index 2bf679ec..32dfbb4d 100644
--- a/relay/client/client.go
+++ b/relay/client/client.go
@@ -181,13 +181,17 @@ func (c *Client) Connect(ctx context.Context) error {
 		return nil
 	}
 
-	if err := c.connect(ctx); err != nil {
+	instanceURL, err := c.connect(ctx)
+	if err != nil {
 		return err
 	}
+	c.muInstanceURL.Lock()
+	c.instanceURL = instanceURL
+	c.muInstanceURL.Unlock()
 
 	c.stateSubscription = NewPeersStateSubscription(c.log, c.relayConn, c.closeConnsByPeerID)
 
-	c.log = c.log.WithField("relay", c.instanceURL.String())
+	c.log = c.log.WithField("relay", instanceURL.String())
 	c.log.Infof("relay connection established")
 
 	c.serviceIsRunning = true
@@ -229,9 +233,18 @@ func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, erro
 
 	c.log.Infof("remote peer is available, prepare the relayed connection: %s", peerID)
 	msgChannel := make(chan Msg, 100)
-	conn := NewConn(c, peerID, msgChannel, c.instanceURL)
 
 	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+
+	c.muInstanceURL.Lock()
+	instanceURL := c.instanceURL
+	c.muInstanceURL.Unlock()
+	conn := NewConn(c, peerID, msgChannel, instanceURL)
+
 	_, ok = c.conns[peerID]
 	if ok {
 		c.mu.Unlock()
@@ -278,69 +291,67 @@ func (c *Client) Close() error {
 	return c.close(true)
 }
 
-func (c *Client) connect(ctx context.Context) error {
+func (c *Client) connect(ctx context.Context) (*RelayAddr, error) {
 	rd := dialer.NewRaceDial(c.log, c.connectionURL, quic.Dialer{}, ws.Dialer{})
 	conn, err := rd.Dial()
 	if err != nil {
-		return err
+		return nil, err
 	}
 	c.relayConn = conn
 
-	if err = c.handShake(ctx); err != nil {
+	instanceURL, err := c.handShake(ctx)
+	if err != nil {
 		cErr := conn.Close()
 		if cErr != nil {
 			c.log.Errorf("failed to close connection: %s", cErr)
 		}
-		return err
+		return nil, err
 	}
 
-	return nil
+	return instanceURL, nil
 }
 
-func (c *Client) handShake(ctx context.Context) error {
+func (c *Client) handShake(ctx context.Context) (*RelayAddr, error) {
 	msg, err := messages.MarshalAuthMsg(c.hashedID, c.authTokenStore.TokenBinary())
 	if err != nil {
 		c.log.Errorf("failed to marshal auth message: %s", err)
-		return err
+		return nil, err
 	}
 
 	_, err = c.relayConn.Write(msg)
 	if err != nil {
 		c.log.Errorf("failed to send auth message: %s", err)
-		return err
+		return nil, err
 	}
 	buf := make([]byte, messages.MaxHandshakeRespSize)
 	n, err := c.readWithTimeout(ctx, buf)
 	if err != nil {
 		c.log.Errorf("failed to read auth response: %s", err)
-		return err
+		return nil, err
 	}
 
 	_, err = messages.ValidateVersion(buf[:n])
 	if err != nil {
-		return fmt.Errorf("validate version: %w", err)
+		return nil, fmt.Errorf("validate version: %w", err)
 	}
 
 	msgType, err := messages.DetermineServerMessageType(buf[:n])
 	if err != nil {
 		c.log.Errorf("failed to determine message type: %s", err)
-		return err
+		return nil, err
 	}
 
 	if msgType != messages.MsgTypeAuthResponse {
 		c.log.Errorf("unexpected message type: %s", msgType)
-		return fmt.Errorf("unexpected message type")
+		return nil, fmt.Errorf("unexpected message type")
 	}
 
 	addr, err := messages.UnmarshalAuthResponse(buf[:n])
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	c.muInstanceURL.Lock()
-	c.instanceURL = &RelayAddr{addr: addr}
-	c.muInstanceURL.Unlock()
-	return nil
+	return &RelayAddr{addr: addr}, nil
 }
 
 func (c *Client) readLoop(hc *healthcheck.Receiver, relayConn net.Conn, internallyStoppedFlag *internalStopFlag) {
@@ -386,10 +397,6 @@ func (c *Client) readLoop(hc *healthcheck.Receiver, relayConn net.Conn, internal
 
 	hc.Stop()
 
-	c.muInstanceURL.Lock()
-	c.instanceURL = nil
-	c.muInstanceURL.Unlock()
-
 	c.stateSubscription.Cleanup()
 	c.wgReadLoop.Done()
 	_ = c.close(false)
@@ -578,8 +585,12 @@ func (c *Client) close(gracefullyExit bool) error {
 		c.log.Warn("relay connection was already marked as not running")
 		return nil
 	}
-
 	c.serviceIsRunning = false
+
+	c.muInstanceURL.Lock()
+	c.instanceURL = nil
+	c.muInstanceURL.Unlock()
+
 	c.log.Infof("closing all peer connections")
 	c.closeAllConns()
 	if gracefullyExit {
diff --git a/relay/client/manager.go b/relay/client/manager.go
index 0fb682d9..b97bc0b9 100644
--- a/relay/client/manager.go
+++ b/relay/client/manager.go
@@ -65,7 +65,7 @@ type Manager struct {
 
 	relayClient *Client
 	// the guard logic can overwrite the relayClient variable, this mutex protect the usage of the variable
-	relayClientMu  sync.Mutex
+	relayClientMu  sync.RWMutex
 	reconnectGuard *Guard
 
 	relayClients      map[string]*RelayTrack
@@ -124,8 +124,8 @@ func (m *Manager) Serve() error {
 // established via the relay server. If the peer is on a different relay server, the manager will establish a new
 // connection to the relay server. It returns back with a net.Conn what represent the remote peer connection.
 func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string) (net.Conn, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
 
 	if m.relayClient == nil {
 		return nil, ErrRelayClientNotConnected
@@ -155,8 +155,8 @@ func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string) (
 
 // Ready returns true if the home Relay client is connected to the relay server.
 func (m *Manager) Ready() bool {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
 
 	if m.relayClient == nil {
 		return false
@@ -174,8 +174,8 @@ func (m *Manager) SetOnReconnectedListener(f func()) {
 // AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection
 // closed.
 func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
 
 	if m.relayClient == nil {
 		return ErrRelayClientNotConnected
@@ -199,8 +199,8 @@ func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServ
 // RelayInstanceAddress returns the address of the permanent relay server. It could change if the network connection is
 // lost. This address will be sent to the target peer to choose the common relay server for the communication.
 func (m *Manager) RelayInstanceAddress() (string, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
 
 	if m.relayClient == nil {
 		return "", ErrRelayClientNotConnected
@@ -300,7 +300,9 @@ func (m *Manager) onServerConnected() {
 func (m *Manager) onServerDisconnected(serverAddress string) {
 	m.relayClientMu.Lock()
 	if serverAddress == m.relayClient.connectionURL {
-		go m.reconnectGuard.StartReconnectTrys(m.ctx, m.relayClient)
+		go func(client *Client) {
+			m.reconnectGuard.StartReconnectTrys(m.ctx, client)
+		}(m.relayClient)
 	}
 	m.relayClientMu.Unlock()
 
diff --git a/relay/client/peer_subscription.go b/relay/client/peer_subscription.go
index 03e7127b..85bd41cb 100644
--- a/relay/client/peer_subscription.go
+++ b/relay/client/peer_subscription.go
@@ -3,6 +3,8 @@ package client
 import (
 	"context"
 	"errors"
+	"fmt"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -28,6 +30,7 @@ type PeersStateSubscription struct {
 
 	listenForOfflinePeers map[messages.PeerID]struct{}
 	waitingPeers          map[messages.PeerID]chan struct{}
+	mu                    sync.Mutex // Mutex to protect access to waitingPeers and listenForOfflinePeers
 }
 
 func NewPeersStateSubscription(log *log.Entry, relayConn relayedConnWriter, offlineCallback func(peerIDs []messages.PeerID)) *PeersStateSubscription {
@@ -43,24 +46,31 @@ func NewPeersStateSubscription(log *log.Entry, relayConn relayedConnWriter, offl
 // OnPeersOnline should be called when a notification is received that certain peers have come online.
 // It checks if any of the peers are being waited on and signals their availability.
 func (s *PeersStateSubscription) OnPeersOnline(peersID []messages.PeerID) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	for _, peerID := range peersID {
 		waitCh, ok := s.waitingPeers[peerID]
 		if !ok {
+			// If meanwhile the peer was unsubscribed, we don't need to signal it
 			continue
 		}
 
-		close(waitCh)
+		waitCh <- struct{}{}
 		delete(s.waitingPeers, peerID)
+		close(waitCh)
 	}
 }
 
 func (s *PeersStateSubscription) OnPeersWentOffline(peersID []messages.PeerID) {
+	s.mu.Lock()
 	relevantPeers := make([]messages.PeerID, 0, len(peersID))
 	for _, peerID := range peersID {
 		if _, ok := s.listenForOfflinePeers[peerID]; ok {
 			relevantPeers = append(relevantPeers, peerID)
 		}
 	}
+	s.mu.Unlock()
 
 	if len(relevantPeers) > 0 {
 		s.offlineCallback(relevantPeers)
@@ -68,36 +78,41 @@ func (s *PeersStateSubscription) OnPeersWentOffline(peersID []messages.PeerID) {
 }
 
 // WaitToBeOnlineAndSubscribe waits for a specific peer to come online and subscribes to its state changes.
-// todo: when we unsubscribe while this is running, this will not return with error
 func (s *PeersStateSubscription) WaitToBeOnlineAndSubscribe(ctx context.Context, peerID messages.PeerID) error {
 	// Check if already waiting for this peer
+	s.mu.Lock()
 	if _, exists := s.waitingPeers[peerID]; exists {
+		s.mu.Unlock()
 		return errors.New("already waiting for peer to come online")
 	}
 
 	// Create a channel to wait for the peer to come online
-	waitCh := make(chan struct{})
+	waitCh := make(chan struct{}, 1)
 	s.waitingPeers[peerID] = waitCh
+	s.listenForOfflinePeers[peerID] = struct{}{}
+	s.mu.Unlock()
 
-	if err := s.subscribeStateChange([]messages.PeerID{peerID}); err != nil {
+	if err := s.subscribeStateChange(peerID); err != nil {
 		s.log.Errorf("failed to subscribe to peer state: %s", err)
-		close(waitCh)
-		delete(s.waitingPeers, peerID)
-		return err
-	}
-
-	defer func() {
+		s.mu.Lock()
 		if ch, exists := s.waitingPeers[peerID]; exists && ch == waitCh {
 			close(waitCh)
 			delete(s.waitingPeers, peerID)
+			delete(s.listenForOfflinePeers, peerID)
 		}
-	}()
+		s.mu.Unlock()
+		return err
+	}
 
 	// Wait for peer to come online or context to be cancelled
 	timeoutCtx, cancel := context.WithTimeout(ctx, OpenConnectionTimeout)
 	defer cancel()
 	select {
-	case <-waitCh:
+	case _, ok := <-waitCh:
+		if !ok {
+			return fmt.Errorf("wait for peer to come online has been cancelled")
+		}
+
 		s.log.Debugf("peer %s is now online", peerID)
 		return nil
 	case <-timeoutCtx.Done():
@@ -105,6 +120,13 @@ func (s *PeersStateSubscription) WaitToBeOnlineAndSubscribe(ctx context.Context,
 		if err := s.unsubscribeStateChange([]messages.PeerID{peerID}); err != nil {
 			s.log.Errorf("failed to unsubscribe from peer state: %s", err)
 		}
+		s.mu.Lock()
+		if ch, exists := s.waitingPeers[peerID]; exists && ch == waitCh {
+			close(waitCh)
+			delete(s.waitingPeers, peerID)
+			delete(s.listenForOfflinePeers, peerID)
+		}
+		s.mu.Unlock()
 		return timeoutCtx.Err()
 	}
 }
@@ -112,6 +134,7 @@ func (s *PeersStateSubscription) WaitToBeOnlineAndSubscribe(ctx context.Context,
 func (s *PeersStateSubscription) UnsubscribeStateChange(peerIDs []messages.PeerID) error {
 	msgErr := s.unsubscribeStateChange(peerIDs)
 
+	s.mu.Lock()
 	for _, peerID := range peerIDs {
 		if wch, ok := s.waitingPeers[peerID]; ok {
 			close(wch)
@@ -120,11 +143,15 @@ func (s *PeersStateSubscription) UnsubscribeStateChange(peerIDs []messages.PeerI
 
 		delete(s.listenForOfflinePeers, peerID)
 	}
+	s.mu.Unlock()
 
 	return msgErr
 }
 
 func (s *PeersStateSubscription) Cleanup() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	for _, waitCh := range s.waitingPeers {
 		close(waitCh)
 	}
@@ -133,16 +160,12 @@ func (s *PeersStateSubscription) Cleanup() {
 	s.listenForOfflinePeers = make(map[messages.PeerID]struct{})
 }
 
-func (s *PeersStateSubscription) subscribeStateChange(peerIDs []messages.PeerID) error {
-	msgs, err := messages.MarshalSubPeerStateMsg(peerIDs)
+func (s *PeersStateSubscription) subscribeStateChange(peerID messages.PeerID) error {
+	msgs, err := messages.MarshalSubPeerStateMsg([]messages.PeerID{peerID})
 	if err != nil {
 		return err
 	}
 
-	for _, peer := range peerIDs {
-		s.listenForOfflinePeers[peer] = struct{}{}
-	}
-
 	for _, msg := range msgs {
 		if _, err := s.relayConn.Write(msg); err != nil {
 			return err
diff --git a/relay/server/store/listener.go b/relay/server/store/listener.go
index e5f45579..b7c5f4ce 100644
--- a/relay/server/store/listener.go
+++ b/relay/server/store/listener.go
@@ -8,6 +8,7 @@ import (
 )
 
 type Listener struct {
+	ctx   context.Context
 	store *Store
 
 	onlineChan                chan messages.PeerID
@@ -15,12 +16,11 @@ type Listener struct {
 	interestedPeersForOffline map[messages.PeerID]struct{}
 	interestedPeersForOnline  map[messages.PeerID]struct{}
 	mu                        sync.RWMutex
-
-	listenerCtx context.Context
 }
 
-func newListener(store *Store) *Listener {
+func newListener(ctx context.Context, store *Store) *Listener {
 	l := &Listener{
+		ctx:   ctx,
 		store: store,
 
 		onlineChan:                make(chan messages.PeerID, 244), //244 is the message size limit in the relay protocol
@@ -65,11 +65,10 @@ func (l *Listener) RemoveInterestedPeer(peerIDs []messages.PeerID) {
 	}
 }
 
-func (l *Listener) listenForEvents(ctx context.Context, onPeersComeOnline, onPeersWentOffline func([]messages.PeerID)) {
-	l.listenerCtx = ctx
+func (l *Listener) listenForEvents(onPeersComeOnline, onPeersWentOffline func([]messages.PeerID)) {
 	for {
 		select {
-		case <-ctx.Done():
+		case <-l.ctx.Done():
 			return
 		case pID := <-l.onlineChan:
 			peers := make([]messages.PeerID, 0)
@@ -102,7 +101,7 @@ func (l *Listener) peerWentOffline(peerID messages.PeerID) {
 	if _, ok := l.interestedPeersForOffline[peerID]; ok {
 		select {
 		case l.offlineChan <- peerID:
-		case <-l.listenerCtx.Done():
+		case <-l.ctx.Done():
 		}
 	}
 }
@@ -114,7 +113,7 @@ func (l *Listener) peerComeOnline(peerID messages.PeerID) {
 	if _, ok := l.interestedPeersForOnline[peerID]; ok {
 		select {
 		case l.onlineChan <- peerID:
-		case <-l.listenerCtx.Done():
+		case <-l.ctx.Done():
 		}
 		delete(l.interestedPeersForOnline, peerID)
 	}
diff --git a/relay/server/store/notifier.go b/relay/server/store/notifier.go
index d04db478..ad2e5354 100644
--- a/relay/server/store/notifier.go
+++ b/relay/server/store/notifier.go
@@ -24,8 +24,8 @@ func NewPeerNotifier(store *Store) *PeerNotifier {
 
 func (pn *PeerNotifier) NewListener(onPeersComeOnline, onPeersWentOffline func([]messages.PeerID)) *Listener {
 	ctx, cancel := context.WithCancel(context.Background())
-	listener := newListener(pn.store)
-	go listener.listenForEvents(ctx, onPeersComeOnline, onPeersWentOffline)
+	listener := newListener(ctx, pn.store)
+	go listener.listenForEvents(onPeersComeOnline, onPeersWentOffline)
 
 	pn.listenersMutex.Lock()
 	pn.listeners[listener] = cancel



Test patch:
diff --git a/client/cmd/testutil_test.go b/client/cmd/testutil_test.go
index b5a80d63..228a5d50 100644
--- a/client/cmd/testutil_test.go
+++ b/client/cmd/testutil_test.go
@@ -109,7 +109,7 @@ func startManagement(t *testing.T, config *types.Config, testFile string) (*grpc
 	}
 
 	secretsManager := mgmt.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
-	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil, &mgmt.MockIntegratedValidator{})
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/client/internal/engine_test.go b/client/internal/engine_test.go
index 4b7a2d60..e75672ed 100644
--- a/client/internal/engine_test.go
+++ b/client/internal/engine_test.go
@@ -1481,6 +1481,10 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 		GetSettings(gomock.Any(), gomock.Any(), gomock.Any()).
 		Return(&types.Settings{}, nil).
 		AnyTimes()
+	settingsMockManager.EXPECT().
+		GetExtraSettings(gomock.Any(), gomock.Any()).
+		Return(&types.ExtraSettings{}, nil).
+		AnyTimes()
 
 	permissionsManager := permissions.NewManager(store)
 
@@ -1490,7 +1494,7 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 	}
 
 	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
-	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil, &server.MockIntegratedValidator{})
 	if err != nil {
 		return nil, "", err
 	}
diff --git a/client/server/server_test.go b/client/server/server_test.go
index 376b7e8b..7c46aac5 100644
--- a/client/server/server_test.go
+++ b/client/server/server_test.go
@@ -212,7 +212,7 @@ func startManagement(t *testing.T, signalAddr string, counter *int) (*grpc.Serve
 	}
 
 	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
-	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil, &server.MockIntegratedValidator{})
 	if err != nil {
 		return nil, "", err
 	}
diff --git a/client/status/status_test.go b/client/status/status_test.go
index 33eda4b9..5b5d23ef 100644
--- a/client/status/status_test.go
+++ b/client/status/status_test.go
@@ -234,7 +234,7 @@ var overview = OutputOverview{
 }
 
 func TestConversionFromFullStatusToOutputOverview(t *testing.T) {
-	convertedResult := ConvertToStatusOutputOverview(resp, false, "", nil, nil, nil)
+	convertedResult := ConvertToStatusOutputOverview(resp, false, "", nil, nil, nil, "")
 
 	assert.Equal(t, overview, convertedResult)
 }
diff --git a/management/client/client_test.go b/management/client/client_test.go
index 1847af73..b59b7c98 100644
--- a/management/client/client_test.go
+++ b/management/client/client_test.go
@@ -112,7 +112,7 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 	}
 
 	secretsManager := mgmt.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
-	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil, mgmt.MockIntegratedValidator{})
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/management/server/account_test.go b/management/server/account_test.go
index fcd40b08..b65dffe6 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -2887,7 +2887,7 @@ func createManager(t testing.TB) (*DefaultAccountManager, error) {
 
 	permissionsManager := permissions.NewManager(store)
 
-	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	if err != nil {
 		return nil, err
 	}
diff --git a/management/server/dns_test.go b/management/server/dns_test.go
index 31c944a2..f2295450 100644
--- a/management/server/dns_test.go
+++ b/management/server/dns_test.go
@@ -219,7 +219,7 @@ func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
 	// return empty extra settings for expected calls to UpdateAccountPeers
 	settingsMockManager.EXPECT().GetExtraSettings(gomock.Any(), gomock.Any()).Return(&types.ExtraSettings{}, nil).AnyTimes()
 	permissionsManager := permissions.NewManager(store)
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 }
 
 func createDNSStore(t *testing.T) (store.Store, error) {
diff --git a/management/server/http/testing/testing_tools/tools.go b/management/server/http/testing/testing_tools/tools.go
index 829bff45..e308f100 100644
--- a/management/server/http/testing/testing_tools/tools.go
+++ b/management/server/http/testing/testing_tools/tools.go
@@ -1,4 +1,5 @@
 package testing_tools
+
 import (
 	"bytes"
 	"context"
@@ -132,7 +133,7 @@ func BuildApiBlackBoxWithDBState(t TB, sqlFile string, expectedPeerUpdate *serve
 	}
 
 	geoMock := &geolocation.Mock{}
-	validatorMock := server.MocIntegratedValidator{}
+	validatorMock := server.MockIntegratedValidator{}
 	proxyController := integrations.NewController(store)
 	userManager := users.NewManager(store)
 	permissionsManager := permissions.NewManager(store)
diff --git a/management/server/management_proto_test.go b/management/server/management_proto_test.go
index 57c00ed9..0d61b3a1 100644
--- a/management/server/management_proto_test.go
+++ b/management/server/management_proto_test.go
@@ -448,7 +448,7 @@ func startManagementForTest(t *testing.T, testFile string, config *types.Config)
 	permissionsManager := permissions.NewManager(store)
 
 	accountManager, err := BuildManager(ctx, store, peersUpdateManager, nil, "", "netbird.selfhosted",
-		eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+		eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 
 	if err != nil {
 		cleanup()
@@ -458,7 +458,7 @@ func startManagementForTest(t *testing.T, testFile string, config *types.Config)
 	secretsManager := NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
 
 	ephemeralMgr := NewEphemeralManager(store, accountManager)
-	mgmtServer, err := NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, ephemeralMgr, nil)
+	mgmtServer, err := NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, ephemeralMgr, nil, MockIntegratedValidator{})
 	if err != nil {
 		return nil, nil, "", cleanup, err
 	}
diff --git a/management/server/management_test.go b/management/server/management_test.go
index 0a6b3f75..ab6f0095 100644
--- a/management/server/management_test.go
+++ b/management/server/management_test.go
@@ -206,7 +206,7 @@ func startServer(
 		eventStore,
 		nil,
 		false,
-		server.MocIntegratedValidator{},
+		server.MockIntegratedValidator{},
 		metrics,
 		port_forwarding.NewControllerMock(),
 		settingsMockManager,
@@ -227,6 +227,7 @@ func startServer(
 		nil,
 		nil,
 		nil,
+		server.MockIntegratedValidator{},
 	)
 	if err != nil {
 		t.Fatalf("failed creating management server: %v", err)
diff --git a/management/server/migration/migration_test.go b/management/server/migration/migration_test.go
index 94377930..ce76bd66 100644
--- a/management/server/migration/migration_test.go
+++ b/management/server/migration/migration_test.go
@@ -4,16 +4,21 @@ import (
 	"context"
 	"encoding/gob"
 	"net"
+	"os"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 
 	"github.com/netbirdio/netbird/management/server/migration"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/testutil"
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/route"
 )
@@ -21,7 +26,41 @@ import (
 func setupDatabase(t *testing.T) *gorm.DB {
 	t.Helper()
 
-	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
+	var db *gorm.DB
+	var err error
+	var dsn string
+	var cleanup func()
+	switch os.Getenv("NETBIRD_STORE_ENGINE") {
+	case "mysql":
+		cleanup, dsn, err = testutil.CreateMysqlTestContainer()
+		if err != nil {
+			t.Fatalf("Failed to create MySQL test container: %v", err)
+		}
+
+		if dsn == "" {
+			t.Fatal("MySQL connection string is empty, ensure the test container is running")
+		}
+
+		db, err = gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), &gorm.Config{})
+	case "postgres":
+		cleanup, dsn, err = testutil.CreatePostgresTestContainer()
+		if err != nil {
+			t.Fatalf("Failed to create PostgreSQL test container: %v", err)
+		}
+
+		if dsn == "" {
+			t.Fatalf("PostgreSQL connection string is empty, ensure the test container is running")
+		}
+
+		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
+	case "sqlite":
+		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
+	default:
+		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
+	}
+	if cleanup != nil {
+		t.Cleanup(cleanup)
+	}
 
 	require.NoError(t, err, "Failed to open database")
 	return db
@@ -34,6 +73,7 @@ func TestMigrateFieldFromGobToJSON_EmptyDB(t *testing.T) {
 }
 
 func TestMigrateFieldFromGobToJSON_WithGobData(t *testing.T) {
+	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
 	db := setupDatabase(t)
 
 	err := db.AutoMigrate(&types.Account{}, &route.Route{})
@@ -97,6 +137,7 @@ func TestMigrateNetIPFieldFromBlobToJSON_EmptyDB(t *testing.T) {
 }
 
 func TestMigrateNetIPFieldFromBlobToJSON_WithBlobData(t *testing.T) {
+	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
 	db := setupDatabase(t)
 
 	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
@@ -117,12 +158,18 @@ func TestMigrateNetIPFieldFromBlobToJSON_WithBlobData(t *testing.T) {
 		Peers []peer `gorm:"foreignKey:AccountID;references:id"`
 	}
 
-	err = db.Save(&account{
+	a := &account{
 		Account: types.Account{Id: "123"},
-		Peers: []peer{
-			{Location: location{ConnectionIP: net.IP{10, 0, 0, 1}}},
-		}},
-	).Error
+	}
+
+	err = db.Save(a).Error
+	require.NoError(t, err, "Failed to insert account")
+
+	a.Peers = []peer{
+		{Location: location{ConnectionIP: net.IP{10, 0, 0, 1}}},
+	}
+
+	err = db.Save(a).Error
 	require.NoError(t, err, "Failed to insert blob data")
 
 	var blobValue string
@@ -143,12 +190,18 @@ func TestMigrateNetIPFieldFromBlobToJSON_WithJSONData(t *testing.T) {
 	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
 	require.NoError(t, err, "Failed to auto-migrate tables")
 
-	err = db.Save(&types.Account{
+	account := &types.Account{
 		Id: "1234",
-		PeersG: []nbpeer.Peer{
-			{Location: nbpeer.Location{ConnectionIP: net.IP{10, 0, 0, 1}}},
-		}},
-	).Error
+	}
+
+	err = db.Save(account).Error
+	require.NoError(t, err, "Failed to insert account")
+
+	account.PeersG = []nbpeer.Peer{
+		{AccountID: "1234", Location: nbpeer.Location{ConnectionIP: net.IP{10, 0, 0, 1}}},
+	}
+
+	err = db.Save(account).Error
 	require.NoError(t, err, "Failed to insert JSON data")
 
 	err = migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "location_connection_ip", "")
@@ -162,12 +215,13 @@ func TestMigrateNetIPFieldFromBlobToJSON_WithJSONData(t *testing.T) {
 func TestMigrateSetupKeyToHashedSetupKey_ForPlainKey(t *testing.T) {
 	db := setupDatabase(t)
 
-	err := db.AutoMigrate(&types.SetupKey{})
+	err := db.AutoMigrate(&types.SetupKey{}, &nbpeer.Peer{})
 	require.NoError(t, err, "Failed to auto-migrate tables")
 
 	err = db.Save(&types.SetupKey{
-		Id:  "1",
-		Key: "EEFDAB47-C1A5-4472-8C05-71DE9A1E8382",
+		Id:        "1",
+		Key:       "EEFDAB47-C1A5-4472-8C05-71DE9A1E8382",
+		UpdatedAt: time.Now(),
 	}).Error
 	require.NoError(t, err, "Failed to insert setup key")
 
@@ -192,6 +246,7 @@ func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case1(t *testing.
 		Id:        "1",
 		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
 		KeySecret: "EEFDA****",
+		UpdatedAt: time.Now(),
 	}).Error
 	require.NoError(t, err, "Failed to insert setup key")
 
@@ -213,8 +268,9 @@ func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case2(t *testing.
 	require.NoError(t, err, "Failed to auto-migrate tables")
 
 	err = db.Save(&types.SetupKey{
-		Id:  "1",
-		Key: "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+		Id:        "1",
+		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+		UpdatedAt: time.Now(),
 	}).Error
 	require.NoError(t, err, "Failed to insert setup key")
 
@@ -235,8 +291,9 @@ func TestDropIndex(t *testing.T) {
 	require.NoError(t, err, "Failed to auto-migrate tables")
 
 	err = db.Save(&types.SetupKey{
-		Id:  "1",
-		Key: "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+		Id:        "1",
+		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+		UpdatedAt: time.Now(),
 	}).Error
 	require.NoError(t, err, "Failed to insert setup key")
 
@@ -249,3 +306,37 @@ func TestDropIndex(t *testing.T) {
 	exist = db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
 	assert.False(t, exist, "Should not have the index")
 }
+
+func TestCreateIndex(t *testing.T) {
+	db := setupDatabase(t)
+	err := db.AutoMigrate(&nbpeer.Peer{})
+	assert.NoError(t, err, "Failed to auto-migrate tables")
+
+	indexName := "idx_account_ip"
+
+	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
+	assert.NoError(t, err, "Migration should not fail to create index")
+
+	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
+	assert.True(t, exist, "Should have the index")
+}
+
+func TestCreateIndexIfExists(t *testing.T) {
+	db := setupDatabase(t)
+	err := db.AutoMigrate(&nbpeer.Peer{})
+	assert.NoError(t, err, "Failed to auto-migrate tables")
+
+	indexName := "idx_account_ip"
+
+	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
+	assert.NoError(t, err, "Migration should not fail to create index")
+
+	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
+	assert.True(t, exist, "Should have the index")
+
+	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
+	assert.NoError(t, err, "Create index should not fail if index exists")
+
+	exist = db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
+	assert.True(t, exist, "Should have the index")
+}
diff --git a/management/server/nameserver_test.go b/management/server/nameserver_test.go
index 8fada742..25eb03b8 100644
--- a/management/server/nameserver_test.go
+++ b/management/server/nameserver_test.go
@@ -785,7 +785,7 @@ func createNSManager(t *testing.T) (*DefaultAccountManager, error) {
 		AnyTimes()
 
 	permissionsManager := permissions.NewManager(store)
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 }
 
 func createNSStore(t *testing.T) (store.Store, error) {
diff --git a/management/server/peer_test.go b/management/server/peer_test.go
index d4102051..4f6ae500 100644
--- a/management/server/peer_test.go
+++ b/management/server/peer_test.go
@@ -1273,7 +1273,7 @@ func Test_RegisterPeerByUser(t *testing.T) {
 	settingsMockManager := settings.NewMockManager(ctrl)
 	permissionsManager := permissions.NewManager(s)
 
-	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
@@ -1353,7 +1353,7 @@ func Test_RegisterPeerBySetupKey(t *testing.T) {
 		AnyTimes()
 	permissionsManager := permissions.NewManager(s)
 
-	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
@@ -1496,7 +1496,7 @@ func Test_RegisterPeerRollbackOnFailure(t *testing.T) {
 
 	permissionsManager := permissions.NewManager(s)
 
-	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
@@ -1570,7 +1570,7 @@ func Test_LoginPeer(t *testing.T) {
 		AnyTimes()
 	permissionsManager := permissions.NewManager(s)
 
-	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	am, err := BuildManager(context.Background(), s, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 	assert.NoError(t, err)
 
 	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
@@ -1848,7 +1848,7 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 			return update, true, nil
 		}
 
-		manager.integratedPeerValidator = MocIntegratedValidator{ValidatePeerFunc: requireUpdateFunc}
+		manager.integratedPeerValidator = MockIntegratedValidator{ValidatePeerFunc: requireUpdateFunc}
 		done := make(chan struct{})
 		go func() {
 			peerShouldReceiveUpdate(t, updMsg)
@@ -1870,7 +1870,7 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 			return update, false, nil
 		}
 
-		manager.integratedPeerValidator = MocIntegratedValidator{ValidatePeerFunc: requireNoUpdateFunc}
+		manager.integratedPeerValidator = MockIntegratedValidator{ValidatePeerFunc: requireNoUpdateFunc}
 		done := make(chan struct{})
 		go func() {
 			peerShouldNotReceiveUpdate(t, updMsg)
diff --git a/management/server/route_test.go b/management/server/route_test.go
index 77cbc75b..37c37f62 100644
--- a/management/server/route_test.go
+++ b/management/server/route_test.go
@@ -1284,7 +1284,7 @@ func createRouterManager(t *testing.T) (*DefaultAccountManager, error) {
 
 	permissionsManager := permissions.NewManager(store)
 
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MockIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager, false)
 }
 
 func createRouterStore(t *testing.T) (store.Store, error) {
diff --git a/management/server/user_test.go b/management/server/user_test.go
index 7508e060..53baf8f7 100644
--- a/management/server/user_test.go
+++ b/management/server/user_test.go
@@ -852,7 +852,7 @@ func TestUser_DeleteUser_RegularUsers(t *testing.T) {
 	am := DefaultAccountManager{
 		Store:                   store,
 		eventStore:              &activity.InMemoryEventStore{},
-		integratedPeerValidator: MocIntegratedValidator{},
+		integratedPeerValidator: MockIntegratedValidator{},
 		permissionsManager:      permissionsManager,
 	}
 
diff --git a/relay/client/manager_test.go b/relay/client/manager_test.go
index d20cdaac..52f2833e 100644
--- a/relay/client/manager_test.go
+++ b/relay/client/manager_test.go
@@ -229,16 +229,14 @@ func TestForeginAutoClose(t *testing.T) {
 	errChan := make(chan error, 1)
 	go func() {
 		t.Log("binding server 1.")
-		err := srv1.Listen(srvCfg1)
-		if err != nil {
+		if err := srv1.Listen(srvCfg1); err != nil {
 			errChan <- err
 		}
 	}()
 
 	defer func() {
 		t.Logf("closing server 1.")
-		err := srv1.Shutdown(ctx)
-		if err != nil {
+		if err := srv1.Shutdown(ctx); err != nil {
 			t.Errorf("failed to close server: %s", err)
 		}
 		t.Logf("server 1. closed")
@@ -287,15 +285,8 @@ func TestForeginAutoClose(t *testing.T) {
 	}
 
 	t.Log("open connection to another peer")
-	conn, err := mgr.OpenConn(ctx, toURL(srvCfg2)[0], "anotherpeer")
-	if err != nil {
-		t.Fatalf("failed to bind channel: %s", err)
-	}
-
-	t.Log("close conn")
-	err = conn.Close()
-	if err != nil {
-		t.Fatalf("failed to close connection: %s", err)
+	if _, err = mgr.OpenConn(ctx, toURL(srvCfg2)[0], "anotherpeer"); err == nil {
+		t.Fatalf("should have failed to open connection to another peer")
 	}
 
 	timeout := relayCleanupInterval + keepUnusedServerTime + 1*time.Second