disable initial sync

.github/ISSUE_TEMPLATE/bug-issue-report.md

@@ -37,21 +37,16 @@ If yes, which one?
 
 **Debug output**
 
-To help us resolve the problem, please attach the following anonymized status output
+To help us resolve the problem, please attach the following debug output
 
   netbird status -dA
 
-Create and upload a debug bundle, and share the returned file key:
-
-  netbird debug for 1m -AS -U
-
-*Uploaded files are automatically deleted after 30 days.*
-
-
-Alternatively, create the file only and attach it here manually:
+As well as the file created by
 
   netbird debug for 1m -AS
 
+  
+We advise reviewing the anonymized output for any remaining personal information.
 
 **Screenshots**
 
@@ -62,10 +57,8 @@ If applicable, add screenshots to help explain your problem.
 Add any other context about the problem here.
 
 **Have you tried these troubleshooting steps?**
-- [ ] Reviewed [client troubleshooting](https://docs.netbird.io/how-to/troubleshooting-client) (if applicable)
 - [ ] Checked for newer NetBird versions
 - [ ] Searched for similar issues on GitHub (including closed ones)
 - [ ] Restarted the NetBird client
 - [ ] Disabled other VPN software
 - [ ] Checked firewall settings
-

.github/pull_request_template.md

@@ -13,5 +13,3 @@
 - [ ] It is a refactor
 - [ ] Created tests that fail without the change (if possible)
 - [ ] Extended the README / documentation, if necessary
-
-> By submitting this pull request, you confirm that you have read and agree to the terms of the [Contributor License Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

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

@@ -146,65 +146,6 @@ jobs:
       - 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)
 
-  test_client_on_docker:
-    name: "Client (Docker) / Unit"
-    needs: [build-cache]
-    runs-on: ubuntu-22.04
-    steps:
-      - name: Install Go
-        uses: actions/setup-go@v5
-        with:
-          go-version: "1.23.x"
-          cache: false
-
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Get Go environment
-        id: go-env
-        run: |
-          echo "cache_dir=$(go env GOCACHE)" >> $GITHUB_OUTPUT
-          echo "modcache_dir=$(go env GOMODCACHE)" >> $GITHUB_OUTPUT
-
-      - name: Cache Go modules
-        uses: actions/cache/restore@v4
-        id: cache-restore
-        with:
-          path: |
-            ${{ steps.go-env.outputs.cache_dir }}
-            ${{ steps.go-env.outputs.modcache_dir }}
-          key: ${{ runner.os }}-gotest-cache-${{ hashFiles('**/go.sum') }}
-          restore-keys: |
-            ${{ runner.os }}-gotest-cache-
-
-      - name: Run tests in container
-        env:
-          HOST_GOCACHE: ${{ steps.go-env.outputs.cache_dir }}
-          HOST_GOMODCACHE: ${{ steps.go-env.outputs.modcache_dir }}
-        run: |
-          CONTAINER_GOCACHE="/root/.cache/go-build"
-          CONTAINER_GOMODCACHE="/go/pkg/mod"
-
-          docker run --rm \
-            --cap-add=NET_ADMIN \
-            --privileged \
-            -v $PWD:/app \
-            -w /app \
-            -v "${HOST_GOCACHE}:${CONTAINER_GOCACHE}" \
-            -v "${HOST_GOMODCACHE}:${CONTAINER_GOMODCACHE}" \
-            -e CGO_ENABLED=1 \
-            -e CI=true \
-            -e DOCKER_CI=true \
-            -e GOARCH=${GOARCH_TARGET} \
-            -e GOCACHE=${CONTAINER_GOCACHE} \
-            -e GOMODCACHE=${CONTAINER_GOMODCACHE} \
-            golang:1.23-alpine \
-            sh -c ' \
-              apk update; apk add --no-cache \
-                ca-certificates iptables ip6tables dbus dbus-dev libpcap-dev build-base; \
-              go test -buildvcs=false -tags devcert -v -timeout 10m -p 1 $(go list -buildvcs=false ./... | grep -v -e /management -e /signal -e /relay -e /client/ui -e /upload-server)
-            '
-
   test_relay:
     name: "Relay / Unit"
     needs: [build-cache]
@@ -223,10 +164,6 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Install dependencies
-        if: steps.cache.outputs.cache-hit != 'true'
-        run: sudo apt update && sudo apt install -y gcc-multilib g++-multilib libc6-dev-i386
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -242,6 +179,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -273,10 +217,6 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Install dependencies
-        if: steps.cache.outputs.cache-hit != 'true'
-        run: sudo apt update && sudo apt install -y gcc-multilib g++-multilib libc6-dev-i386
-
       - name: Get Go environment
         run: |
           echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
@@ -292,6 +232,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -339,6 +286,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -400,6 +354,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -424,7 +385,7 @@ jobs:
           CI=true \
           go test -tags devcert -run=^$ -bench=. \
           -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' \
-          -timeout 20m ./management/...
+          -timeout 20m ./...
 
   api_benchmark:
     name: "Management / Benchmark (API)"
@@ -488,6 +449,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -552,6 +520,13 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gotest-cache-
 
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install 32-bit libpcap
+        if: matrix.arch == '386'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
       - name: Install modules
         run: go mod tidy
 
@@ -566,3 +541,99 @@ jobs:
           go test -tags=integration \
           -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' \
           -timeout 20m ./management/...
+
+  test_client_on_docker:
+    name: "Client (Docker) / Unit"
+    needs: [ build-cache ]
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: "1.23.x"
+          cache: false
+
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Get Go environment
+        run: |
+          echo "cache=$(go env GOCACHE)" >> $GITHUB_ENV
+          echo "modcache=$(go env.GOMODCACHE)" >> $GITHUB_ENV     
+
+      - name: Cache Go modules
+        uses: actions/cache/restore@v4
+        with:
+          path: |
+            ${{ env.cache }}
+            ${{ env.modcache }}
+          key: ${{ runner.os }}-gotest-cache-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.os }}-gotest-cache-
+
+      - name: Install dependencies
+        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
+
+      - name: Install modules
+        run: go mod tidy
+
+      - name: Check git status
+        run: git --no-pager diff --exit-code
+
+      - name: Generate Shared Sock Test bin
+        run: CGO_ENABLED=0 go test -c -o sharedsock-testing.bin ./sharedsock
+
+      - name: Generate RouteManager Test bin
+        run: CGO_ENABLED=0 go test -c -o routemanager-testing.bin ./client/internal/routemanager
+
+      - name: Generate SystemOps Test bin (static via Alpine)
+        run: |
+          docker run --rm -v $PWD:/app -w /app \
+            alpine:latest \
+            sh -c "
+              apk add --no-cache go gcc musl-dev libpcap-dev dbus-dev && \
+              adduser -D -u $(id -u) builder && \
+              su builder -c '\
+                cd /app && \
+                CGO_ENABLED=1 GOOS=linux GOARCH=amd64 \
+                go test -c -o /app/systemops-testing.bin \
+                  -tags netgo \
+                  -ldflags=\"-w -extldflags \\\"-static -ldbus-1 -lpcap\\\"\" \
+                  ./client/internal/routemanager/systemops \
+              '
+            "
+
+      - name: Generate nftables Manager Test bin
+        run: CGO_ENABLED=0 go test -c -o nftablesmanager-testing.bin ./client/firewall/nftables/...
+
+      - name: Generate Engine Test bin
+        run: CGO_ENABLED=1 go test -c -o engine-testing.bin ./client/internal
+
+      - name: Generate Peer Test bin
+        run: CGO_ENABLED=0 go test -c -o peer-testing.bin ./client/internal/peer/
+
+      - run: chmod +x *testing.bin
+
+      - name: Run Shared Sock tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/sharedsock --entrypoint /ci/sharedsock-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run Iface tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/netbird -v /tmp/cache:/tmp/cache -v /tmp/modcache:/tmp/modcache -w /netbird -e GOCACHE=/tmp/cache -e GOMODCACHE=/tmp/modcache -e CGO_ENABLED=0 golang:1.23-alpine go test -test.timeout 5m -test.parallel 1 ./client/iface/...
+
+      - name: Run RouteManager tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/internal/routemanager --entrypoint /ci/routemanager-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run SystemOps tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/internal/routemanager/systemops --entrypoint /ci/systemops-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run nftables Manager tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/firewall --entrypoint /ci/nftablesmanager-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run Engine tests in docker with file store
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/internal -e NETBIRD_STORE_ENGINE="jsonfile" --entrypoint /ci/engine-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run Engine tests in docker with sqlite store
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/internal -e NETBIRD_STORE_ENGINE="sqlite" --entrypoint /ci/engine-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1
+
+      - name: Run Peer tests in docker
+        run: docker run -t --cap-add=NET_ADMIN --privileged --rm -v $PWD:/ci -w /ci/client/internal/peer --entrypoint /ci/peer-testing.bin gcr.io/distroless/base:debug -test.timeout 5m -test.parallel 1

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

@@ -179,7 +179,6 @@ jobs:
           grep -A 7 Relay management.json | grep "rel://$CI_NETBIRD_DOMAIN:33445"
           grep -A 7 Relay management.json | egrep '"Secret": ".+"'
           grep DisablePromptLogin management.json | grep 'true'
-          grep LoginFlag management.json | grep 0
 
       - name: Install modules
         run: go mod tidy

.goreleaser.yaml

@@ -96,20 +96,6 @@ builds:
       - -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
     mod_timestamp: "{{ .CommitTimestamp }}"
 
-  - id: netbird-upload
-    dir: upload-server
-    env: [CGO_ENABLED=0]
-    binary: netbird-upload
-    goos:
-      - linux
-    goarch:
-      - amd64
-      - arm64
-      - arm
-    ldflags:
-      - -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
-    mod_timestamp: "{{ .CommitTimestamp }}"
-
 universal_binaries:
   - id: netbird
 
@@ -423,52 +409,6 @@ dockers:
       - "--label=org.opencontainers.image.revision={{.FullCommit}}"
       - "--label=org.opencontainers.image.version={{.Version}}"
       - "--label=maintainer=dev@netbird.io"
-  - image_templates:
-      - netbirdio/upload:{{ .Version }}-amd64
-    ids:
-      - netbird-upload
-    goarch: amd64
-    use: buildx
-    dockerfile: upload-server/Dockerfile
-    build_flag_templates:
-      - "--platform=linux/amd64"
-      - "--label=org.opencontainers.image.created={{.Date}}"
-      - "--label=org.opencontainers.image.title={{.ProjectName}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=org.opencontainers.image.revision={{.FullCommit}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=maintainer=dev@netbird.io"
-  - image_templates:
-      - netbirdio/upload:{{ .Version }}-arm64v8
-    ids:
-      - netbird-upload
-    goarch: arm64
-    use: buildx
-    dockerfile: upload-server/Dockerfile
-    build_flag_templates:
-      - "--platform=linux/arm64"
-      - "--label=org.opencontainers.image.created={{.Date}}"
-      - "--label=org.opencontainers.image.title={{.ProjectName}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=org.opencontainers.image.revision={{.FullCommit}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=maintainer=dev@netbird.io"
-  - image_templates:
-      - netbirdio/upload:{{ .Version }}-arm
-    ids:
-      - netbird-upload
-    goarch: arm
-    goarm: 6
-    use: buildx
-    dockerfile: upload-server/Dockerfile
-    build_flag_templates:
-      - "--platform=linux/arm"
-      - "--label=org.opencontainers.image.created={{.Date}}"
-      - "--label=org.opencontainers.image.title={{.ProjectName}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=org.opencontainers.image.revision={{.FullCommit}}"
-      - "--label=org.opencontainers.image.version={{.Version}}"
-      - "--label=maintainer=dev@netbird.io"
 docker_manifests:
   - name_template: netbirdio/netbird:{{ .Version }}
     image_templates:
@@ -535,17 +475,7 @@ docker_manifests:
       - netbirdio/management:{{ .Version }}-debug-arm64v8
       - netbirdio/management:{{ .Version }}-debug-arm
       - netbirdio/management:{{ .Version }}-debug-amd64
-  - name_template: netbirdio/upload:{{ .Version }}
-    image_templates:
-      - netbirdio/upload:{{ .Version }}-arm64v8
-      - netbirdio/upload:{{ .Version }}-arm
-      - netbirdio/upload:{{ .Version }}-amd64
 
-  - name_template: netbirdio/upload:latest
-    image_templates:
-      - netbirdio/upload:{{ .Version }}-arm64v8
-      - netbirdio/upload:{{ .Version }}-arm
-      - netbirdio/upload:{{ .Version }}-amd64
 brews:
   - ids:
       - default

client/Dockerfile

@@ -1,6 +1,5 @@
 FROM alpine:3.21.3
-# iproute2: busybox doesn't display ip rules properly
-RUN apk add --no-cache ca-certificates ip6tables iproute2 iptables
+RUN apk add --no-cache ca-certificates iptables ip6tables
 ENV NB_FOREGROUND_MODE=true
 ENTRYPOINT [ "/usr/local/bin/netbird","up"]
-COPY netbird /usr/local/bin/netbird
+COPY netbird /usr/local/bin/netbird

client/cmd/debug.go

@@ -87,27 +87,16 @@ func debugBundle(cmd *cobra.Command, _ []string) error {
 	}()
 
 	client := proto.NewDaemonServiceClient(conn)
-	request := &proto.DebugBundleRequest{
+	resp, err := client.DebugBundle(cmd.Context(), &proto.DebugBundleRequest{
 		Anonymize:  anonymizeFlag,
 		Status:     getStatusOutput(cmd, anonymizeFlag),
 		SystemInfo: debugSystemInfoFlag,
-	}
-	if debugUploadBundle {
-		request.UploadURL = debugUploadBundleURL
-	}
-	resp, err := client.DebugBundle(cmd.Context(), request)
+	})
 	if err != nil {
 		return fmt.Errorf("failed to bundle debug: %v", status.Convert(err).Message())
 	}
-	cmd.Printf("Local file:\n%s\n", resp.GetPath())
 
-	if resp.GetUploadFailureReason() != "" {
-		return fmt.Errorf("upload failed: %s", resp.GetUploadFailureReason())
-	}
-
-	if debugUploadBundle {
-		cmd.Printf("Upload file key:\n%s\n", resp.GetUploadedKey())
-	}
+	cmd.Println(resp.GetPath())
 
 	return nil
 }
@@ -222,19 +211,23 @@ 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{
+
+	resp, err := client.DebugBundle(cmd.Context(), &proto.DebugBundleRequest{
 		Anonymize:  anonymizeFlag,
 		Status:     statusOutput,
 		SystemInfo: debugSystemInfoFlag,
-	}
-	if debugUploadBundle {
-		request.UploadURL = debugUploadBundleURL
-	}
-	resp, err := client.DebugBundle(cmd.Context(), request)
+	})
 	if err != nil {
 		return fmt.Errorf("failed to bundle debug: %v", status.Convert(err).Message())
 	}
 
+	// Disable network map persistence after creating the debug bundle
+	if _, err := client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
+		Enabled: false,
+	}); err != nil {
+		return fmt.Errorf("failed to disable network map persistence: %v", status.Convert(err).Message())
+	}
+
 	if stateWasDown {
 		if _, err := client.Down(cmd.Context(), &proto.DownRequest{}); err != nil {
 			return fmt.Errorf("failed to down: %v", status.Convert(err).Message())
@@ -249,15 +242,7 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		cmd.Println("Log level restored to", initialLogLevel.GetLevel())
 	}
 
-	cmd.Printf("Local file:\n%s\n", resp.GetPath())
-
-	if resp.GetUploadFailureReason() != "" {
-		return fmt.Errorf("upload failed: %s", resp.GetUploadFailureReason())
-	}
-
-	if debugUploadBundle {
-		cmd.Printf("Upload file key:\n%s\n", resp.GetUploadedKey())
-	}
+	cmd.Println(resp.GetPath())
 
 	return nil
 }

client/cmd/login.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"runtime"
 	"strings"
 	"time"
 
@@ -56,9 +55,6 @@ var loginCmd = &cobra.Command{
 				return err
 			}
 
-			// update host's static platform and system information
-			system.UpdateStaticInfo()
-
 			ic := internal.ConfigInput{
 				ManagementURL: managementURL,
 				AdminURL:      adminURL,
@@ -99,11 +95,11 @@ var loginCmd = &cobra.Command{
 		}
 
 		loginRequest := proto.LoginRequest{
-			SetupKey:            providedSetupKey,
-			ManagementUrl:       managementURL,
-			IsUnixDesktopClient: isUnixRunningDesktop(),
-			Hostname:            hostName,
-			DnsLabels:           dnsLabelsReq,
+			SetupKey:             providedSetupKey,
+			ManagementUrl:        managementURL,
+			IsLinuxDesktopClient: isLinuxRunningDesktop(),
+			Hostname:             hostName,
+			DnsLabels:            dnsLabelsReq,
 		}
 
 		if rootCmd.PersistentFlags().Changed(preSharedKeyFlag) {
@@ -196,7 +192,7 @@ func foregroundLogin(ctx context.Context, cmd *cobra.Command, config *internal.C
 }
 
 func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *internal.Config) (*auth.TokenInfo, error) {
-	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isUnixRunningDesktop())
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isLinuxRunningDesktop())
 	if err != nil {
 		return nil, err
 	}
@@ -244,10 +240,7 @@ func openURL(cmd *cobra.Command, verificationURIComplete, userCode string, noBro
 	}
 }
 
-// isUnixRunningDesktop checks if a Linux OS is running desktop environment
-func isUnixRunningDesktop() bool {
-	if runtime.GOOS != "linux" && runtime.GOOS != "freebsd" {
-		return false
-	}
+// isLinuxRunningDesktop checks if a Linux OS is running desktop environment
+func isLinuxRunningDesktop() bool {
 	return os.Getenv("DESKTOP_SESSION") != "" || os.Getenv("XDG_CURRENT_DESKTOP") != ""
 }

client/cmd/root.go

@@ -22,27 +22,23 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 
 	"github.com/netbirdio/netbird/client/internal"
-	"github.com/netbirdio/netbird/upload-server/types"
 )
 
 const (
-	externalIPMapFlag        = "external-ip-map"
-	dnsResolverAddress       = "dns-resolver-address"
-	enableRosenpassFlag      = "enable-rosenpass"
-	rosenpassPermissiveFlag  = "rosenpass-permissive"
-	preSharedKeyFlag         = "preshared-key"
-	interfaceNameFlag        = "interface-name"
-	wireguardPortFlag        = "wireguard-port"
-	networkMonitorFlag       = "network-monitor"
-	disableAutoConnectFlag   = "disable-auto-connect"
-	serverSSHAllowedFlag     = "allow-server-ssh"
-	extraIFaceBlackListFlag  = "extra-iface-blacklist"
-	dnsRouteIntervalFlag     = "dns-router-interval"
-	systemInfoFlag           = "system-info"
-	blockLANAccessFlag       = "block-lan-access"
-	enableLazyConnectionFlag = "enable-lazy-connection"
-	uploadBundle             = "upload-bundle"
-	uploadBundleURL          = "upload-bundle-url"
+	externalIPMapFlag       = "external-ip-map"
+	dnsResolverAddress      = "dns-resolver-address"
+	enableRosenpassFlag     = "enable-rosenpass"
+	rosenpassPermissiveFlag = "rosenpass-permissive"
+	preSharedKeyFlag        = "preshared-key"
+	interfaceNameFlag       = "interface-name"
+	wireguardPortFlag       = "wireguard-port"
+	networkMonitorFlag      = "network-monitor"
+	disableAutoConnectFlag  = "disable-auto-connect"
+	serverSSHAllowedFlag    = "allow-server-ssh"
+	extraIFaceBlackListFlag = "extra-iface-blacklist"
+	dnsRouteIntervalFlag    = "dns-router-interval"
+	systemInfoFlag          = "system-info"
+	blockLANAccessFlag      = "block-lan-access"
 )
 
 var (
@@ -79,9 +75,6 @@ var (
 	debugSystemInfoFlag     bool
 	dnsRouteInterval        time.Duration
 	blockLANAccess          bool
-	debugUploadBundle       bool
-	debugUploadBundleURL    string
-	lazyConnEnabled         bool
 
 	rootCmd = &cobra.Command{
 		Use:          "netbird",
@@ -186,11 +179,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.")
 
 	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/service_controller.go

@@ -16,17 +16,12 @@ import (
 
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/server"
-	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/util"
 )
 
 func (p *program) Start(svc service.Service) error {
 	// Start should not block. Do the actual work async.
 	log.Info("starting Netbird service") //nolint
-
-	// Collect static system and platform information
-	system.UpdateStaticInfo()
-
 	// in any case, even if configuration does not exists we run daemon to serve CLI gRPC API.
 	p.serv = grpc.NewServer()
 

client/cmd/status.go

@@ -44,7 +44,7 @@ func init() {
 	statusCmd.MarkFlagsMutuallyExclusive("detail", "json", "yaml", "ipv4")
 	statusCmd.PersistentFlags().StringSliceVar(&ipsFilter, "filter-by-ips", []string{}, "filters the detailed output by a list of one or more IPs, e.g., --filter-by-ips 100.64.0.100,100.64.0.200")
 	statusCmd.PersistentFlags().StringSliceVar(&prefixNamesFilter, "filter-by-names", []string{}, "filters the detailed output by a list of one or more peer FQDN or hostnames, e.g., --filter-by-names peer-a,peer-b.netbird.cloud")
-	statusCmd.PersistentFlags().StringVar(&statusFilter, "filter-by-status", "", "filters the detailed output by connection status(idle|connecting|connected), e.g., --filter-by-status connected")
+	statusCmd.PersistentFlags().StringVar(&statusFilter, "filter-by-status", "", "filters the detailed output by connection status(connected|disconnected), e.g., --filter-by-status connected")
 }
 
 func statusFunc(cmd *cobra.Command, args []string) error {
@@ -127,12 +127,12 @@ func getStatus(ctx context.Context) (*proto.StatusResponse, error) {
 
 func parseFilters() error {
 	switch strings.ToLower(statusFilter) {
-	case "", "idle", "connecting", "connected":
+	case "", "disconnected", "connected":
 		if strings.ToLower(statusFilter) != "" {
 			enableDetailFlagWhenFilterFlag()
 		}
 	default:
-		return fmt.Errorf("wrong status filter, should be one of connected|connecting|idle, got: %s", statusFilter)
+		return fmt.Errorf("wrong status filter, should be one of connected|disconnected, got: %s", statusFilter)
 	}
 
 	if len(ipsFilter) > 0 {

client/cmd/up.go

@@ -194,10 +194,6 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command) error {
 		ic.BlockLANAccess = &blockLANAccess
 	}
 
-	if cmd.Flag(enableLazyConnectionFlag).Changed {
-		ic.LazyConnectionEnabled = &lazyConnEnabled
-	}
-
 	providedSetupKey, err := getSetupKey()
 	if err != nil {
 		return err
@@ -266,17 +262,17 @@ func runInDaemonMode(ctx context.Context, cmd *cobra.Command) error {
 	}
 
 	loginRequest := proto.LoginRequest{
-		SetupKey:            providedSetupKey,
-		ManagementUrl:       managementURL,
-		AdminURL:            adminURL,
-		NatExternalIPs:      natExternalIPs,
-		CleanNATExternalIPs: natExternalIPs != nil && len(natExternalIPs) == 0,
-		CustomDNSAddress:    customDNSAddressConverted,
-		IsUnixDesktopClient: isUnixRunningDesktop(),
-		Hostname:            hostName,
-		ExtraIFaceBlacklist: extraIFaceBlackList,
-		DnsLabels:           dnsLabels,
-		CleanDNSLabels:      dnsLabels != nil && len(dnsLabels) == 0,
+		SetupKey:             providedSetupKey,
+		ManagementUrl:        managementURL,
+		AdminURL:             adminURL,
+		NatExternalIPs:       natExternalIPs,
+		CleanNATExternalIPs:  natExternalIPs != nil && len(natExternalIPs) == 0,
+		CustomDNSAddress:     customDNSAddressConverted,
+		IsLinuxDesktopClient: isLinuxRunningDesktop(),
+		Hostname:             hostName,
+		ExtraIFaceBlacklist:  extraIFaceBlackList,
+		DnsLabels:            dnsLabels,
+		CleanDNSLabels:       dnsLabels != nil && len(dnsLabels) == 0,
 	}
 
 	if rootCmd.PersistentFlags().Changed(preSharedKeyFlag) {
@@ -336,10 +332,6 @@ func runInDaemonMode(ctx context.Context, cmd *cobra.Command) error {
 		loginRequest.BlockLanAccess = &blockLANAccess
 	}
 
-	if cmd.Flag(enableLazyConnectionFlag).Changed {
-		loginRequest.LazyConnectionEnabled = &lazyConnEnabled
-	}
-
 	var loginErr error
 
 	var loginResp *proto.LoginResponse

client/firewall/iptables/manager_linux.go

@@ -113,16 +113,17 @@ func (m *Manager) AddPeerFiltering(
 func (m *Manager) AddRouteFiltering(
 	id []byte,
 	sources []netip.Prefix,
-	destination firewall.Network,
+	destination netip.Prefix,
 	proto firewall.Protocol,
-	sPort, dPort *firewall.Port,
+	sPort *firewall.Port,
+	dPort *firewall.Port,
 	action firewall.Action,
 ) (firewall.Rule, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if destination.IsPrefix() && !destination.Prefix.Addr().Is4() {
-		return nil, fmt.Errorf("unsupported IP version: %s", destination.Prefix.Addr().String())
+	if !destination.Addr().Is4() {
+		return nil, fmt.Errorf("unsupported IP version: %s", destination.Addr().String())
 	}
 
 	return m.router.AddRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
@@ -242,14 +243,6 @@ func (m *Manager) DeleteDNATRule(rule firewall.Rule) error {
 	return m.router.DeleteDNATRule(rule)
 }
 
-// UpdateSet updates the set with the given prefixes
-func (m *Manager) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.UpdateSet(set, prefixes)
-}
-
 func getConntrackEstablished() []string {
 	return []string{"-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT"}
 }

client/firewall/iptables/router_linux.go

@@ -57,18 +57,18 @@ type ruleInfo struct {
 }
 
 type routeFilteringRuleParams struct {
-	Source      firewall.Network
-	Destination firewall.Network
+	Sources     []netip.Prefix
+	Destination netip.Prefix
 	Proto       firewall.Protocol
 	SPort       *firewall.Port
 	DPort       *firewall.Port
 	Direction   firewall.RuleDirection
 	Action      firewall.Action
+	SetName     string
 }
 
 type routeRules map[string][]string
 
-// the ipset library currently does not support comments, so we use the name only (string)
 type ipsetCounter = refcounter.Counter[string, []netip.Prefix, struct{}]
 
 type router struct {
@@ -129,7 +129,7 @@ func (r *router) init(stateManager *statemanager.Manager) error {
 func (r *router) AddRouteFiltering(
 	id []byte,
 	sources []netip.Prefix,
-	destination firewall.Network,
+	destination netip.Prefix,
 	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
@@ -140,28 +140,27 @@ func (r *router) AddRouteFiltering(
 		return ruleKey, nil
 	}
 
-	var source firewall.Network
+	var setName string
 	if len(sources) > 1 {
-		source.Set = firewall.NewPrefixSet(sources)
-	} else if len(sources) > 0 {
-		source.Prefix = sources[0]
+		setName = firewall.GenerateSetName(sources)
+		if _, err := r.ipsetCounter.Increment(setName, sources); err != nil {
+			return nil, fmt.Errorf("create or get ipset: %w", err)
+		}
 	}
 
 	params := routeFilteringRuleParams{
-		Source:      source,
+		Sources:     sources,
 		Destination: destination,
 		Proto:       proto,
 		SPort:       sPort,
 		DPort:       dPort,
 		Action:      action,
+		SetName:     setName,
 	}
 
-	rule, err := r.genRouteRuleSpec(params, sources)
-	if err != nil {
-		return nil, fmt.Errorf("generate route rule spec: %w", err)
-	}
-
+	rule := genRouteFilteringRuleSpec(params)
 	// Insert DROP rules at the beginning, append ACCEPT rules at the end
+	var err error
 	if action == firewall.ActionDrop {
 		// after the established rule
 		err = r.iptablesClient.Insert(tableFilter, chainRTFWDIN, 2, rule...)
@@ -184,13 +183,17 @@ func (r *router) DeleteRouteRule(rule firewall.Rule) error {
 	ruleKey := rule.ID()
 
 	if rule, exists := r.rules[ruleKey]; exists {
+		setName := r.findSetNameInRule(rule)
+
 		if err := r.iptablesClient.Delete(tableFilter, chainRTFWDIN, rule...); err != nil {
 			return fmt.Errorf("delete route rule: %v", err)
 		}
 		delete(r.rules, ruleKey)
 
-		if err := r.decrementSetCounter(rule); err != nil {
-			return fmt.Errorf("decrement ipset counter: %w", err)
+		if setName != "" {
+			if _, err := r.ipsetCounter.Decrement(setName); err != nil {
+				return fmt.Errorf("failed to remove ipset: %w", err)
+			}
 		}
 	} else {
 		log.Debugf("route rule %s not found", ruleKey)
@@ -201,26 +204,13 @@ func (r *router) DeleteRouteRule(rule firewall.Rule) error {
 	return nil
 }
 
-func (r *router) decrementSetCounter(rule []string) error {
-	sets := r.findSets(rule)
-	var merr *multierror.Error
-	for _, setName := range sets {
-		if _, err := r.ipsetCounter.Decrement(setName); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("decrement counter: %w", err))
-		}
-	}
-
-	return nberrors.FormatErrorOrNil(merr)
-}
-
-func (r *router) findSets(rule []string) []string {
-	var sets []string
+func (r *router) findSetNameInRule(rule []string) string {
 	for i, arg := range rule {
 		if arg == "-m" && i+3 < len(rule) && rule[i+1] == "set" && rule[i+2] == matchSet {
-			sets = append(sets, rule[i+3])
+			return rule[i+3]
 		}
 	}
-	return sets
+	return ""
 }
 
 func (r *router) createIpSet(setName string, sources []netip.Prefix) error {
@@ -241,8 +231,6 @@ func (r *router) deleteIpSet(setName string) error {
 	if err := ipset.Destroy(setName); err != nil {
 		return fmt.Errorf("destroy set %s: %w", setName, err)
 	}
-
-	log.Debugf("Deleted unused ipset %s", setName)
 	return nil
 }
 
@@ -282,14 +270,12 @@ func (r *router) RemoveNatRule(pair firewall.RouterPair) error {
 		log.Errorf("%v", err)
 	}
 
-	if pair.Masquerade {
-		if err := r.removeNatRule(pair); err != nil {
-			return fmt.Errorf("remove nat rule: %w", err)
-		}
+	if err := r.removeNatRule(pair); err != nil {
+		return fmt.Errorf("remove nat rule: %w", err)
+	}
 
-		if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
-			return fmt.Errorf("remove inverse nat rule: %w", err)
-		}
+	if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
+		return fmt.Errorf("remove inverse nat rule: %w", err)
 	}
 
 	if err := r.removeLegacyRouteRule(pair); err != nil {
@@ -327,10 +313,8 @@ func (r *router) removeLegacyRouteRule(pair firewall.RouterPair) error {
 			return fmt.Errorf("remove legacy forwarding rule %s -> %s: %v", pair.Source, pair.Destination, err)
 		}
 		delete(r.rules, ruleKey)
-
-		if err := r.decrementSetCounter(rule); err != nil {
-			return fmt.Errorf("decrement ipset counter: %w", err)
-		}
+	} else {
+		log.Debugf("legacy forwarding rule %s not found", ruleKey)
 	}
 
 	return nil
@@ -615,26 +599,12 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 	rule = append(rule,
 		"-m", "conntrack",
 		"--ctstate", "NEW",
-	)
-	sourceExp, err := r.applyNetwork("-s", pair.Source, nil)
-	if err != nil {
-		return fmt.Errorf("apply network -s: %w", err)
-	}
-	destExp, err := r.applyNetwork("-d", pair.Destination, nil)
-	if err != nil {
-		return fmt.Errorf("apply network -d: %w", err)
-	}
-
-	rule = append(rule, sourceExp...)
-	rule = append(rule, destExp...)
-	rule = append(rule,
+		"-s", pair.Source.String(),
+		"-d", pair.Destination.String(),
 		"-j", "MARK", "--set-mark", fmt.Sprintf("%#x", markValue),
 	)
 
-	// Ensure nat rules come first, so the mark can be overwritten.
-	// Currently overwritten by the dst-type LOCAL rules for redirected traffic.
-	if err := r.iptablesClient.Insert(tableMangle, chainRTPRE, 1, rule...); err != nil {
-		// TODO: rollback ipset counter
+	if err := r.iptablesClient.Append(tableMangle, chainRTPRE, rule...); err != nil {
 		return fmt.Errorf("error while adding marking rule for %s: %v", pair.Destination, err)
 	}
 
@@ -652,10 +622,6 @@ func (r *router) removeNatRule(pair firewall.RouterPair) error {
 			return fmt.Errorf("error while removing marking rule for %s: %v", pair.Destination, err)
 		}
 		delete(r.rules, ruleKey)
-
-		if err := r.decrementSetCounter(rule); err != nil {
-			return fmt.Errorf("decrement ipset counter: %w", err)
-		}
 	} else {
 		log.Debugf("marking rule %s not found", ruleKey)
 	}
@@ -821,21 +787,17 @@ func (r *router) DeleteDNATRule(rule firewall.Rule) error {
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-func (r *router) genRouteRuleSpec(params routeFilteringRuleParams, sources []netip.Prefix) ([]string, error) {
+func genRouteFilteringRuleSpec(params routeFilteringRuleParams) []string {
 	var rule []string
 
-	sourceExp, err := r.applyNetwork("-s", params.Source, sources)
-	if err != nil {
-		return nil, fmt.Errorf("apply network -s: %w", err)
-
-	}
-	destExp, err := r.applyNetwork("-d", params.Destination, nil)
-	if err != nil {
-		return nil, fmt.Errorf("apply network -d: %w", err)
+	if params.SetName != "" {
+		rule = append(rule, "-m", "set", matchSet, params.SetName, "src")
+	} else if len(params.Sources) > 0 {
+		source := params.Sources[0]
+		rule = append(rule, "-s", source.String())
 	}
 
-	rule = append(rule, sourceExp...)
-	rule = append(rule, destExp...)
+	rule = append(rule, "-d", params.Destination.String())
 
 	if params.Proto != firewall.ProtocolALL {
 		rule = append(rule, "-p", strings.ToLower(string(params.Proto)))
@@ -845,47 +807,7 @@ func (r *router) genRouteRuleSpec(params routeFilteringRuleParams, sources []net
 
 	rule = append(rule, "-j", actionToStr(params.Action))
 
-	return rule, nil
-}
-
-func (r *router) applyNetwork(flag string, network firewall.Network, prefixes []netip.Prefix) ([]string, error) {
-	direction := "src"
-	if flag == "-d" {
-		direction = "dst"
-	}
-
-	if network.IsSet() {
-		if _, err := r.ipsetCounter.Increment(network.Set.HashedName(), prefixes); err != nil {
-			return nil, fmt.Errorf("create or get ipset: %w", err)
-		}
-
-		return []string{"-m", "set", matchSet, network.Set.HashedName(), direction}, nil
-	}
-	if network.IsPrefix() {
-		return []string{flag, network.Prefix.String()}, nil
-	}
-
-	// nolint:nilnil
-	return nil, nil
-}
-
-func (r *router) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
-	var merr *multierror.Error
-	for _, prefix := range prefixes {
-		// TODO: Implement IPv6 support
-		if prefix.Addr().Is6() {
-			log.Tracef("skipping IPv6 prefix %s: IPv6 support not yet implemented", prefix)
-			continue
-		}
-		if err := ipset.AddPrefix(set.HashedName(), prefix); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("increment ipset counter: %w", err))
-		}
-	}
-	if merr == nil {
-		log.Debugf("updated set %s with prefixes %v", set.HashedName(), prefixes)
-	}
-
-	return nberrors.FormatErrorOrNil(merr)
+	return rule
 }
 
 func applyPort(flag string, port *firewall.Port) []string {

client/firewall/iptables/router_linux_test.go

@@ -60,8 +60,8 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 
 	pair := firewall.RouterPair{
 		ID:          "abc",
-		Source:      firewall.Network{Prefix: netip.MustParsePrefix("100.100.100.1/32")},
-		Destination: firewall.Network{Prefix: netip.MustParsePrefix("100.100.100.0/24")},
+		Source:      netip.MustParsePrefix("100.100.100.1/32"),
+		Destination: netip.MustParsePrefix("100.100.100.0/24"),
 		Masquerade:  true,
 	}
 
@@ -332,7 +332,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, firewall.Network{Prefix: tt.destination}, tt.proto, tt.sPort, tt.dPort, tt.action)
+			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
 			require.NoError(t, err, "AddRouteFiltering failed")
 
 			// Check if the rule is in the internal map
@@ -347,29 +347,23 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			assert.NoError(t, err, "Failed to check rule existence")
 			assert.True(t, exists, "Rule not found in iptables")
 
-			var source firewall.Network
-			if len(tt.sources) > 1 {
-				source.Set = firewall.NewPrefixSet(tt.sources)
-			} else if len(tt.sources) > 0 {
-				source.Prefix = tt.sources[0]
-			}
 			// Verify rule content
 			params := routeFilteringRuleParams{
-				Source:      source,
-				Destination: firewall.Network{Prefix: tt.destination},
+				Sources:     tt.sources,
+				Destination: tt.destination,
 				Proto:       tt.proto,
 				SPort:       tt.sPort,
 				DPort:       tt.dPort,
 				Action:      tt.action,
+				SetName:     "",
 			}
 
-			expectedRule, err := r.genRouteRuleSpec(params, nil)
-			require.NoError(t, err, "Failed to generate expected rule spec")
+			expectedRule := genRouteFilteringRuleSpec(params)
 
 			if tt.expectSet {
-				setName := firewall.NewPrefixSet(tt.sources).HashedName()
-				expectedRule, err = r.genRouteRuleSpec(params, nil)
-				require.NoError(t, err, "Failed to generate expected rule spec with set")
+				setName := firewall.GenerateSetName(tt.sources)
+				params.SetName = setName
+				expectedRule = genRouteFilteringRuleSpec(params)
 
 				// Check if the set was created
 				_, exists := r.ipsetCounter.Get(setName)
@@ -384,62 +378,3 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 		})
 	}
 }
-
-func TestFindSetNameInRule(t *testing.T) {
-	r := &router{}
-
-	testCases := []struct {
-		name     string
-		rule     []string
-		expected []string
-	}{
-		{
-			name: "Basic rule with two sets",
-			rule: []string{
-				"-A", "NETBIRD-RT-FWD-IN", "-p", "tcp", "-m", "set", "--match-set", "nb-2e5a2a05", "src",
-				"-m", "set", "--match-set", "nb-349ae051", "dst", "-m", "tcp", "--dport", "8080", "-j", "ACCEPT",
-			},
-			expected: []string{"nb-2e5a2a05", "nb-349ae051"},
-		},
-		{
-			name:     "No sets",
-			rule:     []string{"-A", "NETBIRD-RT-FWD-IN", "-p", "tcp", "-j", "ACCEPT"},
-			expected: []string{},
-		},
-		{
-			name: "Multiple sets with different positions",
-			rule: []string{
-				"-m", "set", "--match-set", "set1", "src", "-p", "tcp",
-				"-m", "set", "--match-set", "set-abc123", "dst", "-j", "ACCEPT",
-			},
-			expected: []string{"set1", "set-abc123"},
-		},
-		{
-			name:     "Boundary case - sequence appears at end",
-			rule:     []string{"-p", "tcp", "-m", "set", "--match-set", "final-set"},
-			expected: []string{"final-set"},
-		},
-		{
-			name:     "Incomplete pattern - missing set name",
-			rule:     []string{"-p", "tcp", "-m", "set", "--match-set"},
-			expected: []string{},
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			result := r.findSets(tc.rule)
-
-			if len(result) != len(tc.expected) {
-				t.Errorf("Expected %d sets, got %d. Sets found: %v", len(tc.expected), len(result), result)
-				return
-			}
-
-			for i, set := range result {
-				if set != tc.expected[i] {
-					t.Errorf("Expected set %q at position %d, got %q", tc.expected[i], i, set)
-				}
-			}
-		})
-	}
-}

client/firewall/manager/firewall.go

@@ -1,10 +1,13 @@
 package manager
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"net"
 	"net/netip"
 	"sort"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 
@@ -40,18 +43,6 @@ const (
 // Action is the action to be taken on a rule
 type Action int
 
-// String returns the string representation of the action
-func (a Action) String() string {
-	switch a {
-	case ActionAccept:
-		return "accept"
-	case ActionDrop:
-		return "drop"
-	default:
-		return "unknown"
-	}
-}
-
 const (
 	// ActionAccept is the action to accept a packet
 	ActionAccept Action = iota
@@ -59,33 +50,6 @@ const (
 	ActionDrop
 )
 
-// Network is a rule destination, either a set or a prefix
-type Network struct {
-	Set    Set
-	Prefix netip.Prefix
-}
-
-// String returns the string representation of the destination
-func (d Network) String() string {
-	if d.Prefix.IsValid() {
-		return d.Prefix.String()
-	}
-	if d.IsSet() {
-		return d.Set.HashedName()
-	}
-	return "<invalid network>"
-}
-
-// IsSet returns true if the destination is a set
-func (d Network) IsSet() bool {
-	return d.Set != Set{}
-}
-
-// IsPrefix returns true if the destination is a valid prefix
-func (d Network) IsPrefix() bool {
-	return d.Prefix.IsValid()
-}
-
 // Manager is the high level abstraction of a firewall manager
 //
 // It declares methods which handle actions required by the
@@ -119,9 +83,10 @@ type Manager interface {
 	AddRouteFiltering(
 		id []byte,
 		sources []netip.Prefix,
-		destination Network,
+		destination netip.Prefix,
 		proto Protocol,
-		sPort, dPort *Port,
+		sPort *Port,
+		dPort *Port,
 		action Action,
 	) (Rule, error)
 
@@ -154,9 +119,6 @@ type Manager interface {
 
 	// DeleteDNATRule deletes a DNAT rule
 	DeleteDNATRule(Rule) error
-
-	// UpdateSet updates the set with the given prefixes
-	UpdateSet(hash Set, prefixes []netip.Prefix) error
 }
 
 func GenKey(format string, pair RouterPair) string {
@@ -191,6 +153,22 @@ func SetLegacyManagement(router LegacyManager, isLegacy bool) error {
 	return nil
 }
 
+// GenerateSetName generates a unique name for an ipset based on the given sources.
+func GenerateSetName(sources []netip.Prefix) string {
+	// sort for consistent naming
+	SortPrefixes(sources)
+
+	var sourcesStr strings.Builder
+	for _, src := range sources {
+		sourcesStr.WriteString(src.String())
+	}
+
+	hash := sha256.Sum256([]byte(sourcesStr.String()))
+	shortHash := hex.EncodeToString(hash[:])[:8]
+
+	return fmt.Sprintf("nb-%s", shortHash)
+}
+
 // MergeIPRanges merges overlapping IP ranges and returns a slice of non-overlapping netip.Prefix
 func MergeIPRanges(prefixes []netip.Prefix) []netip.Prefix {
 	if len(prefixes) == 0 {

client/firewall/manager/firewall_test.go

@@ -20,8 +20,8 @@ func TestGenerateSetName(t *testing.T) {
 			netip.MustParsePrefix("192.168.1.0/24"),
 		}
 
-		result1 := manager.NewPrefixSet(prefixes1)
-		result2 := manager.NewPrefixSet(prefixes2)
+		result1 := manager.GenerateSetName(prefixes1)
+		result2 := manager.GenerateSetName(prefixes2)
 
 		if result1 != result2 {
 			t.Errorf("Different orders produced different hashes: %s != %s", result1, result2)
@@ -34,9 +34,9 @@ func TestGenerateSetName(t *testing.T) {
 			netip.MustParsePrefix("10.0.0.0/8"),
 		}
 
-		result := manager.NewPrefixSet(prefixes)
+		result := manager.GenerateSetName(prefixes)
 
-		matched, err := regexp.MatchString(`^nb-[0-9a-f]{8}$`, result.HashedName())
+		matched, err := regexp.MatchString(`^nb-[0-9a-f]{8}$`, result)
 		if err != nil {
 			t.Fatalf("Error matching regex: %v", err)
 		}
@@ -46,8 +46,8 @@ func TestGenerateSetName(t *testing.T) {
 	})
 
 	t.Run("Empty input produces consistent result", func(t *testing.T) {
-		result1 := manager.NewPrefixSet([]netip.Prefix{})
-		result2 := manager.NewPrefixSet([]netip.Prefix{})
+		result1 := manager.GenerateSetName([]netip.Prefix{})
+		result2 := manager.GenerateSetName([]netip.Prefix{})
 
 		if result1 != result2 {
 			t.Errorf("Empty input produced inconsistent results: %s != %s", result1, result2)
@@ -64,8 +64,8 @@ func TestGenerateSetName(t *testing.T) {
 			netip.MustParsePrefix("192.168.1.0/24"),
 		}
 
-		result1 := manager.NewPrefixSet(prefixes1)
-		result2 := manager.NewPrefixSet(prefixes2)
+		result1 := manager.GenerateSetName(prefixes1)
+		result2 := manager.GenerateSetName(prefixes2)
 
 		if result1 != result2 {
 			t.Errorf("Different orders of IPv4 and IPv6 produced different hashes: %s != %s", result1, result2)

client/firewall/manager/routerpair.go

@@ -1,13 +1,15 @@
 package manager
 
 import (
+	"net/netip"
+
 	"github.com/netbirdio/netbird/route"
 )
 
 type RouterPair struct {
 	ID          route.ID
-	Source      Network
-	Destination Network
+	Source      netip.Prefix
+	Destination netip.Prefix
 	Masquerade  bool
 	Inverse     bool
 }

client/firewall/manager/set.go

@@ -1,74 +0,0 @@
-package manager
-
-import (
-	"crypto/sha256"
-	"encoding/hex"
-	"fmt"
-	"net/netip"
-	"slices"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/management/domain"
-)
-
-type Set struct {
-	hash    [4]byte
-	comment string
-}
-
-// String returns the string representation of the set: hashed name and comment
-func (h Set) String() string {
-	if h.comment == "" {
-		return h.HashedName()
-	}
-	return h.HashedName() + ": " + h.comment
-}
-
-// HashedName returns the string representation of the hash
-func (h Set) HashedName() string {
-	return fmt.Sprintf(
-		"nb-%s",
-		hex.EncodeToString(h.hash[:]),
-	)
-}
-
-// Comment returns the comment of the set
-func (h Set) Comment() string {
-	return h.comment
-}
-
-// NewPrefixSet generates a unique name for an ipset based on the given prefixes.
-func NewPrefixSet(prefixes []netip.Prefix) Set {
-	// sort for consistent naming
-	SortPrefixes(prefixes)
-
-	hash := sha256.New()
-	for _, src := range prefixes {
-		bytes, err := src.MarshalBinary()
-		if err != nil {
-			log.Warnf("failed to marshal prefix %s: %v", src, err)
-		}
-		hash.Write(bytes)
-	}
-	var set Set
-	copy(set.hash[:], hash.Sum(nil)[:4])
-
-	return set
-}
-
-// NewDomainSet generates a unique name for an ipset based on the given domains.
-func NewDomainSet(domains domain.List) Set {
-	slices.Sort(domains)
-
-	hash := sha256.New()
-	for _, d := range domains {
-		hash.Write([]byte(d.PunycodeString()))
-	}
-	set := Set{
-		comment: domains.SafeString(),
-	}
-	copy(set.hash[:], hash.Sum(nil)[:4])
-
-	return set
-}

client/firewall/nftables/manager_linux.go

@@ -135,16 +135,17 @@ func (m *Manager) AddPeerFiltering(
 func (m *Manager) AddRouteFiltering(
 	id []byte,
 	sources []netip.Prefix,
-	destination firewall.Network,
+	destination netip.Prefix,
 	proto firewall.Protocol,
-	sPort, dPort *firewall.Port,
+	sPort *firewall.Port,
+	dPort *firewall.Port,
 	action firewall.Action,
 ) (firewall.Rule, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if destination.IsPrefix() && !destination.Prefix.Addr().Is4() {
-		return nil, fmt.Errorf("unsupported IP version: %s", destination.Prefix.Addr().String())
+	if !destination.Addr().Is4() {
+		return nil, fmt.Errorf("unsupported IP version: %s", destination.Addr().String())
 	}
 
 	return m.router.AddRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
@@ -241,7 +242,7 @@ func (m *Manager) SetLegacyManagement(isLegacy bool) error {
 	return firewall.SetLegacyManagement(m.router, isLegacy)
 }
 
-// Close closes the firewall manager
+// Reset firewall to the default state
 func (m *Manager) Close(stateManager *statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
@@ -358,14 +359,6 @@ func (m *Manager) DeleteDNATRule(rule firewall.Rule) error {
 	return m.router.DeleteDNATRule(rule)
 }
 
-// UpdateSet updates the set with the given prefixes
-func (m *Manager) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.UpdateSet(set, prefixes)
-}
-
 func (m *Manager) createWorkTable() (*nftables.Table, error) {
 	tables, err := m.rConn.ListTablesOfFamily(nftables.TableFamilyIPv4)
 	if err != nil {

client/firewall/nftables/manager_linux_test.go

@@ -289,7 +289,7 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 	_, err = manager.AddRouteFiltering(
 		nil,
 		[]netip.Prefix{netip.MustParsePrefix("192.168.2.0/24")},
-		fw.Network{Prefix: netip.MustParsePrefix("10.1.0.0/24")},
+		netip.MustParsePrefix("10.1.0.0/24"),
 		fw.ProtocolTCP,
 		nil,
 		&fw.Port{Values: []uint16{443}},
@@ -298,8 +298,8 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 	require.NoError(t, err, "failed to add route filtering rule")
 
 	pair := fw.RouterPair{
-		Source:      fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-		Destination: fw.Network{Prefix: netip.MustParsePrefix("10.0.0.0/24")},
+		Source:      netip.MustParsePrefix("192.168.1.0/24"),
+		Destination: netip.MustParsePrefix("10.0.0.0/24"),
 		Masquerade:  true,
 	}
 	err = manager.AddNatRule(pair)

client/firewall/nftables/router_linux.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/coreos/go-iptables/iptables"
+	"github.com/davecgh/go-spew/spew"
 	"github.com/google/nftables"
 	"github.com/google/nftables/binaryutil"
 	"github.com/google/nftables/expr"
@@ -43,14 +44,9 @@ const (
 const refreshRulesMapError = "refresh rules map: %w"
 
 var (
-	errFilterTableNotFound = fmt.Errorf("'filter' table not found")
+	errFilterTableNotFound = fmt.Errorf("nftables: 'filter' table not found")
 )
 
-type setInput struct {
-	set      firewall.Set
-	prefixes []netip.Prefix
-}
-
 type router struct {
 	conn        *nftables.Conn
 	workTable   *nftables.Table
@@ -58,7 +54,7 @@ type router struct {
 	chains      map[string]*nftables.Chain
 	// rules is useful to avoid duplicates and to get missing attributes that we don't have when adding new rules
 	rules        map[string]*nftables.Rule
-	ipsetCounter *refcounter.Counter[string, setInput, *nftables.Set]
+	ipsetCounter *refcounter.Counter[string, []netip.Prefix, *nftables.Set]
 
 	wgIface          iFaceMapper
 	ipFwdState       *ipfwdstate.IPForwardingState
@@ -167,7 +163,7 @@ func (r *router) removeNatPreroutingRules() error {
 func (r *router) loadFilterTable() (*nftables.Table, error) {
 	tables, err := r.conn.ListTablesOfFamily(nftables.TableFamilyIPv4)
 	if err != nil {
-		return nil, fmt.Errorf("unable to list tables: %v", err)
+		return nil, fmt.Errorf("nftables: unable to list tables: %v", err)
 	}
 
 	for _, table := range tables {
@@ -320,7 +316,7 @@ func (r *router) setupDataPlaneMark() error {
 func (r *router) AddRouteFiltering(
 	id []byte,
 	sources []netip.Prefix,
-	destination firewall.Network,
+	destination netip.Prefix,
 	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
@@ -335,29 +331,23 @@ func (r *router) AddRouteFiltering(
 	chain := r.chains[chainNameRoutingFw]
 	var exprs []expr.Any
 
-	var source firewall.Network
 	switch {
 	case len(sources) == 1 && sources[0].Bits() == 0:
 		// If it's 0.0.0.0/0, we don't need to add any source matching
 	case len(sources) == 1:
 		// If there's only one source, we can use it directly
-		source.Prefix = sources[0]
+		exprs = append(exprs, generateCIDRMatcherExpressions(true, sources[0])...)
 	default:
-		// If there are multiple sources, use a set
-		source.Set = firewall.NewPrefixSet(sources)
+		// If there are multiple sources, create or get an ipset
+		var err error
+		exprs, err = r.getIpSetExprs(sources, exprs)
+		if err != nil {
+			return nil, fmt.Errorf("get ipset expressions: %w", err)
+		}
 	}
 
-	sourceExp, err := r.applyNetwork(source, sources, true)
-	if err != nil {
-		return nil, fmt.Errorf("apply source: %w", err)
-	}
-	exprs = append(exprs, sourceExp...)
-
-	destExp, err := r.applyNetwork(destination, nil, false)
-	if err != nil {
-		return nil, fmt.Errorf("apply destination: %w", err)
-	}
-	exprs = append(exprs, destExp...)
+	// Handle destination
+	exprs = append(exprs, generateCIDRMatcherExpressions(false, destination)...)
 
 	// Handle protocol
 	if proto != firewall.ProtocolALL {
@@ -401,27 +391,39 @@ func (r *router) AddRouteFiltering(
 		rule = r.conn.AddRule(rule)
 	}
 
+	log.Tracef("Adding route rule %s", spew.Sdump(rule))
 	if err := r.conn.Flush(); err != nil {
 		return nil, fmt.Errorf(flushError, err)
 	}
 
 	r.rules[string(ruleKey)] = rule
 
-	log.Debugf("added route rule: sources=%v, destination=%v, proto=%v, sPort=%v, dPort=%v, action=%v", sources, destination, proto, sPort, dPort, action)
+	log.Debugf("nftables: added route rule: sources=%v, destination=%v, proto=%v, sPort=%v, dPort=%v, action=%v", sources, destination, proto, sPort, dPort, action)
 
 	return ruleKey, nil
 }
 
-func (r *router) getIpSet(set firewall.Set, prefixes []netip.Prefix, isSource bool) ([]expr.Any, error) {
-	ref, err := r.ipsetCounter.Increment(set.HashedName(), setInput{
-		set:      set,
-		prefixes: prefixes,
-	})
+func (r *router) getIpSetExprs(sources []netip.Prefix, exprs []expr.Any) ([]expr.Any, error) {
+	setName := firewall.GenerateSetName(sources)
+	ref, err := r.ipsetCounter.Increment(setName, sources)
 	if err != nil {
-		return nil, fmt.Errorf("create or get ipset: %w", err)
+		return nil, fmt.Errorf("create or get ipset for sources: %w", err)
 	}
 
-	return getIpSetExprs(ref, isSource)
+	exprs = append(exprs,
+		&expr.Payload{
+			DestRegister: 1,
+			Base:         expr.PayloadBaseNetworkHeader,
+			Offset:       12,
+			Len:          4,
+		},
+		&expr.Lookup{
+			SourceRegister: 1,
+			SetName:        ref.Out.Name,
+			SetID:          ref.Out.ID,
+		},
+	)
+	return exprs, nil
 }
 
 func (r *router) DeleteRouteRule(rule firewall.Rule) error {
@@ -440,54 +442,42 @@ func (r *router) DeleteRouteRule(rule firewall.Rule) error {
 		return fmt.Errorf("route rule %s has no handle", ruleKey)
 	}
 
+	setName := r.findSetNameInRule(nftRule)
+
 	if err := r.deleteNftRule(nftRule, ruleKey); err != nil {
 		return fmt.Errorf("delete: %w", err)
 	}
 
+	if setName != "" {
+		if _, err := r.ipsetCounter.Decrement(setName); err != nil {
+			return fmt.Errorf("decrement ipset reference: %w", err)
+		}
+	}
+
 	if err := r.conn.Flush(); err != nil {
 		return fmt.Errorf(flushError, err)
 	}
 
-	if err := r.decrementSetCounter(nftRule); err != nil {
-		return fmt.Errorf("decrement set counter: %w", err)
-	}
-
 	return nil
 }
 
-func (r *router) createIpSet(setName string, input setInput) (*nftables.Set, error) {
+func (r *router) createIpSet(setName string, sources []netip.Prefix) (*nftables.Set, error) {
 	// overlapping prefixes will result in an error, so we need to merge them
-	prefixes := firewall.MergeIPRanges(input.prefixes)
+	sources = firewall.MergeIPRanges(sources)
 
-	nfset := &nftables.Set{
-		Name:    setName,
-		Comment: input.set.Comment(),
-		Table:   r.workTable,
+	set := &nftables.Set{
+		Name:  setName,
+		Table: r.workTable,
 		// required for prefixes
 		Interval: true,
 		KeyType:  nftables.TypeIPAddr,
 	}
 
-	elements := convertPrefixesToSet(prefixes)
-	if err := r.conn.AddSet(nfset, elements); err != nil {
-		return nil, fmt.Errorf("error adding elements to set %s: %w", setName, err)
-	}
-
-	if err := r.conn.Flush(); err != nil {
-		return nil, fmt.Errorf("flush error: %w", err)
-	}
-
-	log.Printf("Created new ipset: %s with %d elements", setName, len(elements)/2)
-
-	return nfset, nil
-}
-
-func convertPrefixesToSet(prefixes []netip.Prefix) []nftables.SetElement {
 	var elements []nftables.SetElement
-	for _, prefix := range prefixes {
+	for _, prefix := range sources {
 		// TODO: Implement IPv6 support
 		if prefix.Addr().Is6() {
-			log.Tracef("skipping IPv6 prefix %s: IPv6 support not yet implemented", prefix)
+			log.Printf("Skipping IPv6 prefix %s: IPv6 support not yet implemented", prefix)
 			continue
 		}
 
@@ -503,7 +493,18 @@ func convertPrefixesToSet(prefixes []netip.Prefix) []nftables.SetElement {
 			nftables.SetElement{Key: lastIP.AsSlice(), IntervalEnd: true},
 		)
 	}
-	return elements
+
+	if err := r.conn.AddSet(set, elements); err != nil {
+		return nil, fmt.Errorf("error adding elements to set %s: %w", setName, err)
+	}
+
+	if err := r.conn.Flush(); err != nil {
+		return nil, fmt.Errorf("flush error: %w", err)
+	}
+
+	log.Printf("Created new ipset: %s with %d elements", setName, len(elements)/2)
+
+	return set, nil
 }
 
 // calculateLastIP determines the last IP in a given prefix.
@@ -527,8 +528,8 @@ func uint32ToBytes(ip uint32) [4]byte {
 	return b
 }
 
-func (r *router) deleteIpSet(setName string, nfset *nftables.Set) error {
-	r.conn.DelSet(nfset)
+func (r *router) deleteIpSet(setName string, set *nftables.Set) error {
+	r.conn.DelSet(set)
 	if err := r.conn.Flush(); err != nil {
 		return fmt.Errorf(flushError, err)
 	}
@@ -537,27 +538,13 @@ func (r *router) deleteIpSet(setName string, nfset *nftables.Set) error {
 	return nil
 }
 
-func (r *router) decrementSetCounter(rule *nftables.Rule) error {
-	sets := r.findSets(rule)
-
-	var merr *multierror.Error
-	for _, setName := range sets {
-		if _, err := r.ipsetCounter.Decrement(setName); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("decrement set counter: %w", err))
-		}
-	}
-
-	return nberrors.FormatErrorOrNil(merr)
-}
-
-func (r *router) findSets(rule *nftables.Rule) []string {
-	var sets []string
+func (r *router) findSetNameInRule(rule *nftables.Rule) string {
 	for _, e := range rule.Exprs {
 		if lookup, ok := e.(*expr.Lookup); ok {
-			sets = append(sets, lookup.SetName)
+			return lookup.SetName
 		}
 	}
-	return sets
+	return ""
 }
 
 func (r *router) deleteNftRule(rule *nftables.Rule, ruleKey string) error {
@@ -599,8 +586,7 @@ func (r *router) AddNatRule(pair firewall.RouterPair) error {
 	}
 
 	if err := r.conn.Flush(); err != nil {
-		// TODO: rollback ipset counter
-		return fmt.Errorf("insert rules for %s: %v", pair.Destination, err)
+		return fmt.Errorf("nftables: insert rules for %s: %v", pair.Destination, err)
 	}
 
 	return nil
@@ -608,22 +594,19 @@ func (r *router) AddNatRule(pair firewall.RouterPair) error {
 
 // addNatRule inserts a nftables rule to the conn client flush queue
 func (r *router) addNatRule(pair firewall.RouterPair) error {
-	sourceExp, err := r.applyNetwork(pair.Source, nil, true)
-	if err != nil {
-		return fmt.Errorf("apply source: %w", err)
-	}
-
-	destExp, err := r.applyNetwork(pair.Destination, nil, false)
-	if err != nil {
-		return fmt.Errorf("apply destination: %w", err)
-	}
+	sourceExp := generateCIDRMatcherExpressions(true, pair.Source)
+	destExp := generateCIDRMatcherExpressions(false, pair.Destination)
 
 	op := expr.CmpOpEq
 	if pair.Inverse {
 		op = expr.CmpOpNeq
 	}
 
-	exprs := []expr.Any{
+	// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
+	// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
+	exprs := getCtNewExprs()
+	exprs = append(exprs,
+		// interface matching
 		&expr.Meta{
 			Key:      expr.MetaKeyIIFNAME,
 			Register: 1,
@@ -633,10 +616,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 			Register: 1,
 			Data:     ifname(r.wgIface.Name()),
 		},
-	}
-	// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
-	// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
-	exprs = append(exprs, getCtNewExprs()...)
+	)
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
@@ -666,9 +646,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 		}
 	}
 
-	// Ensure nat rules come first, so the mark can be overwritten.
-	// Currently overwritten by the dst-type LOCAL rules for redirected traffic.
-	r.rules[ruleKey] = r.conn.InsertRule(&nftables.Rule{
+	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
 		Chain:    r.chains[chainNameManglePrerouting],
 		Exprs:    exprs,
@@ -751,15 +729,8 @@ func (r *router) addPostroutingRules() error {
 
 // addLegacyRouteRule adds a legacy routing rule for mgmt servers pre route acls
 func (r *router) addLegacyRouteRule(pair firewall.RouterPair) error {
-	sourceExp, err := r.applyNetwork(pair.Source, nil, true)
-	if err != nil {
-		return fmt.Errorf("apply source: %w", err)
-	}
-
-	destExp, err := r.applyNetwork(pair.Destination, nil, false)
-	if err != nil {
-		return fmt.Errorf("apply destination: %w", err)
-	}
+	sourceExp := generateCIDRMatcherExpressions(true, pair.Source)
+	destExp := generateCIDRMatcherExpressions(false, pair.Destination)
 
 	exprs := []expr.Any{
 		&expr.Counter{},
@@ -768,8 +739,7 @@ func (r *router) addLegacyRouteRule(pair firewall.RouterPair) error {
 		},
 	}
 
-	exprs = append(exprs, sourceExp...)
-	exprs = append(exprs, destExp...)
+	expression := append(sourceExp, append(destExp, exprs...)...) // nolint:gocritic
 
 	ruleKey := firewall.GenKey(firewall.ForwardingFormat, pair)
 
@@ -782,7 +752,7 @@ func (r *router) addLegacyRouteRule(pair firewall.RouterPair) error {
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
 		Chain:    r.chains[chainNameRoutingFw],
-		Exprs:    exprs,
+		Exprs:    expression,
 		UserData: []byte(ruleKey),
 	})
 	return nil
@@ -797,13 +767,11 @@ func (r *router) removeLegacyRouteRule(pair firewall.RouterPair) error {
 			return fmt.Errorf("remove legacy forwarding rule %s -> %s: %v", pair.Source, pair.Destination, err)
 		}
 
-		log.Debugf("removed legacy forwarding rule %s -> %s", pair.Source, pair.Destination)
+		log.Debugf("nftables: removed legacy forwarding rule %s -> %s", pair.Source, pair.Destination)
 
 		delete(r.rules, ruleKey)
-
-		if err := r.decrementSetCounter(rule); err != nil {
-			return fmt.Errorf("decrement set counter: %w", err)
-		}
+	} else {
+		log.Debugf("nftables: legacy forwarding rule %s not found", ruleKey)
 	}
 
 	return nil
@@ -1014,14 +982,12 @@ func (r *router) RemoveNatRule(pair firewall.RouterPair) error {
 		return fmt.Errorf(refreshRulesMapError, err)
 	}
 
-	if pair.Masquerade {
-		if err := r.removeNatRule(pair); err != nil {
-			return fmt.Errorf("remove prerouting rule: %w", err)
-		}
+	if err := r.removeNatRule(pair); err != nil {
+		return fmt.Errorf("remove prerouting rule: %w", err)
+	}
 
-		if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
-			return fmt.Errorf("remove inverse prerouting rule: %w", err)
-		}
+	if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
+		return fmt.Errorf("remove inverse prerouting rule: %w", err)
 	}
 
 	if err := r.removeLegacyRouteRule(pair); err != nil {
@@ -1029,10 +995,10 @@ func (r *router) RemoveNatRule(pair firewall.RouterPair) error {
 	}
 
 	if err := r.conn.Flush(); err != nil {
-		// TODO: rollback set counter
-		return fmt.Errorf("remove nat rules rule %s: %v", pair.Destination, err)
+		return fmt.Errorf("nftables: received error while applying rule removal for %s: %v", pair.Destination, err)
 	}
 
+	log.Debugf("nftables: removed nat rules for %s", pair.Destination)
 	return nil
 }
 
@@ -1040,19 +1006,16 @@ func (r *router) removeNatRule(pair firewall.RouterPair) error {
 	ruleKey := firewall.GenKey(firewall.PreroutingFormat, pair)
 
 	if rule, exists := r.rules[ruleKey]; exists {
-		if err := r.conn.DelRule(rule); err != nil {
+		err := r.conn.DelRule(rule)
+		if err != nil {
 			return fmt.Errorf("remove prerouting rule %s -> %s: %v", pair.Source, pair.Destination, err)
 		}
 
-		log.Debugf("removed prerouting rule %s -> %s", pair.Source, pair.Destination)
+		log.Debugf("nftables: removed prerouting rule %s -> %s", pair.Source, pair.Destination)
 
 		delete(r.rules, ruleKey)
-
-		if err := r.decrementSetCounter(rule); err != nil {
-			return fmt.Errorf("decrement set counter: %w", err)
-		}
 	} else {
-		log.Debugf("prerouting rule %s not found", ruleKey)
+		log.Debugf("nftables: prerouting rule %s not found", ruleKey)
 	}
 
 	return nil
@@ -1064,7 +1027,7 @@ func (r *router) refreshRulesMap() error {
 	for _, chain := range r.chains {
 		rules, err := r.conn.GetRules(chain.Table, chain)
 		if err != nil {
-			return fmt.Errorf(" unable to list rules: %v", err)
+			return fmt.Errorf("nftables: unable to list rules: %v", err)
 		}
 		for _, rule := range rules {
 			if len(rule.UserData) > 0 {
@@ -1338,54 +1301,13 @@ func (r *router) DeleteDNATRule(rule firewall.Rule) error {
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-func (r *router) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
-	nfset, err := r.conn.GetSetByName(r.workTable, set.HashedName())
-	if err != nil {
-		return fmt.Errorf("get set %s: %w", set.HashedName(), err)
-	}
-
-	elements := convertPrefixesToSet(prefixes)
-	if err := r.conn.SetAddElements(nfset, elements); err != nil {
-		return fmt.Errorf("add elements to set %s: %w", set.HashedName(), err)
-	}
-
-	if err := r.conn.Flush(); err != nil {
-		return fmt.Errorf(flushError, err)
-	}
-
-	log.Debugf("updated set %s with prefixes %v", set.HashedName(), prefixes)
-
-	return nil
-}
-
-// applyNetwork generates nftables expressions for networks (CIDR) or sets
-func (r *router) applyNetwork(
-	network firewall.Network,
-	setPrefixes []netip.Prefix,
-	isSource bool,
-) ([]expr.Any, error) {
-	if network.IsSet() {
-		exprs, err := r.getIpSet(network.Set, setPrefixes, isSource)
-		if err != nil {
-			return nil, fmt.Errorf("source: %w", err)
-		}
-		return exprs, nil
-	}
-
-	if network.IsPrefix() {
-		return applyPrefix(network.Prefix, isSource), nil
-	}
-
-	return nil, nil
-}
-
-// applyPrefix generates nftables expressions for a CIDR prefix
-func applyPrefix(prefix netip.Prefix, isSource bool) []expr.Any {
-	// dst offset
-	offset := uint32(16)
-	if isSource {
-		// src offset
-		offset = 12
+// generateCIDRMatcherExpressions generates nftables expressions that matches a CIDR
+func generateCIDRMatcherExpressions(source bool, prefix netip.Prefix) []expr.Any {
+	var offset uint32
+	if source {
+		offset = 12 // src offset
+	} else {
+		offset = 16 // dst offset
 	}
 
 	ones := prefix.Bits()
@@ -1493,27 +1415,3 @@ func getCtNewExprs() []expr.Any {
 		},
 	}
 }
-
-func getIpSetExprs(ref refcounter.Ref[*nftables.Set], isSource bool) ([]expr.Any, error) {
-
-	// dst offset
-	offset := uint32(16)
-	if isSource {
-		// src offset
-		offset = 12
-	}
-
-	return []expr.Any{
-		&expr.Payload{
-			DestRegister: 1,
-			Base:         expr.PayloadBaseNetworkHeader,
-			Offset:       offset,
-			Len:          4,
-		},
-		&expr.Lookup{
-			SourceRegister: 1,
-			SetName:        ref.Out.Name,
-			SetID:          ref.Out.ID,
-		},
-	}, nil
-}

client/firewall/nftables/router_linux_test.go

@@ -88,8 +88,8 @@ func TestNftablesManager_AddNatRule(t *testing.T) {
 				}
 
 				// Build CIDR matching expressions
-				sourceExp := applyPrefix(testCase.InputPair.Source.Prefix, true)
-				destExp := applyPrefix(testCase.InputPair.Destination.Prefix, false)
+				sourceExp := generateCIDRMatcherExpressions(true, testCase.InputPair.Source)
+				destExp := generateCIDRMatcherExpressions(false, testCase.InputPair.Destination)
 
 				// Combine all expressions in the correct order
 				// nolint:gocritic
@@ -311,7 +311,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, firewall.Network{Prefix: tt.destination}, tt.proto, tt.sPort, tt.dPort, tt.action)
+			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
 			require.NoError(t, err, "AddRouteFiltering failed")
 
 			t.Cleanup(func() {
@@ -441,8 +441,8 @@ func TestNftablesCreateIpSet(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			setName := firewall.NewPrefixSet(tt.sources).HashedName()
-			set, err := r.createIpSet(setName, setInput{prefixes: tt.sources})
+			setName := firewall.GenerateSetName(tt.sources)
+			set, err := r.createIpSet(setName, tt.sources)
 			if err != nil {
 				t.Logf("Failed to create IP set: %v", err)
 				printNftSets()

client/firewall/test/cases_linux.go

@@ -15,8 +15,8 @@ var (
 			Name: "Insert Forwarding IPV4 Rule",
 			InputPair: firewall.RouterPair{
 				ID:          "zxa",
-				Source:      firewall.Network{Prefix: netip.MustParsePrefix("100.100.100.1/32")},
-				Destination: firewall.Network{Prefix: netip.MustParsePrefix("100.100.200.0/24")},
+				Source:      netip.MustParsePrefix("100.100.100.1/32"),
+				Destination: netip.MustParsePrefix("100.100.200.0/24"),
 				Masquerade:  false,
 			},
 		},
@@ -24,8 +24,8 @@ var (
 			Name: "Insert Forwarding And Nat IPV4 Rules",
 			InputPair: firewall.RouterPair{
 				ID:          "zxa",
-				Source:      firewall.Network{Prefix: netip.MustParsePrefix("100.100.100.1/32")},
-				Destination: firewall.Network{Prefix: netip.MustParsePrefix("100.100.200.0/24")},
+				Source:      netip.MustParsePrefix("100.100.100.1/32"),
+				Destination: netip.MustParsePrefix("100.100.200.0/24"),
 				Masquerade:  true,
 			},
 		},
@@ -40,8 +40,8 @@ var (
 			Name: "Remove Forwarding And Nat IPV4 Rules",
 			InputPair: firewall.RouterPair{
 				ID:          "zxa",
-				Source:      firewall.Network{Prefix: netip.MustParsePrefix("100.100.100.1/32")},
-				Destination: firewall.Network{Prefix: netip.MustParsePrefix("100.100.200.0/24")},
+				Source:      netip.MustParsePrefix("100.100.100.1/32"),
+				Destination: netip.MustParsePrefix("100.100.200.0/24"),
 				Masquerade:  true,
 			},
 		},

client/firewall/uspfilter/allow_netbird.go

@@ -12,7 +12,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
-// Close cleans up the firewall manager by removing all rules and closing trackers
+// Reset firewall to the default state
 func (m *Manager) Close(stateManager *statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()

client/firewall/uspfilter/allow_netbird_windows.go

@@ -10,6 +10,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
@@ -21,7 +22,7 @@ const (
 	firewallRuleName        = "Netbird"
 )
 
-// Close cleans up the firewall manager by removing all rules and closing trackers
+// Reset firewall to the default state
 func (m *Manager) Close(*statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
@@ -31,14 +32,17 @@ func (m *Manager) Close(*statemanager.Manager) error {
 
 	if m.udpTracker != nil {
 		m.udpTracker.Close()
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger, m.flowLogger)
 	}
 
 	if m.icmpTracker != nil {
 		m.icmpTracker.Close()
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger, m.flowLogger)
 	}
 
 	if m.tcpTracker != nil {
 		m.tcpTracker.Close()
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger, m.flowLogger)
 	}
 
 	if fwder := m.forwarder.Load(); fwder != nil {

client/firewall/uspfilter/forwarder/forwarder.go

@@ -4,9 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"net/netip"
 	"runtime"
-	"sync"
 
 	log "github.com/sirupsen/logrus"
 	"gvisor.dev/gvisor/pkg/buffer"
@@ -19,7 +17,6 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
 
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/common"
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
@@ -32,10 +29,8 @@ const (
 )
 
 type Forwarder struct {
-	logger     *nblog.Logger
-	flowLogger nftypes.FlowLogger
-	// ruleIdMap is used to store the rule ID for a given connection
-	ruleIdMap    sync.Map
+	logger       *nblog.Logger
+	flowLogger   nftypes.FlowLogger
 	stack        *stack.Stack
 	endpoint     *endpoint
 	udpForwarder *udpForwarder
@@ -172,35 +167,3 @@ func (f *Forwarder) determineDialAddr(addr tcpip.Address) net.IP {
 	}
 	return addr.AsSlice()
 }
-
-func (f *Forwarder) RegisterRuleID(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, ruleID []byte) {
-	key := buildKey(srcIP, dstIP, srcPort, dstPort)
-	f.ruleIdMap.LoadOrStore(key, ruleID)
-}
-
-func (f *Forwarder) getRuleID(srcIP, dstIP netip.Addr, srcPort, dstPort uint16) ([]byte, bool) {
-
-	if value, ok := f.ruleIdMap.Load(buildKey(srcIP, dstIP, srcPort, dstPort)); ok {
-		return value.([]byte), true
-	} else if value, ok := f.ruleIdMap.Load(buildKey(dstIP, srcIP, dstPort, srcPort)); ok {
-		return value.([]byte), true
-	}
-
-	return nil, false
-}
-
-func (f *Forwarder) DeleteRuleID(srcIP, dstIP netip.Addr, srcPort, dstPort uint16) {
-	if _, ok := f.ruleIdMap.LoadAndDelete(buildKey(srcIP, dstIP, srcPort, dstPort)); ok {
-		return
-	}
-	f.ruleIdMap.LoadAndDelete(buildKey(dstIP, srcIP, dstPort, srcPort))
-}
-
-func buildKey(srcIP, dstIP netip.Addr, srcPort, dstPort uint16) conntrack.ConnKey {
-	return conntrack.ConnKey{
-		SrcIP:   srcIP,
-		DstIP:   dstIP,
-		SrcPort: srcPort,
-		DstPort: dstPort,
-	}
-}

client/firewall/uspfilter/forwarder/icmp.go

@@ -25,7 +25,7 @@ func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBuf
 	}
 
 	flowID := uuid.New()
-	f.sendICMPEvent(nftypes.TypeStart, flowID, id, icmpType, icmpCode, 0, 0)
+	f.sendICMPEvent(nftypes.TypeStart, flowID, id, icmpType, icmpCode)
 
 	ctx, cancel := context.WithTimeout(f.ctx, 5*time.Second)
 	defer cancel()
@@ -34,14 +34,14 @@ func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBuf
 	// TODO: support non-root
 	conn, err := lc.ListenPacket(ctx, "ip4:icmp", "0.0.0.0")
 	if err != nil {
-		f.logger.Error("forwarder: Failed to create ICMP socket for %v: %v", epID(id), err)
+		f.logger.Error("Failed to create ICMP socket for %v: %v", epID(id), err)
 
 		// This will make netstack reply on behalf of the original destination, that's ok for now
 		return false
 	}
 	defer func() {
 		if err := conn.Close(); err != nil {
-			f.logger.Debug("forwarder: Failed to close ICMP socket: %v", err)
+			f.logger.Debug("Failed to close ICMP socket: %v", err)
 		}
 	}()
 
@@ -52,37 +52,36 @@ func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBuf
 	payload := fullPacket.AsSlice()
 
 	if _, err = conn.WriteTo(payload, dst); err != nil {
-		f.logger.Error("forwarder: Failed to write ICMP packet for %v: %v", epID(id), err)
+		f.logger.Error("Failed to write ICMP packet for %v: %v", epID(id), err)
 		return true
 	}
 
-	f.logger.Trace("forwarder: Forwarded ICMP packet %v type %v code %v",
+	f.logger.Trace("Forwarded ICMP packet %v type %v code %v",
 		epID(id), icmpHdr.Type(), icmpHdr.Code())
 
 	// For Echo Requests, send and handle response
 	if header.ICMPv4Type(icmpType) == header.ICMPv4Echo {
-		rxBytes := pkt.Size()
-		txBytes := f.handleEchoResponse(icmpHdr, conn, id)
-		f.sendICMPEvent(nftypes.TypeEnd, flowID, id, icmpType, icmpCode, uint64(rxBytes), uint64(txBytes))
+		f.handleEchoResponse(icmpHdr, conn, id)
+		f.sendICMPEvent(nftypes.TypeEnd, flowID, id, icmpType, icmpCode)
 	}
 
 	// For other ICMP types (Time Exceeded, Destination Unreachable, etc) do nothing
 	return true
 }
 
-func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, conn net.PacketConn, id stack.TransportEndpointID) int {
+func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, conn net.PacketConn, id stack.TransportEndpointID) {
 	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
-		f.logger.Error("forwarder: Failed to set read deadline for ICMP response: %v", err)
-		return 0
+		f.logger.Error("Failed to set read deadline for ICMP response: %v", err)
+		return
 	}
 
 	response := make([]byte, f.endpoint.mtu)
 	n, _, err := conn.ReadFrom(response)
 	if err != nil {
 		if !isTimeout(err) {
-			f.logger.Error("forwarder: Failed to read ICMP response: %v", err)
+			f.logger.Error("Failed to read ICMP response: %v", err)
 		}
-		return 0
+		return
 	}
 
 	ipHdr := make([]byte, header.IPv4MinimumSize)
@@ -101,54 +100,28 @@ func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, conn net.PacketCon
 	fullPacket = append(fullPacket, response[:n]...)
 
 	if err := f.InjectIncomingPacket(fullPacket); err != nil {
-		f.logger.Error("forwarder: Failed to inject ICMP response: %v", err)
+		f.logger.Error("Failed to inject ICMP response: %v", err)
 
-		return 0
+		return
 	}
 
-	f.logger.Trace("forwarder: Forwarded ICMP echo reply for %v type %v code %v",
+	f.logger.Trace("Forwarded ICMP echo reply for %v type %v code %v",
 		epID(id), icmpHdr.Type(), icmpHdr.Code())
-
-	return len(fullPacket)
 }
 
 // sendICMPEvent stores flow events for ICMP packets
-func (f *Forwarder) sendICMPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, icmpType, icmpCode uint8, rxBytes, txBytes uint64) {
-	var rxPackets, txPackets uint64
-	if rxBytes > 0 {
-		rxPackets = 1
-	}
-	if txBytes > 0 {
-		txPackets = 1
-	}
-
-	srcIp := netip.AddrFrom4(id.RemoteAddress.As4())
-	dstIp := netip.AddrFrom4(id.LocalAddress.As4())
-
-	fields := nftypes.EventFields{
+func (f *Forwarder) sendICMPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, icmpType, icmpCode uint8) {
+	f.flowLogger.StoreEvent(nftypes.EventFields{
 		FlowID:    flowID,
 		Type:      typ,
 		Direction: nftypes.Ingress,
 		Protocol:  nftypes.ICMP,
 		// TODO: handle ipv6
-		SourceIP: srcIp,
-		DestIP:   dstIp,
+		SourceIP: netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:   netip.AddrFrom4(id.LocalAddress.As4()),
 		ICMPType: icmpType,
 		ICMPCode: icmpCode,
 
-		RxBytes:   rxBytes,
-		TxBytes:   txBytes,
-		RxPackets: rxPackets,
-		TxPackets: txPackets,
-	}
-
-	if typ == nftypes.TypeStart {
-		if ruleId, ok := f.getRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort); ok {
-			fields.RuleID = ruleId
-		}
-	} else {
-		f.DeleteRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort)
-	}
-
-	f.flowLogger.StoreEvent(fields)
+		// TODO: get packets/bytes
+	})
 }

client/firewall/uspfilter/forwarder/tcp.go

@@ -6,10 +6,8 @@ import (
 	"io"
 	"net"
 	"net/netip"
-	"sync"
 
 	"github.com/google/uuid"
-
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
@@ -25,11 +23,11 @@ func (f *Forwarder) handleTCP(r *tcp.ForwarderRequest) {
 
 	flowID := uuid.New()
 
-	f.sendTCPEvent(nftypes.TypeStart, flowID, id, 0, 0, 0, 0)
+	f.sendTCPEvent(nftypes.TypeStart, flowID, id, nil)
 	var success bool
 	defer func() {
 		if !success {
-			f.sendTCPEvent(nftypes.TypeEnd, flowID, id, 0, 0, 0, 0)
+			f.sendTCPEvent(nftypes.TypeEnd, flowID, id, nil)
 		}
 	}()
 
@@ -67,97 +65,67 @@ func (f *Forwarder) handleTCP(r *tcp.ForwarderRequest) {
 }
 
 func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn, outConn net.Conn, ep tcpip.Endpoint, flowID uuid.UUID) {
+	defer func() {
+		if err := inConn.Close(); err != nil {
+			f.logger.Debug("forwarder: inConn close error: %v", err)
+		}
+		if err := outConn.Close(); err != nil {
+			f.logger.Debug("forwarder: outConn close error: %v", err)
+		}
+		ep.Close()
 
+		f.sendTCPEvent(nftypes.TypeEnd, flowID, id, ep)
+	}()
+
+	// Create context for managing the proxy goroutines
 	ctx, cancel := context.WithCancel(f.ctx)
 	defer cancel()
 
-	go func() {
-		<-ctx.Done()
-		// Close connections and endpoint.
-		if err := inConn.Close(); err != nil && !isClosedError(err) {
-			f.logger.Debug("forwarder: inConn close error: %v", err)
-		}
-		if err := outConn.Close(); err != nil && !isClosedError(err) {
-			f.logger.Debug("forwarder: outConn close error: %v", err)
-		}
-
-		ep.Close()
-	}()
-
-	var wg sync.WaitGroup
-	wg.Add(2)
-
-	var (
-		bytesFromInToOut int64 // bytes from client to server (tx for client)
-		bytesFromOutToIn int64 // bytes from server to client (rx for client)
-		errInToOut       error
-		errOutToIn       error
-	)
+	errChan := make(chan error, 2)
 
 	go func() {
-		bytesFromInToOut, errInToOut = io.Copy(outConn, inConn)
-		cancel()
-		wg.Done()
+		_, err := io.Copy(outConn, inConn)
+		errChan <- err
 	}()
 
 	go func() {
-
-		bytesFromOutToIn, errOutToIn = io.Copy(inConn, outConn)
-		cancel()
-		wg.Done()
+		_, err := io.Copy(inConn, outConn)
+		errChan <- err
 	}()
 
-	wg.Wait()
-
-	if errInToOut != nil {
-		if !isClosedError(errInToOut) {
-			f.logger.Error("proxyTCP: copy error (in -> out): %v", errInToOut)
+	select {
+	case <-ctx.Done():
+		f.logger.Trace("forwarder: tearing down TCP connection %v due to context done", epID(id))
+		return
+	case err := <-errChan:
+		if err != nil && !isClosedError(err) {
+			f.logger.Error("proxyTCP: copy error: %v", err)
 		}
+		f.logger.Trace("forwarder: tearing down TCP connection %v", epID(id))
+		return
 	}
-	if errOutToIn != nil {
-		if !isClosedError(errOutToIn) {
-			f.logger.Error("proxyTCP: copy error (out -> in): %v", errOutToIn)
-		}
-	}
-
-	var rxPackets, txPackets uint64
-	if tcpStats, ok := ep.Stats().(*tcp.Stats); ok {
-		// fields are flipped since this is the in conn
-		rxPackets = tcpStats.SegmentsSent.Value()
-		txPackets = tcpStats.SegmentsReceived.Value()
-	}
-
-	f.logger.Trace("forwarder: Removed TCP connection %s [in: %d Pkts/%d B, out: %d Pkts/%d B]", epID(id), rxPackets, bytesFromOutToIn, txPackets, bytesFromInToOut)
-
-	f.sendTCPEvent(nftypes.TypeEnd, flowID, id, uint64(bytesFromOutToIn), uint64(bytesFromInToOut), rxPackets, txPackets)
 }
 
-func (f *Forwarder) sendTCPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, rxBytes, txBytes, rxPackets, txPackets uint64) {
-	srcIp := netip.AddrFrom4(id.RemoteAddress.As4())
-	dstIp := netip.AddrFrom4(id.LocalAddress.As4())
-
+func (f *Forwarder) sendTCPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, ep tcpip.Endpoint) {
 	fields := nftypes.EventFields{
 		FlowID:    flowID,
 		Type:      typ,
 		Direction: nftypes.Ingress,
 		Protocol:  nftypes.TCP,
 		// TODO: handle ipv6
-		SourceIP:   srcIp,
-		DestIP:     dstIp,
+		SourceIP:   netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:     netip.AddrFrom4(id.LocalAddress.As4()),
 		SourcePort: id.RemotePort,
 		DestPort:   id.LocalPort,
-		RxBytes:    rxBytes,
-		TxBytes:    txBytes,
-		RxPackets:  rxPackets,
-		TxPackets:  txPackets,
 	}
 
-	if typ == nftypes.TypeStart {
-		if ruleId, ok := f.getRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort); ok {
-			fields.RuleID = ruleId
+	if ep != nil {
+		if tcpStats, ok := ep.Stats().(*tcp.Stats); ok {
+			// fields are flipped since this is the in conn
+			// TODO: get bytes
+			fields.RxPackets = tcpStats.SegmentsSent.Value()
+			fields.TxPackets = tcpStats.SegmentsReceived.Value()
 		}
-	} else {
-		f.DeleteRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort)
 	}
 
 	f.flowLogger.StoreEvent(fields)

client/firewall/uspfilter/forwarder/udp.go

@@ -149,11 +149,11 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 
 	flowID := uuid.New()
 
-	f.sendUDPEvent(nftypes.TypeStart, flowID, id, 0, 0, 0, 0)
+	f.sendUDPEvent(nftypes.TypeStart, flowID, id, nil)
 	var success bool
 	defer func() {
 		if !success {
-			f.sendUDPEvent(nftypes.TypeEnd, flowID, id, 0, 0, 0, 0)
+			f.sendUDPEvent(nftypes.TypeEnd, flowID, id, nil)
 		}
 	}()
 
@@ -199,6 +199,7 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 		if err := outConn.Close(); err != nil {
 			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
+
 		return
 	}
 	f.udpForwarder.conns[id] = pConn
@@ -211,94 +212,68 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 }
 
 func (f *Forwarder) proxyUDP(ctx context.Context, pConn *udpPacketConn, id stack.TransportEndpointID, ep tcpip.Endpoint) {
-
-	ctx, cancel := context.WithCancel(f.ctx)
-	defer cancel()
-
-	go func() {
-		<-ctx.Done()
-
+	defer func() {
 		pConn.cancel()
-		if err := pConn.conn.Close(); err != nil && !isClosedError(err) {
+		if err := pConn.conn.Close(); err != nil {
 			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", epID(id), err)
 		}
-		if err := pConn.outConn.Close(); err != nil && !isClosedError(err) {
+		if err := pConn.outConn.Close(); err != nil {
 			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
 
 		ep.Close()
+
+		f.udpForwarder.Lock()
+		delete(f.udpForwarder.conns, id)
+		f.udpForwarder.Unlock()
+
+		f.sendUDPEvent(nftypes.TypeEnd, pConn.flowID, id, ep)
 	}()
 
-	var wg sync.WaitGroup
-	wg.Add(2)
+	errChan := make(chan error, 2)
 
-	var txBytes, rxBytes int64
-	var outboundErr, inboundErr error
-
-	// outbound->inbound: copy from pConn.conn to pConn.outConn
 	go func() {
-		defer wg.Done()
-		txBytes, outboundErr = pConn.copy(ctx, pConn.conn, pConn.outConn, &f.udpForwarder.bufPool, "outbound->inbound")
+		errChan <- pConn.copy(ctx, pConn.conn, pConn.outConn, &f.udpForwarder.bufPool, "outbound->inbound")
 	}()
 
-	// inbound->outbound: copy from pConn.outConn to pConn.conn
 	go func() {
-		defer wg.Done()
-		rxBytes, inboundErr = pConn.copy(ctx, pConn.outConn, pConn.conn, &f.udpForwarder.bufPool, "inbound->outbound")
+		errChan <- pConn.copy(ctx, pConn.outConn, pConn.conn, &f.udpForwarder.bufPool, "inbound->outbound")
 	}()
 
-	wg.Wait()
-
-	if outboundErr != nil && !isClosedError(outboundErr) {
-		f.logger.Error("proxyUDP: copy error (outbound->inbound): %v", outboundErr)
+	select {
+	case <-ctx.Done():
+		f.logger.Trace("forwarder: tearing down UDP connection %v due to context done", epID(id))
+		return
+	case err := <-errChan:
+		if err != nil && !isClosedError(err) {
+			f.logger.Error("proxyUDP: copy error: %v", err)
+		}
+		f.logger.Trace("forwarder: tearing down UDP connection %v", epID(id))
+		return
 	}
-	if inboundErr != nil && !isClosedError(inboundErr) {
-		f.logger.Error("proxyUDP: copy error (inbound->outbound): %v", inboundErr)
-	}
-
-	var rxPackets, txPackets uint64
-	if udpStats, ok := ep.Stats().(*tcpip.TransportEndpointStats); ok {
-		// fields are flipped since this is the in conn
-		rxPackets = udpStats.PacketsSent.Value()
-		txPackets = udpStats.PacketsReceived.Value()
-	}
-
-	f.logger.Trace("forwarder: Removed UDP connection %s [in: %d Pkts/%d B, out: %d Pkts/%d B]", epID(id), rxPackets, rxBytes, txPackets, txBytes)
-
-	f.udpForwarder.Lock()
-	delete(f.udpForwarder.conns, id)
-	f.udpForwarder.Unlock()
-
-	f.sendUDPEvent(nftypes.TypeEnd, pConn.flowID, id, uint64(rxBytes), uint64(txBytes), rxPackets, txPackets)
 }
 
 // sendUDPEvent stores flow events for UDP connections
-func (f *Forwarder) sendUDPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, rxBytes, txBytes, rxPackets, txPackets uint64) {
-	srcIp := netip.AddrFrom4(id.RemoteAddress.As4())
-	dstIp := netip.AddrFrom4(id.LocalAddress.As4())
-
+func (f *Forwarder) sendUDPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, ep tcpip.Endpoint) {
 	fields := nftypes.EventFields{
 		FlowID:    flowID,
 		Type:      typ,
 		Direction: nftypes.Ingress,
 		Protocol:  nftypes.UDP,
 		// TODO: handle ipv6
-		SourceIP:   srcIp,
-		DestIP:     dstIp,
+		SourceIP:   netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:     netip.AddrFrom4(id.LocalAddress.As4()),
 		SourcePort: id.RemotePort,
 		DestPort:   id.LocalPort,
-		RxBytes:    rxBytes,
-		TxBytes:    txBytes,
-		RxPackets:  rxPackets,
-		TxPackets:  txPackets,
 	}
 
-	if typ == nftypes.TypeStart {
-		if ruleId, ok := f.getRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort); ok {
-			fields.RuleID = ruleId
+	if ep != nil {
+		if tcpStats, ok := ep.Stats().(*tcpip.TransportEndpointStats); ok {
+			// fields are flipped since this is the in conn
+			// TODO: get bytes
+			fields.RxPackets = tcpStats.PacketsSent.Value()
+			fields.TxPackets = tcpStats.PacketsReceived.Value()
 		}
-	} else {
-		f.DeleteRuleID(srcIp, dstIp, id.RemotePort, id.LocalPort)
 	}
 
 	f.flowLogger.StoreEvent(fields)
@@ -313,20 +288,18 @@ func (c *udpPacketConn) getIdleDuration() time.Duration {
 	return time.Since(lastSeen)
 }
 
-// copy reads from src and writes to dst.
-func (c *udpPacketConn) copy(ctx context.Context, dst net.Conn, src net.Conn, bufPool *sync.Pool, direction string) (int64, error) {
+func (c *udpPacketConn) copy(ctx context.Context, dst net.Conn, src net.Conn, bufPool *sync.Pool, direction string) error {
 	bufp := bufPool.Get().(*[]byte)
 	defer bufPool.Put(bufp)
 	buffer := *bufp
-	var totalBytes int64 = 0
 
 	for {
 		if ctx.Err() != nil {
-			return totalBytes, ctx.Err()
+			return ctx.Err()
 		}
 
 		if err := src.SetDeadline(time.Now().Add(udpTimeout)); err != nil {
-			return totalBytes, fmt.Errorf("set read deadline: %w", err)
+			return fmt.Errorf("set read deadline: %w", err)
 		}
 
 		n, err := src.Read(buffer)
@@ -334,15 +307,14 @@ func (c *udpPacketConn) copy(ctx context.Context, dst net.Conn, src net.Conn, bu
 			if isTimeout(err) {
 				continue
 			}
-			return totalBytes, fmt.Errorf("read from %s: %w", direction, err)
+			return fmt.Errorf("read from %s: %w", direction, err)
 		}
 
-		nWritten, err := dst.Write(buffer[:n])
+		_, err = dst.Write(buffer[:n])
 		if err != nil {
-			return totalBytes, fmt.Errorf("write to %s: %w", direction, err)
+			return fmt.Errorf("write to %s: %w", direction, err)
 		}
 
-		totalBytes += int64(nWritten)
 		c.updateLastSeen()
 	}
 }

client/firewall/uspfilter/rule.go

@@ -29,15 +29,14 @@ func (r *PeerRule) ID() string {
 }
 
 type RouteRule struct {
-	id           string
-	mgmtId       []byte
-	sources      []netip.Prefix
-	dstSet       firewall.Set
-	destinations []netip.Prefix
-	proto        firewall.Protocol
-	srcPort      *firewall.Port
-	dstPort      *firewall.Port
-	action       firewall.Action
+	id          string
+	mgmtId      []byte
+	sources     []netip.Prefix
+	destination netip.Prefix
+	proto       firewall.Protocol
+	srcPort     *firewall.Port
+	dstPort     *firewall.Port
+	action      firewall.Action
 }
 
 // ID returns the rule id

client/firewall/uspfilter/tracer_test.go

@@ -198,12 +198,12 @@ func TestTracePacket(t *testing.T) {
 				m.forwarder.Store(&forwarder.Forwarder{})
 
 				src := netip.PrefixFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 32)
-				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{192, 168, 17, 2}), 32)
-				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, fw.Network{Prefix: dst}, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept)
+				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{172, 17, 0, 2}), 32)
+				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, dst, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept)
 				require.NoError(t, err)
 			},
 			packetBuilder: func() *PacketBuilder {
-				return createPacketBuilder("1.1.1.1", "192.168.17.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
@@ -222,12 +222,12 @@ func TestTracePacket(t *testing.T) {
 				m.nativeRouter.Store(false)
 
 				src := netip.PrefixFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 32)
-				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{192, 168, 17, 2}), 32)
-				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, fw.Network{Prefix: dst}, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionDrop)
+				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{172, 17, 0, 2}), 32)
+				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, dst, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionDrop)
 				require.NoError(t, err)
 			},
 			packetBuilder: func() *PacketBuilder {
-				return createPacketBuilder("1.1.1.1", "192.168.17.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
@@ -245,7 +245,7 @@ func TestTracePacket(t *testing.T) {
 				m.nativeRouter.Store(true)
 			},
 			packetBuilder: func() *PacketBuilder {
-				return createPacketBuilder("1.1.1.1", "192.168.17.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
@@ -263,7 +263,7 @@ func TestTracePacket(t *testing.T) {
 				m.routingEnabled.Store(false)
 			},
 			packetBuilder: func() *PacketBuilder {
-				return createPacketBuilder("1.1.1.1", "192.168.17.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
@@ -425,8 +425,8 @@ func TestTracePacket(t *testing.T) {
 
 			require.True(t, m.localipmanager.IsLocalIP(netip.MustParseAddr("100.10.0.100")),
 				"100.10.0.100 should be recognized as a local IP")
-			require.False(t, m.localipmanager.IsLocalIP(netip.MustParseAddr("192.168.17.2")),
-				"192.168.17.2 should not be recognized as a local IP")
+			require.False(t, m.localipmanager.IsLocalIP(netip.MustParseAddr("172.17.0.2")),
+				"172.17.0.2 should not be recognized as a local IP")
 
 			pb := tc.packetBuilder()
 

client/firewall/uspfilter/uspfilter.go

@@ -49,10 +49,10 @@ var errNatNotSupported = errors.New("nat not supported with userspace firewall")
 // RuleSet is a set of rules grouped by a string key
 type RuleSet map[string]PeerRule
 
-type RouteRules []*RouteRule
+type RouteRules []RouteRule
 
 func (r RouteRules) Sort() {
-	slices.SortStableFunc(r, func(a, b *RouteRule) int {
+	slices.SortStableFunc(r, func(a, b RouteRule) int {
 		// Deny rules come first
 		if a.action == firewall.ActionDrop && b.action != firewall.ActionDrop {
 			return -1
@@ -99,8 +99,6 @@ type Manager struct {
 	forwarder   atomic.Pointer[forwarder.Forwarder]
 	logger      *nblog.Logger
 	flowLogger  nftypes.FlowLogger
-
-	blockRule firewall.Rule
 }
 
 // decoder for packages
@@ -203,35 +201,41 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 		}
 	}
 
+	if err := m.blockInvalidRouted(iface); err != nil {
+		log.Errorf("failed to block invalid routed traffic: %v", err)
+	}
+
 	if err := iface.SetFilter(m); err != nil {
 		return nil, fmt.Errorf("set filter: %w", err)
 	}
 	return m, nil
 }
 
-func (m *Manager) blockInvalidRouted(iface common.IFaceMapper) (firewall.Rule, error) {
+func (m *Manager) blockInvalidRouted(iface common.IFaceMapper) error {
+	if m.forwarder.Load() == nil {
+		return nil
+	}
 	wgPrefix, err := netip.ParsePrefix(iface.Address().Network.String())
 	if err != nil {
-		return nil, fmt.Errorf("parse wireguard network: %w", err)
+		return fmt.Errorf("parse wireguard network: %w", err)
 	}
 	log.Debugf("blocking invalid routed traffic for %s", wgPrefix)
 
-	rule, err := m.addRouteFiltering(
+	if _, err := m.AddRouteFiltering(
 		nil,
 		[]netip.Prefix{netip.PrefixFrom(netip.IPv4Unspecified(), 0)},
-		firewall.Network{Prefix: wgPrefix},
+		wgPrefix,
 		firewall.ProtocolALL,
 		nil,
 		nil,
 		firewall.ActionDrop,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("block wg nte : %w", err)
+	); err != nil {
+		return fmt.Errorf("block wg nte : %w", err)
 	}
 
 	// TODO: Block networks that we're a client of
 
-	return rule, nil
+	return nil
 }
 
 func (m *Manager) determineRouting() error {
@@ -409,23 +413,10 @@ func (m *Manager) AddPeerFiltering(
 func (m *Manager) AddRouteFiltering(
 	id []byte,
 	sources []netip.Prefix,
-	destination firewall.Network,
+	destination netip.Prefix,
 	proto firewall.Protocol,
-	sPort, dPort *firewall.Port,
-	action firewall.Action,
-) (firewall.Rule, error) {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.addRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
-}
-
-func (m *Manager) addRouteFiltering(
-	id []byte,
-	sources []netip.Prefix,
-	destination firewall.Network,
-	proto firewall.Protocol,
-	sPort, dPort *firewall.Port,
+	sPort *firewall.Port,
+	dPort *firewall.Port,
 	action firewall.Action,
 ) (firewall.Rule, error) {
 	if m.nativeRouter.Load() && m.nativeFirewall != nil {
@@ -435,39 +426,34 @@ func (m *Manager) addRouteFiltering(
 	ruleID := uuid.New().String()
 	rule := RouteRule{
 		// TODO: consolidate these IDs
-		id:      ruleID,
-		mgmtId:  id,
-		sources: sources,
-		dstSet:  destination.Set,
-		proto:   proto,
-		srcPort: sPort,
-		dstPort: dPort,
-		action:  action,
-	}
-	if destination.IsPrefix() {
-		rule.destinations = []netip.Prefix{destination.Prefix}
+		id:          ruleID,
+		mgmtId:      id,
+		sources:     sources,
+		destination: destination,
+		proto:       proto,
+		srcPort:     sPort,
+		dstPort:     dPort,
+		action:      action,
 	}
 
-	m.routeRules = append(m.routeRules, &rule)
+	m.mutex.Lock()
+	m.routeRules = append(m.routeRules, rule)
 	m.routeRules.Sort()
+	m.mutex.Unlock()
 
 	return &rule, nil
 }
 
 func (m *Manager) DeleteRouteRule(rule firewall.Rule) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.deleteRouteRule(rule)
-}
-
-func (m *Manager) deleteRouteRule(rule firewall.Rule) error {
 	if m.nativeRouter.Load() && m.nativeFirewall != nil {
 		return m.nativeFirewall.DeleteRouteRule(rule)
 	}
 
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
 	ruleID := rule.ID()
-	idx := slices.IndexFunc(m.routeRules, func(r *RouteRule) bool {
+	idx := slices.IndexFunc(m.routeRules, func(r RouteRule) bool {
 		return r.id == ruleID
 	})
 	if idx < 0 {
@@ -523,52 +509,6 @@ func (m *Manager) DeleteDNATRule(rule firewall.Rule) error {
 	return m.nativeFirewall.DeleteDNATRule(rule)
 }
 
-// UpdateSet updates the rule destinations associated with the given set
-// by merging the existing prefixes with the new ones, then deduplicating.
-func (m *Manager) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
-	if m.nativeRouter.Load() && m.nativeFirewall != nil {
-		return m.nativeFirewall.UpdateSet(set, prefixes)
-	}
-
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	var matches []*RouteRule
-	for _, rule := range m.routeRules {
-		if rule.dstSet == set {
-			matches = append(matches, rule)
-		}
-	}
-
-	if len(matches) == 0 {
-		return fmt.Errorf("no route rule found for set: %s", set)
-	}
-
-	destinations := matches[0].destinations
-	for _, prefix := range prefixes {
-		if prefix.Addr().Is4() {
-			destinations = append(destinations, prefix)
-		}
-	}
-
-	slices.SortFunc(destinations, func(a, b netip.Prefix) int {
-		cmp := a.Addr().Compare(b.Addr())
-		if cmp != 0 {
-			return cmp
-		}
-		return a.Bits() - b.Bits()
-	})
-
-	destinations = slices.Compact(destinations)
-
-	for _, rule := range matches {
-		rule.destinations = destinations
-	}
-	log.Debugf("updated set %s to prefixes %v", set.HashedName(), destinations)
-
-	return nil
-}
-
 // DropOutgoing filter outgoing packets
 func (m *Manager) DropOutgoing(packetData []byte, size int) bool {
 	return m.processOutgoingHooks(packetData, size)
@@ -824,8 +764,7 @@ func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packe
 	proto, pnum := getProtocolFromPacket(d)
 	srcPort, dstPort := getPortsFromPacket(d)
 
-	ruleID, pass := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort)
-	if !pass {
+	if ruleID, pass := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort); !pass {
 		m.logger.Trace("Dropping routed packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
 			ruleID, pnum, srcIP, srcPort, dstIP, dstPort)
 
@@ -851,11 +790,8 @@ func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packe
 	if fwd == nil {
 		m.logger.Trace("failed to forward routed packet (forwarder not initialized)")
 	} else {
-		fwd.RegisterRuleID(srcIP, dstIP, srcPort, dstPort, ruleID)
-
 		if err := fwd.InjectIncomingPacket(packetData); err != nil {
 			m.logger.Error("Failed to inject routed packet: %v", err)
-			fwd.DeleteRuleID(srcIP, dstIP, srcPort, dstPort)
 		}
 	}
 
@@ -1052,15 +988,8 @@ func (m *Manager) routeACLsPass(srcIP, dstIP netip.Addr, proto firewall.Protocol
 	return nil, false
 }
 
-func (m *Manager) ruleMatches(rule *RouteRule, srcAddr, dstAddr netip.Addr, proto firewall.Protocol, srcPort, dstPort uint16) bool {
-	destMatched := false
-	for _, dst := range rule.destinations {
-		if dst.Contains(dstAddr) {
-			destMatched = true
-			break
-		}
-	}
-	if !destMatched {
+func (m *Manager) ruleMatches(rule RouteRule, srcAddr, dstAddr netip.Addr, proto firewall.Protocol, srcPort, dstPort uint16) bool {
+	if !rule.destination.Contains(dstAddr) {
 		return false
 	}
 
@@ -1162,22 +1091,7 @@ func (m *Manager) EnableRouting() error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if err := m.determineRouting(); err != nil {
-		return fmt.Errorf("determine routing: %w", err)
-	}
-
-	if m.forwarder.Load() == nil {
-		return nil
-	}
-
-	rule, err := m.blockInvalidRouted(m.wgIface)
-	if err != nil {
-		return fmt.Errorf("block invalid routed: %w", err)
-	}
-
-	m.blockRule = rule
-
-	return nil
+	return m.determineRouting()
 }
 
 func (m *Manager) DisableRouting() error {
@@ -1202,12 +1116,5 @@ func (m *Manager) DisableRouting() error {
 
 	log.Debug("forwarder stopped")
 
-	if m.blockRule != nil {
-		if err := m.deleteRouteRule(m.blockRule); err != nil {
-			return fmt.Errorf("delete block rule: %w", err)
-		}
-		m.blockRule = nil
-	}
-
 	return nil
 }

client/firewall/uspfilter/uspfilter_filter_test.go

@@ -15,7 +15,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/mocks"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
-	"github.com/netbirdio/netbird/management/domain"
 )
 
 func TestPeerACLFiltering(t *testing.T) {
@@ -189,281 +188,6 @@ func TestPeerACLFiltering(t *testing.T) {
 			ruleAction:      fw.ActionAccept,
 			shouldBeBlocked: true,
 		},
-		{
-			name:            "Allow TCP traffic without port specification",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "Allow UDP traffic without port specification",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolUDP,
-			srcPort:         12345,
-			dstPort:         53,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolUDP,
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "TCP packet doesn't match UDP filter with same port",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolUDP,
-			ruleDstPort:     &fw.Port{Values: []uint16{443}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "UDP packet doesn't match TCP filter with same port",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolUDP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{Values: []uint16{443}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "ICMP packet doesn't match TCP filter",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolICMP,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "ICMP packet doesn't match UDP filter",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolICMP,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolUDP,
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Allow TCP traffic within port range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         8080,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "Block TCP traffic outside port range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         7999,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Edge Case - Port at Range Boundary",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         8100,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "UDP Port Range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolUDP,
-			srcPort:         12345,
-			dstPort:         5060,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolUDP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{5060, 5070}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "Allow multiple destination ports",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         8080,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{Values: []uint16{80, 8080, 443}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "Allow multiple source ports",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         80,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleSrcPort:     &fw.Port{Values: []uint16{12345, 12346, 12347}},
-			ruleAction:      fw.ActionAccept,
-			shouldBeBlocked: false,
-		},
-		// New drop test cases
-		{
-			name:            "Drop TCP traffic from WG peer",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{Values: []uint16{443}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop UDP traffic from WG peer",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolUDP,
-			srcPort:         12345,
-			dstPort:         53,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolUDP,
-			ruleDstPort:     &fw.Port{Values: []uint16{53}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop ICMP traffic from WG peer",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolICMP,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolICMP,
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop all traffic from WG peer",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolALL,
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop traffic from multiple source ports",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         80,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleSrcPort:     &fw.Port{Values: []uint16{12345, 12346, 12347}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop multiple destination ports",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         8080,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{Values: []uint16{80, 8080, 443}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Drop TCP traffic within port range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         8080,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Accept TCP traffic outside drop port range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         7999,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: false,
-		},
-		{
-			name:            "Drop TCP traffic with source port range",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         32100,
-			dstPort:         80,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleSrcPort:     &fw.Port{IsRange: true, Values: []uint16{32000, 33000}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
-		{
-			name:            "Mixed rule - drop specific port but allow other ports",
-			srcIP:           "100.10.0.1",
-			dstIP:           "100.10.0.100",
-			proto:           fw.ProtocolTCP,
-			srcPort:         12345,
-			dstPort:         443,
-			ruleIP:          "100.10.0.1",
-			ruleProto:       fw.ProtocolTCP,
-			ruleDstPort:     &fw.Port{Values: []uint16{443}},
-			ruleAction:      fw.ActionDrop,
-			shouldBeBlocked: true,
-		},
 	}
 
 	t.Run("Implicit DROP (no rules)", func(t *testing.T) {
@@ -474,28 +198,6 @@ func TestPeerACLFiltering(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-
-			if tc.ruleAction == fw.ActionDrop {
-				// add general accept rule to test drop rule
-				// TODO: this only works because 0.0.0.0 is tested last, we need to implement order
-				rules, err := manager.AddPeerFiltering(
-					nil,
-					net.ParseIP("0.0.0.0"),
-					fw.ProtocolALL,
-					nil,
-					nil,
-					fw.ActionAccept,
-					"",
-				)
-				require.NoError(t, err)
-				require.NotEmpty(t, rules)
-				t.Cleanup(func() {
-					for _, rule := range rules {
-						require.NoError(t, manager.DeletePeerRule(rule))
-					}
-				})
-			}
-
 			rules, err := manager.AddPeerFiltering(
 				nil,
 				net.ParseIP(tc.ruleIP),
@@ -601,8 +303,8 @@ func setupRoutedManager(tb testing.TB, network string) *Manager {
 	}
 
 	manager, err := Create(ifaceMock, false, flowLogger)
-	require.NoError(tb, err)
 	require.NoError(tb, manager.EnableRouting())
+	require.NoError(tb, err)
 	require.NotNil(tb, manager)
 	require.True(tb, manager.routingEnabled.Load())
 	require.False(tb, manager.nativeRouter.Load())
@@ -619,7 +321,7 @@ func TestRouteACLFiltering(t *testing.T) {
 
 	type rule struct {
 		sources []netip.Prefix
-		dest    fw.Network
+		dest    netip.Prefix
 		proto   fw.Protocol
 		srcPort *fw.Port
 		dstPort *fw.Port
@@ -645,7 +347,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 443,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{443}},
 				action:  fw.ActionAccept,
@@ -661,7 +363,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 443,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{443}},
 				action:  fw.ActionAccept,
@@ -677,7 +379,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 443,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("0.0.0.0/0")},
+				dest:    netip.MustParsePrefix("0.0.0.0/0"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{443}},
 				action:  fw.ActionAccept,
@@ -693,7 +395,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 53,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolUDP,
 				dstPort: &fw.Port{Values: []uint16{53}},
 				action:  fw.ActionAccept,
@@ -707,7 +409,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			proto: fw.ProtocolICMP,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("0.0.0.0/0")},
+				dest:    netip.MustParsePrefix("0.0.0.0/0"),
 				proto:   fw.ProtocolICMP,
 				action:  fw.ActionAccept,
 			},
@@ -722,7 +424,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
@@ -738,7 +440,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 8080,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
@@ -754,7 +456,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
@@ -770,7 +472,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
@@ -786,7 +488,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				srcPort: &fw.Port{Values: []uint16{12345}},
 				action:  fw.ActionAccept,
@@ -805,7 +507,7 @@ func TestRouteACLFiltering(t *testing.T) {
 					netip.MustParsePrefix("100.10.0.0/16"),
 					netip.MustParsePrefix("172.16.0.0/16"),
 				},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
@@ -819,7 +521,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			proto: fw.ProtocolICMP,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				action:  fw.ActionAccept,
 			},
@@ -834,13 +536,33 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionAccept,
 			},
 			shouldPass: true,
 		},
+		{
+			name:  "Multiple source networks with mismatched protocol",
+			srcIP: "172.16.0.1",
+			dstIP: "192.168.1.100",
+			// Should not match TCP rule
+			proto:   fw.ProtocolUDP,
+			srcPort: 12345,
+			dstPort: 80,
+			rule: rule{
+				sources: []netip.Prefix{
+					netip.MustParsePrefix("100.10.0.0/16"),
+					netip.MustParsePrefix("172.16.0.0/16"),
+				},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
+				proto:   fw.ProtocolTCP,
+				dstPort: &fw.Port{Values: []uint16{80}},
+				action:  fw.ActionAccept,
+			},
+			shouldPass: false,
+		},
 		{
 			name:    "Allow multiple destination ports",
 			srcIP:   "100.10.0.1",
@@ -850,7 +572,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 8080,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80, 8080, 443}},
 				action:  fw.ActionAccept,
@@ -866,7 +588,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				srcPort: &fw.Port{Values: []uint16{12345, 12346, 12347}},
 				action:  fw.ActionAccept,
@@ -882,7 +604,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				srcPort: &fw.Port{Values: []uint16{12345}},
 				dstPort: &fw.Port{Values: []uint16{80}},
@@ -899,7 +621,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 8080,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{
 					IsRange: true,
@@ -918,7 +640,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 7999,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{
 					IsRange: true,
@@ -937,7 +659,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				srcPort: &fw.Port{
 					IsRange: true,
@@ -956,7 +678,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 443,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				srcPort: &fw.Port{
 					IsRange: true,
@@ -978,7 +700,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 8100,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{
 					IsRange: true,
@@ -997,7 +719,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 5060,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolUDP,
 				dstPort: &fw.Port{
 					IsRange: true,
@@ -1016,7 +738,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 8080,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				dstPort: &fw.Port{
 					IsRange: true,
@@ -1035,7 +757,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 443,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{443}},
 				action:  fw.ActionDrop,
@@ -1051,7 +773,7 @@ func TestRouteACLFiltering(t *testing.T) {
 			dstPort: 80,
 			rule: rule{
 				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolALL,
 				action:  fw.ActionDrop,
 			},
@@ -1069,158 +791,17 @@ func TestRouteACLFiltering(t *testing.T) {
 					netip.MustParsePrefix("100.10.0.0/16"),
 					netip.MustParsePrefix("172.16.0.0/16"),
 				},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+				dest:    netip.MustParsePrefix("192.168.1.0/24"),
 				proto:   fw.ProtocolTCP,
 				dstPort: &fw.Port{Values: []uint16{80}},
 				action:  fw.ActionDrop,
 			},
 			shouldPass: false,
 		},
-
-		{
-			name:    "Drop empty destination set",
-			srcIP:   "172.16.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolTCP,
-			srcPort: 12345,
-			dstPort: 80,
-			rule: rule{
-				sources: []netip.Prefix{
-					netip.MustParsePrefix("172.16.0.0/16"),
-				},
-				dest:    fw.Network{Set: fw.Set{}},
-				proto:   fw.ProtocolTCP,
-				dstPort: &fw.Port{Values: []uint16{80}},
-				action:  fw.ActionAccept,
-			},
-			shouldPass: false,
-		},
-		{
-			name:    "Accept TCP traffic outside drop port range",
-			srcIP:   "100.10.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolTCP,
-			srcPort: 12345,
-			dstPort: 7999,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolTCP,
-				dstPort: &fw.Port{IsRange: true, Values: []uint16{8000, 8100}},
-				action:  fw.ActionDrop,
-			},
-			shouldPass: true,
-		},
-		{
-			name:    "Allow TCP traffic without port specification",
-			srcIP:   "100.10.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolTCP,
-			srcPort: 12345,
-			dstPort: 443,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolTCP,
-				action:  fw.ActionAccept,
-			},
-			shouldPass: true,
-		},
-		{
-			name:    "Allow UDP traffic without port specification",
-			srcIP:   "100.10.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolUDP,
-			srcPort: 12345,
-			dstPort: 53,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolUDP,
-				action:  fw.ActionAccept,
-			},
-			shouldPass: true,
-		},
-		{
-			name:    "TCP packet doesn't match UDP filter with same port",
-			srcIP:   "100.10.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolTCP,
-			srcPort: 12345,
-			dstPort: 80,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolUDP,
-				dstPort: &fw.Port{Values: []uint16{80}},
-				action:  fw.ActionAccept,
-			},
-			shouldPass: false,
-		},
-		{
-			name:    "UDP packet doesn't match TCP filter with same port",
-			srcIP:   "100.10.0.1",
-			dstIP:   "192.168.1.100",
-			proto:   fw.ProtocolUDP,
-			srcPort: 12345,
-			dstPort: 80,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolTCP,
-				dstPort: &fw.Port{Values: []uint16{80}},
-				action:  fw.ActionAccept,
-			},
-			shouldPass: false,
-		},
-		{
-			name:  "ICMP packet doesn't match TCP filter",
-			srcIP: "100.10.0.1",
-			dstIP: "192.168.1.100",
-			proto: fw.ProtocolICMP,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolTCP,
-				action:  fw.ActionAccept,
-			},
-			shouldPass: false,
-		},
-		{
-			name:  "ICMP packet doesn't match UDP filter",
-			srcIP: "100.10.0.1",
-			dstIP: "192.168.1.100",
-			proto: fw.ProtocolICMP,
-			rule: rule{
-				sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-				dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
-				proto:   fw.ProtocolUDP,
-				action:  fw.ActionAccept,
-			},
-			shouldPass: false,
-		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			if tc.rule.action == fw.ActionDrop {
-				// add general accept rule to test drop rule
-				rule, err := manager.AddRouteFiltering(
-					nil,
-					[]netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-					fw.Network{Prefix: netip.MustParsePrefix("0.0.0.0/0")},
-					fw.ProtocolALL,
-					nil,
-					nil,
-					fw.ActionAccept,
-				)
-				require.NoError(t, err)
-				require.NotNil(t, rule)
-				t.Cleanup(func() {
-					require.NoError(t, manager.DeleteRouteRule(rule))
-				})
-			}
-
 			rule, err := manager.AddRouteFiltering(
 				nil,
 				tc.rule.sources,
@@ -1255,7 +836,7 @@ func TestRouteACLOrder(t *testing.T) {
 		name  string
 		rules []struct {
 			sources []netip.Prefix
-			dest    fw.Network
+			dest    netip.Prefix
 			proto   fw.Protocol
 			srcPort *fw.Port
 			dstPort *fw.Port
@@ -1276,7 +857,7 @@ func TestRouteACLOrder(t *testing.T) {
 			name: "Drop rules take precedence over accept",
 			rules: []struct {
 				sources []netip.Prefix
-				dest    fw.Network
+				dest    netip.Prefix
 				proto   fw.Protocol
 				srcPort *fw.Port
 				dstPort *fw.Port
@@ -1285,7 +866,7 @@ func TestRouteACLOrder(t *testing.T) {
 				{
 					// Accept rule added first
 					sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-					dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+					dest:    netip.MustParsePrefix("192.168.1.0/24"),
 					proto:   fw.ProtocolTCP,
 					dstPort: &fw.Port{Values: []uint16{80, 443}},
 					action:  fw.ActionAccept,
@@ -1293,7 +874,7 @@ func TestRouteACLOrder(t *testing.T) {
 				{
 					// Drop rule added second but should be evaluated first
 					sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-					dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+					dest:    netip.MustParsePrefix("192.168.1.0/24"),
 					proto:   fw.ProtocolTCP,
 					dstPort: &fw.Port{Values: []uint16{443}},
 					action:  fw.ActionDrop,
@@ -1331,7 +912,7 @@ func TestRouteACLOrder(t *testing.T) {
 			name: "Multiple drop rules take precedence",
 			rules: []struct {
 				sources []netip.Prefix
-				dest    fw.Network
+				dest    netip.Prefix
 				proto   fw.Protocol
 				srcPort *fw.Port
 				dstPort *fw.Port
@@ -1340,14 +921,14 @@ func TestRouteACLOrder(t *testing.T) {
 				{
 					// Accept all
 					sources: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-					dest:    fw.Network{Prefix: netip.MustParsePrefix("0.0.0.0/0")},
+					dest:    netip.MustParsePrefix("0.0.0.0/0"),
 					proto:   fw.ProtocolALL,
 					action:  fw.ActionAccept,
 				},
 				{
 					// Drop specific port
 					sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-					dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+					dest:    netip.MustParsePrefix("192.168.1.0/24"),
 					proto:   fw.ProtocolTCP,
 					dstPort: &fw.Port{Values: []uint16{443}},
 					action:  fw.ActionDrop,
@@ -1355,7 +936,7 @@ func TestRouteACLOrder(t *testing.T) {
 				{
 					// Drop different port
 					sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-					dest:    fw.Network{Prefix: netip.MustParsePrefix("192.168.1.0/24")},
+					dest:    netip.MustParsePrefix("192.168.1.0/24"),
 					proto:   fw.ProtocolTCP,
 					dstPort: &fw.Port{Values: []uint16{80}},
 					action:  fw.ActionDrop,
@@ -1434,53 +1015,3 @@ func TestRouteACLOrder(t *testing.T) {
 		})
 	}
 }
-
-func TestRouteACLSet(t *testing.T) {
-	ifaceMock := &IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-		AddressFunc: func() wgaddr.Address {
-			return wgaddr.Address{
-				IP: net.ParseIP("100.10.0.100"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("100.10.0.0"),
-					Mask: net.CIDRMask(16, 32),
-				},
-			}
-		},
-	}
-
-	manager, err := Create(ifaceMock, false, flowLogger)
-	require.NoError(t, err)
-	t.Cleanup(func() {
-		require.NoError(t, manager.Close(nil))
-	})
-
-	set := fw.NewDomainSet(domain.List{"example.org"})
-
-	// Add rule that uses the set (initially empty)
-	rule, err := manager.AddRouteFiltering(
-		nil,
-		[]netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-		fw.Network{Set: set},
-		fw.ProtocolTCP,
-		nil,
-		nil,
-		fw.ActionAccept,
-	)
-	require.NoError(t, err)
-	require.NotNil(t, rule)
-
-	srcIP := netip.MustParseAddr("100.10.0.1")
-	dstIP := netip.MustParseAddr("192.168.1.100")
-
-	// Check that traffic is dropped (empty set shouldn't match anything)
-	_, isAllowed := manager.routeACLsPass(srcIP, dstIP, fw.ProtocolTCP, 12345, 80)
-	require.False(t, isAllowed, "Empty set should not allow any traffic")
-
-	err = manager.UpdateSet(set, []netip.Prefix{netip.MustParsePrefix("192.168.1.0/24")})
-	require.NoError(t, err)
-
-	// Now the packet should be allowed
-	_, isAllowed = manager.routeACLsPass(srcIP, dstIP, fw.ProtocolTCP, 12345, 80)
-	require.True(t, isAllowed, "After set update, traffic to the added network should be allowed")
-}

client/firewall/uspfilter/uspfilter_test.go

@@ -20,7 +20,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/netflow"
-	"github.com/netbirdio/netbird/management/domain"
 )
 
 var logger = log.NewFromLogrus(logrus.StandardLogger())
@@ -712,203 +711,3 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 		})
 	}
 }
-
-func TestUpdateSetMerge(t *testing.T) {
-	ifaceMock := &IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}
-
-	manager, err := Create(ifaceMock, false, flowLogger)
-	require.NoError(t, err)
-	t.Cleanup(func() {
-		require.NoError(t, manager.Close(nil))
-	})
-
-	set := fw.NewDomainSet(domain.List{"example.org"})
-
-	initialPrefixes := []netip.Prefix{
-		netip.MustParsePrefix("10.0.0.0/24"),
-		netip.MustParsePrefix("192.168.1.0/24"),
-	}
-
-	rule, err := manager.AddRouteFiltering(
-		nil,
-		[]netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-		fw.Network{Set: set},
-		fw.ProtocolTCP,
-		nil,
-		nil,
-		fw.ActionAccept,
-	)
-	require.NoError(t, err)
-	require.NotNil(t, rule)
-
-	// Update the set with initial prefixes
-	err = manager.UpdateSet(set, initialPrefixes)
-	require.NoError(t, err)
-
-	// Test initial prefixes work
-	srcIP := netip.MustParseAddr("100.10.0.1")
-	dstIP1 := netip.MustParseAddr("10.0.0.100")
-	dstIP2 := netip.MustParseAddr("192.168.1.100")
-	dstIP3 := netip.MustParseAddr("172.16.0.100")
-
-	_, isAllowed1 := manager.routeACLsPass(srcIP, dstIP1, fw.ProtocolTCP, 12345, 80)
-	_, isAllowed2 := manager.routeACLsPass(srcIP, dstIP2, fw.ProtocolTCP, 12345, 80)
-	_, isAllowed3 := manager.routeACLsPass(srcIP, dstIP3, fw.ProtocolTCP, 12345, 80)
-
-	require.True(t, isAllowed1, "Traffic to 10.0.0.100 should be allowed")
-	require.True(t, isAllowed2, "Traffic to 192.168.1.100 should be allowed")
-	require.False(t, isAllowed3, "Traffic to 172.16.0.100 should be denied")
-
-	newPrefixes := []netip.Prefix{
-		netip.MustParsePrefix("172.16.0.0/16"),
-		netip.MustParsePrefix("10.1.0.0/24"),
-	}
-
-	err = manager.UpdateSet(set, newPrefixes)
-	require.NoError(t, err)
-
-	// Check that all original prefixes are still included
-	_, isAllowed1 = manager.routeACLsPass(srcIP, dstIP1, fw.ProtocolTCP, 12345, 80)
-	_, isAllowed2 = manager.routeACLsPass(srcIP, dstIP2, fw.ProtocolTCP, 12345, 80)
-	require.True(t, isAllowed1, "Traffic to 10.0.0.100 should still be allowed after update")
-	require.True(t, isAllowed2, "Traffic to 192.168.1.100 should still be allowed after update")
-
-	// Check that new prefixes are included
-	dstIP4 := netip.MustParseAddr("172.16.1.100")
-	dstIP5 := netip.MustParseAddr("10.1.0.50")
-
-	_, isAllowed4 := manager.routeACLsPass(srcIP, dstIP4, fw.ProtocolTCP, 12345, 80)
-	_, isAllowed5 := manager.routeACLsPass(srcIP, dstIP5, fw.ProtocolTCP, 12345, 80)
-
-	require.True(t, isAllowed4, "Traffic to new prefix 172.16.0.0/16 should be allowed")
-	require.True(t, isAllowed5, "Traffic to new prefix 10.1.0.0/24 should be allowed")
-
-	// Verify the rule has all prefixes
-	manager.mutex.RLock()
-	foundRule := false
-	for _, r := range manager.routeRules {
-		if r.id == rule.ID() {
-			foundRule = true
-			require.Len(t, r.destinations, len(initialPrefixes)+len(newPrefixes),
-				"Rule should have all prefixes merged")
-		}
-	}
-	manager.mutex.RUnlock()
-	require.True(t, foundRule, "Rule should be found")
-}
-
-func TestUpdateSetDeduplication(t *testing.T) {
-	ifaceMock := &IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}
-
-	manager, err := Create(ifaceMock, false, flowLogger)
-	require.NoError(t, err)
-	t.Cleanup(func() {
-		require.NoError(t, manager.Close(nil))
-	})
-
-	set := fw.NewDomainSet(domain.List{"example.org"})
-
-	rule, err := manager.AddRouteFiltering(
-		nil,
-		[]netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-		fw.Network{Set: set},
-		fw.ProtocolTCP,
-		nil,
-		nil,
-		fw.ActionAccept,
-	)
-	require.NoError(t, err)
-	require.NotNil(t, rule)
-
-	initialPrefixes := []netip.Prefix{
-		netip.MustParsePrefix("10.0.0.0/24"),
-		netip.MustParsePrefix("10.0.0.0/24"), // Duplicate
-		netip.MustParsePrefix("192.168.1.0/24"),
-		netip.MustParsePrefix("192.168.1.0/24"), // Duplicate
-	}
-
-	err = manager.UpdateSet(set, initialPrefixes)
-	require.NoError(t, err)
-
-	// Check the internal state for deduplication
-	manager.mutex.RLock()
-	foundRule := false
-	for _, r := range manager.routeRules {
-		if r.id == rule.ID() {
-			foundRule = true
-			// Should have deduplicated to 2 prefixes
-			require.Len(t, r.destinations, 2, "Duplicate prefixes should be removed")
-
-			// Check the prefixes are correct
-			expectedPrefixes := []netip.Prefix{
-				netip.MustParsePrefix("10.0.0.0/24"),
-				netip.MustParsePrefix("192.168.1.0/24"),
-			}
-			for i, prefix := range expectedPrefixes {
-				require.True(t, r.destinations[i] == prefix,
-					"Prefix should match expected value")
-			}
-		}
-	}
-	manager.mutex.RUnlock()
-	require.True(t, foundRule, "Rule should be found")
-
-	// Test with overlapping prefixes of different sizes
-	overlappingPrefixes := []netip.Prefix{
-		netip.MustParsePrefix("10.0.0.0/16"),    // More general
-		netip.MustParsePrefix("10.0.0.0/24"),    // More specific (already exists)
-		netip.MustParsePrefix("192.168.0.0/16"), // More general
-		netip.MustParsePrefix("192.168.1.0/24"), // More specific (already exists)
-	}
-
-	err = manager.UpdateSet(set, overlappingPrefixes)
-	require.NoError(t, err)
-
-	// Check that all prefixes are included (no deduplication of overlapping prefixes)
-	manager.mutex.RLock()
-	for _, r := range manager.routeRules {
-		if r.id == rule.ID() {
-			// Should have all 4 prefixes (2 original + 2 new more general ones)
-			require.Len(t, r.destinations, 4,
-				"Overlapping prefixes should not be deduplicated")
-
-			// Verify they're sorted correctly (more specific prefixes should come first)
-			prefixes := make([]string, 0, len(r.destinations))
-			for _, p := range r.destinations {
-				prefixes = append(prefixes, p.String())
-			}
-
-			// Check sorted order
-			require.Equal(t, []string{
-				"10.0.0.0/16",
-				"10.0.0.0/24",
-				"192.168.0.0/16",
-				"192.168.1.0/24",
-			}, prefixes, "Prefixes should be sorted")
-		}
-	}
-	manager.mutex.RUnlock()
-
-	// Test functionality with all prefixes
-	testCases := []struct {
-		dstIP    netip.Addr
-		expected bool
-		desc     string
-	}{
-		{netip.MustParseAddr("10.0.0.100"), true, "IP in both /16 and /24"},
-		{netip.MustParseAddr("10.0.1.100"), true, "IP only in /16"},
-		{netip.MustParseAddr("192.168.1.100"), true, "IP in both /16 and /24"},
-		{netip.MustParseAddr("192.168.2.100"), true, "IP only in /16"},
-		{netip.MustParseAddr("172.16.0.100"), false, "IP not in any prefix"},
-	}
-
-	srcIP := netip.MustParseAddr("100.10.0.1")
-	for _, tc := range testCases {
-		_, isAllowed := manager.routeACLsPass(srcIP, tc.dstIP, fw.ProtocolTCP, 12345, 80)
-		require.Equal(t, tc.expected, isAllowed, tc.desc)
-	}
-}

client/iface/configurer/kernel_unix.go

@@ -201,30 +201,14 @@ func (c *KernelConfigurer) configure(config wgtypes.Config) error {
 func (c *KernelConfigurer) Close() {
 }
 
-func (c *KernelConfigurer) GetStats() (map[string]WGStats, error) {
-	stats := make(map[string]WGStats)
-	wg, err := wgctrl.New()
+func (c *KernelConfigurer) GetStats(peerKey string) (WGStats, error) {
+	peer, err := c.getPeer(c.deviceName, peerKey)
 	if err != nil {
-		return nil, fmt.Errorf("wgctl: %w", err)
+		return WGStats{}, fmt.Errorf("get wireguard stats: %w", err)
 	}
-	defer func() {
-		err = wg.Close()
-		if err != nil {
-			log.Errorf("Got error while closing wgctl: %v", err)
-		}
-	}()
-
-	wgDevice, err := wg.Device(c.deviceName)
-	if err != nil {
-		return nil, fmt.Errorf("get device %s: %w", c.deviceName, err)
-	}
-
-	for _, peer := range wgDevice.Peers {
-		stats[peer.PublicKey.String()] = WGStats{
-			LastHandshake: peer.LastHandshakeTime,
-			TxBytes:       peer.TransmitBytes,
-			RxBytes:       peer.ReceiveBytes,
-		}
-	}
-	return stats, nil
+	return WGStats{
+		LastHandshake: peer.LastHandshakeTime,
+		TxBytes:       peer.TransmitBytes,
+		RxBytes:       peer.ReceiveBytes,
+	}, nil
 }

client/iface/configurer/usp.go

@@ -1,7 +1,6 @@
 package configurer
 
 import (
-	"encoding/base64"
 	"encoding/hex"
 	"fmt"
 	"net"
@@ -18,13 +17,6 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-const (
-	ipcKeyLastHandshakeTimeSec  = "last_handshake_time_sec"
-	ipcKeyLastHandshakeTimeNsec = "last_handshake_time_nsec"
-	ipcKeyTxBytes               = "tx_bytes"
-	ipcKeyRxBytes               = "rx_bytes"
-)
-
 var ErrAllowedIPNotFound = fmt.Errorf("allowed IP not found")
 
 type WGUSPConfigurer struct {
@@ -225,75 +217,91 @@ func (t *WGUSPConfigurer) Close() {
 	}
 }
 
-func (t *WGUSPConfigurer) GetStats() (map[string]WGStats, error) {
+func (t *WGUSPConfigurer) GetStats(peerKey string) (WGStats, error) {
 	ipc, err := t.device.IpcGet()
 	if err != nil {
-		return nil, fmt.Errorf("ipc get: %w", err)
+		return WGStats{}, fmt.Errorf("ipc get: %w", err)
 	}
 
-	return parseTransfers(ipc)
+	stats, err := findPeerInfo(ipc, peerKey, []string{
+		"last_handshake_time_sec",
+		"last_handshake_time_nsec",
+		"tx_bytes",
+		"rx_bytes",
+	})
+	if err != nil {
+		return WGStats{}, fmt.Errorf("find peer info: %w", err)
+	}
+
+	sec, err := strconv.ParseInt(stats["last_handshake_time_sec"], 10, 64)
+	if err != nil {
+		return WGStats{}, fmt.Errorf("parse handshake sec: %w", err)
+	}
+	nsec, err := strconv.ParseInt(stats["last_handshake_time_nsec"], 10, 64)
+	if err != nil {
+		return WGStats{}, fmt.Errorf("parse handshake nsec: %w", err)
+	}
+	txBytes, err := strconv.ParseInt(stats["tx_bytes"], 10, 64)
+	if err != nil {
+		return WGStats{}, fmt.Errorf("parse tx_bytes: %w", err)
+	}
+	rxBytes, err := strconv.ParseInt(stats["rx_bytes"], 10, 64)
+	if err != nil {
+		return WGStats{}, fmt.Errorf("parse rx_bytes: %w", err)
+	}
+
+	return WGStats{
+		LastHandshake: time.Unix(sec, nsec),
+		TxBytes:       txBytes,
+		RxBytes:       rxBytes,
+	}, nil
 }
 
-func parseTransfers(ipc string) (map[string]WGStats, error) {
-	stats := make(map[string]WGStats)
-	var (
-		currentKey   string
-		currentStats WGStats
-		hasPeer      bool
-	)
-	lines := strings.Split(ipc, "\n")
+func findPeerInfo(ipcInput string, peerKey string, searchConfigKeys []string) (map[string]string, error) {
+	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
+	if err != nil {
+		return nil, fmt.Errorf("parse key: %w", err)
+	}
+
+	hexKey := hex.EncodeToString(peerKeyParsed[:])
+
+	lines := strings.Split(ipcInput, "\n")
+
+	configFound := map[string]string{}
+	foundPeer := false
 	for _, line := range lines {
 		line = strings.TrimSpace(line)
 
 		// If we're within the details of the found peer and encounter another public key,
 		// this means we're starting another peer's details. So, stop.
-		if strings.HasPrefix(line, "public_key=") {
-			peerID := strings.TrimPrefix(line, "public_key=")
-			h, err := hex.DecodeString(peerID)
-			if err != nil {
-				return nil, fmt.Errorf("decode peerID: %w", err)
-			}
-			currentKey = base64.StdEncoding.EncodeToString(h)
-			currentStats = WGStats{} // Reset stats for the new peer
-			hasPeer = true
-			stats[currentKey] = currentStats
-			continue
+		if strings.HasPrefix(line, "public_key=") && foundPeer {
+			break
 		}
 
-		if !hasPeer {
-			continue
+		// Identify the peer with the specific public key
+		if line == fmt.Sprintf("public_key=%s", hexKey) {
+			foundPeer = true
 		}
 
-		key := strings.SplitN(line, "=", 2)
-		if len(key) != 2 {
-			continue
-		}
-		switch key[0] {
-		case ipcKeyLastHandshakeTimeSec:
-			hs, err := toLastHandshake(key[1])
-			if err != nil {
-				return nil, err
+		for _, key := range searchConfigKeys {
+			if foundPeer && strings.HasPrefix(line, key+"=") {
+				v := strings.SplitN(line, "=", 2)
+				configFound[v[0]] = v[1]
 			}
-			currentStats.LastHandshake = hs
-			stats[currentKey] = currentStats
-		case ipcKeyRxBytes:
-			rxBytes, err := toBytes(key[1])
-			if err != nil {
-				return nil, fmt.Errorf("parse rx_bytes: %w", err)
-			}
-			currentStats.RxBytes = rxBytes
-			stats[currentKey] = currentStats
-		case ipcKeyTxBytes:
-			TxBytes, err := toBytes(key[1])
-			if err != nil {
-				return nil, fmt.Errorf("parse tx_bytes: %w", err)
-			}
-			currentStats.TxBytes = TxBytes
-			stats[currentKey] = currentStats
 		}
 	}
 
-	return stats, nil
+	// todo: use multierr
+	for _, key := range searchConfigKeys {
+		if _, ok := configFound[key]; !ok {
+			return configFound, fmt.Errorf("config key not found: %s", key)
+		}
+	}
+	if !foundPeer {
+		return nil, fmt.Errorf("%w: %s", ErrPeerNotFound, peerKey)
+	}
+
+	return configFound, nil
 }
 
 func toWgUserspaceString(wgCfg wgtypes.Config) string {
@@ -347,18 +355,6 @@ func toWgUserspaceString(wgCfg wgtypes.Config) string {
 	return sb.String()
 }
 
-func toLastHandshake(stringVar string) (time.Time, error) {
-	sec, err := strconv.ParseInt(stringVar, 10, 64)
-	if err != nil {
-		return time.Time{}, fmt.Errorf("parse handshake sec: %w", err)
-	}
-	return time.Unix(sec, 0), nil
-}
-
-func toBytes(s string) (int64, error) {
-	return strconv.ParseInt(s, 10, 64)
-}
-
 func getFwmark() int {
 	if nbnet.AdvancedRouting() {
 		return nbnet.ControlPlaneMark

client/iface/configurer/usp_test.go

@@ -2,8 +2,10 @@ package configurer
 
 import (
 	"encoding/hex"
+	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
@@ -32,35 +34,58 @@ errno=0
 
 `
 
-func Test_parseTransfers(t *testing.T) {
+func Test_findPeerInfo(t *testing.T) {
 	tests := []struct {
-		name    string
-		peerKey string
-		want    WGStats
+		name       string
+		peerKey    string
+		searchKeys []string
+		want       map[string]string
+		wantErr    bool
 	}{
 		{
-			name:    "single",
-			peerKey: "b85996fecc9c7f1fc6d2572a76eda11d59bcd20be8e543b15ce4bd85a8e75a33",
-			want: WGStats{
-				TxBytes: 0,
-				RxBytes: 0,
+			name:       "single",
+			peerKey:    "58402e695ba1772b1cc9309755f043251ea77fdcf10fbe63989ceb7e19321376",
+			searchKeys: []string{"tx_bytes"},
+			want: map[string]string{
+				"tx_bytes": "38333",
 			},
+			wantErr: false,
 		},
 		{
-			name:    "multiple",
-			peerKey: "58402e695ba1772b1cc9309755f043251ea77fdcf10fbe63989ceb7e19321376",
-			want: WGStats{
-				TxBytes: 38333,
-				RxBytes: 2224,
+			name:       "multiple",
+			peerKey:    "58402e695ba1772b1cc9309755f043251ea77fdcf10fbe63989ceb7e19321376",
+			searchKeys: []string{"tx_bytes", "rx_bytes"},
+			want: map[string]string{
+				"tx_bytes": "38333",
+				"rx_bytes": "2224",
 			},
+			wantErr: false,
 		},
 		{
-			name:    "lastpeer",
-			peerKey: "662e14fd594556f522604703340351258903b64f35553763f19426ab2a515c58",
-			want: WGStats{
-				TxBytes: 1212111,
-				RxBytes: 1929999999,
+			name:       "lastpeer",
+			peerKey:    "662e14fd594556f522604703340351258903b64f35553763f19426ab2a515c58",
+			searchKeys: []string{"tx_bytes", "rx_bytes"},
+			want: map[string]string{
+				"tx_bytes": "1212111",
+				"rx_bytes": "1929999999",
 			},
+			wantErr: false,
+		},
+		{
+			name:       "peer not found",
+			peerKey:    "1111111111111111111111111111111111111111111111111111111111111111",
+			searchKeys: nil,
+			want:       nil,
+			wantErr:    true,
+		},
+		{
+			name:       "key not found",
+			peerKey:    "662e14fd594556f522604703340351258903b64f35553763f19426ab2a515c58",
+			searchKeys: []string{"tx_bytes", "unknown_key"},
+			want: map[string]string{
+				"tx_bytes": "1212111",
+			},
+			wantErr: true,
 		},
 	}
 	for _, tt := range tests {
@@ -71,19 +96,9 @@ func Test_parseTransfers(t *testing.T) {
 			key, err := wgtypes.NewKey(res)
 			require.NoError(t, err)
 
-			stats, err := parseTransfers(ipcFixture)
-			if err != nil {
-				require.NoError(t, err)
-				return
-			}
-
-			stat, ok := stats[key.String()]
-			if !ok {
-				require.True(t, ok)
-				return
-			}
-
-			require.Equal(t, tt.want, stat)
+			got, err := findPeerInfo(ipcFixture, key.String(), tt.searchKeys)
+			assert.Equalf(t, tt.wantErr, err != nil, fmt.Sprintf("findPeerInfo(%v, %v, %v)", ipcFixture, key.String(), tt.searchKeys))
+			assert.Equalf(t, tt.want, got, "findPeerInfo(%v, %v, %v)", ipcFixture, key.String(), tt.searchKeys)
 		})
 	}
 }

client/iface/device/interface.go

@@ -16,5 +16,5 @@ type WGConfigurer interface {
 	AddAllowedIP(peerKey string, allowedIP string) error
 	RemoveAllowedIP(peerKey string, allowedIP string) error
 	Close()
-	GetStats() (map[string]configurer.WGStats, error)
+	GetStats(peerKey string) (configurer.WGStats, error)
 }

client/iface/iface.go

@@ -212,9 +212,9 @@ func (w *WGIface) GetWGDevice() *wgdevice.Device {
 	return w.tun.Device()
 }
 
-// GetStats returns the last handshake time, rx and tx bytes
-func (w *WGIface) GetStats() (map[string]configurer.WGStats, error) {
-	return w.configurer.GetStats()
+// GetStats returns the last handshake time, rx and tx bytes for the given peer
+func (w *WGIface) GetStats(peerKey string) (configurer.WGStats, error) {
+	return w.configurer.GetStats(peerKey)
 }
 
 func (w *WGIface) waitUntilRemoved() error {

client/installer.nsis

@@ -24,8 +24,6 @@
 
 !define AUTOSTART_REG_KEY "Software\Microsoft\Windows\CurrentVersion\Run"
 
-!define NETBIRD_DATA_DIR "$COMMONPROGRAMDATA\Netbird"
-
 Unicode True
 
 ######################################################################
@@ -51,10 +49,6 @@ ShowInstDetails Show
 
 ######################################################################
 
-!include "MUI2.nsh"
-!include LogicLib.nsh
-!include "nsDialogs.nsh"
-
 !define MUI_ICON "${ICON}"
 !define MUI_UNICON "${ICON}"
 !define MUI_WELCOMEFINISHPAGE_BITMAP "${BANNER}"
@@ -64,6 +58,9 @@ ShowInstDetails Show
 !define MUI_FINISHPAGE_RUN_FUNCTION "LaunchLink"
 ######################################################################
 
+!include "MUI2.nsh"
+!include LogicLib.nsh
+
 !define MUI_ABORTWARNING
 !define MUI_UNABORTWARNING
 
@@ -73,16 +70,13 @@ ShowInstDetails Show
 
 !insertmacro MUI_PAGE_DIRECTORY
 
+; Custom page for autostart checkbox
 Page custom AutostartPage AutostartPageLeave
 
 !insertmacro MUI_PAGE_INSTFILES
 
 !insertmacro MUI_PAGE_FINISH
 
-!insertmacro MUI_UNPAGE_WELCOME
-
-UninstPage custom un.DeleteDataPage un.DeleteDataPageLeave
-
 !insertmacro MUI_UNPAGE_CONFIRM
 
 !insertmacro MUI_UNPAGE_INSTFILES
@@ -95,10 +89,6 @@ UninstPage custom un.DeleteDataPage un.DeleteDataPageLeave
 Var AutostartCheckbox
 Var AutostartEnabled
 
-; Variables for uninstall data deletion option
-Var DeleteDataCheckbox
-Var DeleteDataEnabled
-
 ######################################################################
 
 ; Function to create the autostart options page
@@ -114,8 +104,8 @@ Function AutostartPage
 
   ${NSD_CreateCheckbox} 0 20u 100% 10u "Start ${APP_NAME} UI automatically when Windows starts"
   Pop $AutostartCheckbox
-  ${NSD_Check} $AutostartCheckbox
-  StrCpy $AutostartEnabled "1"
+  ${NSD_Check} $AutostartCheckbox ; Default to checked
+  StrCpy $AutostartEnabled "1" ; Default to enabled
 
   nsDialogs::Show
 FunctionEnd
@@ -125,30 +115,6 @@ Function AutostartPageLeave
   ${NSD_GetState} $AutostartCheckbox $AutostartEnabled
 FunctionEnd
 
-; Function to create the uninstall data deletion page
-Function un.DeleteDataPage
-  !insertmacro MUI_HEADER_TEXT "Uninstall Options" "Choose whether to delete ${APP_NAME} data."
-
-  nsDialogs::Create 1018
-  Pop $0
-
-  ${If} $0 == error
-    Abort
-  ${EndIf}
-
-  ${NSD_CreateCheckbox} 0 20u 100% 10u "Delete all ${APP_NAME} configuration and state data (${NETBIRD_DATA_DIR})"
-  Pop $DeleteDataCheckbox
-  ${NSD_Uncheck} $DeleteDataCheckbox
-  StrCpy $DeleteDataEnabled "0"
-
-  nsDialogs::Show
-FunctionEnd
-
-; Function to handle leaving the data deletion page
-Function un.DeleteDataPageLeave
-  ${NSD_GetState} $DeleteDataCheckbox $DeleteDataEnabled
-FunctionEnd
-
 Function GetAppFromCommand
 Exch $1
 Push $2
@@ -210,10 +176,10 @@ ${EndIf}
 FunctionEnd
 ######################################################################
 Section -MainProgram
-    ${INSTALL_TYPE}
-    # SetOverwrite ifnewer
-    SetOutPath "$INSTDIR"
-    File /r "..\\dist\\netbird_windows_amd64\\"
+	${INSTALL_TYPE}
+	# SetOverwrite ifnewer
+	SetOutPath "$INSTDIR"
+	File /r "..\\dist\\netbird_windows_amd64\\"
 SectionEnd
 ######################################################################
 
@@ -259,58 +225,31 @@ SectionEnd
 Section Uninstall
 ${INSTALL_TYPE}
 
-DetailPrint "Stopping Netbird service..."
 ExecWait '"$INSTDIR\${MAIN_APP_EXE}" service stop'
-DetailPrint "Uninstalling Netbird service..."
 ExecWait '"$INSTDIR\${MAIN_APP_EXE}" service uninstall'
 
-DetailPrint "Terminating Netbird UI process..."
+# kill ui client
 ExecWait `taskkill /im ${UI_APP_EXE}.exe /f`
 
 ; Remove autostart registry entry
-DetailPrint "Removing autostart registry entry if exists..."
 DeleteRegValue HKCU "${AUTOSTART_REG_KEY}" "${APP_NAME}"
 
-; Handle data deletion based on checkbox
-DetailPrint "Checking if user requested data deletion..."
-${If} $DeleteDataEnabled == "1"
-    DetailPrint "User opted to delete Netbird data. Removing ${NETBIRD_DATA_DIR}..."
-    ClearErrors
-    RMDir /r "${NETBIRD_DATA_DIR}"
-    IfErrors 0 +2 ; If no errors, jump over the message
-    DetailPrint "Error deleting Netbird data directory. It might be in use or already removed."
-    DetailPrint "Netbird data directory removal complete."
-${Else}
-    DetailPrint "User did not opt to delete Netbird data."
-${EndIf}
-
 # wait the service uninstall take unblock the executable
-DetailPrint "Waiting for service handle to be released..."
 Sleep 3000
-
-DetailPrint "Deleting application files..."
 Delete "$INSTDIR\${UI_APP_EXE}"
 Delete "$INSTDIR\${MAIN_APP_EXE}"
 Delete "$INSTDIR\wintun.dll"
 Delete "$INSTDIR\opengl32.dll"
-DetailPrint "Removing application directory..."
 RmDir /r "$INSTDIR"
 
-DetailPrint "Removing shortcuts..."
 SetShellVarContext all
 Delete "$DESKTOP\${APP_NAME}.lnk"
 Delete "$SMPROGRAMS\${APP_NAME}.lnk"
 
-DetailPrint "Removing registry keys..."
 DeleteRegKey ${REG_ROOT} "${REG_APP_PATH}"
 DeleteRegKey ${REG_ROOT} "${UNINSTALL_PATH}"
-DeleteRegKey ${REG_ROOT} "${UI_REG_APP_PATH}"
-
-DetailPrint "Removing application directory from PATH..."
 EnVar::SetHKLM
 EnVar::DeleteValue "path" "$INSTDIR"
-
-DetailPrint "Uninstallation finished."
 SectionEnd
 
 

client/internal/acl/id/id.go

@@ -18,7 +18,7 @@ func (r RuleID) ID() string {
 
 func GenerateRouteRuleKey(
 	sources []netip.Prefix,
-	destination manager.Network,
+	destination netip.Prefix,
 	proto manager.Protocol,
 	sPort *manager.Port,
 	dPort *manager.Port,

client/internal/acl/manager.go

@@ -18,7 +18,6 @@ import (
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/ssh"
-	"github.com/netbirdio/netbird/management/domain"
 	mgmProto "github.com/netbirdio/netbird/management/proto"
 )
 
@@ -26,7 +25,7 @@ var ErrSourceRangesEmpty = errors.New("sources range is empty")
 
 // Manager is a ACL rules manager
 type Manager interface {
-	ApplyFiltering(networkMap *mgmProto.NetworkMap, dnsRouteFeatureFlag bool)
+	ApplyFiltering(networkMap *mgmProto.NetworkMap)
 }
 
 type protoMatch struct {
@@ -54,7 +53,7 @@ func NewDefaultManager(fm firewall.Manager) *DefaultManager {
 // ApplyFiltering firewall rules to the local firewall manager processed by ACL policy.
 //
 // If allowByDefault is true it appends allow ALL traffic rules to input and output chains.
-func (d *DefaultManager) ApplyFiltering(networkMap *mgmProto.NetworkMap, dnsRouteFeatureFlag bool) {
+func (d *DefaultManager) ApplyFiltering(networkMap *mgmProto.NetworkMap) {
 	d.mutex.Lock()
 	defer d.mutex.Unlock()
 
@@ -76,8 +75,14 @@ func (d *DefaultManager) ApplyFiltering(networkMap *mgmProto.NetworkMap, dnsRout
 
 	d.applyPeerACLs(networkMap)
 
+	// If we got empty rules list but management did not set the networkMap.FirewallRulesIsEmpty flag,
+	// then the mgmt server is older than the client, and we need to allow all traffic for routes
+	isLegacy := len(networkMap.RoutesFirewallRules) == 0 && !networkMap.RoutesFirewallRulesIsEmpty
+	if err := d.firewall.SetLegacyManagement(isLegacy); err != nil {
+		log.Errorf("failed to set legacy management flag: %v", err)
+	}
 
-	if err := d.applyRouteACLs(networkMap.RoutesFirewallRules, dnsRouteFeatureFlag); err != nil {
+	if err := d.applyRouteACLs(networkMap.RoutesFirewallRules); err != nil {
 		log.Errorf("Failed to apply route ACLs: %v", err)
 	}
 
@@ -171,16 +176,16 @@ func (d *DefaultManager) applyPeerACLs(networkMap *mgmProto.NetworkMap) {
 	d.peerRulesPairs = newRulePairs
 }
 
-func (d *DefaultManager) applyRouteACLs(rules []*mgmProto.RouteFirewallRule, dynamicResolver bool) error {
+func (d *DefaultManager) applyRouteACLs(rules []*mgmProto.RouteFirewallRule) error {
 	newRouteRules := make(map[id.RuleID]struct{}, len(rules))
 	var merr *multierror.Error
 
 	// Apply new rules - firewall manager will return existing rule ID if already present
 	for _, rule := range rules {
-		id, err := d.applyRouteACL(rule, dynamicResolver)
+		id, err := d.applyRouteACL(rule)
 		if err != nil {
 			if errors.Is(err, ErrSourceRangesEmpty) {
-				log.Debugf("skipping empty sources rule with destination %s: %v", rule.Destination, err)
+				log.Debugf("skipping empty rule with destination %s: %v", rule.Destination, err)
 			} else {
 				merr = multierror.Append(merr, fmt.Errorf("add route rule: %w", err))
 			}
@@ -203,7 +208,7 @@ func (d *DefaultManager) applyRouteACLs(rules []*mgmProto.RouteFirewallRule, dyn
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-func (d *DefaultManager) applyRouteACL(rule *mgmProto.RouteFirewallRule, dynamicResolver bool) (id.RuleID, error) {
+func (d *DefaultManager) applyRouteACL(rule *mgmProto.RouteFirewallRule) (id.RuleID, error) {
 	if len(rule.SourceRanges) == 0 {
 		return "", ErrSourceRangesEmpty
 	}
@@ -217,9 +222,15 @@ func (d *DefaultManager) applyRouteACL(rule *mgmProto.RouteFirewallRule, dynamic
 		sources = append(sources, source)
 	}
 
-	destination, err := determineDestination(rule, dynamicResolver, sources)
-	if err != nil {
-		return "", fmt.Errorf("determine destination: %w", err)
+	var destination netip.Prefix
+	if rule.IsDynamic {
+		destination = getDefault(sources[0])
+	} else {
+		var err error
+		destination, err = netip.ParsePrefix(rule.Destination)
+		if err != nil {
+			return "", fmt.Errorf("parse destination: %w", err)
+		}
 	}
 
 	protocol, err := convertToFirewallProtocol(rule.Protocol)
@@ -569,33 +580,6 @@ func convertPortInfo(portInfo *mgmProto.PortInfo) *firewall.Port {
 	return nil
 }
 
-func determineDestination(rule *mgmProto.RouteFirewallRule, dynamicResolver bool, sources []netip.Prefix) (firewall.Network, error) {
-	var destination firewall.Network
-
-	if rule.IsDynamic {
-		if dynamicResolver {
-			if len(rule.Domains) > 0 {
-				destination.Set = firewall.NewDomainSet(domain.FromPunycodeList(rule.Domains))
-			} else {
-				// isDynamic is set but no domains = outdated management server
-				log.Warn("connected to an older version of management server (no domains in rules), using default destination")
-				destination.Prefix = getDefault(sources[0])
-			}
-		} else {
-			// client resolves DNS, we (router) don't know the destination
-			destination.Prefix = getDefault(sources[0])
-		}
-		return destination, nil
-	}
-
-	prefix, err := netip.ParsePrefix(rule.Destination)
-	if err != nil {
-		return destination, fmt.Errorf("parse destination: %w", err)
-	}
-	destination.Prefix = prefix
-	return destination, nil
-}
-
 func getDefault(prefix netip.Prefix) netip.Prefix {
 	if prefix.Addr().Is6() {
 		return netip.PrefixFrom(netip.IPv6Unspecified(), 0)

client/internal/acl/manager_test.go

@@ -66,7 +66,7 @@ func TestDefaultManager(t *testing.T) {
 	acl := NewDefaultManager(fw)
 
 	t.Run("apply firewall rules", func(t *testing.T) {
-		acl.ApplyFiltering(networkMap, false)
+		acl.ApplyFiltering(networkMap)
 
 		if len(acl.peerRulesPairs) != 2 {
 			t.Errorf("firewall rules not applied: %v", acl.peerRulesPairs)
@@ -92,7 +92,7 @@ func TestDefaultManager(t *testing.T) {
 			},
 		)
 
-		acl.ApplyFiltering(networkMap, false)
+		acl.ApplyFiltering(networkMap)
 
 		// we should have one old and one new rule in the existed rules
 		if len(acl.peerRulesPairs) != 2 {
@@ -116,13 +116,13 @@ func TestDefaultManager(t *testing.T) {
 		networkMap.FirewallRules = networkMap.FirewallRules[:0]
 
 		networkMap.FirewallRulesIsEmpty = true
-		if acl.ApplyFiltering(networkMap, false); len(acl.peerRulesPairs) != 0 {
+		if acl.ApplyFiltering(networkMap); len(acl.peerRulesPairs) != 0 {
 			t.Errorf("rules should be empty if FirewallRulesIsEmpty is set, got: %v", len(acl.peerRulesPairs))
 			return
 		}
 
 		networkMap.FirewallRulesIsEmpty = false
-		acl.ApplyFiltering(networkMap, false)
+		acl.ApplyFiltering(networkMap)
 		if len(acl.peerRulesPairs) != 1 {
 			t.Errorf("rules should contain 1 rules if FirewallRulesIsEmpty is not set, got: %v", len(acl.peerRulesPairs))
 			return
@@ -359,7 +359,7 @@ func TestDefaultManagerEnableSSHRules(t *testing.T) {
 	}(fw)
 	acl := NewDefaultManager(fw)
 
-	acl.ApplyFiltering(networkMap, false)
+	acl.ApplyFiltering(networkMap)
 
 	if len(acl.peerRulesPairs) != 3 {
 		t.Errorf("expect 3 rules (last must be SSH), got: %d", len(acl.peerRulesPairs))

client/internal/auth/oauth.go

@@ -64,8 +64,13 @@ func (t TokenInfo) GetTokenToUse() string {
 // and if that also fails, the authentication process is deemed unsuccessful
 //
 // On Linux distros without desktop environment support, it only tries to initialize the Device Code Flow
-func NewOAuthFlow(ctx context.Context, config *internal.Config, isUnixDesktopClient bool) (OAuthFlow, error) {
-	if (runtime.GOOS == "linux" || runtime.GOOS == "freebsd") && !isUnixDesktopClient {
+func NewOAuthFlow(ctx context.Context, config *internal.Config, isLinuxDesktopClient bool) (OAuthFlow, error) {
+	if runtime.GOOS == "linux" && !isLinuxDesktopClient {
+		return authenticateWithDeviceCodeFlow(ctx, config)
+	}
+
+	// On FreeBSD we currently do not support desktop environments and offer only Device Code Flow (#2384)
+	if runtime.GOOS == "freebsd" {
 		return authenticateWithDeviceCodeFlow(ctx, config)
 	}
 

client/internal/auth/pkce_flow.go

@@ -101,12 +101,7 @@ func (p *PKCEAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowIn
 		oauth2.SetAuthURLParam("audience", p.providerConfig.Audience),
 	}
 	if !p.providerConfig.DisablePromptLogin {
-		if p.providerConfig.LoginFlag.IsPromptLogin() {
-			params = append(params, oauth2.SetAuthURLParam("prompt", "login"))
-		}
-		if p.providerConfig.LoginFlag.IsMaxAge0Login() {
-			params = append(params, oauth2.SetAuthURLParam("max_age", "0"))
-		}
+		params = append(params, oauth2.SetAuthURLParam("prompt", "login"))
 	}
 
 	authURL := p.oAuthConfig.AuthCodeURL(state, params...)

client/internal/auth/pkce_flow_test.go

@@ -7,36 +7,15 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/client/internal"
-	mgm "github.com/netbirdio/netbird/management/client/common"
 )
 
 func TestPromptLogin(t *testing.T) {
-	const (
-		promptLogin = "prompt=login"
-		maxAge0     = "max_age=0"
-	)
-
 	tt := []struct {
-		name               string
-		loginFlag          mgm.LoginFlag
-		disablePromptLogin bool
-		expect             string
+		name   string
+		prompt bool
 	}{
-		{
-			name:      "Prompt login",
-			loginFlag: mgm.LoginFlagPrompt,
-			expect:    promptLogin,
-		},
-		{
-			name:      "Max age 0 login",
-			loginFlag: mgm.LoginFlagMaxAge0,
-			expect:    maxAge0,
-		},
-		{
-			name:               "Disable prompt login",
-			loginFlag:          mgm.LoginFlagPrompt,
-			disablePromptLogin: true,
-		},
+		{"PromptLogin", true},
+		{"NoPromptLogin", false},
 	}
 
 	for _, tc := range tt {
@@ -49,7 +28,7 @@ func TestPromptLogin(t *testing.T) {
 				AuthorizationEndpoint: "https://test-auth-endpoint.com/authorize",
 				RedirectURLs:          []string{"http://127.0.0.1:33992/"},
 				UseIDToken:            true,
-				LoginFlag:             tc.loginFlag,
+				DisablePromptLogin:    !tc.prompt,
 			}
 			pkce, err := NewPKCEAuthorizationFlow(config)
 			if err != nil {
@@ -59,12 +38,11 @@ func TestPromptLogin(t *testing.T) {
 			if err != nil {
 				t.Fatalf("Failed to request auth info: %v", err)
 			}
-
-			if !tc.disablePromptLogin {
-				require.Contains(t, authInfo.VerificationURIComplete, tc.expect)
+			pattern := "prompt=login"
+			if tc.prompt {
+				require.Contains(t, authInfo.VerificationURIComplete, pattern)
 			} else {
-				require.Contains(t, authInfo.VerificationURIComplete, promptLogin)
-				require.NotContains(t, authInfo.VerificationURIComplete, maxAge0)
+				require.NotContains(t, authInfo.VerificationURIComplete, pattern)
 			}
 		})
 	}

client/internal/config.go

@@ -74,8 +74,6 @@ type ConfigInput struct {
 	DisableNotifications *bool
 
 	DNSLabels domain.List
-
-	LazyConnectionEnabled *bool
 }
 
 // Config Configuration type
@@ -140,8 +138,6 @@ type Config struct {
 	ClientCertKeyPath string
 
 	ClientCertKeyPair *tls.Certificate `json:"-"`
-
-	LazyConnectionEnabled bool
 }
 
 // ReadConfig read config file and return with Config. If it is not exists create a new with default values
@@ -528,12 +524,6 @@ func (config *Config) apply(input ConfigInput) (updated bool, err error) {
 		updated = true
 	}
 
-	if input.LazyConnectionEnabled != nil && *input.LazyConnectionEnabled != config.LazyConnectionEnabled {
-		log.Infof("switching lazy connection to %t", *input.LazyConnectionEnabled)
-		config.LazyConnectionEnabled = *input.LazyConnectionEnabled
-		updated = true
-	}
-
 	return updated, nil
 }
 

client/internal/conn_mgr.go

@@ -1,303 +0,0 @@
-package internal
-
-import (
-	"context"
-	"os"
-	"strconv"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-	"github.com/netbirdio/netbird/client/internal/lazyconn/manager"
-	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
-	"github.com/netbirdio/netbird/client/internal/peerstore"
-)
-
-// ConnMgr coordinates both lazy connections (established on-demand) and permanent peer connections.
-//
-// The connection manager is responsible for:
-// - Managing lazy connections via the lazyConnManager
-// - Maintaining a list of excluded peers that should always have permanent connections
-// - Handling connection establishment based on peer signaling
-//
-// The implementation is not thread-safe; it is protected by engine.syncMsgMux.
-type ConnMgr struct {
-	peerStore      *peerstore.Store
-	statusRecorder *peer.Status
-	iface          lazyconn.WGIface
-	dispatcher     *dispatcher.ConnectionDispatcher
-	enabledLocally bool
-
-	lazyConnMgr *manager.Manager
-
-	wg        sync.WaitGroup
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-}
-
-func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface, dispatcher *dispatcher.ConnectionDispatcher) *ConnMgr {
-	e := &ConnMgr{
-		peerStore:      peerStore,
-		statusRecorder: statusRecorder,
-		iface:          iface,
-		dispatcher:     dispatcher,
-	}
-	if engineConfig.LazyConnectionEnabled || lazyconn.IsLazyConnEnabledByEnv() {
-		e.enabledLocally = true
-	}
-	return e
-}
-
-// Start initializes the connection manager and starts the lazy connection manager if enabled by env var or cmd line option.
-func (e *ConnMgr) Start(ctx context.Context) {
-	if e.lazyConnMgr != nil {
-		log.Errorf("lazy connection manager is already started")
-		return
-	}
-
-	if !e.enabledLocally {
-		log.Infof("lazy connection manager is disabled")
-		return
-	}
-
-	e.initLazyManager(ctx)
-	e.statusRecorder.UpdateLazyConnection(true)
-}
-
-// UpdatedRemoteFeatureFlag is called when the remote feature flag is updated.
-// If enabled, it initializes the lazy connection manager and start it. Do not need to call Start() again.
-// If disabled, then it closes the lazy connection manager and open the connections to all peers.
-func (e *ConnMgr) UpdatedRemoteFeatureFlag(ctx context.Context, enabled bool) error {
-	// do not disable lazy connection manager if it was enabled by env var
-	if e.enabledLocally {
-		return nil
-	}
-
-	if enabled {
-		// if the lazy connection manager is already started, do not start it again
-		if e.lazyConnMgr != nil {
-			return nil
-		}
-
-		log.Infof("lazy connection manager is enabled by management feature flag")
-		e.initLazyManager(ctx)
-		e.statusRecorder.UpdateLazyConnection(true)
-		return e.addPeersToLazyConnManager(ctx)
-	} else {
-		if e.lazyConnMgr == nil {
-			return nil
-		}
-		log.Infof("lazy connection manager is disabled by management feature flag")
-		e.closeManager(ctx)
-		e.statusRecorder.UpdateLazyConnection(false)
-		return nil
-	}
-}
-
-// SetExcludeList sets the list of peer IDs that should always have permanent connections.
-func (e *ConnMgr) SetExcludeList(peerIDs map[string]bool) {
-	if e.lazyConnMgr == nil {
-		return
-	}
-
-	excludedPeers := make([]lazyconn.PeerConfig, 0, len(peerIDs))
-
-	for peerID := range peerIDs {
-		var peerConn *peer.Conn
-		var exists bool
-		if peerConn, exists = e.peerStore.PeerConn(peerID); !exists {
-			log.Warnf("failed to find peer conn for peerID: %s", peerID)
-			continue
-		}
-
-		lazyPeerCfg := lazyconn.PeerConfig{
-			PublicKey:  peerID,
-			AllowedIPs: peerConn.WgConfig().AllowedIps,
-			PeerConnID: peerConn.ConnID(),
-			Log:        peerConn.Log,
-		}
-		excludedPeers = append(excludedPeers, lazyPeerCfg)
-	}
-
-	added := e.lazyConnMgr.ExcludePeer(e.ctx, excludedPeers)
-	for _, peerID := range added {
-		var peerConn *peer.Conn
-		var exists bool
-		if peerConn, exists = e.peerStore.PeerConn(peerID); !exists {
-			// if the peer not exist in the store, it means that the engine will call the AddPeerConn in next step
-			continue
-		}
-
-		peerConn.Log.Infof("peer has been added to lazy connection exclude list, opening permanent connection")
-		if err := peerConn.Open(e.ctx); err != nil {
-			peerConn.Log.Errorf("failed to open connection: %v", err)
-		}
-	}
-}
-
-func (e *ConnMgr) AddPeerConn(ctx context.Context, peerKey string, conn *peer.Conn) (exists bool) {
-	if success := e.peerStore.AddPeerConn(peerKey, conn); !success {
-		return true
-	}
-
-	if !e.isStartedWithLazyMgr() {
-		if err := conn.Open(ctx); err != nil {
-			conn.Log.Errorf("failed to open connection: %v", err)
-		}
-		return
-	}
-
-	if !lazyconn.IsSupported(conn.AgentVersionString()) {
-		conn.Log.Warnf("peer does not support lazy connection (%s), open permanent connection", conn.AgentVersionString())
-		if err := conn.Open(ctx); err != nil {
-			conn.Log.Errorf("failed to open connection: %v", err)
-		}
-		return
-	}
-
-	lazyPeerCfg := lazyconn.PeerConfig{
-		PublicKey:  peerKey,
-		AllowedIPs: conn.WgConfig().AllowedIps,
-		PeerConnID: conn.ConnID(),
-		Log:        conn.Log,
-	}
-	excluded, err := e.lazyConnMgr.AddPeer(lazyPeerCfg)
-	if err != nil {
-		conn.Log.Errorf("failed to add peer to lazyconn manager: %v", err)
-		if err := conn.Open(ctx); err != nil {
-			conn.Log.Errorf("failed to open connection: %v", err)
-		}
-		return
-	}
-
-	if excluded {
-		conn.Log.Infof("peer is on lazy conn manager exclude list, opening connection")
-		if err := conn.Open(ctx); err != nil {
-			conn.Log.Errorf("failed to open connection: %v", err)
-		}
-		return
-	}
-
-	conn.Log.Infof("peer added to lazy conn manager")
-	return
-}
-
-func (e *ConnMgr) RemovePeerConn(peerKey string) {
-	conn, ok := e.peerStore.Remove(peerKey)
-	if !ok {
-		return
-	}
-	defer conn.Close()
-
-	if !e.isStartedWithLazyMgr() {
-		return
-	}
-
-	e.lazyConnMgr.RemovePeer(peerKey)
-	conn.Log.Infof("removed peer from lazy conn manager")
-}
-
-func (e *ConnMgr) OnSignalMsg(ctx context.Context, peerKey string) (*peer.Conn, bool) {
-	conn, ok := e.peerStore.PeerConn(peerKey)
-	if !ok {
-		return nil, false
-	}
-
-	if !e.isStartedWithLazyMgr() {
-		return conn, true
-	}
-
-	if found := e.lazyConnMgr.ActivatePeer(ctx, peerKey); found {
-		conn.Log.Infof("activated peer from inactive state")
-		if err := conn.Open(e.ctx); err != nil {
-			conn.Log.Errorf("failed to open connection: %v", err)
-		}
-	}
-	return conn, true
-}
-
-func (e *ConnMgr) Close() {
-	if !e.isStartedWithLazyMgr() {
-		return
-	}
-
-	e.ctxCancel()
-	e.wg.Wait()
-	e.lazyConnMgr = nil
-}
-
-func (e *ConnMgr) initLazyManager(parentCtx context.Context) {
-	cfg := manager.Config{
-		InactivityThreshold: inactivityThresholdEnv(),
-	}
-	e.lazyConnMgr = manager.NewManager(cfg, e.peerStore, e.iface, e.dispatcher)
-
-	ctx, cancel := context.WithCancel(parentCtx)
-	e.ctx = ctx
-	e.ctxCancel = cancel
-
-	e.wg.Add(1)
-	go func() {
-		defer e.wg.Done()
-		e.lazyConnMgr.Start(ctx)
-	}()
-}
-
-func (e *ConnMgr) addPeersToLazyConnManager(ctx context.Context) error {
-	peers := e.peerStore.PeersPubKey()
-	lazyPeerCfgs := make([]lazyconn.PeerConfig, 0, len(peers))
-	for _, peerID := range peers {
-		var peerConn *peer.Conn
-		var exists bool
-		if peerConn, exists = e.peerStore.PeerConn(peerID); !exists {
-			log.Warnf("failed to find peer conn for peerID: %s", peerID)
-			continue
-		}
-
-		lazyPeerCfg := lazyconn.PeerConfig{
-			PublicKey:  peerID,
-			AllowedIPs: peerConn.WgConfig().AllowedIps,
-			PeerConnID: peerConn.ConnID(),
-			Log:        peerConn.Log,
-		}
-		lazyPeerCfgs = append(lazyPeerCfgs, lazyPeerCfg)
-	}
-
-	return e.lazyConnMgr.AddActivePeers(ctx, lazyPeerCfgs)
-}
-
-func (e *ConnMgr) closeManager(ctx context.Context) {
-	if e.lazyConnMgr == nil {
-		return
-	}
-
-	e.ctxCancel()
-	e.wg.Wait()
-	e.lazyConnMgr = nil
-
-	for _, peerID := range e.peerStore.PeersPubKey() {
-		e.peerStore.PeerConnOpen(ctx, peerID)
-	}
-}
-
-func (e *ConnMgr) isStartedWithLazyMgr() bool {
-	return e.lazyConnMgr != nil && e.ctxCancel != nil
-}
-
-func inactivityThresholdEnv() *time.Duration {
-	envValue := os.Getenv(lazyconn.EnvInactivityThreshold)
-	if envValue == "" {
-		return nil
-	}
-
-	parsedMinutes, err := strconv.Atoi(envValue)
-	if err != nil || parsedMinutes <= 0 {
-		return nil
-	}
-
-	d := time.Duration(parsedMinutes) * time.Minute
-	return &d
-}

client/internal/connect.go

@@ -440,8 +440,7 @@ func createEngineConfig(key wgtypes.Key, config *Config, peerConfig *mgmProto.Pe
 		DisableDNS:          config.DisableDNS,
 		DisableFirewall:     config.DisableFirewall,
 
-		BlockLANAccess:        config.BlockLANAccess,
-		LazyConnectionEnabled: config.LazyConnectionEnabled,
+		BlockLANAccess: config.BlockLANAccess,
 	}
 
 	if config.PreSharedKey != "" {
@@ -482,7 +481,7 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourP
 	return signalClient, nil
 }
 
-// loginToManagement creates Management ServiceDependencies client, establishes a connection, logs-in and gets a global Netbird config (signal, turn, stun hosts, etc)
+// loginToManagement creates Management Services client, establishes a connection, logs-in and gets a global Netbird config (signal, turn, stun hosts, etc)
 func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte, config *Config) (*mgmProto.LoginResponse, error) {
 
 	serverPublicKey, err := client.GetServerPublicKey()

client/internal/debug/debug.go

@@ -4,7 +4,6 @@ import (
 	"archive/zip"
 	"bufio"
 	"bytes"
-	"compress/gzip"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -377,7 +376,6 @@ func (g *BundleGenerator) addCommonConfigFields(configContent *strings.Builder)
 	configContent.WriteString(fmt.Sprintf("DisableFirewall: %v\n", g.internalConfig.DisableFirewall))
 
 	configContent.WriteString(fmt.Sprintf("BlockLANAccess: %v\n", g.internalConfig.BlockLANAccess))
-	configContent.WriteString(fmt.Sprintf("LazyConnectionEnabled: %v\n", g.internalConfig.LazyConnectionEnabled))
 }
 
 func (g *BundleGenerator) addProf() (err error) {
@@ -535,33 +533,6 @@ 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)
-		}
-	}
-
 	stdErrLogPath := filepath.Join(logDir, errorLogFile)
 	stdoutLogPath := filepath.Join(logDir, stdoutLogFile)
 	if runtime.GOOS == "darwin" {
@@ -592,13 +563,16 @@ func (g *BundleGenerator) addSingleLogfile(logPath, targetName string) error {
 		}
 	}()
 
-	var logReader io.Reader = logFile
+	var logReader io.Reader
 	if g.anonymize {
 		var writer *io.PipeWriter
 		logReader, writer = io.Pipe()
 
 		go anonymizeLog(logFile, writer, g.anonymizer)
+	} else {
+		logReader = logFile
 	}
+
 	if err := g.addFileToZip(logReader, targetName); err != nil {
 		return fmt.Errorf("add %s to zip: %w", targetName, err)
 	}
@@ -606,44 +580,6 @@ func (g *BundleGenerator) addSingleLogfile(logPath, targetName string) error {
 	return nil
 }
 
-// addSingleLogFileGz adds a single gzipped log file to the archive
-func (g *BundleGenerator) addSingleLogFileGz(logPath, targetName string) error {
-	f, err := os.Open(logPath)
-	if err != nil {
-		return fmt.Errorf("open gz log file %s: %w", targetName, err)
-	}
-	defer f.Close()
-
-	gzr, err := gzip.NewReader(f)
-	if err != nil {
-		return fmt.Errorf("create gzip reader: %w", err)
-	}
-	defer gzr.Close()
-
-	var logReader io.Reader = gzr
-	if g.anonymize {
-		var pw *io.PipeWriter
-		logReader, pw = io.Pipe()
-		go anonymizeLog(gzr, pw, g.anonymizer)
-	}
-
-	var buf bytes.Buffer
-	gw := gzip.NewWriter(&buf)
-	if _, err := io.Copy(gw, logReader); err != nil {
-		return fmt.Errorf("re-gzip: %w", err)
-	}
-
-	if err := gw.Close(); err != nil {
-		return fmt.Errorf("close gzip writer: %w", err)
-	}
-
-	if err := g.addFileToZip(&buf, targetName); err != nil {
-		return fmt.Errorf("add anonymized gz: %w", err)
-	}
-
-	return nil
-}
-
 func (g *BundleGenerator) addFileToZip(reader io.Reader, filename string) error {
 	header := &zip.FileHeader{
 		Name:     filename,

client/internal/debug/debug_linux.go

@@ -59,16 +59,6 @@ func collectIPTablesRules() (string, error) {
 		builder.WriteString("\n")
 	}
 
-	// Collect ipset information
-	ipsetOutput, err := collectIPSets()
-	if err != nil {
-		log.Warnf("Failed to collect ipset information: %v", err)
-	} else {
-		builder.WriteString("=== ipset list output ===\n")
-		builder.WriteString(ipsetOutput)
-		builder.WriteString("\n")
-	}
-
 	builder.WriteString("=== iptables -v -n -L output ===\n")
 
 	tables := []string{"filter", "nat", "mangle", "raw", "security"}
@@ -88,28 +78,6 @@ func collectIPTablesRules() (string, error) {
 	return builder.String(), nil
 }
 
-// collectIPSets collects information about ipsets
-func collectIPSets() (string, error) {
-	cmd := exec.Command("ipset", "list")
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	if err := cmd.Run(); err != nil {
-		if strings.Contains(err.Error(), "executable file not found") {
-			return "", fmt.Errorf("ipset command not found: %w", err)
-		}
-		return "", fmt.Errorf("execute ipset list: %w (stderr: %s)", err, stderr.String())
-	}
-
-	ipsets := stdout.String()
-	if strings.TrimSpace(ipsets) == "" {
-		return "No ipsets found", nil
-	}
-
-	return ipsets, nil
-}
-
 // collectIPTablesSave uses iptables-save to get rule definitions
 func collectIPTablesSave() (string, error) {
 	cmd := exec.Command("iptables-save")

client/internal/dns/handler_chain_test.go

@@ -1,6 +1,7 @@
 package dns_test
 
 import (
+	"net"
 	"testing"
 
 	"github.com/miekg/dns"
@@ -8,7 +9,6 @@ import (
 	"github.com/stretchr/testify/mock"
 
 	nbdns "github.com/netbirdio/netbird/client/internal/dns"
-	"github.com/netbirdio/netbird/client/internal/dns/test"
 )
 
 // TestHandlerChain_ServeDNS_Priorities tests that handlers are executed in priority order
@@ -30,7 +30,7 @@ func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
 	r.SetQuestion("example.com.", dns.TypeA)
 
 	// Create test writer
-	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 	// Setup expectations - only highest priority handler should be called
 	dnsRouteHandler.On("ServeDNS", mock.Anything, r).Once()
@@ -142,7 +142,7 @@ func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
 
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryDomain, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 			chain.ServeDNS(w, r)
 
@@ -259,7 +259,7 @@ func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
 			// Create and execute request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryDomain, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 			chain.ServeDNS(w, r)
 
 			// Verify expectations
@@ -316,7 +316,7 @@ func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
 	}).Once()
 
 	// Execute
-	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	chain.ServeDNS(w, r)
 
 	// Verify all handlers were called in order
@@ -325,6 +325,20 @@ func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
 	handler3.AssertExpectations(t)
 }
 
+// mockResponseWriter implements dns.ResponseWriter for testing
+type mockResponseWriter struct {
+	mock.Mock
+}
+
+func (m *mockResponseWriter) LocalAddr() net.Addr       { return nil }
+func (m *mockResponseWriter) RemoteAddr() net.Addr      { return nil }
+func (m *mockResponseWriter) WriteMsg(*dns.Msg) error   { return nil }
+func (m *mockResponseWriter) Write([]byte) (int, error) { return 0, nil }
+func (m *mockResponseWriter) Close() error              { return nil }
+func (m *mockResponseWriter) TsigStatus() error         { return nil }
+func (m *mockResponseWriter) TsigTimersOnly(bool)       {}
+func (m *mockResponseWriter) Hijack()                   {}
+
 func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 	tests := []struct {
 		name string
@@ -411,7 +425,7 @@ func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 			// Create test request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 			// Setup expectations
 			for priority, handler := range handlers {
@@ -457,7 +471,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain)
 
 	// Test 1: Initial state
-	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Highest priority handler (routeHandler) should be called
 	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 	matchHandler.On("ServeDNS", mock.Anything, r).Maybe()   // Ensure others are not expected yet
@@ -476,7 +490,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 2: Remove highest priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)
 
-	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Now middle priority handler (matchHandler) should be called
 	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure default is not expected yet
@@ -492,7 +506,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 3: Remove middle priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityMatchDomain)
 
-	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Now lowest priority handler (defaultHandler) should be called
 	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 
@@ -505,7 +519,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 4: Remove last handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
 
-	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	chain.ServeDNS(w4, r) // Call ServeDNS on the now empty chain for this domain
 
 	for _, m := range mocks {
@@ -661,7 +675,7 @@ func TestHandlerChain_CaseSensitivity(t *testing.T) {
 			// Execute request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			chain.ServeDNS(&test.MockResponseWriter{}, r)
+			chain.ServeDNS(&mockResponseWriter{}, r)
 
 			// Verify each handler was called exactly as expected
 			for _, h := range tt.addHandlers {
@@ -805,7 +819,7 @@ func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) {
 
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 			// Setup handler expectations
 			for pattern, handler := range handlers {
@@ -955,7 +969,7 @@ func TestHandlerChain_AddRemoveRoundtrip(t *testing.T) {
 			handler := &nbdns.MockHandler{}
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryPattern, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 			// First verify no handler is called before adding any
 			chain.ServeDNS(w, r)

client/internal/dns/local.go

@@ -0,0 +1,130 @@
+package dns
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+type registrationMap map[string]struct{}
+
+type localResolver struct {
+	registeredMap registrationMap
+	records       sync.Map // key: string (domain_class_type), value: []dns.RR
+}
+
+func (d *localResolver) MatchSubdomains() bool {
+	return true
+}
+
+func (d *localResolver) stop() {
+}
+
+// String returns a string representation of the local resolver
+func (d *localResolver) String() string {
+	return fmt.Sprintf("local resolver [%d records]", len(d.registeredMap))
+}
+
+// ID returns the unique handler ID
+func (d *localResolver) id() handlerID {
+	return "local-resolver"
+}
+
+// ServeDNS handles a DNS request
+func (d *localResolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	if len(r.Question) > 0 {
+		log.Tracef("received local question: domain=%s type=%v class=%v", r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
+	}
+
+	replyMessage := &dns.Msg{}
+	replyMessage.SetReply(r)
+	replyMessage.RecursionAvailable = true
+
+	// lookup all records matching the question
+	records := d.lookupRecords(r)
+	if len(records) > 0 {
+		replyMessage.Rcode = dns.RcodeSuccess
+		replyMessage.Answer = append(replyMessage.Answer, records...)
+	} else {
+		replyMessage.Rcode = dns.RcodeNameError
+	}
+
+	err := w.WriteMsg(replyMessage)
+	if err != nil {
+		log.Debugf("got an error while writing the local resolver response, error: %v", err)
+	}
+}
+
+// lookupRecords fetches *all* DNS records matching the first question in r.
+func (d *localResolver) lookupRecords(r *dns.Msg) []dns.RR {
+	if len(r.Question) == 0 {
+		return nil
+	}
+	question := r.Question[0]
+	question.Name = strings.ToLower(question.Name)
+	key := buildRecordKey(question.Name, question.Qclass, question.Qtype)
+
+	value, found := d.records.Load(key)
+	if !found {
+		// alternatively check if we have a cname
+		if question.Qtype != dns.TypeCNAME {
+			r.Question[0].Qtype = dns.TypeCNAME
+			return d.lookupRecords(r)
+		}
+
+		return nil
+	}
+
+	records, ok := value.([]dns.RR)
+	if !ok {
+		log.Errorf("failed to cast records to []dns.RR, records: %v", value)
+		return nil
+	}
+
+	// if there's more than one record, rotate them (round-robin)
+	if len(records) > 1 {
+		first := records[0]
+		records = append(records[1:], first)
+		d.records.Store(key, records)
+	}
+
+	return records
+}
+
+// registerRecord stores a new record by appending it to any existing list
+func (d *localResolver) registerRecord(record nbdns.SimpleRecord) (string, error) {
+	rr, err := dns.NewRR(record.String())
+	if err != nil {
+		return "", fmt.Errorf("register record: %w", err)
+	}
+
+	rr.Header().Rdlength = record.Len()
+	header := rr.Header()
+	key := buildRecordKey(header.Name, header.Class, header.Rrtype)
+
+	// load any existing slice of records, then append
+	existing, _ := d.records.LoadOrStore(key, []dns.RR{})
+	records := existing.([]dns.RR)
+	records = append(records, rr)
+
+	// store updated slice
+	d.records.Store(key, records)
+	return key, nil
+}
+
+// deleteRecord removes *all* records under the recordKey.
+func (d *localResolver) deleteRecord(recordKey string) {
+	d.records.Delete(dns.Fqdn(recordKey))
+}
+
+// buildRecordKey consistently generates a key: name_class_type
+func buildRecordKey(name string, class, qType uint16) string {
+	return fmt.Sprintf("%s_%d_%d", dns.Fqdn(name), class, qType)
+}
+
+func (d *localResolver) probeAvailability() {}

client/internal/dns/local/local.go

@@ -1,149 +0,0 @@
-package local
-
-import (
-	"fmt"
-	"slices"
-	"strings"
-	"sync"
-
-	"github.com/miekg/dns"
-	log "github.com/sirupsen/logrus"
-	"golang.org/x/exp/maps"
-
-	"github.com/netbirdio/netbird/client/internal/dns/types"
-	nbdns "github.com/netbirdio/netbird/dns"
-)
-
-type Resolver struct {
-	mu      sync.RWMutex
-	records map[dns.Question][]dns.RR
-}
-
-func NewResolver() *Resolver {
-	return &Resolver{
-		records: make(map[dns.Question][]dns.RR),
-	}
-}
-
-func (d *Resolver) MatchSubdomains() bool {
-	return true
-}
-
-// String returns a string representation of the local resolver
-func (d *Resolver) String() string {
-	return fmt.Sprintf("local resolver [%d records]", len(d.records))
-}
-
-func (d *Resolver) Stop() {}
-
-// ID returns the unique handler ID
-func (d *Resolver) ID() types.HandlerID {
-	return "local-resolver"
-}
-
-func (d *Resolver) ProbeAvailability() {}
-
-// ServeDNS handles a DNS request
-func (d *Resolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
-	if len(r.Question) == 0 {
-		log.Debugf("received local resolver request with no question")
-		return
-	}
-	question := r.Question[0]
-	question.Name = strings.ToLower(dns.Fqdn(question.Name))
-
-	log.Tracef("received local question: domain=%s type=%v class=%v", r.Question[0].Name, question.Qtype, question.Qclass)
-
-	replyMessage := &dns.Msg{}
-	replyMessage.SetReply(r)
-	replyMessage.RecursionAvailable = true
-
-	// lookup all records matching the question
-	records := d.lookupRecords(question)
-	if len(records) > 0 {
-		replyMessage.Rcode = dns.RcodeSuccess
-		replyMessage.Answer = append(replyMessage.Answer, records...)
-	} else {
-		// TODO: return success if we have a different record type for the same name, relevant for search domains
-		replyMessage.Rcode = dns.RcodeNameError
-	}
-
-	if err := w.WriteMsg(replyMessage); err != nil {
-		log.Warnf("failed to write the local resolver response: %v", err)
-	}
-}
-
-// lookupRecords fetches *all* DNS records matching the first question in r.
-func (d *Resolver) lookupRecords(question dns.Question) []dns.RR {
-	d.mu.RLock()
-	records, found := d.records[question]
-
-	if !found {
-		d.mu.RUnlock()
-		// alternatively check if we have a cname
-		if question.Qtype != dns.TypeCNAME {
-			question.Qtype = dns.TypeCNAME
-			return d.lookupRecords(question)
-		}
-		return nil
-	}
-
-	recordsCopy := slices.Clone(records)
-	d.mu.RUnlock()
-
-	// if there's more than one record, rotate them (round-robin)
-	if len(recordsCopy) > 1 {
-		d.mu.Lock()
-		records = d.records[question]
-		if len(records) > 1 {
-			first := records[0]
-			records = append(records[1:], first)
-			d.records[question] = records
-		}
-		d.mu.Unlock()
-	}
-
-	return recordsCopy
-}
-
-func (d *Resolver) Update(update []nbdns.SimpleRecord) {
-	d.mu.Lock()
-	defer d.mu.Unlock()
-
-	maps.Clear(d.records)
-
-	for _, rec := range update {
-		if err := d.registerRecord(rec); err != nil {
-			log.Warnf("failed to register the record (%s): %v", rec, err)
-			continue
-		}
-	}
-}
-
-// RegisterRecord stores a new record by appending it to any existing list
-func (d *Resolver) RegisterRecord(record nbdns.SimpleRecord) error {
-	d.mu.Lock()
-	defer d.mu.Unlock()
-
-	return d.registerRecord(record)
-}
-
-// registerRecord performs the registration with the lock already held
-func (d *Resolver) registerRecord(record nbdns.SimpleRecord) error {
-	rr, err := dns.NewRR(record.String())
-	if err != nil {
-		return fmt.Errorf("register record: %w", err)
-	}
-
-	rr.Header().Rdlength = record.Len()
-	header := rr.Header()
-	q := dns.Question{
-		Name:   strings.ToLower(dns.Fqdn(header.Name)),
-		Qtype:  header.Rrtype,
-		Qclass: header.Class,
-	}
-
-	d.records[q] = append(d.records[q], rr)
-
-	return nil
-}

client/internal/dns/local/local_test.go

@@ -1,472 +0,0 @@
-package local
-
-import (
-	"strings"
-	"testing"
-
-	"github.com/miekg/dns"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/internal/dns/test"
-	nbdns "github.com/netbirdio/netbird/dns"
-)
-
-func TestLocalResolver_ServeDNS(t *testing.T) {
-	recordA := nbdns.SimpleRecord{
-		Name:  "peera.netbird.cloud.",
-		Type:  1,
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "1.2.3.4",
-	}
-
-	recordCNAME := nbdns.SimpleRecord{
-		Name:  "peerb.netbird.cloud.",
-		Type:  5,
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "www.netbird.io",
-	}
-
-	testCases := []struct {
-		name                string
-		inputRecord         nbdns.SimpleRecord
-		inputMSG            *dns.Msg
-		responseShouldBeNil bool
-	}{
-		{
-			name:        "Should Resolve A Record",
-			inputRecord: recordA,
-			inputMSG:    new(dns.Msg).SetQuestion(recordA.Name, dns.TypeA),
-		},
-		{
-			name:        "Should Resolve CNAME Record",
-			inputRecord: recordCNAME,
-			inputMSG:    new(dns.Msg).SetQuestion(recordCNAME.Name, dns.TypeCNAME),
-		},
-		{
-			name:                "Should Not Write When Not Found A Record",
-			inputRecord:         recordA,
-			inputMSG:            new(dns.Msg).SetQuestion("not.found.com", dns.TypeA),
-			responseShouldBeNil: true,
-		},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			resolver := NewResolver()
-			_ = resolver.RegisterRecord(testCase.inputRecord)
-			var responseMSG *dns.Msg
-			responseWriter := &test.MockResponseWriter{
-				WriteMsgFunc: func(m *dns.Msg) error {
-					responseMSG = m
-					return nil
-				},
-			}
-
-			resolver.ServeDNS(responseWriter, testCase.inputMSG)
-
-			if responseMSG == nil || len(responseMSG.Answer) == 0 {
-				if testCase.responseShouldBeNil {
-					return
-				}
-				t.Fatalf("should write a response message")
-			}
-
-			answerString := responseMSG.Answer[0].String()
-			if !strings.Contains(answerString, testCase.inputRecord.Name) {
-				t.Fatalf("answer doesn't contain the same domain name: \nWant: %s\nGot:%s", testCase.name, answerString)
-			}
-			if !strings.Contains(answerString, dns.Type(testCase.inputRecord.Type).String()) {
-				t.Fatalf("answer doesn't contain the correct type: \nWant: %s\nGot:%s", dns.Type(testCase.inputRecord.Type).String(), answerString)
-			}
-			if !strings.Contains(answerString, testCase.inputRecord.RData) {
-				t.Fatalf("answer doesn't contain the same address: \nWant: %s\nGot:%s", testCase.inputRecord.RData, answerString)
-			}
-		})
-	}
-}
-
-// TestLocalResolver_Update_StaleRecord verifies that updating
-// a record correctly replaces the old one, preventing stale entries.
-func TestLocalResolver_Update_StaleRecord(t *testing.T) {
-	recordName := "host.example.com."
-	recordType := dns.TypeA
-	recordClass := dns.ClassINET
-
-	record1 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "1.1.1.1",
-	}
-	record2 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "2.2.2.2",
-	}
-
-	recordKey := dns.Question{Name: recordName, Qtype: uint16(recordClass), Qclass: recordType}
-
-	resolver := NewResolver()
-
-	update1 := []nbdns.SimpleRecord{record1}
-	update2 := []nbdns.SimpleRecord{record2}
-
-	// Apply first update
-	resolver.Update(update1)
-
-	// Verify first update
-	resolver.mu.RLock()
-	rrSlice1, found1 := resolver.records[recordKey]
-	resolver.mu.RUnlock()
-
-	require.True(t, found1, "Record key %s not found after first update", recordKey)
-	require.Len(t, rrSlice1, 1, "Should have exactly 1 record after first update")
-	assert.Contains(t, rrSlice1[0].String(), record1.RData, "Record after first update should be %s", record1.RData)
-
-	// Apply second update
-	resolver.Update(update2)
-
-	// Verify second update
-	resolver.mu.RLock()
-	rrSlice2, found2 := resolver.records[recordKey]
-	resolver.mu.RUnlock()
-
-	require.True(t, found2, "Record key %s not found after second update", recordKey)
-	require.Len(t, rrSlice2, 1, "Should have exactly 1 record after update overwriting the key")
-	assert.Contains(t, rrSlice2[0].String(), record2.RData, "The single record should be the updated one (%s)", record2.RData)
-	assert.NotContains(t, rrSlice2[0].String(), record1.RData, "The stale record (%s) should not be present", record1.RData)
-}
-
-// TestLocalResolver_MultipleRecords_SameQuestion verifies that multiple records
-// with the same question are stored properly
-func TestLocalResolver_MultipleRecords_SameQuestion(t *testing.T) {
-	resolver := NewResolver()
-
-	recordName := "multi.example.com."
-	recordType := dns.TypeA
-
-	// Create two records with the same name and type but different IPs
-	record1 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1",
-	}
-	record2 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2",
-	}
-
-	update := []nbdns.SimpleRecord{record1, record2}
-
-	// Apply update with both records
-	resolver.Update(update)
-
-	// Create question that matches both records
-	question := dns.Question{
-		Name:   recordName,
-		Qtype:  recordType,
-		Qclass: dns.ClassINET,
-	}
-
-	// Verify both records are stored
-	resolver.mu.RLock()
-	records, found := resolver.records[question]
-	resolver.mu.RUnlock()
-
-	require.True(t, found, "Records for question %v not found", question)
-	require.Len(t, records, 2, "Should have exactly 2 records for the same question")
-
-	// Verify both record data values are present
-	recordStrings := []string{records[0].String(), records[1].String()}
-	assert.Contains(t, recordStrings[0]+recordStrings[1], record1.RData, "First record data should be present")
-	assert.Contains(t, recordStrings[0]+recordStrings[1], record2.RData, "Second record data should be present")
-}
-
-// TestLocalResolver_RecordRotation verifies that records are rotated in a round-robin fashion
-func TestLocalResolver_RecordRotation(t *testing.T) {
-	resolver := NewResolver()
-
-	recordName := "rotation.example.com."
-	recordType := dns.TypeA
-
-	// Create three records with the same name and type but different IPs
-	record1 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.1",
-	}
-	record2 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.2",
-	}
-	record3 := nbdns.SimpleRecord{
-		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.3",
-	}
-
-	update := []nbdns.SimpleRecord{record1, record2, record3}
-
-	// Apply update with all three records
-	resolver.Update(update)
-
-	msg := new(dns.Msg).SetQuestion(recordName, recordType)
-
-	// First lookup - should return the records in original order
-	var responses [3]*dns.Msg
-
-	// Perform three lookups to verify rotation
-	for i := 0; i < 3; i++ {
-		responseWriter := &test.MockResponseWriter{
-			WriteMsgFunc: func(m *dns.Msg) error {
-				responses[i] = m
-				return nil
-			},
-		}
-
-		resolver.ServeDNS(responseWriter, msg)
-	}
-
-	// Verify all three responses contain answers
-	for i, resp := range responses {
-		require.NotNil(t, resp, "Response %d should not be nil", i)
-		require.Len(t, resp.Answer, 3, "Response %d should have 3 answers", i)
-	}
-
-	// Verify the first record in each response is different due to rotation
-	firstRecordIPs := []string{
-		responses[0].Answer[0].String(),
-		responses[1].Answer[0].String(),
-		responses[2].Answer[0].String(),
-	}
-
-	// Each record should be different (rotated)
-	assert.NotEqual(t, firstRecordIPs[0], firstRecordIPs[1], "First lookup should differ from second lookup due to rotation")
-	assert.NotEqual(t, firstRecordIPs[1], firstRecordIPs[2], "Second lookup should differ from third lookup due to rotation")
-	assert.NotEqual(t, firstRecordIPs[0], firstRecordIPs[2], "First lookup should differ from third lookup due to rotation")
-
-	// After three rotations, we should have cycled through all records
-	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record1.RData)
-	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record2.RData)
-	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record3.RData)
-}
-
-// TestLocalResolver_CaseInsensitiveMatching verifies that DNS record lookups are case-insensitive
-func TestLocalResolver_CaseInsensitiveMatching(t *testing.T) {
-	resolver := NewResolver()
-
-	// Create record with lowercase name
-	lowerCaseRecord := nbdns.SimpleRecord{
-		Name:  "lower.example.com.",
-		Type:  int(dns.TypeA),
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "10.10.10.10",
-	}
-
-	// Create record with mixed case name
-	mixedCaseRecord := nbdns.SimpleRecord{
-		Name:  "MiXeD.ExAmPlE.CoM.",
-		Type:  int(dns.TypeA),
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "20.20.20.20",
-	}
-
-	// Update resolver with the records
-	resolver.Update([]nbdns.SimpleRecord{lowerCaseRecord, mixedCaseRecord})
-
-	testCases := []struct {
-		name          string
-		queryName     string
-		expectedRData string
-		shouldResolve bool
-	}{
-		{
-			name:          "Query lowercase with lowercase record",
-			queryName:     "lower.example.com.",
-			expectedRData: "10.10.10.10",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query uppercase with lowercase record",
-			queryName:     "LOWER.EXAMPLE.COM.",
-			expectedRData: "10.10.10.10",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query mixed case with lowercase record",
-			queryName:     "LoWeR.eXaMpLe.CoM.",
-			expectedRData: "10.10.10.10",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query lowercase with mixed case record",
-			queryName:     "mixed.example.com.",
-			expectedRData: "20.20.20.20",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query uppercase with mixed case record",
-			queryName:     "MIXED.EXAMPLE.COM.",
-			expectedRData: "20.20.20.20",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query with different casing pattern",
-			queryName:     "mIxEd.ExaMpLe.cOm.",
-			expectedRData: "20.20.20.20",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query non-existent domain",
-			queryName:     "nonexistent.example.com.",
-			shouldResolve: false,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			var responseMSG *dns.Msg
-
-			// Create DNS query with the test case name
-			msg := new(dns.Msg).SetQuestion(tc.queryName, dns.TypeA)
-
-			// Create mock response writer to capture the response
-			responseWriter := &test.MockResponseWriter{
-				WriteMsgFunc: func(m *dns.Msg) error {
-					responseMSG = m
-					return nil
-				},
-			}
-
-			// Perform DNS query
-			resolver.ServeDNS(responseWriter, msg)
-
-			// Check if we expect a successful resolution
-			if !tc.shouldResolve {
-				if responseMSG == nil || len(responseMSG.Answer) == 0 {
-					// Expected no answer, test passes
-					return
-				}
-				t.Fatalf("Expected no resolution for %s, but got answer: %v", tc.queryName, responseMSG.Answer)
-			}
-
-			// Verify we got a response
-			require.NotNil(t, responseMSG, "Should have received a response message")
-			require.Greater(t, len(responseMSG.Answer), 0, "Response should contain at least one answer")
-
-			// Verify the response contains the expected data
-			answerString := responseMSG.Answer[0].String()
-			assert.Contains(t, answerString, tc.expectedRData,
-				"Answer should contain the expected IP address %s, got: %s",
-				tc.expectedRData, answerString)
-		})
-	}
-}
-
-// TestLocalResolver_CNAMEFallback verifies that the resolver correctly falls back
-// to checking for CNAME records when the requested record type isn't found
-func TestLocalResolver_CNAMEFallback(t *testing.T) {
-	resolver := NewResolver()
-
-	// Create a CNAME record (but no A record for this name)
-	cnameRecord := nbdns.SimpleRecord{
-		Name:  "alias.example.com.",
-		Type:  int(dns.TypeCNAME),
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "target.example.com.",
-	}
-
-	// Create an A record for the CNAME target
-	targetRecord := nbdns.SimpleRecord{
-		Name:  "target.example.com.",
-		Type:  int(dns.TypeA),
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "192.168.100.100",
-	}
-
-	// Update resolver with both records
-	resolver.Update([]nbdns.SimpleRecord{cnameRecord, targetRecord})
-
-	testCases := []struct {
-		name          string
-		queryName     string
-		queryType     uint16
-		expectedType  string
-		expectedRData string
-		shouldResolve bool
-	}{
-		{
-			name:          "Directly query CNAME record",
-			queryName:     "alias.example.com.",
-			queryType:     dns.TypeCNAME,
-			expectedType:  "CNAME",
-			expectedRData: "target.example.com.",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query A record but get CNAME fallback",
-			queryName:     "alias.example.com.",
-			queryType:     dns.TypeA,
-			expectedType:  "CNAME",
-			expectedRData: "target.example.com.",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query AAAA record but get CNAME fallback",
-			queryName:     "alias.example.com.",
-			queryType:     dns.TypeAAAA,
-			expectedType:  "CNAME",
-			expectedRData: "target.example.com.",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query direct A record",
-			queryName:     "target.example.com.",
-			queryType:     dns.TypeA,
-			expectedType:  "A",
-			expectedRData: "192.168.100.100",
-			shouldResolve: true,
-		},
-		{
-			name:          "Query non-existent name",
-			queryName:     "nonexistent.example.com.",
-			queryType:     dns.TypeA,
-			shouldResolve: false,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			var responseMSG *dns.Msg
-
-			// Create DNS query with the test case parameters
-			msg := new(dns.Msg).SetQuestion(tc.queryName, tc.queryType)
-
-			// Create mock response writer to capture the response
-			responseWriter := &test.MockResponseWriter{
-				WriteMsgFunc: func(m *dns.Msg) error {
-					responseMSG = m
-					return nil
-				},
-			}
-
-			// Perform DNS query
-			resolver.ServeDNS(responseWriter, msg)
-
-			// Check if we expect a successful resolution
-			if !tc.shouldResolve {
-				if responseMSG == nil || len(responseMSG.Answer) == 0 || responseMSG.Rcode != dns.RcodeSuccess {
-					// Expected no resolution, test passes
-					return
-				}
-				t.Fatalf("Expected no resolution for %s, but got answer: %v", tc.queryName, responseMSG.Answer)
-			}
-
-			// Verify we got a successful response
-			require.NotNil(t, responseMSG, "Should have received a response message")
-			require.Equal(t, dns.RcodeSuccess, responseMSG.Rcode, "Response should have success status code")
-			require.Greater(t, len(responseMSG.Answer), 0, "Response should contain at least one answer")
-
-			// Verify the response contains the expected data
-			answerString := responseMSG.Answer[0].String()
-			assert.Contains(t, answerString, tc.expectedType,
-				"Answer should be of type %s, got: %s", tc.expectedType, answerString)
-			assert.Contains(t, answerString, tc.expectedRData,
-				"Answer should contain the expected data %s, got: %s", tc.expectedRData, answerString)
-		})
-	}
-}

client/internal/dns/local_test.go

@@ -0,0 +1,88 @@
+package dns
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/miekg/dns"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+func TestLocalResolver_ServeDNS(t *testing.T) {
+	recordA := nbdns.SimpleRecord{
+		Name:  "peera.netbird.cloud.",
+		Type:  1,
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "1.2.3.4",
+	}
+
+	recordCNAME := nbdns.SimpleRecord{
+		Name:  "peerb.netbird.cloud.",
+		Type:  5,
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "www.netbird.io",
+	}
+
+	testCases := []struct {
+		name                string
+		inputRecord         nbdns.SimpleRecord
+		inputMSG            *dns.Msg
+		responseShouldBeNil bool
+	}{
+		{
+			name:        "Should Resolve A Record",
+			inputRecord: recordA,
+			inputMSG:    new(dns.Msg).SetQuestion(recordA.Name, dns.TypeA),
+		},
+		{
+			name:        "Should Resolve CNAME Record",
+			inputRecord: recordCNAME,
+			inputMSG:    new(dns.Msg).SetQuestion(recordCNAME.Name, dns.TypeCNAME),
+		},
+		{
+			name:                "Should Not Write When Not Found A Record",
+			inputRecord:         recordA,
+			inputMSG:            new(dns.Msg).SetQuestion("not.found.com", dns.TypeA),
+			responseShouldBeNil: true,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			resolver := &localResolver{
+				registeredMap: make(registrationMap),
+			}
+			_, _ = resolver.registerRecord(testCase.inputRecord)
+			var responseMSG *dns.Msg
+			responseWriter := &mockResponseWriter{
+				WriteMsgFunc: func(m *dns.Msg) error {
+					responseMSG = m
+					return nil
+				},
+			}
+
+			resolver.ServeDNS(responseWriter, testCase.inputMSG)
+
+			if responseMSG == nil || len(responseMSG.Answer) == 0 {
+				if testCase.responseShouldBeNil {
+					return
+				}
+				t.Fatalf("should write a response message")
+			}
+
+			answerString := responseMSG.Answer[0].String()
+			if !strings.Contains(answerString, testCase.inputRecord.Name) {
+				t.Fatalf("answer doesn't contain the same domain name: \nWant: %s\nGot:%s", testCase.name, answerString)
+			}
+			if !strings.Contains(answerString, dns.Type(testCase.inputRecord.Type).String()) {
+				t.Fatalf("answer doesn't contain the correct type: \nWant: %s\nGot:%s", dns.Type(testCase.inputRecord.Type).String(), answerString)
+			}
+			if !strings.Contains(answerString, testCase.inputRecord.RData) {
+				t.Fatalf("answer doesn't contain the same address: \nWant: %s\nGot:%s", testCase.inputRecord.RData, answerString)
+			}
+		})
+	}
+}

client/internal/dns/mock_test.go

@@ -0,0 +1,26 @@
+package dns
+
+import (
+	"net"
+
+	"github.com/miekg/dns"
+)
+
+type mockResponseWriter struct {
+	WriteMsgFunc func(m *dns.Msg) error
+}
+
+func (rw *mockResponseWriter) WriteMsg(m *dns.Msg) error {
+	if rw.WriteMsgFunc != nil {
+		return rw.WriteMsgFunc(m)
+	}
+	return nil
+}
+
+func (rw *mockResponseWriter) LocalAddr() net.Addr       { return nil }
+func (rw *mockResponseWriter) RemoteAddr() net.Addr      { return nil }
+func (rw *mockResponseWriter) Write([]byte) (int, error) { return 0, nil }
+func (rw *mockResponseWriter) Close() error              { return nil }
+func (rw *mockResponseWriter) TsigStatus() error         { return nil }
+func (rw *mockResponseWriter) TsigTimersOnly(bool)       {}
+func (rw *mockResponseWriter) Hijack()                   {}

client/internal/dns/server.go

@@ -15,8 +15,6 @@ import (
 	"golang.org/x/exp/maps"
 
 	"github.com/netbirdio/netbird/client/iface/netstack"
-	"github.com/netbirdio/netbird/client/internal/dns/local"
-	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/listener"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -48,6 +46,8 @@ type Server interface {
 	ProbeAvailability()
 }
 
+type handlerID string
+
 type nsGroupsByDomain struct {
 	domain string
 	groups []*nbdns.NameServerGroup
@@ -61,7 +61,7 @@ type DefaultServer struct {
 	mux                sync.Mutex
 	service            service
 	dnsMuxMap          registeredHandlerMap
-	localResolver      *local.Resolver
+	localResolver      *localResolver
 	wgInterface        WGIface
 	hostManager        hostManager
 	updateSerial       uint64
@@ -84,9 +84,9 @@ type DefaultServer struct {
 
 type handlerWithStop interface {
 	dns.Handler
-	Stop()
-	ProbeAvailability()
-	ID() types.HandlerID
+	stop()
+	probeAvailability()
+	id() handlerID
 }
 
 type handlerWrapper struct {
@@ -95,7 +95,7 @@ type handlerWrapper struct {
 	priority int
 }
 
-type registeredHandlerMap map[types.HandlerID]handlerWrapper
+type registeredHandlerMap map[handlerID]handlerWrapper
 
 // NewDefaultServer returns a new dns server
 func NewDefaultServer(
@@ -171,14 +171,16 @@ func newDefaultServer(
 	handlerChain := NewHandlerChain()
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
-		ctx:            ctx,
-		ctxCancel:      stop,
-		disableSys:     disableSys,
-		service:        dnsService,
-		handlerChain:   handlerChain,
-		extraDomains:   make(map[domain.Domain]int),
-		dnsMuxMap:      make(registeredHandlerMap),
-		localResolver:  local.NewResolver(),
+		ctx:          ctx,
+		ctxCancel:    stop,
+		disableSys:   disableSys,
+		service:      dnsService,
+		handlerChain: handlerChain,
+		extraDomains: make(map[domain.Domain]int),
+		dnsMuxMap:    make(registeredHandlerMap),
+		localResolver: &localResolver{
+			registeredMap: make(registrationMap),
+		},
 		wgInterface:    wgInterface,
 		statusRecorder: statusRecorder,
 		stateManager:   stateManager,
@@ -401,7 +403,7 @@ func (s *DefaultServer) ProbeAvailability() {
 		wg.Add(1)
 		go func(mux handlerWithStop) {
 			defer wg.Done()
-			mux.ProbeAvailability()
+			mux.probeAvailability()
 		}(mux.handler)
 	}
 	wg.Wait()
@@ -418,7 +420,7 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 		s.service.Stop()
 	}
 
-	localMuxUpdates, localRecords, err := s.buildLocalHandlerUpdate(update.CustomZones)
+	localMuxUpdates, localRecordsByDomain, err := s.buildLocalHandlerUpdate(update.CustomZones)
 	if err != nil {
 		return fmt.Errorf("local handler updater: %w", err)
 	}
@@ -432,7 +434,7 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	s.updateMux(muxUpdates)
 
 	// register local records
-	s.localResolver.Update(localRecords)
+	s.updateLocalResolver(localRecordsByDomain)
 
 	s.currentConfig = dnsConfigToHostDNSConfig(update, s.service.RuntimeIP(), s.service.RuntimePort())
 
@@ -514,9 +516,11 @@ func (s *DefaultServer) handleErrNoGroupaAll(err error) {
 	)
 }
 
-func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone) ([]handlerWrapper, []nbdns.SimpleRecord, error) {
+func (s *DefaultServer) buildLocalHandlerUpdate(
+	customZones []nbdns.CustomZone,
+) ([]handlerWrapper, map[string][]nbdns.SimpleRecord, error) {
 	var muxUpdates []handlerWrapper
-	var localRecords []nbdns.SimpleRecord
+	localRecords := make(map[string][]nbdns.SimpleRecord)
 
 	for _, customZone := range customZones {
 		if len(customZone.Records) == 0 {
@@ -530,13 +534,17 @@ func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone)
 			priority: PriorityMatchDomain,
 		})
 
+		// group all records under this domain
 		for _, record := range customZone.Records {
+			var class uint16 = dns.ClassINET
 			if record.Class != nbdns.DefaultClass {
 				log.Warnf("received an invalid class type: %s", record.Class)
 				continue
 			}
-			// zone records contain the fqdn, so we can just flatten them
-			localRecords = append(localRecords, record)
+
+			key := buildRecordKey(record.Name, class, uint16(record.Type))
+
+			localRecords[key] = append(localRecords[key], record)
 		}
 	}
 
@@ -619,7 +627,7 @@ func (s *DefaultServer) createHandlersForDomainGroup(domainGroup nsGroupsByDomai
 		}
 
 		if len(handler.upstreamServers) == 0 {
-			handler.Stop()
+			handler.stop()
 			log.Errorf("received a nameserver group with an invalid nameserver list")
 			continue
 		}
@@ -648,7 +656,7 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 	// this will introduce a short period of time when the server is not able to handle DNS requests
 	for _, existing := range s.dnsMuxMap {
 		s.deregisterHandler([]string{existing.domain}, existing.priority)
-		existing.handler.Stop()
+		existing.handler.stop()
 	}
 
 	muxUpdateMap := make(registeredHandlerMap)
@@ -659,7 +667,7 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 			containsRootUpdate = true
 		}
 		s.registerHandler([]string{update.domain}, update.handler, update.priority)
-		muxUpdateMap[update.handler.ID()] = update
+		muxUpdateMap[update.handler.id()] = update
 	}
 
 	// If there's no root update and we had a root handler, restore it
@@ -675,6 +683,33 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 	s.dnsMuxMap = muxUpdateMap
 }
 
+func (s *DefaultServer) updateLocalResolver(update map[string][]nbdns.SimpleRecord) {
+	// remove old records that are no longer present
+	for key := range s.localResolver.registeredMap {
+		_, found := update[key]
+		if !found {
+			s.localResolver.deleteRecord(key)
+		}
+	}
+
+	updatedMap := make(registrationMap)
+	for _, recs := range update {
+		for _, rec := range recs {
+			// convert the record to a dns.RR and register
+			key, err := s.localResolver.registerRecord(rec)
+			if err != nil {
+				log.Warnf("got an error while registering the record (%s), error: %v",
+					rec.String(), err)
+				continue
+			}
+
+			updatedMap[key] = struct{}{}
+		}
+	}
+
+	s.localResolver.registeredMap = updatedMap
+}
+
 func getNSHostPort(ns nbdns.NameServer) string {
 	return fmt.Sprintf("%s:%d", ns.IP.String(), ns.Port)
 }

client/internal/dns/server_test.go

@@ -23,9 +23,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	pfmock "github.com/netbirdio/netbird/client/iface/mocks"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
-	"github.com/netbirdio/netbird/client/internal/dns/local"
-	"github.com/netbirdio/netbird/client/internal/dns/test"
-	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/netflow"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -110,7 +107,6 @@ func generateDummyHandler(domain string, servers []nbdns.NameServer) *upstreamRe
 }
 
 func TestUpdateDNSServer(t *testing.T) {
-
 	nameServers := []nbdns.NameServer{
 		{
 			IP:     netip.MustParseAddr("8.8.8.8"),
@@ -124,21 +120,22 @@ func TestUpdateDNSServer(t *testing.T) {
 		},
 	}
 
-	dummyHandler := local.NewResolver()
+	dummyHandler := &localResolver{}
 
 	testCases := []struct {
 		name                string
 		initUpstreamMap     registeredHandlerMap
-		initLocalRecords    []nbdns.SimpleRecord
+		initLocalMap        registrationMap
 		initSerial          uint64
 		inputSerial         uint64
 		inputUpdate         nbdns.Config
 		shouldFail          bool
 		expectedUpstreamMap registeredHandlerMap
-		expectedLocalQs     []dns.Question
+		expectedLocalMap    registrationMap
 	}{
 		{
 			name:            "Initial Config Should Succeed",
+			initLocalMap:    make(registrationMap),
 			initUpstreamMap: make(registeredHandlerMap),
 			initSerial:      0,
 			inputSerial:     1,
@@ -162,30 +159,30 @@ func TestUpdateDNSServer(t *testing.T) {
 				},
 			},
 			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{
+				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
 					domain:   "netbird.io",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
-				dummyHandler.ID(): handlerWrapper{
+				dummyHandler.id(): handlerWrapper{
 					domain:   "netbird.cloud",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
-				generateDummyHandler(".", nameServers).ID(): handlerWrapper{
+				generateDummyHandler(".", nameServers).id(): handlerWrapper{
 					domain:   nbdns.RootZone,
 					handler:  dummyHandler,
 					priority: PriorityDefault,
 				},
 			},
-			expectedLocalQs: []dns.Question{{Name: "peera.netbird.cloud.", Qtype: dns.TypeA, Qclass: dns.ClassINET}},
+			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
 		},
 		{
-			name:             "New Config Should Succeed",
-			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
+			name:         "New Config Should Succeed",
+			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
-					domain:   "netbird.cloud",
+				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
+					domain:   buildRecordKey(zoneRecords[0].Name, 1, 1),
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
@@ -208,7 +205,7 @@ func TestUpdateDNSServer(t *testing.T) {
 				},
 			},
 			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{
+				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
 					domain:   "netbird.io",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -219,22 +216,22 @@ func TestUpdateDNSServer(t *testing.T) {
 					priority: PriorityMatchDomain,
 				},
 			},
-			expectedLocalQs: []dns.Question{{Name: zoneRecords[0].Name, Qtype: 1, Qclass: 1}},
+			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
 		},
 		{
-			name:             "Smaller Config Serial Should Be Skipped",
-			initLocalRecords: []nbdns.SimpleRecord{},
-			initUpstreamMap:  make(registeredHandlerMap),
-			initSerial:       2,
-			inputSerial:      1,
-			shouldFail:       true,
+			name:            "Smaller Config Serial Should Be Skipped",
+			initLocalMap:    make(registrationMap),
+			initUpstreamMap: make(registeredHandlerMap),
+			initSerial:      2,
+			inputSerial:     1,
+			shouldFail:      true,
 		},
 		{
-			name:             "Empty NS Group Domain Or Not Primary Element Should Fail",
-			initLocalRecords: []nbdns.SimpleRecord{},
-			initUpstreamMap:  make(registeredHandlerMap),
-			initSerial:       0,
-			inputSerial:      1,
+			name:            "Empty NS Group Domain Or Not Primary Element Should Fail",
+			initLocalMap:    make(registrationMap),
+			initUpstreamMap: make(registeredHandlerMap),
+			initSerial:      0,
+			inputSerial:     1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -252,11 +249,11 @@ func TestUpdateDNSServer(t *testing.T) {
 			shouldFail: true,
 		},
 		{
-			name:             "Invalid NS Group Nameservers list Should Fail",
-			initLocalRecords: []nbdns.SimpleRecord{},
-			initUpstreamMap:  make(registeredHandlerMap),
-			initSerial:       0,
-			inputSerial:      1,
+			name:            "Invalid NS Group Nameservers list Should Fail",
+			initLocalMap:    make(registrationMap),
+			initUpstreamMap: make(registeredHandlerMap),
+			initSerial:      0,
+			inputSerial:     1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -274,11 +271,11 @@ func TestUpdateDNSServer(t *testing.T) {
 			shouldFail: true,
 		},
 		{
-			name:             "Invalid Custom Zone Records list Should Skip",
-			initLocalRecords: []nbdns.SimpleRecord{},
-			initUpstreamMap:  make(registeredHandlerMap),
-			initSerial:       0,
-			inputSerial:      1,
+			name:            "Invalid Custom Zone Records list Should Skip",
+			initLocalMap:    make(registrationMap),
+			initUpstreamMap: make(registeredHandlerMap),
+			initSerial:      0,
+			inputSerial:     1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -293,17 +290,17 @@ func TestUpdateDNSServer(t *testing.T) {
 					},
 				},
 			},
-			expectedUpstreamMap: registeredHandlerMap{generateDummyHandler(".", nameServers).ID(): handlerWrapper{
+			expectedUpstreamMap: registeredHandlerMap{generateDummyHandler(".", nameServers).id(): handlerWrapper{
 				domain:   ".",
 				handler:  dummyHandler,
 				priority: PriorityDefault,
 			}},
 		},
 		{
-			name:             "Empty Config Should Succeed and Clean Maps",
-			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
+			name:         "Empty Config Should Succeed and Clean Maps",
+			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
+				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
 					domain:   zoneRecords[0].Name,
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -313,13 +310,13 @@ func TestUpdateDNSServer(t *testing.T) {
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: true},
 			expectedUpstreamMap: make(registeredHandlerMap),
-			expectedLocalQs:     []dns.Question{},
+			expectedLocalMap:    make(registrationMap),
 		},
 		{
-			name:             "Disabled Service Should clean map",
-			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
+			name:         "Disabled Service Should clean map",
+			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
+				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
 					domain:   zoneRecords[0].Name,
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -329,7 +326,7 @@ func TestUpdateDNSServer(t *testing.T) {
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: false},
 			expectedUpstreamMap: make(registeredHandlerMap),
-			expectedLocalQs:     []dns.Question{},
+			expectedLocalMap:    make(registrationMap),
 		},
 	}
 
@@ -380,7 +377,7 @@ func TestUpdateDNSServer(t *testing.T) {
 			}()
 
 			dnsServer.dnsMuxMap = testCase.initUpstreamMap
-			dnsServer.localResolver.Update(testCase.initLocalRecords)
+			dnsServer.localResolver.registeredMap = testCase.initLocalMap
 			dnsServer.updateSerial = testCase.initSerial
 
 			err = dnsServer.UpdateDNSServer(testCase.inputSerial, testCase.inputUpdate)
@@ -402,23 +399,15 @@ func TestUpdateDNSServer(t *testing.T) {
 				}
 			}
 
-			var responseMSG *dns.Msg
-			responseWriter := &test.MockResponseWriter{
-				WriteMsgFunc: func(m *dns.Msg) error {
-					responseMSG = m
-					return nil
-				},
-			}
-			for _, q := range testCase.expectedLocalQs {
-				dnsServer.localResolver.ServeDNS(responseWriter, &dns.Msg{
-					Question: []dns.Question{q},
-				})
+			if len(dnsServer.localResolver.registeredMap) != len(testCase.expectedLocalMap) {
+				t.Fatalf("update local failed, registered map size is different than expected, want %d, got %d", len(testCase.expectedLocalMap), len(dnsServer.localResolver.registeredMap))
 			}
 
-			if len(testCase.expectedLocalQs) > 0 {
-				assert.NotNil(t, responseMSG, "response message should not be nil")
-				assert.Equal(t, dns.RcodeSuccess, responseMSG.Rcode, "response code should be success")
-				assert.NotEmpty(t, responseMSG.Answer, "response message should have answers")
+			for key := range testCase.expectedLocalMap {
+				_, found := dnsServer.localResolver.registeredMap[key]
+				if !found {
+					t.Fatalf("update local failed, key %s was not found in the localResolver.registeredMap: %#v", key, dnsServer.localResolver.registeredMap)
+				}
 			}
 		})
 	}
@@ -502,12 +491,11 @@ func TestDNSFakeResolverHandleUpdates(t *testing.T) {
 	dnsServer.dnsMuxMap = registeredHandlerMap{
 		"id1": handlerWrapper{
 			domain:   zoneRecords[0].Name,
-			handler:  &local.Resolver{},
+			handler:  &localResolver{},
 			priority: PriorityMatchDomain,
 		},
 	}
-	//dnsServer.localResolver.RegisteredMap = local.RegistrationMap{local.BuildRecordKey("netbird.cloud", dns.ClassINET, dns.TypeA): struct{}{}}
-	dnsServer.localResolver.Update([]nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}})
+	dnsServer.localResolver.registeredMap = registrationMap{"netbird.cloud": struct{}{}}
 	dnsServer.updateSerial = 0
 
 	nameServers := []nbdns.NameServer{
@@ -594,7 +582,7 @@ func TestDNSServerStartStop(t *testing.T) {
 			}
 			time.Sleep(100 * time.Millisecond)
 			defer dnsServer.Stop()
-			err = dnsServer.localResolver.RegisterRecord(zoneRecords[0])
+			_, err = dnsServer.localResolver.registerRecord(zoneRecords[0])
 			if err != nil {
 				t.Error(err)
 			}
@@ -642,11 +630,13 @@ func TestDNSServerStartStop(t *testing.T) {
 func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 	hostManager := &mockHostConfigurator{}
 	server := DefaultServer{
-		ctx:           context.Background(),
-		service:       NewServiceViaMemory(&mocWGIface{}),
-		localResolver: local.NewResolver(),
-		handlerChain:  NewHandlerChain(),
-		hostManager:   hostManager,
+		ctx:     context.Background(),
+		service: NewServiceViaMemory(&mocWGIface{}),
+		localResolver: &localResolver{
+			registeredMap: make(registrationMap),
+		},
+		handlerChain: NewHandlerChain(),
+		hostManager:  hostManager,
 		currentConfig: HostDNSConfig{
 			Domains: []DomainConfig{
 				{false, "domain0", false},
@@ -1014,7 +1004,7 @@ func TestHandlerChain_DomainPriorities(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			r := new(dns.Msg)
 			r.SetQuestion(tc.query, dns.TypeA)
-			w := &ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
+			w := &ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 
 			if mh, ok := tc.expectedHandler.(*MockHandler); ok {
 				mh.On("ServeDNS", mock.Anything, r).Once()
@@ -1047,9 +1037,9 @@ type mockHandler struct {
 }
 
 func (m *mockHandler) ServeDNS(dns.ResponseWriter, *dns.Msg) {}
-func (m *mockHandler) Stop()                                 {}
-func (m *mockHandler) ProbeAvailability()                    {}
-func (m *mockHandler) ID() types.HandlerID                   { return types.HandlerID(m.Id) }
+func (m *mockHandler) stop()                                 {}
+func (m *mockHandler) probeAvailability()                    {}
+func (m *mockHandler) id() handlerID                         { return handlerID(m.Id) }
 
 type mockService struct{}
 
@@ -1123,7 +1113,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 		name             string
 		initialHandlers  registeredHandlerMap
 		updates          []handlerWrapper
-		expectedHandlers map[string]string // map[HandlerID]domain
+		expectedHandlers map[string]string // map[handlerID]domain
 		description      string
 	}{
 		{
@@ -1419,7 +1409,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 
 			// Check each expected handler
 			for id, expectedDomain := range tt.expectedHandlers {
-				handler, exists := server.dnsMuxMap[types.HandlerID(id)]
+				handler, exists := server.dnsMuxMap[handlerID(id)]
 				assert.True(t, exists, "Expected handler %s not found", id)
 				if exists {
 					assert.Equal(t, expectedDomain, handler.domain,
@@ -1428,9 +1418,9 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 			}
 
 			// Verify no unexpected handlers exist
-			for HandlerID := range server.dnsMuxMap {
-				_, expected := tt.expectedHandlers[string(HandlerID)]
-				assert.True(t, expected, "Unexpected handler found: %s", HandlerID)
+			for handlerID := range server.dnsMuxMap {
+				_, expected := tt.expectedHandlers[string(handlerID)]
+				assert.True(t, expected, "Unexpected handler found: %s", handlerID)
 			}
 
 			// Verify the handlerChain state and order
@@ -1706,7 +1696,7 @@ func TestExtraDomains(t *testing.T) {
 				handlerChain:   NewHandlerChain(),
 				wgInterface:    &mocWGIface{},
 				hostManager:    mockHostConfig,
-				localResolver:  &local.Resolver{},
+				localResolver:  &localResolver{},
 				service:        mockSvc,
 				statusRecorder: peer.NewRecorder("test"),
 				extraDomains:   make(map[domain.Domain]int),
@@ -1791,7 +1781,7 @@ func TestExtraDomainsRefCounting(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &local.Resolver{},
+		localResolver:  &localResolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),
@@ -1843,7 +1833,7 @@ func TestUpdateConfigWithExistingExtraDomains(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &local.Resolver{},
+		localResolver:  &localResolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),
@@ -1926,7 +1916,7 @@ func TestDomainCaseHandling(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &local.Resolver{},
+		localResolver:  &localResolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),

client/internal/dns/systemd_linux.go

@@ -30,12 +30,9 @@ const (
 	systemdDbusSetDNSMethodSuffix          = systemdDbusLinkInterface + ".SetDNS"
 	systemdDbusSetDefaultRouteMethodSuffix = systemdDbusLinkInterface + ".SetDefaultRoute"
 	systemdDbusSetDomainsMethodSuffix      = systemdDbusLinkInterface + ".SetDomains"
-	systemdDbusSetDNSSECMethodSuffix       = systemdDbusLinkInterface + ".SetDNSSEC"
 	systemdDbusResolvConfModeForeign       = "foreign"
 
 	dbusErrorUnknownObject = "org.freedesktop.DBus.Error.UnknownObject"
-
-	dnsSecDisabled = "no"
 )
 
 type systemdDbusConfigurator struct {
@@ -98,13 +95,9 @@ func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateMana
 		Family:  unix.AF_INET,
 		Address: ipAs4[:],
 	}
-	if err = s.callLinkMethod(systemdDbusSetDNSMethodSuffix, []systemdDbusDNSInput{defaultLinkInput}); err != nil {
-		return fmt.Errorf("set interface DNS server %s:%d: %w", config.ServerIP, config.ServerPort, err)
-	}
-
-	// We don't support dnssec. On some machines this is default on so we explicitly set it to off
-	if err = s.callLinkMethod(systemdDbusSetDNSSECMethodSuffix, dnsSecDisabled); err != nil {
-		log.Warnf("failed to set DNSSEC to 'no': %v", err)
+	err = s.callLinkMethod(systemdDbusSetDNSMethodSuffix, []systemdDbusDNSInput{defaultLinkInput})
+	if err != nil {
+		return fmt.Errorf("setting the interface DNS server %s:%d failed with error: %w", config.ServerIP, config.ServerPort, err)
 	}
 
 	var (

client/internal/dns/test/mock.go

@@ -1,26 +0,0 @@
-package test
-
-import (
-	"net"
-
-	"github.com/miekg/dns"
-)
-
-type MockResponseWriter struct {
-	WriteMsgFunc func(m *dns.Msg) error
-}
-
-func (rw *MockResponseWriter) WriteMsg(m *dns.Msg) error {
-	if rw.WriteMsgFunc != nil {
-		return rw.WriteMsgFunc(m)
-	}
-	return nil
-}
-
-func (rw *MockResponseWriter) LocalAddr() net.Addr       { return nil }
-func (rw *MockResponseWriter) RemoteAddr() net.Addr      { return nil }
-func (rw *MockResponseWriter) Write([]byte) (int, error) { return 0, nil }
-func (rw *MockResponseWriter) Close() error              { return nil }
-func (rw *MockResponseWriter) TsigStatus() error         { return nil }
-func (rw *MockResponseWriter) TsigTimersOnly(bool)       {}
-func (rw *MockResponseWriter) Hijack()                   {}

client/internal/dns/types/types.go

@@ -1,3 +0,0 @@
-package types
-
-type HandlerID string

client/internal/dns/upstream.go

@@ -19,7 +19,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/iface"
-	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/proto"
 )
@@ -82,21 +81,21 @@ func (u *upstreamResolverBase) String() string {
 }
 
 // ID returns the unique handler ID
-func (u *upstreamResolverBase) ID() types.HandlerID {
+func (u *upstreamResolverBase) id() handlerID {
 	servers := slices.Clone(u.upstreamServers)
 	slices.Sort(servers)
 
 	hash := sha256.New()
 	hash.Write([]byte(u.domain + ":"))
 	hash.Write([]byte(strings.Join(servers, ",")))
-	return types.HandlerID("upstream-" + hex.EncodeToString(hash.Sum(nil)[:8]))
+	return handlerID("upstream-" + hex.EncodeToString(hash.Sum(nil)[:8]))
 }
 
 func (u *upstreamResolverBase) MatchSubdomains() bool {
 	return true
 }
 
-func (u *upstreamResolverBase) Stop() {
+func (u *upstreamResolverBase) stop() {
 	log.Debugf("stopping serving DNS for upstreams %s", u.upstreamServers)
 	u.cancel()
 }
@@ -199,9 +198,9 @@ func (u *upstreamResolverBase) checkUpstreamFails(err error) {
 	)
 }
 
-// ProbeAvailability tests all upstream servers simultaneously and
+// probeAvailability tests all upstream servers simultaneously and
 // disables the resolver if none work
-func (u *upstreamResolverBase) ProbeAvailability() {
+func (u *upstreamResolverBase) probeAvailability() {
 	u.mutex.Lock()
 	defer u.mutex.Unlock()
 

client/internal/dns/upstream_test.go

@@ -8,8 +8,6 @@ import (
 	"time"
 
 	"github.com/miekg/dns"
-
-	"github.com/netbirdio/netbird/client/internal/dns/test"
 )
 
 func TestUpstreamResolver_ServeDNS(t *testing.T) {
@@ -68,7 +66,7 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			}
 
 			var responseMSG *dns.Msg
-			responseWriter := &test.MockResponseWriter{
+			responseWriter := &mockResponseWriter{
 				WriteMsgFunc: func(m *dns.Msg) error {
 					responseMSG = m
 					return nil
@@ -132,7 +130,7 @@ func TestUpstreamResolver_DeactivationReactivation(t *testing.T) {
 	resolver.failsTillDeact = 0
 	resolver.reactivatePeriod = time.Microsecond * 100
 
-	responseWriter := &test.MockResponseWriter{
+	responseWriter := &mockResponseWriter{
 		WriteMsgFunc: func(m *dns.Msg) error { return nil },
 	}
 

client/internal/dns/wgiface.go

@@ -5,6 +5,7 @@ package dns
 import (
 	"net"
 
+	"github.com/netbirdio/netbird/client/iface/configurer"
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 )
@@ -17,4 +18,5 @@ type WGIface interface {
 	IsUserspaceBind() bool
 	GetFilter() device.PacketFilter
 	GetDevice() *device.FilteredDevice
+	GetStats(peerKey string) (configurer.WGStats, error)
 }

client/internal/dns/wgiface_windows.go

@@ -1,6 +1,7 @@
 package dns
 
 import (
+	"github.com/netbirdio/netbird/client/iface/configurer"
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 )
@@ -12,5 +13,6 @@ type WGIface interface {
 	IsUserspaceBind() bool
 	GetFilter() device.PacketFilter
 	GetDevice() *device.FilteredDevice
+	GetStats(peerKey string) (configurer.WGStats, error)
 	GetInterfaceGUIDString() (string, error)
 }

client/internal/dnsfwd/forwarder.go

@@ -3,7 +3,6 @@ package dnsfwd
 import (
 	"context"
 	"errors"
-	"fmt"
 	"math"
 	"net"
 	"net/netip"
@@ -11,16 +10,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/hashicorp/go-multierror"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
-	nberrors "github.com/netbirdio/netbird/client/errors"
-	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	nbdns "github.com/netbirdio/netbird/dns"
-	"github.com/netbirdio/netbird/management/domain"
-	"github.com/netbirdio/netbird/route"
 )
 
 const errResolveFailed = "failed to resolve query for domain=%s: %v"
@@ -29,116 +23,79 @@ const upstreamTimeout = 15 * time.Second
 type DNSForwarder struct {
 	listenAddress  string
 	ttl            uint32
+	domains        []string
 	statusRecorder *peer.Status
 
 	dnsServer *dns.Server
 	mux       *dns.ServeMux
-	tcpServer *dns.Server
-	tcpMux    *dns.ServeMux
 
-	mutex      sync.RWMutex
-	fwdEntries []*ForwarderEntry
-	firewall   firewall.Manager
+	resId sync.Map
 }
 
-func NewDNSForwarder(listenAddress string, ttl uint32, firewall firewall.Manager, statusRecorder *peer.Status) *DNSForwarder {
+func NewDNSForwarder(listenAddress string, ttl uint32, statusRecorder *peer.Status) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
 		listenAddress:  listenAddress,
 		ttl:            ttl,
-		firewall:       firewall,
 		statusRecorder: statusRecorder,
 	}
 }
 
-func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
-	log.Infof("starting DNS forwarder on address=%s", f.listenAddress)
-
-	// UDP server
+func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error {
+	log.Infof("listen DNS forwarder on address=%s", f.listenAddress)
 	mux := dns.NewServeMux()
-	f.mux = mux
-	f.dnsServer = &dns.Server{
+
+	dnsServer := &dns.Server{
 		Addr:    f.listenAddress,
 		Net:     "udp",
 		Handler: mux,
 	}
-	// TCP server
-	tcpMux := dns.NewServeMux()
-	f.tcpMux = tcpMux
-	f.tcpServer = &dns.Server{
-		Addr:    f.listenAddress,
-		Net:     "tcp",
-		Handler: tcpMux,
-	}
+	f.dnsServer = dnsServer
+	f.mux = mux
 
-	f.UpdateDomains(entries)
+	f.UpdateDomains(domains, resIds)
 
-	errCh := make(chan error, 2)
-
-	go func() {
-		log.Infof("DNS UDP listener running on %s", f.listenAddress)
-		errCh <- f.dnsServer.ListenAndServe()
-	}()
-	go func() {
-		log.Infof("DNS TCP listener running on %s", f.listenAddress)
-		errCh <- f.tcpServer.ListenAndServe()
-	}()
-
-	// return the first error we get (e.g. bind failure or shutdown)
-	return <-errCh
+	return dnsServer.ListenAndServe()
 }
-func (f *DNSForwarder) UpdateDomains(entries []*ForwarderEntry) {
-	f.mutex.Lock()
-	defer f.mutex.Unlock()
 
-	if f.mux == nil {
-		log.Debug("DNS mux is nil, skipping domain update")
-		f.fwdEntries = entries
-		return
+func (f *DNSForwarder) UpdateDomains(domains []string, resIds map[string]string) {
+	log.Debugf("Updating domains from %v to %v", f.domains, domains)
+
+	for _, d := range f.domains {
+		f.mux.HandleRemove(d)
 	}
+	f.resId.Clear()
 
-	oldDomains := filterDomains(f.fwdEntries)
-	for _, d := range oldDomains {
-		f.mux.HandleRemove(d.PunycodeString())
-		f.tcpMux.HandleRemove(d.PunycodeString())
-	}
-
-	newDomains := filterDomains(entries)
+	newDomains := filterDomains(domains)
 	for _, d := range newDomains {
-		f.mux.HandleFunc(d.PunycodeString(), f.handleDNSQueryUDP)
-		f.tcpMux.HandleFunc(d.PunycodeString(), f.handleDNSQueryTCP)
+		f.mux.HandleFunc(d, f.handleDNSQuery)
 	}
 
-	f.fwdEntries = entries
-	log.Debugf("Updated domains from %v to %v", oldDomains, newDomains)
+	for domain, resId := range resIds {
+		if domain != "" {
+			f.resId.Store(domain, resId)
+		}
+	}
+
+	f.domains = newDomains
 }
 
 func (f *DNSForwarder) Close(ctx context.Context) error {
-	var result *multierror.Error
-
-	if f.dnsServer != nil {
-		if err := f.dnsServer.ShutdownContext(ctx); err != nil {
-			result = multierror.Append(result, fmt.Errorf("UDP shutdown: %w", err))
-		}
-	}
-	if f.tcpServer != nil {
-		if err := f.tcpServer.ShutdownContext(ctx); err != nil {
-			result = multierror.Append(result, fmt.Errorf("TCP shutdown: %w", err))
-		}
-	}
-
-	return nberrors.FormatErrorOrNil(result)
-}
-
-func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) *dns.Msg {
-	if len(query.Question) == 0 {
+	if f.dnsServer == nil {
 		return nil
 	}
-	question := query.Question[0]
-	log.Tracef("received DNS request for DNS forwarder: domain=%v type=%v class=%v",
-		question.Name, question.Qtype, question.Qclass)
+	return f.dnsServer.ShutdownContext(ctx)
+}
 
-	domain := strings.ToLower(question.Name)
+func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
+	if len(query.Question) == 0 {
+		return
+	}
+	log.Tracef("received DNS request for DNS forwarder: domain=%v type=%v class=%v",
+		query.Question[0].Name, query.Question[0].Qtype, query.Question[0].Qclass)
+
+	question := query.Question[0]
+	domain := question.Name
 
 	resp := query.SetReply(query)
 	var network string
@@ -154,96 +111,41 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) *dns
 		if err := w.WriteMsg(resp); err != nil {
 			log.Errorf("failed to write DNS response: %v", err)
 		}
-		return nil
+		return
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), upstreamTimeout)
 	defer cancel()
 	ips, err := net.DefaultResolver.LookupNetIP(ctx, network, domain)
 	if err != nil {
-		f.handleDNSError(w, query, resp, domain, err)
-		return nil
+		f.handleDNSError(w, resp, domain, err)
+		return
+	}
+
+	resId := f.getResIdForDomain(strings.TrimSuffix(domain, "."))
+	if resId != "" {
+		for _, ip := range ips {
+			var ipWithSuffix string
+			if ip.Is4() {
+				ipWithSuffix = ip.String() + "/32"
+				log.Tracef("resolved domain=%s to IPv4=%s", domain, ipWithSuffix)
+			} else {
+				ipWithSuffix = ip.String() + "/128"
+				log.Tracef("resolved domain=%s to IPv6=%s", domain, ipWithSuffix)
+			}
+			f.statusRecorder.AddResolvedIPLookupEntry(ipWithSuffix, resId)
+		}
 	}
 
-	f.updateInternalState(domain, ips)
 	f.addIPsToResponse(resp, domain, ips)
 
-	return resp
-}
-
-func (f *DNSForwarder) handleDNSQueryUDP(w dns.ResponseWriter, query *dns.Msg) {
-
-	resp := f.handleDNSQuery(w, query)
-	if resp == nil {
-		return
-	}
-
-	opt := query.IsEdns0()
-	maxSize := dns.MinMsgSize
-	if opt != nil {
-		// client advertised a larger EDNS0 buffer
-		maxSize = int(opt.UDPSize())
-	}
-
-	// if our response is too big, truncate and set the TC bit
-	if resp.Len() > maxSize {
-		resp.Truncate(maxSize)
-	}
-
 	if err := w.WriteMsg(resp); err != nil {
 		log.Errorf("failed to write DNS response: %v", err)
 	}
 }
 
-func (f *DNSForwarder) handleDNSQueryTCP(w dns.ResponseWriter, query *dns.Msg) {
-	resp := f.handleDNSQuery(w, query)
-	if resp == nil {
-		return
-	}
-
-	if err := w.WriteMsg(resp); err != nil {
-		log.Errorf("failed to write DNS response: %v", err)
-	}
-}
-
-func (f *DNSForwarder) updateInternalState(domain string, ips []netip.Addr) {
-	var prefixes []netip.Prefix
-	mostSpecificResId, matchingEntries := f.getMatchingEntries(strings.TrimSuffix(domain, "."))
-	if mostSpecificResId != "" {
-		for _, ip := range ips {
-			var prefix netip.Prefix
-			if ip.Is4() {
-				prefix = netip.PrefixFrom(ip, 32)
-			} else {
-				prefix = netip.PrefixFrom(ip, 128)
-			}
-			prefixes = append(prefixes, prefix)
-			f.statusRecorder.AddResolvedIPLookupEntry(prefix, mostSpecificResId)
-		}
-	}
-
-	if f.firewall != nil {
-		f.updateFirewall(matchingEntries, prefixes)
-	}
-}
-
-func (f *DNSForwarder) updateFirewall(matchingEntries []*ForwarderEntry, prefixes []netip.Prefix) {
-	var merr *multierror.Error
-	for _, entry := range matchingEntries {
-		if err := f.firewall.UpdateSet(entry.Set, prefixes); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("update set for domain=%s: %w", entry.Domain, err))
-		}
-	}
-	if merr != nil {
-		log.Errorf("failed to update firewall sets (%d/%d): %v",
-			len(merr.Errors),
-			len(matchingEntries),
-			nberrors.FormatErrorOrNil(merr))
-	}
-}
-
 // handleDNSError processes DNS lookup errors and sends an appropriate error response
-func (f *DNSForwarder) handleDNSError(w dns.ResponseWriter, query, resp *dns.Msg, domain string, err error) {
+func (f *DNSForwarder) handleDNSError(w dns.ResponseWriter, resp *dns.Msg, domain string, err error) {
 	var dnsErr *net.DNSError
 
 	switch {
@@ -255,7 +157,7 @@ func (f *DNSForwarder) handleDNSError(w dns.ResponseWriter, query, resp *dns.Msg
 		}
 
 		if dnsErr.Server != "" {
-			log.Warnf("failed to resolve query for type=%s domain=%s server=%s: %v", dns.TypeToString[query.Question[0].Qtype], domain, dnsErr.Server, err)
+			log.Warnf("failed to resolve query for domain=%s server=%s: %v", domain, dnsErr.Server, err)
 		} else {
 			log.Warnf(errResolveFailed, domain, err)
 		}
@@ -302,53 +204,45 @@ func (f *DNSForwarder) addIPsToResponse(resp *dns.Msg, domain string, ips []neti
 	}
 }
 
-// getMatchingEntries retrieves the resource IDs for a given domain.
-// It returns the most specific match and all matching resource IDs.
-func (f *DNSForwarder) getMatchingEntries(domain string) (route.ResID, []*ForwarderEntry) {
-	var selectedResId route.ResID
+func (f *DNSForwarder) getResIdForDomain(domain string) string {
+	var selectedResId string
 	var bestScore int
-	var matches []*ForwarderEntry
 
-	f.mutex.RLock()
-	defer f.mutex.RUnlock()
-
-	for _, entry := range f.fwdEntries {
+	f.resId.Range(func(key, value interface{}) bool {
 		var score int
-		pattern := entry.Domain.PunycodeString()
+		pattern := key.(string)
 
 		switch {
 		case strings.HasPrefix(pattern, "*."):
 			baseDomain := strings.TrimPrefix(pattern, "*.")
-
-			if strings.EqualFold(domain, baseDomain) || strings.HasSuffix(domain, "."+baseDomain) {
+			if domain == baseDomain || strings.HasSuffix(domain, "."+baseDomain) {
 				score = len(baseDomain)
-				matches = append(matches, entry)
 			}
 		case domain == pattern:
 			score = math.MaxInt
-			matches = append(matches, entry)
 		default:
-			continue
+			return true
 		}
 
 		if score > bestScore {
 			bestScore = score
-			selectedResId = entry.ResID
+			selectedResId = value.(string)
 		}
-	}
+		return true
+	})
 
-	return selectedResId, matches
+	return selectedResId
 }
 
 // filterDomains returns a list of normalized domains
-func filterDomains(entries []*ForwarderEntry) domain.List {
-	newDomains := make(domain.List, 0, len(entries))
-	for _, d := range entries {
-		if d.Domain == "" {
+func filterDomains(domains []string) []string {
+	newDomains := make([]string, 0, len(domains))
+	for _, d := range domains {
+		if d == "" {
 			log.Warn("empty domain in DNS forwarder")
 			continue
 		}
-		newDomains = append(newDomains, domain.Domain(nbdns.NormalizeZone(d.Domain.PunycodeString())))
+		newDomains = append(newDomains, nbdns.NormalizeZone(d))
 	}
 	return newDomains
 }

client/internal/dnsfwd/forwarder_test.go

@@ -1,61 +1,56 @@
 package dnsfwd
 
 import (
+	"sync"
 	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/management/domain"
-	"github.com/netbirdio/netbird/route"
 )
 
-func Test_getMatchingEntries(t *testing.T) {
+func TestGetResIdForDomain(t *testing.T) {
 	testCases := []struct {
 		name           string
-		storedMappings map[string]route.ResID // key: domain pattern, value: resId
+		storedMappings map[string]string // key: domain pattern, value: resId
 		queryDomain    string
-		expectedResId  route.ResID
+		expectedResId  string
 	}{
 		{
 			name:           "Empty map returns empty string",
-			storedMappings: map[string]route.ResID{},
+			storedMappings: map[string]string{},
 			queryDomain:    "example.com",
 			expectedResId:  "",
 		},
 		{
 			name:           "Exact match returns stored resId",
-			storedMappings: map[string]route.ResID{"example.com": "res1"},
+			storedMappings: map[string]string{"example.com": "res1"},
 			queryDomain:    "example.com",
 			expectedResId:  "res1",
 		},
 		{
 			name:           "Wildcard pattern matches base domain",
-			storedMappings: map[string]route.ResID{"*.example.com": "res2"},
+			storedMappings: map[string]string{"*.example.com": "res2"},
 			queryDomain:    "example.com",
 			expectedResId:  "res2",
 		},
 		{
 			name:           "Wildcard pattern matches subdomain",
-			storedMappings: map[string]route.ResID{"*.example.com": "res3"},
+			storedMappings: map[string]string{"*.example.com": "res3"},
 			queryDomain:    "foo.example.com",
 			expectedResId:  "res3",
 		},
 		{
 			name:           "Wildcard pattern does not match different domain",
-			storedMappings: map[string]route.ResID{"*.example.com": "res4"},
+			storedMappings: map[string]string{"*.example.com": "res4"},
 			queryDomain:    "foo.notexample.com",
 			expectedResId:  "",
 		},
 		{
 			name:           "Non-wildcard pattern does not match subdomain",
-			storedMappings: map[string]route.ResID{"example.com": "res5"},
+			storedMappings: map[string]string{"example.com": "res5"},
 			queryDomain:    "foo.example.com",
 			expectedResId:  "",
 		},
 		{
 			name: "Exact match over overlapping wildcard",
-			storedMappings: map[string]route.ResID{
+			storedMappings: map[string]string{
 				"*.example.com":   "resWildcard",
 				"foo.example.com": "resExact",
 			},
@@ -64,7 +59,7 @@ func Test_getMatchingEntries(t *testing.T) {
 		},
 		{
 			name: "Overlapping wildcards: Select more specific wildcard",
-			storedMappings: map[string]route.ResID{
+			storedMappings: map[string]string{
 				"*.example.com":     "resA",
 				"*.sub.example.com": "resB",
 			},
@@ -73,7 +68,7 @@ func Test_getMatchingEntries(t *testing.T) {
 		},
 		{
 			name: "Wildcard multi-level subdomain match",
-			storedMappings: map[string]route.ResID{
+			storedMappings: map[string]string{
 				"*.example.com": "resMulti",
 			},
 			queryDomain:   "a.b.example.com",
@@ -83,21 +78,18 @@ func Test_getMatchingEntries(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			fwd := &DNSForwarder{}
-
-			var entries []*ForwarderEntry
-			for domainPattern, resId := range tc.storedMappings {
-				d, err := domain.FromString(domainPattern)
-				require.NoError(t, err)
-				entries = append(entries, &ForwarderEntry{
-					Domain: d,
-					ResID:  resId,
-				})
+			fwd := &DNSForwarder{
+				resId: sync.Map{},
 			}
-			fwd.UpdateDomains(entries)
 
-			got, _ := fwd.getMatchingEntries(tc.queryDomain)
-			assert.Equal(t, got, tc.expectedResId)
+			for domainPattern, resId := range tc.storedMappings {
+				fwd.resId.Store(domainPattern, resId)
+			}
+
+			got := fwd.getResIdForDomain(tc.queryDomain)
+			if got != tc.expectedResId {
+				t.Errorf("For query domain %q, expected resId %q, but got %q", tc.queryDomain, tc.expectedResId, got)
+			}
 		})
 	}
 }

client/internal/dnsfwd/manager.go

@@ -11,8 +11,6 @@ import (
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/management/domain"
-	"github.com/netbirdio/netbird/route"
 )
 
 const (
@@ -21,19 +19,11 @@ const (
 	dnsTTL     = 60 //seconds
 )
 
-// ForwarderEntry is a mapping from a domain to a resource ID and a hash of the parent domain list.
-type ForwarderEntry struct {
-	Domain domain.Domain
-	ResID  route.ResID
-	Set    firewall.Set
-}
-
 type Manager struct {
 	firewall       firewall.Manager
 	statusRecorder *peer.Status
 
 	fwRules      []firewall.Rule
-	tcpRules     []firewall.Rule
 	dnsForwarder *DNSForwarder
 }
 
@@ -44,7 +34,7 @@ func NewManager(fw firewall.Manager, statusRecorder *peer.Status) *Manager {
 	}
 }
 
-func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
+func (m *Manager) Start(domains []string, resIds map[string]string) error {
 	log.Infof("starting DNS forwarder")
 	if m.dnsForwarder != nil {
 		return nil
@@ -54,9 +44,9 @@ func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
 		return err
 	}
 
-	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.firewall, m.statusRecorder)
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.statusRecorder)
 	go func() {
-		if err := m.dnsForwarder.Listen(fwdEntries); err != nil {
+		if err := m.dnsForwarder.Listen(domains, resIds); err != nil {
 			// todo handle close error if it is exists
 			log.Errorf("failed to start DNS forwarder, err: %v", err)
 		}
@@ -65,12 +55,12 @@ func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
 	return nil
 }
 
-func (m *Manager) UpdateDomains(entries []*ForwarderEntry) {
+func (m *Manager) UpdateDomains(domains []string, resIds map[string]string) {
 	if m.dnsForwarder == nil {
 		return
 	}
 
-	m.dnsForwarder.UpdateDomains(entries)
+	m.dnsForwarder.UpdateDomains(domains, resIds)
 }
 
 func (m *Manager) Stop(ctx context.Context) error {
@@ -91,47 +81,34 @@ func (m *Manager) Stop(ctx context.Context) error {
 	return nberrors.FormatErrorOrNil(mErr)
 }
 
-func (m *Manager) allowDNSFirewall() error {
+func (h *Manager) allowDNSFirewall() error {
 	dport := &firewall.Port{
 		IsRange: false,
 		Values:  []uint16{ListenPort},
 	}
 
-	if m.firewall == nil {
+	if h.firewall == nil {
 		return nil
 	}
 
-	dnsRules, err := m.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
+	dnsRules, err := h.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
 	if err != nil {
 		log.Errorf("failed to add allow DNS router rules, err: %v", err)
 		return err
 	}
-	m.fwRules = dnsRules
-
-	tcpRules, err := m.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolTCP, nil, dport, firewall.ActionAccept, "")
-	if err != nil {
-		log.Errorf("failed to add allow DNS router rules, err: %v", err)
-		return err
-	}
-	m.tcpRules = tcpRules
+	h.fwRules = dnsRules
 
 	return nil
 }
 
-func (m *Manager) dropDNSFirewall() error {
+func (h *Manager) dropDNSFirewall() error {
 	var mErr *multierror.Error
-	for _, rule := range m.fwRules {
-		if err := m.firewall.DeletePeerRule(rule); err != nil {
-			mErr = multierror.Append(mErr, fmt.Errorf("failed to delete DNS router rules, err: %v", err))
-		}
-	}
-	for _, rule := range m.tcpRules {
-		if err := m.firewall.DeletePeerRule(rule); err != nil {
+	for _, rule := range h.fwRules {
+		if err := h.firewall.DeletePeerRule(rule); err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf("failed to delete DNS router rules, err: %v", err))
 		}
 	}
 
-	m.fwRules = nil
-	m.tcpRules = nil
+	h.fwRules = nil
 	return nberrors.FormatErrorOrNil(mErr)
 }

client/internal/engine.go

@@ -38,7 +38,6 @@ import (
 	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/networkmonitor"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
@@ -123,8 +122,6 @@ type EngineConfig struct {
 	DisableFirewall     bool
 
 	BlockLANAccess bool
-
-	LazyConnectionEnabled bool
 }
 
 // Engine is a mechanism responsible for reacting on Signal and Management stream events and managing connections to the remote peers.
@@ -137,8 +134,6 @@ type Engine struct {
 	// peerConns is a map that holds all the peers that are known to this peer
 	peerStore *peerstore.Store
 
-	connMgr *ConnMgr
-
 	beforePeerHook nbnet.AddHookFunc
 	afterPeerHook  nbnet.RemoveHookFunc
 
@@ -175,8 +170,7 @@ type Engine struct {
 	sshServerFunc func(hostKeyPEM []byte, addr string) (nbssh.Server, error)
 	sshServer     nbssh.Server
 
-	statusRecorder     *peer.Status
-	peerConnDispatcher *dispatcher.ConnectionDispatcher
+	statusRecorder *peer.Status
 
 	firewall          firewallManager.Manager
 	routeManager      routemanager.Manager
@@ -241,8 +235,6 @@ func NewEngine(
 		checks:         checks,
 		connSemaphore:  semaphoregroup.NewSemaphoreGroup(connInitLimit),
 	}
-
-	path := statemanager.GetDefaultStatePath()
 	if runtime.GOOS == "ios" {
 		if !fileExists(mobileDep.StateFilePath) {
 			err := createFile(mobileDep.StateFilePath)
@@ -252,9 +244,11 @@ func NewEngine(
 			}
 		}
 
-		path = mobileDep.StateFilePath
+		engine.stateManager = statemanager.New(mobileDep.StateFilePath)
+	}
+	if path := statemanager.GetDefaultStatePath(); path != "" {
+		engine.stateManager = statemanager.New(path)
 	}
-	engine.stateManager = statemanager.New(path)
 
 	return engine
 }
@@ -268,10 +262,6 @@ func (e *Engine) Stop() error {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 
-	if e.connMgr != nil {
-		e.connMgr.Close()
-	}
-
 	// stopping network monitor first to avoid starting the engine again
 	if e.networkMonitor != nil {
 		e.networkMonitor.Stop()
@@ -307,7 +297,8 @@ func (e *Engine) Stop() error {
 	e.statusRecorder.UpdateDNSStates([]peer.NSGroupState{})
 	e.statusRecorder.UpdateRelayStates([]relay.ProbeResult{})
 
-	if err := e.removeAllPeers(); err != nil {
+	err := e.removeAllPeers()
+	if err != nil {
 		return fmt.Errorf("failed to remove all peers: %s", err)
 	}
 
@@ -414,7 +405,8 @@ func (e *Engine) Start() error {
 
 	e.routeManager.SetRouteChangeListener(e.mobileDep.NetworkChangeListener)
 
-	if err = e.wgInterfaceCreate(); err != nil {
+	err = e.wgInterfaceCreate()
+	if err != nil {
 		log.Errorf("failed creating tunnel interface %s: [%s]", e.config.WgIfaceName, err.Error())
 		e.close()
 		return fmt.Errorf("create wg interface: %w", err)
@@ -450,11 +442,6 @@ func (e *Engine) Start() error {
 		NATExternalIPs:       e.parseNATExternalIPMappings(),
 	}
 
-	e.peerConnDispatcher = dispatcher.NewConnectionDispatcher()
-
-	e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface, e.peerConnDispatcher)
-	e.connMgr.Start(e.ctx)
-
 	e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg)
 	e.srWatcher.Start()
 
@@ -463,6 +450,7 @@ func (e *Engine) Start() error {
 
 	// starting network monitor at the very last to avoid disruptions
 	e.startNetworkMonitor()
+
 	return nil
 }
 
@@ -539,7 +527,7 @@ func (e *Engine) blockLanAccess() {
 		if _, err := e.firewall.AddRouteFiltering(
 			nil,
 			[]netip.Prefix{v4},
-			firewallManager.Network{Prefix: network},
+			network,
 			firewallManager.ProtocolALL,
 			nil,
 			nil,
@@ -562,16 +550,6 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 	var modified []*mgmProto.RemotePeerConfig
 	for _, p := range peersUpdate {
 		peerPubKey := p.GetWgPubKey()
-		currentPeer, ok := e.peerStore.PeerConn(peerPubKey)
-		if !ok {
-			continue
-		}
-
-		if currentPeer.AgentVersionString() != p.AgentVersion {
-			modified = append(modified, p)
-			continue
-		}
-
 		allowedIPs, ok := e.peerStore.AllowedIPs(peerPubKey)
 		if !ok {
 			continue
@@ -581,7 +559,8 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 			continue
 		}
 
-		if err := e.statusRecorder.UpdatePeerFQDN(peerPubKey, p.GetFqdn()); err != nil {
+		err := e.statusRecorder.UpdatePeerFQDN(peerPubKey, p.GetFqdn())
+		if err != nil {
 			log.Warnf("error updating peer's %s fqdn in the status recorder, got error: %v", peerPubKey, err)
 		}
 	}
@@ -642,11 +621,16 @@ func (e *Engine) removePeer(peerKey string) error {
 		e.sshServer.RemoveAuthorizedKey(peerKey)
 	}
 
-	e.connMgr.RemovePeerConn(peerKey)
+	defer func() {
+		err := e.statusRecorder.RemovePeer(peerKey)
+		if err != nil {
+			log.Warnf("received error when removing peer %s from status recorder: %v", peerKey, err)
+		}
+	}()
 
-	err := e.statusRecorder.RemovePeer(peerKey)
-	if err != nil {
-		log.Warnf("received error when removing peer %s from status recorder: %v", peerKey, err)
+	conn, exists := e.peerStore.Remove(peerKey)
+	if exists {
+		conn.Close()
 	}
 	return nil
 }
@@ -968,44 +952,31 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		return nil
 	}
 
-	if err := e.connMgr.UpdatedRemoteFeatureFlag(e.ctx, networkMap.GetPeerConfig().GetLazyConnectionEnabled()); err != nil {
-		log.Errorf("failed to update lazy connection feature flag: %v", err)
-	}
-
 	if e.firewall != nil {
 		if localipfw, ok := e.firewall.(localIpUpdater); ok {
 			if err := localipfw.UpdateLocalIPs(); err != nil {
 				log.Errorf("failed to update local IPs: %v", err)
 			}
 		}
-
-		// If we got empty rules list but management did not set the networkMap.FirewallRulesIsEmpty flag,
-		// then the mgmt server is older than the client, and we need to allow all traffic for routes.
-		// This needs to be toggled before applying routes.
-		isLegacy := len(networkMap.RoutesFirewallRules) == 0 && !networkMap.RoutesFirewallRulesIsEmpty
-		if err := e.firewall.SetLegacyManagement(isLegacy); err != nil {
-			log.Errorf("failed to set legacy management flag: %v", err)
-		}
 	}
 
+	// DNS forwarder
 	dnsRouteFeatureFlag := toDNSFeatureFlag(networkMap)
+	dnsRouteDomains, resourceIds := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
+	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains, resourceIds)
 
-	// apply routes first, route related actions might depend on routing being enabled
 	routes := toRoutes(networkMap.GetRoutes())
 	if err := e.routeManager.UpdateRoutes(serial, routes, dnsRouteFeatureFlag); err != nil {
 		log.Errorf("failed to update clientRoutes, err: %v", err)
 	}
 
+	// acls might need routing to be enabled, so we apply after routes
 	if e.acl != nil {
-		e.acl.ApplyFiltering(networkMap, dnsRouteFeatureFlag)
+		e.acl.ApplyFiltering(networkMap)
 	}
 
-	fwdEntries := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), routes)
-	e.updateDNSForwarder(dnsRouteFeatureFlag, fwdEntries)
-
 	// Ingress forward rules
-	forwardingRules, err := e.updateForwardRules(networkMap.GetForwardingRules())
-	if err != nil {
+	if err := e.updateForwardRules(networkMap.GetForwardingRules()); err != nil {
 		log.Errorf("failed to update forward rules, err: %v", err)
 	}
 
@@ -1051,10 +1022,6 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		}
 	}
 
-	// must set the exclude list after the peers are added. Without it the manager can not figure out the peers parameters from the store
-	excludedLazyPeers := e.toExcludedLazyPeers(routes, forwardingRules, networkMap.GetRemotePeers())
-	e.connMgr.SetExcludeList(excludedLazyPeers)
-
 	protoDNSConfig := networkMap.GetDNSConfig()
 	if protoDNSConfig == nil {
 		protoDNSConfig = &mgmProto.DNSConfig{}
@@ -1112,24 +1079,29 @@ func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
 	return routes
 }
 
-func toRouteDomains(myPubKey string, routes []*route.Route) []*dnsfwd.ForwarderEntry {
-	var entries []*dnsfwd.ForwarderEntry
-	for _, route := range routes {
-		if len(route.Domains) == 0 {
+func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) ([]string, map[string]string) {
+	if protoRoutes == nil {
+		protoRoutes = []*mgmProto.Route{}
+	}
+
+	var dnsRoutes []string
+	resIds := make(map[string]string)
+	for _, protoRoute := range protoRoutes {
+		if len(protoRoute.Domains) == 0 {
 			continue
 		}
-		if route.Peer == myPubKey {
-			domainSet := firewallManager.NewDomainSet(route.Domains)
-			for _, d := range route.Domains {
-				entries = append(entries, &dnsfwd.ForwarderEntry{
-					Domain: d,
-					Set:    domainSet,
-					ResID:  route.GetResourceID(),
-				})
+		if protoRoute.Peer == myPubKey {
+			dnsRoutes = append(dnsRoutes, protoRoute.Domains...)
+			// resource ID is the first part of the ID
+			resId := strings.Split(protoRoute.ID, ":")
+			for _, domain := range protoRoute.Domains {
+				if len(resId) > 0 {
+					resIds[domain] = resId[0]
+				}
 			}
 		}
 	}
-	return entries
+	return dnsRoutes, resIds
 }
 
 func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig, network *net.IPNet) nbdns.Config {
@@ -1188,7 +1160,7 @@ func (e *Engine) updateOfflinePeers(offlinePeers []*mgmProto.RemotePeerConfig) {
 			IP:               strings.Join(offlinePeer.GetAllowedIps(), ","),
 			PubKey:           offlinePeer.GetWgPubKey(),
 			FQDN:             offlinePeer.GetFqdn(),
-			ConnStatus:       peer.StatusIdle,
+			ConnStatus:       peer.StatusDisconnected,
 			ConnStatusUpdate: time.Now(),
 			Mux:              new(sync.RWMutex),
 		}
@@ -1224,17 +1196,12 @@ func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error {
 		peerIPs = append(peerIPs, allowedNetIP)
 	}
 
-	conn, err := e.createPeerConn(peerKey, peerIPs, peerConfig.AgentVersion)
+	conn, err := e.createPeerConn(peerKey, peerIPs)
 	if err != nil {
 		return fmt.Errorf("create peer connection: %w", err)
 	}
 
-	err = e.statusRecorder.AddPeer(peerKey, peerConfig.Fqdn, peerIPs[0].Addr().String())
-	if err != nil {
-		log.Warnf("error adding peer %s to status recorder, got error: %v", peerKey, err)
-	}
-
-	if exists := e.connMgr.AddPeerConn(e.ctx, peerKey, conn); exists {
+	if ok := e.peerStore.AddPeerConn(peerKey, conn); !ok {
 		conn.Close()
 		return fmt.Errorf("peer already exists: %s", peerKey)
 	}
@@ -1243,10 +1210,17 @@ func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error {
 		conn.AddBeforeAddPeerHook(e.beforePeerHook)
 		conn.AddAfterRemovePeerHook(e.afterPeerHook)
 	}
+
+	err = e.statusRecorder.AddPeer(peerKey, peerConfig.Fqdn)
+	if err != nil {
+		log.Warnf("error adding peer %s to status recorder, got error: %v", peerKey, err)
+	}
+
+	conn.Open()
 	return nil
 }
 
-func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentVersion string) (*peer.Conn, error) {
+func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix) (*peer.Conn, error) {
 	log.Debugf("creating peer connection %s", pubKey)
 
 	wgConfig := peer.WgConfig{
@@ -1260,12 +1234,11 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 	// randomize connection timeout
 	timeout := time.Duration(rand.Intn(PeerConnectionTimeoutMax-PeerConnectionTimeoutMin)+PeerConnectionTimeoutMin) * time.Millisecond
 	config := peer.ConnConfig{
-		Key:          pubKey,
-		LocalKey:     e.config.WgPrivateKey.PublicKey().String(),
-		AgentVersion: agentVersion,
-		Timeout:      timeout,
-		WgConfig:     wgConfig,
-		LocalWgPort:  e.config.WgPort,
+		Key:         pubKey,
+		LocalKey:    e.config.WgPrivateKey.PublicKey().String(),
+		Timeout:     timeout,
+		WgConfig:    wgConfig,
+		LocalWgPort: e.config.WgPort,
 		RosenpassConfig: peer.RosenpassConfig{
 			PubKey:         e.getRosenpassPubKey(),
 			Addr:           e.getRosenpassAddr(),
@@ -1281,16 +1254,7 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 		},
 	}
 
-	serviceDependencies := peer.ServiceDependencies{
-		StatusRecorder:     e.statusRecorder,
-		Signaler:           e.signaler,
-		IFaceDiscover:      e.mobileDep.IFaceDiscover,
-		RelayManager:       e.relayManager,
-		SrWatcher:          e.srWatcher,
-		Semaphore:          e.connSemaphore,
-		PeerConnDispatcher: e.peerConnDispatcher,
-	}
-	peerConn, err := peer.NewConn(config, serviceDependencies)
+	peerConn, err := peer.NewConn(e.ctx, config, e.statusRecorder, e.signaler, e.mobileDep.IFaceDiscover, e.relayManager, e.srWatcher, e.connSemaphore)
 	if err != nil {
 		return nil, err
 	}
@@ -1311,7 +1275,7 @@ func (e *Engine) receiveSignalEvents() {
 			e.syncMsgMux.Lock()
 			defer e.syncMsgMux.Unlock()
 
-			conn, ok := e.connMgr.OnSignalMsg(e.ctx, msg.Key)
+			conn, ok := e.peerStore.PeerConn(msg.Key)
 			if !ok {
 				return fmt.Errorf("wrongly addressed message %s", msg.Key)
 			}
@@ -1619,39 +1583,13 @@ func (e *Engine) getRosenpassAddr() string {
 // RunHealthProbes executes health checks for Signal, Management, Relay and WireGuard services
 // and updates the status recorder with the latest states.
 func (e *Engine) RunHealthProbes() bool {
-	e.syncMsgMux.Lock()
-
 	signalHealthy := e.signal.IsHealthy()
 	log.Debugf("signal health check: healthy=%t", signalHealthy)
 
 	managementHealthy := e.mgmClient.IsHealthy()
 	log.Debugf("management health check: healthy=%t", managementHealthy)
 
-	stuns := slices.Clone(e.STUNs)
-	turns := slices.Clone(e.TURNs)
-
-	if e.wgInterface != nil {
-		stats, err := e.wgInterface.GetStats()
-		if err != nil {
-			log.Warnf("failed to get wireguard stats: %v", err)
-			e.syncMsgMux.Unlock()
-			return false
-		}
-		for _, key := range e.peerStore.PeersPubKey() {
-			// wgStats could be zero value, in which case we just reset the stats
-			wgStats, ok := stats[key]
-			if !ok {
-				continue
-			}
-			if err := e.statusRecorder.UpdateWireGuardPeerState(key, wgStats); err != nil {
-				log.Debugf("failed to update wg stats for peer %s: %s", key, err)
-			}
-		}
-	}
-
-	e.syncMsgMux.Unlock()
-
-	results := e.probeICE(stuns, turns)
+	results := append(e.probeSTUNs(), e.probeTURNs()...)
 	e.statusRecorder.UpdateRelayStates(results)
 
 	relayHealthy := true
@@ -1663,16 +1601,37 @@ func (e *Engine) RunHealthProbes() bool {
 	}
 	log.Debugf("relay health check: healthy=%t", relayHealthy)
 
+	for _, key := range e.peerStore.PeersPubKey() {
+		wgStats, err := e.wgInterface.GetStats(key)
+		if err != nil {
+			log.Debugf("failed to get wg stats for peer %s: %s", key, err)
+			continue
+		}
+		// wgStats could be zero value, in which case we just reset the stats
+		if err := e.statusRecorder.UpdateWireGuardPeerState(key, wgStats); err != nil {
+			log.Debugf("failed to update wg stats for peer %s: %s", key, err)
+		}
+	}
+
 	allHealthy := signalHealthy && managementHealthy && relayHealthy
 	log.Debugf("all health checks completed: healthy=%t", allHealthy)
 	return allHealthy
 }
 
-func (e *Engine) probeICE(stuns, turns []*stun.URI) []relay.ProbeResult {
-	return append(
-		relay.ProbeAll(e.ctx, relay.ProbeSTUN, stuns),
-		relay.ProbeAll(e.ctx, relay.ProbeSTUN, turns)...,
-	)
+func (e *Engine) probeSTUNs() []relay.ProbeResult {
+	e.syncMsgMux.Lock()
+	stuns := slices.Clone(e.STUNs)
+	e.syncMsgMux.Unlock()
+
+	return relay.ProbeAll(e.ctx, relay.ProbeSTUN, stuns)
+}
+
+func (e *Engine) probeTURNs() []relay.ProbeResult {
+	e.syncMsgMux.Lock()
+	turns := slices.Clone(e.TURNs)
+	e.syncMsgMux.Unlock()
+
+	return relay.ProbeAll(e.ctx, relay.ProbeTURN, turns)
 }
 
 // restartEngine restarts the engine by cancelling the client context
@@ -1792,10 +1751,7 @@ func (e *Engine) GetWgAddr() net.IP {
 }
 
 // updateDNSForwarder start or stop the DNS forwarder based on the domains and the feature flag
-func (e *Engine) updateDNSForwarder(
-	enabled bool,
-	fwdEntries []*dnsfwd.ForwarderEntry,
-) {
+func (e *Engine) updateDNSForwarder(enabled bool, domains []string, resIds map[string]string) {
 	if !enabled {
 		if e.dnsForwardMgr == nil {
 			return
@@ -1806,18 +1762,18 @@ func (e *Engine) updateDNSForwarder(
 		return
 	}
 
-	if len(fwdEntries) > 0 {
+	if len(domains) > 0 {
+		log.Infof("enable domain router service for domains: %v", domains)
 		if e.dnsForwardMgr == nil {
 			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder)
 
-			if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
+			if err := e.dnsForwardMgr.Start(domains, resIds); err != nil {
 				log.Errorf("failed to start DNS forward: %v", err)
 				e.dnsForwardMgr = nil
 			}
-
-			log.Infof("started domain router service with %d entries", len(fwdEntries))
 		} else {
-			e.dnsForwardMgr.UpdateDomains(fwdEntries)
+			log.Infof("update domain router service for domains: %v", domains)
+			e.dnsForwardMgr.UpdateDomains(domains, resIds)
 		}
 	} else if e.dnsForwardMgr != nil {
 		log.Infof("disable domain router service")
@@ -1859,21 +1815,21 @@ func (e *Engine) Address() (netip.Addr, error) {
 	return ip.Unmap(), nil
 }
 
-func (e *Engine) updateForwardRules(rules []*mgmProto.ForwardingRule) ([]firewallManager.ForwardRule, error) {
+func (e *Engine) updateForwardRules(rules []*mgmProto.ForwardingRule) error {
 	if e.firewall == nil {
 		log.Warn("firewall is disabled, not updating forwarding rules")
-		return nil, nil
+		return nil
 	}
 
 	if len(rules) == 0 {
 		if e.ingressGatewayMgr == nil {
-			return nil, nil
+			return nil
 		}
 
 		err := e.ingressGatewayMgr.Close()
 		e.ingressGatewayMgr = nil
 		e.statusRecorder.SetIngressGwMgr(nil)
-		return nil, err
+		return err
 	}
 
 	if e.ingressGatewayMgr == nil {
@@ -1924,35 +1880,7 @@ func (e *Engine) updateForwardRules(rules []*mgmProto.ForwardingRule) ([]firewal
 		log.Errorf("failed to update forwarding rules: %v", err)
 	}
 
-	return forwardingRules, nberrors.FormatErrorOrNil(merr)
-}
-
-func (e *Engine) toExcludedLazyPeers(routes []*route.Route, rules []firewallManager.ForwardRule, peers []*mgmProto.RemotePeerConfig) map[string]bool {
-	excludedPeers := make(map[string]bool)
-	for _, r := range routes {
-		if r.Peer == "" {
-			continue
-		}
-		if !excludedPeers[r.Peer] {
-			log.Infof("exclude router peer from lazy connection: %s", r.Peer)
-			excludedPeers[r.Peer] = true
-		}
-	}
-
-	for _, r := range rules {
-		ip := r.TranslatedAddress
-		for _, p := range peers {
-			for _, allowedIP := range p.GetAllowedIps() {
-				if allowedIP != ip.String() {
-					continue
-				}
-				log.Infof("exclude forwarder peer from lazy connection: %s", p.GetWgPubKey())
-				excludedPeers[p.GetWgPubKey()] = true
-			}
-		}
-	}
-
-	return excludedPeers
+	return nberrors.FormatErrorOrNil(merr)
 }
 
 // isChecksEqual checks if two slices of checks are equal.

client/internal/engine_test.go

@@ -28,6 +28,8 @@ import (
 
 	"github.com/netbirdio/management-integrations/integrations"
 
+	"github.com/netbirdio/netbird/management/server/types"
+
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/configurer"
@@ -36,7 +38,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/wgproxy"
 	"github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/routemanager"
@@ -52,7 +53,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
-	"github.com/netbirdio/netbird/management/server/types"
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	"github.com/netbirdio/netbird/route"
 	signal "github.com/netbirdio/netbird/signal/client"
@@ -93,7 +93,7 @@ type MockWGIface struct {
 	GetFilterFunc              func() device.PacketFilter
 	GetDeviceFunc              func() *device.FilteredDevice
 	GetWGDeviceFunc            func() *wgdevice.Device
-	GetStatsFunc               func() (map[string]configurer.WGStats, error)
+	GetStatsFunc               func(peerKey string) (configurer.WGStats, error)
 	GetInterfaceGUIDStringFunc func() (string, error)
 	GetProxyFunc               func() wgproxy.Proxy
 	GetNetFunc                 func() *netstack.Net
@@ -171,8 +171,8 @@ func (m *MockWGIface) GetWGDevice() *wgdevice.Device {
 	return m.GetWGDeviceFunc()
 }
 
-func (m *MockWGIface) GetStats() (map[string]configurer.WGStats, error) {
-	return m.GetStatsFunc()
+func (m *MockWGIface) GetStats(peerKey string) (configurer.WGStats, error) {
+	return m.GetStatsFunc(peerKey)
 }
 
 func (m *MockWGIface) GetProxy() wgproxy.Proxy {
@@ -378,9 +378,6 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 				},
 			}
 		},
-		UpdatePeerFunc: func(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error {
-			return nil
-		},
 	}
 	engine.wgInterface = wgIface
 	engine.routeManager = routemanager.NewManager(routemanager.ManagerConfig{
@@ -403,8 +400,6 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 	engine.udpMux = bind.NewUniversalUDPMuxDefault(bind.UniversalUDPMuxParams{UDPConn: conn})
 	engine.ctx = ctx
 	engine.srWatcher = guard.NewSRWatcher(nil, nil, nil, icemaker.Config{})
-	engine.connMgr = NewConnMgr(engine.config, engine.statusRecorder, engine.peerStore, wgIface, dispatcher.NewConnectionDispatcher())
-	engine.connMgr.Start(ctx)
 
 	type testCase struct {
 		name       string
@@ -775,8 +770,6 @@ func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {
 
 			engine.routeManager = mockRouteManager
 			engine.dnsServer = &dns.MockServer{}
-			engine.connMgr = NewConnMgr(engine.config, engine.statusRecorder, engine.peerStore, engine.wgInterface, dispatcher.NewConnectionDispatcher())
-			engine.connMgr.Start(ctx)
 
 			defer func() {
 				exitErr := engine.Stop()
@@ -973,8 +966,6 @@ func TestEngine_UpdateNetworkMapWithDNSUpdate(t *testing.T) {
 			}
 
 			engine.dnsServer = mockDNSServer
-			engine.connMgr = NewConnMgr(engine.config, engine.statusRecorder, engine.peerStore, engine.wgInterface, dispatcher.NewConnectionDispatcher())
-			engine.connMgr.Start(ctx)
 
 			defer func() {
 				exitErr := engine.Stop()
@@ -1485,7 +1476,7 @@ func getConnectedPeers(e *Engine) int {
 	i := 0
 	for _, id := range e.peerStore.PeersPubKey() {
 		conn, _ := e.peerStore.PeerConn(id)
-		if conn.IsConnected() {
+		if conn.Status() == peer.StatusConnected {
 			i++
 		}
 	}

client/internal/iface_common.go

@@ -35,6 +35,6 @@ type wgIfaceBase interface {
 	GetFilter() device.PacketFilter
 	GetDevice() *device.FilteredDevice
 	GetWGDevice() *wgdevice.Device
-	GetStats() (map[string]configurer.WGStats, error)
+	GetStats(peerKey string) (configurer.WGStats, error)
 	GetNet() *netstack.Net
 }

client/internal/lazyconn/activity/listen_ip.go

@@ -1,9 +0,0 @@
-//go:build !linux || android
-
-package activity
-
-import "net"
-
-var (
-	listenIP = net.IP{127, 0, 0, 1}
-)

client/internal/lazyconn/activity/listen_ip_linux.go

@@ -1,10 +0,0 @@
-//go:build !android
-
-package activity
-
-import "net"
-
-var (
-	// use this ip to avoid eBPF proxy congestion
-	listenIP = net.IP{127, 0, 1, 1}
-)

client/internal/lazyconn/activity/listener.go

@@ -1,106 +0,0 @@
-package activity
-
-import (
-	"fmt"
-	"net"
-	"sync"
-	"sync/atomic"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-)
-
-// Listener it is not a thread safe implementation, do not call Close before ReadPackets. It will cause blocking
-type Listener struct {
-	wgIface  lazyconn.WGIface
-	peerCfg  lazyconn.PeerConfig
-	conn     *net.UDPConn
-	endpoint *net.UDPAddr
-	done     sync.Mutex
-
-	isClosed atomic.Bool // use to avoid error log when closing the listener
-}
-
-func NewListener(wgIface lazyconn.WGIface, cfg lazyconn.PeerConfig) (*Listener, error) {
-	d := &Listener{
-		wgIface: wgIface,
-		peerCfg: cfg,
-	}
-
-	conn, err := d.newConn()
-	if err != nil {
-		return nil, fmt.Errorf("failed to creating activity listener: %v", err)
-	}
-	d.conn = conn
-	d.endpoint = conn.LocalAddr().(*net.UDPAddr)
-
-	if err := d.createEndpoint(); err != nil {
-		return nil, err
-	}
-	d.done.Lock()
-	cfg.Log.Infof("created activity listener: %s", conn.LocalAddr().(*net.UDPAddr).String())
-	return d, nil
-}
-
-func (d *Listener) ReadPackets() {
-	for {
-		n, remoteAddr, err := d.conn.ReadFromUDP(make([]byte, 1))
-		if err != nil {
-			if d.isClosed.Load() {
-				d.peerCfg.Log.Debugf("exit from activity listener")
-			} else {
-				d.peerCfg.Log.Errorf("failed to read from activity listener: %s", err)
-			}
-			break
-		}
-
-		if n < 1 {
-			d.peerCfg.Log.Warnf("received %d bytes from %s, too short", n, remoteAddr)
-			continue
-		}
-		break
-	}
-
-	if err := d.removeEndpoint(); err != nil {
-		d.peerCfg.Log.Errorf("failed to remove endpoint: %s", err)
-	}
-
-	_ = d.conn.Close() // do not care err because some cases it will return "use of closed network connection"
-	d.done.Unlock()
-}
-
-func (d *Listener) Close() {
-	d.peerCfg.Log.Infof("closing listener: %s", d.conn.LocalAddr().String())
-	d.isClosed.Store(true)
-
-	if err := d.conn.Close(); err != nil {
-		d.peerCfg.Log.Errorf("failed to close UDP listener: %s", err)
-	}
-	d.done.Lock()
-}
-
-func (d *Listener) removeEndpoint() error {
-	d.peerCfg.Log.Debugf("removing lazy endpoint: %s", d.endpoint.String())
-	return d.wgIface.RemovePeer(d.peerCfg.PublicKey)
-}
-
-func (d *Listener) createEndpoint() error {
-	d.peerCfg.Log.Debugf("creating lazy endpoint: %s", d.endpoint.String())
-	return d.wgIface.UpdatePeer(d.peerCfg.PublicKey, d.peerCfg.AllowedIPs, 0, d.endpoint, nil)
-}
-
-func (d *Listener) newConn() (*net.UDPConn, error) {
-	addr := &net.UDPAddr{
-		Port: 0,
-		IP:   listenIP,
-	}
-
-	conn, err := net.ListenUDP("udp", addr)
-	if err != nil {
-		log.Errorf("failed to create activity listener on %s: %s", addr, err)
-		return nil, err
-	}
-
-	return conn, nil
-}

client/internal/lazyconn/activity/listener_test.go

@@ -1,41 +0,0 @@
-package activity
-
-import (
-	"testing"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-)
-
-func TestNewListener(t *testing.T) {
-	peer := &MocPeer{
-		PeerID: "examplePublicKey1",
-	}
-
-	cfg := lazyconn.PeerConfig{
-		PublicKey:  peer.PeerID,
-		PeerConnID: peer.ConnID(),
-		Log:        log.WithField("peer", "examplePublicKey1"),
-	}
-
-	l, err := NewListener(MocWGIface{}, cfg)
-	if err != nil {
-		t.Fatalf("failed to create listener: %v", err)
-	}
-
-	chanClosed := make(chan struct{})
-	go func() {
-		defer close(chanClosed)
-		l.ReadPackets()
-	}()
-
-	time.Sleep(1 * time.Second)
-	l.Close()
-
-	select {
-	case <-chanClosed:
-	case <-time.After(time.Second):
-	}
-}

client/internal/lazyconn/activity/manager.go

@@ -1,95 +0,0 @@
-package activity
-
-import (
-	"sync"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-	peerid "github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-type Manager struct {
-	OnActivityChan chan peerid.ConnID
-
-	wgIface lazyconn.WGIface
-
-	peers map[peerid.ConnID]*Listener
-	done  chan struct{}
-
-	mu sync.Mutex
-}
-
-func NewManager(wgIface lazyconn.WGIface) *Manager {
-	m := &Manager{
-		OnActivityChan: make(chan peerid.ConnID, 1),
-		wgIface:        wgIface,
-		peers:          make(map[peerid.ConnID]*Listener),
-		done:           make(chan struct{}),
-	}
-	return m
-}
-
-func (m *Manager) MonitorPeerActivity(peerCfg lazyconn.PeerConfig) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if _, ok := m.peers[peerCfg.PeerConnID]; ok {
-		log.Warnf("activity listener already exists for: %s", peerCfg.PublicKey)
-		return nil
-	}
-
-	listener, err := NewListener(m.wgIface, peerCfg)
-	if err != nil {
-		return err
-	}
-	m.peers[peerCfg.PeerConnID] = listener
-
-	go m.waitForTraffic(listener, peerCfg.PeerConnID)
-	return nil
-}
-
-func (m *Manager) RemovePeer(log *log.Entry, peerConnID peerid.ConnID) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	listener, ok := m.peers[peerConnID]
-	if !ok {
-		return
-	}
-	log.Debugf("removing activity listener")
-	delete(m.peers, peerConnID)
-	listener.Close()
-}
-
-func (m *Manager) Close() {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	close(m.done)
-	for peerID, listener := range m.peers {
-		delete(m.peers, peerID)
-		listener.Close()
-	}
-}
-
-func (m *Manager) waitForTraffic(listener *Listener, peerConnID peerid.ConnID) {
-	listener.ReadPackets()
-
-	m.mu.Lock()
-	if _, ok := m.peers[peerConnID]; !ok {
-		m.mu.Unlock()
-		return
-	}
-	delete(m.peers, peerConnID)
-	m.mu.Unlock()
-
-	m.notify(peerConnID)
-}
-
-func (m *Manager) notify(peerConnID peerid.ConnID) {
-	select {
-	case <-m.done:
-	case m.OnActivityChan <- peerConnID:
-	}
-}

client/internal/lazyconn/activity/manager_test.go

@@ -1,162 +0,0 @@
-package activity
-
-import (
-	"net"
-	"net/netip"
-	"testing"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-	peerid "github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-type MocPeer struct {
-	PeerID string
-}
-
-func (m *MocPeer) ConnID() peerid.ConnID {
-	return peerid.ConnID(m)
-}
-
-type MocWGIface struct {
-}
-
-func (m MocWGIface) RemovePeer(string) error {
-	return nil
-}
-
-func (m MocWGIface) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error {
-	return nil
-
-}
-
-func TestManager_MonitorPeerActivity(t *testing.T) {
-	mocWgInterface := &MocWGIface{}
-
-	peer1 := &MocPeer{
-		PeerID: "examplePublicKey1",
-	}
-	mgr := NewManager(mocWgInterface)
-	defer mgr.Close()
-	peerCfg1 := lazyconn.PeerConfig{
-		PublicKey:  peer1.PeerID,
-		PeerConnID: peer1.ConnID(),
-		Log:        log.WithField("peer", "examplePublicKey1"),
-	}
-
-	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 {
-		t.Fatalf("failed to trigger activity: %v", err)
-	}
-
-	select {
-	case peerConnID := <-mgr.OnActivityChan:
-		if peerConnID != peerCfg1.PeerConnID {
-			t.Fatalf("unexpected peerConnID: %v", peerConnID)
-		}
-	case <-time.After(1 * time.Second):
-	}
-}
-
-func TestManager_RemovePeerActivity(t *testing.T) {
-	mocWgInterface := &MocWGIface{}
-
-	peer1 := &MocPeer{
-		PeerID: "examplePublicKey1",
-	}
-	mgr := NewManager(mocWgInterface)
-	defer mgr.Close()
-
-	peerCfg1 := lazyconn.PeerConfig{
-		PublicKey:  peer1.PeerID,
-		PeerConnID: peer1.ConnID(),
-		Log:        log.WithField("peer", "examplePublicKey1"),
-	}
-
-	if err := mgr.MonitorPeerActivity(peerCfg1); err != nil {
-		t.Fatalf("failed to monitor peer activity: %v", err)
-	}
-
-	addr := mgr.peers[peerCfg1.PeerConnID].conn.LocalAddr().String()
-
-	mgr.RemovePeer(peerCfg1.Log, peerCfg1.PeerConnID)
-
-	if err := trigger(addr); err != nil {
-		t.Fatalf("failed to trigger activity: %v", err)
-	}
-
-	select {
-	case <-mgr.OnActivityChan:
-		t.Fatal("should not have active activity")
-	case <-time.After(1 * time.Second):
-	}
-}
-
-func TestManager_MultiPeerActivity(t *testing.T) {
-	mocWgInterface := &MocWGIface{}
-
-	peer1 := &MocPeer{
-		PeerID: "examplePublicKey1",
-	}
-	mgr := NewManager(mocWgInterface)
-	defer mgr.Close()
-
-	peerCfg1 := lazyconn.PeerConfig{
-		PublicKey:  peer1.PeerID,
-		PeerConnID: peer1.ConnID(),
-		Log:        log.WithField("peer", "examplePublicKey1"),
-	}
-
-	peer2 := &MocPeer{}
-	peerCfg2 := lazyconn.PeerConfig{
-		PublicKey:  peer2.PeerID,
-		PeerConnID: peer2.ConnID(),
-		Log:        log.WithField("peer", "examplePublicKey2"),
-	}
-
-	if err := mgr.MonitorPeerActivity(peerCfg1); err != nil {
-		t.Fatalf("failed to monitor peer activity: %v", err)
-	}
-
-	if err := mgr.MonitorPeerActivity(peerCfg2); err != nil {
-		t.Fatalf("failed to monitor peer activity: %v", err)
-	}
-
-	if err := trigger(mgr.peers[peerCfg1.PeerConnID].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 {
-		t.Fatalf("failed to trigger activity: %v", err)
-	}
-
-	for i := 0; i < 2; i++ {
-		select {
-		case <-mgr.OnActivityChan:
-		case <-time.After(1 * time.Second):
-			t.Fatal("timed out waiting for activity")
-		}
-	}
-}
-
-func trigger(addr string) error {
-	// Create a connection to the destination UDP address and port
-	conn, err := net.Dial("udp", addr)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-
-	// Write the bytes to the UDP connection
-	_, err = conn.Write([]byte{0x01, 0x02, 0x03, 0x04, 0x05})
-	if err != nil {
-		return err
-	}
-	return nil
-}

client/internal/lazyconn/doc.go

@@ -1,32 +0,0 @@
-/*
-Package lazyconn provides mechanisms for managing lazy connections, which activate on demand to optimize resource usage and establish connections efficiently.
-
-## Overview
-
-The package includes a `Manager` component responsible for:
-- Managing lazy connections activated on-demand
-- Managing inactivity monitors for lazy connections (based on peer disconnection events)
-- Maintaining a list of excluded peers that should always have permanent connections
-- Handling remote peer connection initiatives based on peer signaling
-
-## Thread-Safe Operations
-
-The `Manager` ensures thread safety across multiple operations, categorized by caller:
-
-- **Engine (single goroutine)**:
-  - `AddPeer`: Adds a peer to the connection manager.
-  - `RemovePeer`: Removes a peer from the connection manager.
-  - `ActivatePeer`: Activates a lazy connection for a peer. This come from Signal client
-  - `ExcludePeer`: Marks peers for a permanent connection. Like router peers and other peers that should always have a connection.
-
-- **Connection Dispatcher (any peer routine)**:
-  - `onPeerConnected`: Suspend the inactivity monitor for an active peer connection.
-  - `onPeerDisconnected`: Starts the inactivity monitor for a disconnected peer.
-
-- **Activity Manager**:
-  - `onPeerActivity`: Run peer.Open(context).
-
-- **Inactivity Monitor**:
-  - `onPeerInactivityTimedOut`: Close peer connection and restart activity monitor.
-*/
-package lazyconn

client/internal/lazyconn/env.go

@@ -1,26 +0,0 @@
-package lazyconn
-
-import (
-	"os"
-	"strconv"
-
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	EnvEnableLazyConn      = "NB_ENABLE_EXPERIMENTAL_LAZY_CONN"
-	EnvInactivityThreshold = "NB_LAZY_CONN_INACTIVITY_THRESHOLD"
-)
-
-func IsLazyConnEnabledByEnv() bool {
-	val := os.Getenv(EnvEnableLazyConn)
-	if val == "" {
-		return false
-	}
-	enabled, err := strconv.ParseBool(val)
-	if err != nil {
-		log.Warnf("failed to parse %s: %v", EnvEnableLazyConn, err)
-		return false
-	}
-	return enabled
-}

client/internal/lazyconn/inactivity/inactivity.go

@@ -1,70 +0,0 @@
-package inactivity
-
-import (
-	"context"
-	"time"
-
-	peer "github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-const (
-	DefaultInactivityThreshold = 60 * time.Minute // idle after 1 hour inactivity
-	MinimumInactivityThreshold = 3 * time.Minute
-)
-
-type Monitor struct {
-	id                  peer.ConnID
-	timer               *time.Timer
-	cancel              context.CancelFunc
-	inactivityThreshold time.Duration
-}
-
-func NewInactivityMonitor(peerID peer.ConnID, threshold time.Duration) *Monitor {
-	i := &Monitor{
-		id:                  peerID,
-		timer:               time.NewTimer(0),
-		inactivityThreshold: threshold,
-	}
-	i.timer.Stop()
-	return i
-}
-
-func (i *Monitor) Start(ctx context.Context, timeoutChan chan peer.ConnID) {
-	i.timer.Reset(i.inactivityThreshold)
-	defer i.timer.Stop()
-
-	ctx, i.cancel = context.WithCancel(ctx)
-	defer func() {
-		defer i.cancel()
-		select {
-		case <-i.timer.C:
-		default:
-		}
-	}()
-
-	select {
-	case <-i.timer.C:
-		select {
-		case timeoutChan <- i.id:
-		case <-ctx.Done():
-			return
-		}
-	case <-ctx.Done():
-		return
-	}
-}
-
-func (i *Monitor) Stop() {
-	if i.cancel == nil {
-		return
-	}
-	i.cancel()
-}
-
-func (i *Monitor) PauseTimer() {
-	i.timer.Stop()
-}
-
-func (i *Monitor) ResetTimer() {
-	i.timer.Reset(i.inactivityThreshold)
-}

client/internal/lazyconn/inactivity/inactivity_test.go

@@ -1,156 +0,0 @@
-package inactivity
-
-import (
-	"context"
-	"testing"
-	"time"
-
-	peerid "github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-type MocPeer struct {
-}
-
-func (m *MocPeer) ConnID() peerid.ConnID {
-	return peerid.ConnID(m)
-}
-
-func TestInactivityMonitor(t *testing.T) {
-	tCtx, testTimeoutCancel := context.WithTimeout(context.Background(), time.Second*5)
-	defer testTimeoutCancel()
-
-	p := &MocPeer{}
-	im := NewInactivityMonitor(p.ConnID(), time.Second*2)
-
-	timeoutChan := make(chan peerid.ConnID)
-
-	exitChan := make(chan struct{})
-
-	go func() {
-		defer close(exitChan)
-		im.Start(tCtx, timeoutChan)
-	}()
-
-	select {
-	case <-timeoutChan:
-	case <-tCtx.Done():
-		t.Fatal("timeout")
-	}
-
-	select {
-	case <-exitChan:
-	case <-tCtx.Done():
-		t.Fatal("timeout")
-	}
-}
-
-func TestReuseInactivityMonitor(t *testing.T) {
-	p := &MocPeer{}
-	im := NewInactivityMonitor(p.ConnID(), time.Second*2)
-
-	timeoutChan := make(chan peerid.ConnID)
-
-	for i := 2; i > 0; i-- {
-		exitChan := make(chan struct{})
-
-		testTimeoutCtx, testTimeoutCancel := context.WithTimeout(context.Background(), time.Second*5)
-
-		go func() {
-			defer close(exitChan)
-			im.Start(testTimeoutCtx, timeoutChan)
-		}()
-
-		select {
-		case <-timeoutChan:
-		case <-testTimeoutCtx.Done():
-			t.Fatal("timeout")
-		}
-
-		select {
-		case <-exitChan:
-		case <-testTimeoutCtx.Done():
-			t.Fatal("timeout")
-		}
-		testTimeoutCancel()
-	}
-}
-
-func TestStopInactivityMonitor(t *testing.T) {
-	tCtx, testTimeoutCancel := context.WithTimeout(context.Background(), time.Second*5)
-	defer testTimeoutCancel()
-
-	p := &MocPeer{}
-	im := NewInactivityMonitor(p.ConnID(), DefaultInactivityThreshold)
-
-	timeoutChan := make(chan peerid.ConnID)
-
-	exitChan := make(chan struct{})
-
-	go func() {
-		defer close(exitChan)
-		im.Start(tCtx, timeoutChan)
-	}()
-
-	go func() {
-		time.Sleep(3 * time.Second)
-		im.Stop()
-	}()
-
-	select {
-	case <-timeoutChan:
-		t.Fatal("unexpected timeout")
-	case <-exitChan:
-	case <-tCtx.Done():
-		t.Fatal("timeout")
-	}
-}
-
-func TestPauseInactivityMonitor(t *testing.T) {
-	tCtx, testTimeoutCancel := context.WithTimeout(context.Background(), time.Second*10)
-	defer testTimeoutCancel()
-
-	p := &MocPeer{}
-	trashHold := time.Second * 3
-	im := NewInactivityMonitor(p.ConnID(), trashHold)
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	timeoutChan := make(chan peerid.ConnID)
-
-	exitChan := make(chan struct{})
-
-	go func() {
-		defer close(exitChan)
-		im.Start(ctx, timeoutChan)
-	}()
-
-	time.Sleep(1 * time.Second) // grant time to start the monitor
-	im.PauseTimer()
-
-	// check to do not receive timeout
-	thresholdCtx, thresholdCancel := context.WithTimeout(context.Background(), trashHold+time.Second)
-	defer thresholdCancel()
-	select {
-	case <-exitChan:
-		t.Fatal("unexpected exit")
-	case <-timeoutChan:
-		t.Fatal("unexpected timeout")
-	case <-thresholdCtx.Done():
-		// test ok
-	case <-tCtx.Done():
-		t.Fatal("test timed out")
-	}
-
-	// test reset timer
-	im.ResetTimer()
-
-	select {
-	case <-tCtx.Done():
-		t.Fatal("test timed out")
-	case <-exitChan:
-		t.Fatal("unexpected exit")
-	case <-timeoutChan:
-		// expected timeout
-	}
-}

client/internal/lazyconn/manager/manager.go

@@ -1,404 +0,0 @@
-package manager
-
-import (
-	"context"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/lazyconn"
-	"github.com/netbirdio/netbird/client/internal/lazyconn/activity"
-	"github.com/netbirdio/netbird/client/internal/lazyconn/inactivity"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
-	peerid "github.com/netbirdio/netbird/client/internal/peer/id"
-	"github.com/netbirdio/netbird/client/internal/peerstore"
-)
-
-const (
-	watcherActivity watcherType = iota
-	watcherInactivity
-)
-
-type watcherType int
-
-type managedPeer struct {
-	peerCfg         *lazyconn.PeerConfig
-	expectedWatcher watcherType
-}
-
-type Config struct {
-	InactivityThreshold *time.Duration
-}
-
-// Manager manages lazy connections
-// It is responsible for:
-// - Managing lazy connections activated on-demand
-// - Managing inactivity monitors for lazy connections (based on peer disconnection events)
-// - Maintaining a list of excluded peers that should always have permanent connections
-// - Handling connection establishment based on peer signaling
-type Manager struct {
-	peerStore           *peerstore.Store
-	connStateDispatcher *dispatcher.ConnectionDispatcher
-	inactivityThreshold time.Duration
-
-	connStateListener    *dispatcher.ConnectionListener
-	managedPeers         map[string]*lazyconn.PeerConfig
-	managedPeersByConnID map[peerid.ConnID]*managedPeer
-	excludes             map[string]lazyconn.PeerConfig
-	managedPeersMu       sync.Mutex
-
-	activityManager    *activity.Manager
-	inactivityMonitors map[peerid.ConnID]*inactivity.Monitor
-
-	cancel     context.CancelFunc
-	onInactive chan peerid.ConnID
-}
-
-func NewManager(config Config, peerStore *peerstore.Store, wgIface lazyconn.WGIface, connStateDispatcher *dispatcher.ConnectionDispatcher) *Manager {
-	log.Infof("setup lazy connection service")
-	m := &Manager{
-		peerStore:            peerStore,
-		connStateDispatcher:  connStateDispatcher,
-		inactivityThreshold:  inactivity.DefaultInactivityThreshold,
-		managedPeers:         make(map[string]*lazyconn.PeerConfig),
-		managedPeersByConnID: make(map[peerid.ConnID]*managedPeer),
-		excludes:             make(map[string]lazyconn.PeerConfig),
-		activityManager:      activity.NewManager(wgIface),
-		inactivityMonitors:   make(map[peerid.ConnID]*inactivity.Monitor),
-		onInactive:           make(chan peerid.ConnID),
-	}
-
-	if config.InactivityThreshold != nil {
-		if *config.InactivityThreshold >= inactivity.MinimumInactivityThreshold {
-			m.inactivityThreshold = *config.InactivityThreshold
-		} else {
-			log.Warnf("inactivity threshold is too low, using %v", m.inactivityThreshold)
-		}
-	}
-
-	m.connStateListener = &dispatcher.ConnectionListener{
-		OnConnected:    m.onPeerConnected,
-		OnDisconnected: m.onPeerDisconnected,
-	}
-
-	connStateDispatcher.AddListener(m.connStateListener)
-
-	return m
-}
-
-// Start starts the manager and listens for peer activity and inactivity events
-func (m *Manager) Start(ctx context.Context) {
-	defer m.close()
-
-	ctx, m.cancel = context.WithCancel(ctx)
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case peerConnID := <-m.activityManager.OnActivityChan:
-			m.onPeerActivity(ctx, peerConnID)
-		case peerConnID := <-m.onInactive:
-			m.onPeerInactivityTimedOut(peerConnID)
-		}
-	}
-}
-
-// ExcludePeer marks peers for a permanent connection
-// It removes peers from the managed list if they are added to the exclude list
-// Adds them back to the managed list and start the inactivity listener if they are removed from the exclude list. In
-// this case, we suppose that the connection status is connected or connecting.
-// If the peer is not exists yet in the managed list then the responsibility is the upper layer to call the AddPeer function
-func (m *Manager) ExcludePeer(ctx context.Context, peerConfigs []lazyconn.PeerConfig) []string {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	added := make([]string, 0)
-	excludes := make(map[string]lazyconn.PeerConfig, len(peerConfigs))
-
-	for _, peerCfg := range peerConfigs {
-		log.Infof("update excluded lazy connection list with peer: %s", peerCfg.PublicKey)
-		excludes[peerCfg.PublicKey] = peerCfg
-	}
-
-	// if a peer is newly added to the exclude list, remove from the managed peers list
-	for pubKey, peerCfg := range excludes {
-		if _, wasExcluded := m.excludes[pubKey]; wasExcluded {
-			continue
-		}
-
-		added = append(added, pubKey)
-		peerCfg.Log.Infof("peer newly added to lazy connection exclude list")
-		m.removePeer(pubKey)
-	}
-
-	// if a peer has been removed from exclude list then it should be added to the managed peers
-	for pubKey, peerCfg := range m.excludes {
-		if _, stillExcluded := excludes[pubKey]; stillExcluded {
-			continue
-		}
-
-		peerCfg.Log.Infof("peer removed from lazy connection exclude list")
-
-		if err := m.addActivePeer(ctx, peerCfg); err != nil {
-			log.Errorf("failed to add peer to lazy connection manager: %s", err)
-			continue
-		}
-	}
-
-	m.excludes = excludes
-	return added
-}
-
-func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	peerCfg.Log.Debugf("adding peer to lazy connection manager")
-
-	_, exists := m.excludes[peerCfg.PublicKey]
-	if exists {
-		return true, nil
-	}
-
-	if _, ok := m.managedPeers[peerCfg.PublicKey]; ok {
-		peerCfg.Log.Warnf("peer already managed")
-		return false, nil
-	}
-
-	if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
-		return false, err
-	}
-
-	im := inactivity.NewInactivityMonitor(peerCfg.PeerConnID, m.inactivityThreshold)
-	m.inactivityMonitors[peerCfg.PeerConnID] = im
-
-	m.managedPeers[peerCfg.PublicKey] = &peerCfg
-	m.managedPeersByConnID[peerCfg.PeerConnID] = &managedPeer{
-		peerCfg:         &peerCfg,
-		expectedWatcher: watcherActivity,
-	}
-	return false, nil
-}
-
-// AddActivePeers adds a list of peers to the lazy connection manager
-// suppose these peers was in connected or in connecting states
-func (m *Manager) AddActivePeers(ctx context.Context, peerCfg []lazyconn.PeerConfig) error {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	for _, cfg := range peerCfg {
-		if _, ok := m.managedPeers[cfg.PublicKey]; ok {
-			cfg.Log.Errorf("peer already managed")
-			continue
-		}
-
-		if err := m.addActivePeer(ctx, cfg); err != nil {
-			cfg.Log.Errorf("failed to add peer to lazy connection manager: %v", err)
-			return err
-		}
-	}
-	return nil
-}
-
-func (m *Manager) RemovePeer(peerID string) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	m.removePeer(peerID)
-}
-
-// ActivatePeer activates a peer connection when a signal message is received
-func (m *Manager) ActivatePeer(ctx context.Context, peerID string) (found bool) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	cfg, ok := m.managedPeers[peerID]
-	if !ok {
-		return false
-	}
-
-	mp, ok := m.managedPeersByConnID[cfg.PeerConnID]
-	if !ok {
-		return false
-	}
-
-	// signal messages coming continuously after success activation, with this avoid the multiple activation
-	if mp.expectedWatcher == watcherInactivity {
-		return false
-	}
-
-	mp.expectedWatcher = watcherInactivity
-
-	m.activityManager.RemovePeer(cfg.Log, cfg.PeerConnID)
-
-	im, ok := m.inactivityMonitors[cfg.PeerConnID]
-	if !ok {
-		cfg.Log.Errorf("inactivity monitor not found for peer")
-		return false
-	}
-
-	mp.peerCfg.Log.Infof("starting inactivity monitor")
-	go im.Start(ctx, m.onInactive)
-
-	return true
-}
-
-func (m *Manager) addActivePeer(ctx context.Context, peerCfg lazyconn.PeerConfig) error {
-	if _, ok := m.managedPeers[peerCfg.PublicKey]; ok {
-		peerCfg.Log.Warnf("peer already managed")
-		return nil
-	}
-
-	im := inactivity.NewInactivityMonitor(peerCfg.PeerConnID, m.inactivityThreshold)
-	m.inactivityMonitors[peerCfg.PeerConnID] = im
-
-	m.managedPeers[peerCfg.PublicKey] = &peerCfg
-	m.managedPeersByConnID[peerCfg.PeerConnID] = &managedPeer{
-		peerCfg:         &peerCfg,
-		expectedWatcher: watcherInactivity,
-	}
-
-	peerCfg.Log.Infof("starting inactivity monitor on peer that has been removed from exclude list")
-	go im.Start(ctx, m.onInactive)
-	return nil
-}
-
-func (m *Manager) removePeer(peerID string) {
-	cfg, ok := m.managedPeers[peerID]
-	if !ok {
-		return
-	}
-
-	cfg.Log.Infof("removing lazy peer")
-
-	if im, ok := m.inactivityMonitors[cfg.PeerConnID]; ok {
-		im.Stop()
-		delete(m.inactivityMonitors, cfg.PeerConnID)
-		cfg.Log.Debugf("inactivity monitor stopped")
-	}
-
-	m.activityManager.RemovePeer(cfg.Log, cfg.PeerConnID)
-	delete(m.managedPeers, peerID)
-	delete(m.managedPeersByConnID, cfg.PeerConnID)
-}
-
-func (m *Manager) close() {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	m.cancel()
-
-	m.connStateDispatcher.RemoveListener(m.connStateListener)
-	m.activityManager.Close()
-	for _, iw := range m.inactivityMonitors {
-		iw.Stop()
-	}
-	m.inactivityMonitors = make(map[peerid.ConnID]*inactivity.Monitor)
-	m.managedPeers = make(map[string]*lazyconn.PeerConfig)
-	m.managedPeersByConnID = make(map[peerid.ConnID]*managedPeer)
-	log.Infof("lazy connection manager closed")
-}
-
-func (m *Manager) onPeerActivity(ctx context.Context, peerConnID peerid.ConnID) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	mp, ok := m.managedPeersByConnID[peerConnID]
-	if !ok {
-		log.Errorf("peer not found by conn id: %v", peerConnID)
-		return
-	}
-
-	if mp.expectedWatcher != watcherActivity {
-		mp.peerCfg.Log.Warnf("ignore activity event")
-		return
-	}
-
-	mp.peerCfg.Log.Infof("detected peer activity")
-
-	mp.expectedWatcher = watcherInactivity
-
-	mp.peerCfg.Log.Infof("starting inactivity monitor")
-	go m.inactivityMonitors[peerConnID].Start(ctx, m.onInactive)
-
-	m.peerStore.PeerConnOpen(ctx, mp.peerCfg.PublicKey)
-}
-
-func (m *Manager) onPeerInactivityTimedOut(peerConnID peerid.ConnID) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	mp, ok := m.managedPeersByConnID[peerConnID]
-	if !ok {
-		log.Errorf("peer not found by id: %v", peerConnID)
-		return
-	}
-
-	if mp.expectedWatcher != watcherInactivity {
-		mp.peerCfg.Log.Warnf("ignore inactivity event")
-		return
-	}
-
-	mp.peerCfg.Log.Infof("connection timed out")
-
-	// this is blocking operation, potentially can be optimized
-	m.peerStore.PeerConnClose(mp.peerCfg.PublicKey)
-
-	mp.peerCfg.Log.Infof("start activity monitor")
-
-	mp.expectedWatcher = watcherActivity
-
-	// just in case free up
-	m.inactivityMonitors[peerConnID].PauseTimer()
-
-	if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
-		mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
-		return
-	}
-}
-
-func (m *Manager) onPeerConnected(peerConnID peerid.ConnID) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	mp, ok := m.managedPeersByConnID[peerConnID]
-	if !ok {
-		return
-	}
-
-	if mp.expectedWatcher != watcherInactivity {
-		return
-	}
-
-	iw, ok := m.inactivityMonitors[mp.peerCfg.PeerConnID]
-	if !ok {
-		mp.peerCfg.Log.Errorf("inactivity monitor not found for peer")
-		return
-	}
-
-	mp.peerCfg.Log.Infof("peer connected, pausing inactivity monitor while connection is not disconnected")
-	iw.PauseTimer()
-}
-
-func (m *Manager) onPeerDisconnected(peerConnID peerid.ConnID) {
-	m.managedPeersMu.Lock()
-	defer m.managedPeersMu.Unlock()
-
-	mp, ok := m.managedPeersByConnID[peerConnID]
-	if !ok {
-		return
-	}
-
-	if mp.expectedWatcher != watcherInactivity {
-		return
-	}
-
-	iw, ok := m.inactivityMonitors[mp.peerCfg.PeerConnID]
-	if !ok {
-		return
-	}
-
-	mp.peerCfg.Log.Infof("reset inactivity monitor timer")
-	iw.ResetTimer()
-}

client/internal/lazyconn/peercfg.go

@@ -1,16 +0,0 @@
-package lazyconn
-
-import (
-	"net/netip"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-type PeerConfig struct {
-	PublicKey  string
-	AllowedIPs []netip.Prefix
-	PeerConnID id.ConnID
-	Log        *log.Entry
-}

client/internal/lazyconn/support.go

@@ -1,41 +0,0 @@
-package lazyconn
-
-import (
-	"strings"
-
-	"github.com/hashicorp/go-version"
-)
-
-var (
-	minVersion = version.Must(version.NewVersion("0.45.0"))
-)
-
-func IsSupported(agentVersion string) bool {
-	if agentVersion == "development" {
-		return true
-	}
-
-	// filter out versions like this: a6c5960, a7d5c522, d47be154
-	if !strings.Contains(agentVersion, ".") {
-		return false
-	}
-
-	normalizedVersion := normalizeVersion(agentVersion)
-	inputVer, err := version.NewVersion(normalizedVersion)
-	if err != nil {
-		return false
-	}
-
-	return inputVer.GreaterThanOrEqual(minVersion)
-}
-
-func normalizeVersion(version string) string {
-	// Remove prefixes like 'v' or 'a'
-	if len(version) > 0 && (version[0] == 'v' || version[0] == 'a') {
-		version = version[1:]
-	}
-
-	// Remove any suffixes like '-dirty', '-dev', '-SNAPSHOT', etc.
-	parts := strings.Split(version, "-")
-	return parts[0]
-}

client/internal/lazyconn/support_test.go

@@ -1,31 +0,0 @@
-package lazyconn
-
-import "testing"
-
-func TestIsSupported(t *testing.T) {
-	tests := []struct {
-		version string
-		want    bool
-	}{
-		{"development", true},
-		{"0.45.0", true},
-		{"v0.45.0", true},
-		{"0.45.1", true},
-		{"0.45.1-SNAPSHOT-559e6731", true},
-		{"v0.45.1-dev", true},
-		{"a7d5c522", false},
-		{"0.9.6", false},
-		{"0.9.6-SNAPSHOT", false},
-		{"0.9.6-SNAPSHOT-2033650", false},
-		{"meta_wt_version", false},
-		{"v0.31.1-dev", false},
-		{"", false},
-	}
-	for _, tt := range tests {
-		t.Run(tt.version, func(t *testing.T) {
-			if got := IsSupported(tt.version); got != tt.want {
-				t.Errorf("IsSupported() = %v, want %v", got, tt.want)
-			}
-		})
-	}
-}

client/internal/lazyconn/wgiface.go

@@ -1,14 +0,0 @@
-package lazyconn
-
-import (
-	"net"
-	"net/netip"
-	"time"
-
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-)
-
-type WGIface interface {
-	RemovePeer(peerKey string) error
-	UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error
-}

client/internal/netflow/manager.go

@@ -123,14 +123,8 @@ func (m *Manager) disableFlow() error {
 
 	m.logger.Close()
 
-	if m.receiverClient == nil {
-		return nil
-	}
-
-	err := m.receiverClient.Close()
-	m.receiverClient = nil
-	if err != nil {
-		return fmt.Errorf("close: %w", err)
+	if m.receiverClient != nil {
+		return m.receiverClient.Close()
 	}
 
 	return nil

client/internal/networkmonitor/check_change_bsd.go

@@ -19,7 +19,7 @@ import (
 func checkChange(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
 	fd, err := unix.Socket(syscall.AF_ROUTE, syscall.SOCK_RAW, syscall.AF_UNSPEC)
 	if err != nil {
-		return fmt.Errorf("open routing socket: %v", err)
+		return fmt.Errorf("failed to open routing socket: %v", err)
 	}
 	defer func() {
 		err := unix.Close(fd)

client/internal/networkmonitor/check_change_windows.go

@@ -13,7 +13,7 @@ import (
 func checkChange(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
 	routeMonitor, err := systemops.NewRouteMonitor(ctx)
 	if err != nil {
-		return fmt.Errorf("create route monitor: %w", err)
+		return fmt.Errorf("failed to create route monitor: %w", err)
 	}
 	defer func() {
 		if err := routeMonitor.Stop(); err != nil {
@@ -38,46 +38,32 @@ func checkChange(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) er
 }
 
 func routeChanged(route systemops.RouteUpdate, nexthopv4, nexthopv6 systemops.Nexthop) bool {
-	if intf := route.NextHop.Intf; intf != nil && isSoftInterface(intf.Name) {
-		log.Debugf("Network monitor: ignoring default route change for next hop with soft interface %s", route.NextHop)
-		return false
-	}
-
-	// TODO: for the empty nexthop ip (on-link), determine the family differently
-	nexthop := nexthopv4
-	if route.NextHop.IP.Is6() {
-		nexthop = nexthopv6
+	intf := "<nil>"
+	if route.Interface != nil {
+		intf = route.Interface.Name
+		if isSoftInterface(intf) {
+			log.Debugf("Network monitor: ignoring default route change for soft interface %s", intf)
+			return false
+		}
 	}
 
 	switch route.Type {
-	case systemops.RouteModified, systemops.RouteAdded:
-		return handleRouteAddedOrModified(route, nexthop)
-	case systemops.RouteDeleted:
-		return handleRouteDeleted(route, nexthop)
-	}
-
-	return false
-}
-
-func handleRouteAddedOrModified(route systemops.RouteUpdate, nexthop systemops.Nexthop) bool {
-	// For added/modified routes, we care about different next hops
-	if !nexthop.Equal(route.NextHop) {
-		action := "changed"
-		if route.Type == systemops.RouteAdded {
-			action = "added"
+	case systemops.RouteModified:
+		// TODO: get routing table to figure out if our route is affected for modified routes
+		log.Infof("Network monitor: default route changed: via %s, interface %s", route.NextHop, intf)
+		return true
+	case systemops.RouteAdded:
+		if route.NextHop.Is4() && route.NextHop != nexthopv4.IP || route.NextHop.Is6() && route.NextHop != nexthopv6.IP {
+			log.Infof("Network monitor: default route added: via %s, interface %s", route.NextHop, intf)
+			return true
+		}
+	case systemops.RouteDeleted:
+		if nexthopv4.Intf != nil && route.NextHop == nexthopv4.IP || nexthopv6.Intf != nil && route.NextHop == nexthopv6.IP {
+			log.Infof("Network monitor: default route removed: via %s, interface %s", route.NextHop, intf)
+			return true
 		}
-		log.Infof("Network monitor: default route %s: via %s", action, route.NextHop)
-		return true
 	}
-	return false
-}
 
-func handleRouteDeleted(route systemops.RouteUpdate, nexthop systemops.Nexthop) bool {
-	// For deleted routes, we care about our tracked next hop being deleted
-	if nexthop.Equal(route.NextHop) {
-		log.Infof("Network monitor: default route removed: via %s", route.NextHop)
-		return true
-	}
 	return false
 }
 

client/internal/networkmonitor/check_change_windows_test.go

@@ -1,404 +0,0 @@
-package networkmonitor
-
-import (
-	"net"
-	"net/netip"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
-)
-
-func TestRouteChanged(t *testing.T) {
-	tests := []struct {
-		name      string
-		route     systemops.RouteUpdate
-		nexthopv4 systemops.Nexthop
-		nexthopv6 systemops.Nexthop
-		expected  bool
-	}{
-		{
-			name: "soft interface should be ignored",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{
-						Name: "ISATAP-Interface", // isSoftInterface checks name
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.2"),
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-			},
-			expected: false,
-		},
-		{
-			name: "modified route with different v4 nexthop IP should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.2"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-			},
-			expected: true,
-		},
-		{
-			name: "modified route with same v4 nexthop (IP and Intf Index) should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-			},
-			expected: false,
-		},
-		{
-			name: "added route with different v6 nexthop IP should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteAdded,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("2001:db8::2"),
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.1"),
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			expected: true,
-		},
-		{
-			name: "added route with same v6 nexthop (IP and Intf Index) should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteAdded,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("2001:db8::1"),
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.1"),
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			expected: false,
-		},
-		{
-			name: "deleted route matching tracked v4 nexthop (IP and Intf Index) should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-			},
-			expected: true,
-		},
-		{
-			name: "deleted route not matching tracked v4 nexthop (different IP) should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP: netip.MustParseAddr("192.168.1.3"), // Different IP
-					Intf: &net.Interface{
-						Index: 1, Name: "eth0",
-					},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP: netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{
-					Index: 1, Name: "eth0",
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP: netip.MustParseAddr("2001:db8::1"),
-			},
-			expected: false,
-		},
-		{
-			name: "modified v4 route with same IP, different Intf Index should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{Index: 2, Name: "eth1"}, // Different Intf Index
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP:   netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: true,
-		},
-		{
-			name: "modified v4 route with same IP, one Intf nil, other non-nil should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("192.168.1.1"),
-					Intf: nil, // Intf is nil
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP:   netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"}, // Tracked Intf is not nil
-			},
-			expected: true,
-		},
-		{
-			name: "added v4 route with same IP, different Intf Index should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteAdded,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{Index: 2, Name: "eth1"}, // Different Intf Index
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP:   netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: true,
-		},
-		{
-			name: "deleted v4 route with same IP, different Intf Index should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{ // This is the route being deleted
-					IP:   netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv4: systemops.Nexthop{ // This is our tracked nexthop
-				IP:   netip.MustParseAddr("192.168.1.1"),
-				Intf: &net.Interface{Index: 2, Name: "eth1"}, // Different Intf Index
-			},
-			expected: false, // Because nexthopv4.Equal(route.NextHop) will be false
-		},
-		{
-			name: "modified v6 route with different IP, same Intf Index should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("2001:db8::3"), // Different IP
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: true,
-		},
-		{
-			name: "modified v6 route with same IP, different Intf Index should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("2001:db8::1"),
-					Intf: &net.Interface{Index: 2, Name: "eth1"}, // Different Intf Index
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: true,
-		},
-		{
-			name: "modified v6 route with same IP, same Intf Index should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteModified,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("2001:db8::1"),
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: false,
-		},
-		{
-			name: "deleted v6 route matching tracked nexthop (IP and Intf Index) should return true",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("2001:db8::1"),
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: true,
-		},
-		{
-			name: "deleted v6 route not matching tracked nexthop (different IP) should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("2001:db8::3"), // Different IP
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv6: systemops.Nexthop{
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: false,
-		},
-		{
-			name: "deleted v6 route not matching tracked nexthop (same IP, different Intf Index) should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteDeleted,
-				Destination: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
-				NextHop: systemops.Nexthop{ // This is the route being deleted
-					IP:   netip.MustParseAddr("2001:db8::1"),
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv6: systemops.Nexthop{ // This is our tracked nexthop
-				IP:   netip.MustParseAddr("2001:db8::1"),
-				Intf: &net.Interface{Index: 2, Name: "eth1"}, // Different Intf Index
-			},
-			expected: false,
-		},
-		{
-			name: "unknown route type should return false",
-			route: systemops.RouteUpdate{
-				Type:        systemops.RouteUpdateType(99), // Unknown type
-				Destination: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
-				NextHop: systemops.Nexthop{
-					IP:   netip.MustParseAddr("192.168.1.1"),
-					Intf: &net.Interface{Index: 1, Name: "eth0"},
-				},
-			},
-			nexthopv4: systemops.Nexthop{
-				IP:   netip.MustParseAddr("192.168.1.2"), // Different from route.NextHop
-				Intf: &net.Interface{Index: 1, Name: "eth0"},
-			},
-			expected: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := routeChanged(tt.route, tt.nexthopv4, tt.nexthopv6)
-			assert.Equal(t, tt.expected, result)
-		})
-	}
-}
-
-func TestIsSoftInterface(t *testing.T) {
-	tests := []struct {
-		name     string
-		ifname   string
-		expected bool
-	}{
-		{
-			name:     "ISATAP interface should be detected",
-			ifname:   "ISATAP tunnel adapter",
-			expected: true,
-		},
-		{
-			name:     "lowercase soft interface should be detected",
-			ifname:   "isatap.{14A5CF17-CA72-43EC-B4EA-B4B093641B7D}",
-			expected: true,
-		},
-		{
-			name:     "Teredo interface should be detected",
-			ifname:   "Teredo Tunneling Pseudo-Interface",
-			expected: true,
-		},
-		{
-			name:     "regular interface should not be detected as soft",
-			ifname:   "eth0",
-			expected: false,
-		},
-		{
-			name:     "another regular interface should not be detected as soft",
-			ifname:   "wlan0",
-			expected: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := isSoftInterface(tt.ifname)
-			assert.Equal(t, tt.expected, result)
-		})
-	}
-}

client/internal/networkmonitor/monitor.go

@@ -118,12 +118,9 @@ 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 {
-			if !errors.Is(err, context.Canceled) {
-				log.Errorf("Network monitor: failed to check for changes: %v", err)
-			}
+			close(event)
 			return
 		}
 		// prevent blocking

client/internal/peer/conn.go

@@ -17,12 +17,8 @@ import (
 
 	"github.com/netbirdio/netbird/client/iface/configurer"
 	"github.com/netbirdio/netbird/client/iface/wgproxy"
-	"github.com/netbirdio/netbird/client/internal/peer/conntype"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
-	"github.com/netbirdio/netbird/client/internal/peer/id"
-	"github.com/netbirdio/netbird/client/internal/peer/worker"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	"github.com/netbirdio/netbird/route"
@@ -30,19 +26,31 @@ import (
 	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 )
 
+type ConnPriority int
+
+func (cp ConnPriority) String() string {
+	switch cp {
+	case connPriorityNone:
+		return "None"
+	case connPriorityRelay:
+		return "PriorityRelay"
+	case connPriorityICETurn:
+		return "PriorityICETurn"
+	case connPriorityICEP2P:
+		return "PriorityICEP2P"
+	default:
+		return fmt.Sprintf("ConnPriority(%d)", cp)
+	}
+}
+
 const (
 	defaultWgKeepAlive = 25 * time.Second
-)
 
-type ServiceDependencies struct {
-	StatusRecorder     *Status
-	Signaler           *Signaler
-	IFaceDiscover      stdnet.ExternalIFaceDiscover
-	RelayManager       *relayClient.Manager
-	SrWatcher          *guard.SRWatcher
-	Semaphore          *semaphoregroup.SemaphoreGroup
-	PeerConnDispatcher *dispatcher.ConnectionDispatcher
-}
+	connPriorityNone    ConnPriority = 0
+	connPriorityRelay   ConnPriority = 1
+	connPriorityICETurn ConnPriority = 2
+	connPriorityICEP2P  ConnPriority = 3
+)
 
 type WgConfig struct {
 	WgListenPort int
@@ -68,8 +76,6 @@ type ConnConfig struct {
 	// LocalKey is a public key of a local peer
 	LocalKey string
 
-	AgentVersion string
-
 	Timeout time.Duration
 
 	WgConfig WgConfig
@@ -83,23 +89,22 @@ type ConnConfig struct {
 }
 
 type Conn struct {
-	Log            *log.Entry
+	log            *log.Entry
 	mu             sync.Mutex
 	ctx            context.Context
 	ctxCancel      context.CancelFunc
 	config         ConnConfig
 	statusRecorder *Status
 	signaler       *Signaler
-	iFaceDiscover  stdnet.ExternalIFaceDiscover
 	relayManager   *relayClient.Manager
-	srWatcher      *guard.SRWatcher
+	handshaker     *Handshaker
 
 	onConnected    func(remoteWireGuardKey string, remoteRosenpassPubKey []byte, wireGuardIP string, remoteRosenpassAddr string)
 	onDisconnected func(remotePeer string)
 
-	statusRelay         *worker.AtomicWorkerStatus
-	statusICE           *worker.AtomicWorkerStatus
-	currentConnPriority conntype.ConnPriority
+	statusRelay         *AtomicConnStatus
+	statusICE           *AtomicConnStatus
+	currentConnPriority ConnPriority
 	opened              bool // this flag is used to prevent close in case of not opened connection
 
 	workerICE   *WorkerICE
@@ -115,12 +120,9 @@ type Conn struct {
 
 	wgProxyICE   wgproxy.Proxy
 	wgProxyRelay wgproxy.Proxy
-	handshaker   *Handshaker
 
-	guard              *guard.Guard
-	semaphore          *semaphoregroup.SemaphoreGroup
-	peerConnDispatcher *dispatcher.ConnectionDispatcher
-	wg                 sync.WaitGroup
+	guard     *guard.Guard
+	semaphore *semaphoregroup.SemaphoreGroup
 
 	// debug purpose
 	dumpState *stateDump
@@ -128,101 +130,91 @@ type Conn struct {
 
 // NewConn creates a new not opened Conn to the remote peer.
 // To establish a connection run Conn.Open
-func NewConn(config ConnConfig, services ServiceDependencies) (*Conn, error) {
+func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Status, signaler *Signaler, iFaceDiscover stdnet.ExternalIFaceDiscover, relayManager *relayClient.Manager, srWatcher *guard.SRWatcher, semaphore *semaphoregroup.SemaphoreGroup) (*Conn, error) {
 	if len(config.WgConfig.AllowedIps) == 0 {
 		return nil, fmt.Errorf("allowed IPs is empty")
 	}
 
+	ctx, ctxCancel := context.WithCancel(engineCtx)
 	connLog := log.WithField("peer", config.Key)
 
 	var conn = &Conn{
-		Log:                connLog,
-		config:             config,
-		statusRecorder:     services.StatusRecorder,
-		signaler:           services.Signaler,
-		iFaceDiscover:      services.IFaceDiscover,
-		relayManager:       services.RelayManager,
-		srWatcher:          services.SrWatcher,
-		semaphore:          services.Semaphore,
-		peerConnDispatcher: services.PeerConnDispatcher,
-		statusRelay:        worker.NewAtomicStatus(),
-		statusICE:          worker.NewAtomicStatus(),
-		dumpState:          newStateDump(config.Key, connLog, services.StatusRecorder),
+		log:            connLog,
+		ctx:            ctx,
+		ctxCancel:      ctxCancel,
+		config:         config,
+		statusRecorder: statusRecorder,
+		signaler:       signaler,
+		relayManager:   relayManager,
+		statusRelay:    NewAtomicConnStatus(),
+		statusICE:      NewAtomicConnStatus(),
+		semaphore:      semaphore,
+		dumpState:      newStateDump(config.Key, connLog, statusRecorder),
 	}
 
-	return conn, nil
-}
-
-// Open opens connection to the remote peer
-// It will try to establish a connection using ICE and in parallel with relay. The higher priority connection type will
-// be used.
-func (conn *Conn) Open(engineCtx context.Context) error {
-	conn.semaphore.Add(engineCtx)
-
-	conn.mu.Lock()
-	defer conn.mu.Unlock()
-
-	if conn.opened {
-		conn.semaphore.Done(engineCtx)
-		return nil
-	}
-
-	conn.ctx, conn.ctxCancel = context.WithCancel(engineCtx)
-
-	conn.workerRelay = NewWorkerRelay(conn.Log, isController(conn.config), conn.config, conn, conn.relayManager, conn.dumpState)
+	ctrl := isController(config)
+	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, 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)
+	workerICE, err := NewWorkerICE(ctx, connLog, config, conn, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	conn.workerICE = workerICE
 
-	conn.handshaker = NewHandshaker(conn.Log, conn.config, conn.signaler, conn.workerICE, conn.workerRelay)
+	conn.handshaker = NewHandshaker(ctx, connLog, config, signaler, conn.workerICE, conn.workerRelay)
 
 	conn.handshaker.AddOnNewOfferListener(conn.workerRelay.OnNewOffer)
 	if os.Getenv("NB_FORCE_RELAY") != "true" {
 		conn.handshaker.AddOnNewOfferListener(conn.workerICE.OnNewOffer)
 	}
 
-	conn.guard = guard.NewGuard(conn.Log, conn.isConnectedOnAllWay, conn.config.Timeout, conn.srWatcher)
+	conn.guard = guard.NewGuard(connLog, ctrl, conn.isConnectedOnAllWay, config.Timeout, srWatcher)
 
-	conn.wg.Add(1)
-	go func() {
-		defer conn.wg.Done()
-		conn.handshaker.Listen(conn.ctx)
-	}()
-	go conn.dumpState.Start(conn.ctx)
+	go conn.handshaker.Listen()
+
+	go conn.dumpState.Start(ctx)
+	return conn, nil
+}
+
+// Open opens connection to the remote peer
+// It will try to establish a connection using ICE and in parallel with relay. The higher priority connection type will
+// be used.
+func (conn *Conn) Open() {
+	conn.semaphore.Add(conn.ctx)
+	conn.log.Debugf("open connection to peer")
+
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+	conn.opened = true
 
 	peerState := State{
 		PubKey:           conn.config.Key,
+		IP:               conn.config.WgConfig.AllowedIps[0].Addr().String(),
 		ConnStatusUpdate: time.Now(),
-		ConnStatus:       StatusConnecting,
+		ConnStatus:       StatusDisconnected,
 		Mux:              new(sync.RWMutex),
 	}
-	if err := conn.statusRecorder.UpdatePeerState(peerState); err != nil {
-		conn.Log.Warnf("error while updating the state err: %v", err)
+	err := conn.statusRecorder.UpdatePeerState(peerState)
+	if err != nil {
+		conn.log.Warnf("error while updating the state err: %v", err)
 	}
 
-	conn.wg.Add(1)
-	go func() {
-		defer conn.wg.Done()
-		conn.waitInitialRandomSleepTime(conn.ctx)
-		conn.semaphore.Done(conn.ctx)
+	go conn.startHandshakeAndReconnect(conn.ctx)
+}
 
-		conn.dumpState.SendOffer()
-		if err := conn.handshaker.sendOffer(); err != nil {
-			conn.Log.Errorf("failed to send initial offer: %v", err)
-		}
+func (conn *Conn) startHandshakeAndReconnect(ctx context.Context) {
+	defer conn.semaphore.Done(conn.ctx)
+	conn.waitInitialRandomSleepTime(ctx)
 
-		conn.wg.Add(1)
-		go func() {
-			conn.guard.Start(conn.ctx, conn.onGuardEvent)
-			conn.wg.Done()
-		}()
-	}()
-	conn.opened = true
-	return nil
+	conn.dumpState.SendOffer()
+	err := conn.handshaker.sendOffer()
+	if err != nil {
+		conn.log.Errorf("failed to send initial offer: %v", err)
+	}
+
+	go conn.guard.Start(ctx)
+	go conn.listenGuardEvent(ctx)
 }
 
 // Close closes this peer Conn issuing a close event to the Conn closeCh
@@ -231,14 +223,14 @@ func (conn *Conn) Close() {
 	defer conn.wgWatcherWg.Wait()
 	defer conn.mu.Unlock()
 
+	conn.log.Infof("close peer connection")
+	conn.ctxCancel()
+
 	if !conn.opened {
-		conn.Log.Debugf("ignore close connection to peer")
+		conn.log.Debugf("ignore close connection to peer")
 		return
 	}
 
-	conn.Log.Infof("close peer connection")
-	conn.ctxCancel()
-
 	conn.workerRelay.DisableWgWatcher()
 	conn.workerRelay.CloseConn()
 	conn.workerICE.Close()
@@ -246,7 +238,7 @@ func (conn *Conn) Close() {
 	if conn.wgProxyRelay != nil {
 		err := conn.wgProxyRelay.CloseConn()
 		if err != nil {
-			conn.Log.Errorf("failed to close wg proxy for relay: %v", err)
+			conn.log.Errorf("failed to close wg proxy for relay: %v", err)
 		}
 		conn.wgProxyRelay = nil
 	}
@@ -254,13 +246,13 @@ func (conn *Conn) Close() {
 	if conn.wgProxyICE != nil {
 		err := conn.wgProxyICE.CloseConn()
 		if err != nil {
-			conn.Log.Errorf("failed to close wg proxy for ice: %v", err)
+			conn.log.Errorf("failed to close wg proxy for ice: %v", err)
 		}
 		conn.wgProxyICE = nil
 	}
 
 	if err := conn.removeWgPeer(); err != nil {
-		conn.Log.Errorf("failed to remove wg endpoint: %v", err)
+		conn.log.Errorf("failed to remove wg endpoint: %v", err)
 	}
 
 	conn.freeUpConnID()
@@ -270,16 +262,14 @@ func (conn *Conn) Close() {
 	}
 
 	conn.setStatusToDisconnected()
-	conn.opened = false
-	conn.wg.Wait()
-	conn.Log.Infof("peer connection closed")
+	conn.log.Infof("peer connection has been closed")
 }
 
 // OnRemoteAnswer handles an offer from the remote peer and returns true if the message was accepted, false otherwise
 // 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, conn.statusICE, conn.statusRelay)
 	return conn.handshaker.OnRemoteAnswer(answer)
 }
 
@@ -308,7 +298,7 @@ func (conn *Conn) SetOnDisconnected(handler func(remotePeer string)) {
 
 func (conn *Conn) OnRemoteOffer(offer OfferAnswer) bool {
 	conn.dumpState.RemoteOffer()
-	conn.Log.Infof("OnRemoteOffer, on status ICE: %s, status Relay: %s", conn.statusICE, conn.statusRelay)
+	conn.log.Infof("OnRemoteOffer, on status ICE: %s, status Relay: %s", conn.statusICE, conn.statusRelay)
 	return conn.handshaker.OnRemoteOffer(offer)
 }
 
@@ -317,24 +307,19 @@ func (conn *Conn) WgConfig() WgConfig {
 	return conn.config.WgConfig
 }
 
-// IsConnected unit tests only
-// refactor unit test to use status recorder use refactor status recorded to manage connection status in peer.Conn
-func (conn *Conn) IsConnected() bool {
+// Status returns current status of the Conn
+func (conn *Conn) Status() ConnStatus {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
-	return conn.currentConnPriority != conntype.None
+	return conn.evalStatus()
 }
 
 func (conn *Conn) GetKey() string {
 	return conn.config.Key
 }
 
-func (conn *Conn) ConnID() id.ConnID {
-	return id.ConnID(conn)
-}
-
 // configureConnection starts proxying traffic from/to local Wireguard and sets connection status to StatusConnected
-func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConnInfo ICEConnInfo) {
+func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEConnInfo) {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -342,21 +327,21 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 		return
 	}
 
-	if remoteConnNil(conn.Log, iceConnInfo.RemoteConn) {
-		conn.Log.Errorf("remote ICE connection is nil")
+	if remoteConnNil(conn.log, iceConnInfo.RemoteConn) {
+		conn.log.Errorf("remote ICE connection is nil")
 		return
 	}
 
 	// 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)
-		conn.statusICE.SetConnected()
+		conn.log.Infof("current connection priority (%s) is higher than the new one (%s), do not upgrade connection", conn.currentConnPriority, priority)
+		conn.statusICE.Set(StatusConnected)
 		conn.updateIceState(iceConnInfo)
 		return
 	}
 
-	conn.Log.Infof("set ICE to active connection")
+	conn.log.Infof("set ICE to active connection")
 	conn.dumpState.P2PConnected()
 
 	var (
@@ -368,7 +353,7 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 		conn.dumpState.NewLocalProxy()
 		wgProxy, err = conn.newProxy(iceConnInfo.RemoteConn)
 		if err != nil {
-			conn.Log.Errorf("failed to add turn net.Conn to local proxy: %v", err)
+			conn.log.Errorf("failed to add turn net.Conn to local proxy: %v", err)
 			return
 		}
 		ep = wgProxy.EndpointAddr()
@@ -384,7 +369,7 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 	}
 
 	if err := conn.runBeforeAddPeerHooks(ep.IP); err != nil {
-		conn.Log.Errorf("Before add peer hook failed: %v", err)
+		conn.log.Errorf("Before add peer hook failed: %v", err)
 	}
 
 	conn.workerRelay.DisableWgWatcher()
@@ -403,16 +388,10 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 		return
 	}
 	wgConfigWorkaround()
-
-	oldState := conn.currentConnPriority
 	conn.currentConnPriority = priority
-	conn.statusICE.SetConnected()
+	conn.statusICE.Set(StatusConnected)
 	conn.updateIceState(iceConnInfo)
 	conn.doOnConnected(iceConnInfo.RosenpassPubKey, iceConnInfo.RosenpassAddr)
-
-	if oldState == conntype.None {
-		conn.peerConnDispatcher.NotifyConnected(conn.ConnID())
-	}
 }
 
 func (conn *Conn) onICEStateDisconnected() {
@@ -423,22 +402,22 @@ func (conn *Conn) onICEStateDisconnected() {
 		return
 	}
 
-	conn.Log.Tracef("ICE connection state changed to disconnected")
+	conn.log.Tracef("ICE connection state changed to disconnected")
 
 	if conn.wgProxyICE != nil {
 		if err := conn.wgProxyICE.CloseConn(); err != nil {
-			conn.Log.Warnf("failed to close deprecated wg proxy conn: %v", err)
+			conn.log.Warnf("failed to close deprecated wg proxy conn: %v", err)
 		}
 	}
 
 	// switch back to relay connection
 	if conn.isReadyToUpgrade() {
-		conn.Log.Infof("ICE disconnected, set Relay to active connection")
+		conn.log.Infof("ICE disconnected, set Relay to active connection")
 		conn.dumpState.SwitchToRelay()
 		conn.wgProxyRelay.Work()
 
 		if err := conn.configureWGEndpoint(conn.wgProxyRelay.EndpointAddr(), conn.rosenpassRemoteKey); err != nil {
-			conn.Log.Errorf("failed to switch to relay conn: %v", err)
+			conn.log.Errorf("failed to switch to relay conn: %v", err)
 		}
 
 		conn.wgWatcherWg.Add(1)
@@ -446,18 +425,17 @@ func (conn *Conn) onICEStateDisconnected() {
 			defer conn.wgWatcherWg.Done()
 			conn.workerRelay.EnableWgWatcher(conn.ctx)
 		}()
-		conn.currentConnPriority = conntype.Relay
+		conn.currentConnPriority = connPriorityRelay
 	} else {
-		conn.Log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", conntype.None.String())
-		conn.currentConnPriority = conntype.None
-		conn.peerConnDispatcher.NotifyDisconnected(conn.ConnID())
+		conn.log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", connPriorityNone.String())
+		conn.currentConnPriority = connPriorityNone
 	}
 
-	changed := conn.statusICE.Get() != worker.StatusDisconnected
+	changed := conn.statusICE.Get() != StatusDisconnected
 	if changed {
 		conn.guard.SetICEConnDisconnected()
 	}
-	conn.statusICE.SetDisconnected()
+	conn.statusICE.Set(StatusDisconnected)
 
 	peerState := State{
 		PubKey:           conn.config.Key,
@@ -468,7 +446,7 @@ func (conn *Conn) onICEStateDisconnected() {
 
 	err := conn.statusRecorder.UpdatePeerICEStateToDisconnected(peerState)
 	if err != nil {
-		conn.Log.Warnf("unable to set peer's state to disconnected ice, got error: %v", err)
+		conn.log.Warnf("unable to set peer's state to disconnected ice, got error: %v", err)
 	}
 }
 
@@ -478,41 +456,41 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 
 	if conn.ctx.Err() != nil {
 		if err := rci.relayedConn.Close(); err != nil {
-			conn.Log.Warnf("failed to close unnecessary relayed connection: %v", err)
+			conn.log.Warnf("failed to close unnecessary relayed connection: %v", err)
 		}
 		return
 	}
 
 	conn.dumpState.RelayConnected()
-	conn.Log.Debugf("Relay connection has been established, setup the WireGuard")
+	conn.log.Debugf("Relay connection has been established, setup the WireGuard")
 
 	wgProxy, err := conn.newProxy(rci.relayedConn)
 	if err != nil {
-		conn.Log.Errorf("failed to add relayed net.Conn to local proxy: %v", err)
+		conn.log.Errorf("failed to add relayed net.Conn to local proxy: %v", err)
 		return
 	}
 	conn.dumpState.NewLocalProxy()
 
-	conn.Log.Infof("created new wgProxy for relay connection: %s", wgProxy.EndpointAddr().String())
+	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.Infof("do not switch to relay because current priority is: %s", conn.currentConnPriority.String())
 		conn.setRelayedProxy(wgProxy)
-		conn.statusRelay.SetConnected()
+		conn.statusRelay.Set(StatusConnected)
 		conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
 		return
 	}
 
 	if err := conn.runBeforeAddPeerHooks(wgProxy.EndpointAddr().IP); err != nil {
-		conn.Log.Errorf("Before add peer hook failed: %v", err)
+		conn.log.Errorf("Before add peer hook failed: %v", err)
 	}
 
 	wgProxy.Work()
 	if err := conn.configureWGEndpoint(wgProxy.EndpointAddr(), rci.rosenpassPubKey); err != nil {
 		if err := wgProxy.CloseConn(); err != nil {
-			conn.Log.Warnf("Failed to close relay connection: %v", err)
+			conn.log.Warnf("Failed to close relay connection: %v", err)
 		}
-		conn.Log.Errorf("Failed to update WireGuard peer configuration: %v", err)
+		conn.log.Errorf("Failed to update WireGuard peer configuration: %v", err)
 		return
 	}
 
@@ -524,13 +502,12 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 
 	wgConfigWorkaround()
 	conn.rosenpassRemoteKey = rci.rosenpassPubKey
-	conn.currentConnPriority = conntype.Relay
-	conn.statusRelay.SetConnected()
+	conn.currentConnPriority = connPriorityRelay
+	conn.statusRelay.Set(StatusConnected)
 	conn.setRelayedProxy(wgProxy)
 	conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
-	conn.Log.Infof("start to communicate with peer via relay")
+	conn.log.Infof("start to communicate with peer via relay")
 	conn.doOnConnected(rci.rosenpassPubKey, rci.rosenpassAddr)
-	conn.peerConnDispatcher.NotifyConnected(conn.ConnID())
 }
 
 func (conn *Conn) onRelayDisconnected() {
@@ -541,15 +518,14 @@ func (conn *Conn) onRelayDisconnected() {
 		return
 	}
 
-	conn.Log.Debugf("relay connection is disconnected")
+	conn.log.Infof("relay connection is disconnected")
 
-	if conn.currentConnPriority == conntype.Relay {
-		conn.Log.Debugf("clean up WireGuard config")
+	if conn.currentConnPriority == connPriorityRelay {
+		conn.log.Infof("clean up WireGuard config")
 		if err := conn.removeWgPeer(); err != nil {
-			conn.Log.Errorf("failed to remove wg endpoint: %v", err)
+			conn.log.Errorf("failed to remove wg endpoint: %v", err)
 		}
-		conn.currentConnPriority = conntype.None
-		conn.peerConnDispatcher.NotifyDisconnected(conn.ConnID())
+		conn.currentConnPriority = connPriorityNone
 	}
 
 	if conn.wgProxyRelay != nil {
@@ -557,11 +533,11 @@ func (conn *Conn) onRelayDisconnected() {
 		conn.wgProxyRelay = nil
 	}
 
-	changed := conn.statusRelay.Get() != worker.StatusDisconnected
+	changed := conn.statusRelay.Get() != StatusDisconnected
 	if changed {
 		conn.guard.SetRelayedConnDisconnected()
 	}
-	conn.statusRelay.SetDisconnected()
+	conn.statusRelay.Set(StatusDisconnected)
 
 	peerState := State{
 		PubKey:           conn.config.Key,
@@ -570,15 +546,22 @@ func (conn *Conn) onRelayDisconnected() {
 		ConnStatusUpdate: time.Now(),
 	}
 	if err := conn.statusRecorder.UpdatePeerRelayedStateToDisconnected(peerState); err != nil {
-		conn.Log.Warnf("unable to save peer's state to Relay disconnected, got error: %v", err)
+		conn.log.Warnf("unable to save peer's state to Relay disconnected, got error: %v", err)
 	}
 }
 
-func (conn *Conn) onGuardEvent() {
-	conn.Log.Debugf("send offer to peer")
-	conn.dumpState.SendOffer()
-	if err := conn.handshaker.SendOffer(); err != nil {
-		conn.Log.Errorf("failed to send offer: %v", err)
+func (conn *Conn) listenGuardEvent(ctx context.Context) {
+	for {
+		select {
+		case <-conn.guard.Reconnect:
+			conn.log.Infof("send offer to peer")
+			conn.dumpState.SendOffer()
+			if err := conn.handshaker.SendOffer(); err != nil {
+				conn.log.Errorf("failed to send offer: %v", err)
+			}
+		case <-ctx.Done():
+			return
+		}
 	}
 }
 
@@ -605,7 +588,7 @@ func (conn *Conn) updateRelayStatus(relayServerAddr string, rosenpassPubKey []by
 
 	err := conn.statusRecorder.UpdatePeerRelayedState(peerState)
 	if err != nil {
-		conn.Log.Warnf("unable to save peer's Relay state, got error: %v", err)
+		conn.log.Warnf("unable to save peer's Relay state, got error: %v", err)
 	}
 }
 
@@ -624,18 +607,17 @@ func (conn *Conn) updateIceState(iceConnInfo ICEConnInfo) {
 
 	err := conn.statusRecorder.UpdatePeerICEState(peerState)
 	if err != nil {
-		conn.Log.Warnf("unable to save peer's ICE state, got error: %v", err)
+		conn.log.Warnf("unable to save peer's ICE state, got error: %v", err)
 	}
 }
 
 func (conn *Conn) setStatusToDisconnected() {
-	conn.statusRelay.SetDisconnected()
-	conn.statusICE.SetDisconnected()
-	conn.currentConnPriority = conntype.None
+	conn.statusRelay.Set(StatusDisconnected)
+	conn.statusICE.Set(StatusDisconnected)
 
 	peerState := State{
 		PubKey:           conn.config.Key,
-		ConnStatus:       StatusIdle,
+		ConnStatus:       StatusDisconnected,
 		ConnStatusUpdate: time.Now(),
 		Mux:              new(sync.RWMutex),
 	}
@@ -643,10 +625,10 @@ func (conn *Conn) setStatusToDisconnected() {
 	if err != nil {
 		// pretty common error because by that time Engine can already remove the peer and status won't be available.
 		// todo rethink status updates
-		conn.Log.Debugf("error while updating peer's state, err: %v", err)
+		conn.log.Debugf("error while updating peer's state, err: %v", err)
 	}
 	if err := conn.statusRecorder.UpdateWireGuardPeerState(conn.config.Key, configurer.WGStats{}); err != nil {
-		conn.Log.Debugf("failed to reset wireguard stats for peer: %s", err)
+		conn.log.Debugf("failed to reset wireguard stats for peer: %s", err)
 	}
 }
 
@@ -674,24 +656,32 @@ func (conn *Conn) waitInitialRandomSleepTime(ctx context.Context) {
 }
 
 func (conn *Conn) isRelayed() bool {
-	switch conn.currentConnPriority {
-	case conntype.Relay, conntype.ICETurn:
-		return true
-	default:
+	if conn.statusRelay.Get() == StatusDisconnected && (conn.statusICE.Get() == StatusDisconnected || conn.statusICE.Get() == StatusConnecting) {
 		return false
 	}
+
+	if conn.currentConnPriority == connPriorityICEP2P {
+		return false
+	}
+
+	return true
 }
 
 func (conn *Conn) evalStatus() ConnStatus {
-	if conn.statusRelay.Get() == worker.StatusConnected || conn.statusICE.Get() == worker.StatusConnected {
+	if conn.statusRelay.Get() == StatusConnected || conn.statusICE.Get() == StatusConnected {
 		return StatusConnected
 	}
 
-	return StatusConnecting
+	if conn.statusRelay.Get() == StatusConnecting || conn.statusICE.Get() == StatusConnecting {
+		return StatusConnecting
+	}
+
+	return StatusDisconnected
 }
 
 func (conn *Conn) isConnectedOnAllWay() (connected bool) {
-	// would be better to protect this with a mutex, but it could cause deadlock with Close function
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
 
 	defer func() {
 		if !connected {
@@ -699,12 +689,12 @@ func (conn *Conn) isConnectedOnAllWay() (connected bool) {
 		}
 	}()
 
-	if conn.statusICE.Get() == worker.StatusDisconnected {
+	if conn.statusICE.Get() == StatusDisconnected {
 		return false
 	}
 
 	if conn.workerRelay.IsRelayConnectionSupportedWithPeer() {
-		if conn.statusRelay.Get() == worker.StatusDisconnected {
+		if conn.statusRelay.Get() != StatusConnected {
 			return false
 		}
 	}
@@ -726,7 +716,7 @@ func (conn *Conn) freeUpConnID() {
 	if conn.connIDRelay != "" {
 		for _, hook := range conn.afterRemovePeerHooks {
 			if err := hook(conn.connIDRelay); err != nil {
-				conn.Log.Errorf("After remove peer hook failed: %v", err)
+				conn.log.Errorf("After remove peer hook failed: %v", err)
 			}
 		}
 		conn.connIDRelay = ""
@@ -735,7 +725,7 @@ func (conn *Conn) freeUpConnID() {
 	if conn.connIDICE != "" {
 		for _, hook := range conn.afterRemovePeerHooks {
 			if err := hook(conn.connIDICE); err != nil {
-				conn.Log.Errorf("After remove peer hook failed: %v", err)
+				conn.log.Errorf("After remove peer hook failed: %v", err)
 			}
 		}
 		conn.connIDICE = ""
@@ -743,7 +733,7 @@ func (conn *Conn) freeUpConnID() {
 }
 
 func (conn *Conn) newProxy(remoteConn net.Conn) (wgproxy.Proxy, error) {
-	conn.Log.Debugf("setup proxied WireGuard connection")
+	conn.log.Debugf("setup proxied WireGuard connection")
 	udpAddr := &net.UDPAddr{
 		IP:   conn.config.WgConfig.AllowedIps[0].Addr().AsSlice(),
 		Port: conn.config.WgConfig.WgListenPort,
@@ -751,18 +741,18 @@ func (conn *Conn) newProxy(remoteConn net.Conn) (wgproxy.Proxy, error) {
 
 	wgProxy := conn.config.WgConfig.WgInterface.GetProxy()
 	if err := wgProxy.AddTurnConn(conn.ctx, udpAddr, remoteConn); err != nil {
-		conn.Log.Errorf("failed to add turn net.Conn to local proxy: %v", err)
+		conn.log.Errorf("failed to add turn net.Conn to local proxy: %v", err)
 		return nil, err
 	}
 	return wgProxy, nil
 }
 
 func (conn *Conn) isReadyToUpgrade() bool {
-	return conn.wgProxyRelay != nil && conn.currentConnPriority != conntype.Relay
+	return conn.wgProxyRelay != nil && conn.currentConnPriority != connPriorityRelay
 }
 
 func (conn *Conn) isICEActive() bool {
-	return (conn.currentConnPriority == conntype.ICEP2P || conn.currentConnPriority == conntype.ICETurn) && conn.statusICE.Get() == worker.StatusConnected
+	return (conn.currentConnPriority == connPriorityICEP2P || conn.currentConnPriority == connPriorityICETurn) && conn.statusICE.Get() == StatusConnected
 }
 
 func (conn *Conn) removeWgPeer() error {
@@ -770,10 +760,10 @@ func (conn *Conn) removeWgPeer() error {
 }
 
 func (conn *Conn) handleConfigurationFailure(err error, wgProxy wgproxy.Proxy) {
-	conn.Log.Warnf("Failed to update wg peer configuration: %v", err)
+	conn.log.Warnf("Failed to update wg peer configuration: %v", err)
 	if wgProxy != nil {
 		if ierr := wgProxy.CloseConn(); ierr != nil {
-			conn.Log.Warnf("Failed to close wg proxy: %v", ierr)
+			conn.log.Warnf("Failed to close wg proxy: %v", ierr)
 		}
 	}
 	if conn.wgProxyRelay != nil {
@@ -783,16 +773,16 @@ func (conn *Conn) handleConfigurationFailure(err error, wgProxy wgproxy.Proxy) {
 
 func (conn *Conn) logTraceConnState() {
 	if conn.workerRelay.IsRelayConnectionSupportedWithPeer() {
-		conn.Log.Tracef("connectivity guard check, relay state: %s, ice state: %s", conn.statusRelay, conn.statusICE)
+		conn.log.Tracef("connectivity guard check, relay state: %s, ice state: %s", conn.statusRelay, conn.statusICE)
 	} else {
-		conn.Log.Tracef("connectivity guard check, ice state: %s", conn.statusICE)
+		conn.log.Tracef("connectivity guard check, ice state: %s", conn.statusICE)
 	}
 }
 
 func (conn *Conn) setRelayedProxy(proxy wgproxy.Proxy) {
 	if conn.wgProxyRelay != nil {
 		if err := conn.wgProxyRelay.CloseConn(); err != nil {
-			conn.Log.Warnf("failed to close deprecated wg proxy conn: %v", err)
+			conn.log.Warnf("failed to close deprecated wg proxy conn: %v", err)
 		}
 	}
 	conn.wgProxyRelay = proxy
@@ -803,10 +793,6 @@ func (conn *Conn) AllowedIP() netip.Addr {
 	return conn.config.WgConfig.AllowedIps[0].Addr()
 }
 
-func (conn *Conn) AgentVersionString() string {
-	return conn.config.AgentVersion
-}
-
 func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key {
 	if conn.config.RosenpassConfig.PubKey == nil {
 		return conn.config.WgConfig.PreSharedKey
@@ -818,7 +804,7 @@ func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key {
 
 	determKey, err := conn.rosenpassDetermKey()
 	if err != nil {
-		conn.Log.Errorf("failed to generate Rosenpass initial key: %v", err)
+		conn.log.Errorf("failed to generate Rosenpass initial key: %v", err)
 		return conn.config.WgConfig.PreSharedKey
 	}
 

client/internal/peer/conn_status.go

@@ -1,29 +1,58 @@
 package peer
 
 import (
+	"sync/atomic"
+
 	log "github.com/sirupsen/logrus"
 )
 
 const (
-	// StatusIdle indicate the peer is in disconnected state
-	StatusIdle ConnStatus = iota
+	// StatusConnected indicate the peer is in connected state
+	StatusConnected ConnStatus = iota
 	// StatusConnecting indicate the peer is in connecting state
 	StatusConnecting
-	// StatusConnected indicate the peer is in connected state
-	StatusConnected
+	// StatusDisconnected indicate the peer is in disconnected state
+	StatusDisconnected
 )
 
 // ConnStatus describe the status of a peer's connection
 type ConnStatus int32
 
+// AtomicConnStatus is a thread-safe wrapper for ConnStatus
+type AtomicConnStatus struct {
+	status atomic.Int32
+}
+
+// NewAtomicConnStatus creates a new AtomicConnStatus with the given initial status
+func NewAtomicConnStatus() *AtomicConnStatus {
+	acs := &AtomicConnStatus{}
+	acs.Set(StatusDisconnected)
+	return acs
+}
+
+// Get returns the current connection status
+func (acs *AtomicConnStatus) Get() ConnStatus {
+	return ConnStatus(acs.status.Load())
+}
+
+// Set updates the connection status
+func (acs *AtomicConnStatus) Set(status ConnStatus) {
+	acs.status.Store(int32(status))
+}
+
+// String returns the string representation of the current status
+func (acs *AtomicConnStatus) String() string {
+	return acs.Get().String()
+}
+
 func (s ConnStatus) String() string {
 	switch s {
 	case StatusConnecting:
 		return "Connecting"
 	case StatusConnected:
 		return "Connected"
-	case StatusIdle:
-		return "Idle"
+	case StatusDisconnected:
+		return "Disconnected"
 	default:
 		log.Errorf("unknown status: %d", s)
 		return "INVALID_PEER_CONNECTION_STATUS"

client/internal/peer/conn_status_test.go

@@ -14,7 +14,7 @@ func TestConnStatus_String(t *testing.T) {
 		want   string
 	}{
 		{"StatusConnected", StatusConnected, "Connected"},
-		{"StatusIdle", StatusIdle, "Idle"},
+		{"StatusDisconnected", StatusDisconnected, "Disconnected"},
 		{"StatusConnecting", StatusConnecting, "Connecting"},
 	}
 
@@ -24,4 +24,5 @@ func TestConnStatus_String(t *testing.T) {
 			assert.Equal(t, got, table.want, "they should be equal")
 		})
 	}
+
 }

client/internal/peer/conn_test.go

@@ -1,6 +1,7 @@
 package peer
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"sync"
@@ -10,7 +11,6 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/netbirdio/netbird/client/iface"
-	"github.com/netbirdio/netbird/client/internal/peer/dispatcher"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	"github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
@@ -18,8 +18,6 @@ import (
 	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 )
 
-var testDispatcher = dispatcher.NewConnectionDispatcher()
-
 var connConf = ConnConfig{
 	Key:         "LLHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
 	LocalKey:    "RRHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
@@ -50,13 +48,7 @@ func TestNewConn_interfaceFilter(t *testing.T) {
 
 func TestConn_GetKey(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-
-	sd := ServiceDependencies{
-		SrWatcher:          swWatcher,
-		Semaphore:          semaphoregroup.NewSemaphoreGroup(1),
-		PeerConnDispatcher: testDispatcher,
-	}
-	conn, err := NewConn(connConf, sd)
+	conn, err := NewConn(context.Background(), connConf, nil, nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -68,13 +60,7 @@ func TestConn_GetKey(t *testing.T) {
 
 func TestConn_OnRemoteOffer(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-	sd := ServiceDependencies{
-		StatusRecorder:     NewRecorder("https://mgm"),
-		SrWatcher:          swWatcher,
-		Semaphore:          semaphoregroup.NewSemaphoreGroup(1),
-		PeerConnDispatcher: testDispatcher,
-	}
-	conn, err := NewConn(connConf, sd)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -108,13 +94,7 @@ func TestConn_OnRemoteOffer(t *testing.T) {
 
 func TestConn_OnRemoteAnswer(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-	sd := ServiceDependencies{
-		StatusRecorder:     NewRecorder("https://mgm"),
-		SrWatcher:          swWatcher,
-		Semaphore:          semaphoregroup.NewSemaphoreGroup(1),
-		PeerConnDispatcher: testDispatcher,
-	}
-	conn, err := NewConn(connConf, sd)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -145,6 +125,43 @@ func TestConn_OnRemoteAnswer(t *testing.T) {
 
 	wg.Wait()
 }
+func TestConn_Status(t *testing.T) {
+	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
+	if err != nil {
+		return
+	}
+
+	tables := []struct {
+		name        string
+		statusIce   ConnStatus
+		statusRelay ConnStatus
+		want        ConnStatus
+	}{
+		{"StatusConnected", StatusConnected, StatusConnected, StatusConnected},
+		{"StatusDisconnected", StatusDisconnected, StatusDisconnected, StatusDisconnected},
+		{"StatusConnecting", StatusConnecting, StatusConnecting, StatusConnecting},
+		{"StatusConnectingIce", StatusConnecting, StatusDisconnected, StatusConnecting},
+		{"StatusConnectingIceAlternative", StatusConnecting, StatusConnected, StatusConnected},
+		{"StatusConnectingRelay", StatusDisconnected, StatusConnecting, StatusConnecting},
+		{"StatusConnectingRelayAlternative", StatusConnected, StatusConnecting, StatusConnected},
+	}
+
+	for _, table := range tables {
+		t.Run(table.name, func(t *testing.T) {
+			si := NewAtomicConnStatus()
+			si.Set(table.statusIce)
+			conn.statusICE = si
+
+			sr := NewAtomicConnStatus()
+			sr.Set(table.statusRelay)
+			conn.statusRelay = sr
+
+			got := conn.Status()
+			assert.Equal(t, got, table.want, "they should be equal")
+		})
+	}
+}
 
 func TestConn_presharedKey(t *testing.T) {
 	conn1 := Conn{

client/internal/peer/conntype/priority.go

@@ -1,29 +0,0 @@
-package conntype
-
-import (
-	"fmt"
-)
-
-const (
-	None    ConnPriority = 0
-	Relay   ConnPriority = 1
-	ICETurn ConnPriority = 2
-	ICEP2P  ConnPriority = 3
-)
-
-type ConnPriority int
-
-func (cp ConnPriority) String() string {
-	switch cp {
-	case None:
-		return "None"
-	case Relay:
-		return "PriorityRelay"
-	case ICETurn:
-		return "PriorityICETurn"
-	case ICEP2P:
-		return "PriorityICEP2P"
-	default:
-		return fmt.Sprintf("ConnPriority(%d)", cp)
-	}
-}

client/internal/peer/dispatcher/dispatcher.go

@@ -1,52 +0,0 @@
-package dispatcher
-
-import (
-	"sync"
-
-	"github.com/netbirdio/netbird/client/internal/peer/id"
-)
-
-type ConnectionListener struct {
-	OnConnected    func(peerID id.ConnID)
-	OnDisconnected func(peerID id.ConnID)
-}
-
-type ConnectionDispatcher struct {
-	listeners map[*ConnectionListener]struct{}
-	mu        sync.Mutex
-}
-
-func NewConnectionDispatcher() *ConnectionDispatcher {
-	return &ConnectionDispatcher{
-		listeners: make(map[*ConnectionListener]struct{}),
-	}
-}
-
-func (e *ConnectionDispatcher) AddListener(listener *ConnectionListener) {
-	e.mu.Lock()
-	defer e.mu.Unlock()
-	e.listeners[listener] = struct{}{}
-}
-
-func (e *ConnectionDispatcher) RemoveListener(listener *ConnectionListener) {
-	e.mu.Lock()
-	defer e.mu.Unlock()
-
-	delete(e.listeners, listener)
-}
-
-func (e *ConnectionDispatcher) NotifyConnected(peerConnID id.ConnID) {
-	e.mu.Lock()
-	defer e.mu.Unlock()
-	for listener := range e.listeners {
-		listener.OnConnected(peerConnID)
-	}
-}
-
-func (e *ConnectionDispatcher) NotifyDisconnected(peerConnID id.ConnID) {
-	e.mu.Lock()
-	defer e.mu.Unlock()
-	for listener := range e.listeners {
-		listener.OnDisconnected(peerConnID)
-	}
-}

client/internal/peer/guard/guard.go

@@ -8,6 +8,10 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const (
+	reconnectMaxElapsedTime = 30 * time.Minute
+)
+
 type isConnectedFunc func() bool
 
 // Guard is responsible for the reconnection logic.
@@ -21,6 +25,7 @@ type isConnectedFunc func() bool
 type Guard struct {
 	Reconnect               chan struct{}
 	log                     *log.Entry
+	isController            bool
 	isConnectedOnAllWay     isConnectedFunc
 	timeout                 time.Duration
 	srWatcher               *SRWatcher
@@ -28,10 +33,11 @@ type Guard struct {
 	iCEConnDisconnected     chan struct{}
 }
 
-func NewGuard(log *log.Entry, isConnectedFn isConnectedFunc, timeout time.Duration, srWatcher *SRWatcher) *Guard {
+func NewGuard(log *log.Entry, isController bool, isConnectedFn isConnectedFunc, timeout time.Duration, srWatcher *SRWatcher) *Guard {
 	return &Guard{
 		Reconnect:               make(chan struct{}, 1),
 		log:                     log,
+		isController:            isController,
 		isConnectedOnAllWay:     isConnectedFn,
 		timeout:                 timeout,
 		srWatcher:               srWatcher,
@@ -40,8 +46,12 @@ func NewGuard(log *log.Entry, isConnectedFn isConnectedFunc, timeout time.Durati
 	}
 }
 
-func (g *Guard) Start(ctx context.Context, eventCallback func()) {
-	g.reconnectLoopWithRetry(ctx, eventCallback)
+func (g *Guard) Start(ctx context.Context) {
+	if g.isController {
+		g.reconnectLoopWithRetry(ctx)
+	} else {
+		g.listenForDisconnectEvents(ctx)
+	}
 }
 
 func (g *Guard) SetRelayedConnDisconnected() {
@@ -58,9 +68,9 @@ func (g *Guard) SetICEConnDisconnected() {
 	}
 }
 
-// reconnectLoopWithRetry periodically check the connection status.
+// reconnectLoopWithRetry periodically check (max 30 min) the connection status.
 // Try to send offer while the P2P is not established or while the Relay is not connected if is it supported
-func (g *Guard) reconnectLoopWithRetry(ctx context.Context, callback func()) {
+func (g *Guard) reconnectLoopWithRetry(ctx context.Context) {
 	waitForInitialConnectionTry(ctx)
 
 	srReconnectedChan := g.srWatcher.NewListener()
@@ -83,7 +93,7 @@ func (g *Guard) reconnectLoopWithRetry(ctx context.Context, callback func()) {
 			}
 
 			if !g.isConnectedOnAllWay() {
-				callback()
+				g.triggerOfferSending()
 			}
 
 		case <-g.relayedConnDisconnected:
@@ -111,12 +121,39 @@ func (g *Guard) reconnectLoopWithRetry(ctx context.Context, callback func()) {
 	}
 }
 
+// listenForDisconnectEvents is used when the peer is not a controller and it should reconnect to the peer
+// when the connection is lost. It will try to establish a connection only once time if before the connection was established
+// It track separately the ice and relay connection status. Just because a lower priority connection reestablished it does not
+// mean that to switch to it. We always force to use the higher priority connection.
+func (g *Guard) listenForDisconnectEvents(ctx context.Context) {
+	srReconnectedChan := g.srWatcher.NewListener()
+	defer g.srWatcher.RemoveListener(srReconnectedChan)
+
+	g.log.Infof("start listen for reconnect events...")
+	for {
+		select {
+		case <-g.relayedConnDisconnected:
+			g.log.Debugf("Relay connection changed, triggering reconnect")
+			g.triggerOfferSending()
+		case <-g.iCEConnDisconnected:
+			g.log.Debugf("ICE state changed, try to send new offer")
+			g.triggerOfferSending()
+		case <-srReconnectedChan:
+			g.triggerOfferSending()
+		case <-ctx.Done():
+			g.log.Debugf("context is done, stop reconnect loop")
+			return
+		}
+	}
+}
+
 func (g *Guard) prepareExponentTicker(ctx context.Context) *backoff.Ticker {
 	bo := backoff.WithContext(&backoff.ExponentialBackOff{
 		InitialInterval:     800 * time.Millisecond,
 		RandomizationFactor: 0.1,
 		Multiplier:          2,
 		MaxInterval:         g.timeout,
+		MaxElapsedTime:      reconnectMaxElapsedTime,
 		Stop:                backoff.Stop,
 		Clock:               backoff.SystemClock,
 	}, ctx)
@@ -127,6 +164,13 @@ func (g *Guard) prepareExponentTicker(ctx context.Context) *backoff.Ticker {
 	return ticker
 }
 
+func (g *Guard) triggerOfferSending() {
+	select {
+	case g.Reconnect <- struct{}{}:
+	default:
+	}
+}
+
 // Give chance to the peer to establish the initial connection.
 // With it, we can decrease to send necessary offer
 func waitForInitialConnectionTry(ctx context.Context) {