run windows iface tests without race flag

using for now a temp fixed for ice

.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
+          skip-single-stacks: true

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

@@ -42,5 +42,5 @@ jobs:
         run: git --no-pager diff --exit-code
 
       - name: Test
-        run: NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -tags=devcert -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 5m -p 1 $(go list ./... | grep -v /management)
+        run: NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -tags=devcert -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -race -timeout 5m -p 1 $(go list ./... | grep -v /management)
 

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

@@ -103,7 +103,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        arch: [ '386','amd64' ]
+        include:
+          - arch: "386"
+            raceFlag: ""
+          - arch: "amd64"
+            raceFlag: "-race"
     runs-on: ubuntu-22.04
     steps:
       - name: Install Go
@@ -144,7 +148,7 @@ jobs:
         run: git --no-pager diff --exit-code
 
       - name: Test
-        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} CI=true go test -tags devcert -exec 'sudo' -timeout 10m -p 1 $(go list ./... | grep -v -e /management -e /signal -e /relay)
+        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} CI=true go test -tags devcert -exec 'sudo' ${{ matrix.raceFlag }} -timeout 10m -p 1 $(go list ./... | grep -v -e /management -e /signal -e /relay)
 
   test_client_on_docker:
     name: "Client (Docker) / Unit"

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

@@ -63,10 +63,16 @@ jobs:
       - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOMODCACHE=${{ env.cache }}
       - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOCACHE=${{ env.modcache }}
       - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe mod tidy
-      - run: echo "files=$(go list ./... | ForEach-Object { $_ } | Where-Object { $_ -notmatch '/management' })" >> $env:GITHUB_ENV
+      - run: echo "files=$(go list ./... | ForEach-Object { $_ } | Where-Object { $_ -notmatch '/management' } | Where-Object { $_ -notmatch '/client/iface' } )" >> $env:GITHUB_ENV
 
-      - name: test
-        run: PsExec64 -s -w ${{ github.workspace }} cmd.exe /c "C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe test -tags=devcert -timeout 10m -p 1 ${{ env.files }} > test-out.txt 2>&1"
-      - name: test output
+      - name: test without iface
+        run: PsExec64 -s -w ${{ github.workspace }} cmd.exe /c "C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe test -tags=devcert -race -timeout 10m -p 1 ${{ env.files }} > test-out.txt 2>&1"
+      - name: test output without iface
+        if: ${{ always() }}
+        run: Get-Content test-out.txt
+      # todo: remove this once iface tests are stable with race flag
+      - name: test iface
+        run: PsExec64 -s -w ${{ github.workspace }} cmd.exe /c "C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe test -tags=devcert -timeout 10m -p 1 ./client/iface/... > test-out.txt 2>&1"
+      - name: test output iface
         if: ${{ always() }}
         run: Get-Content test-out.txt

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

@@ -43,7 +43,7 @@ jobs:
       - name: gomobile init
         run: gomobile init
       - name: build android netbird lib
-        run: PATH=$PATH:$(go env GOPATH) gomobile bind -o $GITHUB_WORKSPACE/netbird.aar -javapkg=io.netbird.gomobile -ldflags="-X golang.zx2c4.com/wireguard/ipc.socketDirectory=/data/data/io.netbird.client/cache/wireguard -X github.com/netbirdio/netbird/version.version=buildtest"  $GITHUB_WORKSPACE/client/android
+        run: PATH=$PATH:$(go env GOPATH) gomobile bind -o $GITHUB_WORKSPACE/netbird.aar -javapkg=io.netbird.gomobile -ldflags="-checklinkname=0 -X golang.zx2c4.com/wireguard/ipc.socketDirectory=/data/data/io.netbird.client/cache/wireguard -X github.com/netbirdio/netbird/version.version=buildtest"  $GITHUB_WORKSPACE/client/android
         env:
           CGO_ENABLED: 0
           ANDROID_NDK_HOME: /usr/local/lib/android/sdk/ndk/23.1.7779620

.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          

README.md

@@ -50,10 +50,9 @@
 
 **Secure.** NetBird enables secure remote access by applying granular access policies while allowing you to manage them intuitively from a single place. Works universally on any infrastructure.
 
-### Open-Source Network Security in a Single Platform
+### Open Source Network Security in a Single Platform
 
-
-![netbird_2](https://github.com/netbirdio/netbird/assets/700848/46bc3b73-508d-4a0e-bb9a-f465d68646ab)
+<img width="1188" alt="centralized-network-management 1" src="https://github.com/user-attachments/assets/c28cc8e4-15d2-4d2f-bb97-a6433db39d56" />
 
 ### NetBird on Lawrence Systems (Video)
 [![Watch the video](https://img.youtube.com/vi/Kwrff6h0rEw/0.jpg)](https://www.youtube.com/watch?v=Kwrff6h0rEw)

client/android/client.go

@@ -64,7 +64,9 @@ type Client struct {
 }
 
 // NewClient instantiate a new Client
-func NewClient(cfgFile, deviceName string, uiVersion string, tunAdapter TunAdapter, iFaceDiscover IFaceDiscover, networkChangeListener NetworkChangeListener) *Client {
+func NewClient(cfgFile string, androidSDKVersion int, deviceName string, uiVersion string, tunAdapter TunAdapter, iFaceDiscover IFaceDiscover, networkChangeListener NetworkChangeListener) *Client {
+	execWorkaround(androidSDKVersion)
+
 	net.SetAndroidProtectSocketFn(tunAdapter.ProtectSocket)
 	return &Client{
 		cfgFile:               cfgFile,

client/android/exec.go

@@ -0,0 +1,26 @@
+//go:build android
+
+package android
+
+import (
+	"fmt"
+	_ "unsafe"
+)
+
+// https://github.com/golang/go/pull/69543/commits/aad6b3b32c81795f86bc4a9e81aad94899daf520
+// In Android version 11 and earlier, pidfd-related system calls
+// are not allowed by the seccomp policy, which causes crashes due
+// to SIGSYS signals.
+
+//go:linkname checkPidfdOnce os.checkPidfdOnce
+var checkPidfdOnce func() error
+
+func execWorkaround(androidSDKVersion int) {
+	if androidSDKVersion > 30 { // above Android 11
+		return
+	}
+
+	checkPidfdOnce = func() error {
+		return fmt.Errorf("unsupported Android version")
+	}
+}

client/cmd/debug.go

@@ -17,10 +17,18 @@ import (
 	"github.com/netbirdio/netbird/client/server"
 	nbstatus "github.com/netbirdio/netbird/client/status"
 	mgmProto "github.com/netbirdio/netbird/management/proto"
+	"github.com/netbirdio/netbird/upload-server/types"
 )
 
 const errCloseConnection = "Failed to close connection: %v"
 
+var (
+	logFileCount        uint32
+	systemInfoFlag      bool
+	uploadBundleFlag    bool
+	uploadBundleURLFlag string
+)
+
 var debugCmd = &cobra.Command{
 	Use:   "debug",
 	Short: "Debugging commands",
@@ -88,12 +96,13 @@ func debugBundle(cmd *cobra.Command, _ []string) error {
 
 	client := proto.NewDaemonServiceClient(conn)
 	request := &proto.DebugBundleRequest{
-		Anonymize:  anonymizeFlag,
-		Status:     getStatusOutput(cmd, anonymizeFlag),
-		SystemInfo: debugSystemInfoFlag,
+		Anonymize:    anonymizeFlag,
+		Status:       getStatusOutput(cmd, anonymizeFlag),
+		SystemInfo:   systemInfoFlag,
+		LogFileCount: logFileCount,
 	}
-	if debugUploadBundle {
-		request.UploadURL = debugUploadBundleURL
+	if uploadBundleFlag {
+		request.UploadURL = uploadBundleURLFlag
 	}
 	resp, err := client.DebugBundle(cmd.Context(), request)
 	if err != nil {
@@ -105,7 +114,7 @@ func debugBundle(cmd *cobra.Command, _ []string) error {
 		return fmt.Errorf("upload failed: %s", resp.GetUploadFailureReason())
 	}
 
-	if debugUploadBundle {
+	if uploadBundleFlag {
 		cmd.Printf("Upload file key:\n%s\n", resp.GetUploadedKey())
 	}
 
@@ -223,12 +232,13 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 	headerPreDown := fmt.Sprintf("----- Netbird pre-down - Timestamp: %s - Duration: %s", time.Now().Format(time.RFC3339), duration)
 	statusOutput = fmt.Sprintf("%s\n%s\n%s", statusOutput, headerPreDown, getStatusOutput(cmd, anonymizeFlag))
 	request := &proto.DebugBundleRequest{
-		Anonymize:  anonymizeFlag,
-		Status:     statusOutput,
-		SystemInfo: debugSystemInfoFlag,
+		Anonymize:    anonymizeFlag,
+		Status:       statusOutput,
+		SystemInfo:   systemInfoFlag,
+		LogFileCount: logFileCount,
 	}
-	if debugUploadBundle {
-		request.UploadURL = debugUploadBundleURL
+	if uploadBundleFlag {
+		request.UploadURL = uploadBundleURLFlag
 	}
 	resp, err := client.DebugBundle(cmd.Context(), request)
 	if err != nil {
@@ -255,7 +265,7 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("upload failed: %s", resp.GetUploadFailureReason())
 	}
 
-	if debugUploadBundle {
+	if uploadBundleFlag {
 		cmd.Printf("Upload file key:\n%s\n", resp.GetUploadedKey())
 	}
 
@@ -375,3 +385,15 @@ func generateDebugBundle(config *internal.Config, recorder *peer.Status, connect
 	}
 	log.Infof("Generated debug bundle from SIGUSR1 at: %s", path)
 }
+
+func init() {
+	debugBundleCmd.Flags().Uint32VarP(&logFileCount, "log-file-count", "C", 1, "Number of rotated log files to include in debug bundle")
+	debugBundleCmd.Flags().BoolVarP(&systemInfoFlag, "system-info", "S", true, "Adds system information to the debug bundle")
+	debugBundleCmd.Flags().BoolVarP(&uploadBundleFlag, "upload-bundle", "U", false, "Uploads the debug bundle to a server")
+	debugBundleCmd.Flags().StringVar(&uploadBundleURLFlag, "upload-bundle-url", types.DefaultBundleURL, "Service URL to get an URL to upload the debug bundle")
+
+	forCmd.Flags().Uint32VarP(&logFileCount, "log-file-count", "C", 1, "Number of rotated log files to include in debug bundle")
+	forCmd.Flags().BoolVarP(&systemInfoFlag, "system-info", "S", true, "Adds system information to the debug bundle")
+	forCmd.Flags().BoolVarP(&uploadBundleFlag, "upload-bundle", "U", false, "Uploads the debug bundle to a server")
+	forCmd.Flags().StringVar(&uploadBundleURLFlag, "upload-bundle-url", types.DefaultBundleURL, "Service URL to get an URL to upload the debug bundle")
+}

client/cmd/login_test.go

@@ -1,3 +1,5 @@
+//go:build !race
+
 package cmd
 
 import (

client/cmd/root.go

@@ -22,7 +22,6 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 
 	"github.com/netbirdio/netbird/client/internal"
-	"github.com/netbirdio/netbird/upload-server/types"
 )
 
 const (
@@ -38,10 +37,7 @@ const (
 	serverSSHAllowedFlag     = "allow-server-ssh"
 	extraIFaceBlackListFlag  = "extra-iface-blacklist"
 	dnsRouteIntervalFlag     = "dns-router-interval"
-	systemInfoFlag           = "system-info"
 	enableLazyConnectionFlag = "enable-lazy-connection"
-	uploadBundle             = "upload-bundle"
-	uploadBundleURL          = "upload-bundle-url"
 )
 
 var (
@@ -75,10 +71,7 @@ var (
 	autoConnectDisabled     bool
 	extraIFaceBlackList     []string
 	anonymizeFlag           bool
-	debugSystemInfoFlag     bool
 	dnsRouteInterval        time.Duration
-	debugUploadBundle       bool
-	debugUploadBundleURL    string
 	lazyConnEnabled         bool
 
 	rootCmd = &cobra.Command{
@@ -184,11 +177,8 @@ func init() {
 	upCmd.PersistentFlags().BoolVar(&rosenpassPermissive, rosenpassPermissiveFlag, false, "[Experimental] Enable Rosenpass in permissive mode to allow this peer to accept WireGuard connections without requiring Rosenpass functionality from peers that do not have Rosenpass enabled.")
 	upCmd.PersistentFlags().BoolVar(&serverSSHAllowed, serverSSHAllowedFlag, false, "Allow SSH server on peer. If enabled, the SSH server will be permitted")
 	upCmd.PersistentFlags().BoolVar(&autoConnectDisabled, disableAutoConnectFlag, false, "Disables auto-connect feature. If enabled, then the client won't connect automatically when the service starts.")
-	upCmd.PersistentFlags().BoolVar(&lazyConnEnabled, enableLazyConnectionFlag, false, "[Experimental] Enable the lazy connection feature. If enabled, the client will establish connections on-demand.")
+	upCmd.PersistentFlags().BoolVar(&lazyConnEnabled, enableLazyConnectionFlag, false, "[Experimental] Enable the lazy connection feature. If enabled, the client will establish connections on-demand. Note: this setting may be overridden by management configuration.")
 
-	debugCmd.PersistentFlags().BoolVarP(&debugSystemInfoFlag, systemInfoFlag, "S", true, "Adds system information to the debug bundle")
-	debugCmd.PersistentFlags().BoolVarP(&debugUploadBundle, uploadBundle, "U", false, fmt.Sprintf("Uploads the debug bundle to a server from URL defined by %s", uploadBundleURL))
-	debugCmd.PersistentFlags().StringVar(&debugUploadBundleURL, uploadBundleURL, types.DefaultBundleURL, "Service URL to get an URL to upload the debug bundle")
 }
 
 // SetupCloseHandler handles SIGTERM signal and exits with success

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

client/firewall/uspfilter/conntrack/udp.go

@@ -218,3 +218,9 @@ func (t *UDPTracker) sendEvent(typ nftypes.Type, conn *UDPConnTrack, ruleID []by
 		TxBytes:    conn.BytesTx.Load(),
 	})
 }
+
+func (t *UDPTracker) getConnectionsLen() int {
+	t.mutex.RLock()
+	defer t.mutex.RUnlock()
+	return len(t.connections)
+}

client/firewall/uspfilter/conntrack/udp_test.go

@@ -202,13 +202,13 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 	}
 
 	// Verify initial connections
-	assert.Len(t, tracker.connections, 2)
+	assert.Equal(t, 2, tracker.getConnectionsLen())
 
 	// Wait for connection timeout and cleanup interval
 	time.Sleep(timeout + 2*cleanupInterval)
 
 	tracker.mutex.RLock()
-	connCount := len(tracker.connections)
+	connCount := tracker.getConnectionsLen()
 	tracker.mutex.RUnlock()
 
 	// Verify connections were cleaned up

client/iface/bind/activity.go

@@ -34,14 +34,14 @@ func NewActivityRecorder() *ActivityRecorder {
 }
 
 // GetLastActivities returns a snapshot of peer last activity
-func (r *ActivityRecorder) GetLastActivities() map[string]time.Time {
+func (r *ActivityRecorder) GetLastActivities() map[string]monotime.Time {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
 
-	activities := make(map[string]time.Time, len(r.peers))
+	activities := make(map[string]monotime.Time, len(r.peers))
 	for key, record := range r.peers {
-		unixNano := record.LastActivity.Load()
-		activities[key] = time.Unix(0, unixNano)
+		monoTime := record.LastActivity.Load()
+		activities[key] = monotime.Time(monoTime)
 	}
 	return activities
 }
@@ -51,18 +51,20 @@ func (r *ActivityRecorder) UpsertAddress(publicKey string, address netip.AddrPor
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	if pr, exists := r.peers[publicKey]; exists {
-		delete(r.addrToPeer, pr.Address)
-		pr.Address = address
+	var record *PeerRecord
+	record, exists := r.peers[publicKey]
+	if exists {
+		delete(r.addrToPeer, record.Address)
+		record.Address = address
 	} else {
-		record := &PeerRecord{
+		record = &PeerRecord{
 			Address: address,
 		}
-		record.LastActivity.Store(monotime.Now())
+		record.LastActivity.Store(int64(monotime.Now()))
 		r.peers[publicKey] = record
 	}
 
-	r.addrToPeer[address] = r.peers[publicKey]
+	r.addrToPeer[address] = record
 }
 
 func (r *ActivityRecorder) Remove(publicKey string) {
@@ -84,7 +86,7 @@ func (r *ActivityRecorder) record(address netip.AddrPort) {
 		return
 	}
 
-	now := monotime.Now()
+	now := int64(monotime.Now())
 	last := record.LastActivity.Load()
 	if now-last < saveFrequency {
 		return

client/iface/bind/activity_test.go

@@ -4,6 +4,8 @@ import (
 	"net/netip"
 	"testing"
 	"time"
+
+	"github.com/netbirdio/netbird/monotime"
 )
 
 func TestActivityRecorder_GetLastActivities(t *testing.T) {
@@ -17,11 +19,7 @@ func TestActivityRecorder_GetLastActivities(t *testing.T) {
 		t.Fatalf("Expected activity for peer %s, but got none", peer)
 	}
 
-	if p.IsZero() {
-		t.Fatalf("Expected activity for peer %s, but got zero", peer)
-	}
-
-	if p.Before(time.Now().Add(-2 * time.Minute)) {
+	if monotime.Since(p) > 5*time.Second {
 		t.Fatalf("Expected activity for peer %s to be recent, but got %v", peer, p)
 	}
 }

client/iface/configurer/kernel_unix.go

@@ -11,6 +11,8 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/monotime"
 )
 
 var zeroKey wgtypes.Key
@@ -277,6 +279,6 @@ func (c *KernelConfigurer) GetStats() (map[string]WGStats, error) {
 	return stats, nil
 }
 
-func (c *KernelConfigurer) LastActivities() map[string]time.Time {
+func (c *KernelConfigurer) LastActivities() map[string]monotime.Time {
 	return nil
 }

client/iface/configurer/usp.go

@@ -17,6 +17,7 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
+	"github.com/netbirdio/netbird/monotime"
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
@@ -223,7 +224,7 @@ func (c *WGUSPConfigurer) FullStats() (*Stats, error) {
 	return parseStatus(c.deviceName, ipcStr)
 }
 
-func (c *WGUSPConfigurer) LastActivities() map[string]time.Time {
+func (c *WGUSPConfigurer) LastActivities() map[string]monotime.Time {
 	return c.activityRecorder.GetLastActivities()
 }
 

client/iface/device/interface.go

@@ -8,6 +8,7 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/iface/configurer"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 type WGConfigurer interface {
@@ -19,5 +20,5 @@ type WGConfigurer interface {
 	Close()
 	GetStats() (map[string]configurer.WGStats, error)
 	FullStats() (*configurer.Stats, error)
-	LastActivities() map[string]time.Time
+	LastActivities() map[string]monotime.Time
 }

client/iface/iface.go

@@ -21,6 +21,7 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/iface/wgproxy"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 const (
@@ -29,6 +30,11 @@ const (
 	WgInterfaceDefault = configurer.WgInterfaceDefault
 )
 
+var (
+	// ErrIfaceNotFound is returned when the WireGuard interface is not found
+	ErrIfaceNotFound = fmt.Errorf("wireguard interface not found")
+)
+
 type wgProxyFactory interface {
 	GetProxy() wgproxy.Proxy
 	Free() error
@@ -117,6 +123,9 @@ func (w *WGIface) UpdateAddr(newAddr string) error {
 func (w *WGIface) UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	if w.configurer == nil {
+		return ErrIfaceNotFound
+	}
 
 	log.Debugf("updating interface %s peer %s, endpoint %s, allowedIPs %v", w.tun.DeviceName(), peerKey, endpoint, allowedIps)
 	return w.configurer.UpdatePeer(peerKey, allowedIps, keepAlive, endpoint, preSharedKey)
@@ -126,6 +135,9 @@ func (w *WGIface) UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAliv
 func (w *WGIface) RemovePeer(peerKey string) error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	if w.configurer == nil {
+		return ErrIfaceNotFound
+	}
 
 	log.Debugf("Removing peer %s from interface %s ", peerKey, w.tun.DeviceName())
 	return w.configurer.RemovePeer(peerKey)
@@ -135,6 +147,9 @@ func (w *WGIface) RemovePeer(peerKey string) error {
 func (w *WGIface) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	if w.configurer == nil {
+		return ErrIfaceNotFound
+	}
 
 	log.Debugf("Adding allowed IP to interface %s and peer %s: allowed IP %s ", w.tun.DeviceName(), peerKey, allowedIP)
 	return w.configurer.AddAllowedIP(peerKey, allowedIP)
@@ -144,6 +159,9 @@ func (w *WGIface) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
 func (w *WGIface) RemoveAllowedIP(peerKey string, allowedIP netip.Prefix) error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	if w.configurer == nil {
+		return ErrIfaceNotFound
+	}
 
 	log.Debugf("Removing allowed IP from interface %s and peer %s: allowed IP %s ", w.tun.DeviceName(), peerKey, allowedIP)
 	return w.configurer.RemoveAllowedIP(peerKey, allowedIP)
@@ -214,18 +232,29 @@ func (w *WGIface) GetWGDevice() *wgdevice.Device {
 
 // GetStats returns the last handshake time, rx and tx bytes
 func (w *WGIface) GetStats() (map[string]configurer.WGStats, error) {
+	if w.configurer == nil {
+		return nil, ErrIfaceNotFound
+	}
 	return w.configurer.GetStats()
 }
 
-func (w *WGIface) LastActivities() map[string]time.Time {
+func (w *WGIface) LastActivities() map[string]monotime.Time {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
+	if w.configurer == nil {
+		return nil
+	}
+
 	return w.configurer.LastActivities()
 
 }
 
 func (w *WGIface) FullStats() (*configurer.Stats, error) {
+	if w.configurer == nil {
+		return nil, ErrIfaceNotFound
+	}
+
 	return w.configurer.FullStats()
 }
 

client/iface/netstack/tun.go

@@ -41,9 +41,12 @@ func (t *NetStackTun) Create() (tun.Device, *netstack.Net, error) {
 	}
 	t.tundev = nsTunDev
 
-	skipProxy, err := strconv.ParseBool(os.Getenv(EnvSkipProxy))
-	if err != nil {
-		log.Errorf("failed to parse %s: %s", EnvSkipProxy, err)
+	var skipProxy bool
+	if val := os.Getenv(EnvSkipProxy); val != "" {
+		skipProxy, err = strconv.ParseBool(val)
+		if err != nil {
+			log.Errorf("failed to parse %s: %s", EnvSkipProxy, err)
+		}
 	}
 	if skipProxy {
 		return nsTunDev, tunNet, nil

client/iface/wgproxy/bind/proxy.go

@@ -12,6 +12,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
+	"github.com/netbirdio/netbird/client/iface/wgproxy/listener"
 )
 
 type ProxyBind struct {
@@ -28,6 +29,17 @@ type ProxyBind struct {
 	pausedMu  sync.Mutex
 	paused    bool
 	isStarted bool
+
+	closeListener *listener.CloseListener
+}
+
+func NewProxyBind(bind *bind.ICEBind) *ProxyBind {
+	p := &ProxyBind{
+		Bind:          bind,
+		closeListener: listener.NewCloseListener(),
+	}
+
+	return p
 }
 
 // AddTurnConn adds a new connection to the bind.
@@ -54,6 +66,10 @@ func (p *ProxyBind) EndpointAddr() *net.UDPAddr {
 	}
 }
 
+func (p *ProxyBind) SetDisconnectListener(disconnected func()) {
+	p.closeListener.SetCloseListener(disconnected)
+}
+
 func (p *ProxyBind) Work() {
 	if p.remoteConn == nil {
 		return
@@ -96,6 +112,9 @@ func (p *ProxyBind) close() error {
 	if p.closed {
 		return nil
 	}
+
+	p.closeListener.SetCloseListener(nil)
+
 	p.closed = true
 
 	p.cancel()
@@ -122,6 +141,7 @@ func (p *ProxyBind) proxyToLocal(ctx context.Context) {
 			if ctx.Err() != nil {
 				return
 			}
+			p.closeListener.Notify()
 			log.Errorf("failed to read from remote conn: %s, %s", p.remoteConn.RemoteAddr(), err)
 			return
 		}

client/iface/wgproxy/ebpf/wrapper.go

@@ -11,6 +11,8 @@ import (
 	"sync"
 
 	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/iface/wgproxy/listener"
 )
 
 // ProxyWrapper help to keep the remoteConn instance for net.Conn.Close function call
@@ -26,6 +28,15 @@ type ProxyWrapper struct {
 	pausedMu  sync.Mutex
 	paused    bool
 	isStarted bool
+
+	closeListener *listener.CloseListener
+}
+
+func NewProxyWrapper(WgeBPFProxy *WGEBPFProxy) *ProxyWrapper {
+	return &ProxyWrapper{
+		WgeBPFProxy:   WgeBPFProxy,
+		closeListener: listener.NewCloseListener(),
+	}
 }
 
 func (p *ProxyWrapper) AddTurnConn(ctx context.Context, endpoint *net.UDPAddr, remoteConn net.Conn) error {
@@ -43,6 +54,10 @@ func (p *ProxyWrapper) EndpointAddr() *net.UDPAddr {
 	return p.wgEndpointAddr
 }
 
+func (p *ProxyWrapper) SetDisconnectListener(disconnected func()) {
+	p.closeListener.SetCloseListener(disconnected)
+}
+
 func (p *ProxyWrapper) Work() {
 	if p.remoteConn == nil {
 		return
@@ -77,6 +92,8 @@ func (e *ProxyWrapper) CloseConn() error {
 
 	e.cancel()
 
+	e.closeListener.SetCloseListener(nil)
+
 	if err := e.remoteConn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
 		return fmt.Errorf("failed to close remote conn: %w", err)
 	}
@@ -117,6 +134,7 @@ func (p *ProxyWrapper) readFromRemote(ctx context.Context, buf []byte) (int, err
 		if ctx.Err() != nil {
 			return 0, ctx.Err()
 		}
+		p.closeListener.Notify()
 		if !errors.Is(err, io.EOF) {
 			log.Errorf("failed to read from turn conn (endpoint: :%d): %s", p.wgEndpointAddr.Port, err)
 		}

client/iface/wgproxy/factory_kernel.go

@@ -36,9 +36,8 @@ func (w *KernelFactory) GetProxy() Proxy {
 		return udpProxy.NewWGUDPProxy(w.wgPort)
 	}
 
-	return &ebpf.ProxyWrapper{
-		WgeBPFProxy: w.ebpfProxy,
-	}
+	return ebpf.NewProxyWrapper(w.ebpfProxy)
+
 }
 
 func (w *KernelFactory) Free() error {

client/iface/wgproxy/factory_usp.go

@@ -20,9 +20,7 @@ func NewUSPFactory(iceBind *bind.ICEBind) *USPFactory {
 }
 
 func (w *USPFactory) GetProxy() Proxy {
-	return &proxyBind.ProxyBind{
-		Bind: w.bind,
-	}
+	return proxyBind.NewProxyBind(w.bind)
 }
 
 func (w *USPFactory) Free() error {

client/iface/wgproxy/listener/listener.go

@@ -0,0 +1,19 @@
+package listener
+
+type CloseListener struct {
+	listener func()
+}
+
+func NewCloseListener() *CloseListener {
+	return &CloseListener{}
+}
+
+func (c *CloseListener) SetCloseListener(listener func()) {
+	c.listener = listener
+}
+
+func (c *CloseListener) Notify() {
+	if c.listener != nil {
+		c.listener()
+	}
+}

client/iface/wgproxy/proxy.go

@@ -12,4 +12,5 @@ type Proxy interface {
 	Work()                      // Work start or resume the proxy
 	Pause()                     // Pause to forward the packages from remote connection to WireGuard. The opposite way still works.
 	CloseConn() error
+	SetDisconnectListener(disconnected func())
 }

client/iface/wgproxy/proxy_test.go

@@ -98,9 +98,7 @@ func TestProxyCloseByRemoteConn(t *testing.T) {
 				t.Errorf("failed to free ebpf proxy: %s", err)
 			}
 		}()
-		proxyWrapper := &ebpf.ProxyWrapper{
-			WgeBPFProxy: ebpfProxy,
-		}
+		proxyWrapper := ebpf.NewProxyWrapper(ebpfProxy)
 
 		tests = append(tests, struct {
 			name  string

client/iface/wgproxy/udp/proxy.go

@@ -12,6 +12,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	cerrors "github.com/netbirdio/netbird/client/errors"
+	"github.com/netbirdio/netbird/client/iface/wgproxy/listener"
 )
 
 // WGUDPProxy proxies
@@ -28,6 +29,8 @@ type WGUDPProxy struct {
 	pausedMu  sync.Mutex
 	paused    bool
 	isStarted bool
+
+	closeListener *listener.CloseListener
 }
 
 // NewWGUDPProxy instantiate a UDP based WireGuard proxy. This is not a thread safe implementation
@@ -35,6 +38,7 @@ func NewWGUDPProxy(wgPort int) *WGUDPProxy {
 	log.Debugf("Initializing new user space proxy with port %d", wgPort)
 	p := &WGUDPProxy{
 		localWGListenPort: wgPort,
+		closeListener:     listener.NewCloseListener(),
 	}
 	return p
 }
@@ -67,6 +71,10 @@ func (p *WGUDPProxy) EndpointAddr() *net.UDPAddr {
 	return endpointUdpAddr
 }
 
+func (p *WGUDPProxy) SetDisconnectListener(disconnected func()) {
+	p.closeListener.SetCloseListener(disconnected)
+}
+
 // Work starts the proxy or resumes it if it was paused
 func (p *WGUDPProxy) Work() {
 	if p.remoteConn == nil {
@@ -111,6 +119,8 @@ func (p *WGUDPProxy) close() error {
 	if p.closed {
 		return nil
 	}
+
+	p.closeListener.SetCloseListener(nil)
 	p.closed = true
 
 	p.cancel()
@@ -141,6 +151,7 @@ func (p *WGUDPProxy) proxyToRemote(ctx context.Context) {
 			if ctx.Err() != nil {
 				return
 			}
+			p.closeListener.Notify()
 			log.Debugf("failed to read from wg interface conn: %s", err)
 			return
 		}

client/internal/conn_mgr.go

@@ -226,7 +226,6 @@ func (e *ConnMgr) ActivatePeer(ctx context.Context, conn *peer.Conn) {
 	}
 
 	if found := e.lazyConnMgr.ActivatePeer(conn.GetKey()); found {
-		conn.Log.Infof("activated peer from inactive state")
 		if err := conn.Open(ctx); err != nil {
 			conn.Log.Errorf("failed to open connection: %v", err)
 		}

client/internal/debug/debug.go

@@ -167,6 +167,7 @@ type BundleGenerator struct {
 	anonymize         bool
 	clientStatus      string
 	includeSystemInfo bool
+	logFileCount      uint32
 
 	archive *zip.Writer
 }
@@ -175,6 +176,7 @@ type BundleConfig struct {
 	Anonymize         bool
 	ClientStatus      string
 	IncludeSystemInfo bool
+	LogFileCount      uint32
 }
 
 type GeneratorDependencies struct {
@@ -185,6 +187,12 @@ type GeneratorDependencies struct {
 }
 
 func NewBundleGenerator(deps GeneratorDependencies, cfg BundleConfig) *BundleGenerator {
+	// Default to 1 log file for backward compatibility when 0 is provided
+	logFileCount := cfg.LogFileCount
+	if logFileCount == 0 {
+		logFileCount = 1
+	}
+
 	return &BundleGenerator{
 		anonymizer: anonymize.NewAnonymizer(anonymize.DefaultAddresses()),
 
@@ -196,6 +204,7 @@ func NewBundleGenerator(deps GeneratorDependencies, cfg BundleConfig) *BundleGen
 		anonymize:         cfg.Anonymize,
 		clientStatus:      cfg.ClientStatus,
 		includeSystemInfo: cfg.IncludeSystemInfo,
+		logFileCount:      logFileCount,
 	}
 }
 
@@ -561,32 +570,8 @@ func (g *BundleGenerator) addLogfile() error {
 		return fmt.Errorf("add client log file to zip: %w", err)
 	}
 
-	// add latest rotated log file
-	pattern := filepath.Join(logDir, "client-*.log.gz")
-	files, err := filepath.Glob(pattern)
-	if err != nil {
-		log.Warnf("failed to glob rotated logs: %v", err)
-	} else if len(files) > 0 {
-		// pick the file with the latest ModTime
-		sort.Slice(files, func(i, j int) bool {
-			fi, err := os.Stat(files[i])
-			if err != nil {
-				log.Warnf("failed to stat rotated log %s: %v", files[i], err)
-				return false
-			}
-			fj, err := os.Stat(files[j])
-			if err != nil {
-				log.Warnf("failed to stat rotated log %s: %v", files[j], err)
-				return false
-			}
-			return fi.ModTime().Before(fj.ModTime())
-		})
-		latest := files[len(files)-1]
-		name := filepath.Base(latest)
-		if err := g.addSingleLogFileGz(latest, name); err != nil {
-			log.Warnf("failed to add rotated log %s: %v", name, err)
-		}
-	}
+	// add rotated log files based on logFileCount
+	g.addRotatedLogFiles(logDir)
 
 	stdErrLogPath := filepath.Join(logDir, errorLogFile)
 	stdoutLogPath := filepath.Join(logDir, stdoutLogFile)
@@ -670,6 +655,52 @@ func (g *BundleGenerator) addSingleLogFileGz(logPath, targetName string) error {
 	return nil
 }
 
+// addRotatedLogFiles adds rotated log files to the bundle based on logFileCount
+func (g *BundleGenerator) addRotatedLogFiles(logDir string) {
+	if g.logFileCount == 0 {
+		return
+	}
+
+	pattern := filepath.Join(logDir, "client-*.log.gz")
+	files, err := filepath.Glob(pattern)
+	if err != nil {
+		log.Warnf("failed to glob rotated logs: %v", err)
+		return
+	}
+
+	if len(files) == 0 {
+		return
+	}
+
+	// sort files by modification time (newest first)
+	sort.Slice(files, func(i, j int) bool {
+		fi, err := os.Stat(files[i])
+		if err != nil {
+			log.Warnf("failed to stat rotated log %s: %v", files[i], err)
+			return false
+		}
+		fj, err := os.Stat(files[j])
+		if err != nil {
+			log.Warnf("failed to stat rotated log %s: %v", files[j], err)
+			return false
+		}
+		return fi.ModTime().After(fj.ModTime())
+	})
+
+	// include up to logFileCount rotated files
+	maxFiles := int(g.logFileCount)
+	if maxFiles > len(files) {
+		maxFiles = len(files)
+	}
+
+	for i := 0; i < maxFiles; i++ {
+		name := filepath.Base(files[i])
+		if err := g.addSingleLogFileGz(files[i], name); err != nil {
+			log.Warnf("failed to add rotated log %s: %v", name, err)
+		}
+	}
+}
+
 func (g *BundleGenerator) addFileToZip(reader io.Reader, filename string) error {
 	header := &zip.FileHeader{
 		Name:     filename,

client/internal/dns/upstream.go

@@ -50,7 +50,7 @@ type upstreamResolverBase struct {
 	upstreamClient   upstreamClient
 	upstreamServers  []string
 	domain           string
-	disabled         bool
+	disabled         atomic.Bool
 	failsCount       atomic.Int32
 	successCount     atomic.Int32
 	failsTillDeact   int32
@@ -176,7 +176,7 @@ func (u *upstreamResolverBase) checkUpstreamFails(err error) {
 	u.mutex.Lock()
 	defer u.mutex.Unlock()
 
-	if u.failsCount.Load() < u.failsTillDeact || u.disabled {
+	if u.failsCount.Load() < u.failsTillDeact || u.disabled.Load() {
 		return
 	}
 
@@ -305,7 +305,7 @@ func (u *upstreamResolverBase) waitUntilResponse() {
 	u.failsCount.Store(0)
 	u.successCount.Add(1)
 	u.reactivate()
-	u.disabled = false
+	u.disabled.Store(false)
 }
 
 // isTimeout returns true if the given error is a network timeout error.
@@ -320,14 +320,14 @@ func isTimeout(err error) bool {
 }
 
 func (u *upstreamResolverBase) disable(err error) {
-	if u.disabled {
+	if u.disabled.Load() {
 		return
 	}
 
 	log.Warnf("Upstream resolving is Disabled for %v", reactivatePeriod)
 	u.successCount.Store(0)
 	u.deactivate(err)
-	u.disabled = true
+	u.disabled.Store(true)
 	go u.waitUntilResponse()
 }
 

client/internal/dns/upstream_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/netip"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -135,32 +136,41 @@ func TestUpstreamResolver_DeactivationReactivation(t *testing.T) {
 	responseWriter := &test.MockResponseWriter{
 		WriteMsgFunc: func(m *dns.Msg) error { return nil },
 	}
-
+	lmux := sync.Mutex{}
 	failed := false
 	resolver.deactivate = func(error) {
+		lmux.Lock()
 		failed = true
+		lmux.Unlock()
 	}
 
 	reactivated := false
 	resolver.reactivate = func() {
+		lmux.Lock()
 		reactivated = true
+		lmux.Unlock()
 	}
 
 	resolver.ServeDNS(responseWriter, new(dns.Msg).SetQuestion("one.one.one.one.", dns.TypeA))
-
-	if !failed {
+	lmux.Lock()
+	failedCheck := failed
+	lmux.Unlock()
+	if !failedCheck {
 		t.Errorf("expected that resolving was deactivated")
 		return
 	}
 
-	if !resolver.disabled {
+	if !resolver.disabled.Load() {
 		t.Errorf("resolver should be Disabled")
 		return
 	}
 
 	time.Sleep(time.Millisecond * 200)
 
-	if !reactivated {
+	lmux.Lock()
+	checkReactivated := reactivated
+	lmux.Unlock()
+	if !checkReactivated {
 		t.Errorf("expected that resolving was reactivated")
 		return
 	}
@@ -170,7 +180,7 @@ func TestUpstreamResolver_DeactivationReactivation(t *testing.T) {
 		return
 	}
 
-	if resolver.disabled {
+	if resolver.disabled.Load() {
 		t.Errorf("should be enabled")
 	}
 }

client/internal/engine.go

@@ -836,7 +836,10 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 			}
 			go func() {
 				// blocking
-				err = e.sshServer.Start()
+				e.syncMsgMux.Lock()
+				sshServer := e.sshServer
+				e.syncMsgMux.Unlock()
+				err = sshServer.Start()
 				if err != nil {
 					// will throw error when we stop it even if it is a graceful stop
 					log.Debugf("stopped SSH server with error %v", err)
@@ -851,6 +854,8 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 		}
 	} else if !isNil(e.sshServer) {
 		// Disable SSH server request, so stop it if it was running
+		e.syncMsgMux.Lock()
+		defer e.syncMsgMux.Unlock()
 		err := e.sshServer.Stop()
 		if err != nil {
 			log.Warnf("failed to stop SSH server %v", err)

client/internal/engine_test.go

@@ -52,6 +52,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/monotime"
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	"github.com/netbirdio/netbird/route"
 	signal "github.com/netbirdio/netbird/signal/client"
@@ -96,7 +97,7 @@ type MockWGIface struct {
 	GetInterfaceGUIDStringFunc func() (string, error)
 	GetProxyFunc               func() wgproxy.Proxy
 	GetNetFunc                 func() *netstack.Net
-	LastActivitiesFunc         func() map[string]time.Time
+	LastActivitiesFunc         func() map[string]monotime.Time
 }
 
 func (m *MockWGIface) FullStats() (*configurer.Stats, error) {
@@ -187,7 +188,7 @@ func (m *MockWGIface) GetNet() *netstack.Net {
 	return m.GetNetFunc()
 }
 
-func (m *MockWGIface) LastActivities() map[string]time.Time {
+func (m *MockWGIface) LastActivities() map[string]monotime.Time {
 	if m.LastActivitiesFunc != nil {
 		return m.LastActivitiesFunc()
 	}
@@ -1480,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)
 
@@ -1489,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
 	}

client/internal/iface_common.go

@@ -14,6 +14,7 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/iface/wgproxy"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 type wgIfaceBase interface {
@@ -38,5 +39,5 @@ type wgIfaceBase interface {
 	GetStats() (map[string]configurer.WGStats, error)
 	GetNet() *netstack.Net
 	FullStats() (*configurer.Stats, error)
-	LastActivities() map[string]time.Time
+	LastActivities() map[string]monotime.Time
 }

client/internal/lazyconn/activity/listener.go

@@ -48,7 +48,7 @@ func (d *Listener) ReadPackets() {
 		n, remoteAddr, err := d.conn.ReadFromUDP(make([]byte, 1))
 		if err != nil {
 			if d.isClosed.Load() {
-				d.peerCfg.Log.Debugf("exit from activity listener")
+				d.peerCfg.Log.Infof("exit from activity listener")
 			} else {
 				d.peerCfg.Log.Errorf("failed to read from activity listener: %s", err)
 			}
@@ -59,9 +59,11 @@ func (d *Listener) ReadPackets() {
 			d.peerCfg.Log.Warnf("received %d bytes from %s, too short", n, remoteAddr)
 			continue
 		}
+		d.peerCfg.Log.Infof("activity detected")
 		break
 	}
 
+	d.peerCfg.Log.Debugf("removing lazy endpoint: %s", d.endpoint.String())
 	if err := d.removeEndpoint(); err != nil {
 		d.peerCfg.Log.Errorf("failed to remove endpoint: %s", err)
 	}
@@ -71,7 +73,7 @@ func (d *Listener) ReadPackets() {
 }
 
 func (d *Listener) Close() {
-	d.peerCfg.Log.Infof("closing listener: %s", d.conn.LocalAddr().String())
+	d.peerCfg.Log.Infof("closing activity listener: %s", d.conn.LocalAddr().String())
 	d.isClosed.Store(true)
 
 	if err := d.conn.Close(); err != nil {
@@ -81,7 +83,6 @@ func (d *Listener) Close() {
 }
 
 func (d *Listener) removeEndpoint() error {
-	d.peerCfg.Log.Debugf("removing lazy endpoint: %s", d.endpoint.String())
 	return d.wgIface.RemovePeer(d.peerCfg.PublicKey)
 }
 

client/internal/lazyconn/activity/manager.go

@@ -102,3 +102,11 @@ func (m *Manager) notify(peerConnID peerid.ConnID) {
 	case m.OnActivityChan <- peerConnID:
 	}
 }
+
+func (m *Manager) getPeerListener(peerConnID peerid.ConnID) (*Listener, bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	listener, ok := m.peers[peerConnID]
+	return listener, ok
+}

client/internal/lazyconn/activity/manager_test.go

@@ -50,8 +50,11 @@ func TestManager_MonitorPeerActivity(t *testing.T) {
 	if err := mgr.MonitorPeerActivity(peerCfg1); err != nil {
 		t.Fatalf("failed to monitor peer activity: %v", err)
 	}
-
-	if err := trigger(mgr.peers[peerCfg1.PeerConnID].conn.LocalAddr().String()); err != nil {
+	listener, ok := mgr.getPeerListener(peerCfg1.PeerConnID)
+	if !ok {
+		t.Fatalf("failed to get peer listener: %s", peerCfg1.PublicKey)
+	}
+	if err := trigger(listener.conn.LocalAddr().String()); err != nil {
 		t.Fatalf("failed to trigger activity: %v", err)
 	}
 
@@ -83,7 +86,12 @@ func TestManager_RemovePeerActivity(t *testing.T) {
 		t.Fatalf("failed to monitor peer activity: %v", err)
 	}
 
-	addr := mgr.peers[peerCfg1.PeerConnID].conn.LocalAddr().String()
+	peer1Listener, ok := mgr.getPeerListener(peerCfg1.PeerConnID)
+	if !ok {
+		t.Fatalf("failed to get peer listener: %s", peerCfg1.PublicKey)
+	}
+
+	addr := peer1Listener.conn.LocalAddr().String()
 
 	mgr.RemovePeer(peerCfg1.Log, peerCfg1.PeerConnID)
 
@@ -128,11 +136,20 @@ func TestManager_MultiPeerActivity(t *testing.T) {
 		t.Fatalf("failed to monitor peer activity: %v", err)
 	}
 
-	if err := trigger(mgr.peers[peerCfg1.PeerConnID].conn.LocalAddr().String()); err != nil {
+	peer1Listener, ok := mgr.getPeerListener(peerCfg1.PeerConnID)
+	if !ok {
+		t.Fatalf("failed to get peer listener: %s", peerCfg1.PublicKey)
+	}
+	if err := trigger(peer1Listener.conn.LocalAddr().String()); err != nil {
 		t.Fatalf("failed to trigger activity: %v", err)
 	}
 
-	if err := trigger(mgr.peers[peerCfg2.PeerConnID].conn.LocalAddr().String()); err != nil {
+	peer2Listener, ok := mgr.getPeerListener(peerCfg2.PeerConnID)
+	if !ok {
+		t.Fatalf("failed to get peer listener: %s", peerCfg2.PublicKey)
+	}
+
+	if err := trigger(peer2Listener.conn.LocalAddr().String()); err != nil {
 		t.Fatalf("failed to trigger activity: %v", err)
 	}
 

client/internal/lazyconn/inactivity/manager.go

@@ -8,6 +8,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/internal/lazyconn"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 const (
@@ -18,7 +19,7 @@ const (
 )
 
 type WgInterface interface {
-	LastActivities() map[string]time.Time
+	LastActivities() map[string]monotime.Time
 }
 
 type Manager struct {
@@ -124,6 +125,7 @@ func (m *Manager) checkStats() (map[string]struct{}, error) {
 
 	idlePeers := make(map[string]struct{})
 
+	checkTime := time.Now()
 	for peerID, peerCfg := range m.interestedPeers {
 		lastActive, ok := lastActivities[peerID]
 		if !ok {
@@ -132,8 +134,9 @@ func (m *Manager) checkStats() (map[string]struct{}, error) {
 			continue
 		}
 
-		if time.Since(lastActive) > m.inactivityThreshold {
-			peerCfg.Log.Infof("peer is inactive since: %v", lastActive)
+		since := monotime.Since(lastActive)
+		if since > m.inactivityThreshold {
+			peerCfg.Log.Infof("peer is inactive since time: %s", checkTime.Add(-since).String())
 			idlePeers[peerID] = struct{}{}
 		}
 	}

client/internal/lazyconn/inactivity/manager_test.go

@@ -9,13 +9,14 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/netbirdio/netbird/client/internal/lazyconn"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 type mockWgInterface struct {
-	lastActivities map[string]time.Time
+	lastActivities map[string]monotime.Time
 }
 
-func (m *mockWgInterface) LastActivities() map[string]time.Time {
+func (m *mockWgInterface) LastActivities() map[string]monotime.Time {
 	return m.lastActivities
 }
 
@@ -23,8 +24,8 @@ func TestPeerTriggersInactivity(t *testing.T) {
 	peerID := "peer1"
 
 	wgMock := &mockWgInterface{
-		lastActivities: map[string]time.Time{
-			peerID: time.Now().Add(-20 * time.Minute),
+		lastActivities: map[string]monotime.Time{
+			peerID: monotime.Time(int64(monotime.Now()) - int64(20*time.Minute)),
 		},
 	}
 
@@ -64,8 +65,8 @@ func TestPeerTriggersActivity(t *testing.T) {
 	peerID := "peer1"
 
 	wgMock := &mockWgInterface{
-		lastActivities: map[string]time.Time{
-			peerID: time.Now().Add(-5 * time.Minute),
+		lastActivities: map[string]monotime.Time{
+			peerID: monotime.Time(int64(monotime.Now()) - int64(5*time.Minute)),
 		},
 	}
 

client/internal/lazyconn/manager/manager.go

@@ -258,12 +258,13 @@ func (m *Manager) ActivatePeer(peerID string) (found bool) {
 		return false
 	}
 
+	cfg.Log.Infof("activate peer from inactive state by remote signal message")
+
 	if !m.activateSinglePeer(cfg, mp) {
 		return false
 	}
 
 	m.activateHAGroupPeers(cfg)
-
 	return true
 }
 
@@ -571,12 +572,12 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
 		// this is blocking operation, potentially can be optimized
 		m.peerStore.PeerConnIdle(mp.peerCfg.PublicKey)
 
-		mp.peerCfg.Log.Infof("start activity monitor")
-
 		mp.expectedWatcher = watcherActivity
 
 		m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
 
+		mp.peerCfg.Log.Infof("start activity monitor")
+
 		if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
 			mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
 			continue

client/internal/lazyconn/wgiface.go

@@ -6,11 +6,13 @@ import (
 	"time"
 
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/monotime"
 )
 
 type WGIface interface {
 	RemovePeer(peerKey string) error
 	UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error
 	IsUserspaceBind() bool
-	LastActivities() map[string]time.Time
+	LastActivities() map[string]monotime.Time
 }

client/internal/networkmonitor/monitor.go

@@ -22,6 +22,19 @@ const (
 )
 
 var checkChangeFn = checkChange
+var mux sync.Mutex
+
+func getCheckChangeFn() func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
+	mux.Lock()
+	defer mux.Unlock()
+	return checkChangeFn
+}
+
+func setCheckChangeFn(fn func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error) {
+	mux.Lock()
+	defer mux.Unlock()
+	checkChangeFn = fn
+}
 
 // NetworkMonitor watches for changes in network configuration.
 type NetworkMonitor struct {
@@ -120,7 +133,8 @@ func (nw *NetworkMonitor) Stop() {
 func (nw *NetworkMonitor) checkChanges(ctx context.Context, event chan struct{}, nexthop4 systemops.Nexthop, nexthop6 systemops.Nexthop) {
 	defer close(event)
 	for {
-		if err := checkChangeFn(ctx, nexthop4, nexthop6); err != nil {
+		checkFn := getCheckChangeFn()
+		if err := checkFn(ctx, nexthop4, nexthop6); err != nil {
 			if !errors.Is(err, context.Canceled) {
 				log.Errorf("Network monitor: failed to check for changes: %v", err)
 			}

client/internal/networkmonitor/monitor_test.go

@@ -25,10 +25,10 @@ func (m *MocMultiEvent) checkChange(ctx context.Context, nexthopv4, nexthopv6 sy
 }
 
 func TestNetworkMonitor_Close(t *testing.T) {
-	checkChangeFn = func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
+	setCheckChangeFn(func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
 		<-ctx.Done()
 		return ctx.Err()
-	}
+	})
 	nw := New()
 
 	var resErr error
@@ -48,7 +48,7 @@ func TestNetworkMonitor_Close(t *testing.T) {
 }
 
 func TestNetworkMonitor_Event(t *testing.T) {
-	checkChangeFn = func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
+	setCheckChangeFn(func(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
 		timeout, cancel := context.WithTimeout(ctx, 3*time.Second)
 		defer cancel()
 		select {
@@ -57,7 +57,7 @@ func TestNetworkMonitor_Event(t *testing.T) {
 		case <-timeout.Done():
 			return nil
 		}
-	}
+	})
 	nw := New()
 	defer nw.Stop()
 
@@ -77,7 +77,7 @@ func TestNetworkMonitor_Event(t *testing.T) {
 func TestNetworkMonitor_MultiEvent(t *testing.T) {
 	eventsRepeated := 3
 	me := &MocMultiEvent{counter: eventsRepeated}
-	checkChangeFn = me.checkChange
+	setCheckChangeFn(me.checkChange)
 
 	nw := New()
 	defer nw.Stop()

client/internal/peer/conn.go

@@ -99,7 +99,7 @@ type Conn struct {
 
 	statusRelay         *worker.AtomicWorkerStatus
 	statusICE           *worker.AtomicWorkerStatus
-	currentConnPriority conntype.ConnPriority
+	currentConnPriority conntype.ConnPriorityStore
 	opened              bool // this flag is used to prevent close in case of not opened connection
 
 	workerICE   *WorkerICE
@@ -167,7 +167,7 @@ func (conn *Conn) Open(engineCtx context.Context) error {
 
 	conn.ctx, conn.ctxCancel = context.WithCancel(engineCtx)
 
-	conn.workerRelay = NewWorkerRelay(conn.Log, isController(conn.config), conn.config, conn, conn.relayManager, conn.dumpState)
+	conn.workerRelay = NewWorkerRelay(conn.ctx, conn.Log, isController(conn.config), conn.config, conn, conn.relayManager, conn.dumpState)
 
 	relayIsSupportedLocally := conn.workerRelay.RelayIsSupportedLocally()
 	workerICE, err := NewWorkerICE(conn.ctx, conn.Log, conn.config, conn, conn.signaler, conn.iFaceDiscover, conn.statusRecorder, relayIsSupportedLocally)
@@ -283,7 +283,7 @@ func (conn *Conn) Close(signalToRemote bool) {
 // doesn't block, discards the message if connection wasn't ready
 func (conn *Conn) OnRemoteAnswer(answer OfferAnswer) bool {
 	conn.dumpState.RemoteAnswer()
-	conn.Log.Infof("OnRemoteAnswer, priority: %s, status ICE: %s, status relay: %s", conn.currentConnPriority, conn.statusICE, conn.statusRelay)
+	conn.Log.Infof("OnRemoteAnswer, priority: %s, status ICE: %s, status relay: %s", conn.currentConnPriority.Get(), conn.statusICE, conn.statusRelay)
 	return conn.handshaker.OnRemoteAnswer(answer)
 }
 
@@ -353,8 +353,8 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 
 	// this never should happen, because Relay is the lower priority and ICE always close the deprecated connection before upgrade
 	// todo consider to remove this check
-	if conn.currentConnPriority > priority {
-		conn.Log.Infof("current connection priority (%s) is higher than the new one (%s), do not upgrade connection", conn.currentConnPriority, priority)
+	if conn.currentConnPriority.Get() > priority {
+		conn.Log.Infof("current connection priority (%s) is higher than the new one (%s), do not upgrade connection", conn.currentConnPriority.Get(), priority)
 		conn.statusICE.SetConnected()
 		conn.updateIceState(iceConnInfo)
 		return
@@ -408,7 +408,7 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 	}
 	wgConfigWorkaround()
 
-	conn.currentConnPriority = priority
+	conn.currentConnPriority.Set(priority)
 	conn.statusICE.SetConnected()
 	conn.updateIceState(iceConnInfo)
 	conn.doOnConnected(iceConnInfo.RosenpassPubKey, iceConnInfo.RosenpassAddr)
@@ -445,10 +445,10 @@ func (conn *Conn) onICEStateDisconnected() {
 			defer conn.wgWatcherWg.Done()
 			conn.workerRelay.EnableWgWatcher(conn.ctx)
 		}()
-		conn.currentConnPriority = conntype.Relay
+		conn.currentConnPriority.Set(conntype.Relay)
 	} else {
 		conn.Log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", conntype.None.String())
-		conn.currentConnPriority = conntype.None
+		conn.currentConnPriority.Set(conntype.None)
 	}
 
 	changed := conn.statusICE.Get() != worker.StatusDisconnected
@@ -489,12 +489,14 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		conn.Log.Errorf("failed to add relayed net.Conn to local proxy: %v", err)
 		return
 	}
+	wgProxy.SetDisconnectListener(conn.onRelayDisconnected)
+
 	conn.dumpState.NewLocalProxy()
 
 	conn.Log.Infof("created new wgProxy for relay connection: %s", wgProxy.EndpointAddr().String())
 
 	if conn.isICEActive() {
-		conn.Log.Debugf("do not switch to relay because current priority is: %s", conn.currentConnPriority.String())
+		conn.Log.Debugf("do not switch to relay because current priority is: %s", conn.currentConnPriority.Get().String())
 		conn.setRelayedProxy(wgProxy)
 		conn.statusRelay.SetConnected()
 		conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
@@ -522,7 +524,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 
 	wgConfigWorkaround()
 	conn.rosenpassRemoteKey = rci.rosenpassPubKey
-	conn.currentConnPriority = conntype.Relay
+	conn.currentConnPriority.Set(conntype.Relay)
 	conn.statusRelay.SetConnected()
 	conn.setRelayedProxy(wgProxy)
 	conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
@@ -540,9 +542,9 @@ func (conn *Conn) onRelayDisconnected() {
 
 	conn.Log.Debugf("relay connection is disconnected")
 
-	if conn.currentConnPriority == conntype.Relay {
+	if conn.currentConnPriority.Get() == conntype.Relay {
 		conn.Log.Debugf("clean up WireGuard config")
-		conn.currentConnPriority = conntype.None
+		conn.currentConnPriority.Set(conntype.None)
 	}
 
 	if conn.wgProxyRelay != nil {
@@ -624,7 +626,7 @@ func (conn *Conn) updateIceState(iceConnInfo ICEConnInfo) {
 func (conn *Conn) setStatusToDisconnected() {
 	conn.statusRelay.SetDisconnected()
 	conn.statusICE.SetDisconnected()
-	conn.currentConnPriority = conntype.None
+	conn.currentConnPriority.Set(conntype.None)
 
 	peerState := State{
 		PubKey:           conn.config.Key,
@@ -667,7 +669,7 @@ func (conn *Conn) waitInitialRandomSleepTime(ctx context.Context) {
 }
 
 func (conn *Conn) isRelayed() bool {
-	switch conn.currentConnPriority {
+	switch conn.currentConnPriority.Get() {
 	case conntype.Relay, conntype.ICETurn:
 		return true
 	default:
@@ -751,11 +753,11 @@ func (conn *Conn) newProxy(remoteConn net.Conn) (wgproxy.Proxy, error) {
 }
 
 func (conn *Conn) isReadyToUpgrade() bool {
-	return conn.wgProxyRelay != nil && conn.currentConnPriority != conntype.Relay
+	return conn.wgProxyRelay != nil && conn.currentConnPriority.Get() != conntype.Relay
 }
 
 func (conn *Conn) isICEActive() bool {
-	return (conn.currentConnPriority == conntype.ICEP2P || conn.currentConnPriority == conntype.ICETurn) && conn.statusICE.Get() == worker.StatusConnected
+	return (conn.currentConnPriority.Get() == conntype.ICEP2P || conn.currentConnPriority.Get() == conntype.ICETurn) && conn.statusICE.Get() == worker.StatusConnected
 }
 
 func (conn *Conn) removeWgPeer() error {

client/internal/peer/conntype/priority.go

@@ -2,6 +2,7 @@ package conntype
 
 import (
 	"fmt"
+	"sync/atomic"
 )
 
 const (
@@ -11,7 +12,7 @@ const (
 	ICEP2P  ConnPriority = 3
 )
 
-type ConnPriority int
+type ConnPriority int32
 
 func (cp ConnPriority) String() string {
 	switch cp {
@@ -27,3 +28,15 @@ func (cp ConnPriority) String() string {
 		return fmt.Sprintf("ConnPriority(%d)", cp)
 	}
 }
+
+type ConnPriorityStore struct {
+	store atomic.Int32
+}
+
+func (cps *ConnPriorityStore) Get() ConnPriority {
+	return ConnPriority(cps.store.Load())
+}
+
+func (cps *ConnPriorityStore) Set(cp ConnPriority) {
+	cps.store.Store(int32(cp))
+}

client/internal/peer/notifier_test.go

@@ -9,30 +9,54 @@ type mocListener struct {
 	lastState int
 	wg        sync.WaitGroup
 	peers     int
+	mux       sync.Mutex
 }
 
 func (l *mocListener) OnConnected() {
+	l.mux.Lock()
+	defer l.mux.Unlock()
 	l.lastState = stateConnected
 	l.wg.Done()
 }
 func (l *mocListener) OnDisconnected() {
+	l.mux.Lock()
+	defer l.mux.Unlock()
 	l.lastState = stateDisconnected
 	l.wg.Done()
 }
 func (l *mocListener) OnConnecting() {
+	l.mux.Lock()
+	defer l.mux.Unlock()
 	l.lastState = stateConnecting
 	l.wg.Done()
 }
 func (l *mocListener) OnDisconnecting() {
+	l.mux.Lock()
+	defer l.mux.Unlock()
 	l.lastState = stateDisconnecting
 	l.wg.Done()
+
+}
+
+func (l *mocListener) getLastState() int {
+	l.mux.Lock()
+	defer l.mux.Unlock()
+	return l.lastState
 }
 
 func (l *mocListener) OnAddressChanged(host, addr string) {
 
 }
 func (l *mocListener) OnPeersListChanged(size int) {
+	l.mux.Lock()
 	l.peers = size
+	l.mux.Unlock()
+}
+
+func (l *mocListener) getPeers() int {
+	l.mux.Lock()
+	defer l.mux.Unlock()
+	return l.peers
 }
 
 func (l *mocListener) setWaiter() {
@@ -77,7 +101,7 @@ func Test_notifier_SetListener(t *testing.T) {
 	n.lastNotification = stateConnecting
 	n.setListener(listener)
 	listener.wait()
-	if listener.lastState != n.lastNotification {
+	if listener.getLastState() != n.lastNotification {
 		t.Errorf("invalid state: %d, expected: %d", listener.lastState, n.lastNotification)
 	}
 }
@@ -91,7 +115,7 @@ func Test_notifier_RemoveListener(t *testing.T) {
 	n.removeListener()
 	n.peerListChanged(1)
 
-	if listener.peers != 0 {
+	if listener.getPeers() != 0 {
 		t.Errorf("invalid state: %d", listener.peers)
 	}
 }

client/internal/peer/worker_relay.go

@@ -19,6 +19,7 @@ type RelayConnInfo struct {
 }
 
 type WorkerRelay struct {
+	peerCtx      context.Context
 	log          *log.Entry
 	isController bool
 	config       ConnConfig
@@ -33,8 +34,9 @@ type WorkerRelay struct {
 	wgWatcher *WGWatcher
 }
 
-func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, conn *Conn, relayManager relayClient.ManagerService, stateDump *stateDump) *WorkerRelay {
+func NewWorkerRelay(ctx context.Context, log *log.Entry, ctrl bool, config ConnConfig, conn *Conn, relayManager relayClient.ManagerService, stateDump *stateDump) *WorkerRelay {
 	r := &WorkerRelay{
+		peerCtx:      ctx,
 		log:          log,
 		isController: ctrl,
 		config:       config,
@@ -62,7 +64,7 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 
 	srv := w.preferredRelayServer(currentRelayAddress, remoteOfferAnswer.RelaySrvAddress)
 
-	relayedConn, err := w.relayManager.OpenConn(srv, w.config.Key)
+	relayedConn, err := w.relayManager.OpenConn(w.peerCtx, srv, w.config.Key)
 	if err != nil {
 		if errors.Is(err, relayClient.ErrConnAlreadyExists) {
 			w.log.Debugf("handled offer by reusing existing relay connection")

client/proto/daemon.pb.go

@@ -2290,6 +2290,7 @@ type DebugBundleRequest struct {
 	Status        string                 `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"`
 	SystemInfo    bool                   `protobuf:"varint,3,opt,name=systemInfo,proto3" json:"systemInfo,omitempty"`
 	UploadURL     string                 `protobuf:"bytes,4,opt,name=uploadURL,proto3" json:"uploadURL,omitempty"`
+	LogFileCount  uint32                 `protobuf:"varint,5,opt,name=logFileCount,proto3" json:"logFileCount,omitempty"`
 	unknownFields protoimpl.UnknownFields
 	sizeCache     protoimpl.SizeCache
 }
@@ -2352,6 +2353,13 @@ func (x *DebugBundleRequest) GetUploadURL() string {
 	return ""
 }
 
+func (x *DebugBundleRequest) GetLogFileCount() uint32 {
+	if x != nil {
+		return x.LogFileCount
+	}
+	return 0
+}
+
 type DebugBundleResponse struct {
 	state               protoimpl.MessageState `protogen:"open.v1"`
 	Path                string                 `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"`
@@ -3746,14 +3754,15 @@ const file_daemon_proto_rawDesc = "" +
 	"\x12translatedHostname\x18\x04 \x01(\tR\x12translatedHostname\x128\n" +
 	"\x0etranslatedPort\x18\x05 \x01(\v2\x10.daemon.PortInfoR\x0etranslatedPort\"G\n" +
 	"\x17ForwardingRulesResponse\x12,\n" +
-	"\x05rules\x18\x01 \x03(\v2\x16.daemon.ForwardingRuleR\x05rules\"\x88\x01\n" +
+	"\x05rules\x18\x01 \x03(\v2\x16.daemon.ForwardingRuleR\x05rules\"\xac\x01\n" +
 	"\x12DebugBundleRequest\x12\x1c\n" +
 	"\tanonymize\x18\x01 \x01(\bR\tanonymize\x12\x16\n" +
 	"\x06status\x18\x02 \x01(\tR\x06status\x12\x1e\n" +
 	"\n" +
 	"systemInfo\x18\x03 \x01(\bR\n" +
 	"systemInfo\x12\x1c\n" +
-	"\tuploadURL\x18\x04 \x01(\tR\tuploadURL\"}\n" +
+	"\tuploadURL\x18\x04 \x01(\tR\tuploadURL\x12\"\n" +
+	"\flogFileCount\x18\x05 \x01(\rR\flogFileCount\"}\n" +
 	"\x13DebugBundleResponse\x12\x12\n" +
 	"\x04path\x18\x01 \x01(\tR\x04path\x12 \n" +
 	"\vuploadedKey\x18\x02 \x01(\tR\vuploadedKey\x120\n" +

client/proto/daemon.proto

@@ -356,6 +356,7 @@ message DebugBundleRequest {
   string status = 2;
   bool systemInfo = 3;
   string uploadURL = 4;
+  uint32 logFileCount = 5;
 }
 
 message DebugBundleResponse {

client/proto/daemon_grpc.pb.go

@@ -1,8 +1,4 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
-// versions:
-// - protoc-gen-go-grpc v1.5.1
-// - protoc             v5.29.3
-// source: daemon.proto
 
 package proto
 
@@ -15,31 +11,8 @@ import (
 
 // This is a compile-time assertion to ensure that this generated file
 // is compatible with the grpc package it is being compiled against.
-// Requires gRPC-Go v1.64.0 or later.
-const _ = grpc.SupportPackageIsVersion9
-
-const (
-	DaemonService_Login_FullMethodName                    = "/daemon.DaemonService/Login"
-	DaemonService_WaitSSOLogin_FullMethodName             = "/daemon.DaemonService/WaitSSOLogin"
-	DaemonService_Up_FullMethodName                       = "/daemon.DaemonService/Up"
-	DaemonService_Status_FullMethodName                   = "/daemon.DaemonService/Status"
-	DaemonService_Down_FullMethodName                     = "/daemon.DaemonService/Down"
-	DaemonService_GetConfig_FullMethodName                = "/daemon.DaemonService/GetConfig"
-	DaemonService_ListNetworks_FullMethodName             = "/daemon.DaemonService/ListNetworks"
-	DaemonService_SelectNetworks_FullMethodName           = "/daemon.DaemonService/SelectNetworks"
-	DaemonService_DeselectNetworks_FullMethodName         = "/daemon.DaemonService/DeselectNetworks"
-	DaemonService_ForwardingRules_FullMethodName          = "/daemon.DaemonService/ForwardingRules"
-	DaemonService_DebugBundle_FullMethodName              = "/daemon.DaemonService/DebugBundle"
-	DaemonService_GetLogLevel_FullMethodName              = "/daemon.DaemonService/GetLogLevel"
-	DaemonService_SetLogLevel_FullMethodName              = "/daemon.DaemonService/SetLogLevel"
-	DaemonService_ListStates_FullMethodName               = "/daemon.DaemonService/ListStates"
-	DaemonService_CleanState_FullMethodName               = "/daemon.DaemonService/CleanState"
-	DaemonService_DeleteState_FullMethodName              = "/daemon.DaemonService/DeleteState"
-	DaemonService_SetNetworkMapPersistence_FullMethodName = "/daemon.DaemonService/SetNetworkMapPersistence"
-	DaemonService_TracePacket_FullMethodName              = "/daemon.DaemonService/TracePacket"
-	DaemonService_SubscribeEvents_FullMethodName          = "/daemon.DaemonService/SubscribeEvents"
-	DaemonService_GetEvents_FullMethodName                = "/daemon.DaemonService/GetEvents"
-)
+// Requires gRPC-Go v1.32.0 or later.
+const _ = grpc.SupportPackageIsVersion7
 
 // DaemonServiceClient is the client API for DaemonService service.
 //
@@ -80,7 +53,7 @@ type DaemonServiceClient interface {
 	// SetNetworkMapPersistence enables or disables network map persistence
 	SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error)
 	TracePacket(ctx context.Context, in *TracePacketRequest, opts ...grpc.CallOption) (*TracePacketResponse, error)
-	SubscribeEvents(ctx context.Context, in *SubscribeRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[SystemEvent], error)
+	SubscribeEvents(ctx context.Context, in *SubscribeRequest, opts ...grpc.CallOption) (DaemonService_SubscribeEventsClient, error)
 	GetEvents(ctx context.Context, in *GetEventsRequest, opts ...grpc.CallOption) (*GetEventsResponse, error)
 }
 
@@ -93,9 +66,8 @@ func NewDaemonServiceClient(cc grpc.ClientConnInterface) DaemonServiceClient {
 }
 
 func (c *daemonServiceClient) Login(ctx context.Context, in *LoginRequest, opts ...grpc.CallOption) (*LoginResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(LoginResponse)
-	err := c.cc.Invoke(ctx, DaemonService_Login_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/Login", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -103,9 +75,8 @@ func (c *daemonServiceClient) Login(ctx context.Context, in *LoginRequest, opts
 }
 
 func (c *daemonServiceClient) WaitSSOLogin(ctx context.Context, in *WaitSSOLoginRequest, opts ...grpc.CallOption) (*WaitSSOLoginResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(WaitSSOLoginResponse)
-	err := c.cc.Invoke(ctx, DaemonService_WaitSSOLogin_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/WaitSSOLogin", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -113,9 +84,8 @@ func (c *daemonServiceClient) WaitSSOLogin(ctx context.Context, in *WaitSSOLogin
 }
 
 func (c *daemonServiceClient) Up(ctx context.Context, in *UpRequest, opts ...grpc.CallOption) (*UpResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(UpResponse)
-	err := c.cc.Invoke(ctx, DaemonService_Up_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/Up", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -123,9 +93,8 @@ func (c *daemonServiceClient) Up(ctx context.Context, in *UpRequest, opts ...grp
 }
 
 func (c *daemonServiceClient) Status(ctx context.Context, in *StatusRequest, opts ...grpc.CallOption) (*StatusResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(StatusResponse)
-	err := c.cc.Invoke(ctx, DaemonService_Status_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/Status", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -133,9 +102,8 @@ func (c *daemonServiceClient) Status(ctx context.Context, in *StatusRequest, opt
 }
 
 func (c *daemonServiceClient) Down(ctx context.Context, in *DownRequest, opts ...grpc.CallOption) (*DownResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(DownResponse)
-	err := c.cc.Invoke(ctx, DaemonService_Down_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/Down", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -143,9 +111,8 @@ func (c *daemonServiceClient) Down(ctx context.Context, in *DownRequest, opts ..
 }
 
 func (c *daemonServiceClient) GetConfig(ctx context.Context, in *GetConfigRequest, opts ...grpc.CallOption) (*GetConfigResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(GetConfigResponse)
-	err := c.cc.Invoke(ctx, DaemonService_GetConfig_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/GetConfig", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -153,9 +120,8 @@ func (c *daemonServiceClient) GetConfig(ctx context.Context, in *GetConfigReques
 }
 
 func (c *daemonServiceClient) ListNetworks(ctx context.Context, in *ListNetworksRequest, opts ...grpc.CallOption) (*ListNetworksResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(ListNetworksResponse)
-	err := c.cc.Invoke(ctx, DaemonService_ListNetworks_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -163,9 +129,8 @@ func (c *daemonServiceClient) ListNetworks(ctx context.Context, in *ListNetworks
 }
 
 func (c *daemonServiceClient) SelectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(SelectNetworksResponse)
-	err := c.cc.Invoke(ctx, DaemonService_SelectNetworks_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SelectNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -173,9 +138,8 @@ func (c *daemonServiceClient) SelectNetworks(ctx context.Context, in *SelectNetw
 }
 
 func (c *daemonServiceClient) DeselectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(SelectNetworksResponse)
-	err := c.cc.Invoke(ctx, DaemonService_DeselectNetworks_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeselectNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -183,9 +147,8 @@ func (c *daemonServiceClient) DeselectNetworks(ctx context.Context, in *SelectNe
 }
 
 func (c *daemonServiceClient) ForwardingRules(ctx context.Context, in *EmptyRequest, opts ...grpc.CallOption) (*ForwardingRulesResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(ForwardingRulesResponse)
-	err := c.cc.Invoke(ctx, DaemonService_ForwardingRules_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ForwardingRules", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -193,9 +156,8 @@ func (c *daemonServiceClient) ForwardingRules(ctx context.Context, in *EmptyRequ
 }
 
 func (c *daemonServiceClient) DebugBundle(ctx context.Context, in *DebugBundleRequest, opts ...grpc.CallOption) (*DebugBundleResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(DebugBundleResponse)
-	err := c.cc.Invoke(ctx, DaemonService_DebugBundle_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DebugBundle", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -203,9 +165,8 @@ func (c *daemonServiceClient) DebugBundle(ctx context.Context, in *DebugBundleRe
 }
 
 func (c *daemonServiceClient) GetLogLevel(ctx context.Context, in *GetLogLevelRequest, opts ...grpc.CallOption) (*GetLogLevelResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(GetLogLevelResponse)
-	err := c.cc.Invoke(ctx, DaemonService_GetLogLevel_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/GetLogLevel", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -213,9 +174,8 @@ func (c *daemonServiceClient) GetLogLevel(ctx context.Context, in *GetLogLevelRe
 }
 
 func (c *daemonServiceClient) SetLogLevel(ctx context.Context, in *SetLogLevelRequest, opts ...grpc.CallOption) (*SetLogLevelResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(SetLogLevelResponse)
-	err := c.cc.Invoke(ctx, DaemonService_SetLogLevel_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SetLogLevel", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -223,9 +183,8 @@ func (c *daemonServiceClient) SetLogLevel(ctx context.Context, in *SetLogLevelRe
 }
 
 func (c *daemonServiceClient) ListStates(ctx context.Context, in *ListStatesRequest, opts ...grpc.CallOption) (*ListStatesResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(ListStatesResponse)
-	err := c.cc.Invoke(ctx, DaemonService_ListStates_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListStates", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -233,9 +192,8 @@ func (c *daemonServiceClient) ListStates(ctx context.Context, in *ListStatesRequ
 }
 
 func (c *daemonServiceClient) CleanState(ctx context.Context, in *CleanStateRequest, opts ...grpc.CallOption) (*CleanStateResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(CleanStateResponse)
-	err := c.cc.Invoke(ctx, DaemonService_CleanState_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/CleanState", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -243,9 +201,8 @@ func (c *daemonServiceClient) CleanState(ctx context.Context, in *CleanStateRequ
 }
 
 func (c *daemonServiceClient) DeleteState(ctx context.Context, in *DeleteStateRequest, opts ...grpc.CallOption) (*DeleteStateResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(DeleteStateResponse)
-	err := c.cc.Invoke(ctx, DaemonService_DeleteState_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeleteState", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -253,9 +210,8 @@ func (c *daemonServiceClient) DeleteState(ctx context.Context, in *DeleteStateRe
 }
 
 func (c *daemonServiceClient) SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(SetNetworkMapPersistenceResponse)
-	err := c.cc.Invoke(ctx, DaemonService_SetNetworkMapPersistence_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SetNetworkMapPersistence", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -263,22 +219,20 @@ func (c *daemonServiceClient) SetNetworkMapPersistence(ctx context.Context, in *
 }
 
 func (c *daemonServiceClient) TracePacket(ctx context.Context, in *TracePacketRequest, opts ...grpc.CallOption) (*TracePacketResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(TracePacketResponse)
-	err := c.cc.Invoke(ctx, DaemonService_TracePacket_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/TracePacket", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
 	return out, nil
 }
 
-func (c *daemonServiceClient) SubscribeEvents(ctx context.Context, in *SubscribeRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[SystemEvent], error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
-	stream, err := c.cc.NewStream(ctx, &DaemonService_ServiceDesc.Streams[0], DaemonService_SubscribeEvents_FullMethodName, cOpts...)
+func (c *daemonServiceClient) SubscribeEvents(ctx context.Context, in *SubscribeRequest, opts ...grpc.CallOption) (DaemonService_SubscribeEventsClient, error) {
+	stream, err := c.cc.NewStream(ctx, &DaemonService_ServiceDesc.Streams[0], "/daemon.DaemonService/SubscribeEvents", opts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &grpc.GenericClientStream[SubscribeRequest, SystemEvent]{ClientStream: stream}
+	x := &daemonServiceSubscribeEventsClient{stream}
 	if err := x.ClientStream.SendMsg(in); err != nil {
 		return nil, err
 	}
@@ -288,13 +242,26 @@ func (c *daemonServiceClient) SubscribeEvents(ctx context.Context, in *Subscribe
 	return x, nil
 }
 
-// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
-type DaemonService_SubscribeEventsClient = grpc.ServerStreamingClient[SystemEvent]
+type DaemonService_SubscribeEventsClient interface {
+	Recv() (*SystemEvent, error)
+	grpc.ClientStream
+}
+
+type daemonServiceSubscribeEventsClient struct {
+	grpc.ClientStream
+}
+
+func (x *daemonServiceSubscribeEventsClient) Recv() (*SystemEvent, error) {
+	m := new(SystemEvent)
+	if err := x.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
 
 func (c *daemonServiceClient) GetEvents(ctx context.Context, in *GetEventsRequest, opts ...grpc.CallOption) (*GetEventsResponse, error) {
-	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(GetEventsResponse)
-	err := c.cc.Invoke(ctx, DaemonService_GetEvents_FullMethodName, in, out, cOpts...)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/GetEvents", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -303,7 +270,7 @@ func (c *daemonServiceClient) GetEvents(ctx context.Context, in *GetEventsReques
 
 // DaemonServiceServer is the server API for DaemonService service.
 // All implementations must embed UnimplementedDaemonServiceServer
-// for forward compatibility.
+// for forward compatibility
 type DaemonServiceServer interface {
 	// Login uses setup key to prepare configuration for the daemon.
 	Login(context.Context, *LoginRequest) (*LoginResponse, error)
@@ -340,17 +307,14 @@ type DaemonServiceServer interface {
 	// SetNetworkMapPersistence enables or disables network map persistence
 	SetNetworkMapPersistence(context.Context, *SetNetworkMapPersistenceRequest) (*SetNetworkMapPersistenceResponse, error)
 	TracePacket(context.Context, *TracePacketRequest) (*TracePacketResponse, error)
-	SubscribeEvents(*SubscribeRequest, grpc.ServerStreamingServer[SystemEvent]) error
+	SubscribeEvents(*SubscribeRequest, DaemonService_SubscribeEventsServer) error
 	GetEvents(context.Context, *GetEventsRequest) (*GetEventsResponse, error)
 	mustEmbedUnimplementedDaemonServiceServer()
 }
 
-// UnimplementedDaemonServiceServer must be embedded to have
-// forward compatible implementations.
-//
-// NOTE: this should be embedded by value instead of pointer to avoid a nil
-// pointer dereference when methods are called.
-type UnimplementedDaemonServiceServer struct{}
+// UnimplementedDaemonServiceServer must be embedded to have forward compatible implementations.
+type UnimplementedDaemonServiceServer struct {
+}
 
 func (UnimplementedDaemonServiceServer) Login(context.Context, *LoginRequest) (*LoginResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method Login not implemented")
@@ -406,14 +370,13 @@ func (UnimplementedDaemonServiceServer) SetNetworkMapPersistence(context.Context
 func (UnimplementedDaemonServiceServer) TracePacket(context.Context, *TracePacketRequest) (*TracePacketResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method TracePacket not implemented")
 }
-func (UnimplementedDaemonServiceServer) SubscribeEvents(*SubscribeRequest, grpc.ServerStreamingServer[SystemEvent]) error {
+func (UnimplementedDaemonServiceServer) SubscribeEvents(*SubscribeRequest, DaemonService_SubscribeEventsServer) error {
 	return status.Errorf(codes.Unimplemented, "method SubscribeEvents not implemented")
 }
 func (UnimplementedDaemonServiceServer) GetEvents(context.Context, *GetEventsRequest) (*GetEventsResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method GetEvents not implemented")
 }
 func (UnimplementedDaemonServiceServer) mustEmbedUnimplementedDaemonServiceServer() {}
-func (UnimplementedDaemonServiceServer) testEmbeddedByValue()                       {}
 
 // UnsafeDaemonServiceServer may be embedded to opt out of forward compatibility for this service.
 // Use of this interface is not recommended, as added methods to DaemonServiceServer will
@@ -423,13 +386,6 @@ type UnsafeDaemonServiceServer interface {
 }
 
 func RegisterDaemonServiceServer(s grpc.ServiceRegistrar, srv DaemonServiceServer) {
-	// If the following call pancis, it indicates UnimplementedDaemonServiceServer was
-	// embedded by pointer and is nil.  This will cause panics if an
-	// unimplemented method is ever invoked, so we test this at initialization
-	// time to prevent it from happening at runtime later due to I/O.
-	if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
-		t.testEmbeddedByValue()
-	}
 	s.RegisterService(&DaemonService_ServiceDesc, srv)
 }
 
@@ -443,7 +399,7 @@ func _DaemonService_Login_Handler(srv interface{}, ctx context.Context, dec func
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_Login_FullMethodName,
+		FullMethod: "/daemon.DaemonService/Login",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).Login(ctx, req.(*LoginRequest))
@@ -461,7 +417,7 @@ func _DaemonService_WaitSSOLogin_Handler(srv interface{}, ctx context.Context, d
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_WaitSSOLogin_FullMethodName,
+		FullMethod: "/daemon.DaemonService/WaitSSOLogin",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).WaitSSOLogin(ctx, req.(*WaitSSOLoginRequest))
@@ -479,7 +435,7 @@ func _DaemonService_Up_Handler(srv interface{}, ctx context.Context, dec func(in
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_Up_FullMethodName,
+		FullMethod: "/daemon.DaemonService/Up",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).Up(ctx, req.(*UpRequest))
@@ -497,7 +453,7 @@ func _DaemonService_Status_Handler(srv interface{}, ctx context.Context, dec fun
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_Status_FullMethodName,
+		FullMethod: "/daemon.DaemonService/Status",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).Status(ctx, req.(*StatusRequest))
@@ -515,7 +471,7 @@ func _DaemonService_Down_Handler(srv interface{}, ctx context.Context, dec func(
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_Down_FullMethodName,
+		FullMethod: "/daemon.DaemonService/Down",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).Down(ctx, req.(*DownRequest))
@@ -533,7 +489,7 @@ func _DaemonService_GetConfig_Handler(srv interface{}, ctx context.Context, dec
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_GetConfig_FullMethodName,
+		FullMethod: "/daemon.DaemonService/GetConfig",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).GetConfig(ctx, req.(*GetConfigRequest))
@@ -551,7 +507,7 @@ func _DaemonService_ListNetworks_Handler(srv interface{}, ctx context.Context, d
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_ListNetworks_FullMethodName,
+		FullMethod: "/daemon.DaemonService/ListNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).ListNetworks(ctx, req.(*ListNetworksRequest))
@@ -569,7 +525,7 @@ func _DaemonService_SelectNetworks_Handler(srv interface{}, ctx context.Context,
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_SelectNetworks_FullMethodName,
+		FullMethod: "/daemon.DaemonService/SelectNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).SelectNetworks(ctx, req.(*SelectNetworksRequest))
@@ -587,7 +543,7 @@ func _DaemonService_DeselectNetworks_Handler(srv interface{}, ctx context.Contex
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_DeselectNetworks_FullMethodName,
+		FullMethod: "/daemon.DaemonService/DeselectNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).DeselectNetworks(ctx, req.(*SelectNetworksRequest))
@@ -605,7 +561,7 @@ func _DaemonService_ForwardingRules_Handler(srv interface{}, ctx context.Context
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_ForwardingRules_FullMethodName,
+		FullMethod: "/daemon.DaemonService/ForwardingRules",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).ForwardingRules(ctx, req.(*EmptyRequest))
@@ -623,7 +579,7 @@ func _DaemonService_DebugBundle_Handler(srv interface{}, ctx context.Context, de
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_DebugBundle_FullMethodName,
+		FullMethod: "/daemon.DaemonService/DebugBundle",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).DebugBundle(ctx, req.(*DebugBundleRequest))
@@ -641,7 +597,7 @@ func _DaemonService_GetLogLevel_Handler(srv interface{}, ctx context.Context, de
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_GetLogLevel_FullMethodName,
+		FullMethod: "/daemon.DaemonService/GetLogLevel",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).GetLogLevel(ctx, req.(*GetLogLevelRequest))
@@ -659,7 +615,7 @@ func _DaemonService_SetLogLevel_Handler(srv interface{}, ctx context.Context, de
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_SetLogLevel_FullMethodName,
+		FullMethod: "/daemon.DaemonService/SetLogLevel",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).SetLogLevel(ctx, req.(*SetLogLevelRequest))
@@ -677,7 +633,7 @@ func _DaemonService_ListStates_Handler(srv interface{}, ctx context.Context, dec
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_ListStates_FullMethodName,
+		FullMethod: "/daemon.DaemonService/ListStates",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).ListStates(ctx, req.(*ListStatesRequest))
@@ -695,7 +651,7 @@ func _DaemonService_CleanState_Handler(srv interface{}, ctx context.Context, dec
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_CleanState_FullMethodName,
+		FullMethod: "/daemon.DaemonService/CleanState",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).CleanState(ctx, req.(*CleanStateRequest))
@@ -713,7 +669,7 @@ func _DaemonService_DeleteState_Handler(srv interface{}, ctx context.Context, de
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_DeleteState_FullMethodName,
+		FullMethod: "/daemon.DaemonService/DeleteState",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).DeleteState(ctx, req.(*DeleteStateRequest))
@@ -731,7 +687,7 @@ func _DaemonService_SetNetworkMapPersistence_Handler(srv interface{}, ctx contex
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_SetNetworkMapPersistence_FullMethodName,
+		FullMethod: "/daemon.DaemonService/SetNetworkMapPersistence",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).SetNetworkMapPersistence(ctx, req.(*SetNetworkMapPersistenceRequest))
@@ -749,7 +705,7 @@ func _DaemonService_TracePacket_Handler(srv interface{}, ctx context.Context, de
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_TracePacket_FullMethodName,
+		FullMethod: "/daemon.DaemonService/TracePacket",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).TracePacket(ctx, req.(*TracePacketRequest))
@@ -762,11 +718,21 @@ func _DaemonService_SubscribeEvents_Handler(srv interface{}, stream grpc.ServerS
 	if err := stream.RecvMsg(m); err != nil {
 		return err
 	}
-	return srv.(DaemonServiceServer).SubscribeEvents(m, &grpc.GenericServerStream[SubscribeRequest, SystemEvent]{ServerStream: stream})
+	return srv.(DaemonServiceServer).SubscribeEvents(m, &daemonServiceSubscribeEventsServer{stream})
 }
 
-// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
-type DaemonService_SubscribeEventsServer = grpc.ServerStreamingServer[SystemEvent]
+type DaemonService_SubscribeEventsServer interface {
+	Send(*SystemEvent) error
+	grpc.ServerStream
+}
+
+type daemonServiceSubscribeEventsServer struct {
+	grpc.ServerStream
+}
+
+func (x *daemonServiceSubscribeEventsServer) Send(m *SystemEvent) error {
+	return x.ServerStream.SendMsg(m)
+}
 
 func _DaemonService_GetEvents_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
 	in := new(GetEventsRequest)
@@ -778,7 +744,7 @@ func _DaemonService_GetEvents_Handler(srv interface{}, ctx context.Context, dec
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: DaemonService_GetEvents_FullMethodName,
+		FullMethod: "/daemon.DaemonService/GetEvents",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return srv.(DaemonServiceServer).GetEvents(ctx, req.(*GetEventsRequest))

client/proto/generate.sh

@@ -11,7 +11,7 @@ fi
 old_pwd=$(pwd)
 script_path=$(dirname $(realpath "$0"))
 cd "$script_path"
-go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.26
+go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.36.6
 go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.1
 protoc -I ./ ./daemon.proto --go_out=../ --go-grpc_out=../ --experimental_allow_proto3_optional
 cd "$old_pwd"

client/server/debug.go

@@ -42,6 +42,7 @@ func (s *Server) DebugBundle(_ context.Context, req *proto.DebugBundleRequest) (
 			Anonymize:         req.GetAnonymize(),
 			ClientStatus:      req.GetStatus(),
 			IncludeSystemInfo: req.GetSystemInfo(),
+			LogFileCount:      req.GetLogFileCount(),
 		},
 	)
 

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
 	}

go.mod

@@ -25,7 +25,7 @@ require (
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	golang.zx2c4.com/wireguard/windows v0.5.3
 	google.golang.org/grpc v1.64.1
-	google.golang.org/protobuf v1.36.5
+	google.golang.org/protobuf v1.36.6
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 )
 
@@ -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
@@ -257,6 +257,6 @@ replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-2
 
 replace github.com/cloudflare/circl => github.com/cunicu/circl v0.0.0-20230801113412-fec58fc7b5f6
 
-replace github.com/pion/ice/v3 => github.com/netbirdio/ice/v3 v3.0.0-20240315174635-e72a50fcb64e
+replace github.com/pion/ice/v3 => github.com/netbirdio/ice/v3 v3.0.0-20250718163601-725c8ac53a31
 
 replace github.com/libp2p/go-netroute => github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944

go.sum

@@ -501,10 +501,10 @@ github.com/neelance/astrewrite v0.0.0-20160511093645-99348263ae86/go.mod h1:kHJE
 github.com/neelance/sourcemap v0.0.0-20200213170602-2833bce08e4c/go.mod h1:Qr6/a/Q4r9LP1IltGz7tA7iOK1WonHEYhu1HRBA7ZiM=
 github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944 h1:TDtJKmM6Sf8uYFx/dMeqNOL90KUoRscdfpFZ3Im89uk=
 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/ice/v3 v3.0.0-20250718163601-725c8ac53a31 h1:lr/CnQ9NnlHr4yjDaCqy3V1FW+y9DDpzqxu1+YXzXtc=
+github.com/netbirdio/ice/v3 v3.0.0-20250718163601-725c8ac53a31/go.mod h1:YMLU7qbKfVjmEv7EoZPIVEI+kNYxWCdPK3VS0BU+U4Q=
+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=
@@ -1164,8 +1164,8 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
-google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
+google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
+google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

management/client/client_test.go

@@ -87,6 +87,12 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 		).
 		Return(&types.Settings{}, nil).
 		AnyTimes()
+	settingsMockManager.
+		EXPECT().
+		GetExtraSettings(gomock.Any(), gomock.Any()).
+		Return(&types.ExtraSettings{}, nil).
+		AnyTimes()
+
 	permissionsManagerMock := permissions.NewMockManager(ctrl)
 	permissionsManagerMock.
 		EXPECT().
@@ -106,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)
 	}

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

management/server/account/manager.go

@@ -112,6 +112,7 @@ type Manager interface {
 	GetAccountSettings(ctx context.Context, accountID string, userID string) (*types.Settings, error)
 	DeleteSetupKey(ctx context.Context, accountID, userID, keyID string) error
 	UpdateAccountPeers(ctx context.Context, accountID string)
+	BufferUpdateAccountPeers(ctx context.Context, accountID string)
 	BuildUserInfosForAccount(ctx context.Context, accountID, initiatorUserID string, accountUsers []*types.User) (map[string]*types.UserInfo, error)
 	SyncUserJWTGroups(ctx context.Context, userAuth nbcontext.UserAuth) error
 	GetStore() store.Store

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
 	}

management/server/dns_test.go

@@ -216,8 +216,10 @@ func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
 	t.Cleanup(ctrl.Finish)
 
 	settingsMockManager := settings.NewMockManager(ctrl)
+	// 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) {

management/server/ephemeral.go

@@ -15,6 +15,8 @@ import (
 
 const (
 	ephemeralLifeTime = 10 * time.Minute
+	// cleanupWindow is the time window to wait after nearest peer deadline to start the cleanup procedure.
+	cleanupWindow = 1 * time.Minute
 )
 
 var (
@@ -41,6 +43,9 @@ type EphemeralManager struct {
 	tailPeer  *ephemeralPeer
 	peersLock sync.Mutex
 	timer     *time.Timer
+
+	lifeTime      time.Duration
+	cleanupWindow time.Duration
 }
 
 // NewEphemeralManager instantiate new EphemeralManager
@@ -48,6 +53,9 @@ func NewEphemeralManager(store store.Store, accountManager nbAccount.Manager) *E
 	return &EphemeralManager{
 		store:          store,
 		accountManager: accountManager,
+
+		lifeTime:      ephemeralLifeTime,
+		cleanupWindow: cleanupWindow,
 	}
 }
 
@@ -60,7 +68,7 @@ func (e *EphemeralManager) LoadInitialPeers(ctx context.Context) {
 
 	e.loadEphemeralPeers(ctx)
 	if e.headPeer != nil {
-		e.timer = time.AfterFunc(ephemeralLifeTime, func() {
+		e.timer = time.AfterFunc(e.lifeTime, func() {
 			e.cleanup(ctx)
 		})
 	}
@@ -113,9 +121,13 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
 		return
 	}
 
-	e.addPeer(peer.AccountID, peer.ID, newDeadLine())
+	e.addPeer(peer.AccountID, peer.ID, e.newDeadLine())
 	if e.timer == nil {
-		e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() {
+		delay := e.headPeer.deadline.Sub(timeNow()) + e.cleanupWindow
+		if delay < 0 {
+			delay = 0
+		}
+		e.timer = time.AfterFunc(delay, func() {
 			e.cleanup(ctx)
 		})
 	}
@@ -128,7 +140,7 @@ func (e *EphemeralManager) loadEphemeralPeers(ctx context.Context) {
 		return
 	}
 
-	t := newDeadLine()
+	t := e.newDeadLine()
 	for _, p := range peers {
 		e.addPeer(p.AccountID, p.ID, t)
 	}
@@ -155,7 +167,11 @@ func (e *EphemeralManager) cleanup(ctx context.Context) {
 	}
 
 	if e.headPeer != nil {
-		e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() {
+		delay := e.headPeer.deadline.Sub(timeNow()) + e.cleanupWindow
+		if delay < 0 {
+			delay = 0
+		}
+		e.timer = time.AfterFunc(delay, func() {
 			e.cleanup(ctx)
 		})
 	} else {
@@ -164,13 +180,20 @@ func (e *EphemeralManager) cleanup(ctx context.Context) {
 
 	e.peersLock.Unlock()
 
+	bufferAccountCall := make(map[string]struct{})
+
 	for id, p := range deletePeers {
 		log.WithContext(ctx).Debugf("delete ephemeral peer: %s", id)
 		err := e.accountManager.DeletePeer(ctx, p.accountID, id, activity.SystemInitiator)
 		if err != nil {
 			log.WithContext(ctx).Errorf("failed to delete ephemeral peer: %s", err)
+		} else {
+			bufferAccountCall[p.accountID] = struct{}{}
 		}
 	}
+	for accountID := range bufferAccountCall {
+		e.accountManager.BufferUpdateAccountPeers(ctx, accountID)
+	}
 }
 
 func (e *EphemeralManager) addPeer(accountID string, peerID string, deadline time.Time) {
@@ -223,6 +246,6 @@ func (e *EphemeralManager) isPeerOnList(id string) bool {
 	return false
 }
 
-func newDeadLine() time.Time {
-	return timeNow().Add(ephemeralLifeTime)
+func (e *EphemeralManager) newDeadLine() time.Time {
+	return timeNow().Add(e.lifeTime)
 }

management/server/ephemeral_test.go

@@ -3,9 +3,12 @@ package server
 import (
 	"context"
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -27,28 +30,65 @@ func (s *MockStore) GetAllEphemeralPeers(_ context.Context, _ store.LockingStren
 	return peers, nil
 }
 
-type MocAccountManager struct {
+type MockAccountManager struct {
+	mu sync.Mutex
 	nbAccount.Manager
-	store *MockStore
+	store             *MockStore
+	deletePeerCalls   int
+	bufferUpdateCalls map[string]int
+	wg                *sync.WaitGroup
 }
 
-func (a MocAccountManager) DeletePeer(_ context.Context, accountID, peerID, userID string) error {
+func (a *MockAccountManager) DeletePeer(_ context.Context, accountID, peerID, userID string) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	a.deletePeerCalls++
+	if a.wg != nil {
+		a.wg.Done()
+	}
 	delete(a.store.account.Peers, peerID)
-	return nil //nolint:nil
+	return nil
 }
 
-func (a MocAccountManager) GetStore() store.Store {
+func (a *MockAccountManager) GetDeletePeerCalls() int {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return a.deletePeerCalls
+}
+
+func (a *MockAccountManager) BufferUpdateAccountPeers(ctx context.Context, accountID string) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	if a.bufferUpdateCalls == nil {
+		a.bufferUpdateCalls = make(map[string]int)
+	}
+	a.bufferUpdateCalls[accountID]++
+}
+
+func (a *MockAccountManager) GetBufferUpdateCalls(accountID string) int {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	if a.bufferUpdateCalls == nil {
+		return 0
+	}
+	return a.bufferUpdateCalls[accountID]
+}
+
+func (a *MockAccountManager) GetStore() store.Store {
 	return a.store
 }
 
 func TestNewManager(t *testing.T) {
+	t.Cleanup(func() {
+		timeNow = time.Now
+	})
 	startTime := time.Now()
 	timeNow = func() time.Time {
 		return startTime
 	}
 
 	store := &MockStore{}
-	am := MocAccountManager{
+	am := MockAccountManager{
 		store: store,
 	}
 
@@ -56,7 +96,7 @@ func TestNewManager(t *testing.T) {
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	mgr := NewEphemeralManager(store, am)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
 	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
@@ -67,13 +107,16 @@ func TestNewManager(t *testing.T) {
 }
 
 func TestNewManagerPeerConnected(t *testing.T) {
+	t.Cleanup(func() {
+		timeNow = time.Now
+	})
 	startTime := time.Now()
 	timeNow = func() time.Time {
 		return startTime
 	}
 
 	store := &MockStore{}
-	am := MocAccountManager{
+	am := MockAccountManager{
 		store: store,
 	}
 
@@ -81,7 +124,7 @@ func TestNewManagerPeerConnected(t *testing.T) {
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	mgr := NewEphemeralManager(store, am)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
 	mgr.OnPeerConnected(context.Background(), store.account.Peers["ephemeral_peer_0"])
 
@@ -95,13 +138,16 @@ func TestNewManagerPeerConnected(t *testing.T) {
 }
 
 func TestNewManagerPeerDisconnected(t *testing.T) {
+	t.Cleanup(func() {
+		timeNow = time.Now
+	})
 	startTime := time.Now()
 	timeNow = func() time.Time {
 		return startTime
 	}
 
 	store := &MockStore{}
-	am := MocAccountManager{
+	am := MockAccountManager{
 		store: store,
 	}
 
@@ -109,7 +155,7 @@ func TestNewManagerPeerDisconnected(t *testing.T) {
 	numberOfEphemeralPeers := 3
 	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
 
-	mgr := NewEphemeralManager(store, am)
+	mgr := NewEphemeralManager(store, &am)
 	mgr.loadEphemeralPeers(context.Background())
 	for _, v := range store.account.Peers {
 		mgr.OnPeerConnected(context.Background(), v)
@@ -126,6 +172,36 @@ func TestNewManagerPeerDisconnected(t *testing.T) {
 	}
 }
 
+func TestCleanupSchedulingBehaviorIsBatched(t *testing.T) {
+	const (
+		ephemeralPeers    = 10
+		testLifeTime      = 1 * time.Second
+		testCleanupWindow = 100 * time.Millisecond
+	)
+	mockStore := &MockStore{}
+	mockAM := &MockAccountManager{
+		store: mockStore,
+	}
+	mockAM.wg = &sync.WaitGroup{}
+	mockAM.wg.Add(ephemeralPeers)
+	mgr := NewEphemeralManager(mockStore, mockAM)
+	mgr.lifeTime = testLifeTime
+	mgr.cleanupWindow = testCleanupWindow
+
+	account := newAccountWithId(context.Background(), "account", "", "", false)
+	mockStore.account = account
+	for i := range ephemeralPeers {
+		p := &nbpeer.Peer{ID: fmt.Sprintf("peer-%d", i), AccountID: account.Id, Ephemeral: true}
+		mockStore.account.Peers[p.ID] = p
+		time.Sleep(testCleanupWindow / ephemeralPeers)
+		mgr.OnPeerDisconnected(context.Background(), p)
+	}
+	mockAM.wg.Wait()
+	assert.Len(t, mockStore.account.Peers, 0, "all ephemeral peers should be cleaned up after the lifetime")
+	assert.Equal(t, 1, mockAM.GetBufferUpdateCalls(account.Id), "buffer update should be called once")
+	assert.Equal(t, ephemeralPeers, mockAM.GetDeletePeerCalls(), "should have deleted all peers")
+}
+
 func seedPeers(store *MockStore, numberOfPeers int, numberOfEphemeralPeers int) {
 	store.account = newAccountWithId(context.Background(), "my account", "", "", false)
 

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")

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)

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
 }

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
 }

management/server/management_proto_test.go

@@ -440,11 +440,15 @@ func startManagementForTest(t *testing.T, testFile string, config *types.Config)
 		GetSettings(gomock.Any(), gomock.Any(), gomock.Any()).
 		AnyTimes().
 		Return(&types.Settings{}, nil)
-
+	settingsMockManager.
+		EXPECT().
+		GetExtraSettings(gomock.Any(), gomock.Any()).
+		Return(&types.ExtraSettings{}, nil).
+		AnyTimes()
 	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()
@@ -454,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
 	}

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)

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

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")
+}

management/server/mock_server/account_mock.go

@@ -120,10 +120,20 @@ type MockAccountManager struct {
 	GetAccountOnboardingFunc              func(ctx context.Context, accountID, userID string) (*types.AccountOnboarding, error)
 	UpdateAccountOnboardingFunc           func(ctx context.Context, accountID, userID string, onboarding *types.AccountOnboarding) (*types.AccountOnboarding, error)
 	GetOrCreateAccountByPrivateDomainFunc func(ctx context.Context, initiatorId, domain string) (*types.Account, bool, error)
+	UpdateAccountPeersFunc                func(ctx context.Context, accountID string)
+	BufferUpdateAccountPeersFunc          func(ctx context.Context, accountID string)
 }
 
 func (am *MockAccountManager) UpdateAccountPeers(ctx context.Context, accountID string) {
-	// do nothing
+	if am.UpdateAccountPeersFunc != nil {
+		am.UpdateAccountPeersFunc(ctx, accountID)
+	}
+}
+
+func (am *MockAccountManager) BufferUpdateAccountPeers(ctx context.Context, accountID string) {
+	if am.BufferUpdateAccountPeersFunc != nil {
+		am.BufferUpdateAccountPeersFunc(ctx, accountID)
+	}
 }
 
 func (am *MockAccountManager) DeleteSetupKey(ctx context.Context, accountID, userID, keyID string) error {

management/server/nameserver_test.go

@@ -778,8 +778,14 @@ func createNSManager(t *testing.T) (*DefaultAccountManager, error) {
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
+	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.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) {

management/server/peer.go

@@ -9,6 +9,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/rs/xid"
@@ -236,11 +237,23 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 
 		if peer.Name != update.Name {
 			var newLabel string
-			newLabel, err = getPeerIPDNSLabel(ctx, transaction, peer.IP, accountID, update.Name)
+
+			newLabel, err = nbdns.GetParsedDomainLabel(update.Name)
 			if err != nil {
-				return fmt.Errorf("failed to get free DNS label: %w", err)
+				newLabel = ""
+			} else {
+				_, err := transaction.GetPeerIdByLabel(ctx, store.LockingStrengthNone, accountID, update.Name)
+				if err == nil {
+					newLabel = ""
+				}
 			}
 
+			if newLabel == "" {
+				newLabel, err = getPeerIPDNSLabel(peer.IP, update.Name)
+				if err != nil {
+					return fmt.Errorf("failed to get free DNS label: %w", err)
+				}
+			}
 			peer.Name = update.Name
 			peer.DNSLabel = newLabel
 			peerLabelChanged = true
@@ -375,7 +388,7 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 		storeEvent()
 	}
 
-	if updateAccountPeers {
+	if updateAccountPeers && userID != activity.SystemInitiator {
 		am.BufferUpdateAccountPeers(ctx, accountID)
 	}
 
@@ -472,6 +485,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 	var groupsToAdd []string
 	var allowExtraDNSLabels bool
 	var accountID string
+	var isEphemeral bool
 	if addedByUser {
 		user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthNone, userID)
 		if err != nil {
@@ -501,7 +515,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		setupKeyName = sk.Name
 		allowExtraDNSLabels = sk.AllowExtraDNSLabels
 		accountID = sk.AccountID
-
+		isEphemeral = sk.Ephemeral
 		if !sk.AllowExtraDNSLabels && len(peer.ExtraDNSLabels) > 0 {
 			return nil, nil, nil, status.Errorf(status.PreconditionFailed, "couldn't add peer: setup key doesn't allow extra DNS labels")
 		}
@@ -573,11 +587,17 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		}
 
 		var freeLabel string
-		freeLabel, err = getPeerIPDNSLabel(ctx, am.Store, freeIP, accountID, peer.Meta.Hostname)
-		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to get free DNS label: %w", err)
+		if isEphemeral || attempt > 1 {
+			freeLabel, err = getPeerIPDNSLabel(freeIP, peer.Meta.Hostname)
+			if err != nil {
+				return nil, nil, nil, fmt.Errorf("failed to get free DNS label: %w", err)
+			}
+		} else {
+			freeLabel, err = nbdns.GetParsedDomainLabel(peer.Meta.Hostname)
+			if err != nil {
+				return nil, nil, nil, fmt.Errorf("failed to get free DNS label: %w", err)
+			}
 		}
-
 		newPeer.DNSLabel = freeLabel
 		newPeer.IP = freeIP
 
@@ -647,7 +667,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		if isUniqueConstraintError(err) {
 			unlock()
 			unlock = nil
-			log.WithContext(ctx).Debugf("Failed to add peer in attempt %d, retrying: %v", attempt, err)
+			log.WithContext(ctx).WithFields(log.Fields{"dns_label": freeLabel, "ip": freeIP}).Tracef("Failed to add peer in attempt %d, retrying: %v", attempt, err)
 			continue
 		}
 
@@ -681,7 +701,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 	return am.getValidatedPeerWithMap(ctx, false, accountID, newPeer)
 }
 
-func getPeerIPDNSLabel(ctx context.Context, tx store.Store, ip net.IP, accountID, peerHostName string) (string, error) {
+func getPeerIPDNSLabel(ip net.IP, peerHostName string) (string, error) {
 	ip = ip.To4()
 
 	dnsName, err := nbdns.GetParsedDomainLabel(peerHostName)
@@ -689,12 +709,6 @@ func getPeerIPDNSLabel(ctx context.Context, tx store.Store, ip net.IP, accountID
 		return "", fmt.Errorf("failed to parse peer host name %s: %w", peerHostName, err)
 	}
 
-	_, err = tx.GetPeerIdByLabel(ctx, store.LockingStrengthNone, accountID, dnsName)
-	if err != nil {
-		//nolint:nilerr
-		return dnsName, nil
-	}
-
 	return fmt.Sprintf("%s-%d-%d", dnsName, ip[2], ip[3]), nil
 }
 
@@ -1177,6 +1191,19 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 
 	globalStart := time.Now()
 
+	hasPeersConnected := false
+	for _, peer := range account.Peers {
+		if am.peersUpdateManager.HasChannel(peer.ID) {
+			hasPeersConnected = true
+			break
+		}
+
+	}
+
+	if !hasPeersConnected {
+		return
+	}
+
 	approvedPeersMap, err := am.integratedPeerValidator.GetValidatedPeers(account.Id, maps.Values(account.Groups), maps.Values(account.Peers), account.Settings.Extra)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to send out updates to peers, failed to get validate peers: %v", err)
@@ -1198,6 +1225,12 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 		return
 	}
 
+	extraSetting, err := am.settingsManager.GetExtraSettings(ctx, accountID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to get flow enabled status: %v", err)
+		return
+	}
+
 	for _, peer := range account.Peers {
 		if !am.peersUpdateManager.HasChannel(peer.ID) {
 			log.WithContext(ctx).Tracef("peer %s doesn't have a channel, skipping network map update", peer.ID)
@@ -1232,12 +1265,6 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 			}
 			am.metrics.UpdateChannelMetrics().CountMergeNetworkMapDuration(time.Since(start))
 
-			extraSetting, err := am.settingsManager.GetExtraSettings(ctx, accountID)
-			if err != nil {
-				log.WithContext(ctx).Errorf("failed to get flow enabled status: %v", err)
-				return
-			}
-
 			start = time.Now()
 			update := toSyncResponse(ctx, nil, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting)
 			am.metrics.UpdateChannelMetrics().CountToSyncResponseDuration(time.Since(start))
@@ -1254,18 +1281,39 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 	}
 }
 
-func (am *DefaultAccountManager) BufferUpdateAccountPeers(ctx context.Context, accountID string) {
-	mu, _ := am.accountUpdateLocks.LoadOrStore(accountID, &sync.Mutex{})
-	lock := mu.(*sync.Mutex)
+type bufferUpdate struct {
+	mu     sync.Mutex
+	next   *time.Timer
+	update atomic.Bool
+}
 
-	if !lock.TryLock() {
+func (am *DefaultAccountManager) BufferUpdateAccountPeers(ctx context.Context, accountID string) {
+	bufUpd, _ := am.accountUpdateLocks.LoadOrStore(accountID, &bufferUpdate{})
+	b := bufUpd.(*bufferUpdate)
+
+	if !b.mu.TryLock() {
+		b.update.Store(true)
 		return
 	}
 
+	if b.next != nil {
+		b.next.Stop()
+	}
+
 	go func() {
-		time.Sleep(time.Duration(am.updateAccountPeersBufferInterval.Load()))
-		lock.Unlock()
+		defer b.mu.Unlock()
 		am.UpdateAccountPeers(ctx, accountID)
+		if !b.update.Load() {
+			return
+		}
+		b.update.Store(false)
+		if b.next == nil {
+			b.next = time.AfterFunc(time.Duration(am.updateAccountPeersBufferInterval.Load()), func() {
+				am.UpdateAccountPeers(ctx, accountID)
+			})
+			return
+		}
+		b.next.Reset(time.Duration(am.updateAccountPeersBufferInterval.Load()))
 	}()
 }
 

management/server/peer_test.go

@@ -13,6 +13,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -25,6 +26,7 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/mock_server"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/status"
@@ -1271,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"
@@ -1344,9 +1346,14 @@ func Test_RegisterPeerBySetupKey(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
+	settingsMockManager.
+		EXPECT().
+		GetExtraSettings(gomock.Any(), gomock.Any()).
+		Return(&types.ExtraSettings{}, nil).
+		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"
@@ -1489,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"
@@ -1556,9 +1563,14 @@ func Test_LoginPeer(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
+	settingsMockManager.
+		EXPECT().
+		GetExtraSettings(gomock.Any(), gomock.Any()).
+		Return(&types.ExtraSettings{}, nil).
+		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"
@@ -1836,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)
@@ -1858,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)
@@ -2241,3 +2253,131 @@ func Test_AddPeer(t *testing.T) {
 	assert.Equal(t, totalPeers, maps.Values(account.SetupKeys)[0].UsedTimes)
 	assert.Equal(t, uint64(totalPeers), account.Network.Serial)
 }
+
+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/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) {

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

monotime/time.go

@@ -9,6 +9,8 @@ var (
 	baseWallNano int64
 )
 
+type Time int64
+
 func init() {
 	baseWallTime = time.Now()
 	baseWallNano = baseWallTime.UnixNano()
@@ -23,7 +25,11 @@ func init() {
 // and using time.Since() for elapsed calculation, this avoids repeated
 // time.Now() calls and leverages Go's internal monotonic clock for
 // efficient duration measurement.
-func Now() int64 {
+func Now() Time {
 	elapsed := time.Since(baseWallTime)
-	return baseWallNano + int64(elapsed)
+	return Time(baseWallNano + int64(elapsed))
+}
+
+func Since(t Time) time.Duration {
+	return time.Duration(Now() - t)
 }

relay/auth/validator.go

@@ -7,13 +7,6 @@ import (
 	authv2 "github.com/netbirdio/netbird/relay/auth/hmac/v2"
 )
 
-// Validator is an interface that defines the Validate method.
-type Validator interface {
-	Validate(any) error
-	// Deprecated: Use Validate instead.
-	ValidateHelloMsgType(any) error
-}
-
 type TimedHMACValidator struct {
 	authenticatorV2 *authv2.Validator
 	authenticator   *auth.TimedHMACValidator

relay/client/client.go

@@ -124,15 +124,14 @@ func (cc *connContainer) close() {
 // While the Connect is in progress, the OpenConn function will block until the connection is established with relay server.
 type Client struct {
 	log            *log.Entry
-	parentCtx      context.Context
 	connectionURL  string
 	authTokenStore *auth.TokenStore
-	hashedID       []byte
+	hashedID       messages.PeerID
 
 	bufPool *sync.Pool
 
 	relayConn        net.Conn
-	conns            map[string]*connContainer
+	conns            map[messages.PeerID]*connContainer
 	serviceIsRunning bool
 	mu               sync.Mutex // protect serviceIsRunning and conns
 	readLoopMutex    sync.Mutex
@@ -142,14 +141,17 @@ type Client struct {
 
 	onDisconnectListener func(string)
 	listenerMutex        sync.Mutex
+
+	stateSubscription *PeersStateSubscription
 }
 
 // NewClient creates a new client for the relay server. The client is not connected to the server until the Connect
-func NewClient(ctx context.Context, serverURL string, authTokenStore *auth.TokenStore, peerID string) *Client {
-	hashedID, hashedStringId := messages.HashID(peerID)
+func NewClient(serverURL string, authTokenStore *auth.TokenStore, peerID string) *Client {
+	hashedID := messages.HashID(peerID)
+	relayLog := log.WithFields(log.Fields{"relay": serverURL})
+
 	c := &Client{
-		log:            log.WithFields(log.Fields{"relay": serverURL}),
-		parentCtx:      ctx,
+		log:            relayLog,
 		connectionURL:  serverURL,
 		authTokenStore: authTokenStore,
 		hashedID:       hashedID,
@@ -159,14 +161,15 @@ func NewClient(ctx context.Context, serverURL string, authTokenStore *auth.Token
 				return &buf
 			},
 		},
-		conns: make(map[string]*connContainer),
+		conns: make(map[messages.PeerID]*connContainer),
 	}
-	c.log.Infof("create new relay connection: local peerID: %s, local peer hashedID: %s", peerID, hashedStringId)
+
+	c.log.Infof("create new relay connection: local peerID: %s, local peer hashedID: %s", peerID, hashedID)
 	return c
 }
 
 // Connect establishes a connection to the relay server. It blocks until the connection is established or an error occurs.
-func (c *Client) Connect() error {
+func (c *Client) Connect(ctx context.Context) error {
 	c.log.Infof("connecting to relay server")
 	c.readLoopMutex.Lock()
 	defer c.readLoopMutex.Unlock()
@@ -178,17 +181,27 @@ func (c *Client) Connect() error {
 		return nil
 	}
 
-	if err := c.connect(); err != nil {
+	instanceURL, err := c.connect(ctx)
+	if err != nil {
 		return err
 	}
+	c.muInstanceURL.Lock()
+	c.instanceURL = instanceURL
+	c.muInstanceURL.Unlock()
 
-	c.log = c.log.WithField("relay", c.instanceURL.String())
+	c.stateSubscription = NewPeersStateSubscription(c.log, c.relayConn, c.closeConnsByPeerID)
+
+	c.log = c.log.WithField("relay", instanceURL.String())
 	c.log.Infof("relay connection established")
 
 	c.serviceIsRunning = true
 
+	internallyStoppedFlag := newInternalStopFlag()
+	hc := healthcheck.NewReceiver(c.log)
+	go c.listenForStopEvents(ctx, hc, c.relayConn, internallyStoppedFlag)
+
 	c.wgReadLoop.Add(1)
-	go c.readLoop(c.relayConn)
+	go c.readLoop(hc, c.relayConn, internallyStoppedFlag)
 
 	return nil
 }
@@ -196,26 +209,50 @@ func (c *Client) Connect() error {
 // OpenConn create a new net.Conn for the destination peer ID. In case if the connection is in progress
 // to the relay server, the function will block until the connection is established or timed out. Otherwise,
 // it will return immediately.
+// It block until the server confirm the peer is online.
 // todo: what should happen if call with the same peerID with multiple times?
-func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
+func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, error) {
+	peerID := messages.HashID(dstPeerID)
 
+	c.mu.Lock()
 	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+	_, ok := c.conns[peerID]
+	if ok {
+		c.mu.Unlock()
+		return nil, ErrConnAlreadyExists
+	}
+	c.mu.Unlock()
+
+	if err := c.stateSubscription.WaitToBeOnlineAndSubscribe(ctx, peerID); err != nil {
+		c.log.Errorf("peer not available: %s, %s", peerID, err)
+		return nil, err
+	}
+
+	c.log.Infof("remote peer is available, prepare the relayed connection: %s", peerID)
+	msgChannel := make(chan Msg, 100)
+
+	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
 		return nil, fmt.Errorf("relay connection is not established")
 	}
 
-	hashedID, hashedStringID := messages.HashID(dstPeerID)
-	_, ok := c.conns[hashedStringID]
+	c.muInstanceURL.Lock()
+	instanceURL := c.instanceURL
+	c.muInstanceURL.Unlock()
+	conn := NewConn(c, peerID, msgChannel, instanceURL)
+
+	_, ok = c.conns[peerID]
 	if ok {
+		c.mu.Unlock()
+		_ = conn.Close()
 		return nil, ErrConnAlreadyExists
 	}
-
-	c.log.Infof("open connection to peer: %s", hashedStringID)
-	msgChannel := make(chan Msg, 100)
-	conn := NewConn(c, hashedID, hashedStringID, msgChannel, c.instanceURL)
-
-	c.conns[hashedStringID] = newConnContainer(c.log, conn, msgChannel)
+	c.conns[peerID] = newConnContainer(c.log, conn, msgChannel)
+	c.mu.Unlock()
 	return conn, nil
 }
 
@@ -254,76 +291,70 @@ func (c *Client) Close() error {
 	return c.close(true)
 }
 
-func (c *Client) connect() 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(); 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() 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(buf)
+	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(relayConn net.Conn) {
-	internallyStoppedFlag := newInternalStopFlag()
-	hc := healthcheck.NewReceiver(c.log)
-	go c.listenForStopEvents(hc, relayConn, internallyStoppedFlag)
-
+func (c *Client) readLoop(hc *healthcheck.Receiver, relayConn net.Conn, internallyStoppedFlag *internalStopFlag) {
 	var (
 		errExit error
 		n       int
@@ -366,10 +397,7 @@ func (c *Client) readLoop(relayConn net.Conn) {
 
 	hc.Stop()
 
-	c.muInstanceURL.Lock()
-	c.instanceURL = nil
-	c.muInstanceURL.Unlock()
-
+	c.stateSubscription.Cleanup()
 	c.wgReadLoop.Done()
 	_ = c.close(false)
 	c.notifyDisconnected()
@@ -382,6 +410,14 @@ func (c *Client) handleMsg(msgType messages.MsgType, buf []byte, bufPtr *[]byte,
 		c.bufPool.Put(bufPtr)
 	case messages.MsgTypeTransport:
 		return c.handleTransportMsg(buf, bufPtr, internallyStoppedFlag)
+	case messages.MsgTypePeersOnline:
+		c.handlePeersOnlineMsg(buf)
+		c.bufPool.Put(bufPtr)
+		return true
+	case messages.MsgTypePeersWentOffline:
+		c.handlePeersWentOfflineMsg(buf)
+		c.bufPool.Put(bufPtr)
+		return true
 	case messages.MsgTypeClose:
 		c.log.Debugf("relay connection close by server")
 		c.bufPool.Put(bufPtr)
@@ -413,18 +449,16 @@ func (c *Client) handleTransportMsg(buf []byte, bufPtr *[]byte, internallyStoppe
 		return true
 	}
 
-	stringID := messages.HashIDToString(peerID)
-
 	c.mu.Lock()
 	if !c.serviceIsRunning {
 		c.mu.Unlock()
 		c.bufPool.Put(bufPtr)
 		return false
 	}
-	container, ok := c.conns[stringID]
+	container, ok := c.conns[*peerID]
 	c.mu.Unlock()
 	if !ok {
-		c.log.Errorf("peer not found: %s", stringID)
+		c.log.Errorf("peer not found: %s", peerID.String())
 		c.bufPool.Put(bufPtr)
 		return true
 	}
@@ -437,9 +471,9 @@ func (c *Client) handleTransportMsg(buf []byte, bufPtr *[]byte, internallyStoppe
 	return true
 }
 
-func (c *Client) writeTo(connReference *Conn, id string, dstID []byte, payload []byte) (int, error) {
+func (c *Client) writeTo(connReference *Conn, dstID messages.PeerID, payload []byte) (int, error) {
 	c.mu.Lock()
-	conn, ok := c.conns[id]
+	conn, ok := c.conns[dstID]
 	c.mu.Unlock()
 	if !ok {
 		return 0, net.ErrClosed
@@ -464,7 +498,7 @@ func (c *Client) writeTo(connReference *Conn, id string, dstID []byte, payload [
 	return len(payload), err
 }
 
-func (c *Client) listenForStopEvents(hc *healthcheck.Receiver, conn net.Conn, internalStopFlag *internalStopFlag) {
+func (c *Client) listenForStopEvents(ctx context.Context, hc *healthcheck.Receiver, conn net.Conn, internalStopFlag *internalStopFlag) {
 	for {
 		select {
 		case _, ok := <-hc.OnTimeout:
@@ -478,7 +512,7 @@ func (c *Client) listenForStopEvents(hc *healthcheck.Receiver, conn net.Conn, in
 				c.log.Warnf("failed to close connection: %s", err)
 			}
 			return
-		case <-c.parentCtx.Done():
+		case <-ctx.Done():
 			err := c.close(true)
 			if err != nil {
 				c.log.Errorf("failed to teardown connection: %s", err)
@@ -492,10 +526,31 @@ func (c *Client) closeAllConns() {
 	for _, container := range c.conns {
 		container.close()
 	}
-	c.conns = make(map[string]*connContainer)
+	c.conns = make(map[messages.PeerID]*connContainer)
 }
 
-func (c *Client) closeConn(connReference *Conn, id string) error {
+func (c *Client) closeConnsByPeerID(peerIDs []messages.PeerID) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	for _, peerID := range peerIDs {
+		container, ok := c.conns[peerID]
+		if !ok {
+			c.log.Warnf("can not close connection, peer not found: %s", peerID)
+			continue
+		}
+
+		container.log.Infof("remote peer has been disconnected, free up connection: %s", peerID)
+		container.close()
+		delete(c.conns, peerID)
+	}
+
+	if err := c.stateSubscription.UnsubscribeStateChange(peerIDs); err != nil {
+		c.log.Errorf("failed to unsubscribe from peer state change: %s, %s", peerIDs, err)
+	}
+}
+
+func (c *Client) closeConn(connReference *Conn, id messages.PeerID) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -507,6 +562,11 @@ func (c *Client) closeConn(connReference *Conn, id string) error {
 	if container.conn != connReference {
 		return fmt.Errorf("conn reference mismatch")
 	}
+
+	if err := c.stateSubscription.UnsubscribeStateChange([]messages.PeerID{id}); err != nil {
+		container.log.Errorf("failed to unsubscribe from peer state change: %s", err)
+	}
+
 	c.log.Infof("free up connection to peer: %s", id)
 	delete(c.conns, id)
 	container.close()
@@ -525,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 {
@@ -559,8 +623,8 @@ func (c *Client) writeCloseMsg() {
 	}
 }
 
-func (c *Client) readWithTimeout(buf []byte) (int, error) {
-	ctx, cancel := context.WithTimeout(c.parentCtx, serverResponseTimeout)
+func (c *Client) readWithTimeout(ctx context.Context, buf []byte) (int, error) {
+	ctx, cancel := context.WithTimeout(ctx, serverResponseTimeout)
 	defer cancel()
 
 	readDone := make(chan struct{})
@@ -581,3 +645,21 @@ func (c *Client) readWithTimeout(buf []byte) (int, error) {
 		return n, err
 	}
 }
+
+func (c *Client) handlePeersOnlineMsg(buf []byte) {
+	peersID, err := messages.UnmarshalPeersOnlineMsg(buf)
+	if err != nil {
+		c.log.Errorf("failed to unmarshal peers online msg: %s", err)
+		return
+	}
+	c.stateSubscription.OnPeersOnline(peersID)
+}
+
+func (c *Client) handlePeersWentOfflineMsg(buf []byte) {
+	peersID, err := messages.UnMarshalPeersWentOffline(buf)
+	if err != nil {
+		c.log.Errorf("failed to unmarshal peers went offline msg: %s", err)
+		return
+	}
+	c.stateSubscription.OnPeersWentOffline(peersID)
+}

relay/client/client_test.go

@@ -18,14 +18,19 @@ import (
 )
 
 var (
-	av               = &allow.Auth{}
 	hmacTokenStore   = &hmac.TokenStore{}
 	serverListenAddr = "127.0.0.1:1234"
 	serverURL        = "rel://127.0.0.1:1234"
+	serverCfg        = server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: serverURL,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	}
 )
 
 func TestMain(m *testing.M) {
-	_ = util.InitLog("error", "console")
+	_ = util.InitLog("debug", "console")
 	code := m.Run()
 	os.Exit(code)
 }
@@ -33,7 +38,7 @@ func TestMain(m *testing.M) {
 func TestClient(t *testing.T) {
 	ctx := context.Background()
 
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -58,37 +63,37 @@ func TestClient(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 	t.Log("alice connecting to server")
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
 	defer clientAlice.Close()
 
 	t.Log("placeholder connecting to server")
-	clientPlaceHolder := NewClient(ctx, serverURL, hmacTokenStore, "clientPlaceHolder")
-	err = clientPlaceHolder.Connect()
+	clientPlaceHolder := NewClient(serverURL, hmacTokenStore, "clientPlaceHolder")
+	err = clientPlaceHolder.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
 	defer clientPlaceHolder.Close()
 
 	t.Log("Bob connecting to server")
-	clientBob := NewClient(ctx, serverURL, hmacTokenStore, "bob")
-	err = clientBob.Connect()
+	clientBob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = clientBob.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
 	defer clientBob.Close()
 
 	t.Log("Alice open connection to Bob")
-	connAliceToBob, err := clientAlice.OpenConn("bob")
+	connAliceToBob, err := clientAlice.OpenConn(ctx, "bob")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
 
 	t.Log("Bob open connection to Alice")
-	connBobToAlice, err := clientBob.OpenConn("alice")
+	connBobToAlice, err := clientBob.OpenConn(ctx, "alice")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -115,7 +120,7 @@ func TestClient(t *testing.T) {
 func TestRegistration(t *testing.T) {
 	ctx := context.Background()
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -132,8 +137,8 @@ func TestRegistration(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		_ = srv.Shutdown(ctx)
 		t.Fatalf("failed to connect to server: %s", err)
@@ -172,8 +177,8 @@ func TestRegistrationTimeout(t *testing.T) {
 		_ = fakeTCPListener.Close()
 	}(fakeTCPListener)
 
-	clientAlice := NewClient(ctx, "127.0.0.1:1234", hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice := NewClient("127.0.0.1:1234", hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
 	if err == nil {
 		t.Errorf("failed to connect to server: %s", err)
 	}
@@ -189,7 +194,7 @@ func TestEcho(t *testing.T) {
 	idAlice := "alice"
 	idBob := "bob"
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -213,8 +218,8 @@ func TestEcho(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, idAlice)
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
@@ -225,8 +230,8 @@ func TestEcho(t *testing.T) {
 		}
 	}()
 
-	clientBob := NewClient(ctx, serverURL, hmacTokenStore, idBob)
-	err = clientBob.Connect()
+	clientBob := NewClient(serverURL, hmacTokenStore, idBob)
+	err = clientBob.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
@@ -237,12 +242,12 @@ func TestEcho(t *testing.T) {
 		}
 	}()
 
-	connAliceToBob, err := clientAlice.OpenConn(idBob)
+	connAliceToBob, err := clientAlice.OpenConn(ctx, idBob)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
 
-	connBobToAlice, err := clientBob.OpenConn(idAlice)
+	connBobToAlice, err := clientBob.OpenConn(ctx, idAlice)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -278,7 +283,7 @@ func TestBindToUnavailabePeer(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -303,14 +308,14 @@ func TestBindToUnavailabePeer(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		t.Errorf("failed to connect to server: %s", err)
 	}
-	_, err = clientAlice.OpenConn("bob")
-	if err != nil {
-		t.Errorf("failed to bind channel: %s", err)
+	_, err = clientAlice.OpenConn(ctx, "bob")
+	if err == nil {
+		t.Errorf("expected error when binding to unavailable peer, got nil")
 	}
 
 	log.Infof("closing client")
@@ -324,7 +329,7 @@ func TestBindReconnect(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -349,24 +354,24 @@ func TestBindReconnect(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	clientBob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = clientBob.Connect(ctx)
 	if err != nil {
 		t.Errorf("failed to connect to server: %s", err)
 	}
 
-	_, err = clientAlice.OpenConn("bob")
+	_, err = clientAlice.OpenConn(ctx, "bob")
 	if err != nil {
-		t.Errorf("failed to bind channel: %s", err)
+		t.Fatalf("failed to bind channel: %s", err)
 	}
 
-	clientBob := NewClient(ctx, serverURL, hmacTokenStore, "bob")
-	err = clientBob.Connect()
-	if err != nil {
-		t.Errorf("failed to connect to server: %s", err)
-	}
-
-	chBob, err := clientBob.OpenConn("alice")
+	chBob, err := clientBob.OpenConn(ctx, "alice")
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -377,18 +382,28 @@ func TestBindReconnect(t *testing.T) {
 		t.Errorf("failed to close client: %s", err)
 	}
 
-	clientAlice = NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	clientAlice = NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		t.Errorf("failed to connect to server: %s", err)
 	}
 
-	chAlice, err := clientAlice.OpenConn("bob")
+	chAlice, err := clientAlice.OpenConn(ctx, "bob")
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
 
 	testString := "hello alice, I am bob"
+	_, err = chBob.Write([]byte(testString))
+	if err == nil {
+		t.Errorf("expected error when writing to channel, got nil")
+	}
+
+	chBob, err = clientBob.OpenConn(ctx, "alice")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
 	_, err = chBob.Write([]byte(testString))
 	if err != nil {
 		t.Errorf("failed to write to channel: %s", err)
@@ -415,7 +430,7 @@ func TestCloseConn(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -440,13 +455,19 @@ func TestCloseConn(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	bob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = bob.Connect(ctx)
 	if err != nil {
 		t.Errorf("failed to connect to server: %s", err)
 	}
 
-	conn, err := clientAlice.OpenConn("bob")
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn(ctx, "bob")
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -472,7 +493,7 @@ func TestCloseRelayConn(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -496,13 +517,19 @@ func TestCloseRelayConn(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, "alice")
-	err = clientAlice.Connect()
+	bob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = bob.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
 
-	conn, err := clientAlice.OpenConn("bob")
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn(ctx, "bob")
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -514,7 +541,7 @@ func TestCloseRelayConn(t *testing.T) {
 		t.Errorf("unexpected reading from closed connection")
 	}
 
-	_, err = clientAlice.OpenConn("bob")
+	_, err = clientAlice.OpenConn(ctx, "bob")
 	if err == nil {
 		t.Errorf("unexpected opening connection to closed server")
 	}
@@ -524,7 +551,7 @@ func TestCloseByServer(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv1, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv1, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -544,8 +571,8 @@ func TestCloseByServer(t *testing.T) {
 
 	idAlice := "alice"
 	log.Debugf("connect by alice")
-	relayClient := NewClient(ctx, serverURL, hmacTokenStore, idAlice)
-	err = relayClient.Connect()
+	relayClient := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = relayClient.Connect(ctx)
 	if err != nil {
 		log.Fatalf("failed to connect to server: %s", err)
 	}
@@ -567,7 +594,7 @@ func TestCloseByServer(t *testing.T) {
 		log.Fatalf("timeout waiting for client to disconnect")
 	}
 
-	_, err = relayClient.OpenConn("bob")
+	_, err = relayClient.OpenConn(ctx, "bob")
 	if err == nil {
 		t.Errorf("unexpected opening connection to closed server")
 	}
@@ -577,7 +604,7 @@ func TestCloseByClient(t *testing.T) {
 	ctx := context.Background()
 
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -596,8 +623,8 @@ func TestCloseByClient(t *testing.T) {
 
 	idAlice := "alice"
 	log.Debugf("connect by alice")
-	relayClient := NewClient(ctx, serverURL, hmacTokenStore, idAlice)
-	err = relayClient.Connect()
+	relayClient := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = relayClient.Connect(ctx)
 	if err != nil {
 		log.Fatalf("failed to connect to server: %s", err)
 	}
@@ -607,7 +634,7 @@ func TestCloseByClient(t *testing.T) {
 		t.Errorf("failed to close client: %s", err)
 	}
 
-	_, err = relayClient.OpenConn("bob")
+	_, err = relayClient.OpenConn(ctx, "bob")
 	if err == nil {
 		t.Errorf("unexpected opening connection to closed server")
 	}
@@ -623,7 +650,7 @@ func TestCloseNotDrainedChannel(t *testing.T) {
 	idAlice := "alice"
 	idBob := "bob"
 	srvCfg := server.ListenerConfig{Address: serverListenAddr}
-	srv, err := server.NewServer(otel.Meter(""), serverURL, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -647,8 +674,8 @@ func TestCloseNotDrainedChannel(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	clientAlice := NewClient(ctx, serverURL, hmacTokenStore, idAlice)
-	err = clientAlice.Connect()
+	clientAlice := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = clientAlice.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
@@ -659,8 +686,8 @@ func TestCloseNotDrainedChannel(t *testing.T) {
 		}
 	}()
 
-	clientBob := NewClient(ctx, serverURL, hmacTokenStore, idBob)
-	err = clientBob.Connect()
+	clientBob := NewClient(serverURL, hmacTokenStore, idBob)
+	err = clientBob.Connect(ctx)
 	if err != nil {
 		t.Fatalf("failed to connect to server: %s", err)
 	}
@@ -671,12 +698,12 @@ func TestCloseNotDrainedChannel(t *testing.T) {
 		}
 	}()
 
-	connAliceToBob, err := clientAlice.OpenConn(idBob)
+	connAliceToBob, err := clientAlice.OpenConn(ctx, idBob)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
 
-	connBobToAlice, err := clientBob.OpenConn(idAlice)
+	connBobToAlice, err := clientBob.OpenConn(ctx, idAlice)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}

relay/client/conn.go

@@ -3,13 +3,14 @@ package client
 import (
 	"net"
 	"time"
+
+	"github.com/netbirdio/netbird/relay/messages"
 )
 
 // Conn represent a connection to a relayed remote peer.
 type Conn struct {
 	client      *Client
-	dstID       []byte
-	dstStringID string
+	dstID       messages.PeerID
 	messageChan chan Msg
 	instanceURL *RelayAddr
 }
@@ -17,14 +18,12 @@ type Conn struct {
 // NewConn creates a new connection to a relayed remote peer.
 // client: the client instance, it used to send messages to the destination peer
 // dstID: the destination peer ID
-// dstStringID: the destination peer ID in string format
 // messageChan: the channel where the messages will be received
 // instanceURL: the relay instance URL, it used to get the proper server instance address for the remote peer
-func NewConn(client *Client, dstID []byte, dstStringID string, messageChan chan Msg, instanceURL *RelayAddr) *Conn {
+func NewConn(client *Client, dstID messages.PeerID, messageChan chan Msg, instanceURL *RelayAddr) *Conn {
 	c := &Conn{
 		client:      client,
 		dstID:       dstID,
-		dstStringID: dstStringID,
 		messageChan: messageChan,
 		instanceURL: instanceURL,
 	}
@@ -33,7 +32,7 @@ func NewConn(client *Client, dstID []byte, dstStringID string, messageChan chan
 }
 
 func (c *Conn) Write(p []byte) (n int, err error) {
-	return c.client.writeTo(c, c.dstStringID, c.dstID, p)
+	return c.client.writeTo(c, c.dstID, p)
 }
 
 func (c *Conn) Read(b []byte) (n int, err error) {
@@ -48,7 +47,7 @@ func (c *Conn) Read(b []byte) (n int, err error) {
 }
 
 func (c *Conn) Close() error {
-	return c.client.closeConn(c, c.dstStringID)
+	return c.client.closeConn(c, c.dstID)
 }
 
 func (c *Conn) LocalAddr() net.Addr {

relay/client/guard.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
@@ -10,8 +11,21 @@ import (
 
 var (
 	reconnectingTimeout = 60 * time.Second
+	mux                 sync.Mutex
 )
 
+func getReconnectingTimeout() time.Duration {
+	mux.Lock()
+	defer mux.Unlock()
+	return reconnectingTimeout
+}
+
+func setReconnectingTimeout(timeout time.Duration) {
+	mux.Lock()
+	defer mux.Unlock()
+	reconnectingTimeout = timeout
+}
+
 // Guard manage the reconnection tries to the Relay server in case of disconnection event.
 type Guard struct {
 	// OnNewRelayClient is a channel that is used to notify the relay manager about a new relay client instance.
@@ -80,7 +94,7 @@ func (g *Guard) tryToQuickReconnect(parentCtx context.Context, rc *Client) bool
 
 	log.Infof("try to reconnect to Relay server: %s", rc.connectionURL)
 
-	if err := rc.Connect(); err != nil {
+	if err := rc.Connect(parentCtx); err != nil {
 		log.Errorf("failed to reconnect to relay server: %s", err)
 		return false
 	}
@@ -128,7 +142,7 @@ func exponentTicker(ctx context.Context) *backoff.Ticker {
 	bo := backoff.WithContext(&backoff.ExponentialBackOff{
 		InitialInterval: 2 * time.Second,
 		Multiplier:      2,
-		MaxInterval:     reconnectingTimeout,
+		MaxInterval:     getReconnectingTimeout(),
 		Clock:           backoff.SystemClock,
 	}, ctx)
 

relay/client/manager.go

@@ -42,7 +42,7 @@ type OnServerCloseListener func()
 // ManagerService is the interface for the relay manager.
 type ManagerService interface {
 	Serve() error
-	OpenConn(serverAddress, peerKey string) (net.Conn, error)
+	OpenConn(ctx context.Context, serverAddress, peerKey string) (net.Conn, error)
 	AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error
 	RelayInstanceAddress() (string, error)
 	ServerURLs() []string
@@ -123,7 +123,7 @@ func (m *Manager) Serve() error {
 // OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
 // 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(serverAddress, peerKey string) (net.Conn, error) {
+func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string) (net.Conn, error) {
 	m.relayClientMu.Lock()
 	defer m.relayClientMu.Unlock()
 
@@ -141,10 +141,10 @@ func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 	)
 	if !foreign {
 		log.Debugf("open peer connection via permanent server: %s", peerKey)
-		netConn, err = m.relayClient.OpenConn(peerKey)
+		netConn, err = m.relayClient.OpenConn(ctx, peerKey)
 	} else {
 		log.Debugf("open peer connection via foreign server: %s", serverAddress)
-		netConn, err = m.openConnVia(serverAddress, peerKey)
+		netConn, err = m.openConnVia(ctx, serverAddress, peerKey)
 	}
 	if err != nil {
 		return nil, err
@@ -229,7 +229,7 @@ func (m *Manager) UpdateToken(token *relayAuth.Token) error {
 	return m.tokenStore.UpdateToken(token)
 }
 
-func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
+func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string) (net.Conn, error) {
 	// check if already has a connection to the desired relay server
 	m.relayClientsMutex.RLock()
 	rt, ok := m.relayClients[serverAddress]
@@ -240,7 +240,7 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 		if rt.err != nil {
 			return nil, rt.err
 		}
-		return rt.relayClient.OpenConn(peerKey)
+		return rt.relayClient.OpenConn(ctx, peerKey)
 	}
 	m.relayClientsMutex.RUnlock()
 
@@ -255,7 +255,7 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 		if rt.err != nil {
 			return nil, rt.err
 		}
-		return rt.relayClient.OpenConn(peerKey)
+		return rt.relayClient.OpenConn(ctx, peerKey)
 	}
 
 	// create a new relay client and store it in the relayClients map
@@ -264,8 +264,8 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 	m.relayClients[serverAddress] = rt
 	m.relayClientsMutex.Unlock()
 
-	relayClient := NewClient(m.ctx, serverAddress, m.tokenStore, m.peerID)
-	err := relayClient.Connect()
+	relayClient := NewClient(serverAddress, m.tokenStore, m.peerID)
+	err := relayClient.Connect(m.ctx)
 	if err != nil {
 		rt.err = err
 		rt.Unlock()
@@ -279,7 +279,7 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 	rt.relayClient = relayClient
 	rt.Unlock()
 
-	conn, err := relayClient.OpenConn(peerKey)
+	conn, err := relayClient.OpenConn(ctx, peerKey)
 	if err != nil {
 		return nil, err
 	}
@@ -350,6 +350,12 @@ func (m *Manager) startCleanupLoop() {
 	}
 }
 
+func (m *Manager) getClientLen() int {
+	m.relayClientsMutex.Lock()
+	defer m.relayClientsMutex.Unlock()
+	return len(m.relayClients)
+}
+
 func (m *Manager) cleanUpUnusedRelays() {
 	m.relayClientsMutex.Lock()
 	defer m.relayClientsMutex.Unlock()

relay/client/manager_test.go

@@ -8,6 +8,7 @@ import (
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel"
 
+	"github.com/netbirdio/netbird/relay/auth/allow"
 	"github.com/netbirdio/netbird/relay/server"
 )
 
@@ -22,16 +23,22 @@ func TestEmptyURL(t *testing.T) {
 func TestForeignConn(t *testing.T) {
 	ctx := context.Background()
 
-	srvCfg1 := server.ListenerConfig{
+	lstCfg1 := server.ListenerConfig{
 		Address: "localhost:1234",
 	}
-	srv1, err := server.NewServer(otel.Meter(""), srvCfg1.Address, false, av)
+
+	srv1, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: lstCfg1.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
 	errChan := make(chan error, 1)
 	go func() {
-		err := srv1.Listen(srvCfg1)
+		err := srv1.Listen(lstCfg1)
 		if err != nil {
 			errChan <- err
 		}
@@ -51,7 +58,12 @@ func TestForeignConn(t *testing.T) {
 	srvCfg2 := server.ListenerConfig{
 		Address: "localhost:2234",
 	}
-	srv2, err := server.NewServer(otel.Meter(""), srvCfg2.Address, false, av)
+	srv2, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: srvCfg2.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -74,32 +86,26 @@ func TestForeignConn(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	idAlice := "alice"
-	log.Debugf("connect by alice")
 	mCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	clientAlice := NewManager(mCtx, toURL(srvCfg1), idAlice)
-	err = clientAlice.Serve()
-	if err != nil {
+	clientAlice := NewManager(mCtx, toURL(lstCfg1), "alice")
+	if err := clientAlice.Serve(); err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
 
-	idBob := "bob"
-	log.Debugf("connect by bob")
-	clientBob := NewManager(mCtx, toURL(srvCfg2), idBob)
-	err = clientBob.Serve()
-	if err != nil {
+	clientBob := NewManager(mCtx, toURL(srvCfg2), "bob")
+	if err := clientBob.Serve(); err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
 	bobsSrvAddr, err := clientBob.RelayInstanceAddress()
 	if err != nil {
 		t.Fatalf("failed to get relay address: %s", err)
 	}
-	connAliceToBob, err := clientAlice.OpenConn(bobsSrvAddr, idBob)
+	connAliceToBob, err := clientAlice.OpenConn(ctx, bobsSrvAddr, "bob")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
-	connBobToAlice, err := clientBob.OpenConn(bobsSrvAddr, idAlice)
+	connBobToAlice, err := clientBob.OpenConn(ctx, bobsSrvAddr, "alice")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -137,7 +143,7 @@ func TestForeginConnClose(t *testing.T) {
 	srvCfg1 := server.ListenerConfig{
 		Address: "localhost:1234",
 	}
-	srv1, err := server.NewServer(otel.Meter(""), srvCfg1.Address, false, av)
+	srv1, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -163,7 +169,7 @@ func TestForeginConnClose(t *testing.T) {
 	srvCfg2 := server.ListenerConfig{
 		Address: "localhost:2234",
 	}
-	srv2, err := server.NewServer(otel.Meter(""), srvCfg2.Address, false, av)
+	srv2, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -186,16 +192,20 @@ func TestForeginConnClose(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	idAlice := "alice"
-	log.Debugf("connect by alice")
 	mCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	mgr := NewManager(mCtx, toURL(srvCfg1), idAlice)
+
+	mgrBob := NewManager(mCtx, toURL(srvCfg2), "bob")
+	if err := mgrBob.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	mgr := NewManager(mCtx, toURL(srvCfg1), "alice")
 	err = mgr.Serve()
 	if err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
-	conn, err := mgr.OpenConn(toURL(srvCfg2)[0], "anotherpeer")
+	conn, err := mgr.OpenConn(ctx, toURL(srvCfg2)[0], "bob")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -212,23 +222,21 @@ func TestForeginAutoClose(t *testing.T) {
 	srvCfg1 := server.ListenerConfig{
 		Address: "localhost:1234",
 	}
-	srv1, err := server.NewServer(otel.Meter(""), srvCfg1.Address, false, av)
+	srv1, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
 	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")
@@ -241,7 +249,7 @@ func TestForeginAutoClose(t *testing.T) {
 	srvCfg2 := server.ListenerConfig{
 		Address: "localhost:2234",
 	}
-	srv2, err := server.NewServer(otel.Meter(""), srvCfg2.Address, false, av)
+	srv2, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -277,22 +285,15 @@ func TestForeginAutoClose(t *testing.T) {
 	}
 
 	t.Log("open connection to another peer")
-	conn, err := mgr.OpenConn(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
 	t.Logf("waiting for relay cleanup: %s", timeout)
 	time.Sleep(timeout)
-	if len(mgr.relayClients) != 0 {
-		t.Errorf("expected 0, got %d", len(mgr.relayClients))
+	if mgr.getClientLen() != 0 {
+		t.Errorf("expected 0, got %d", mgr.getClientLen())
 	}
 
 	t.Logf("closing manager")
@@ -300,12 +301,12 @@ func TestForeginAutoClose(t *testing.T) {
 
 func TestAutoReconnect(t *testing.T) {
 	ctx := context.Background()
-	reconnectingTimeout = 2 * time.Second
+	setReconnectingTimeout(2 * time.Second)
 
 	srvCfg := server.ListenerConfig{
 		Address: "localhost:1234",
 	}
-	srv, err := server.NewServer(otel.Meter(""), srvCfg.Address, false, av)
+	srv, err := server.NewServer(serverCfg)
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
@@ -330,6 +331,13 @@ func TestAutoReconnect(t *testing.T) {
 
 	mCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
+
+	clientBob := NewManager(mCtx, toURL(srvCfg), "bob")
+	err = clientBob.Serve()
+	if err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
 	clientAlice := NewManager(mCtx, toURL(srvCfg), "alice")
 	err = clientAlice.Serve()
 	if err != nil {
@@ -339,7 +347,7 @@ func TestAutoReconnect(t *testing.T) {
 	if err != nil {
 		t.Errorf("failed to get relay address: %s", err)
 	}
-	conn, err := clientAlice.OpenConn(ra, "bob")
+	conn, err := clientAlice.OpenConn(ctx, ra, "bob")
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -354,10 +362,10 @@ func TestAutoReconnect(t *testing.T) {
 	}
 
 	log.Infof("waiting for reconnection")
-	time.Sleep(reconnectingTimeout + 1*time.Second)
+	time.Sleep(getReconnectingTimeout() + 1*time.Second)
 
 	log.Infof("reopent the connection")
-	_, err = clientAlice.OpenConn(ra, "bob")
+	_, err = clientAlice.OpenConn(ctx, ra, "bob")
 	if err != nil {
 		t.Errorf("failed to open channel: %s", err)
 	}
@@ -366,24 +374,27 @@ func TestAutoReconnect(t *testing.T) {
 func TestNotifierDoubleAdd(t *testing.T) {
 	ctx := context.Background()
 
-	srvCfg1 := server.ListenerConfig{
+	listenerCfg1 := server.ListenerConfig{
 		Address: "localhost:1234",
 	}
-	srv1, err := server.NewServer(otel.Meter(""), srvCfg1.Address, false, av)
+	srv, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: listenerCfg1.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
 	if err != nil {
 		t.Fatalf("failed to create server: %s", err)
 	}
 	errChan := make(chan error, 1)
 	go func() {
-		err := srv1.Listen(srvCfg1)
-		if err != nil {
+		if err := srv.Listen(listenerCfg1); err != nil {
 			errChan <- err
 		}
 	}()
 
 	defer func() {
-		err := srv1.Shutdown(ctx)
-		if err != nil {
+		if err := srv.Shutdown(ctx); err != nil {
 			t.Errorf("failed to close server: %s", err)
 		}
 	}()
@@ -392,17 +403,21 @@ func TestNotifierDoubleAdd(t *testing.T) {
 		t.Fatalf("failed to start server: %s", err)
 	}
 
-	idAlice := "alice"
 	log.Debugf("connect by alice")
 	mCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	clientAlice := NewManager(mCtx, toURL(srvCfg1), idAlice)
-	err = clientAlice.Serve()
-	if err != nil {
+
+	clientBob := NewManager(mCtx, toURL(listenerCfg1), "bob")
+	if err = clientBob.Serve(); err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
 
-	conn1, err := clientAlice.OpenConn(clientAlice.ServerURLs()[0], "idBob")
+	clientAlice := NewManager(mCtx, toURL(listenerCfg1), "alice")
+	if err = clientAlice.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	conn1, err := clientAlice.OpenConn(ctx, clientAlice.ServerURLs()[0], "bob")
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}

relay/client/peer_subscription.go

@@ -0,0 +1,191 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+const (
+	OpenConnectionTimeout = 30 * time.Second
+)
+
+type relayedConnWriter interface {
+	Write(p []byte) (n int, err error)
+}
+
+// PeersStateSubscription manages subscriptions to peer state changes (online/offline)
+// over a relay connection. It allows tracking peers' availability and handling offline
+// events via a callback. We get online notification from the server only once.
+type PeersStateSubscription struct {
+	log             *log.Entry
+	relayConn       relayedConnWriter
+	offlineCallback func(peerIDs []messages.PeerID)
+
+	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 {
+	return &PeersStateSubscription{
+		log:                   log,
+		relayConn:             relayConn,
+		offlineCallback:       offlineCallback,
+		listenForOfflinePeers: make(map[messages.PeerID]struct{}),
+		waitingPeers:          make(map[messages.PeerID]chan struct{}),
+	}
+}
+
+// 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
+		}
+
+		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)
+	}
+}
+
+// WaitToBeOnlineAndSubscribe waits for a specific peer to come online and subscribes to its state changes.
+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{}, 1)
+	s.waitingPeers[peerID] = waitCh
+	s.listenForOfflinePeers[peerID] = struct{}{}
+	s.mu.Unlock()
+
+	if err := s.subscribeStateChange(peerID); err != nil {
+		s.log.Errorf("failed to subscribe to 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 err
+	}
+
+	// Wait for peer to come online or context to be cancelled
+	timeoutCtx, cancel := context.WithTimeout(ctx, OpenConnectionTimeout)
+	defer cancel()
+	select {
+	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():
+		s.log.Debugf("context timed out while waiting for peer %s to come online", peerID)
+		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()
+	}
+}
+
+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)
+			delete(s.waitingPeers, peerID)
+		}
+
+		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)
+	}
+
+	s.waitingPeers = make(map[messages.PeerID]chan struct{})
+	s.listenForOfflinePeers = make(map[messages.PeerID]struct{})
+}
+
+func (s *PeersStateSubscription) subscribeStateChange(peerID messages.PeerID) error {
+	msgs, err := messages.MarshalSubPeerStateMsg([]messages.PeerID{peerID})
+	if err != nil {
+		return err
+	}
+
+	for _, msg := range msgs {
+		if _, err := s.relayConn.Write(msg); err != nil {
+			return err
+		}
+
+	}
+	return nil
+}
+
+func (s *PeersStateSubscription) unsubscribeStateChange(peerIDs []messages.PeerID) error {
+	msgs, err := messages.MarshalUnsubPeerStateMsg(peerIDs)
+	if err != nil {
+		return err
+	}
+
+	var connWriteErr error
+	for _, msg := range msgs {
+		if _, err := s.relayConn.Write(msg); err != nil {
+			connWriteErr = err
+		}
+	}
+	return connWriteErr
+}

relay/client/peer_subscription_test.go

@@ -0,0 +1,99 @@
+package client
+
+import (
+	"bytes"
+	"context"
+	"testing"
+	"time"
+
+	"github.com/netbirdio/netbird/relay/messages"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type mockRelayedConn struct {
+}
+
+func (m *mockRelayedConn) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func TestWaitToBeOnlineAndSubscribe_Success(t *testing.T) {
+	peerID := messages.HashID("peer1")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{}) // discard log output
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Launch wait in background
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		sub.OnPeersOnline([]messages.PeerID{peerID})
+	}()
+
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	assert.NoError(t, err)
+}
+
+func TestWaitToBeOnlineAndSubscribe_Timeout(t *testing.T) {
+	peerID := messages.HashID("peer2")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	assert.Error(t, err)
+	assert.Equal(t, context.DeadlineExceeded, err)
+}
+
+func TestWaitToBeOnlineAndSubscribe_Duplicate(t *testing.T) {
+	peerID := messages.HashID("peer3")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx := context.Background()
+	go func() {
+		_ = sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+
+	}()
+	time.Sleep(100 * time.Millisecond)
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "already waiting")
+}
+
+func TestUnsubscribeStateChange(t *testing.T) {
+	peerID := messages.HashID("peer4")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	doneChan := make(chan struct{})
+	go func() {
+		_ = sub.WaitToBeOnlineAndSubscribe(context.Background(), peerID)
+		close(doneChan)
+	}()
+	time.Sleep(100 * time.Millisecond)
+
+	err := sub.UnsubscribeStateChange([]messages.PeerID{peerID})
+	assert.NoError(t, err)
+
+	select {
+	case <-doneChan:
+	case <-time.After(200 * time.Millisecond):
+		// Expected timeout, meaning the subscription was successfully unsubscribed
+		t.Errorf("timeout")
+	}
+}

relay/client/picker.go

@@ -70,8 +70,8 @@ func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
 
 func (sp *ServerPicker) startConnection(ctx context.Context, resultChan chan connResult, url string) {
 	log.Infof("try to connecting to relay server: %s", url)
-	relayClient := NewClient(ctx, url, sp.TokenStore, sp.PeerID)
-	err := relayClient.Connect()
+	relayClient := NewClient(url, sp.TokenStore, sp.PeerID)
+	err := relayClient.Connect(ctx)
 	resultChan <- connResult{
 		RelayClient: relayClient,
 		Url:         url,

relay/cmd/root.go

@@ -141,7 +141,14 @@ func execute(cmd *cobra.Command, args []string) error {
 	hashedSecret := sha256.Sum256([]byte(cobraConfig.AuthSecret))
 	authenticator := auth.NewTimedHMACValidator(hashedSecret[:], 24*time.Hour)
 
-	srv, err := server.NewServer(metricsServer.Meter, cobraConfig.ExposedAddress, tlsSupport, authenticator)
+	cfg := server.Config{
+		Meter:          metricsServer.Meter,
+		ExposedAddress: cobraConfig.ExposedAddress,
+		AuthValidator:  authenticator,
+		TLSSupport:     tlsSupport,
+	}
+
+	srv, err := server.NewServer(cfg)
 	if err != nil {
 		log.Debugf("failed to create relay server: %v", err)
 		return fmt.Errorf("failed to create relay server: %v", err)

relay/healthcheck/receiver.go

@@ -2,14 +2,26 @@ package healthcheck
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 )
 
-var (
-	heartbeatTimeout = healthCheckInterval + 10*time.Second
-)
+var heartbeatTimeout = getHealthCheckInterval() + 10*time.Second
+var mux sync.Mutex
+
+func getHeartBeatTimeout() time.Duration {
+	mux.Lock()
+	defer mux.Unlock()
+	return heartbeatTimeout
+}
+
+func setHeartBeatTimeout(interval time.Duration) {
+	mux.Lock()
+	defer mux.Unlock()
+	heartbeatTimeout = interval
+}
 
 // Receiver is a healthcheck receiver
 // It will listen for heartbeat and check if the heartbeat is not received in a certain time
@@ -56,7 +68,7 @@ func (r *Receiver) Stop() {
 }
 
 func (r *Receiver) waitForHealthcheck() {
-	ticker := time.NewTicker(heartbeatTimeout)
+	ticker := time.NewTicker(getHeartBeatTimeout())
 	defer ticker.Stop()
 	defer r.ctxCancel()
 	defer close(r.OnTimeout)

relay/healthcheck/receiver_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 func TestNewReceiver(t *testing.T) {
-	heartbeatTimeout = 5 * time.Second
+	setHeartBeatTimeout(5 * time.Second)
 	r := NewReceiver(log.WithContext(context.Background()))
 
 	select {
@@ -23,7 +23,7 @@ func TestNewReceiver(t *testing.T) {
 }
 
 func TestNewReceiverNotReceive(t *testing.T) {
-	heartbeatTimeout = 1 * time.Second
+	setHeartBeatTimeout(1 * time.Second)
 	r := NewReceiver(log.WithContext(context.Background()))
 
 	select {
@@ -34,7 +34,7 @@ func TestNewReceiverNotReceive(t *testing.T) {
 }
 
 func TestNewReceiverAck(t *testing.T) {
-	heartbeatTimeout = 2 * time.Second
+	setHeartBeatTimeout(2 * time.Second)
 	r := NewReceiver(log.WithContext(context.Background()))
 
 	r.Heartbeat()
@@ -59,13 +59,13 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 
 	for _, tc := range testsCases {
 		t.Run(tc.name, func(t *testing.T) {
-			originalInterval := healthCheckInterval
-			originalTimeout := heartbeatTimeout
-			healthCheckInterval = 1 * time.Second
-			heartbeatTimeout = healthCheckInterval + 500*time.Millisecond
+			originalInterval := getHealthCheckInterval()
+			originalTimeout := getHeartBeatTimeout()
+			setHealthCheckInterval(1 * time.Second)
+			setHeartBeatTimeout(getHealthCheckInterval() + 500*time.Millisecond)
 			defer func() {
-				healthCheckInterval = originalInterval
-				heartbeatTimeout = originalTimeout
+				setHealthCheckInterval(originalInterval)
+				setHeartBeatTimeout(originalTimeout)
 			}()
 			//nolint:tenv
 			os.Setenv(defaultAttemptThresholdEnv, fmt.Sprintf("%d", tc.threshold))
@@ -73,7 +73,7 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 
 			receiver := NewReceiver(log.WithField("test_name", tc.name))
 
-			testTimeout := heartbeatTimeout*time.Duration(tc.threshold) + healthCheckInterval
+			testTimeout := getHeartBeatTimeout()*time.Duration(tc.threshold) + getHealthCheckInterval()
 
 			if tc.resetCounterOnce {
 				receiver.Heartbeat()

relay/healthcheck/sender.go

@@ -19,6 +19,30 @@ var (
 	healthCheckTimeout  = 20 * time.Second
 )
 
+func getHealthCheckInterval() time.Duration {
+	mux.Lock()
+	defer mux.Unlock()
+	return healthCheckInterval
+}
+
+func setHealthCheckInterval(interval time.Duration) {
+	mux.Lock()
+	defer mux.Unlock()
+	healthCheckInterval = interval
+}
+
+func getHealthCheckTimeout() time.Duration {
+	mux.Lock()
+	defer mux.Unlock()
+	return healthCheckTimeout
+}
+
+func setHealthCheckTimeout(timeout time.Duration) {
+	mux.Lock()
+	defer mux.Unlock()
+	healthCheckTimeout = timeout
+}
+
 // Sender is a healthcheck sender
 // It will send healthcheck signal to the receiver
 // If the receiver does not receive the signal in a certain time, it will send a timeout signal and stop to work
@@ -57,7 +81,7 @@ func (hc *Sender) OnHCResponse() {
 }
 
 func (hc *Sender) StartHealthCheck(ctx context.Context) {
-	ticker := time.NewTicker(healthCheckInterval)
+	ticker := time.NewTicker(getHealthCheckInterval())
 	defer ticker.Stop()
 
 	timeoutTicker := time.NewTicker(hc.getTimeoutTime())
@@ -94,7 +118,7 @@ func (hc *Sender) StartHealthCheck(ctx context.Context) {
 }
 
 func (hc *Sender) getTimeoutTime() time.Duration {
-	return healthCheckInterval + healthCheckTimeout
+	return getHealthCheckInterval() + getHealthCheckTimeout()
 }
 
 func getAttemptThresholdFromEnv() int {

relay/healthcheck/sender_test.go

@@ -12,8 +12,8 @@ import (
 
 func TestMain(m *testing.M) {
 	// override the health check interval to speed up the test
-	healthCheckInterval = 2 * time.Second
-	healthCheckTimeout = 100 * time.Millisecond
+	setHealthCheckInterval(2 * time.Second)
+	setHealthCheckTimeout(100 * time.Millisecond)
 	code := m.Run()
 	os.Exit(code)
 }
@@ -32,7 +32,7 @@ func TestNewHealthPeriod(t *testing.T) {
 			hc.OnHCResponse()
 		case <-hc.Timeout:
 			t.Fatalf("health check is timed out")
-		case <-time.After(healthCheckInterval + 100*time.Millisecond):
+		case <-time.After(getHealthCheckInterval() + 100*time.Millisecond):
 			t.Fatalf("health check not received")
 		}
 	}
@@ -46,7 +46,7 @@ func TestNewHealthFailed(t *testing.T) {
 
 	select {
 	case <-hc.Timeout:
-	case <-time.After(healthCheckInterval + healthCheckTimeout + 100*time.Millisecond):
+	case <-time.After(getHealthCheckInterval() + getHealthCheckTimeout() + 100*time.Millisecond):
 		t.Fatalf("health check is not timed out")
 	}
 }
@@ -89,7 +89,7 @@ func TestTimeoutReset(t *testing.T) {
 			hc.OnHCResponse()
 		case <-hc.Timeout:
 			t.Fatalf("health check is timed out")
-		case <-time.After(healthCheckInterval + 100*time.Millisecond):
+		case <-time.After(getHealthCheckInterval() + 100*time.Millisecond):
 			t.Fatalf("health check not received")
 		}
 	}
@@ -118,13 +118,13 @@ func TestSenderHealthCheckAttemptThreshold(t *testing.T) {
 
 	for _, tc := range testsCases {
 		t.Run(tc.name, func(t *testing.T) {
-			originalInterval := healthCheckInterval
-			originalTimeout := healthCheckTimeout
-			healthCheckInterval = 1 * time.Second
-			healthCheckTimeout = 500 * time.Millisecond
+			originalInterval := getHealthCheckInterval()
+			originalTimeout := getHealthCheckTimeout()
+			setHealthCheckInterval(1 * time.Second)
+			setHealthCheckTimeout(500 * time.Millisecond)
 			defer func() {
-				healthCheckInterval = originalInterval
-				healthCheckTimeout = originalTimeout
+				setHealthCheckInterval(originalInterval)
+				setHealthCheckTimeout(originalTimeout)
 			}()
 
 			//nolint:tenv
@@ -155,7 +155,7 @@ func TestSenderHealthCheckAttemptThreshold(t *testing.T) {
 				}
 			}()
 
-			testTimeout := sender.getTimeoutTime()*time.Duration(tc.threshold) + healthCheckInterval
+			testTimeout := sender.getTimeoutTime()*time.Duration(tc.threshold) + getHealthCheckInterval()
 
 			select {
 			case <-sender.Timeout:

relay/messages/id.go

@@ -8,24 +8,24 @@ import (
 
 const (
 	prefixLength = 4
-	IDSize       = prefixLength + sha256.Size
+	peerIDSize   = prefixLength + sha256.Size
 )
 
 var (
 	prefix = []byte("sha-") // 4 bytes
 )
 
-// HashID generates a sha256 hash from the peerID and returns the hash and the human-readable string
-func HashID(peerID string) ([]byte, string) {
-	idHash := sha256.Sum256([]byte(peerID))
-	idHashString := string(prefix) + base64.StdEncoding.EncodeToString(idHash[:])
-	var prefixedHash []byte
-	prefixedHash = append(prefixedHash, prefix...)
-	prefixedHash = append(prefixedHash, idHash[:]...)
-	return prefixedHash, idHashString
+type PeerID [peerIDSize]byte
+
+func (p PeerID) String() string {
+	return fmt.Sprintf("%s%s", p[:prefixLength], base64.StdEncoding.EncodeToString(p[prefixLength:]))
 }
 
-// HashIDToString converts a hash to a human-readable string
-func HashIDToString(idHash []byte) string {
-	return fmt.Sprintf("%s%s", idHash[:prefixLength], base64.StdEncoding.EncodeToString(idHash[prefixLength:]))
+// HashID generates a sha256 hash from the peerID and returns the hash and the human-readable string
+func HashID(peerID string) PeerID {
+	idHash := sha256.Sum256([]byte(peerID))
+	var prefixedHash [peerIDSize]byte
+	copy(prefixedHash[:prefixLength], prefix)
+	copy(prefixedHash[prefixLength:], idHash[:])
+	return prefixedHash
 }

relay/messages/id_test.go

@@ -1,13 +0,0 @@
-package messages
-
-import (
-	"testing"
-)
-
-func TestHashID(t *testing.T) {
-	hashedID, hashedStringId := HashID("alice")
-	enc := HashIDToString(hashedID)
-	if enc != hashedStringId {
-		t.Errorf("expected %s, got %s", hashedStringId, enc)
-	}
-}

relay/messages/message.go

@@ -9,19 +9,26 @@ import (
 const (
 	MaxHandshakeSize     = 212
 	MaxHandshakeRespSize = 8192
+	MaxMessageSize       = 8820
 
 	CurrentProtocolVersion = 1
 
 	MsgTypeUnknown MsgType = 0
 	// Deprecated: Use MsgTypeAuth instead.
-	MsgTypeHello MsgType = 1
+	MsgTypeHello = 1
 	// Deprecated: Use MsgTypeAuthResponse instead.
-	MsgTypeHelloResponse MsgType = 2
-	MsgTypeTransport     MsgType = 3
-	MsgTypeClose         MsgType = 4
-	MsgTypeHealthCheck   MsgType = 5
-	MsgTypeAuth                  = 6
-	MsgTypeAuthResponse          = 7
+	MsgTypeHelloResponse = 2
+	MsgTypeTransport     = 3
+	MsgTypeClose         = 4
+	MsgTypeHealthCheck   = 5
+	MsgTypeAuth          = 6
+	MsgTypeAuthResponse  = 7
+
+	// Peers state messages
+	MsgTypeSubscribePeerState   = 8
+	MsgTypeUnsubscribePeerState = 9
+	MsgTypePeersOnline          = 10
+	MsgTypePeersWentOffline     = 11
 
 	// base size of the message
 	sizeOfVersionByte = 1
@@ -30,17 +37,17 @@ const (
 
 	// auth message
 	sizeOfMagicByte     = 4
-	headerSizeAuth      = sizeOfMagicByte + IDSize
+	headerSizeAuth      = sizeOfMagicByte + peerIDSize
 	offsetMagicByte     = sizeOfProtoHeader
 	offsetAuthPeerID    = sizeOfProtoHeader + sizeOfMagicByte
 	headerTotalSizeAuth = sizeOfProtoHeader + headerSizeAuth
 
 	// hello message
-	headerSizeHello     = sizeOfMagicByte + IDSize
+	headerSizeHello     = sizeOfMagicByte + peerIDSize
 	headerSizeHelloResp = 0
 
 	// transport
-	headerSizeTransport      = IDSize
+	headerSizeTransport      = peerIDSize
 	offsetTransportID        = sizeOfProtoHeader
 	headerTotalSizeTransport = sizeOfProtoHeader + headerSizeTransport
 )
@@ -72,6 +79,14 @@ func (m MsgType) String() string {
 		return "close"
 	case MsgTypeHealthCheck:
 		return "health check"
+	case MsgTypeSubscribePeerState:
+		return "subscribe peer state"
+	case MsgTypeUnsubscribePeerState:
+		return "unsubscribe peer state"
+	case MsgTypePeersOnline:
+		return "peers online"
+	case MsgTypePeersWentOffline:
+		return "peers went offline"
 	default:
 		return "unknown"
 	}
@@ -102,7 +117,9 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) {
 		MsgTypeAuth,
 		MsgTypeTransport,
 		MsgTypeClose,
-		MsgTypeHealthCheck:
+		MsgTypeHealthCheck,
+		MsgTypeSubscribePeerState,
+		MsgTypeUnsubscribePeerState:
 		return msgType, nil
 	default:
 		return MsgTypeUnknown, fmt.Errorf("invalid msg type %d", msgType)
@@ -122,7 +139,9 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 		MsgTypeAuthResponse,
 		MsgTypeTransport,
 		MsgTypeClose,
-		MsgTypeHealthCheck:
+		MsgTypeHealthCheck,
+		MsgTypePeersOnline,
+		MsgTypePeersWentOffline:
 		return msgType, nil
 	default:
 		return MsgTypeUnknown, fmt.Errorf("invalid msg type %d", msgType)
@@ -135,11 +154,7 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 // message is used to authenticate the client with the server. The authentication is done using an HMAC method.
 // The protocol does not limit to use HMAC, it can be any other method. If the authentication failed the server will
 // close the network connection without any response.
-func MarshalHelloMsg(peerID []byte, additions []byte) ([]byte, error) {
-	if len(peerID) != IDSize {
-		return nil, fmt.Errorf("invalid peerID length: %d", len(peerID))
-	}
-
+func MarshalHelloMsg(peerID PeerID, additions []byte) ([]byte, error) {
 	msg := make([]byte, sizeOfProtoHeader+sizeOfMagicByte, sizeOfProtoHeader+headerSizeHello+len(additions))
 
 	msg[0] = byte(CurrentProtocolVersion)
@@ -147,7 +162,7 @@ func MarshalHelloMsg(peerID []byte, additions []byte) ([]byte, error) {
 
 	copy(msg[sizeOfProtoHeader:sizeOfProtoHeader+sizeOfMagicByte], magicHeader)
 
-	msg = append(msg, peerID...)
+	msg = append(msg, peerID[:]...)
 	msg = append(msg, additions...)
 
 	return msg, nil
@@ -156,7 +171,7 @@ func MarshalHelloMsg(peerID []byte, additions []byte) ([]byte, error) {
 // Deprecated: Use UnmarshalAuthMsg instead.
 // UnmarshalHelloMsg extracts peerID and the additional data from the hello message. The Additional data is used to
 // authenticate the client with the server.
-func UnmarshalHelloMsg(msg []byte) ([]byte, []byte, error) {
+func UnmarshalHelloMsg(msg []byte) (*PeerID, []byte, error) {
 	if len(msg) < sizeOfProtoHeader+headerSizeHello {
 		return nil, nil, ErrInvalidMessageLength
 	}
@@ -164,7 +179,9 @@ func UnmarshalHelloMsg(msg []byte) ([]byte, []byte, error) {
 		return nil, nil, errors.New("invalid magic header")
 	}
 
-	return msg[sizeOfProtoHeader+sizeOfMagicByte : sizeOfProtoHeader+headerSizeHello], msg[headerSizeHello:], nil
+	peerID := PeerID(msg[sizeOfProtoHeader+sizeOfMagicByte : sizeOfProtoHeader+headerSizeHello])
+
+	return &peerID, msg[headerSizeHello:], nil
 }
 
 // Deprecated: Use MarshalAuthResponse instead.
@@ -197,34 +214,33 @@ func UnmarshalHelloResponse(msg []byte) ([]byte, error) {
 // message is used to authenticate the client with the server. The authentication is done using an HMAC method.
 // The protocol does not limit to use HMAC, it can be any other method. If the authentication failed the server will
 // close the network connection without any response.
-func MarshalAuthMsg(peerID []byte, authPayload []byte) ([]byte, error) {
-	if len(peerID) != IDSize {
-		return nil, fmt.Errorf("invalid peerID length: %d", len(peerID))
+func MarshalAuthMsg(peerID PeerID, authPayload []byte) ([]byte, error) {
+	if headerTotalSizeAuth+len(authPayload) > MaxHandshakeSize {
+		return nil, fmt.Errorf("too large auth payload")
 	}
 
-	msg := make([]byte, sizeOfProtoHeader+sizeOfMagicByte, headerTotalSizeAuth+len(authPayload))
-
+	msg := make([]byte, headerTotalSizeAuth+len(authPayload))
 	msg[0] = byte(CurrentProtocolVersion)
 	msg[1] = byte(MsgTypeAuth)
-
 	copy(msg[sizeOfProtoHeader:], magicHeader)
-
-	msg = append(msg, peerID...)
-	msg = append(msg, authPayload...)
-
+	copy(msg[offsetAuthPeerID:], peerID[:])
+	copy(msg[headerTotalSizeAuth:], authPayload)
 	return msg, nil
 }
 
 // UnmarshalAuthMsg extracts peerID and the auth payload from the message
-func UnmarshalAuthMsg(msg []byte) ([]byte, []byte, error) {
+func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) {
 	if len(msg) < headerTotalSizeAuth {
 		return nil, nil, ErrInvalidMessageLength
 	}
+
+	// Validate the magic header
 	if !bytes.Equal(msg[offsetMagicByte:offsetMagicByte+sizeOfMagicByte], magicHeader) {
 		return nil, nil, errors.New("invalid magic header")
 	}
 
-	return msg[offsetAuthPeerID:headerTotalSizeAuth], msg[headerTotalSizeAuth:], nil
+	peerID := PeerID(msg[offsetAuthPeerID:headerTotalSizeAuth])
+	return &peerID, msg[headerTotalSizeAuth:], nil
 }
 
 // MarshalAuthResponse creates a response message to the auth.
@@ -268,45 +284,48 @@ func MarshalCloseMsg() []byte {
 // MarshalTransportMsg creates a transport message.
 // The transport message is used to exchange data between peers. The message contains the data to be exchanged and the
 // destination peer hashed ID.
-func MarshalTransportMsg(peerID, payload []byte) ([]byte, error) {
-	if len(peerID) != IDSize {
-		return nil, fmt.Errorf("invalid peerID length: %d", len(peerID))
-	}
-
-	msg := make([]byte, headerTotalSizeTransport, headerTotalSizeTransport+len(payload))
+func MarshalTransportMsg(peerID PeerID, payload []byte) ([]byte, error) {
+	// todo validate size
+	msg := make([]byte, headerTotalSizeTransport+len(payload))
 	msg[0] = byte(CurrentProtocolVersion)
 	msg[1] = byte(MsgTypeTransport)
-	copy(msg[sizeOfProtoHeader:], peerID)
-	msg = append(msg, payload...)
-
+	copy(msg[sizeOfProtoHeader:], peerID[:])
+	copy(msg[sizeOfProtoHeader+peerIDSize:], payload)
 	return msg, nil
 }
 
 // UnmarshalTransportMsg extracts the peerID and the payload from the transport message.
-func UnmarshalTransportMsg(buf []byte) ([]byte, []byte, error) {
+func UnmarshalTransportMsg(buf []byte) (*PeerID, []byte, error) {
 	if len(buf) < headerTotalSizeTransport {
 		return nil, nil, ErrInvalidMessageLength
 	}
 
-	return buf[offsetTransportID:headerTotalSizeTransport], buf[headerTotalSizeTransport:], nil
+	const offsetEnd = offsetTransportID + peerIDSize
+	var peerID PeerID
+	copy(peerID[:], buf[offsetTransportID:offsetEnd])
+	return &peerID, buf[headerTotalSizeTransport:], nil
 }
 
 // UnmarshalTransportID extracts the peerID from the transport message.
-func UnmarshalTransportID(buf []byte) ([]byte, error) {
+func UnmarshalTransportID(buf []byte) (*PeerID, error) {
 	if len(buf) < headerTotalSizeTransport {
 		return nil, ErrInvalidMessageLength
 	}
-	return buf[offsetTransportID:headerTotalSizeTransport], nil
+
+	const offsetEnd = offsetTransportID + peerIDSize
+	var id PeerID
+	copy(id[:], buf[offsetTransportID:offsetEnd])
+	return &id, nil
 }
 
 // UpdateTransportMsg updates the peerID in the transport message.
 // With this function the server can reuse the given byte slice to update the peerID in the transport message. So do
 // need to allocate a new byte slice.
-func UpdateTransportMsg(msg []byte, peerID []byte) error {
-	if len(msg) < offsetTransportID+len(peerID) {
+func UpdateTransportMsg(msg []byte, peerID PeerID) error {
+	if len(msg) < offsetTransportID+peerIDSize {
 		return ErrInvalidMessageLength
 	}
-	copy(msg[offsetTransportID:], peerID)
+	copy(msg[offsetTransportID:], peerID[:])
 	return nil
 }
 

relay/messages/message_test.go

@@ -5,7 +5,7 @@ import (
 )
 
 func TestMarshalHelloMsg(t *testing.T) {
-	peerID := []byte("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
+	peerID := HashID("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
 	msg, err := MarshalHelloMsg(peerID, nil)
 	if err != nil {
 		t.Fatalf("error: %v", err)
@@ -24,13 +24,13 @@ func TestMarshalHelloMsg(t *testing.T) {
 	if err != nil {
 		t.Fatalf("error: %v", err)
 	}
-	if string(receivedPeerID) != string(peerID) {
+	if receivedPeerID.String() != peerID.String() {
 		t.Errorf("expected %s, got %s", peerID, receivedPeerID)
 	}
 }
 
 func TestMarshalAuthMsg(t *testing.T) {
-	peerID := []byte("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
+	peerID := HashID("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
 	msg, err := MarshalAuthMsg(peerID, []byte{})
 	if err != nil {
 		t.Fatalf("error: %v", err)
@@ -49,7 +49,7 @@ func TestMarshalAuthMsg(t *testing.T) {
 	if err != nil {
 		t.Fatalf("error: %v", err)
 	}
-	if string(receivedPeerID) != string(peerID) {
+	if receivedPeerID.String() != peerID.String() {
 		t.Errorf("expected %s, got %s", peerID, receivedPeerID)
 	}
 }
@@ -80,7 +80,7 @@ func TestMarshalAuthResponse(t *testing.T) {
 }
 
 func TestMarshalTransportMsg(t *testing.T) {
-	peerID := []byte("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
+	peerID := HashID("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
 	payload := []byte("payload")
 	msg, err := MarshalTransportMsg(peerID, payload)
 	if err != nil {
@@ -101,7 +101,7 @@ func TestMarshalTransportMsg(t *testing.T) {
 		t.Fatalf("failed to unmarshal transport id: %v", err)
 	}
 
-	if string(uPeerID) != string(peerID) {
+	if uPeerID.String() != peerID.String() {
 		t.Errorf("expected %s, got %s", peerID, uPeerID)
 	}
 
@@ -110,8 +110,8 @@ func TestMarshalTransportMsg(t *testing.T) {
 		t.Fatalf("error: %v", err)
 	}
 
-	if string(id) != string(peerID) {
-		t.Errorf("expected %s, got %s", peerID, id)
+	if id.String() != peerID.String() {
+		t.Errorf("expected: '%s', got: '%s'", peerID, id)
 	}
 
 	if string(respPayload) != string(payload) {

relay/messages/peer_state.go

@@ -0,0 +1,92 @@
+package messages
+
+import (
+	"fmt"
+)
+
+func MarshalSubPeerStateMsg(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(ids, byte(MsgTypeSubscribePeerState))
+}
+
+func UnmarshalSubPeerStateMsg(buf []byte) ([]PeerID, error) {
+	return unmarshalPeerIDs(buf)
+}
+
+func MarshalUnsubPeerStateMsg(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(ids, byte(MsgTypeUnsubscribePeerState))
+}
+
+func UnmarshalUnsubPeerStateMsg(buf []byte) ([]PeerID, error) {
+	return unmarshalPeerIDs(buf)
+}
+
+func MarshalPeersOnline(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(ids, byte(MsgTypePeersOnline))
+}
+
+func UnmarshalPeersOnlineMsg(buf []byte) ([]PeerID, error) {
+	return unmarshalPeerIDs(buf)
+}
+
+func MarshalPeersWentOffline(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(ids, byte(MsgTypePeersWentOffline))
+}
+
+func UnMarshalPeersWentOffline(buf []byte) ([]PeerID, error) {
+	return unmarshalPeerIDs(buf)
+}
+
+// marshalPeerIDs is a generic function to marshal peer IDs with a specific message type
+func marshalPeerIDs(ids []PeerID, msgType byte) ([][]byte, error) {
+	if len(ids) == 0 {
+		return nil, fmt.Errorf("no list of peer ids provided")
+	}
+
+	const maxPeersPerMessage = (MaxMessageSize - sizeOfProtoHeader) / peerIDSize
+	var messages [][]byte
+
+	for i := 0; i < len(ids); i += maxPeersPerMessage {
+		end := i + maxPeersPerMessage
+		if end > len(ids) {
+			end = len(ids)
+		}
+		chunk := ids[i:end]
+
+		totalSize := sizeOfProtoHeader + len(chunk)*peerIDSize
+		buf := make([]byte, totalSize)
+		buf[0] = byte(CurrentProtocolVersion)
+		buf[1] = msgType
+
+		offset := sizeOfProtoHeader
+		for _, id := range chunk {
+			copy(buf[offset:], id[:])
+			offset += peerIDSize
+		}
+
+		messages = append(messages, buf)
+	}
+
+	return messages, nil
+}
+
+// unmarshalPeerIDs is a generic function to unmarshal peer IDs from a buffer
+func unmarshalPeerIDs(buf []byte) ([]PeerID, error) {
+	if len(buf) < sizeOfProtoHeader {
+		return nil, fmt.Errorf("invalid message format")
+	}
+
+	if (len(buf)-sizeOfProtoHeader)%peerIDSize != 0 {
+		return nil, fmt.Errorf("invalid peer list size: %d", len(buf)-sizeOfProtoHeader)
+	}
+
+	numIDs := (len(buf) - sizeOfProtoHeader) / peerIDSize
+
+	ids := make([]PeerID, numIDs)
+	offset := sizeOfProtoHeader
+	for i := 0; i < numIDs; i++ {
+		copy(ids[i][:], buf[offset:offset+peerIDSize])
+		offset += peerIDSize
+	}
+
+	return ids, nil
+}