Simplify debug code

- Add stateful firewall functionality for UDP/TCP/ICMP in userspace firewalll
- Removes all egress drop rules/filters, still needs refactoring so we don't add output rules to any chains/filters.
- on Linux, if the OUTPUT policy is DROP  then we don't do anything about it (no extra allow rules). This is up to the user, if they don't want anything leaving their machine they'll have to manage these rules explicitly.

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+
+github: [netbirdio]

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

@@ -21,6 +21,7 @@ jobs:
         uses: actions/setup-go@v5
         with:
           go-version: "1.23.x"
+          cache: false
       - name: Checkout code
         uses: actions/checkout@v4
 
@@ -28,8 +29,9 @@ jobs:
         uses: actions/cache@v4
         with:
           path: ~/go/pkg/mod
-          key: macos-go-${{ hashFiles('**/go.sum') }}
+          key: macos-gotest-${{ hashFiles('**/go.sum') }}
           restore-keys: |
+            macos-gotest-
             macos-go-
 
       - name: Install libpcap
@@ -42,4 +44,4 @@ jobs:
         run: git --no-pager diff --exit-code
 
       - name: Test
-        run: NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 5m -p 1 ./...
+        run: NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 5m -p 1 $(go list ./... | grep -v /management)

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

@@ -11,30 +11,115 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  build-cache:
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: "1.23.x"
+          cache: false
+
+      - 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@v4
+        id: cache
+        with:
+          path: |
+            ${{ env.cache }}
+            ${{ env.modcache }}
+          key: ${{ runner.os }}-gotest-cache-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.os }}-gotest-cache-${{ hashFiles('**/go.sum') }}
+            
+
+      - name: Install dependencies
+        if: steps.cache.outputs.cache-hit != 'true'
+        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: steps.cache.outputs.cache-hit != 'true'
+        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
+
+      - name: Build client
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: client
+        run: CGO_ENABLED=1 go build .
+
+      - name: Build client 386
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: client
+        run: CGO_ENABLED=1 GOARCH=386 go build -o client-386 .
+
+      - name: Build management
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: management
+        run: CGO_ENABLED=1 go build .
+
+      - name: Build management 386
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: management
+        run: CGO_ENABLED=1 GOARCH=386 go build -o management-386 .
+
+      - name: Build signal
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: signal
+        run: CGO_ENABLED=1 go build .
+
+      - name: Build signal 386
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: signal
+        run: CGO_ENABLED=1 GOARCH=386 go build -o signal-386 .
+
+      - name: Build relay
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: relay
+        run: CGO_ENABLED=1 go build .
+
+      - name: Build relay 386
+        if: steps.cache.outputs.cache-hit != 'true'
+        working-directory: relay
+        run: CGO_ENABLED=1 GOARCH=386 go build -o relay-386 .
+
   test:
+    needs: [build-cache]
     strategy:
+      fail-fast: false
       matrix:
         arch: [ '386','amd64' ]
-        store: [ 'sqlite', 'postgres']
     runs-on: ubuntu-22.04
     steps:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
           go-version: "1.23.x"
-
-
-      - name: Cache Go modules
-        uses: actions/cache@v4
-        with:
-          path: ~/go/pkg/mod
-          key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
-          restore-keys: |
-            ${{ runner.os }}-go-
+          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
 
@@ -49,27 +134,134 @@ jobs:
         run: git --no-pager diff --exit-code
 
       - name: Test
-        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 10m -p 1 ./...
+        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} CI=true go test -exec 'sudo' -timeout 10m -p 1 $(go list ./... | grep -v /management)
+
+  test_management:
+    needs: [ build-cache ]
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: [ '386','amd64' ]
+        store: [ 'sqlite', 'postgres']
+    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 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
+
+      - name: check git status
+        run: git --no-pager diff --exit-code
+
+      - name: Test
+        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -p 1 -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 10m $(go list ./... | grep /management)
+
+  benchmark:
+    needs: [ build-cache ]
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: [ '386','amd64' ]
+        store: [ 'sqlite', 'postgres' ]
+    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 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
+
+      - name: check git status
+        run: git --no-pager diff --exit-code
+
+      - name: Test
+        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -run=^$ -bench=. -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 10m -p 1 ./...
 
   test_client_on_docker:
+    needs: [ build-cache ]
     runs-on: ubuntu-20.04
     steps:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
           go-version: "1.23.x"
-
-      - name: Cache Go modules
-        uses: actions/cache@v4
-        with:
-          path: ~/go/pkg/mod
-          key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
-          restore-keys: |
-            ${{ runner.os }}-go-
+          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
 

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

@@ -24,6 +24,23 @@ jobs:
         id: go
         with:
           go-version: "1.23.x"
+          cache: false
+
+      - name: Get Go environment
+        run: |
+          echo "cache=$(go env GOCACHE)" >> $env:GITHUB_ENV
+          echo "modcache=$(go env GOMODCACHE)" >> $env:GITHUB_ENV
+
+      - name: Cache Go modules
+        uses: actions/cache@v4
+        with:
+          path: |
+            ${{ env.cache }}
+            ${{ env.modcache }}
+          key: ${{ runner.os }}-gotest-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.os }}-gotest-
+            ${{ runner.os }}-go-
 
       - name: Download wintun
         uses: carlosperate/download-file-action@v2
@@ -42,11 +59,13 @@ jobs:
       - run: choco install -y sysinternals --ignore-checksums
       - run: choco install -y mingw
 
-      - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOMODCACHE=C:\Users\runneradmin\go\pkg\mod
-      - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOCACHE=C:\Users\runneradmin\AppData\Local\go-build
+      - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOMODCACHE=${{ env.cache }}
+      - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe env -w GOCACHE=${{ env.modcache }}
+      - run: PsExec64 -s -w ${{ github.workspace }} C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe mod tidy
+      - run: echo "files=$(go list ./... | ForEach-Object { $_ } | Where-Object { $_ -notmatch '/management' })" >> $env:GITHUB_ENV
 
       - name: test
-        run: PsExec64 -s -w ${{ github.workspace }} cmd.exe /c "C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe test -timeout 10m -p 1 ./... > test-out.txt 2>&1"
+        run: PsExec64 -s -w ${{ github.workspace }} cmd.exe /c "C:\hostedtoolcache\windows\go\${{ steps.go.outputs.go-version }}\x64\bin\go.exe test -timeout 10m -p 1 ${{ env.files }} > test-out.txt 2>&1"
       - name: test output
         if: ${{ always() }}
         run: Get-Content test-out.txt

.github/workflows/golangci-lint.yml

@@ -19,7 +19,7 @@ jobs:
       - name: codespell
         uses: codespell-project/actions-codespell@v2
         with:
-          ignore_words_list: erro,clienta,hastable,iif,groupd
+          ignore_words_list: erro,clienta,hastable,iif,groupd,testIn
           skip: go.mod,go.sum
           only_warn: 1
   golangci:
@@ -46,7 +46,7 @@ jobs:
         if: matrix.os == 'ubuntu-latest'
         run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev libpcap-dev
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v3
+        uses: golangci/golangci-lint-action@v4
         with:
           version: latest
-          args: --timeout=12m
+          args: --timeout=12m --out-format colored-line-number

.github/workflows/release.yml

@@ -9,7 +9,7 @@ on:
   pull_request:
 
 env:
-  SIGN_PIPE_VER: "v0.0.16"
+  SIGN_PIPE_VER: "v0.0.17"
   GORELEASER_VER: "v2.3.2"
   PRODUCT_NAME: "NetBird"
   COPYRIGHT: "Wiretrustee UG (haftungsbeschreankt)"

README.md

@@ -17,8 +17,12 @@
        <img src="https://img.shields.io/badge/license-BSD--3-blue" />
      </a> 
     <br>
-    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2p5zwhm4g-8fHollzrQa5y4PZF5AEpvQ">
+    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">
         <img src="https://img.shields.io/badge/slack-@netbird-red.svg?logo=slack"/>
+     </a>  
+     <br>
+    <a href="https://gurubase.io/g/netbird">
+        <img src="https://img.shields.io/badge/Gurubase-Ask%20NetBird%20Guru-006BFF"/>
      </a>    
   </p>
 </div>
@@ -30,7 +34,7 @@
   <br/>
   See <a href="https://netbird.io/docs/">Documentation</a>
   <br/>
-   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2p5zwhm4g-8fHollzrQa5y4PZF5AEpvQ">Slack channel</a>
+   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">Slack channel</a>
   <br/>
  
 </strong>

client/Dockerfile

@@ -1,4 +1,4 @@
-FROM alpine:3.20
+FROM alpine:3.21.0
 RUN apk add --no-cache ca-certificates iptables ip6tables
 ENV NB_FOREGROUND_MODE=true
 ENTRYPOINT [ "/usr/local/bin/netbird","up"]

client/anonymize/anonymize.go

@@ -12,6 +12,8 @@ import (
 	"strings"
 )
 
+const anonTLD = ".domain"
+
 type Anonymizer struct {
 	ipAnonymizer     map[netip.Addr]netip.Addr
 	domainAnonymizer map[string]string
@@ -19,6 +21,8 @@ type Anonymizer struct {
 	currentAnonIPv6  netip.Addr
 	startAnonIPv4    netip.Addr
 	startAnonIPv6    netip.Addr
+
+	domainKeyRegex *regexp.Regexp
 }
 
 func DefaultAddresses() (netip.Addr, netip.Addr) {
@@ -34,6 +38,8 @@ func NewAnonymizer(startIPv4, startIPv6 netip.Addr) *Anonymizer {
 		currentAnonIPv6:  startIPv6,
 		startAnonIPv4:    startIPv4,
 		startAnonIPv6:    startIPv6,
+
+		domainKeyRegex: regexp.MustCompile(`\bdomain=([^\s,:"]+)`),
 	}
 }
 
@@ -83,29 +89,39 @@ func (a *Anonymizer) AnonymizeIPString(ip string) string {
 }
 
 func (a *Anonymizer) AnonymizeDomain(domain string) string {
-	if strings.HasSuffix(domain, "netbird.io") ||
-		strings.HasSuffix(domain, "netbird.selfhosted") ||
-		strings.HasSuffix(domain, "netbird.cloud") ||
-		strings.HasSuffix(domain, "netbird.stage") ||
-		strings.HasSuffix(domain, ".domain") {
+	baseDomain := domain
+	hasDot := strings.HasSuffix(domain, ".")
+	if hasDot {
+		baseDomain = domain[:len(domain)-1]
+	}
+
+	if strings.HasSuffix(baseDomain, "netbird.io") ||
+		strings.HasSuffix(baseDomain, "netbird.selfhosted") ||
+		strings.HasSuffix(baseDomain, "netbird.cloud") ||
+		strings.HasSuffix(baseDomain, "netbird.stage") ||
+		strings.HasSuffix(baseDomain, anonTLD) {
 		return domain
 	}
 
-	parts := strings.Split(domain, ".")
+	parts := strings.Split(baseDomain, ".")
 	if len(parts) < 2 {
 		return domain
 	}
 
-	baseDomain := parts[len(parts)-2] + "." + parts[len(parts)-1]
+	baseForLookup := parts[len(parts)-2] + "." + parts[len(parts)-1]
 
-	anonymized, ok := a.domainAnonymizer[baseDomain]
+	anonymized, ok := a.domainAnonymizer[baseForLookup]
 	if !ok {
-		anonymizedBase := "anon-" + generateRandomString(5) + ".domain"
-		a.domainAnonymizer[baseDomain] = anonymizedBase
+		anonymizedBase := "anon-" + generateRandomString(5) + anonTLD
+		a.domainAnonymizer[baseForLookup] = anonymizedBase
 		anonymized = anonymizedBase
 	}
 
-	return strings.Replace(domain, baseDomain, anonymized, 1)
+	result := strings.Replace(baseDomain, baseForLookup, anonymized, 1)
+	if hasDot {
+		result += "."
+	}
+	return result
 }
 
 func (a *Anonymizer) AnonymizeURI(uri string) string {
@@ -152,27 +168,22 @@ func (a *Anonymizer) AnonymizeString(str string) string {
 	return str
 }
 
-// AnonymizeSchemeURI finds and anonymizes URIs with stun, stuns, turn, and turns schemes.
+// AnonymizeSchemeURI finds and anonymizes URIs with ws, wss, rel, rels, stun, stuns, turn, and turns schemes.
 func (a *Anonymizer) AnonymizeSchemeURI(text string) string {
-	re := regexp.MustCompile(`(?i)\b(stuns?:|turns?:|https?://)\S+\b`)
+	re := regexp.MustCompile(`(?i)\b(wss?://|rels?://|stuns?:|turns?:|https?://)\S+\b`)
 
 	return re.ReplaceAllStringFunc(text, a.AnonymizeURI)
 }
 
-// AnonymizeDNSLogLine anonymizes domain names in DNS log entries by replacing them with a random string.
 func (a *Anonymizer) AnonymizeDNSLogLine(logEntry string) string {
-	domainPattern := `dns\.Question{Name:"([^"]+)",`
-	domainRegex := regexp.MustCompile(domainPattern)
-
-	return domainRegex.ReplaceAllStringFunc(logEntry, func(match string) string {
-		parts := strings.Split(match, `"`)
+	return a.domainKeyRegex.ReplaceAllStringFunc(logEntry, func(match string) string {
+		parts := strings.SplitN(match, "=", 2)
 		if len(parts) >= 2 {
 			domain := parts[1]
-			if strings.HasSuffix(domain, ".domain") {
+			if strings.HasSuffix(domain, anonTLD) {
 				return match
 			}
-			randomDomain := generateRandomString(10) + ".domain"
-			return strings.Replace(match, domain, randomDomain, 1)
+			return "domain=" + a.AnonymizeDomain(domain)
 		}
 		return match
 	})
@@ -201,6 +212,8 @@ func isWellKnown(addr netip.Addr) bool {
 		"2606:4700:4700::1111", "2606:4700:4700::1001", // Cloudflare DNS IPv6
 		"9.9.9.9", "149.112.112.112", // Quad9 DNS IPv4
 		"2620:fe::fe", "2620:fe::9", // Quad9 DNS IPv6
+
+		"128.0.0.0", "8000::", // 2nd split subnet for default routes
 	}
 
 	if slices.Contains(wellKnown, addr.String()) {

client/anonymize/anonymize_test.go

@@ -46,11 +46,59 @@ func TestAnonymizeIP(t *testing.T) {
 
 func TestAnonymizeDNSLogLine(t *testing.T) {
 	anonymizer := anonymize.NewAnonymizer(netip.Addr{}, netip.Addr{})
-	testLog := `2024-04-23T20:01:11+02:00 TRAC client/internal/dns/local.go:25: received question: dns.Question{Name:"example.com", Qtype:0x1c, Qclass:0x1}`
+	tests := []struct {
+		name     string
+		input    string
+		original string
+		expect   string
+	}{
+		{
+			name:     "Basic domain with trailing content",
+			input:    "received DNS request for DNS forwarder: domain=example.com: something happened with code=123",
+			original: "example.com",
+			expect:   `received DNS request for DNS forwarder: domain=anon-[a-zA-Z0-9]+\.domain: something happened with code=123`,
+		},
+		{
+			name:     "Domain with trailing dot",
+			input:    "domain=example.com. processing request with status=pending",
+			original: "example.com",
+			expect:   `domain=anon-[a-zA-Z0-9]+\.domain\. processing request with status=pending`,
+		},
+		{
+			name:     "Multiple domains in log",
+			input:    "forward domain=first.com status=ok, redirect to domain=second.com port=443",
+			original: "first.com", // testing just one is sufficient as AnonymizeDomain is tested separately
+			expect:   `forward domain=anon-[a-zA-Z0-9]+\.domain status=ok, redirect to domain=anon-[a-zA-Z0-9]+\.domain port=443`,
+		},
+		{
+			name:     "Already anonymized domain",
+			input:    "got request domain=anon-xyz123.domain from=client1 to=server2",
+			original: "", // nothing should be anonymized
+			expect:   `got request domain=anon-xyz123\.domain from=client1 to=server2`,
+		},
+		{
+			name:     "Subdomain with trailing dot",
+			input:    "domain=sub.example.com. next_hop=10.0.0.1 proto=udp",
+			original: "example.com",
+			expect:   `domain=sub\.anon-[a-zA-Z0-9]+\.domain\. next_hop=10\.0\.0\.1 proto=udp`,
+		},
+		{
+			name:     "Handler chain pattern log",
+			input:    "pattern: domain=example.com. original: domain=*.example.com. wildcard=true priority=100",
+			original: "example.com",
+			expect:   `pattern: domain=anon-[a-zA-Z0-9]+\.domain\. original: domain=\*\.anon-[a-zA-Z0-9]+\.domain\. wildcard=true priority=100`,
+		},
+	}
 
-	result := anonymizer.AnonymizeDNSLogLine(testLog)
-	require.NotEqual(t, testLog, result)
-	assert.NotContains(t, result, "example.com")
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			result := anonymizer.AnonymizeDNSLogLine(tc.input)
+			if tc.original != "" {
+				assert.NotContains(t, result, tc.original)
+			}
+			assert.Regexp(t, tc.expect, result)
+		})
+	}
 }
 
 func TestAnonymizeDomain(t *testing.T) {
@@ -67,18 +115,36 @@ func TestAnonymizeDomain(t *testing.T) {
 			`^anon-[a-zA-Z0-9]+\.domain$`,
 			true,
 		},
+		{
+			"Domain with Trailing Dot",
+			"example.com.",
+			`^anon-[a-zA-Z0-9]+\.domain.$`,
+			true,
+		},
 		{
 			"Subdomain",
 			"sub.example.com",
 			`^sub\.anon-[a-zA-Z0-9]+\.domain$`,
 			true,
 		},
+		{
+			"Subdomain with Trailing Dot",
+			"sub.example.com.",
+			`^sub\.anon-[a-zA-Z0-9]+\.domain.$`,
+			true,
+		},
 		{
 			"Protected Domain",
 			"netbird.io",
 			`^netbird\.io$`,
 			false,
 		},
+		{
+			"Protected Domain with Trailing Dot",
+			"netbird.io.",
+			`^netbird\.io.$`,
+			false,
+		},
 	}
 
 	for _, tc := range tests {
@@ -140,8 +206,16 @@ func TestAnonymizeSchemeURI(t *testing.T) {
 		expect string
 	}{
 		{"STUN URI in text", "Connection made via stun:example.com", `Connection made via stun:anon-[a-zA-Z0-9]+\.domain`},
+		{"STUNS URI in message", "Secure connection to stuns:example.com:443", `Secure connection to stuns:anon-[a-zA-Z0-9]+\.domain:443`},
 		{"TURN URI in log", "Failed attempt turn:some.example.com:3478?transport=tcp: retrying", `Failed attempt turn:some.anon-[a-zA-Z0-9]+\.domain:3478\?transport=tcp: retrying`},
+		{"TURNS URI in message", "Secure connection to turns:example.com:5349", `Secure connection to turns:anon-[a-zA-Z0-9]+\.domain:5349`},
+		{"HTTP URI in text", "Visit http://example.com for more", `Visit http://anon-[a-zA-Z0-9]+\.domain for more`},
+		{"HTTPS URI in CAPS", "Visit HTTPS://example.com for more", `Visit https://anon-[a-zA-Z0-9]+\.domain for more`},
 		{"HTTPS URI in message", "Visit https://example.com for more", `Visit https://anon-[a-zA-Z0-9]+\.domain for more`},
+		{"WS URI in log", "Connection established to ws://example.com:8080", `Connection established to ws://anon-[a-zA-Z0-9]+\.domain:8080`},
+		{"WSS URI in message", "Secure connection to wss://example.com", `Secure connection to wss://anon-[a-zA-Z0-9]+\.domain`},
+		{"Rel URI in text", "Relaying to rel://example.com", `Relaying to rel://anon-[a-zA-Z0-9]+\.domain`},
+		{"Rels URI in message", "Relaying to rels://example.com", `Relaying to rels://anon-[a-zA-Z0-9]+\.domain`},
 	}
 
 	for _, tc := range tests {

client/cmd/debug.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -61,6 +62,15 @@ var forCmd = &cobra.Command{
 	RunE:    runForDuration,
 }
 
+var persistenceCmd = &cobra.Command{
+	Use:     "persistence [on|off]",
+	Short:   "Set network map memory persistence",
+	Long:    `Configure whether the latest network map should persist in memory. When enabled, the last known network map will be kept in memory.`,
+	Example: "  netbird debug persistence on",
+	Args:    cobra.ExactArgs(1),
+	RunE:    setNetworkMapPersistence,
+}
+
 func debugBundle(cmd *cobra.Command, _ []string) error {
 	conn, err := getClient(cmd)
 	if err != nil {
@@ -171,6 +181,13 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 
 	time.Sleep(1 * time.Second)
 
+	// Enable network map persistence before bringing the service up
+	if _, err := client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
+		Enabled: true,
+	}); err != nil {
+		return fmt.Errorf("failed to enable network map persistence: %v", status.Convert(err).Message())
+	}
+
 	if _, err := client.Up(cmd.Context(), &proto.UpRequest{}); err != nil {
 		return fmt.Errorf("failed to up: %v", status.Convert(err).Message())
 	}
@@ -200,6 +217,13 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		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())
@@ -219,6 +243,34 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+func setNetworkMapPersistence(cmd *cobra.Command, args []string) error {
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Errorf(errCloseConnection, err)
+		}
+	}()
+
+	persistence := strings.ToLower(args[0])
+	if persistence != "on" && persistence != "off" {
+		return fmt.Errorf("invalid persistence value: %s. Use 'on' or 'off'", args[0])
+	}
+
+	client := proto.NewDaemonServiceClient(conn)
+	_, err = client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
+		Enabled: persistence == "on",
+	})
+	if err != nil {
+		return fmt.Errorf("failed to set network map persistence: %v", status.Convert(err).Message())
+	}
+
+	cmd.Printf("Network map persistence set to: %s\n", persistence)
+	return nil
+}
+
 func getStatusOutput(cmd *cobra.Command) string {
 	var statusOutputString string
 	statusResp, err := getStatus(cmd.Context())

client/cmd/networks.go

@@ -0,0 +1,173 @@
+package cmd
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"google.golang.org/grpc/status"
+
+	"github.com/netbirdio/netbird/client/proto"
+)
+
+var appendFlag bool
+
+var networksCMD = &cobra.Command{
+	Use:     "networks",
+	Aliases: []string{"routes"},
+	Short:   "Manage networks",
+	Long:    `Commands to list, select, or deselect networks. Replaces the "routes" command.`,
+}
+
+var routesListCmd = &cobra.Command{
+	Use:     "list",
+	Aliases: []string{"ls"},
+	Short:   "List networks",
+	Example: "  netbird networks list",
+	Long:    "List all available network routes.",
+	RunE:    networksList,
+}
+
+var routesSelectCmd = &cobra.Command{
+	Use:     "select network...|all",
+	Short:   "Select network",
+	Long:    "Select a list of networks by identifiers or 'all' to clear all selections and to accept all (including new) networks.\nDefault mode is replace, use -a to append to already selected networks.",
+	Example: "  netbird networks select all\n  netbird networks select route1 route2\n  netbird routes select -a route3",
+	Args:    cobra.MinimumNArgs(1),
+	RunE:    networksSelect,
+}
+
+var routesDeselectCmd = &cobra.Command{
+	Use:     "deselect network...|all",
+	Short:   "Deselect networks",
+	Long:    "Deselect previously selected networks by identifiers or 'all' to disable accepting any networks.",
+	Example: "  netbird networks deselect all\n  netbird networks deselect route1 route2",
+	Args:    cobra.MinimumNArgs(1),
+	RunE:    networksDeselect,
+}
+
+func init() {
+	routesSelectCmd.PersistentFlags().BoolVarP(&appendFlag, "append", "a", false, "Append to current network selection instead of replacing")
+}
+
+func networksList(cmd *cobra.Command, _ []string) error {
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	client := proto.NewDaemonServiceClient(conn)
+	resp, err := client.ListNetworks(cmd.Context(), &proto.ListNetworksRequest{})
+	if err != nil {
+		return fmt.Errorf("failed to list network: %v", status.Convert(err).Message())
+	}
+
+	if len(resp.Routes) == 0 {
+		cmd.Println("No networks available.")
+		return nil
+	}
+
+	printNetworks(cmd, resp)
+
+	return nil
+}
+
+func printNetworks(cmd *cobra.Command, resp *proto.ListNetworksResponse) {
+	cmd.Println("Available Networks:")
+	for _, route := range resp.Routes {
+		printNetwork(cmd, route)
+	}
+}
+
+func printNetwork(cmd *cobra.Command, route *proto.Network) {
+	selectedStatus := getSelectedStatus(route)
+	domains := route.GetDomains()
+
+	if len(domains) > 0 {
+		printDomainRoute(cmd, route, domains, selectedStatus)
+	} else {
+		printNetworkRoute(cmd, route, selectedStatus)
+	}
+}
+
+func getSelectedStatus(route *proto.Network) string {
+	if route.GetSelected() {
+		return "Selected"
+	}
+	return "Not Selected"
+}
+
+func printDomainRoute(cmd *cobra.Command, route *proto.Network, domains []string, selectedStatus string) {
+	cmd.Printf("\n  - ID: %s\n    Domains: %s\n    Status: %s\n", route.GetID(), strings.Join(domains, ", "), selectedStatus)
+	resolvedIPs := route.GetResolvedIPs()
+
+	if len(resolvedIPs) > 0 {
+		printResolvedIPs(cmd, domains, resolvedIPs)
+	} else {
+		cmd.Printf("    Resolved IPs: -\n")
+	}
+}
+
+func printNetworkRoute(cmd *cobra.Command, route *proto.Network, selectedStatus string) {
+	cmd.Printf("\n  - ID: %s\n    Network: %s\n    Status: %s\n", route.GetID(), route.GetRange(), selectedStatus)
+}
+
+func printResolvedIPs(cmd *cobra.Command, _ []string, resolvedIPs map[string]*proto.IPList) {
+	cmd.Printf("    Resolved IPs:\n")
+	for resolvedDomain, ipList := range resolvedIPs {
+		cmd.Printf("      [%s]: %s\n", resolvedDomain, strings.Join(ipList.GetIps(), ", "))
+	}
+}
+
+func networksSelect(cmd *cobra.Command, args []string) error {
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	client := proto.NewDaemonServiceClient(conn)
+	req := &proto.SelectNetworksRequest{
+		NetworkIDs: args,
+	}
+
+	if len(args) == 1 && args[0] == "all" {
+		req.All = true
+	} else if appendFlag {
+		req.Append = true
+	}
+
+	if _, err := client.SelectNetworks(cmd.Context(), req); err != nil {
+		return fmt.Errorf("failed to select networks: %v", status.Convert(err).Message())
+	}
+
+	cmd.Println("Networks selected successfully.")
+
+	return nil
+}
+
+func networksDeselect(cmd *cobra.Command, args []string) error {
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	client := proto.NewDaemonServiceClient(conn)
+	req := &proto.SelectNetworksRequest{
+		NetworkIDs: args,
+	}
+
+	if len(args) == 1 && args[0] == "all" {
+		req.All = true
+	}
+
+	if _, err := client.DeselectNetworks(cmd.Context(), req); err != nil {
+		return fmt.Errorf("failed to deselect networks: %v", status.Convert(err).Message())
+	}
+
+	cmd.Println("Networks deselected successfully.")
+
+	return nil
+}

client/cmd/pprof.go

@@ -0,0 +1,33 @@
+//go:build pprof
+// +build pprof
+
+package cmd
+
+import (
+	"net/http"
+	_ "net/http/pprof"
+	"os"
+
+	log "github.com/sirupsen/logrus"
+)
+
+func init() {
+	addr := pprofAddr()
+	go pprof(addr)
+}
+
+func pprofAddr() string {
+	listenAddr := os.Getenv("NB_PPROF_ADDR")
+	if listenAddr == "" {
+		return "localhost:6969"
+	}
+
+	return listenAddr
+}
+
+func pprof(listenAddr string) {
+	log.Infof("listening pprof on: %s\n", listenAddr)
+	if err := http.ListenAndServe(listenAddr, nil); err != nil {
+		log.Fatalf("Failed to start pprof: %v", err)
+	}
+}

client/cmd/root.go

@@ -142,19 +142,20 @@ func init() {
 	rootCmd.AddCommand(loginCmd)
 	rootCmd.AddCommand(versionCmd)
 	rootCmd.AddCommand(sshCmd)
-	rootCmd.AddCommand(routesCmd)
+	rootCmd.AddCommand(networksCMD)
 	rootCmd.AddCommand(debugCmd)
 
 	serviceCmd.AddCommand(runCmd, startCmd, stopCmd, restartCmd) // service control commands are subcommands of service
 	serviceCmd.AddCommand(installCmd, uninstallCmd)              // service installer commands are subcommands of service
 
-	routesCmd.AddCommand(routesListCmd)
-	routesCmd.AddCommand(routesSelectCmd, routesDeselectCmd)
+	networksCMD.AddCommand(routesListCmd)
+	networksCMD.AddCommand(routesSelectCmd, routesDeselectCmd)
 
 	debugCmd.AddCommand(debugBundleCmd)
 	debugCmd.AddCommand(logCmd)
 	logCmd.AddCommand(logLevelCmd)
 	debugCmd.AddCommand(forCmd)
+	debugCmd.AddCommand(persistenceCmd)
 
 	upCmd.PersistentFlags().StringSliceVar(&natExternalIPs, externalIPMapFlag, nil,
 		`Sets external IPs maps between local addresses and interfaces.`+

client/cmd/route.go

@@ -1,174 +0,0 @@
-package cmd
-
-import (
-	"fmt"
-	"strings"
-
-	"github.com/spf13/cobra"
-	"google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/proto"
-)
-
-var appendFlag bool
-
-var routesCmd = &cobra.Command{
-	Use:   "routes",
-	Short: "Manage network routes",
-	Long:  `Commands to list, select, or deselect network routes.`,
-}
-
-var routesListCmd = &cobra.Command{
-	Use:     "list",
-	Aliases: []string{"ls"},
-	Short:   "List routes",
-	Example: "  netbird routes list",
-	Long:    "List all available network routes.",
-	RunE:    routesList,
-}
-
-var routesSelectCmd = &cobra.Command{
-	Use:     "select route...|all",
-	Short:   "Select routes",
-	Long:    "Select a list of routes by identifiers or 'all' to clear all selections and to accept all (including new) routes.\nDefault mode is replace, use -a to append to already selected routes.",
-	Example: "  netbird routes select all\n  netbird routes select route1 route2\n  netbird routes select -a route3",
-	Args:    cobra.MinimumNArgs(1),
-	RunE:    routesSelect,
-}
-
-var routesDeselectCmd = &cobra.Command{
-	Use:     "deselect route...|all",
-	Short:   "Deselect routes",
-	Long:    "Deselect previously selected routes by identifiers or 'all' to disable accepting any routes.",
-	Example: "  netbird routes deselect all\n  netbird routes deselect route1 route2",
-	Args:    cobra.MinimumNArgs(1),
-	RunE:    routesDeselect,
-}
-
-func init() {
-	routesSelectCmd.PersistentFlags().BoolVarP(&appendFlag, "append", "a", false, "Append to current route selection instead of replacing")
-}
-
-func routesList(cmd *cobra.Command, _ []string) error {
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-
-	client := proto.NewDaemonServiceClient(conn)
-	resp, err := client.ListRoutes(cmd.Context(), &proto.ListRoutesRequest{})
-	if err != nil {
-		return fmt.Errorf("failed to list routes: %v", status.Convert(err).Message())
-	}
-
-	if len(resp.Routes) == 0 {
-		cmd.Println("No routes available.")
-		return nil
-	}
-
-	printRoutes(cmd, resp)
-
-	return nil
-}
-
-func printRoutes(cmd *cobra.Command, resp *proto.ListRoutesResponse) {
-	cmd.Println("Available Routes:")
-	for _, route := range resp.Routes {
-		printRoute(cmd, route)
-	}
-}
-
-func printRoute(cmd *cobra.Command, route *proto.Route) {
-	selectedStatus := getSelectedStatus(route)
-	domains := route.GetDomains()
-
-	if len(domains) > 0 {
-		printDomainRoute(cmd, route, domains, selectedStatus)
-	} else {
-		printNetworkRoute(cmd, route, selectedStatus)
-	}
-}
-
-func getSelectedStatus(route *proto.Route) string {
-	if route.GetSelected() {
-		return "Selected"
-	}
-	return "Not Selected"
-}
-
-func printDomainRoute(cmd *cobra.Command, route *proto.Route, domains []string, selectedStatus string) {
-	cmd.Printf("\n  - ID: %s\n    Domains: %s\n    Status: %s\n", route.GetID(), strings.Join(domains, ", "), selectedStatus)
-	resolvedIPs := route.GetResolvedIPs()
-
-	if len(resolvedIPs) > 0 {
-		printResolvedIPs(cmd, domains, resolvedIPs)
-	} else {
-		cmd.Printf("    Resolved IPs: -\n")
-	}
-}
-
-func printNetworkRoute(cmd *cobra.Command, route *proto.Route, selectedStatus string) {
-	cmd.Printf("\n  - ID: %s\n    Network: %s\n    Status: %s\n", route.GetID(), route.GetNetwork(), selectedStatus)
-}
-
-func printResolvedIPs(cmd *cobra.Command, domains []string, resolvedIPs map[string]*proto.IPList) {
-	cmd.Printf("    Resolved IPs:\n")
-	for _, domain := range domains {
-		if ipList, exists := resolvedIPs[domain]; exists {
-			cmd.Printf("      [%s]: %s\n", domain, strings.Join(ipList.GetIps(), ", "))
-		}
-	}
-}
-
-func routesSelect(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-
-	client := proto.NewDaemonServiceClient(conn)
-	req := &proto.SelectRoutesRequest{
-		RouteIDs: args,
-	}
-
-	if len(args) == 1 && args[0] == "all" {
-		req.All = true
-	} else if appendFlag {
-		req.Append = true
-	}
-
-	if _, err := client.SelectRoutes(cmd.Context(), req); err != nil {
-		return fmt.Errorf("failed to select routes: %v", status.Convert(err).Message())
-	}
-
-	cmd.Println("Routes selected successfully.")
-
-	return nil
-}
-
-func routesDeselect(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-
-	client := proto.NewDaemonServiceClient(conn)
-	req := &proto.SelectRoutesRequest{
-		RouteIDs: args,
-	}
-
-	if len(args) == 1 && args[0] == "all" {
-		req.All = true
-	}
-
-	if _, err := client.DeselectRoutes(cmd.Context(), req); err != nil {
-		return fmt.Errorf("failed to deselect routes: %v", status.Convert(err).Message())
-	}
-
-	cmd.Println("Routes deselected successfully.")
-
-	return nil
-}

client/cmd/service.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"sync"
 
 	"github.com/kardianos/service"
 	log "github.com/sirupsen/logrus"
@@ -13,10 +14,11 @@ import (
 )
 
 type program struct {
-	ctx            context.Context
-	cancel         context.CancelFunc
-	serv           *grpc.Server
-	serverInstance *server.Server
+	ctx              context.Context
+	cancel           context.CancelFunc
+	serv             *grpc.Server
+	serverInstance   *server.Server
+	serverInstanceMu sync.Mutex
 }
 
 func newProgram(ctx context.Context, cancel context.CancelFunc) *program {

client/cmd/service_controller.go

@@ -61,7 +61,9 @@ func (p *program) Start(svc service.Service) error {
 		}
 		proto.RegisterDaemonServiceServer(p.serv, serverInstance)
 
+		p.serverInstanceMu.Lock()
 		p.serverInstance = serverInstance
+		p.serverInstanceMu.Unlock()
 
 		log.Printf("started daemon server: %v", split[1])
 		if err := p.serv.Serve(listen); err != nil {
@@ -72,6 +74,7 @@ func (p *program) Start(svc service.Service) error {
 }
 
 func (p *program) Stop(srv service.Service) error {
+	p.serverInstanceMu.Lock()
 	if p.serverInstance != nil {
 		in := new(proto.DownRequest)
 		_, err := p.serverInstance.Down(p.ctx, in)
@@ -79,6 +82,7 @@ func (p *program) Stop(srv service.Service) error {
 			log.Errorf("failed to stop daemon: %v", err)
 		}
 	}
+	p.serverInstanceMu.Unlock()
 
 	p.cancel()
 

client/cmd/state.go

@@ -0,0 +1,181 @@
+package cmd
+
+import (
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"google.golang.org/grpc/status"
+
+	"github.com/netbirdio/netbird/client/proto"
+)
+
+var (
+	allFlag bool
+)
+
+var stateCmd = &cobra.Command{
+	Use:   "state",
+	Short: "Manage daemon state",
+	Long:  "Provides commands for managing and inspecting the Netbird daemon state.",
+}
+
+var stateListCmd = &cobra.Command{
+	Use:     "list",
+	Aliases: []string{"ls"},
+	Short:   "List all stored states",
+	Long:    "Lists all registered states with their status and basic information.",
+	Example: "  netbird state list",
+	RunE:    stateList,
+}
+
+var stateCleanCmd = &cobra.Command{
+	Use:   "clean [state-name]",
+	Short: "Clean stored states",
+	Long: `Clean specific state or all states. The daemon must not be running.
+This will perform cleanup operations and remove the state.`,
+	Example: `  netbird state clean dns_state
+  netbird state clean --all`,
+	RunE: stateClean,
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		// Check mutual exclusivity between --all flag and state-name argument
+		if allFlag && len(args) > 0 {
+			return fmt.Errorf("cannot specify both --all flag and state name")
+		}
+		if !allFlag && len(args) != 1 {
+			return fmt.Errorf("requires a state name argument or --all flag")
+		}
+		return nil
+	},
+}
+
+var stateDeleteCmd = &cobra.Command{
+	Use:   "delete [state-name]",
+	Short: "Delete stored states",
+	Long: `Delete specific state or all states from storage. The daemon must not be running.
+This will remove the state without performing any cleanup operations.`,
+	Example: `  netbird state delete dns_state
+  netbird state delete --all`,
+	RunE: stateDelete,
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		// Check mutual exclusivity between --all flag and state-name argument
+		if allFlag && len(args) > 0 {
+			return fmt.Errorf("cannot specify both --all flag and state name")
+		}
+		if !allFlag && len(args) != 1 {
+			return fmt.Errorf("requires a state name argument or --all flag")
+		}
+		return nil
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(stateCmd)
+	stateCmd.AddCommand(stateListCmd, stateCleanCmd, stateDeleteCmd)
+
+	stateCleanCmd.Flags().BoolVarP(&allFlag, "all", "a", false, "Clean all states")
+	stateDeleteCmd.Flags().BoolVarP(&allFlag, "all", "a", false, "Delete all states")
+}
+
+func stateList(cmd *cobra.Command, _ []string) error {
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Errorf(errCloseConnection, err)
+		}
+	}()
+
+	client := proto.NewDaemonServiceClient(conn)
+	resp, err := client.ListStates(cmd.Context(), &proto.ListStatesRequest{})
+	if err != nil {
+		return fmt.Errorf("failed to list states: %v", status.Convert(err).Message())
+	}
+
+	cmd.Printf("\nStored states:\n\n")
+	for _, state := range resp.States {
+		cmd.Printf("- %s\n", state.Name)
+	}
+
+	return nil
+}
+
+func stateClean(cmd *cobra.Command, args []string) error {
+	var stateName string
+	if !allFlag {
+		stateName = args[0]
+	}
+
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Errorf(errCloseConnection, err)
+		}
+	}()
+
+	client := proto.NewDaemonServiceClient(conn)
+	resp, err := client.CleanState(cmd.Context(), &proto.CleanStateRequest{
+		StateName: stateName,
+		All:       allFlag,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to clean state: %v", status.Convert(err).Message())
+	}
+
+	if resp.CleanedStates == 0 {
+		cmd.Println("No states were cleaned")
+		return nil
+	}
+
+	if allFlag {
+		cmd.Printf("Successfully cleaned %d states\n", resp.CleanedStates)
+	} else {
+		cmd.Printf("Successfully cleaned state %q\n", stateName)
+	}
+
+	return nil
+}
+
+func stateDelete(cmd *cobra.Command, args []string) error {
+	var stateName string
+	if !allFlag {
+		stateName = args[0]
+	}
+
+	conn, err := getClient(cmd)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Errorf(errCloseConnection, err)
+		}
+	}()
+
+	client := proto.NewDaemonServiceClient(conn)
+	resp, err := client.DeleteState(cmd.Context(), &proto.DeleteStateRequest{
+		StateName: stateName,
+		All:       allFlag,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to delete state: %v", status.Convert(err).Message())
+	}
+
+	if resp.DeletedStates == 0 {
+		cmd.Println("No states were deleted")
+		return nil
+	}
+
+	if allFlag {
+		cmd.Printf("Successfully deleted %d states\n", resp.DeletedStates)
+	} else {
+		cmd.Printf("Successfully deleted state %q\n", stateName)
+	}
+
+	return nil
+}

client/cmd/status.go

@@ -40,6 +40,7 @@ type peerStateDetailOutput struct {
 	Latency                time.Duration    `json:"latency" yaml:"latency"`
 	RosenpassEnabled       bool             `json:"quantumResistance" yaml:"quantumResistance"`
 	Routes                 []string         `json:"routes" yaml:"routes"`
+	Networks               []string         `json:"networks" yaml:"networks"`
 }
 
 type peersStateOutput struct {
@@ -98,6 +99,7 @@ type statusOutputOverview struct {
 	RosenpassEnabled    bool                       `json:"quantumResistance" yaml:"quantumResistance"`
 	RosenpassPermissive bool                       `json:"quantumResistancePermissive" yaml:"quantumResistancePermissive"`
 	Routes              []string                   `json:"routes" yaml:"routes"`
+	Networks            []string                   `json:"networks" yaml:"networks"`
 	NSServerGroups      []nsServerGroupStateOutput `json:"dnsServers" yaml:"dnsServers"`
 }
 
@@ -282,7 +284,8 @@ func convertToStatusOutputOverview(resp *proto.StatusResponse) statusOutputOverv
 		FQDN:                pbFullStatus.GetLocalPeerState().GetFqdn(),
 		RosenpassEnabled:    pbFullStatus.GetLocalPeerState().GetRosenpassEnabled(),
 		RosenpassPermissive: pbFullStatus.GetLocalPeerState().GetRosenpassPermissive(),
-		Routes:              pbFullStatus.GetLocalPeerState().GetRoutes(),
+		Routes:              pbFullStatus.GetLocalPeerState().GetNetworks(),
+		Networks:            pbFullStatus.GetLocalPeerState().GetNetworks(),
 		NSServerGroups:      mapNSGroups(pbFullStatus.GetDnsServers()),
 	}
 
@@ -390,7 +393,8 @@ func mapPeers(peers []*proto.PeerState) peersStateOutput {
 			TransferSent:           transferSent,
 			Latency:                pbPeerState.GetLatency().AsDuration(),
 			RosenpassEnabled:       pbPeerState.GetRosenpassEnabled(),
-			Routes:                 pbPeerState.GetRoutes(),
+			Routes:                 pbPeerState.GetNetworks(),
+			Networks:               pbPeerState.GetNetworks(),
 		}
 
 		peersStateDetail = append(peersStateDetail, peerState)
@@ -491,10 +495,10 @@ func parseGeneralSummary(overview statusOutputOverview, showURL bool, showRelays
 		relaysString = fmt.Sprintf("%d/%d Available", overview.Relays.Available, overview.Relays.Total)
 	}
 
-	routes := "-"
-	if len(overview.Routes) > 0 {
-		sort.Strings(overview.Routes)
-		routes = strings.Join(overview.Routes, ", ")
+	networks := "-"
+	if len(overview.Networks) > 0 {
+		sort.Strings(overview.Networks)
+		networks = strings.Join(overview.Networks, ", ")
 	}
 
 	var dnsServersString string
@@ -556,6 +560,7 @@ func parseGeneralSummary(overview statusOutputOverview, showURL bool, showRelays
 			"Interface type: %s\n"+
 			"Quantum resistance: %s\n"+
 			"Routes: %s\n"+
+			"Networks: %s\n"+
 			"Peers count: %s\n",
 		fmt.Sprintf("%s/%s%s", goos, goarch, goarm),
 		overview.DaemonVersion,
@@ -568,7 +573,8 @@ func parseGeneralSummary(overview statusOutputOverview, showURL bool, showRelays
 		interfaceIP,
 		interfaceTypeString,
 		rosenpassEnabledStatus,
-		routes,
+		networks,
+		networks,
 		peersCountString,
 	)
 	return summary
@@ -631,10 +637,10 @@ func parsePeers(peers peersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 			}
 		}
 
-		routes := "-"
-		if len(peerState.Routes) > 0 {
-			sort.Strings(peerState.Routes)
-			routes = strings.Join(peerState.Routes, ", ")
+		networks := "-"
+		if len(peerState.Networks) > 0 {
+			sort.Strings(peerState.Networks)
+			networks = strings.Join(peerState.Networks, ", ")
 		}
 
 		peerString := fmt.Sprintf(
@@ -652,6 +658,7 @@ func parsePeers(peers peersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 				"  Transfer status (received/sent) %s/%s\n"+
 				"  Quantum resistance: %s\n"+
 				"  Routes: %s\n"+
+				"  Networks: %s\n"+
 				"  Latency: %s\n",
 			peerState.FQDN,
 			peerState.IP,
@@ -668,7 +675,8 @@ func parsePeers(peers peersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 			toIEC(peerState.TransferReceived),
 			toIEC(peerState.TransferSent),
 			rosenpassEnabledStatus,
-			routes,
+			networks,
+			networks,
 			peerState.Latency.String(),
 		)
 
@@ -680,7 +688,7 @@ func parsePeers(peers peersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 func skipDetailByFilters(peerState *proto.PeerState, isConnected bool) bool {
 	statusEval := false
 	ipEval := false
-	nameEval := false
+	nameEval := true
 
 	if statusFilter != "" {
 		lowerStatusFilter := strings.ToLower(statusFilter)
@@ -700,11 +708,13 @@ func skipDetailByFilters(peerState *proto.PeerState, isConnected bool) bool {
 
 	if len(prefixNamesFilter) > 0 {
 		for prefixNameFilter := range prefixNamesFilterMap {
-			if !strings.HasPrefix(peerState.Fqdn, prefixNameFilter) {
-				nameEval = true
+			if strings.HasPrefix(peerState.Fqdn, prefixNameFilter) {
+				nameEval = false
 				break
 			}
 		}
+	} else {
+		nameEval = false
 	}
 
 	return statusEval || ipEval || nameEval
@@ -808,6 +818,14 @@ func anonymizePeerDetail(a *anonymize.Anonymizer, peer *peerStateDetailOutput) {
 
 	peer.RelayAddress = a.AnonymizeURI(peer.RelayAddress)
 
+	for i, route := range peer.Networks {
+		peer.Networks[i] = a.AnonymizeIPString(route)
+	}
+
+	for i, route := range peer.Networks {
+		peer.Networks[i] = a.AnonymizeRoute(route)
+	}
+
 	for i, route := range peer.Routes {
 		peer.Routes[i] = a.AnonymizeIPString(route)
 	}
@@ -848,6 +866,10 @@ func anonymizeOverview(a *anonymize.Anonymizer, overview *statusOutputOverview)
 		}
 	}
 
+	for i, route := range overview.Networks {
+		overview.Networks[i] = a.AnonymizeRoute(route)
+	}
+
 	for i, route := range overview.Routes {
 		overview.Routes[i] = a.AnonymizeRoute(route)
 	}

client/cmd/status_test.go

@@ -44,7 +44,7 @@ var resp = &proto.StatusResponse{
 				LastWireguardHandshake:     timestamppb.New(time.Date(2001, time.Month(1), 1, 1, 1, 2, 0, time.UTC)),
 				BytesRx:                    200,
 				BytesTx:                    100,
-				Routes: []string{
+				Networks: []string{
 					"10.1.0.0/24",
 				},
 				Latency: durationpb.New(time.Duration(10000000)),
@@ -93,7 +93,7 @@ var resp = &proto.StatusResponse{
 			PubKey:          "Some-Pub-Key",
 			KernelInterface: true,
 			Fqdn:            "some-localhost.awesome-domain.com",
-			Routes: []string{
+			Networks: []string{
 				"10.10.0.0/24",
 			},
 		},
@@ -149,6 +149,9 @@ var overview = statusOutputOverview{
 				Routes: []string{
 					"10.1.0.0/24",
 				},
+				Networks: []string{
+					"10.1.0.0/24",
+				},
 				Latency: time.Duration(10000000),
 			},
 			{
@@ -230,6 +233,9 @@ var overview = statusOutputOverview{
 	Routes: []string{
 		"10.10.0.0/24",
 	},
+	Networks: []string{
+		"10.10.0.0/24",
+	},
 }
 
 func TestConversionFromFullStatusToOutputOverview(t *testing.T) {
@@ -295,6 +301,9 @@ func TestParsingToJSON(t *testing.T) {
                 "quantumResistance": false,
                 "routes": [
                   "10.1.0.0/24"
+                ],
+                "networks": [
+                  "10.1.0.0/24"
                 ]
               },
               {
@@ -318,7 +327,8 @@ func TestParsingToJSON(t *testing.T) {
                 "transferSent": 1000,
 				"latency": 10000000,
                 "quantumResistance": false,
-                "routes": null
+                "routes": null,
+                "networks": null
               }
             ]
           },
@@ -359,6 +369,9 @@ func TestParsingToJSON(t *testing.T) {
           "routes": [
             "10.10.0.0/24"
           ],
+          "networks": [
+            "10.10.0.0/24"
+          ],
           "dnsServers": [
             {
               "servers": [
@@ -418,6 +431,8 @@ func TestParsingToYAML(t *testing.T) {
           quantumResistance: false
           routes:
             - 10.1.0.0/24
+          networks:
+            - 10.1.0.0/24
         - fqdn: peer-2.awesome-domain.com
           netbirdIp: 192.168.178.102
           publicKey: Pubkey2
@@ -437,6 +452,7 @@ func TestParsingToYAML(t *testing.T) {
           latency: 10ms
           quantumResistance: false
           routes: []
+          networks: []
 cliVersion: development
 daemonVersion: 0.14.1
 management:
@@ -465,6 +481,8 @@ quantumResistance: false
 quantumResistancePermissive: false
 routes:
     - 10.10.0.0/24
+networks:
+    - 10.10.0.0/24
 dnsServers:
     - servers:
         - 8.8.8.8:53
@@ -509,6 +527,7 @@ func TestParsingToDetail(t *testing.T) {
   Transfer status (received/sent) 200 B/100 B
   Quantum resistance: false
   Routes: 10.1.0.0/24
+  Networks: 10.1.0.0/24
   Latency: 10ms
 
  peer-2.awesome-domain.com:
@@ -525,6 +544,7 @@ func TestParsingToDetail(t *testing.T) {
   Transfer status (received/sent) 2.0 KiB/1000 B
   Quantum resistance: false
   Routes: -
+  Networks: -
   Latency: 10ms
 
 OS: %s/%s
@@ -543,6 +563,7 @@ NetBird IP: 192.168.178.100/16
 Interface type: Kernel
 Quantum resistance: false
 Routes: 10.10.0.0/24
+Networks: 10.10.0.0/24
 Peers count: 2/2 Connected
 `, lastConnectionUpdate1, lastHandshake1, lastConnectionUpdate2, lastHandshake2, runtime.GOOS, runtime.GOARCH, overview.CliVersion)
 
@@ -564,6 +585,7 @@ NetBird IP: 192.168.178.100/16
 Interface type: Kernel
 Quantum resistance: false
 Routes: 10.10.0.0/24
+Networks: 10.10.0.0/24
 Peers count: 2/2 Connected
 `
 

client/cmd/testutil_test.go

@@ -10,6 +10,8 @@ import (
 	"go.opentelemetry.io/otel"
 
 	"github.com/netbirdio/netbird/management/server/activity"
+	"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/util"
@@ -71,7 +73,7 @@ func startManagement(t *testing.T, config *mgmt.Config, testFile string) (*grpc.
 		t.Fatal(err)
 	}
 	s := grpc.NewServer()
-	store, cleanUp, err := mgmt.NewTestStoreFromSQL(context.Background(), testFile, t.TempDir())
+	store, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), testFile, t.TempDir())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -93,7 +95,7 @@ func startManagement(t *testing.T, config *mgmt.Config, testFile string) (*grpc.
 	}
 
 	secretsManager := mgmt.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay)
-	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, peersUpdateManager, secretsManager, nil, nil)
+	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settings.NewManager(store), peersUpdateManager, secretsManager, nil, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

client/firewall/iptables/acl_linux.go

@@ -332,18 +332,12 @@ func (m *aclManager) createDefaultChains() error {
 
 // The OUTPUT chain gets an extra rule to allow traffic to any set up routes, the return traffic is handled by the INPUT related/established rule.
 func (m *aclManager) seedInitialEntries() {
-
 	established := getConntrackEstablished()
 
 	m.appendToEntries("INPUT", []string{"-i", m.wgIface.Name(), "-j", "DROP"})
 	m.appendToEntries("INPUT", []string{"-i", m.wgIface.Name(), "-j", chainNameInputRules})
 	m.appendToEntries("INPUT", append([]string{"-i", m.wgIface.Name()}, established...))
 
-	m.appendToEntries("OUTPUT", []string{"-o", m.wgIface.Name(), "-j", "DROP"})
-	m.appendToEntries("OUTPUT", []string{"-o", m.wgIface.Name(), "-j", chainNameOutputRules})
-	m.appendToEntries("OUTPUT", []string{"-o", m.wgIface.Name(), "!", "-d", m.wgIface.Address().String(), "-j", "ACCEPT"})
-	m.appendToEntries("OUTPUT", append([]string{"-o", m.wgIface.Name()}, established...))
-
 	m.appendToEntries("FORWARD", []string{"-i", m.wgIface.Name(), "-j", "DROP"})
 	m.appendToEntries("FORWARD", []string{"-i", m.wgIface.Name(), "-j", m.routingFwChainName})
 	m.appendToEntries("FORWARD", append([]string{"-o", m.wgIface.Name()}, established...))
@@ -352,14 +346,14 @@ func (m *aclManager) seedInitialEntries() {
 func (m *aclManager) seedInitialOptionalEntries() {
 	m.optionalEntries["FORWARD"] = []entry{
 		{
-			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark), "-j", chainNameInputRules},
+			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected), "-j", chainNameInputRules},
 			position: 2,
 		},
 	}
 
 	m.optionalEntries["PREROUTING"] = []entry{
 		{
-			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark)},
+			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected)},
 			position: 1,
 		},
 	}

client/firewall/iptables/manager_linux.go

@@ -83,9 +83,11 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early to ensure cleanup of chains
-	if err := stateManager.PersistState(context.Background()); err != nil {
-		log.Errorf("failed to persist state: %v", err)
-	}
+	go func() {
+		if err := stateManager.PersistState(context.Background()); err != nil {
+			log.Errorf("failed to persist state: %v", err)
+		}
+	}()
 
 	return nil
 }
@@ -205,24 +207,19 @@ func (m *Manager) AllowNetbird() error {
 		"",
 	)
 	if err != nil {
-		return fmt.Errorf("failed to allow netbird interface traffic: %w", err)
+		return fmt.Errorf("allow netbird interface traffic: %w", err)
 	}
-	_, err = m.AddPeerFiltering(
-		net.ParseIP("0.0.0.0"),
-		"all",
-		nil,
-		nil,
-		firewall.RuleDirectionOUT,
-		firewall.ActionAccept,
-		"",
-		"",
-	)
-	return err
+	return nil
 }
 
 // Flush doesn't need to be implemented for this manager
 func (m *Manager) Flush() error { return nil }
 
+// CollectStats returns connection tracking statistics
+func (m *Manager) CollectStats() []*firewall.FlowStats {
+	return nil
+}
+
 func getConntrackEstablished() []string {
 	return []string{"-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT"}
 }

client/firewall/iptables/router_linux.go

@@ -18,22 +18,24 @@ import (
 	"github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
-)
-
-const (
-	ipv4Nat = "netbird-rt-nat"
+	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
 // constants needed to manage and create iptable rules
 const (
 	tableFilter             = "filter"
 	tableNat                = "nat"
+	tableMangle             = "mangle"
 	chainPOSTROUTING        = "POSTROUTING"
+	chainPREROUTING         = "PREROUTING"
 	chainRTNAT              = "NETBIRD-RT-NAT"
 	chainRTFWD              = "NETBIRD-RT-FWD"
+	chainRTPRE              = "NETBIRD-RT-PRE"
 	routingFinalForwardJump = "ACCEPT"
 	routingFinalNatJump     = "MASQUERADE"
 
+	jumpPre  = "jump-pre"
+	jumpNat  = "jump-nat"
 	matchSet = "--match-set"
 )
 
@@ -323,24 +325,25 @@ func (r *router) Reset() error {
 }
 
 func (r *router) cleanUpDefaultForwardRules() error {
-	err := r.cleanJumpRules()
-	if err != nil {
-		return err
+	if err := r.cleanJumpRules(); err != nil {
+		return fmt.Errorf("clean jump rules: %w", err)
 	}
 
 	log.Debug("flushing routing related tables")
-	for _, chain := range []string{chainRTFWD, chainRTNAT} {
-		table := r.getTableForChain(chain)
-
-		ok, err := r.iptablesClient.ChainExists(table, chain)
+	for _, chainInfo := range []struct {
+		chain string
+		table string
+	}{
+		{chainRTFWD, tableFilter},
+		{chainRTNAT, tableNat},
+		{chainRTPRE, tableMangle},
+	} {
+		ok, err := r.iptablesClient.ChainExists(chainInfo.table, chainInfo.chain)
 		if err != nil {
-			log.Errorf("failed check chain %s, error: %v", chain, err)
-			return err
+			return fmt.Errorf("check chain %s in table %s: %w", chainInfo.chain, chainInfo.table, err)
 		} else if ok {
-			err = r.iptablesClient.ClearAndDeleteChain(table, chain)
-			if err != nil {
-				log.Errorf("failed cleaning chain %s, error: %v", chain, err)
-				return err
+			if err = r.iptablesClient.ClearAndDeleteChain(chainInfo.table, chainInfo.chain); err != nil {
+				return fmt.Errorf("clear and delete chain %s in table %s: %w", chainInfo.chain, chainInfo.table, err)
 			}
 		}
 	}
@@ -349,9 +352,16 @@ func (r *router) cleanUpDefaultForwardRules() error {
 }
 
 func (r *router) createContainers() error {
-	for _, chain := range []string{chainRTFWD, chainRTNAT} {
-		if err := r.createAndSetupChain(chain); err != nil {
-			return fmt.Errorf("create chain %s: %w", chain, err)
+	for _, chainInfo := range []struct {
+		chain string
+		table string
+	}{
+		{chainRTFWD, tableFilter},
+		{chainRTPRE, tableMangle},
+		{chainRTNAT, tableNat},
+	} {
+		if err := r.createAndSetupChain(chainInfo.chain); err != nil {
+			return fmt.Errorf("create chain %s in table %s: %w", chainInfo.chain, chainInfo.table, err)
 		}
 	}
 
@@ -359,6 +369,10 @@ func (r *router) createContainers() error {
 		return fmt.Errorf("insert established rule: %w", err)
 	}
 
+	if err := r.addPostroutingRules(); err != nil {
+		return fmt.Errorf("add static nat rules: %w", err)
+	}
+
 	if err := r.addJumpRules(); err != nil {
 		return fmt.Errorf("add jump rules: %w", err)
 	}
@@ -366,6 +380,32 @@ func (r *router) createContainers() error {
 	return nil
 }
 
+func (r *router) addPostroutingRules() error {
+	// First rule for outbound masquerade
+	rule1 := []string{
+		"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasquerade),
+		"!", "-o", "lo",
+		"-j", routingFinalNatJump,
+	}
+	if err := r.iptablesClient.Append(tableNat, chainRTNAT, rule1...); err != nil {
+		return fmt.Errorf("add outbound masquerade rule: %v", err)
+	}
+	r.rules["static-nat-outbound"] = rule1
+
+	// Second rule for return traffic masquerade
+	rule2 := []string{
+		"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasqueradeReturn),
+		"-o", r.wgIface.Name(),
+		"-j", routingFinalNatJump,
+	}
+	if err := r.iptablesClient.Append(tableNat, chainRTNAT, rule2...); err != nil {
+		return fmt.Errorf("add return masquerade rule: %v", err)
+	}
+	r.rules["static-nat-return"] = rule2
+
+	return nil
+}
+
 func (r *router) createAndSetupChain(chain string) error {
 	table := r.getTableForChain(chain)
 
@@ -377,10 +417,14 @@ func (r *router) createAndSetupChain(chain string) error {
 }
 
 func (r *router) getTableForChain(chain string) string {
-	if chain == chainRTNAT {
+	switch chain {
+	case chainRTNAT:
 		return tableNat
+	case chainRTPRE:
+		return tableMangle
+	default:
+		return tableFilter
 	}
-	return tableFilter
 }
 
 func (r *router) insertEstablishedRule(chain string) error {
@@ -398,25 +442,39 @@ func (r *router) insertEstablishedRule(chain string) error {
 }
 
 func (r *router) addJumpRules() error {
-	rule := []string{"-j", chainRTNAT}
-	err := r.iptablesClient.Insert(tableNat, chainPOSTROUTING, 1, rule...)
-	if err != nil {
-		return err
+	// Jump to NAT chain
+	natRule := []string{"-j", chainRTNAT}
+	if err := r.iptablesClient.Insert(tableNat, chainPOSTROUTING, 1, natRule...); err != nil {
+		return fmt.Errorf("add nat jump rule: %v", err)
 	}
-	r.rules[ipv4Nat] = rule
+	r.rules[jumpNat] = natRule
+
+	// Jump to prerouting chain
+	preRule := []string{"-j", chainRTPRE}
+	if err := r.iptablesClient.Insert(tableMangle, chainPREROUTING, 1, preRule...); err != nil {
+		return fmt.Errorf("add prerouting jump rule: %v", err)
+	}
+	r.rules[jumpPre] = preRule
 
 	return nil
 }
 
 func (r *router) cleanJumpRules() error {
-	rule, found := r.rules[ipv4Nat]
-	if found {
-		err := r.iptablesClient.DeleteIfExists(tableNat, chainPOSTROUTING, rule...)
-		if err != nil {
-			return fmt.Errorf("failed cleaning rule from chain %s, err: %v", chainPOSTROUTING, err)
+	for _, ruleKey := range []string{jumpNat, jumpPre} {
+		if rule, exists := r.rules[ruleKey]; exists {
+			table := tableNat
+			chain := chainPOSTROUTING
+			if ruleKey == jumpPre {
+				table = tableMangle
+				chain = chainPREROUTING
+			}
+
+			if err := r.iptablesClient.DeleteIfExists(table, chain, rule...); err != nil {
+				return fmt.Errorf("delete rule from chain %s in table %s, err: %v", chain, table, err)
+			}
+			delete(r.rules, ruleKey)
 		}
 	}
-
 	return nil
 }
 
@@ -424,19 +482,35 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
 
 	if rule, exists := r.rules[ruleKey]; exists {
-		if err := r.iptablesClient.DeleteIfExists(tableNat, chainRTNAT, rule...); err != nil {
-			return fmt.Errorf("error while removing existing NAT rule for %s: %v", pair.Destination, err)
+		if err := r.iptablesClient.DeleteIfExists(tableMangle, chainRTPRE, rule...); err != nil {
+			return fmt.Errorf("error while removing existing marking rule for %s: %v", pair.Destination, err)
 		}
 		delete(r.rules, ruleKey)
 	}
 
-	rule := genRuleSpec(routingFinalNatJump, pair.Source, pair.Destination, r.wgIface.Name(), pair.Inverse)
-	if err := r.iptablesClient.Append(tableNat, chainRTNAT, rule...); err != nil {
-		return fmt.Errorf("error while appending new NAT rule for %s: %v", pair.Destination, err)
+	markValue := nbnet.PreroutingFwmarkMasquerade
+	if pair.Inverse {
+		markValue = nbnet.PreroutingFwmarkMasqueradeReturn
+	}
+
+	rule := []string{"-i", r.wgIface.Name()}
+	if pair.Inverse {
+		rule = []string{"!", "-i", r.wgIface.Name()}
+	}
+
+	rule = append(rule,
+		"-m", "conntrack",
+		"--ctstate", "NEW",
+		"-s", pair.Source.String(),
+		"-d", pair.Destination.String(),
+		"-j", "MARK", "--set-mark", fmt.Sprintf("%#x", markValue),
+	)
+
+	if err := r.iptablesClient.Append(tableMangle, chainRTPRE, rule...); err != nil {
+		return fmt.Errorf("error while adding marking rule for %s: %v", pair.Destination, err)
 	}
 
 	r.rules[ruleKey] = rule
-
 	return nil
 }
 
@@ -444,13 +518,12 @@ func (r *router) removeNatRule(pair firewall.RouterPair) error {
 	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
 
 	if rule, exists := r.rules[ruleKey]; exists {
-		if err := r.iptablesClient.DeleteIfExists(tableNat, chainRTNAT, rule...); err != nil {
-			return fmt.Errorf("error while removing existing nat rule for %s: %v", pair.Destination, err)
+		if err := r.iptablesClient.DeleteIfExists(tableMangle, chainRTPRE, rule...); err != nil {
+			return fmt.Errorf("error while removing marking rule for %s: %v", pair.Destination, err)
 		}
-
 		delete(r.rules, ruleKey)
 	} else {
-		log.Debugf("nat rule %s not found", ruleKey)
+		log.Debugf("marking rule %s not found", ruleKey)
 	}
 
 	return nil
@@ -482,16 +555,6 @@ func (r *router) updateState() {
 	}
 }
 
-func genRuleSpec(jump string, source, destination netip.Prefix, intf string, inverse bool) []string {
-	intdir := "-i"
-	lointdir := "-o"
-	if inverse {
-		intdir = "-o"
-		lointdir = "-i"
-	}
-	return []string{intdir, intf, "!", lointdir, "lo", "-s", source.String(), "-d", destination.String(), "-j", jump}
-}
-
 func genRouteFilteringRuleSpec(params routeFilteringRuleParams) []string {
 	var rule []string
 

client/firewall/iptables/router_linux_test.go

@@ -3,17 +3,18 @@
 package iptables
 
 import (
+	"fmt"
 	"net/netip"
 	"os/exec"
 	"testing"
 
 	"github.com/coreos/go-iptables/iptables"
-	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/firewall/test"
+	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
 func isIptablesSupported() bool {
@@ -34,14 +35,24 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 	require.NoError(t, manager.init(nil))
 
 	defer func() {
-		_ = manager.Reset()
+		assert.NoError(t, manager.Reset(), "shouldn't return error")
 	}()
 
-	require.Len(t, manager.rules, 2, "should have created rules map")
+	// Now 5 rules:
+	// 1. established rule in forward chain
+	// 2. jump rule to NAT chain
+	// 3. jump rule to PRE chain
+	// 4. static outbound masquerade rule
+	// 5. static return masquerade rule
+	require.Len(t, manager.rules, 5, "should have created rules map")
 
-	exists, err := manager.iptablesClient.Exists(tableNat, chainPOSTROUTING, manager.rules[ipv4Nat]...)
+	exists, err := manager.iptablesClient.Exists(tableNat, chainPOSTROUTING, "-j", chainRTNAT)
 	require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainPOSTROUTING)
-	require.True(t, exists, "postrouting rule should exist")
+	require.True(t, exists, "postrouting jump rule should exist")
+
+	exists, err = manager.iptablesClient.Exists(tableMangle, chainPREROUTING, "-j", chainRTPRE)
+	require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainPREROUTING)
+	require.True(t, exists, "prerouting jump rule should exist")
 
 	pair := firewall.RouterPair{
 		ID:          "abc",
@@ -49,22 +60,15 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 		Destination: netip.MustParsePrefix("100.100.100.0/24"),
 		Masquerade:  true,
 	}
-	forward4Rule := []string{"-s", pair.Source.String(), "-d", pair.Destination.String(), "-j", routingFinalForwardJump}
 
-	err = manager.iptablesClient.Insert(tableFilter, chainRTFWD, 1, forward4Rule...)
-	require.NoError(t, err, "inserting rule should not return error")
-
-	nat4Rule := genRuleSpec(routingFinalNatJump, pair.Source, pair.Destination, ifaceMock.Name(), false)
-
-	err = manager.iptablesClient.Insert(tableNat, chainRTNAT, 1, nat4Rule...)
-	require.NoError(t, err, "inserting rule should not return error")
+	err = manager.AddNatRule(pair)
+	require.NoError(t, err, "adding NAT rule should not return error")
 
 	err = manager.Reset()
 	require.NoError(t, err, "shouldn't return error")
 }
 
 func TestIptablesManager_AddNatRule(t *testing.T) {
-
 	if !isIptablesSupported() {
 		t.SkipNow()
 	}
@@ -79,52 +83,66 @@ func TestIptablesManager_AddNatRule(t *testing.T) {
 			require.NoError(t, manager.init(nil))
 
 			defer func() {
-				err := manager.Reset()
-				if err != nil {
-					log.Errorf("failed to reset iptables manager: %s", err)
-				}
+				assert.NoError(t, manager.Reset(), "shouldn't return error")
 			}()
 
 			err = manager.AddNatRule(testCase.InputPair)
-			require.NoError(t, err, "forwarding pair should be inserted")
+			require.NoError(t, err, "marking rule should be inserted")
 
 			natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
-			natRule := genRuleSpec(routingFinalNatJump, testCase.InputPair.Source, testCase.InputPair.Destination, ifaceMock.Name(), false)
-
-			exists, err := iptablesClient.Exists(tableNat, chainRTNAT, natRule...)
-			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainRTNAT)
-			if testCase.InputPair.Masquerade {
-				require.True(t, exists, "nat rule should be created")
-				foundNatRule, foundNat := manager.rules[natRuleKey]
-				require.True(t, foundNat, "nat rule should exist in the map")
-				require.Equal(t, natRule[:4], foundNatRule[:4], "stored nat rule should match")
-			} else {
-				require.False(t, exists, "nat rule should not be created")
-				_, foundNat := manager.rules[natRuleKey]
-				require.False(t, foundNat, "nat rule should not exist in the map")
+			markingRule := []string{
+				"-i", ifaceMock.Name(),
+				"-m", "conntrack",
+				"--ctstate", "NEW",
+				"-s", testCase.InputPair.Source.String(),
+				"-d", testCase.InputPair.Destination.String(),
+				"-j", "MARK", "--set-mark",
+				fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasquerade),
 			}
 
-			inNatRuleKey := firewall.GenKey(firewall.NatFormat, firewall.GetInversePair(testCase.InputPair))
-			inNatRule := genRuleSpec(routingFinalNatJump, firewall.GetInversePair(testCase.InputPair).Source, firewall.GetInversePair(testCase.InputPair).Destination, ifaceMock.Name(), true)
-
-			exists, err = iptablesClient.Exists(tableNat, chainRTNAT, inNatRule...)
-			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainRTNAT)
+			exists, err := iptablesClient.Exists(tableMangle, chainRTPRE, markingRule...)
+			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainRTPRE)
 			if testCase.InputPair.Masquerade {
-				require.True(t, exists, "income nat rule should be created")
-				foundNatRule, foundNat := manager.rules[inNatRuleKey]
-				require.True(t, foundNat, "income nat rule should exist in the map")
-				require.Equal(t, inNatRule[:4], foundNatRule[:4], "stored income nat rule should match")
+				require.True(t, exists, "marking rule should be created")
+				foundRule, found := manager.rules[natRuleKey]
+				require.True(t, found, "marking rule should exist in the map")
+				require.Equal(t, markingRule, foundRule, "stored marking rule should match")
 			} else {
-				require.False(t, exists, "nat rule should not be created")
-				_, foundNat := manager.rules[inNatRuleKey]
-				require.False(t, foundNat, "income nat rule should not exist in the map")
+				require.False(t, exists, "marking rule should not be created")
+				_, found := manager.rules[natRuleKey]
+				require.False(t, found, "marking rule should not exist in the map")
+			}
+
+			// Check inverse rule
+			inversePair := firewall.GetInversePair(testCase.InputPair)
+			inverseRuleKey := firewall.GenKey(firewall.NatFormat, inversePair)
+			inverseMarkingRule := []string{
+				"!", "-i", ifaceMock.Name(),
+				"-m", "conntrack",
+				"--ctstate", "NEW",
+				"-s", inversePair.Source.String(),
+				"-d", inversePair.Destination.String(),
+				"-j", "MARK", "--set-mark",
+				fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasqueradeReturn),
+			}
+
+			exists, err = iptablesClient.Exists(tableMangle, chainRTPRE, inverseMarkingRule...)
+			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainRTPRE)
+			if testCase.InputPair.Masquerade {
+				require.True(t, exists, "inverse marking rule should be created")
+				foundRule, found := manager.rules[inverseRuleKey]
+				require.True(t, found, "inverse marking rule should exist in the map")
+				require.Equal(t, inverseMarkingRule, foundRule, "stored inverse marking rule should match")
+			} else {
+				require.False(t, exists, "inverse marking rule should not be created")
+				_, found := manager.rules[inverseRuleKey]
+				require.False(t, found, "inverse marking rule should not exist in the map")
 			}
 		})
 	}
 }
 
 func TestIptablesManager_RemoveNatRule(t *testing.T) {
-
 	if !isIptablesSupported() {
 		t.SkipNow()
 	}
@@ -137,42 +155,52 @@ func TestIptablesManager_RemoveNatRule(t *testing.T) {
 			require.NoError(t, err, "shouldn't return error")
 			require.NoError(t, manager.init(nil))
 			defer func() {
-				_ = manager.Reset()
+				assert.NoError(t, manager.Reset(), "shouldn't return error")
 			}()
 
-			require.NoError(t, err, "shouldn't return error")
-
-			natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
-			natRule := genRuleSpec(routingFinalNatJump, testCase.InputPair.Source, testCase.InputPair.Destination, ifaceMock.Name(), false)
-
-			err = iptablesClient.Insert(tableNat, chainRTNAT, 1, natRule...)
-			require.NoError(t, err, "inserting rule should not return error")
-
-			inNatRuleKey := firewall.GenKey(firewall.NatFormat, firewall.GetInversePair(testCase.InputPair))
-			inNatRule := genRuleSpec(routingFinalNatJump, firewall.GetInversePair(testCase.InputPair).Source, firewall.GetInversePair(testCase.InputPair).Destination, ifaceMock.Name(), true)
-
-			err = iptablesClient.Insert(tableNat, chainRTNAT, 1, inNatRule...)
-			require.NoError(t, err, "inserting rule should not return error")
-
-			err = manager.Reset()
-			require.NoError(t, err, "shouldn't return error")
+			err = manager.AddNatRule(testCase.InputPair)
+			require.NoError(t, err, "should add NAT rule without error")
 
 			err = manager.RemoveNatRule(testCase.InputPair)
 			require.NoError(t, err, "shouldn't return error")
 
-			exists, err := iptablesClient.Exists(tableNat, chainRTNAT, natRule...)
-			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainRTNAT)
-			require.False(t, exists, "nat rule should not exist")
+			natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
+			markingRule := []string{
+				"-i", ifaceMock.Name(),
+				"-m", "conntrack",
+				"--ctstate", "NEW",
+				"-s", testCase.InputPair.Source.String(),
+				"-d", testCase.InputPair.Destination.String(),
+				"-j", "MARK", "--set-mark",
+				fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasquerade),
+			}
+
+			exists, err := iptablesClient.Exists(tableMangle, chainRTPRE, markingRule...)
+			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainRTPRE)
+			require.False(t, exists, "marking rule should not exist")
 
 			_, found := manager.rules[natRuleKey]
-			require.False(t, found, "nat rule should exist in the manager map")
+			require.False(t, found, "marking rule should not exist in the manager map")
 
-			exists, err = iptablesClient.Exists(tableNat, chainRTNAT, inNatRule...)
-			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainRTNAT)
-			require.False(t, exists, "income nat rule should not exist")
+			// Check inverse rule removal
+			inversePair := firewall.GetInversePair(testCase.InputPair)
+			inverseRuleKey := firewall.GenKey(firewall.NatFormat, inversePair)
+			inverseMarkingRule := []string{
+				"!", "-i", ifaceMock.Name(),
+				"-m", "conntrack",
+				"--ctstate", "NEW",
+				"-s", inversePair.Source.String(),
+				"-d", inversePair.Destination.String(),
+				"-j", "MARK", "--set-mark",
+				fmt.Sprintf("%#x", nbnet.PreroutingFwmarkMasqueradeReturn),
+			}
 
-			_, found = manager.rules[inNatRuleKey]
-			require.False(t, found, "income nat rule should exist in the manager map")
+			exists, err = iptablesClient.Exists(tableMangle, chainRTPRE, inverseMarkingRule...)
+			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainRTPRE)
+			require.False(t, exists, "inverse marking rule should not exist")
+
+			_, found = manager.rules[inverseRuleKey]
+			require.False(t, found, "inverse marking rule should not exist in the map")
 		})
 	}
 }

client/firewall/iptables/rulestore_linux.go

@@ -37,6 +37,11 @@ func (s *ipList) UnmarshalJSON(data []byte) error {
 		return err
 	}
 	s.ips = temp.IPs
+
+	if temp.IPs == nil {
+		temp.IPs = make(map[string]struct{})
+	}
+
 	return nil
 }
 
@@ -89,5 +94,10 @@ func (s *ipsetStore) UnmarshalJSON(data []byte) error {
 		return err
 	}
 	s.ipsets = temp.IPSets
+
+	if temp.IPSets == nil {
+		temp.IPSets = make(map[string]*ipList)
+	}
+
 	return nil
 }

client/firewall/manager/firewall.go

@@ -17,6 +17,7 @@ import (
 const (
 	ForwardingFormatPrefix = "netbird-fwd-"
 	ForwardingFormat       = "netbird-fwd-%s-%t"
+	PreroutingFormat       = "netbird-prerouting-%s-%t"
 	NatFormat              = "netbird-nat-%s-%t"
 )
 
@@ -99,6 +100,9 @@ type Manager interface {
 
 	// Flush the changes to firewall controller
 	Flush() error
+
+	// CollectStats returns the statistics of the firewall manager
+	CollectStats() []*FlowStats
 }
 
 func GenKey(format string, pair RouterPair) string {

client/firewall/manager/stats.go

@@ -0,0 +1,107 @@
+package manager
+
+import (
+	"encoding/json"
+	"net"
+	"slices"
+	"strconv"
+	"sync/atomic"
+	"time"
+)
+
+const (
+	DirectionInbound  Direction = 0
+	DirectionOutbound Direction = 1
+)
+
+type Direction uint8
+
+func (d Direction) String() string {
+	switch d {
+	case DirectionInbound:
+		return "inbound"
+	case DirectionOutbound:
+		return "outbound"
+	default:
+		return "unknown"
+	}
+}
+
+// FlowStats tracks statistics for an individual connection
+type FlowStats struct {
+	StartTime  time.Time
+	LastSeen   time.Time
+	BytesIn    atomic.Uint64
+	BytesOut   atomic.Uint64
+	PacketsIn  atomic.Uint64
+	PacketsOut atomic.Uint64
+	Protocol   uint8
+	Direction  Direction
+	SourceIP   net.IP
+	DestIP     net.IP
+	SourcePort uint16
+	DestPort   uint16
+}
+
+func (f *FlowStats) Clone() *FlowStats {
+	flowCopy := FlowStats{
+		StartTime:  f.StartTime,
+		LastSeen:   f.LastSeen,
+		Protocol:   f.Protocol,
+		Direction:  f.Direction,
+		SourceIP:   slices.Clone(f.SourceIP),
+		DestIP:     slices.Clone(f.DestIP),
+		SourcePort: f.SourcePort,
+		DestPort:   f.DestPort,
+	}
+	flowCopy.BytesIn.Store(f.BytesIn.Load())
+	flowCopy.BytesOut.Store(f.BytesOut.Load())
+	flowCopy.PacketsIn.Store(f.PacketsIn.Load())
+	flowCopy.PacketsOut.Store(f.PacketsOut.Load())
+
+	return &flowCopy
+}
+
+// MarshalJSON implements json.Marshaler interface
+func (f *FlowStats) MarshalJSON() ([]byte, error) {
+	return json.Marshal(&struct {
+		StartTime  time.Time
+		LastSeen   time.Time
+		BytesIn    uint64
+		BytesOut   uint64
+		PacketsIn  uint64
+		PacketsOut uint64
+		Protocol   Protocol
+		Direction  string
+		SourceIP   net.IP
+		DestIP     net.IP
+		SourcePort uint16
+		DestPort   uint16
+	}{
+		StartTime:  f.StartTime,
+		LastSeen:   f.LastSeen,
+		BytesIn:    f.BytesIn.Load(),
+		BytesOut:   f.BytesOut.Load(),
+		PacketsIn:  f.PacketsIn.Load(),
+		PacketsOut: f.PacketsOut.Load(),
+		Protocol:   protoFromInt(f.Protocol),
+		Direction:  f.Direction.String(),
+		SourceIP:   f.SourceIP,
+		DestIP:     f.DestIP,
+		SourcePort: f.SourcePort,
+		DestPort:   f.DestPort,
+	})
+}
+
+func protoFromInt(p uint8) Protocol {
+	switch p {
+	case 6:
+		return ProtocolTCP
+	case 17:
+		return ProtocolUDP
+	case 1:
+		return ProtocolICMP
+	default:
+		return Protocol(strconv.Itoa(int(p)))
+	}
+}

client/firewall/nftables/acl_linux.go

@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
-	"net/netip"
 	"strconv"
 	"strings"
 	"time"
@@ -28,7 +27,6 @@ const (
 
 	// filter chains contains the rules that jump to the rules chains
 	chainNameInputFilter   = "netbird-acl-input-filter"
-	chainNameOutputFilter  = "netbird-acl-output-filter"
 	chainNameForwardFilter = "netbird-acl-forward-filter"
 	chainNamePrerouting    = "netbird-rt-prerouting"
 
@@ -441,18 +439,6 @@ func (m *AclManager) createDefaultChains() (err error) {
 		return err
 	}
 
-	// netbird-acl-output-filter
-	// type filter hook output priority filter; policy accept;
-	chain = m.createFilterChainWithHook(chainNameOutputFilter, nftables.ChainHookOutput)
-	m.addFwdAllow(chain, expr.MetaKeyOIFNAME)
-	m.addJumpRule(chain, m.chainOutputRules.Name, expr.MetaKeyOIFNAME) // to netbird-acl-output-rules
-	m.addDropExpressions(chain, expr.MetaKeyOIFNAME)
-	err = m.rConn.Flush()
-	if err != nil {
-		log.Debugf("failed to create chain (%s): %s", chainNameOutputFilter, err)
-		return err
-	}
-
 	// netbird-acl-forward-filter
 	chainFwFilter := m.createFilterChainWithHook(chainNameForwardFilter, nftables.ChainHookForward)
 	m.addJumpRulesToRtForward(chainFwFilter) // to netbird-rt-fwd
@@ -520,7 +506,7 @@ func (m *AclManager) addPreroutingRule(preroutingChain *nftables.Chain) {
 			},
 			&expr.Immediate{
 				Register: 1,
-				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
 			},
 			&expr.Meta{
 				Key:            expr.MetaKeyMARK,
@@ -543,7 +529,7 @@ func (m *AclManager) addFwmarkToForward(chainFwFilter *nftables.Chain) {
 			&expr.Cmp{
 				Op:       expr.CmpOpEq,
 				Register: 1,
-				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
 			},
 			&expr.Verdict{
 				Kind:  expr.VerdictJump,
@@ -619,45 +605,6 @@ func (m *AclManager) addDropExpressions(chain *nftables.Chain, ifaceKey expr.Met
 	return nil
 }
 
-func (m *AclManager) addFwdAllow(chain *nftables.Chain, iifname expr.MetaKey) {
-	ip, _ := netip.AddrFromSlice(m.wgIface.Address().Network.IP.To4())
-	dstOp := expr.CmpOpNeq
-	expressions := []expr.Any{
-		&expr.Meta{Key: iifname, Register: 1},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 1,
-			Data:     ifname(m.wgIface.Name()),
-		},
-		&expr.Payload{
-			DestRegister: 2,
-			Base:         expr.PayloadBaseNetworkHeader,
-			Offset:       16,
-			Len:          4,
-		},
-		&expr.Bitwise{
-			SourceRegister: 2,
-			DestRegister:   2,
-			Len:            4,
-			Xor:            []byte{0x0, 0x0, 0x0, 0x0},
-			Mask:           m.wgIface.Address().Network.Mask,
-		},
-		&expr.Cmp{
-			Op:       dstOp,
-			Register: 2,
-			Data:     ip.Unmap().AsSlice(),
-		},
-		&expr.Verdict{
-			Kind: expr.VerdictAccept,
-		},
-	}
-	_ = m.rConn.AddRule(&nftables.Rule{
-		Table: chain.Table,
-		Chain: chain,
-		Exprs: expressions,
-	})
-}
-
 func (m *AclManager) addJumpRule(chain *nftables.Chain, to string, ifaceKey expr.MetaKey) {
 	expressions := []expr.Any{
 		&expr.Meta{Key: ifaceKey, Register: 1},

client/firewall/nftables/manager_linux.go

@@ -99,9 +99,11 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early
-	if err := stateManager.PersistState(context.Background()); err != nil {
-		log.Errorf("failed to persist state: %v", err)
-	}
+	go func() {
+		if err := stateManager.PersistState(context.Background()); err != nil {
+			log.Errorf("failed to persist state: %v", err)
+		}
+	}()
 
 	return nil
 }
@@ -197,7 +199,7 @@ func (m *Manager) AllowNetbird() error {
 
 	var chain *nftables.Chain
 	for _, c := range chains {
-		if c.Table.Name == tableNameFilter && c.Name == chainNameForward {
+		if c.Table.Name == tableNameFilter && c.Name == chainNameInput {
 			chain = c
 			break
 		}
@@ -274,7 +276,7 @@ func (m *Manager) resetNetbirdInputRules() error {
 
 func (m *Manager) deleteNetbirdInputRules(chains []*nftables.Chain) {
 	for _, c := range chains {
-		if c.Table.Name == "filter" && c.Name == "INPUT" {
+		if c.Table.Name == tableNameFilter && c.Name == chainNameInput {
 			rules, err := m.rConn.GetRules(c.Table, c)
 			if err != nil {
 				log.Errorf("get rules for chain %q: %v", c.Name, err)
@@ -321,6 +323,11 @@ func (m *Manager) Flush() error {
 	return m.aclManager.Flush()
 }
 
+// CollectStats returns connection tracking statistics
+func (m *Manager) CollectStats() []*firewall.FlowStats {
+	return nil
+}
+
 func (m *Manager) createWorkTable() (*nftables.Table, error) {
 	tables, err := m.rConn.ListTablesOfFamily(nftables.TableFamilyIPv4)
 	if err != nil {
@@ -349,7 +356,9 @@ func (m *Manager) applyAllowNetbirdRules(chain *nftables.Chain) {
 				Register: 1,
 				Data:     ifname(m.wgIface.Name()),
 			},
-			&expr.Verdict{},
+			&expr.Verdict{
+				Kind: expr.VerdictAccept,
+			},
 		},
 		UserData: []byte(allowNetbirdInputRuleID),
 	}

client/firewall/nftables/manager_linux_test.go

@@ -1,9 +1,11 @@
 package nftables
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"net/netip"
+	"os/exec"
 	"testing"
 	"time"
 
@@ -225,3 +227,105 @@ func TestNFtablesCreatePerformance(t *testing.T) {
 		})
 	}
 }
+
+func runIptablesSave(t *testing.T) (string, string) {
+	t.Helper()
+	var stdout, stderr bytes.Buffer
+	cmd := exec.Command("iptables-save")
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	err := cmd.Run()
+	require.NoError(t, err, "iptables-save failed to run")
+
+	return stdout.String(), stderr.String()
+}
+
+func verifyIptablesOutput(t *testing.T, stdout, stderr string) {
+	t.Helper()
+	// Check for any incompatibility warnings
+	require.NotContains(t,
+		stderr,
+		"incompatible",
+		"iptables-save produced compatibility warning. Full stderr: %s",
+		stderr,
+	)
+
+	// Verify standard tables are present
+	expectedTables := []string{
+		"*filter",
+		"*nat",
+		"*mangle",
+	}
+
+	for _, table := range expectedTables {
+		require.Contains(t,
+			stdout,
+			table,
+			"iptables-save output missing expected table: %s\nFull stdout: %s",
+			table,
+			stdout,
+		)
+	}
+}
+
+func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
+	if check() != NFTABLES {
+		t.Skip("nftables not supported on this system")
+	}
+
+	if _, err := exec.LookPath("iptables-save"); err != nil {
+		t.Skipf("iptables-save not available on this system: %v", err)
+	}
+
+	// First ensure iptables-nft tables exist by running iptables-save
+	stdout, stderr := runIptablesSave(t)
+	verifyIptablesOutput(t, stdout, stderr)
+
+	manager, err := Create(ifaceMock)
+	require.NoError(t, err, "failed to create manager")
+	require.NoError(t, manager.Init(nil))
+
+	t.Cleanup(func() {
+		err := manager.Reset(nil)
+		require.NoError(t, err, "failed to reset manager state")
+
+		// Verify iptables output after reset
+		stdout, stderr := runIptablesSave(t)
+		verifyIptablesOutput(t, stdout, stderr)
+	})
+
+	ip := net.ParseIP("100.96.0.1")
+	_, err = manager.AddPeerFiltering(
+		ip,
+		fw.ProtocolTCP,
+		nil,
+		&fw.Port{Values: []int{80}},
+		fw.RuleDirectionIN,
+		fw.ActionAccept,
+		"",
+		"test rule",
+	)
+	require.NoError(t, err, "failed to add peer filtering rule")
+
+	_, err = manager.AddRouteFiltering(
+		[]netip.Prefix{netip.MustParsePrefix("192.168.2.0/24")},
+		netip.MustParsePrefix("10.1.0.0/24"),
+		fw.ProtocolTCP,
+		nil,
+		&fw.Port{Values: []int{443}},
+		fw.ActionAccept,
+	)
+	require.NoError(t, err, "failed to add route filtering rule")
+
+	pair := fw.RouterPair{
+		Source:      netip.MustParsePrefix("192.168.1.0/24"),
+		Destination: netip.MustParsePrefix("10.0.0.0/24"),
+		Masquerade:  true,
+	}
+	err = manager.AddNatRule(pair)
+	require.NoError(t, err, "failed to add NAT rule")
+
+	stdout, stderr = runIptablesSave(t)
+	verifyIptablesOutput(t, stdout, stderr)
+}

client/firewall/nftables/router_linux.go

@@ -21,6 +21,7 @@ import (
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
+	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
 const (
@@ -124,7 +125,6 @@ func (r *router) createContainers() error {
 	insertReturnTrafficRule(r.conn, r.workTable, r.chains[chainNameRoutingFw])
 
 	prio := *nftables.ChainPriorityNATSource - 1
-
 	r.chains[chainNameRoutingNat] = r.conn.AddChain(&nftables.Chain{
 		Name:     chainNameRoutingNat,
 		Table:    r.workTable,
@@ -133,6 +133,21 @@ func (r *router) createContainers() error {
 		Type:     nftables.ChainTypeNAT,
 	})
 
+	// Chain is created by acl manager
+	// TODO: move creation to a common place
+	r.chains[chainNamePrerouting] = &nftables.Chain{
+		Name:     chainNamePrerouting,
+		Table:    r.workTable,
+		Type:     nftables.ChainTypeFilter,
+		Hooknum:  nftables.ChainHookPrerouting,
+		Priority: nftables.ChainPriorityMangle,
+	}
+
+	// Add the single NAT rule that matches on mark
+	if err := r.addPostroutingRules(); err != nil {
+		return fmt.Errorf("add single nat rule: %v", err)
+	}
+
 	if err := r.acceptForwardRules(); err != nil {
 		log.Errorf("failed to add accept rules for the forward chain: %s", err)
 	}
@@ -422,59 +437,149 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 	sourceExp := generateCIDRMatcherExpressions(true, pair.Source)
 	destExp := generateCIDRMatcherExpressions(false, pair.Destination)
 
-	dir := expr.MetaKeyIIFNAME
-	notDir := expr.MetaKeyOIFNAME
+	op := expr.CmpOpEq
 	if pair.Inverse {
-		dir = expr.MetaKeyOIFNAME
-		notDir = expr.MetaKeyIIFNAME
+		op = expr.CmpOpNeq
 	}
 
-	lo := ifname("lo")
-	intf := ifname(r.wgIface.Name())
-
 	exprs := []expr.Any{
-		&expr.Meta{
-			Key:      dir,
+		// 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.
+		&expr.Ct{
+			Key:      expr.CtKeySTATE,
 			Register: 1,
 		},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 1,
-			Data:     intf,
-		},
-
-		// We need to exclude the loopback interface as this changes the ebpf proxy port
-		&expr.Meta{
-			Key:      notDir,
-			Register: 1,
+		&expr.Bitwise{
+			SourceRegister: 1,
+			DestRegister:   1,
+			Len:            4,
+			Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
+			Xor:            binaryutil.NativeEndian.PutUint32(0),
 		},
 		&expr.Cmp{
 			Op:       expr.CmpOpNeq,
 			Register: 1,
-			Data:     lo,
+			Data:     []byte{0, 0, 0, 0},
+		},
+
+		// interface matching
+		&expr.Meta{
+			Key:      expr.MetaKeyIIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       op,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
 		},
 	}
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
+
+	var markValue uint32 = nbnet.PreroutingFwmarkMasquerade
+	if pair.Inverse {
+		markValue = nbnet.PreroutingFwmarkMasqueradeReturn
+	}
+
 	exprs = append(exprs,
-		&expr.Counter{}, &expr.Masq{},
+		&expr.Immediate{
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(markValue),
+		},
+		&expr.Meta{
+			Key:            expr.MetaKeyMARK,
+			SourceRegister: true,
+			Register:       1,
+		},
 	)
 
-	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
+	ruleKey := firewall.GenKey(firewall.PreroutingFormat, pair)
 
 	if _, exists := r.rules[ruleKey]; exists {
 		if err := r.removeNatRule(pair); err != nil {
-			return fmt.Errorf("remove routing rule: %w", err)
+			return fmt.Errorf("remove prerouting rule: %w", err)
 		}
 	}
 
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
-		Chain:    r.chains[chainNameRoutingNat],
+		Chain:    r.chains[chainNamePrerouting],
 		Exprs:    exprs,
 		UserData: []byte(ruleKey),
 	})
+
+	return nil
+}
+
+// addPostroutingRules adds the masquerade rules
+func (r *router) addPostroutingRules() error {
+	// First masquerade rule for traffic coming in from WireGuard interface
+	exprs := []expr.Any{
+		// Match on the first fwmark
+		&expr.Meta{
+			Key:      expr.MetaKeyMARK,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkMasquerade),
+		},
+
+		// We need to exclude the loopback interface as this changes the ebpf proxy port
+		&expr.Meta{
+			Key:      expr.MetaKeyOIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpNeq,
+			Register: 1,
+			Data:     ifname("lo"),
+		},
+		&expr.Counter{},
+		&expr.Masq{},
+	}
+
+	r.conn.AddRule(&nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameRoutingNat],
+		Exprs: exprs,
+	})
+
+	// Second masquerade rule for traffic going out through WireGuard interface
+	exprs2 := []expr.Any{
+		// Match on the second fwmark
+		&expr.Meta{
+			Key:      expr.MetaKeyMARK,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkMasqueradeReturn),
+		},
+
+		// Match WireGuard interface
+		&expr.Meta{
+			Key:      expr.MetaKeyOIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
+		},
+		&expr.Counter{},
+		&expr.Masq{},
+	}
+
+	r.conn.AddRule(&nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameRoutingNat],
+		Exprs: exprs2,
+	})
+
 	return nil
 }
 
@@ -723,18 +828,18 @@ func (r *router) removeAcceptForwardRulesIptables(ipt *iptables.IPTables) error
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-// RemoveNatRule removes a nftables rule pair from nat chains
+// RemoveNatRule removes the prerouting mark rule
 func (r *router) RemoveNatRule(pair firewall.RouterPair) error {
 	if err := r.refreshRulesMap(); err != nil {
 		return fmt.Errorf(refreshRulesMapError, err)
 	}
 
 	if err := r.removeNatRule(pair); err != nil {
-		return fmt.Errorf("remove nat rule: %w", err)
+		return fmt.Errorf("remove prerouting rule: %w", err)
 	}
 
 	if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
-		return fmt.Errorf("remove inverse nat rule: %w", err)
+		return fmt.Errorf("remove inverse prerouting rule: %w", err)
 	}
 
 	if err := r.removeLegacyRouteRule(pair); err != nil {
@@ -749,21 +854,20 @@ func (r *router) RemoveNatRule(pair firewall.RouterPair) error {
 	return nil
 }
 
-// removeNatRule adds a nftables rule to the removal queue and deletes it from the rules map
 func (r *router) removeNatRule(pair firewall.RouterPair) error {
-	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
+	ruleKey := firewall.GenKey(firewall.PreroutingFormat, pair)
 
 	if rule, exists := r.rules[ruleKey]; exists {
 		err := r.conn.DelRule(rule)
 		if err != nil {
-			return fmt.Errorf("remove nat rule %s -> %s: %v", pair.Source, pair.Destination, err)
+			return fmt.Errorf("remove prerouting rule %s -> %s: %v", pair.Source, pair.Destination, err)
 		}
 
-		log.Debugf("nftables: removed nat rule %s -> %s", pair.Source, pair.Destination)
+		log.Debugf("nftables: removed prerouting rule %s -> %s", pair.Source, pair.Destination)
 
 		delete(r.rules, ruleKey)
 	} else {
-		log.Debugf("nftables: nat rule %s not found", ruleKey)
+		log.Debugf("nftables: prerouting rule %s not found", ruleKey)
 	}
 
 	return nil

client/firewall/nftables/router_linux_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/coreos/go-iptables/iptables"
 	"github.com/google/nftables"
+	"github.com/google/nftables/binaryutil"
 	"github.com/google/nftables/expr"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -32,100 +33,87 @@ func TestNftablesManager_AddNatRule(t *testing.T) {
 		t.Skip("nftables not supported on this OS")
 	}
 
-	table, err := createWorkTable()
-	require.NoError(t, err, "Failed to create work table")
-
-	defer deleteWorkTable()
-
 	for _, testCase := range test.InsertRuleTestCases {
 		t.Run(testCase.Name, func(t *testing.T) {
-			manager, err := newRouter(table, ifaceMock)
-			require.NoError(t, err, "failed to create router")
-			require.NoError(t, manager.init(table))
+			// need fw manager to init both acl mgr and router for all chains to be present
+			manager, err := Create(ifaceMock)
+			t.Cleanup(func() {
+				require.NoError(t, manager.Reset(nil))
+			})
+			require.NoError(t, err)
+			require.NoError(t, manager.Init(nil))
 
 			nftablesTestingClient := &nftables.Conn{}
 
-			defer func(manager *router) {
-				require.NoError(t, manager.Reset(), "failed to reset rules")
-			}(manager)
-
-			require.NoError(t, err, "shouldn't return error")
-
-			err = manager.AddNatRule(testCase.InputPair)
+			rtr := manager.router
+			err = rtr.AddNatRule(testCase.InputPair)
 			require.NoError(t, err, "pair should be inserted")
 
-			defer func(manager *router, pair firewall.RouterPair) {
-				require.NoError(t, manager.RemoveNatRule(pair), "failed to remove rule")
-			}(manager, testCase.InputPair)
+			t.Cleanup(func() {
+				require.NoError(t, rtr.RemoveNatRule(testCase.InputPair), "failed to remove rule")
+			})
 
 			if testCase.InputPair.Masquerade {
-				sourceExp := generateCIDRMatcherExpressions(true, testCase.InputPair.Source)
-				destExp := generateCIDRMatcherExpressions(false, testCase.InputPair.Destination)
-				testingExpression := append(sourceExp, destExp...) //nolint:gocritic
-				testingExpression = append(testingExpression,
-					&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
+				// Build expected expressions for connection tracking
+				conntrackExprs := []expr.Any{
+					&expr.Ct{
+						Key:      expr.CtKeySTATE,
+						Register: 1,
+					},
+					&expr.Bitwise{
+						SourceRegister: 1,
+						DestRegister:   1,
+						Len:            4,
+						Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
+						Xor:            binaryutil.NativeEndian.PutUint32(0),
+					},
+					&expr.Cmp{
+						Op:       expr.CmpOpNeq,
+						Register: 1,
+						Data:     []byte{0, 0, 0, 0},
+					},
+				}
+
+				// Build interface matching expression
+				ifaceExprs := []expr.Any{
+					&expr.Meta{
+						Key:      expr.MetaKeyIIFNAME,
+						Register: 1,
+					},
 					&expr.Cmp{
 						Op:       expr.CmpOpEq,
 						Register: 1,
 						Data:     ifname(ifaceMock.Name()),
 					},
-					&expr.Meta{Key: expr.MetaKeyOIFNAME, Register: 1},
-					&expr.Cmp{
-						Op:       expr.CmpOpNeq,
-						Register: 1,
-						Data:     ifname("lo"),
-					},
-				)
-
-				natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
-				found := 0
-				for _, chain := range manager.chains {
-					rules, err := nftablesTestingClient.GetRules(chain.Table, chain)
-					require.NoError(t, err, "should list rules for %s table and %s chain", chain.Table.Name, chain.Name)
-					for _, rule := range rules {
-						if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
-							require.ElementsMatchf(t, rule.Exprs[:len(testingExpression)], testingExpression, "nat rule elements should match")
-							found = 1
-						}
-					}
 				}
-				require.Equal(t, 1, found, "should find at least 1 rule to test")
-			}
 
-			if testCase.InputPair.Masquerade {
+				// Build CIDR matching expressions
 				sourceExp := generateCIDRMatcherExpressions(true, testCase.InputPair.Source)
 				destExp := generateCIDRMatcherExpressions(false, testCase.InputPair.Destination)
-				testingExpression := append(sourceExp, destExp...) //nolint:gocritic
-				testingExpression = append(testingExpression,
-					&expr.Meta{Key: expr.MetaKeyOIFNAME, Register: 1},
-					&expr.Cmp{
-						Op:       expr.CmpOpEq,
-						Register: 1,
-						Data:     ifname(ifaceMock.Name()),
-					},
-					&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
-					&expr.Cmp{
-						Op:       expr.CmpOpNeq,
-						Register: 1,
-						Data:     ifname("lo"),
-					},
-				)
 
-				inNatRuleKey := firewall.GenKey(firewall.NatFormat, firewall.GetInversePair(testCase.InputPair))
+				// Combine all expressions in the correct order
+				// nolint:gocritic
+				testingExpression := append(conntrackExprs, ifaceExprs...)
+				testingExpression = append(testingExpression, sourceExp...)
+				testingExpression = append(testingExpression, destExp...)
+
+				natRuleKey := firewall.GenKey(firewall.PreroutingFormat, testCase.InputPair)
 				found := 0
-				for _, chain := range manager.chains {
-					rules, err := nftablesTestingClient.GetRules(chain.Table, chain)
-					require.NoError(t, err, "should list rules for %s table and %s chain", chain.Table.Name, chain.Name)
-					for _, rule := range rules {
-						if len(rule.UserData) > 0 && string(rule.UserData) == inNatRuleKey {
-							require.ElementsMatchf(t, rule.Exprs[:len(testingExpression)], testingExpression, "income nat rule elements should match")
-							found = 1
+				for _, chain := range rtr.chains {
+					if chain.Name == chainNamePrerouting {
+						rules, err := nftablesTestingClient.GetRules(chain.Table, chain)
+						require.NoError(t, err, "should list rules for %s table and %s chain", chain.Table.Name, chain.Name)
+						for _, rule := range rules {
+							if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
+								// Compare expressions up to the mark setting expressions
+								require.ElementsMatchf(t, rule.Exprs[:len(testingExpression)], testingExpression, "prerouting nat rule elements should match")
+								found = 1
+							}
 						}
 					}
 				}
-				require.Equal(t, 1, found, "should find at least 1 rule to test")
+				require.Equal(t, 1, found, "should find at least 1 rule in prerouting chain")
 			}
-
 		})
 	}
 }
@@ -135,68 +123,66 @@ func TestNftablesManager_RemoveNatRule(t *testing.T) {
 		t.Skip("nftables not supported on this OS")
 	}
 
-	table, err := createWorkTable()
-	require.NoError(t, err, "Failed to create work table")
-
-	defer deleteWorkTable()
-
 	for _, testCase := range test.RemoveRuleTestCases {
 		t.Run(testCase.Name, func(t *testing.T) {
-			manager, err := newRouter(table, ifaceMock)
-			require.NoError(t, err, "failed to create router")
-			require.NoError(t, manager.init(table))
-
-			nftablesTestingClient := &nftables.Conn{}
-
-			defer func(manager *router) {
-				require.NoError(t, manager.Reset(), "failed to reset rules")
-			}(manager)
-
-			sourceExp := generateCIDRMatcherExpressions(true, testCase.InputPair.Source)
-			destExp := generateCIDRMatcherExpressions(false, testCase.InputPair.Destination)
-
-			natExp := append(sourceExp, append(destExp, &expr.Counter{}, &expr.Masq{})...) //nolint:gocritic
-			natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
-
-			insertedNat := nftablesTestingClient.InsertRule(&nftables.Rule{
-				Table:    manager.workTable,
-				Chain:    manager.chains[chainNameRoutingNat],
-				Exprs:    natExp,
-				UserData: []byte(natRuleKey),
+			manager, err := Create(ifaceMock)
+			t.Cleanup(func() {
+				require.NoError(t, manager.Reset(nil))
 			})
+			require.NoError(t, err)
+			require.NoError(t, manager.Init(nil))
 
-			sourceExp = generateCIDRMatcherExpressions(true, firewall.GetInversePair(testCase.InputPair).Source)
-			destExp = generateCIDRMatcherExpressions(false, firewall.GetInversePair(testCase.InputPair).Destination)
+			rtr := manager.router
 
-			natExp = append(sourceExp, append(destExp, &expr.Counter{}, &expr.Masq{})...) //nolint:gocritic
-			inNatRuleKey := firewall.GenKey(firewall.NatFormat, firewall.GetInversePair(testCase.InputPair))
+			// First add the NAT rule using the router's method
+			err = rtr.AddNatRule(testCase.InputPair)
+			require.NoError(t, err, "should add NAT rule")
 
-			insertedInNat := nftablesTestingClient.InsertRule(&nftables.Rule{
-				Table:    manager.workTable,
-				Chain:    manager.chains[chainNameRoutingNat],
-				Exprs:    natExp,
-				UserData: []byte(inNatRuleKey),
-			})
-
-			err = nftablesTestingClient.Flush()
-			require.NoError(t, err, "shouldn't return error")
-
-			err = manager.Reset()
-			require.NoError(t, err, "shouldn't return error")
-
-			err = manager.RemoveNatRule(testCase.InputPair)
-			require.NoError(t, err, "shouldn't return error")
-
-			for _, chain := range manager.chains {
-				rules, err := nftablesTestingClient.GetRules(chain.Table, chain)
-				require.NoError(t, err, "should list rules for %s table and %s chain", chain.Table.Name, chain.Name)
-				for _, rule := range rules {
-					if len(rule.UserData) > 0 {
-						require.NotEqual(t, insertedNat.UserData, rule.UserData, "nat rule should not exist")
-						require.NotEqual(t, insertedInNat.UserData, rule.UserData, "income nat rule should not exist")
-					}
+			// Verify the rule was added
+			natRuleKey := firewall.GenKey(firewall.PreroutingFormat, testCase.InputPair)
+			found := false
+			rules, err := rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
+			require.NoError(t, err, "should list rules")
+			for _, rule := range rules {
+				if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
+					found = true
+					break
 				}
 			}
+			require.True(t, found, "NAT rule should exist before removal")
+
+			// Now remove the rule
+			err = rtr.RemoveNatRule(testCase.InputPair)
+			require.NoError(t, err, "shouldn't return error when removing rule")
+
+			// Verify the rule was removed
+			found = false
+			rules, err = rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
+			require.NoError(t, err, "should list rules after removal")
+			for _, rule := range rules {
+				if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
+					found = true
+					break
+				}
+			}
+			require.False(t, found, "NAT rule should not exist after removal")
+
+			// Verify the static postrouting rules still exist
+			rules, err = rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNameRoutingNat])
+			require.NoError(t, err, "should list postrouting rules")
+			foundCounter := false
+			for _, rule := range rules {
+				for _, e := range rule.Exprs {
+					if _, ok := e.(*expr.Counter); ok {
+						foundCounter = true
+						break
+					}
+				}
+				if foundCounter {
+					break
+				}
+			}
+			require.True(t, foundCounter, "static postrouting rule should remain")
 		})
 	}
 }

client/firewall/nftables/state.go

@@ -1 +0,0 @@
-package nftables

client/firewall/uspfilter/allow_netbird.go

@@ -2,7 +2,10 @@
 
 package uspfilter
 
-import "github.com/netbirdio/netbird/client/internal/statemanager"
+import (
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
+	"github.com/netbirdio/netbird/client/internal/statemanager"
+)
 
 // Reset firewall to the default state
 func (m *Manager) Reset(stateManager *statemanager.Manager) error {
@@ -12,6 +15,21 @@ func (m *Manager) Reset(stateManager *statemanager.Manager) error {
 	m.outgoingRules = make(map[string]RuleSet)
 	m.incomingRules = make(map[string]RuleSet)
 
+	if m.udpTracker != nil {
+		m.udpTracker.Close()
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, nil)
+	}
+
+	if m.icmpTracker != nil {
+		m.icmpTracker.Close()
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, nil)
+	}
+
+	if m.tcpTracker != nil {
+		m.tcpTracker.Close()
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, nil)
+	}
+
 	if m.nativeFirewall != nil {
 		return m.nativeFirewall.Reset(stateManager)
 	}

client/firewall/uspfilter/allow_netbird_windows.go

@@ -7,6 +7,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
@@ -26,6 +27,21 @@ func (m *Manager) Reset(*statemanager.Manager) error {
 	m.outgoingRules = make(map[string]RuleSet)
 	m.incomingRules = make(map[string]RuleSet)
 
+	if m.udpTracker != nil {
+		m.udpTracker.Close()
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, nil)
+	}
+
+	if m.icmpTracker != nil {
+		m.icmpTracker.Close()
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, nil)
+	}
+
+	if m.tcpTracker != nil {
+		m.tcpTracker.Close()
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, nil)
+	}
+
 	if !isWindowsFirewallReachable() {
 		return nil
 	}

client/firewall/uspfilter/conntrack/common.go

@@ -0,0 +1,138 @@
+// common.go
+package conntrack
+
+import (
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// BaseConnTrack provides common fields and locking for all connection types
+type BaseConnTrack struct {
+	sync.RWMutex
+	SourceIP    net.IP
+	DestIP      net.IP
+	SourcePort  uint16
+	DestPort    uint16
+	lastSeen    atomic.Int64 // Unix nano for atomic access
+	established atomic.Bool
+}
+
+// these small methods will be inlined by the compiler
+
+// UpdateLastSeen safely updates the last seen timestamp
+func (b *BaseConnTrack) UpdateLastSeen() {
+	b.lastSeen.Store(time.Now().UnixNano())
+}
+
+// IsEstablished safely checks if connection is established
+func (b *BaseConnTrack) IsEstablished() bool {
+	return b.established.Load()
+}
+
+// SetEstablished safely sets the established state
+func (b *BaseConnTrack) SetEstablished(state bool) {
+	b.established.Store(state)
+}
+
+// GetLastSeen safely gets the last seen timestamp
+func (b *BaseConnTrack) GetLastSeen() time.Time {
+	return time.Unix(0, b.lastSeen.Load())
+}
+
+// timeoutExceeded checks if the connection has exceeded the given timeout
+func (b *BaseConnTrack) timeoutExceeded(timeout time.Duration) bool {
+	lastSeen := time.Unix(0, b.lastSeen.Load())
+	return time.Since(lastSeen) > timeout
+}
+
+// IPAddr is a fixed-size IP address to avoid allocations
+type IPAddr [16]byte
+
+// MakeIPAddr creates an IPAddr from net.IP
+func MakeIPAddr(ip net.IP) (addr IPAddr) {
+	// Optimization: check for v4 first as it's more common
+	if ip4 := ip.To4(); ip4 != nil {
+		copy(addr[12:], ip4)
+	} else {
+		copy(addr[:], ip.To16())
+	}
+	return addr
+}
+
+// ConnKey uniquely identifies a connection
+type ConnKey struct {
+	SrcIP   IPAddr
+	DstIP   IPAddr
+	SrcPort uint16
+	DstPort uint16
+}
+
+// makeConnKey creates a connection key
+func makeConnKey(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16) ConnKey {
+	return ConnKey{
+		SrcIP:   MakeIPAddr(srcIP),
+		DstIP:   MakeIPAddr(dstIP),
+		SrcPort: srcPort,
+		DstPort: dstPort,
+	}
+}
+
+// ValidateIPs checks if IPs match without allocation
+func ValidateIPs(connIP IPAddr, pktIP net.IP) bool {
+	if ip4 := pktIP.To4(); ip4 != nil {
+		// Compare IPv4 addresses (last 4 bytes)
+		for i := 0; i < 4; i++ {
+			if connIP[12+i] != ip4[i] {
+				return false
+			}
+		}
+		return true
+	}
+	// Compare full IPv6 addresses
+	ip6 := pktIP.To16()
+	for i := 0; i < 16; i++ {
+		if connIP[i] != ip6[i] {
+			return false
+		}
+	}
+	return true
+}
+
+// PreallocatedIPs is a pool of IP byte slices to reduce allocations
+type PreallocatedIPs struct {
+	sync.Pool
+}
+
+// NewPreallocatedIPs creates a new IP pool
+func NewPreallocatedIPs() *PreallocatedIPs {
+	return &PreallocatedIPs{
+		Pool: sync.Pool{
+			New: func() interface{} {
+				ip := make(net.IP, 16)
+				return &ip
+			},
+		},
+	}
+}
+
+// Get retrieves an IP from the pool
+func (p *PreallocatedIPs) Get() net.IP {
+	return *p.Pool.Get().(*net.IP)
+}
+
+// Put returns an IP to the pool
+func (p *PreallocatedIPs) Put(ip net.IP) {
+	p.Pool.Put(&ip)
+}
+
+// copyIP copies an IP address efficiently
+func copyIP(dst, src net.IP) {
+	if len(src) == 16 {
+		copy(dst, src)
+	} else {
+		// Handle IPv4
+		copy(dst[12:], src.To4())
+	}
+}

client/firewall/uspfilter/conntrack/common_test.go

@@ -0,0 +1,115 @@
+package conntrack
+
+import (
+	"net"
+	"testing"
+)
+
+func BenchmarkIPOperations(b *testing.B) {
+	b.Run("MakeIPAddr", func(b *testing.B) {
+		ip := net.ParseIP("192.168.1.1")
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = MakeIPAddr(ip)
+		}
+	})
+
+	b.Run("ValidateIPs", func(b *testing.B) {
+		ip1 := net.ParseIP("192.168.1.1")
+		ip2 := net.ParseIP("192.168.1.1")
+		addr := MakeIPAddr(ip1)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = ValidateIPs(addr, ip2)
+		}
+	})
+
+	b.Run("IPPool", func(b *testing.B) {
+		pool := NewPreallocatedIPs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			ip := pool.Get()
+			pool.Put(ip)
+		}
+	})
+
+}
+func BenchmarkAtomicOperations(b *testing.B) {
+	conn := &BaseConnTrack{}
+	b.Run("UpdateLastSeen", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			conn.UpdateLastSeen()
+		}
+	})
+
+	b.Run("IsEstablished", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_ = conn.IsEstablished()
+		}
+	})
+
+	b.Run("SetEstablished", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			conn.SetEstablished(i%2 == 0)
+		}
+	})
+
+	b.Run("GetLastSeen", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_ = conn.GetLastSeen()
+		}
+	})
+}
+
+// Memory pressure tests
+func BenchmarkMemoryPressure(b *testing.B) {
+	b.Run("TCPHighLoad", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+		defer tracker.Close()
+
+		// Generate different IPs
+		srcIPs := make([]net.IP, 100)
+		dstIPs := make([]net.IP, 100)
+		for i := 0; i < 100; i++ {
+			srcIPs[i] = net.IPv4(192, 168, byte(i/256), byte(i%256))
+			dstIPs[i] = net.IPv4(10, 0, byte(i/256), byte(i%256))
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			srcIdx := i % len(srcIPs)
+			dstIdx := (i + 1) % len(dstIPs)
+			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80, TCPSyn, nil)
+
+			// Simulate some valid inbound packets
+			if i%3 == 0 {
+				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535), TCPAck, nil)
+			}
+		}
+	})
+
+	b.Run("UDPHighLoad", func(b *testing.B) {
+		tracker := NewUDPTracker(DefaultUDPTimeout, nil)
+		defer tracker.Close()
+
+		// Generate different IPs
+		srcIPs := make([]net.IP, 100)
+		dstIPs := make([]net.IP, 100)
+		for i := 0; i < 100; i++ {
+			srcIPs[i] = net.IPv4(192, 168, byte(i/256), byte(i%256))
+			dstIPs[i] = net.IPv4(10, 0, byte(i/256), byte(i%256))
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			srcIdx := i % len(srcIPs)
+			dstIdx := (i + 1) % len(dstIPs)
+			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80, nil)
+
+			// Simulate some valid inbound packets
+			if i%3 == 0 {
+				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535), nil)
+			}
+		}
+	})
+}

client/firewall/uspfilter/conntrack/icmp.go

@@ -0,0 +1,187 @@
+package conntrack
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/google/gopacket/layers"
+
+	fw "github.com/netbirdio/netbird/client/firewall/manager"
+)
+
+const (
+	// DefaultICMPTimeout is the default timeout for ICMP connections
+	DefaultICMPTimeout = 30 * time.Second
+	// ICMPCleanupInterval is how often we check for stale ICMP connections
+	ICMPCleanupInterval = 15 * time.Second
+)
+
+// ICMPConnKey uniquely identifies an ICMP connection
+type ICMPConnKey struct {
+	// Supports both IPv4 and IPv6
+	SrcIP    [16]byte
+	DstIP    [16]byte
+	Sequence uint16 // ICMP sequence number
+	ID       uint16 // ICMP identifier
+}
+
+// ICMPConnTrack represents an ICMP connection state
+type ICMPConnTrack struct {
+	BaseConnTrack
+	Sequence uint16
+	ID       uint16
+}
+
+// ICMPTracker manages ICMP connection states
+type ICMPTracker struct {
+	connections   map[ICMPConnKey]*ICMPConnTrack
+	timeout       time.Duration
+	cleanupTicker *time.Ticker
+	mutex         sync.RWMutex
+	done          chan struct{}
+	ipPool        *PreallocatedIPs
+	stats         *Stats
+}
+
+// NewICMPTracker creates a new ICMP connection tracker
+func NewICMPTracker(timeout time.Duration, stats *Stats) *ICMPTracker {
+	if timeout == 0 {
+		timeout = DefaultICMPTimeout
+	}
+
+	tracker := &ICMPTracker{
+		connections:   make(map[ICMPConnKey]*ICMPConnTrack),
+		timeout:       timeout,
+		cleanupTicker: time.NewTicker(ICMPCleanupInterval),
+		done:          make(chan struct{}),
+		ipPool:        NewPreallocatedIPs(),
+		stats:         stats,
+	}
+
+	go tracker.cleanupRoutine()
+	return tracker
+}
+
+// TrackOutbound records an outbound ICMP Echo Request
+func (t *ICMPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, id uint16, seq uint16, packetData []byte) {
+	key := makeICMPKey(srcIP, dstIP, id, seq)
+	now := time.Now().UnixNano()
+
+	t.mutex.Lock()
+	conn, exists := t.connections[key]
+	if !exists {
+		srcIPCopy := t.ipPool.Get()
+		dstIPCopy := t.ipPool.Get()
+		copyIP(srcIPCopy, srcIP)
+		copyIP(dstIPCopy, dstIP)
+
+		conn = &ICMPConnTrack{
+			BaseConnTrack: BaseConnTrack{
+				SourceIP: srcIPCopy,
+				DestIP:   dstIPCopy,
+			},
+			ID:       id,
+			Sequence: seq,
+		}
+		conn.lastSeen.Store(now)
+		conn.established.Store(true)
+		t.connections[key] = conn
+
+		if t.stats != nil {
+			t.stats.TrackNewConnection(1, srcIP, dstIP, 0, 0, fw.DirectionOutbound)
+		}
+	}
+	t.mutex.Unlock()
+
+	if t.stats != nil {
+		key := makeConnKey(srcIP, dstIP, 0, 0)
+		t.stats.TrackPacket(1, false, uint64(len(packetData)), false, key)
+	}
+	conn.lastSeen.Store(now)
+}
+
+// IsValidInbound checks if an inbound ICMP Echo Reply matches a tracked request
+func (t *ICMPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, id uint16, seq uint16, icmpType uint8, packetData []byte) bool {
+	switch icmpType {
+	case uint8(layers.ICMPv4TypeDestinationUnreachable),
+		uint8(layers.ICMPv4TypeTimeExceeded):
+		return true
+	case uint8(layers.ICMPv4TypeEchoReply):
+		// continue processing
+	default:
+		return false
+	}
+
+	key := makeICMPKey(dstIP, srcIP, id, seq)
+
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if !exists {
+		return false
+	}
+
+	if conn.timeoutExceeded(t.timeout) {
+		return false
+	}
+
+	if t.stats != nil {
+		key := makeConnKey(srcIP, dstIP, 0, 0)
+		t.stats.TrackPacket(1, false, uint64(len(packetData)), true, key)
+	}
+
+	return conn.IsEstablished() &&
+		ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
+		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
+		conn.ID == id &&
+		conn.Sequence == seq
+}
+
+func (t *ICMPTracker) cleanupRoutine() {
+	for {
+		select {
+		case <-t.cleanupTicker.C:
+			t.cleanup()
+		case <-t.done:
+			return
+		}
+	}
+}
+func (t *ICMPTracker) cleanup() {
+	t.mutex.Lock()
+	defer t.mutex.Unlock()
+
+	for key, conn := range t.connections {
+		if conn.timeoutExceeded(t.timeout) {
+			t.ipPool.Put(conn.SourceIP)
+			t.ipPool.Put(conn.DestIP)
+			delete(t.connections, key)
+		}
+	}
+}
+
+// Close stops the cleanup routine and releases resources
+func (t *ICMPTracker) Close() {
+	t.cleanupTicker.Stop()
+	close(t.done)
+
+	t.mutex.Lock()
+	for _, conn := range t.connections {
+		t.ipPool.Put(conn.SourceIP)
+		t.ipPool.Put(conn.DestIP)
+	}
+	t.connections = nil
+	t.mutex.Unlock()
+}
+
+// makeICMPKey creates an ICMP connection key
+func makeICMPKey(srcIP net.IP, dstIP net.IP, id uint16, seq uint16) ICMPConnKey {
+	return ICMPConnKey{
+		SrcIP:    MakeIPAddr(srcIP),
+		DstIP:    MakeIPAddr(dstIP),
+		ID:       id,
+		Sequence: seq,
+	}
+}

client/firewall/uspfilter/conntrack/icmp_test.go

@@ -0,0 +1,39 @@
+package conntrack
+
+import (
+	"net"
+	"testing"
+)
+
+func BenchmarkICMPTracker(b *testing.B) {
+	b.Run("TrackOutbound", func(b *testing.B) {
+		tracker := NewICMPTracker(DefaultICMPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), uint16(i%65535), nil)
+		}
+	})
+
+	b.Run("IsValidInbound", func(b *testing.B) {
+		tracker := NewICMPTracker(DefaultICMPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		// Pre-populate some connections
+		for i := 0; i < 1000; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), uint16(i), nil)
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.IsValidInbound(dstIP, srcIP, uint16(i%1000), uint16(i%1000), 0, nil)
+		}
+	})
+}

client/firewall/uspfilter/conntrack/stats.go

@@ -0,0 +1,172 @@
+package conntrack
+
+import (
+	"net"
+	"slices"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	fw "github.com/netbirdio/netbird/client/firewall/manager"
+)
+
+// Stats represents connection tracking statistics
+type Stats struct {
+	TotalConnsCreated   atomic.Uint64
+	TotalConnsTimedOut  atomic.Uint64
+	TotalPacketsDropped atomic.Uint64
+	ActiveConns         atomic.Int64
+
+	TCPConns  atomic.Int64
+	UDPConns  atomic.Int64
+	ICMPConns atomic.Int64
+
+	TCPStateStats struct {
+		SynReceived   atomic.Uint64
+		Established   atomic.Uint64
+		FinWait       atomic.Uint64
+		TimeWait      atomic.Uint64
+		InvalidStates atomic.Uint64
+	}
+
+	PacketStats struct {
+		TCPPackets  atomic.Uint64
+		UDPPackets  atomic.Uint64
+		ICMPPackets atomic.Uint64
+	}
+
+	flowMutex sync.RWMutex
+	flows     map[ConnKey]*fw.FlowStats
+}
+
+// NewStats creates a new Stats instance
+func NewStats() *Stats {
+	return &Stats{
+		flows: make(map[ConnKey]*fw.FlowStats),
+	}
+}
+
+// TrackNewConnection records a new connection
+func (s *Stats) TrackNewConnection(proto uint8, srcIP net.IP, dstIP net.IP, srcPort, dstPort uint16, direction fw.Direction) {
+	s.TotalConnsCreated.Add(1)
+	s.ActiveConns.Add(1)
+
+	switch proto {
+	case 6: // TCP
+		s.TCPConns.Add(1)
+	case 17: // UDP
+		s.UDPConns.Add(1)
+	case 1: // ICMP
+		s.ICMPConns.Add(1)
+	}
+
+	flow := &fw.FlowStats{
+		StartTime:  time.Now(),
+		LastSeen:   time.Now(),
+		Protocol:   proto,
+		Direction:  direction,
+		SourceIP:   slices.Clone(srcIP),
+		DestIP:     slices.Clone(dstIP),
+		SourcePort: srcPort,
+		DestPort:   dstPort,
+	}
+
+	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	s.flowMutex.Lock()
+	s.flows[key] = flow
+	s.flowMutex.Unlock()
+}
+
+// TrackConnectionClosed records a connection closure
+func (s *Stats) TrackConnectionClosed(proto uint8, timedOut bool, key ConnKey) {
+	s.ActiveConns.Add(-1)
+
+	if timedOut {
+		s.TotalConnsTimedOut.Add(1)
+	}
+
+	switch proto {
+	case 6: // TCP
+		s.TCPConns.Add(-1)
+	case 17: // UDP
+		s.UDPConns.Add(-1)
+	case 1: // ICMP
+		s.ICMPConns.Add(-1)
+	}
+
+	s.flowMutex.Lock()
+	delete(s.flows, key)
+	s.flowMutex.Unlock()
+}
+
+// TrackPacket records packet statistics
+func (s *Stats) TrackPacket(proto uint8, dropped bool, bytes uint64, isInbound bool, key ConnKey) {
+	if dropped {
+		s.TotalPacketsDropped.Add(1)
+		return
+	}
+
+	switch proto {
+	case 6: // TCP
+		s.PacketStats.TCPPackets.Add(1)
+	case 17: // UDP
+		s.PacketStats.UDPPackets.Add(1)
+	case 1: // ICMP
+		s.PacketStats.ICMPPackets.Add(1)
+	}
+
+	s.flowMutex.RLock()
+	if flow, exists := s.flows[key]; exists {
+		if isInbound {
+			flow.BytesIn.Add(bytes)
+			flow.PacketsIn.Add(1)
+		} else {
+			flow.BytesOut.Add(bytes)
+			flow.PacketsOut.Add(1)
+		}
+		flow.LastSeen = time.Now()
+	}
+	s.flowMutex.RUnlock()
+}
+
+// TrackTCPState updates TCP state statistics
+func (s *Stats) TrackTCPState(newState TCPState) {
+	switch newState {
+	case TCPStateSynReceived:
+		s.TCPStateStats.SynReceived.Add(1)
+	case TCPStateEstablished:
+		s.TCPStateStats.Established.Add(1)
+	case TCPStateFinWait1, TCPStateFinWait2:
+		s.TCPStateStats.FinWait.Add(1)
+	case TCPStateTimeWait:
+		s.TCPStateStats.TimeWait.Add(1)
+	default:
+		s.TCPStateStats.InvalidStates.Add(1)
+	}
+}
+
+// GetFlowSnapshot returns a copy of current flow statistics if enabled
+func (s *Stats) GetFlowSnapshot() []*fw.FlowStats {
+	s.flowMutex.RLock()
+	defer s.flowMutex.RUnlock()
+
+	snapshot := make([]*fw.FlowStats, 0, len(s.flows))
+	for _, flow := range s.flows {
+		snapshot = append(snapshot, flow.Clone())
+	}
+	return snapshot
+}
+
+// CleanupFlows removes flow entries older than the specified duration if enabled
+func (s *Stats) CleanupFlows(maxAge time.Duration) {
+	threshold := time.Now().Add(-maxAge)
+
+	s.flowMutex.Lock()
+	defer s.flowMutex.Unlock()
+
+	for key, flow := range s.flows {
+		if flow.LastSeen.Before(threshold) {
+			delete(s.flows, key)
+		}
+	}
+}

client/firewall/uspfilter/conntrack/tcp.go

@@ -0,0 +1,399 @@
+package conntrack
+
+// TODO: Send RST packets for invalid/timed-out connections
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	fw "github.com/netbirdio/netbird/client/firewall/manager"
+)
+
+const (
+	// MSL (Maximum Segment Lifetime) is typically 2 minutes
+	MSL = 2 * time.Minute
+	// TimeWaitTimeout (TIME-WAIT) should last 2*MSL
+	TimeWaitTimeout = 2 * MSL
+)
+
+const (
+	TCPSyn  uint8 = 0x02
+	TCPAck  uint8 = 0x10
+	TCPFin  uint8 = 0x01
+	TCPRst  uint8 = 0x04
+	TCPPush uint8 = 0x08
+	TCPUrg  uint8 = 0x20
+)
+
+const (
+	// DefaultTCPTimeout is the default timeout for established TCP connections
+	DefaultTCPTimeout = 3 * time.Hour
+	// TCPHandshakeTimeout is timeout for TCP handshake completion
+	TCPHandshakeTimeout = 60 * time.Second
+	// TCPCleanupInterval is how often we check for stale connections
+	TCPCleanupInterval = 5 * time.Minute
+)
+
+// TCPState represents the state of a TCP connection
+type TCPState int
+
+const (
+	TCPStateNew TCPState = iota
+	TCPStateSynSent
+	TCPStateSynReceived
+	TCPStateEstablished
+	TCPStateFinWait1
+	TCPStateFinWait2
+	TCPStateClosing
+	TCPStateTimeWait
+	TCPStateCloseWait
+	TCPStateLastAck
+	TCPStateClosed
+)
+
+// TCPConnKey uniquely identifies a TCP connection
+type TCPConnKey struct {
+	SrcIP   [16]byte
+	DstIP   [16]byte
+	SrcPort uint16
+	DstPort uint16
+}
+
+// TCPConnTrack represents a TCP connection state
+type TCPConnTrack struct {
+	BaseConnTrack
+	State TCPState
+}
+
+// TCPTracker manages TCP connection states
+type TCPTracker struct {
+	connections   map[ConnKey]*TCPConnTrack
+	mutex         sync.RWMutex
+	cleanupTicker *time.Ticker
+	done          chan struct{}
+	timeout       time.Duration
+	ipPool        *PreallocatedIPs
+	stats         *Stats
+}
+
+// NewTCPTracker creates a new TCP connection tracker
+func NewTCPTracker(timeout time.Duration, stats *Stats) *TCPTracker {
+	tracker := &TCPTracker{
+		connections:   make(map[ConnKey]*TCPConnTrack),
+		cleanupTicker: time.NewTicker(TCPCleanupInterval),
+		done:          make(chan struct{}),
+		timeout:       timeout,
+		ipPool:        NewPreallocatedIPs(),
+		stats:         stats,
+	}
+
+	go tracker.cleanupRoutine()
+	return tracker
+}
+
+// TrackOutbound processes an outbound TCP packet and updates connection state
+func (t *TCPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, flags uint8, packetData []byte) {
+	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	now := time.Now().UnixNano()
+
+	t.mutex.Lock()
+	conn, exists := t.connections[key]
+	if !exists {
+		srcIPCopy := t.ipPool.Get()
+		dstIPCopy := t.ipPool.Get()
+		copyIP(srcIPCopy, srcIP)
+		copyIP(dstIPCopy, dstIP)
+
+		conn = &TCPConnTrack{
+			BaseConnTrack: BaseConnTrack{
+				SourceIP:   srcIPCopy,
+				DestIP:     dstIPCopy,
+				SourcePort: srcPort,
+				DestPort:   dstPort,
+			},
+			State: TCPStateNew,
+		}
+		conn.lastSeen.Store(now)
+		conn.established.Store(false)
+		t.connections[key] = conn
+
+		if t.stats != nil {
+			t.stats.TrackNewConnection(6, srcIP, dstIP, srcPort, dstPort, fw.DirectionOutbound)
+		}
+	}
+	t.mutex.Unlock()
+
+	conn.Lock()
+	oldState := conn.State
+	t.updateState(conn, flags, true)
+	if oldState != conn.State && t.stats != nil {
+		t.stats.TrackTCPState(conn.State)
+	}
+	conn.Unlock()
+
+	if t.stats != nil {
+		t.stats.TrackPacket(6, false, uint64(len(packetData)), false, key)
+	}
+	conn.lastSeen.Store(now)
+}
+
+// IsValidInbound checks if an inbound TCP packet matches a tracked connection
+func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, flags uint8, packetData []byte) bool {
+
+	if !isValidFlagCombination(flags) {
+		return false
+	}
+
+	// Handle new SYN packets
+	if flags&TCPSyn != 0 && flags&TCPAck == 0 {
+		key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
+		t.mutex.Lock()
+		if _, exists := t.connections[key]; !exists {
+			// Use preallocated IPs
+			srcIPCopy := t.ipPool.Get()
+			dstIPCopy := t.ipPool.Get()
+			copyIP(srcIPCopy, dstIP)
+			copyIP(dstIPCopy, srcIP)
+
+			conn := &TCPConnTrack{
+				BaseConnTrack: BaseConnTrack{
+					SourceIP:   srcIPCopy,
+					DestIP:     dstIPCopy,
+					SourcePort: dstPort,
+					DestPort:   srcPort,
+				},
+				State: TCPStateSynReceived,
+			}
+			conn.lastSeen.Store(time.Now().UnixNano())
+			conn.established.Store(false)
+			t.connections[key] = conn
+		}
+		t.mutex.Unlock()
+
+		if t.stats != nil {
+			t.stats.TrackPacket(6, false, uint64(len(packetData)), true, key)
+		}
+
+		return true
+	}
+
+	// Look up existing connection
+	key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if !exists {
+		return false
+	}
+
+	if t.stats != nil {
+		t.stats.TrackPacket(6, false, uint64(len(packetData)), true, key)
+	}
+
+	// Handle RST packets
+	if flags&TCPRst != 0 {
+		conn.Lock()
+		isEstablished := conn.IsEstablished()
+		if isEstablished || conn.State == TCPStateSynSent || conn.State == TCPStateSynReceived {
+			conn.State = TCPStateClosed
+			conn.SetEstablished(false)
+			conn.Unlock()
+			return true
+		}
+		conn.Unlock()
+		return false
+	}
+
+	// Update state
+	conn.Lock()
+	t.updateState(conn, flags, false)
+	conn.UpdateLastSeen()
+	isEstablished := conn.IsEstablished()
+	isValidState := t.isValidStateForFlags(conn.State, flags)
+	conn.Unlock()
+
+	return isEstablished || isValidState
+}
+
+// updateState updates the TCP connection state based on flags
+func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound bool) {
+	// Handle RST flag specially - it always causes transition to closed
+	if flags&TCPRst != 0 {
+		conn.State = TCPStateClosed
+		conn.SetEstablished(false)
+		return
+	}
+
+	switch conn.State {
+	case TCPStateNew:
+		if flags&TCPSyn != 0 && flags&TCPAck == 0 {
+			conn.State = TCPStateSynSent
+		}
+
+	case TCPStateSynSent:
+		if flags&TCPSyn != 0 && flags&TCPAck != 0 {
+			if isOutbound {
+				conn.State = TCPStateSynReceived
+			} else {
+				// Simultaneous open
+				conn.State = TCPStateEstablished
+				conn.SetEstablished(true)
+			}
+		}
+
+	case TCPStateSynReceived:
+		if flags&TCPAck != 0 && flags&TCPSyn == 0 {
+			conn.State = TCPStateEstablished
+			conn.SetEstablished(true)
+		}
+
+	case TCPStateEstablished:
+		if flags&TCPFin != 0 {
+			if isOutbound {
+				conn.State = TCPStateFinWait1
+			} else {
+				conn.State = TCPStateCloseWait
+			}
+			conn.SetEstablished(false)
+		}
+
+	case TCPStateFinWait1:
+		switch {
+		case flags&TCPFin != 0 && flags&TCPAck != 0:
+			// Simultaneous close - both sides sent FIN
+			conn.State = TCPStateClosing
+		case flags&TCPFin != 0:
+			conn.State = TCPStateFinWait2
+		case flags&TCPAck != 0:
+			conn.State = TCPStateFinWait2
+		}
+
+	case TCPStateFinWait2:
+		if flags&TCPFin != 0 {
+			conn.State = TCPStateTimeWait
+		}
+
+	case TCPStateClosing:
+		if flags&TCPAck != 0 {
+			conn.State = TCPStateTimeWait
+			// Keep established = false from previous state
+		}
+
+	case TCPStateCloseWait:
+		if flags&TCPFin != 0 {
+			conn.State = TCPStateLastAck
+		}
+
+	case TCPStateLastAck:
+		if flags&TCPAck != 0 {
+			conn.State = TCPStateClosed
+		}
+
+	case TCPStateTimeWait:
+		// Stay in TIME-WAIT for 2MSL before transitioning to closed
+		// This is handled by the cleanup routine
+	}
+}
+
+// isValidStateForFlags checks if the TCP flags are valid for the current connection state
+func (t *TCPTracker) isValidStateForFlags(state TCPState, flags uint8) bool {
+	if !isValidFlagCombination(flags) {
+		return false
+	}
+
+	switch state {
+	case TCPStateNew:
+		return flags&TCPSyn != 0 && flags&TCPAck == 0
+	case TCPStateSynSent:
+		return flags&TCPSyn != 0 && flags&TCPAck != 0
+	case TCPStateSynReceived:
+		return flags&TCPAck != 0
+	case TCPStateEstablished:
+		if flags&TCPRst != 0 {
+			return true
+		}
+		return flags&TCPAck != 0
+	case TCPStateFinWait1:
+		return flags&TCPFin != 0 || flags&TCPAck != 0
+	case TCPStateFinWait2:
+		return flags&TCPFin != 0 || flags&TCPAck != 0
+	case TCPStateClosing:
+		// In CLOSING state, we should accept the final ACK
+		return flags&TCPAck != 0
+	case TCPStateTimeWait:
+		// In TIME_WAIT, we might see retransmissions
+		return flags&TCPAck != 0
+	case TCPStateCloseWait:
+		return flags&TCPFin != 0 || flags&TCPAck != 0
+	case TCPStateLastAck:
+		return flags&TCPAck != 0
+	}
+	return false
+}
+
+func (t *TCPTracker) cleanupRoutine() {
+	for {
+		select {
+		case <-t.cleanupTicker.C:
+			t.cleanup()
+		case <-t.done:
+			return
+		}
+	}
+}
+
+func (t *TCPTracker) cleanup() {
+	t.mutex.Lock()
+	defer t.mutex.Unlock()
+
+	for key, conn := range t.connections {
+		var timeout time.Duration
+		switch {
+		case conn.State == TCPStateTimeWait:
+			timeout = TimeWaitTimeout
+		case conn.IsEstablished():
+			timeout = t.timeout
+		default:
+			timeout = TCPHandshakeTimeout
+		}
+
+		lastSeen := conn.GetLastSeen()
+		if time.Since(lastSeen) > timeout {
+			// Return IPs to pool
+			t.ipPool.Put(conn.SourceIP)
+			t.ipPool.Put(conn.DestIP)
+			delete(t.connections, key)
+		}
+	}
+}
+
+// Close stops the cleanup routine and releases resources
+func (t *TCPTracker) Close() {
+	t.cleanupTicker.Stop()
+	close(t.done)
+
+	// Clean up all remaining IPs
+	t.mutex.Lock()
+	for _, conn := range t.connections {
+		t.ipPool.Put(conn.SourceIP)
+		t.ipPool.Put(conn.DestIP)
+	}
+	t.connections = nil
+	t.mutex.Unlock()
+}
+
+func isValidFlagCombination(flags uint8) bool {
+	// Invalid: SYN+FIN
+	if flags&TCPSyn != 0 && flags&TCPFin != 0 {
+		return false
+	}
+
+	// Invalid: RST with SYN or FIN
+	if flags&TCPRst != 0 && (flags&TCPSyn != 0 || flags&TCPFin != 0) {
+		return false
+	}
+
+	return true
+}

client/firewall/uspfilter/conntrack/tcp_test.go

@@ -0,0 +1,311 @@
+package conntrack
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTCPStateMachine(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+	defer tracker.Close()
+
+	srcIP := net.ParseIP("100.64.0.1")
+	dstIP := net.ParseIP("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	t.Run("Security Tests", func(t *testing.T) {
+		tests := []struct {
+			name     string
+			flags    uint8
+			wantDrop bool
+			desc     string
+		}{
+			{
+				name:     "Block unsolicited SYN-ACK",
+				flags:    TCPSyn | TCPAck,
+				wantDrop: true,
+				desc:     "Should block SYN-ACK without prior SYN",
+			},
+			{
+				name:     "Block invalid SYN-FIN",
+				flags:    TCPSyn | TCPFin,
+				wantDrop: true,
+				desc:     "Should block invalid SYN-FIN combination",
+			},
+			{
+				name:     "Block unsolicited RST",
+				flags:    TCPRst,
+				wantDrop: true,
+				desc:     "Should block RST without connection",
+			},
+			{
+				name:     "Block unsolicited ACK",
+				flags:    TCPAck,
+				wantDrop: true,
+				desc:     "Should block ACK without connection",
+			},
+			{
+				name:     "Block data without connection",
+				flags:    TCPAck | TCPPush,
+				wantDrop: true,
+				desc:     "Should block data without established connection",
+			},
+		}
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				isValid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, tt.flags, nil)
+				require.Equal(t, !tt.wantDrop, isValid, tt.desc)
+			})
+		}
+	})
+
+	t.Run("Connection Flow Tests", func(t *testing.T) {
+		tests := []struct {
+			name string
+			test func(*testing.T)
+			desc string
+		}{
+			{
+				name: "Normal Handshake",
+				test: func(t *testing.T) {
+					t.Helper()
+
+					// Send initial SYN
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, nil)
+
+					// Receive SYN-ACK
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, nil)
+					require.True(t, valid, "SYN-ACK should be allowed")
+
+					// Send ACK
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, nil)
+
+					// Test data transfer
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck, nil)
+					require.True(t, valid, "Data should be allowed after handshake")
+				},
+			},
+			{
+				name: "Normal Close",
+				test: func(t *testing.T) {
+					t.Helper()
+
+					// First establish connection
+					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+					// Send FIN
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, nil)
+
+					// Receive ACK for FIN
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, nil)
+					require.True(t, valid, "ACK for FIN should be allowed")
+
+					// Receive FIN from other side
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, nil)
+					require.True(t, valid, "FIN should be allowed")
+
+					// Send final ACK
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, nil)
+				},
+			},
+			{
+				name: "RST During Connection",
+				test: func(t *testing.T) {
+					t.Helper()
+
+					// First establish connection
+					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+					// Receive RST
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, nil)
+					require.True(t, valid, "RST should be allowed for established connection")
+
+					// Verify connection is closed
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck, nil)
+					t.Helper()
+
+					require.False(t, valid, "Data should be blocked after RST")
+				},
+			},
+			{
+				name: "Simultaneous Close",
+				test: func(t *testing.T) {
+					t.Helper()
+
+					// First establish connection
+					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+					// Both sides send FIN+ACK
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, nil)
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, nil)
+					require.True(t, valid, "Simultaneous FIN should be allowed")
+
+					// Both sides send final ACK
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, nil)
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, nil)
+					require.True(t, valid, "Final ACKs should be allowed")
+				},
+			},
+		}
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				t.Helper()
+
+				tracker = NewTCPTracker(DefaultTCPTimeout, nil)
+				tt.test(t)
+			})
+		}
+	})
+}
+
+func TestRSTHandling(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+	defer tracker.Close()
+
+	srcIP := net.ParseIP("100.64.0.1")
+	dstIP := net.ParseIP("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	tests := []struct {
+		name       string
+		setupState func()
+		sendRST    func()
+		wantValid  bool
+		desc       string
+	}{
+		{
+			name: "RST in established",
+			setupState: func() {
+				// Establish connection first
+				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, nil)
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, nil)
+				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, nil)
+			},
+			sendRST: func() {
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, nil)
+			},
+			wantValid: true,
+			desc:      "Should accept RST for established connection",
+		},
+		{
+			name:       "RST without connection",
+			setupState: func() {},
+			sendRST: func() {
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, nil)
+			},
+			wantValid: false,
+			desc:      "Should reject RST without connection",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.setupState()
+			tt.sendRST()
+
+			// Verify connection state is as expected
+			key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+			conn := tracker.connections[key]
+			if tt.wantValid {
+				require.NotNil(t, conn)
+				require.Equal(t, TCPStateClosed, conn.State)
+				require.False(t, conn.IsEstablished())
+			}
+		})
+	}
+}
+
+// Helper to establish a TCP connection
+func establishConnection(t *testing.T, tracker *TCPTracker, srcIP, dstIP net.IP, srcPort, dstPort uint16) {
+	t.Helper()
+
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, nil)
+
+	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, nil)
+	require.True(t, valid, "SYN-ACK should be allowed")
+
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, nil)
+}
+
+func BenchmarkTCPTracker(b *testing.B) {
+	b.Run("TrackOutbound", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn, nil)
+		}
+	})
+
+	b.Run("IsValidInbound", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		// Pre-populate some connections
+		for i := 0; i < 1000; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn, nil)
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), TCPAck, nil)
+		}
+	})
+
+	b.Run("ConcurrentAccess", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		b.RunParallel(func(pb *testing.PB) {
+			i := 0
+			for pb.Next() {
+				if i%2 == 0 {
+					tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn, nil)
+				} else {
+					tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%65535), TCPAck, nil)
+				}
+				i++
+			}
+		})
+	})
+}
+
+// Benchmark connection cleanup
+func BenchmarkCleanup(b *testing.B) {
+	b.Run("TCPCleanup", func(b *testing.B) {
+		tracker := NewTCPTracker(100*time.Millisecond, nil) // Short timeout for testing
+		defer tracker.Close()
+
+		// Pre-populate with expired connections
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+		for i := 0; i < 10000; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn, nil)
+		}
+
+		// Wait for connections to expire
+		time.Sleep(200 * time.Millisecond)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.cleanup()
+		}
+	})
+}

client/firewall/uspfilter/conntrack/udp.go

@@ -0,0 +1,173 @@
+package conntrack
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	fw "github.com/netbirdio/netbird/client/firewall/manager"
+)
+
+const (
+	// DefaultUDPTimeout is the default timeout for UDP connections
+	DefaultUDPTimeout = 30 * time.Second
+	// UDPCleanupInterval is how often we check for stale connections
+	UDPCleanupInterval = 15 * time.Second
+)
+
+// UDPConnTrack represents a UDP connection state
+type UDPConnTrack struct {
+	BaseConnTrack
+}
+
+// UDPTracker manages UDP connection states
+type UDPTracker struct {
+	connections   map[ConnKey]*UDPConnTrack
+	timeout       time.Duration
+	cleanupTicker *time.Ticker
+	mutex         sync.RWMutex
+	done          chan struct{}
+	ipPool        *PreallocatedIPs
+	stats         *Stats
+}
+
+// NewUDPTracker creates a new UDP connection tracker
+func NewUDPTracker(timeout time.Duration, stats *Stats) *UDPTracker {
+	if timeout == 0 {
+		timeout = DefaultUDPTimeout
+	}
+
+	tracker := &UDPTracker{
+		connections:   make(map[ConnKey]*UDPConnTrack),
+		timeout:       timeout,
+		cleanupTicker: time.NewTicker(UDPCleanupInterval),
+		done:          make(chan struct{}),
+		ipPool:        NewPreallocatedIPs(),
+		stats:         stats,
+	}
+
+	go tracker.cleanupRoutine()
+	return tracker
+}
+
+// TrackOutbound records an outbound UDP connection
+func (t *UDPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, packetData []byte) {
+	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	now := time.Now().UnixNano()
+
+	t.mutex.Lock()
+	conn, exists := t.connections[key]
+	if !exists {
+		srcIPCopy := t.ipPool.Get()
+		dstIPCopy := t.ipPool.Get()
+		copyIP(srcIPCopy, srcIP)
+		copyIP(dstIPCopy, dstIP)
+
+		conn = &UDPConnTrack{
+			BaseConnTrack: BaseConnTrack{
+				SourceIP:   srcIPCopy,
+				DestIP:     dstIPCopy,
+				SourcePort: srcPort,
+				DestPort:   dstPort,
+			},
+		}
+		conn.lastSeen.Store(now)
+		conn.established.Store(true)
+		t.connections[key] = conn
+
+		if t.stats != nil {
+			t.stats.TrackNewConnection(17, srcIP, dstIP, srcPort, dstPort, fw.DirectionOutbound)
+		}
+	}
+	t.mutex.Unlock()
+
+	if t.stats != nil {
+		t.stats.TrackPacket(17, false, uint64(len(packetData)), false, key)
+	}
+	conn.lastSeen.Store(now)
+}
+
+// IsValidInbound checks if an inbound packet matches a tracked connection
+func (t *UDPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, packetData []byte) bool {
+	key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
+
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if !exists {
+		return false
+	}
+
+	if conn.timeoutExceeded(t.timeout) {
+		return false
+	}
+
+	if t.stats != nil {
+		t.stats.TrackPacket(17, false, uint64(len(packetData)), true, key)
+	}
+
+	return conn.IsEstablished() &&
+		ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
+		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
+		conn.DestPort == srcPort &&
+		conn.SourcePort == dstPort
+}
+
+// cleanupRoutine periodically removes stale connections
+func (t *UDPTracker) cleanupRoutine() {
+	for {
+		select {
+		case <-t.cleanupTicker.C:
+			t.cleanup()
+		case <-t.done:
+			return
+		}
+	}
+}
+
+func (t *UDPTracker) cleanup() {
+	t.mutex.Lock()
+	defer t.mutex.Unlock()
+
+	for key, conn := range t.connections {
+		if conn.timeoutExceeded(t.timeout) {
+			t.ipPool.Put(conn.SourceIP)
+			t.ipPool.Put(conn.DestIP)
+			delete(t.connections, key)
+		}
+	}
+}
+
+// Close stops the cleanup routine and releases resources
+func (t *UDPTracker) Close() {
+	t.cleanupTicker.Stop()
+	close(t.done)
+
+	t.mutex.Lock()
+	for _, conn := range t.connections {
+		t.ipPool.Put(conn.SourceIP)
+		t.ipPool.Put(conn.DestIP)
+	}
+	t.connections = nil
+	t.mutex.Unlock()
+}
+
+// GetConnection safely retrieves a connection state
+func (t *UDPTracker) GetConnection(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16) (*UDPConnTrack, bool) {
+	t.mutex.RLock()
+	defer t.mutex.RUnlock()
+
+	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	conn, exists := t.connections[key]
+	if !exists {
+		return nil, false
+	}
+
+	return conn, true
+}
+
+// Timeout returns the configured timeout duration for the tracker
+func (t *UDPTracker) Timeout() time.Duration {
+	return t.timeout
+}

client/firewall/uspfilter/conntrack/udp_test.go

@@ -0,0 +1,243 @@
+package conntrack
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewUDPTracker(t *testing.T) {
+	tests := []struct {
+		name        string
+		timeout     time.Duration
+		wantTimeout time.Duration
+	}{
+		{
+			name:        "with custom timeout",
+			timeout:     1 * time.Minute,
+			wantTimeout: 1 * time.Minute,
+		},
+		{
+			name:        "with zero timeout uses default",
+			timeout:     0,
+			wantTimeout: DefaultUDPTimeout,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tracker := NewUDPTracker(tt.timeout, nil)
+			assert.NotNil(t, tracker)
+			assert.Equal(t, tt.wantTimeout, tracker.timeout)
+			assert.NotNil(t, tracker.connections)
+			assert.NotNil(t, tracker.cleanupTicker)
+			assert.NotNil(t, tracker.done)
+		})
+	}
+}
+
+func TestUDPTracker_TrackOutbound(t *testing.T) {
+	tracker := NewUDPTracker(DefaultUDPTimeout, nil)
+	defer tracker.Close()
+
+	srcIP := net.ParseIP("192.168.1.2")
+	dstIP := net.ParseIP("192.168.1.3")
+	srcPort := uint16(12345)
+	dstPort := uint16(53)
+
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, nil)
+
+	// Verify connection was tracked
+	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	conn, exists := tracker.connections[key]
+	require.True(t, exists)
+	assert.True(t, conn.SourceIP.Equal(srcIP))
+	assert.True(t, conn.DestIP.Equal(dstIP))
+	assert.Equal(t, srcPort, conn.SourcePort)
+	assert.Equal(t, dstPort, conn.DestPort)
+	assert.True(t, conn.IsEstablished())
+	assert.WithinDuration(t, time.Now(), conn.GetLastSeen(), 1*time.Second)
+}
+
+func TestUDPTracker_IsValidInbound(t *testing.T) {
+	tracker := NewUDPTracker(1*time.Second, nil)
+	defer tracker.Close()
+
+	srcIP := net.ParseIP("192.168.1.2")
+	dstIP := net.ParseIP("192.168.1.3")
+	srcPort := uint16(12345)
+	dstPort := uint16(53)
+
+	// Track outbound connection
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, nil)
+
+	tests := []struct {
+		name    string
+		srcIP   net.IP
+		dstIP   net.IP
+		srcPort uint16
+		dstPort uint16
+		sleep   time.Duration
+		want    bool
+	}{
+		{
+			name:    "valid inbound response",
+			srcIP:   dstIP,   // Original destination is now source
+			dstIP:   srcIP,   // Original source is now destination
+			srcPort: dstPort, // Original destination port is now source
+			dstPort: srcPort, // Original source port is now destination
+			sleep:   0,
+			want:    true,
+		},
+		{
+			name:    "invalid source IP",
+			srcIP:   net.ParseIP("192.168.1.4"),
+			dstIP:   srcIP,
+			srcPort: dstPort,
+			dstPort: srcPort,
+			sleep:   0,
+			want:    false,
+		},
+		{
+			name:    "invalid destination IP",
+			srcIP:   dstIP,
+			dstIP:   net.ParseIP("192.168.1.4"),
+			srcPort: dstPort,
+			dstPort: srcPort,
+			sleep:   0,
+			want:    false,
+		},
+		{
+			name:    "invalid source port",
+			srcIP:   dstIP,
+			dstIP:   srcIP,
+			srcPort: 54321,
+			dstPort: srcPort,
+			sleep:   0,
+			want:    false,
+		},
+		{
+			name:    "invalid destination port",
+			srcIP:   dstIP,
+			dstIP:   srcIP,
+			srcPort: dstPort,
+			dstPort: 54321,
+			sleep:   0,
+			want:    false,
+		},
+		{
+			name:    "expired connection",
+			srcIP:   dstIP,
+			dstIP:   srcIP,
+			srcPort: dstPort,
+			dstPort: srcPort,
+			sleep:   2 * time.Second, // Longer than tracker timeout
+			want:    false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.sleep > 0 {
+				time.Sleep(tt.sleep)
+			}
+			got := tracker.IsValidInbound(tt.srcIP, tt.dstIP, tt.srcPort, tt.dstPort, nil)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func TestUDPTracker_Cleanup(t *testing.T) {
+	// Use shorter intervals for testing
+	timeout := 50 * time.Millisecond
+	cleanupInterval := 25 * time.Millisecond
+
+	// Create tracker with custom cleanup interval
+	tracker := &UDPTracker{
+		connections:   make(map[ConnKey]*UDPConnTrack),
+		timeout:       timeout,
+		cleanupTicker: time.NewTicker(cleanupInterval),
+		done:          make(chan struct{}),
+		ipPool:        NewPreallocatedIPs(),
+	}
+
+	// Start cleanup routine
+	go tracker.cleanupRoutine()
+
+	// Add some connections
+	connections := []struct {
+		srcIP   net.IP
+		dstIP   net.IP
+		srcPort uint16
+		dstPort uint16
+	}{
+		{
+			srcIP:   net.ParseIP("192.168.1.2"),
+			dstIP:   net.ParseIP("192.168.1.3"),
+			srcPort: 12345,
+			dstPort: 53,
+		},
+		{
+			srcIP:   net.ParseIP("192.168.1.4"),
+			dstIP:   net.ParseIP("192.168.1.5"),
+			srcPort: 12346,
+			dstPort: 53,
+		},
+	}
+
+	for _, conn := range connections {
+		tracker.TrackOutbound(conn.srcIP, conn.dstIP, conn.srcPort, conn.dstPort, nil)
+	}
+
+	// Verify initial connections
+	assert.Len(t, tracker.connections, 2)
+
+	// Wait for connection timeout and cleanup interval
+	time.Sleep(timeout + 2*cleanupInterval)
+
+	tracker.mutex.RLock()
+	connCount := len(tracker.connections)
+	tracker.mutex.RUnlock()
+
+	// Verify connections were cleaned up
+	assert.Equal(t, 0, connCount, "Expected all connections to be cleaned up")
+
+	// Properly close the tracker
+	tracker.Close()
+}
+
+func BenchmarkUDPTracker(b *testing.B) {
+	b.Run("TrackOutbound", func(b *testing.B) {
+		tracker := NewUDPTracker(DefaultUDPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, nil)
+		}
+	})
+
+	b.Run("IsValidInbound", func(b *testing.B) {
+		tracker := NewUDPTracker(DefaultUDPTimeout, nil)
+		defer tracker.Close()
+
+		srcIP := net.ParseIP("192.168.1.1")
+		dstIP := net.ParseIP("192.168.1.2")
+
+		// Pre-populate some connections
+		for i := 0; i < 1000; i++ {
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, nil)
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), nil)
+		}
+	})
+}

client/firewall/uspfilter/uspfilter.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
+	"os"
+	"strconv"
 	"sync"
 
 	"github.com/google/gopacket"
@@ -12,6 +14,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -19,6 +22,11 @@ import (
 
 const layerTypeAll = 0
 
+const (
+	EnvDisableConntrack = "NB_DISABLE_CONNTRACK"
+	EnvEnableStats      = "NB_ENABLE_CONNTRACK_STATS"
+)
+
 var (
 	errRouteNotSupported = fmt.Errorf("route not supported with userspace firewall")
 )
@@ -42,6 +50,14 @@ type Manager struct {
 	nativeFirewall firewall.Manager
 
 	mutex sync.RWMutex
+
+	stateful    bool
+	udpTracker  *conntrack.UDPTracker
+	icmpTracker *conntrack.ICMPTracker
+	tcpTracker  *conntrack.TCPTracker
+
+	statsEnabled bool
+	stats        *conntrack.Stats
 }
 
 // decoder for packages
@@ -73,6 +89,9 @@ func CreateWithNativeFirewall(iface IFaceMapper, nativeFirewall firewall.Manager
 }
 
 func create(iface IFaceMapper) (*Manager, error) {
+	disableConntrack, _ := strconv.ParseBool(os.Getenv(EnvDisableConntrack))
+	enableStats, _ := strconv.ParseBool(os.Getenv(EnvEnableStats))
+
 	m := &Manager{
 		decoders: sync.Pool{
 			New: func() any {
@@ -90,6 +109,22 @@ func create(iface IFaceMapper) (*Manager, error) {
 		outgoingRules: make(map[string]RuleSet),
 		incomingRules: make(map[string]RuleSet),
 		wgIface:       iface,
+		stateful:      !disableConntrack,
+		statsEnabled:  enableStats,
+	}
+
+	if enableStats {
+		m.stats = conntrack.NewStats()
+		log.Info("connection tracking statistics enabled")
+	}
+
+	// Only initialize trackers if stateful mode is enabled
+	if disableConntrack {
+		log.Info("conntrack is disabled")
+	} else {
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.stats)
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.stats)
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.stats)
 	}
 
 	if err := iface.SetFilter(m); err != nil {
@@ -239,7 +274,7 @@ func (m *Manager) DeletePeerRule(rule firewall.Rule) error {
 // SetLegacyManagement doesn't need to be implemented for this manager
 func (m *Manager) SetLegacyManagement(isLegacy bool) error {
 	if m.nativeFirewall == nil {
-		return errRouteNotSupported
+		return nil
 	}
 	return m.nativeFirewall.SetLegacyManagement(isLegacy)
 }
@@ -249,16 +284,16 @@ func (m *Manager) Flush() error { return nil }
 
 // DropOutgoing filter outgoing packets
 func (m *Manager) DropOutgoing(packetData []byte) bool {
-	return m.dropFilter(packetData, m.outgoingRules, false)
+	return m.processOutgoingHooks(packetData)
 }
 
 // DropIncoming filter incoming packets
 func (m *Manager) DropIncoming(packetData []byte) bool {
-	return m.dropFilter(packetData, m.incomingRules, true)
+	return m.dropFilter(packetData, m.incomingRules)
 }
 
-// dropFilter implements same logic for booth direction of the traffic
-func (m *Manager) dropFilter(packetData []byte, rules map[string]RuleSet, isIncomingPacket bool) bool {
+// processOutgoingHooks processes UDP hooks for outgoing packets and tracks TCP/UDP/ICMP
+func (m *Manager) processOutgoingHooks(packetData []byte) bool {
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
@@ -266,61 +301,195 @@ func (m *Manager) dropFilter(packetData []byte, rules map[string]RuleSet, isInco
 	defer m.decoders.Put(d)
 
 	if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-		log.Tracef("couldn't decode layer, err: %s", err)
-		return true
+		return false
 	}
 
 	if len(d.decoded) < 2 {
-		log.Tracef("not enough levels in network packet")
+		return false
+	}
+
+	srcIP, dstIP := m.extractIPs(d)
+	if srcIP == nil {
+		return false
+	}
+
+	// Always process UDP hooks
+	if d.decoded[1] == layers.LayerTypeUDP {
+		// Track UDP state only if enabled
+		if m.stateful {
+			m.udpTracker.TrackOutbound(srcIP, dstIP,
+				uint16(d.udp.SrcPort),
+				uint16(d.udp.DstPort),
+				packetData)
+		}
+		return m.checkUDPHooks(d, dstIP, packetData)
+	}
+
+	// Track other protocols only if stateful mode is enabled
+	if m.stateful {
+		switch d.decoded[1] {
+		case layers.LayerTypeTCP:
+			m.tcpTracker.TrackOutbound(srcIP, dstIP,
+				uint16(d.tcp.SrcPort),
+				uint16(d.tcp.DstPort),
+				getTCPFlags(&d.tcp),
+				packetData)
+		case layers.LayerTypeICMPv4:
+			m.icmpTracker.TrackOutbound(srcIP, dstIP,
+				d.icmp4.Id,
+				d.icmp4.Seq,
+				packetData)
+		}
+	}
+
+	return false
+}
+
+func (m *Manager) extractIPs(d *decoder) (srcIP, dstIP net.IP) {
+	switch d.decoded[0] {
+	case layers.LayerTypeIPv4:
+		return d.ip4.SrcIP, d.ip4.DstIP
+	case layers.LayerTypeIPv6:
+		return d.ip6.SrcIP, d.ip6.DstIP
+	default:
+		return nil, nil
+	}
+}
+
+func getTCPFlags(tcp *layers.TCP) uint8 {
+	var flags uint8
+	if tcp.SYN {
+		flags |= conntrack.TCPSyn
+	}
+	if tcp.ACK {
+		flags |= conntrack.TCPAck
+	}
+	if tcp.FIN {
+		flags |= conntrack.TCPFin
+	}
+	if tcp.RST {
+		flags |= conntrack.TCPRst
+	}
+	if tcp.PSH {
+		flags |= conntrack.TCPPush
+	}
+	if tcp.URG {
+		flags |= conntrack.TCPUrg
+	}
+	return flags
+}
+
+func (m *Manager) checkUDPHooks(d *decoder, dstIP net.IP, packetData []byte) bool {
+	for _, ipKey := range []string{dstIP.String(), "0.0.0.0", "::"} {
+		if rules, exists := m.outgoingRules[ipKey]; exists {
+			for _, rule := range rules {
+				if rule.udpHook != nil && (rule.dPort == 0 || rule.dPort == uint16(d.udp.DstPort)) {
+					return rule.udpHook(packetData)
+				}
+			}
+		}
+	}
+	return false
+}
+
+// dropFilter implements filtering logic for incoming packets
+func (m *Manager) dropFilter(packetData []byte, rules map[string]RuleSet) bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	d := m.decoders.Get().(*decoder)
+	defer m.decoders.Put(d)
+
+	if !m.isValidPacket(d, packetData) {
 		return true
 	}
 
-	ipLayer := d.decoded[0]
-
-	switch ipLayer {
-	case layers.LayerTypeIPv4:
-		if !m.wgNetwork.Contains(d.ip4.SrcIP) || !m.wgNetwork.Contains(d.ip4.DstIP) {
-			return false
-		}
-	case layers.LayerTypeIPv6:
-		if !m.wgNetwork.Contains(d.ip6.SrcIP) || !m.wgNetwork.Contains(d.ip6.DstIP) {
-			return false
-		}
-	default:
+	srcIP, dstIP := m.extractIPs(d)
+	if srcIP == nil {
 		log.Errorf("unknown layer: %v", d.decoded[0])
 		return true
 	}
 
-	var ip net.IP
-	switch ipLayer {
-	case layers.LayerTypeIPv4:
-		if isIncomingPacket {
-			ip = d.ip4.SrcIP
-		} else {
-			ip = d.ip4.DstIP
-		}
-	case layers.LayerTypeIPv6:
-		if isIncomingPacket {
-			ip = d.ip6.SrcIP
-		} else {
-			ip = d.ip6.DstIP
-		}
+	if !m.isWireguardTraffic(srcIP, dstIP) {
+		return false
 	}
 
-	filter, ok := validateRule(ip, packetData, rules[ip.String()], d)
-	if ok {
-		return filter
+	// Check connection state only if enabled
+	if m.stateful && m.isValidTrackedConnection(d, srcIP, dstIP, packetData) {
+		return false
 	}
-	filter, ok = validateRule(ip, packetData, rules["0.0.0.0"], d)
-	if ok {
-		return filter
+
+	return m.applyRules(srcIP, packetData, rules, d)
+}
+
+func (m *Manager) isValidPacket(d *decoder, packetData []byte) bool {
+	if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
+		log.Tracef("couldn't decode layer, err: %s", err)
+		return false
 	}
-	filter, ok = validateRule(ip, packetData, rules["::"], d)
-	if ok {
+
+	if len(d.decoded) < 2 {
+		log.Tracef("not enough levels in network packet")
+		return false
+	}
+	return true
+}
+
+func (m *Manager) isWireguardTraffic(srcIP, dstIP net.IP) bool {
+	return m.wgNetwork.Contains(srcIP) && m.wgNetwork.Contains(dstIP)
+}
+
+func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP, packetData []byte) bool {
+	switch d.decoded[1] {
+	case layers.LayerTypeTCP:
+		return m.tcpTracker.IsValidInbound(
+			srcIP,
+			dstIP,
+			uint16(d.tcp.SrcPort),
+			uint16(d.tcp.DstPort),
+			getTCPFlags(&d.tcp),
+			packetData,
+		)
+
+	case layers.LayerTypeUDP:
+		return m.udpTracker.IsValidInbound(
+			srcIP,
+			dstIP,
+			uint16(d.udp.SrcPort),
+			uint16(d.udp.DstPort),
+			packetData,
+		)
+
+	case layers.LayerTypeICMPv4:
+		return m.icmpTracker.IsValidInbound(
+			srcIP,
+			dstIP,
+			d.icmp4.Id,
+			d.icmp4.Seq,
+			d.icmp4.TypeCode.Type(),
+			packetData,
+		)
+
+		// TODO: ICMPv6
+	}
+
+	return false
+}
+
+func (m *Manager) applyRules(srcIP net.IP, packetData []byte, rules map[string]RuleSet, d *decoder) bool {
+	if filter, ok := validateRule(srcIP, packetData, rules[srcIP.String()], d); ok {
 		return filter
 	}
 
-	// default policy is DROP ALL
+	if filter, ok := validateRule(srcIP, packetData, rules["0.0.0.0"], d); ok {
+		return filter
+	}
+
+	if filter, ok := validateRule(srcIP, packetData, rules["::"], d); ok {
+		return filter
+	}
+
+	// Default policy: DROP ALL
 	return true
 }
 
@@ -438,3 +607,11 @@ func (m *Manager) RemovePacketHook(hookID string) error {
 	}
 	return fmt.Errorf("hook with given id not found")
 }
+
+// CollectStats returns connection tracking statistics
+func (m *Manager) CollectStats() []*firewall.FlowStats {
+	if m.stats == nil {
+		return nil
+	}
+	return m.stats.GetFlowSnapshot()
+}

client/firewall/uspfilter/uspfilter_test.go

@@ -3,6 +3,7 @@ package uspfilter
 import (
 	"fmt"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
@@ -11,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	fw "github.com/netbirdio/netbird/client/firewall/manager"
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
@@ -185,10 +187,10 @@ func TestAddUDPPacketHook(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			manager := &Manager{
-				incomingRules: map[string]RuleSet{},
-				outgoingRules: map[string]RuleSet{},
-			}
+			manager, err := Create(&IFaceMock{
+				SetFilterFunc: func(device.PacketFilter) error { return nil },
+			})
+			require.NoError(t, err)
 
 			manager.AddUDPPacketHook(tt.in, tt.ip, tt.dPort, tt.hook)
 
@@ -313,7 +315,7 @@ func TestNotMatchByIP(t *testing.T) {
 		t.Errorf("failed to set network layer for checksum: %v", err)
 		return
 	}
-	payload := gopacket.Payload([]byte("test"))
+	payload := gopacket.Payload("test")
 
 	buf := gopacket.NewSerializeBuffer()
 	opts := gopacket.SerializeOptions{
@@ -325,7 +327,7 @@ func TestNotMatchByIP(t *testing.T) {
 		return
 	}
 
-	if m.dropFilter(buf.Bytes(), m.outgoingRules, false) {
+	if m.dropFilter(buf.Bytes(), m.outgoingRules) {
 		t.Errorf("expected packet to be accepted")
 		return
 	}
@@ -348,6 +350,9 @@ func TestRemovePacketHook(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create Manager: %s", err)
 	}
+	defer func() {
+		require.NoError(t, manager.Reset(nil))
+	}()
 
 	// Add a UDP packet hook
 	hookFunc := func(data []byte) bool { return true }
@@ -384,6 +389,88 @@ func TestRemovePacketHook(t *testing.T) {
 	}
 }
 
+func TestProcessOutgoingHooks(t *testing.T) {
+	manager, err := Create(&IFaceMock{
+		SetFilterFunc: func(device.PacketFilter) error { return nil },
+	})
+	require.NoError(t, err)
+
+	manager.wgNetwork = &net.IPNet{
+		IP:   net.ParseIP("100.10.0.0"),
+		Mask: net.CIDRMask(16, 32),
+	}
+	manager.udpTracker.Close()
+	manager.udpTracker = conntrack.NewUDPTracker(100*time.Millisecond, nil)
+	defer func() {
+		require.NoError(t, manager.Reset(nil))
+	}()
+
+	manager.decoders = sync.Pool{
+		New: func() any {
+			d := &decoder{
+				decoded: []gopacket.LayerType{},
+			}
+			d.parser = gopacket.NewDecodingLayerParser(
+				layers.LayerTypeIPv4,
+				&d.eth, &d.ip4, &d.ip6, &d.icmp4, &d.icmp6, &d.tcp, &d.udp,
+			)
+			d.parser.IgnoreUnsupported = true
+			return d
+		},
+	}
+
+	hookCalled := false
+	hookID := manager.AddUDPPacketHook(
+		false,
+		net.ParseIP("100.10.0.100"),
+		53,
+		func([]byte) bool {
+			hookCalled = true
+			return true
+		},
+	)
+	require.NotEmpty(t, hookID)
+
+	// Create test UDP packet
+	ipv4 := &layers.IPv4{
+		TTL:      64,
+		Version:  4,
+		SrcIP:    net.ParseIP("100.10.0.1"),
+		DstIP:    net.ParseIP("100.10.0.100"),
+		Protocol: layers.IPProtocolUDP,
+	}
+	udp := &layers.UDP{
+		SrcPort: 51334,
+		DstPort: 53,
+	}
+
+	err = udp.SetNetworkLayerForChecksum(ipv4)
+	require.NoError(t, err)
+	payload := gopacket.Payload("test")
+
+	buf := gopacket.NewSerializeBuffer()
+	opts := gopacket.SerializeOptions{
+		ComputeChecksums: true,
+		FixLengths:       true,
+	}
+	err = gopacket.SerializeLayers(buf, opts, ipv4, udp, payload)
+	require.NoError(t, err)
+
+	// Test hook gets called
+	result := manager.processOutgoingHooks(buf.Bytes())
+	require.True(t, result)
+	require.True(t, hookCalled)
+
+	// Test non-UDP packet is ignored
+	ipv4.Protocol = layers.IPProtocolTCP
+	buf = gopacket.NewSerializeBuffer()
+	err = gopacket.SerializeLayers(buf, opts, ipv4)
+	require.NoError(t, err)
+
+	result = manager.processOutgoingHooks(buf.Bytes())
+	require.False(t, result)
+}
+
 func TestUSPFilterCreatePerformance(t *testing.T) {
 	for _, testMax := range []int{10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000} {
 		t.Run(fmt.Sprintf("Testing %d rules", testMax), func(t *testing.T) {
@@ -418,3 +505,213 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 		})
 	}
 }
+
+func TestStatefulFirewall_UDPTracking(t *testing.T) {
+	manager, err := Create(&IFaceMock{
+		SetFilterFunc: func(device.PacketFilter) error { return nil },
+	})
+	require.NoError(t, err)
+
+	manager.wgNetwork = &net.IPNet{
+		IP:   net.ParseIP("100.10.0.0"),
+		Mask: net.CIDRMask(16, 32),
+	}
+
+	manager.udpTracker.Close() // Close the existing tracker
+	manager.udpTracker = conntrack.NewUDPTracker(200*time.Millisecond, nil)
+	manager.decoders = sync.Pool{
+		New: func() any {
+			d := &decoder{
+				decoded: []gopacket.LayerType{},
+			}
+			d.parser = gopacket.NewDecodingLayerParser(
+				layers.LayerTypeIPv4,
+				&d.eth, &d.ip4, &d.ip6, &d.icmp4, &d.icmp6, &d.tcp, &d.udp,
+			)
+			d.parser.IgnoreUnsupported = true
+			return d
+		},
+	}
+	defer func() {
+		require.NoError(t, manager.Reset(nil))
+	}()
+
+	// Set up packet parameters
+	srcIP := net.ParseIP("100.10.0.1")
+	dstIP := net.ParseIP("100.10.0.100")
+	srcPort := uint16(51334)
+	dstPort := uint16(53)
+
+	// Create outbound packet
+	outboundIPv4 := &layers.IPv4{
+		TTL:      64,
+		Version:  4,
+		SrcIP:    srcIP,
+		DstIP:    dstIP,
+		Protocol: layers.IPProtocolUDP,
+	}
+	outboundUDP := &layers.UDP{
+		SrcPort: layers.UDPPort(srcPort),
+		DstPort: layers.UDPPort(dstPort),
+	}
+
+	err = outboundUDP.SetNetworkLayerForChecksum(outboundIPv4)
+	require.NoError(t, err)
+
+	outboundBuf := gopacket.NewSerializeBuffer()
+	opts := gopacket.SerializeOptions{
+		ComputeChecksums: true,
+		FixLengths:       true,
+	}
+
+	err = gopacket.SerializeLayers(outboundBuf, opts,
+		outboundIPv4,
+		outboundUDP,
+		gopacket.Payload("test"),
+	)
+	require.NoError(t, err)
+
+	// Process outbound packet and verify connection tracking
+	drop := manager.DropOutgoing(outboundBuf.Bytes())
+	require.False(t, drop, "Initial outbound packet should not be dropped")
+
+	// Verify connection was tracked
+	conn, exists := manager.udpTracker.GetConnection(srcIP, srcPort, dstIP, dstPort)
+
+	require.True(t, exists, "Connection should be tracked after outbound packet")
+	require.True(t, conntrack.ValidateIPs(conntrack.MakeIPAddr(srcIP), conn.SourceIP), "Source IP should match")
+	require.True(t, conntrack.ValidateIPs(conntrack.MakeIPAddr(dstIP), conn.DestIP), "Destination IP should match")
+	require.Equal(t, srcPort, conn.SourcePort, "Source port should match")
+	require.Equal(t, dstPort, conn.DestPort, "Destination port should match")
+
+	// Create valid inbound response packet
+	inboundIPv4 := &layers.IPv4{
+		TTL:      64,
+		Version:  4,
+		SrcIP:    dstIP, // Original destination is now source
+		DstIP:    srcIP, // Original source is now destination
+		Protocol: layers.IPProtocolUDP,
+	}
+	inboundUDP := &layers.UDP{
+		SrcPort: layers.UDPPort(dstPort), // Original destination port is now source
+		DstPort: layers.UDPPort(srcPort), // Original source port is now destination
+	}
+
+	err = inboundUDP.SetNetworkLayerForChecksum(inboundIPv4)
+	require.NoError(t, err)
+
+	inboundBuf := gopacket.NewSerializeBuffer()
+	err = gopacket.SerializeLayers(inboundBuf, opts,
+		inboundIPv4,
+		inboundUDP,
+		gopacket.Payload("response"),
+	)
+	require.NoError(t, err)
+	// Test roundtrip response handling over time
+	checkPoints := []struct {
+		sleep       time.Duration
+		shouldAllow bool
+		description string
+	}{
+		{
+			sleep:       0,
+			shouldAllow: true,
+			description: "Immediate response should be allowed",
+		},
+		{
+			sleep:       50 * time.Millisecond,
+			shouldAllow: true,
+			description: "Response within timeout should be allowed",
+		},
+		{
+			sleep:       100 * time.Millisecond,
+			shouldAllow: true,
+			description: "Response at half timeout should be allowed",
+		},
+		{
+			// tracker hasn't updated conn for 250ms -> greater than 200ms timeout
+			sleep:       250 * time.Millisecond,
+			shouldAllow: false,
+			description: "Response after timeout should be dropped",
+		},
+	}
+
+	for _, cp := range checkPoints {
+		time.Sleep(cp.sleep)
+
+		drop = manager.dropFilter(inboundBuf.Bytes(), manager.incomingRules)
+		require.Equal(t, cp.shouldAllow, !drop, cp.description)
+
+		// If the connection should still be valid, verify it exists
+		if cp.shouldAllow {
+			conn, exists := manager.udpTracker.GetConnection(srcIP, srcPort, dstIP, dstPort)
+			require.True(t, exists, "Connection should still exist during valid window")
+			require.True(t, time.Since(conn.GetLastSeen()) < manager.udpTracker.Timeout(),
+				"LastSeen should be updated for valid responses")
+		}
+	}
+
+	// Test invalid response packets (while connection is expired)
+	invalidCases := []struct {
+		name        string
+		modifyFunc  func(*layers.IPv4, *layers.UDP)
+		description string
+	}{
+		{
+			name: "wrong source IP",
+			modifyFunc: func(ip *layers.IPv4, udp *layers.UDP) {
+				ip.SrcIP = net.ParseIP("100.10.0.101")
+			},
+			description: "Response from wrong IP should be dropped",
+		},
+		{
+			name: "wrong destination IP",
+			modifyFunc: func(ip *layers.IPv4, udp *layers.UDP) {
+				ip.DstIP = net.ParseIP("100.10.0.2")
+			},
+			description: "Response to wrong IP should be dropped",
+		},
+		{
+			name: "wrong source port",
+			modifyFunc: func(ip *layers.IPv4, udp *layers.UDP) {
+				udp.SrcPort = 54
+			},
+			description: "Response from wrong port should be dropped",
+		},
+		{
+			name: "wrong destination port",
+			modifyFunc: func(ip *layers.IPv4, udp *layers.UDP) {
+				udp.DstPort = 51335
+			},
+			description: "Response to wrong port should be dropped",
+		},
+	}
+
+	// Create a new outbound connection for invalid tests
+	drop = manager.processOutgoingHooks(outboundBuf.Bytes())
+	require.False(t, drop, "Second outbound packet should not be dropped")
+
+	for _, tc := range invalidCases {
+		t.Run(tc.name, func(t *testing.T) {
+			testIPv4 := *inboundIPv4
+			testUDP := *inboundUDP
+
+			tc.modifyFunc(&testIPv4, &testUDP)
+
+			err = testUDP.SetNetworkLayerForChecksum(&testIPv4)
+			require.NoError(t, err)
+
+			testBuf := gopacket.NewSerializeBuffer()
+			err = gopacket.SerializeLayers(testBuf, opts,
+				&testIPv4,
+				&testUDP,
+				gopacket.Payload("response"),
+			)
+			require.NoError(t, err)
+
+			// Verify the invalid packet is dropped
+			drop = manager.dropFilter(testBuf.Bytes(), manager.incomingRules)
+			require.True(t, drop, tc.description)
+		})
+	}
+}

client/iface/bind/control_android.go

@@ -0,0 +1,12 @@
+package bind
+
+import (
+	wireguard "golang.zx2c4.com/wireguard/conn"
+
+	nbnet "github.com/netbirdio/netbird/util/net"
+)
+
+func init() {
+	// ControlFns is not thread safe and should only be modified during init.
+	*wireguard.ControlFns = append(*wireguard.ControlFns, nbnet.ControlProtectSocket)
+}

client/iface/bind/ice_bind.go

@@ -12,6 +12,7 @@ import (
 	"github.com/pion/transport/v3"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/ipv4"
+	"golang.org/x/net/ipv6"
 	wgConn "golang.zx2c4.com/wireguard/conn"
 )
 
@@ -24,8 +25,8 @@ type receiverCreator struct {
 	iceBind *ICEBind
 }
 
-func (rc receiverCreator) CreateIPv4ReceiverFn(msgPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) wgConn.ReceiveFunc {
-	return rc.iceBind.createIPv4ReceiverFn(msgPool, pc, conn)
+func (rc receiverCreator) CreateIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool, msgPool *sync.Pool) wgConn.ReceiveFunc {
+	return rc.iceBind.createIPv4ReceiverFn(pc, conn, rxOffload, msgPool)
 }
 
 // ICEBind is a bind implementation with two main features:
@@ -154,7 +155,7 @@ func (b *ICEBind) Send(bufs [][]byte, ep wgConn.Endpoint) error {
 	return nil
 }
 
-func (s *ICEBind) createIPv4ReceiverFn(ipv4MsgsPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) wgConn.ReceiveFunc {
+func (s *ICEBind) createIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool, msgsPool *sync.Pool) wgConn.ReceiveFunc {
 	s.muUDPMux.Lock()
 	defer s.muUDPMux.Unlock()
 
@@ -166,16 +167,30 @@ func (s *ICEBind) createIPv4ReceiverFn(ipv4MsgsPool *sync.Pool, pc *ipv4.PacketC
 		},
 	)
 	return func(bufs [][]byte, sizes []int, eps []wgConn.Endpoint) (n int, err error) {
-		msgs := ipv4MsgsPool.Get().(*[]ipv4.Message)
-		defer ipv4MsgsPool.Put(msgs)
+		msgs := getMessages(msgsPool)
 		for i := range bufs {
 			(*msgs)[i].Buffers[0] = bufs[i]
+			(*msgs)[i].OOB = (*msgs)[i].OOB[:cap((*msgs)[i].OOB)]
 		}
+		defer putMessages(msgs, msgsPool)
 		var numMsgs int
-		if runtime.GOOS == "linux" {
-			numMsgs, err = pc.ReadBatch(*msgs, 0)
-			if err != nil {
-				return 0, err
+		if runtime.GOOS == "linux" || runtime.GOOS == "android" {
+			if rxOffload {
+				readAt := len(*msgs) - (wgConn.IdealBatchSize / wgConn.UdpSegmentMaxDatagrams)
+				//nolint
+				numMsgs, err = pc.ReadBatch((*msgs)[readAt:], 0)
+				if err != nil {
+					return 0, err
+				}
+				numMsgs, err = wgConn.SplitCoalescedMessages(*msgs, readAt, wgConn.GetGSOSize)
+				if err != nil {
+					return 0, err
+				}
+			} else {
+				numMsgs, err = pc.ReadBatch(*msgs, 0)
+				if err != nil {
+					return 0, err
+				}
 			}
 		} else {
 			msg := &(*msgs)[0]
@@ -191,11 +206,12 @@ func (s *ICEBind) createIPv4ReceiverFn(ipv4MsgsPool *sync.Pool, pc *ipv4.PacketC
 			// todo: handle err
 			ok, _ := s.filterOutStunMessages(msg.Buffers, msg.N, msg.Addr)
 			if ok {
-				sizes[i] = 0
-			} else {
-				sizes[i] = msg.N
+				continue
+			}
+			sizes[i] = msg.N
+			if sizes[i] == 0 {
+				continue
 			}
-
 			addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
 			ep := &wgConn.StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation
 			wgConn.GetSrcFromControl(msg.OOB[:msg.NN], ep)
@@ -273,3 +289,15 @@ func fakeAddress(peerAddress *net.UDPAddr) (*net.UDPAddr, error) {
 	}
 	return newAddr, nil
 }
+
+func getMessages(msgsPool *sync.Pool) *[]ipv6.Message {
+	return msgsPool.Get().(*[]ipv6.Message)
+}
+
+func putMessages(msgs *[]ipv6.Message, msgsPool *sync.Pool) {
+	for i := range *msgs {
+		(*msgs)[i].OOB = (*msgs)[i].OOB[:0]
+		(*msgs)[i] = ipv6.Message{Buffers: (*msgs)[i].Buffers, OOB: (*msgs)[i].OOB}
+	}
+	msgsPool.Put(msgs)
+}

client/iface/bind/udp_mux.go

@@ -162,12 +162,13 @@ func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
 		params.Logger.Warn("UDPMuxDefault should not listening on unspecified address, use NewMultiUDPMuxFromPort instead")
 		var networks []ice.NetworkType
 		switch {
-		case addr.IP.To4() != nil:
-			networks = []ice.NetworkType{ice.NetworkTypeUDP4}
 
 		case addr.IP.To16() != nil:
 			networks = []ice.NetworkType{ice.NetworkTypeUDP4, ice.NetworkTypeUDP6}
 
+		case addr.IP.To4() != nil:
+			networks = []ice.NetworkType{ice.NetworkTypeUDP4}
+
 		default:
 			params.Logger.Errorf("LocalAddr expected IPV4 or IPV6, got %T", params.UDPConn.LocalAddr())
 		}

client/iface/device/kernel_module_linux.go

@@ -27,14 +27,14 @@ import (
 type status int
 
 const (
-	defaultModuleDir        = "/lib/modules"
-	unknown          status = iota
-	unloaded
-	unloading
-	loading
-	live
-	inuse
-	envDisableWireGuardKernel = "NB_WG_KERNEL_DISABLED"
+	unknown                   status = 1
+	unloaded                  status = 2
+	unloading                 status = 3
+	loading                   status = 4
+	live                      status = 5
+	inuse                     status = 6
+	defaultModuleDir                 = "/lib/modules"
+	envDisableWireGuardKernel        = "NB_WG_KERNEL_DISABLED"
 )
 
 type module struct {

client/iface/wgproxy/bind/proxy.go

@@ -2,6 +2,7 @@ package bind
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
@@ -94,7 +95,10 @@ func (p *ProxyBind) close() error {
 
 	p.Bind.RemoveEndpoint(p.wgAddr)
 
-	return p.remoteConn.Close()
+	if rErr := p.remoteConn.Close(); rErr != nil && !errors.Is(rErr, net.ErrClosed) {
+		return rErr
+	}
+	return nil
 }
 
 func (p *ProxyBind) proxyToLocal(ctx context.Context) {
@@ -104,8 +108,8 @@ func (p *ProxyBind) proxyToLocal(ctx context.Context) {
 		}
 	}()
 
-	buf := make([]byte, 1500)
 	for {
+		buf := make([]byte, 1500)
 		n, err := p.remoteConn.Read(buf)
 		if err != nil {
 			if ctx.Err() != nil {

client/iface/wgproxy/ebpf/wrapper.go

@@ -77,7 +77,7 @@ func (e *ProxyWrapper) CloseConn() error {
 
 	e.cancel()
 
-	if err := e.remoteConn.Close(); err != nil {
+	if err := e.remoteConn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
 		return fmt.Errorf("failed to close remote conn: %w", err)
 	}
 	return nil

client/iface/wgproxy/udp/proxy.go

@@ -116,7 +116,7 @@ func (p *WGUDPProxy) close() error {
 	p.cancel()
 
 	var result *multierror.Error
-	if err := p.remoteConn.Close(); err != nil {
+	if err := p.remoteConn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
 		result = multierror.Append(result, fmt.Errorf("remote conn: %s", err))
 	}
 

client/internal/config.go

@@ -46,6 +46,7 @@ type ConfigInput struct {
 	ManagementURL       string
 	AdminURL            string
 	ConfigPath          string
+	StateFilePath       string
 	PreSharedKey        *string
 	ServerSSHAllowed    *bool
 	NATExternalIPs      []string
@@ -105,10 +106,10 @@ type Config struct {
 
 	// DNSRouteInterval is the interval in which the DNS routes are updated
 	DNSRouteInterval time.Duration
-	//Path to a certificate used for mTLS authentication
+	// Path to a certificate used for mTLS authentication
 	ClientCertPath string
 
-	//Path to corresponding private key of ClientCertPath
+	// Path to corresponding private key of ClientCertPath
 	ClientCertKeyPath string
 
 	ClientCertKeyPair *tls.Certificate `json:"-"`
@@ -116,7 +117,7 @@ type Config struct {
 
 // ReadConfig read config file and return with Config. If it is not exists create a new with default values
 func ReadConfig(configPath string) (*Config, error) {
-	if configFileIsExists(configPath) {
+	if fileExists(configPath) {
 		err := util.EnforcePermission(configPath)
 		if err != nil {
 			log.Errorf("failed to enforce permission on config dir: %v", err)
@@ -149,7 +150,7 @@ func ReadConfig(configPath string) (*Config, error) {
 
 // UpdateConfig update existing configuration according to input configuration and return with the configuration
 func UpdateConfig(input ConfigInput) (*Config, error) {
-	if !configFileIsExists(input.ConfigPath) {
+	if !fileExists(input.ConfigPath) {
 		return nil, status.Errorf(codes.NotFound, "config file doesn't exist")
 	}
 
@@ -158,13 +159,13 @@ func UpdateConfig(input ConfigInput) (*Config, error) {
 
 // UpdateOrCreateConfig reads existing config or generates a new one
 func UpdateOrCreateConfig(input ConfigInput) (*Config, error) {
-	if !configFileIsExists(input.ConfigPath) {
+	if !fileExists(input.ConfigPath) {
 		log.Infof("generating new config %s", input.ConfigPath)
 		cfg, err := createNewConfig(input)
 		if err != nil {
 			return nil, err
 		}
-		err = util.WriteJsonWithRestrictedPermission(input.ConfigPath, cfg)
+		err = util.WriteJsonWithRestrictedPermission(context.Background(), input.ConfigPath, cfg)
 		return cfg, err
 	}
 
@@ -185,7 +186,7 @@ func CreateInMemoryConfig(input ConfigInput) (*Config, error) {
 
 // WriteOutConfig write put the prepared config to the given path
 func WriteOutConfig(path string, config *Config) error {
-	return util.WriteJson(path, config)
+	return util.WriteJson(context.Background(), path, config)
 }
 
 // createNewConfig creates a new config generating a new Wireguard key and saving to file
@@ -215,7 +216,7 @@ func update(input ConfigInput) (*Config, error) {
 	}
 
 	if updated {
-		if err := util.WriteJson(input.ConfigPath, config); err != nil {
+		if err := util.WriteJson(context.Background(), input.ConfigPath, config); err != nil {
 			return nil, err
 		}
 	}
@@ -472,11 +473,19 @@ func isPreSharedKeyHidden(preSharedKey *string) bool {
 	return false
 }
 
-func configFileIsExists(path string) bool {
+func fileExists(path string) bool {
 	_, err := os.Stat(path)
 	return !os.IsNotExist(err)
 }
 
+func createFile(path string) error {
+	file, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	return file.Close()
+}
+
 // UpdateOldManagementURL checks whether client can switch to the new Management URL with port 443 and the management domain.
 // If it can switch, then it updates the config and returns a new one. Otherwise, it returns the provided config.
 // The check is performed only for the NetBird's managed version.

client/internal/connect.go

@@ -40,6 +40,8 @@ type ConnectClient struct {
 	statusRecorder *peer.Status
 	engine         *Engine
 	engineMutex    sync.Mutex
+
+	persistNetworkMap bool
 }
 
 func NewConnectClient(
@@ -89,6 +91,7 @@ func (c *ConnectClient) RunOniOS(
 	fileDescriptor int32,
 	networkChangeListener listener.NetworkChangeListener,
 	dnsManager dns.IosDnsManager,
+	stateFilePath string,
 ) error {
 	// Set GC percent to 5% to reduce memory usage as iOS only allows 50MB of memory for the extension.
 	debug.SetGCPercent(5)
@@ -97,6 +100,7 @@ func (c *ConnectClient) RunOniOS(
 		FileDescriptor:        fileDescriptor,
 		NetworkChangeListener: networkChangeListener,
 		DnsManager:            dnsManager,
+		StateFilePath:         stateFilePath,
 	}
 	return c.run(mobileDependency, nil, nil)
 }
@@ -157,7 +161,8 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		engineCtx, cancel := context.WithCancel(c.ctx)
 		defer func() {
-			c.statusRecorder.MarkManagementDisconnected(state.err)
+			_, err := state.Status()
+			c.statusRecorder.MarkManagementDisconnected(err)
 			c.statusRecorder.CleanLocalPeerState()
 			cancel()
 		}()
@@ -207,7 +212,8 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		c.statusRecorder.MarkSignalDisconnected(nil)
 		defer func() {
-			c.statusRecorder.MarkSignalDisconnected(state.err)
+			_, err := state.Status()
+			c.statusRecorder.MarkSignalDisconnected(err)
 		}()
 
 		// with the global Wiretrustee config in hand connect (just a connection, no stream yet) Signal
@@ -230,6 +236,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		relayURLs, token := parseRelayInfo(loginResp)
 		relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String())
+		c.statusRecorder.SetRelayMgr(relayManager)
 		if len(relayURLs) > 0 {
 			if token != nil {
 				if err := relayManager.UpdateToken(token); err != nil {
@@ -240,9 +247,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 			log.Infof("connecting to the Relay service(s): %s", strings.Join(relayURLs, ", "))
 			if err = relayManager.Serve(); err != nil {
 				log.Error(err)
-				return wrapErr(err)
 			}
-			c.statusRecorder.SetRelayMgr(relayManager)
 		}
 
 		peerConfig := loginResp.GetPeerConfig()
@@ -257,7 +262,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		c.engineMutex.Lock()
 		c.engine = NewEngineWithProbes(engineCtx, cancel, signalClient, mgmClient, relayManager, engineConfig, mobileDependency, c.statusRecorder, probes, checks)
-
+		c.engine.SetNetworkMapPersistence(c.persistNetworkMap)
 		c.engineMutex.Unlock()
 
 		if err := c.engine.Start(); err != nil {
@@ -335,6 +340,19 @@ func (c *ConnectClient) Engine() *Engine {
 	return e
 }
 
+// Status returns the current client status
+func (c *ConnectClient) Status() StatusType {
+	if c == nil {
+		return StatusIdle
+	}
+	status, err := CtxGetState(c.ctx).Status()
+	if err != nil {
+		return StatusIdle
+	}
+
+	return status
+}
+
 func (c *ConnectClient) Stop() error {
 	if c == nil {
 		return nil
@@ -361,6 +379,22 @@ func (c *ConnectClient) isContextCancelled() bool {
 	}
 }
 
+// SetNetworkMapPersistence enables or disables network map persistence.
+// When enabled, the last received network map will be stored and can be retrieved
+// through the Engine's getLatestNetworkMap method. When disabled, any stored
+// network map will be cleared. This functionality is primarily used for debugging
+// and should not be enabled during normal operation.
+func (c *ConnectClient) SetNetworkMapPersistence(enabled bool) {
+	c.engineMutex.Lock()
+	c.persistNetworkMap = enabled
+	c.engineMutex.Unlock()
+
+	engine := c.Engine()
+	if engine != nil {
+		engine.SetNetworkMapPersistence(enabled)
+	}
+}
+
 // createEngineConfig converts configuration received from Management Service to EngineConfig
 func createEngineConfig(key wgtypes.Key, config *Config, peerConfig *mgmProto.PeerConfig) (*EngineConfig, error) {
 	nm := false

client/internal/dns/handler_chain.go

@@ -0,0 +1,226 @@
+package dns
+
+import (
+	"slices"
+	"strings"
+	"sync"
+
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	PriorityDNSRoute    = 100
+	PriorityMatchDomain = 50
+	PriorityDefault     = 0
+)
+
+type SubdomainMatcher interface {
+	dns.Handler
+	MatchSubdomains() bool
+}
+
+type HandlerEntry struct {
+	Handler         dns.Handler
+	Priority        int
+	Pattern         string
+	OrigPattern     string
+	IsWildcard      bool
+	StopHandler     handlerWithStop
+	MatchSubdomains bool
+}
+
+// HandlerChain represents a prioritized chain of DNS handlers
+type HandlerChain struct {
+	mu       sync.RWMutex
+	handlers []HandlerEntry
+}
+
+// ResponseWriterChain wraps a dns.ResponseWriter to track if handler wants to continue chain
+type ResponseWriterChain struct {
+	dns.ResponseWriter
+	origPattern    string
+	shouldContinue bool
+}
+
+func (w *ResponseWriterChain) WriteMsg(m *dns.Msg) error {
+	// Check if this is a continue signal (NXDOMAIN with Zero bit set)
+	if m.Rcode == dns.RcodeNameError && m.MsgHdr.Zero {
+		w.shouldContinue = true
+		return nil
+	}
+	return w.ResponseWriter.WriteMsg(m)
+}
+
+func NewHandlerChain() *HandlerChain {
+	return &HandlerChain{
+		handlers: make([]HandlerEntry, 0),
+	}
+}
+
+// GetOrigPattern returns the original pattern of the handler that wrote the response
+func (w *ResponseWriterChain) GetOrigPattern() string {
+	return w.origPattern
+}
+
+// AddHandler adds a new handler to the chain, replacing any existing handler with the same pattern and priority
+func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority int, stopHandler handlerWithStop) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	origPattern := pattern
+	isWildcard := strings.HasPrefix(pattern, "*.")
+	if isWildcard {
+		pattern = pattern[2:]
+	}
+	pattern = dns.Fqdn(pattern)
+	origPattern = dns.Fqdn(origPattern)
+
+	// First remove any existing handler with same original pattern and priority
+	for i := len(c.handlers) - 1; i >= 0; i-- {
+		if c.handlers[i].OrigPattern == origPattern && c.handlers[i].Priority == priority {
+			if c.handlers[i].StopHandler != nil {
+				c.handlers[i].StopHandler.stop()
+			}
+			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
+			break
+		}
+	}
+
+	// Check if handler implements SubdomainMatcher interface
+	matchSubdomains := false
+	if matcher, ok := handler.(SubdomainMatcher); ok {
+		matchSubdomains = matcher.MatchSubdomains()
+	}
+
+	log.Debugf("adding handler pattern: domain=%s original: domain=%s wildcard=%v match_subdomain=%v priority=%d",
+		pattern, origPattern, isWildcard, matchSubdomains, priority)
+
+	entry := HandlerEntry{
+		Handler:         handler,
+		Priority:        priority,
+		Pattern:         pattern,
+		OrigPattern:     origPattern,
+		IsWildcard:      isWildcard,
+		StopHandler:     stopHandler,
+		MatchSubdomains: matchSubdomains,
+	}
+
+	// Insert handler in priority order
+	pos := 0
+	for i, h := range c.handlers {
+		if h.Priority < priority {
+			pos = i
+			break
+		}
+		pos = i + 1
+	}
+
+	c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...)
+}
+
+// RemoveHandler removes a handler for the given pattern and priority
+func (c *HandlerChain) RemoveHandler(pattern string, priority int) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	pattern = dns.Fqdn(pattern)
+
+	// Find and remove handlers matching both original pattern and priority
+	for i := len(c.handlers) - 1; i >= 0; i-- {
+		entry := c.handlers[i]
+		if entry.OrigPattern == pattern && entry.Priority == priority {
+			if entry.StopHandler != nil {
+				entry.StopHandler.stop()
+			}
+			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
+			return
+		}
+	}
+}
+
+// HasHandlers returns true if there are any handlers remaining for the given pattern
+func (c *HandlerChain) HasHandlers(pattern string) bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	pattern = dns.Fqdn(pattern)
+	for _, entry := range c.handlers {
+		if entry.Pattern == pattern {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *HandlerChain) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	if len(r.Question) == 0 {
+		return
+	}
+
+	qname := r.Question[0].Name
+	log.Tracef("handling DNS request for domain=%s", qname)
+
+	c.mu.RLock()
+	handlers := slices.Clone(c.handlers)
+	c.mu.RUnlock()
+
+	if log.IsLevelEnabled(log.TraceLevel) {
+		log.Tracef("current handlers (%d):", len(handlers))
+		for _, h := range handlers {
+			log.Tracef("  - pattern: domain=%s original: domain=%s wildcard=%v priority=%d",
+				h.Pattern, h.OrigPattern, h.IsWildcard, h.Priority)
+		}
+	}
+
+	// Try handlers in priority order
+	for _, entry := range handlers {
+		var matched bool
+		switch {
+		case entry.Pattern == ".":
+			matched = true
+		case entry.IsWildcard:
+			parts := strings.Split(strings.TrimSuffix(qname, entry.Pattern), ".")
+			matched = len(parts) >= 2 && strings.HasSuffix(qname, entry.Pattern)
+		default:
+			// For non-wildcard patterns:
+			// If handler wants subdomain matching, allow suffix match
+			// Otherwise require exact match
+			if entry.MatchSubdomains {
+				matched = qname == entry.Pattern || strings.HasSuffix(qname, "."+entry.Pattern)
+			} else {
+				matched = qname == entry.Pattern
+			}
+		}
+
+		if !matched {
+			log.Tracef("trying domain match: request: domain=%s pattern: domain=%s wildcard=%v match_subdomain=%v matched=false",
+				qname, entry.OrigPattern, entry.MatchSubdomains, entry.IsWildcard)
+			continue
+		}
+
+		log.Tracef("handler matched: request: domain=%s pattern: domain=%s wildcard=%v match_subdomain=%v",
+			qname, entry.OrigPattern, entry.IsWildcard, entry.MatchSubdomains)
+
+		chainWriter := &ResponseWriterChain{
+			ResponseWriter: w,
+			origPattern:    entry.OrigPattern,
+		}
+		entry.Handler.ServeDNS(chainWriter, r)
+
+		// If handler wants to continue, try next handler
+		if chainWriter.shouldContinue {
+			log.Tracef("handler requested continue to next handler")
+			continue
+		}
+		return
+	}
+
+	// No handler matched or all handlers passed
+	log.Tracef("no handler found for domain=%s", qname)
+	resp := &dns.Msg{}
+	resp.SetRcode(r, dns.RcodeNameError)
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed to write DNS response: %v", err)
+	}
+}

client/internal/dns/handler_chain_test.go

@@ -0,0 +1,511 @@
+package dns_test
+
+import (
+	"net"
+	"testing"
+
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	nbdns "github.com/netbirdio/netbird/client/internal/dns"
+)
+
+// TestHandlerChain_ServeDNS_Priorities tests that handlers are executed in priority order
+func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	// Create mock handlers for different priorities
+	defaultHandler := &nbdns.MockHandler{}
+	matchDomainHandler := &nbdns.MockHandler{}
+	dnsRouteHandler := &nbdns.MockHandler{}
+
+	// Setup handlers with different priorities
+	chain.AddHandler("example.com.", defaultHandler, nbdns.PriorityDefault, nil)
+	chain.AddHandler("example.com.", matchDomainHandler, nbdns.PriorityMatchDomain, nil)
+	chain.AddHandler("example.com.", dnsRouteHandler, nbdns.PriorityDNSRoute, nil)
+
+	// Create test request
+	r := new(dns.Msg)
+	r.SetQuestion("example.com.", dns.TypeA)
+
+	// Create test writer
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+	// Setup expectations - only highest priority handler should be called
+	dnsRouteHandler.On("ServeDNS", mock.Anything, r).Once()
+	matchDomainHandler.On("ServeDNS", mock.Anything, r).Maybe()
+	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe()
+
+	// Execute
+	chain.ServeDNS(w, r)
+
+	// Verify all expectations were met
+	dnsRouteHandler.AssertExpectations(t)
+	matchDomainHandler.AssertExpectations(t)
+	defaultHandler.AssertExpectations(t)
+}
+
+// TestHandlerChain_ServeDNS_DomainMatching tests various domain matching scenarios
+func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
+	tests := []struct {
+		name            string
+		handlerDomain   string
+		queryDomain     string
+		isWildcard      bool
+		matchSubdomains bool
+		shouldMatch     bool
+	}{
+		{
+			name:            "exact match",
+			handlerDomain:   "example.com.",
+			queryDomain:     "example.com.",
+			isWildcard:      false,
+			matchSubdomains: false,
+			shouldMatch:     true,
+		},
+		{
+			name:            "subdomain with non-wildcard and MatchSubdomains true",
+			handlerDomain:   "example.com.",
+			queryDomain:     "sub.example.com.",
+			isWildcard:      false,
+			matchSubdomains: true,
+			shouldMatch:     true,
+		},
+		{
+			name:            "subdomain with non-wildcard and MatchSubdomains false",
+			handlerDomain:   "example.com.",
+			queryDomain:     "sub.example.com.",
+			isWildcard:      false,
+			matchSubdomains: false,
+			shouldMatch:     false,
+		},
+		{
+			name:            "wildcard match",
+			handlerDomain:   "*.example.com.",
+			queryDomain:     "sub.example.com.",
+			isWildcard:      true,
+			matchSubdomains: false,
+			shouldMatch:     true,
+		},
+		{
+			name:            "wildcard no match on apex",
+			handlerDomain:   "*.example.com.",
+			queryDomain:     "example.com.",
+			isWildcard:      true,
+			matchSubdomains: false,
+			shouldMatch:     false,
+		},
+		{
+			name:            "root zone match",
+			handlerDomain:   ".",
+			queryDomain:     "anything.com.",
+			isWildcard:      false,
+			matchSubdomains: false,
+			shouldMatch:     true,
+		},
+		{
+			name:            "no match different domain",
+			handlerDomain:   "example.com.",
+			queryDomain:     "example.org.",
+			isWildcard:      false,
+			matchSubdomains: false,
+			shouldMatch:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			var handler dns.Handler
+
+			if tt.matchSubdomains {
+				mockSubHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
+				handler = mockSubHandler
+				if tt.shouldMatch {
+					mockSubHandler.On("ServeDNS", mock.Anything, mock.Anything).Once()
+				}
+			} else {
+				mockHandler := &nbdns.MockHandler{}
+				handler = mockHandler
+				if tt.shouldMatch {
+					mockHandler.On("ServeDNS", mock.Anything, mock.Anything).Once()
+				}
+			}
+
+			pattern := tt.handlerDomain
+			if tt.isWildcard {
+				pattern = "*." + tt.handlerDomain[2:]
+			}
+
+			chain.AddHandler(pattern, handler, nbdns.PriorityDefault, nil)
+
+			r := new(dns.Msg)
+			r.SetQuestion(tt.queryDomain, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			chain.ServeDNS(w, r)
+
+			if h, ok := handler.(*nbdns.MockHandler); ok {
+				h.AssertExpectations(t)
+			} else if h, ok := handler.(*nbdns.MockSubdomainHandler); ok {
+				h.AssertExpectations(t)
+			}
+		})
+	}
+}
+
+// TestHandlerChain_ServeDNS_OverlappingDomains tests behavior with overlapping domain patterns
+func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
+	tests := []struct {
+		name     string
+		handlers []struct {
+			pattern  string
+			priority int
+		}
+		queryDomain     string
+		expectedCalls   int
+		expectedHandler int // index of the handler that should be called
+	}{
+		{
+			name: "wildcard and exact same priority - exact should win",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "example.com.", priority: nbdns.PriorityDefault},
+			},
+			queryDomain:     "example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // exact match handler should be called
+		},
+		{
+			name: "higher priority wildcard over lower priority exact",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // higher priority wildcard handler should be called
+		},
+		{
+			name: "multiple wildcards different priorities",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "*.example.com.", priority: nbdns.PriorityMatchDomain},
+				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 2, // highest priority handler should be called
+		},
+		{
+			name: "subdomain with mix of patterns",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "test.example.com.", priority: nbdns.PriorityMatchDomain},
+				{pattern: "*.test.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "sub.test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 2, // highest priority matching handler should be called
+		},
+		{
+			name: "root zone with specific domain",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: ".", priority: nbdns.PriorityDefault},
+				{pattern: "example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // higher priority specific domain should win over root
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			var handlers []*nbdns.MockHandler
+
+			// Setup handlers and expectations
+			for i := range tt.handlers {
+				handler := &nbdns.MockHandler{}
+				handlers = append(handlers, handler)
+
+				// Set expectation based on whether this handler should be called
+				if i == tt.expectedHandler {
+					handler.On("ServeDNS", mock.Anything, mock.Anything).Once()
+				} else {
+					handler.On("ServeDNS", mock.Anything, mock.Anything).Maybe()
+				}
+
+				chain.AddHandler(tt.handlers[i].pattern, handler, tt.handlers[i].priority, nil)
+			}
+
+			// Create and execute request
+			r := new(dns.Msg)
+			r.SetQuestion(tt.queryDomain, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			chain.ServeDNS(w, r)
+
+			// Verify expectations
+			for _, handler := range handlers {
+				handler.AssertExpectations(t)
+			}
+		})
+	}
+}
+
+// TestHandlerChain_ServeDNS_ChainContinuation tests the chain continuation functionality
+func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	// Create handlers
+	handler1 := &nbdns.MockHandler{}
+	handler2 := &nbdns.MockHandler{}
+	handler3 := &nbdns.MockHandler{}
+
+	// Add handlers in priority order
+	chain.AddHandler("example.com.", handler1, nbdns.PriorityDNSRoute, nil)
+	chain.AddHandler("example.com.", handler2, nbdns.PriorityMatchDomain, nil)
+	chain.AddHandler("example.com.", handler3, nbdns.PriorityDefault, nil)
+
+	// Create test request
+	r := new(dns.Msg)
+	r.SetQuestion("example.com.", dns.TypeA)
+
+	// Setup mock responses to simulate chain continuation
+	handler1.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// First handler signals continue
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeNameError)
+		resp.MsgHdr.Zero = true // Signal to continue
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	handler2.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// Second handler signals continue
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeNameError)
+		resp.MsgHdr.Zero = true
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	handler3.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// Last handler responds normally
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeSuccess)
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	// Execute
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	chain.ServeDNS(w, r)
+
+	// Verify all handlers were called in order
+	handler1.AssertExpectations(t)
+	handler2.AssertExpectations(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
+		ops  []struct {
+			action   string // "add" or "remove"
+			pattern  string
+			priority int
+		}
+		query         string
+		expectedCalls map[int]bool // map[priority]shouldBeCalled
+	}{
+		{
+			name: "remove high priority keeps lower priority handler",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"remove", "example.com.", nbdns.PriorityDNSRoute},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    false,
+				nbdns.PriorityMatchDomain: true,
+			},
+		},
+		{
+			name: "remove lower priority keeps high priority handler",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"remove", "example.com.", nbdns.PriorityMatchDomain},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    true,
+				nbdns.PriorityMatchDomain: false,
+			},
+		},
+		{
+			name: "remove all handlers in order",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"add", "example.com.", nbdns.PriorityDefault},
+				{"remove", "example.com.", nbdns.PriorityDNSRoute},
+				{"remove", "example.com.", nbdns.PriorityMatchDomain},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    false,
+				nbdns.PriorityMatchDomain: false,
+				nbdns.PriorityDefault:     true,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			handlers := make(map[int]*nbdns.MockHandler)
+
+			// Execute operations
+			for _, op := range tt.ops {
+				if op.action == "add" {
+					handler := &nbdns.MockHandler{}
+					handlers[op.priority] = handler
+					chain.AddHandler(op.pattern, handler, op.priority, nil)
+				} else {
+					chain.RemoveHandler(op.pattern, op.priority)
+				}
+			}
+
+			// Create test request
+			r := new(dns.Msg)
+			r.SetQuestion(tt.query, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			// Setup expectations
+			for priority, handler := range handlers {
+				if shouldCall, exists := tt.expectedCalls[priority]; exists && shouldCall {
+					handler.On("ServeDNS", mock.Anything, r).Once()
+				} else {
+					handler.On("ServeDNS", mock.Anything, r).Maybe()
+				}
+			}
+
+			// Execute request
+			chain.ServeDNS(w, r)
+
+			// Verify expectations
+			for _, handler := range handlers {
+				handler.AssertExpectations(t)
+			}
+
+			// Verify handler exists check
+			for priority, shouldExist := range tt.expectedCalls {
+				if shouldExist {
+					assert.True(t, chain.HasHandlers(tt.ops[0].pattern),
+						"Handler chain should have handlers for pattern after removing priority %d", priority)
+				}
+			}
+		})
+	}
+}
+
+func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	testDomain := "example.com."
+	testQuery := "test.example.com."
+
+	// Create handlers with MatchSubdomains enabled
+	routeHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
+	matchHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
+	defaultHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
+
+	// Create test request that will be reused
+	r := new(dns.Msg)
+	r.SetQuestion(testQuery, dns.TypeA)
+
+	// Add handlers in mixed order
+	chain.AddHandler(testDomain, defaultHandler, nbdns.PriorityDefault, nil)
+	chain.AddHandler(testDomain, routeHandler, nbdns.PriorityDNSRoute, nil)
+	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain, nil)
+
+	// Test 1: Initial state with all three handlers
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Highest priority handler (routeHandler) should be called
+	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	routeHandler.AssertExpectations(t)
+
+	// Test 2: Remove highest priority handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)
+	assert.True(t, chain.HasHandlers(testDomain))
+
+	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Now middle priority handler (matchHandler) should be called
+	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	matchHandler.AssertExpectations(t)
+
+	// Test 3: Remove middle priority handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityMatchDomain)
+	assert.True(t, chain.HasHandlers(testDomain))
+
+	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Now lowest priority handler (defaultHandler) should be called
+	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	defaultHandler.AssertExpectations(t)
+
+	// Test 4: Remove last handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
+	assert.False(t, chain.HasHandlers(testDomain))
+}

client/internal/dns/local.go

@@ -17,12 +17,24 @@ type localResolver struct {
 	records       sync.Map
 }
 
+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))
+}
+
 // ServeDNS handles a DNS request
 func (d *localResolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
-	log.Tracef("received question: %#v", r.Question[0])
+	if len(r.Question) > 0 {
+		log.Tracef("received 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

client/internal/dns/mock_server.go

@@ -3,14 +3,30 @@ package dns
 import (
 	"fmt"
 
+	"github.com/miekg/dns"
+
 	nbdns "github.com/netbirdio/netbird/dns"
 )
 
 // MockServer is the mock instance of a dns server
 type MockServer struct {
-	InitializeFunc      func() error
-	StopFunc            func()
-	UpdateDNSServerFunc func(serial uint64, update nbdns.Config) error
+	InitializeFunc        func() error
+	StopFunc              func()
+	UpdateDNSServerFunc   func(serial uint64, update nbdns.Config) error
+	RegisterHandlerFunc   func([]string, dns.Handler, int)
+	DeregisterHandlerFunc func([]string, int)
+}
+
+func (m *MockServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+	if m.RegisterHandlerFunc != nil {
+		m.RegisterHandlerFunc(domains, handler, priority)
+	}
+}
+
+func (m *MockServer) DeregisterHandler(domains []string, priority int) {
+	if m.DeregisterHandlerFunc != nil {
+		m.DeregisterHandlerFunc(domains, priority)
+	}
 }
 
 // Initialize mock implementation of Initialize from Server interface

client/internal/dns/server.go

@@ -7,7 +7,6 @@ import (
 	"runtime"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/miekg/dns"
 	"github.com/mitchellh/hashstructure/v2"
@@ -31,6 +30,8 @@ type IosDnsManager interface {
 
 // Server is a dns server interface
 type Server interface {
+	RegisterHandler(domains []string, handler dns.Handler, priority int)
+	DeregisterHandler(domains []string, priority int)
 	Initialize() error
 	Stop()
 	DnsIP() string
@@ -49,12 +50,14 @@ type DefaultServer struct {
 	mux                sync.Mutex
 	service            service
 	dnsMuxMap          registeredHandlerMap
+	handlerPriorities  map[string]int
 	localResolver      *localResolver
 	wgInterface        WGIface
 	hostManager        hostManager
 	updateSerial       uint64
 	previousConfigHash uint64
 	currentConfig      HostDNSConfig
+	handlerChain       *HandlerChain
 
 	// permanent related properties
 	permanent      bool
@@ -75,8 +78,9 @@ type handlerWithStop interface {
 }
 
 type muxUpdate struct {
-	domain  string
-	handler handlerWithStop
+	domain   string
+	handler  handlerWithStop
+	priority int
 }
 
 // NewDefaultServer returns a new dns server
@@ -136,10 +140,12 @@ func NewDefaultServerIos(
 func newDefaultServer(ctx context.Context, wgInterface WGIface, dnsService service, statusRecorder *peer.Status, stateManager *statemanager.Manager) *DefaultServer {
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
-		ctx:       ctx,
-		ctxCancel: stop,
-		service:   dnsService,
-		dnsMuxMap: make(registeredHandlerMap),
+		ctx:               ctx,
+		ctxCancel:         stop,
+		service:           dnsService,
+		handlerChain:      NewHandlerChain(),
+		dnsMuxMap:         make(registeredHandlerMap),
+		handlerPriorities: make(map[string]int),
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
@@ -152,6 +158,51 @@ func newDefaultServer(ctx context.Context, wgInterface WGIface, dnsService servi
 	return defaultServer
 }
 
+func (s *DefaultServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	s.registerHandler(domains, handler, priority)
+}
+
+func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, priority int) {
+	log.Debugf("registering handler %s with priority %d", handler, priority)
+
+	for _, domain := range domains {
+		if domain == "" {
+			log.Warn("skipping empty domain")
+			continue
+		}
+		s.handlerChain.AddHandler(domain, handler, priority, nil)
+		s.handlerPriorities[domain] = priority
+		s.service.RegisterMux(nbdns.NormalizeZone(domain), s.handlerChain)
+	}
+}
+
+func (s *DefaultServer) DeregisterHandler(domains []string, priority int) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	s.deregisterHandler(domains, priority)
+}
+
+func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
+	log.Debugf("deregistering handler %v with priority %d", domains, priority)
+
+	for _, domain := range domains {
+		s.handlerChain.RemoveHandler(domain, priority)
+
+		// Only deregister from service if no handlers remain
+		if !s.handlerChain.HasHandlers(domain) {
+			if domain == "" {
+				log.Warn("skipping empty domain")
+				continue
+			}
+			s.service.DeregisterMux(nbdns.NormalizeZone(domain))
+		}
+	}
+}
+
 // Initialize instantiate host manager and the dns service
 func (s *DefaultServer) Initialize() (err error) {
 	s.mux.Lock()
@@ -323,12 +374,12 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 		log.Error(err)
 	}
 
-	// persist dns state right away
-	ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
-	defer cancel()
-	if err := s.stateManager.PersistState(ctx); err != nil {
-		log.Errorf("Failed to persist dns state: %v", err)
-	}
+	go func() {
+		// persist dns state right away
+		if err := s.stateManager.PersistState(s.ctx); err != nil {
+			log.Errorf("Failed to persist dns state: %v", err)
+		}
+	}()
 
 	if s.searchDomainNotifier != nil {
 		s.searchDomainNotifier.onNewSearchDomains(s.SearchDomains())
@@ -344,14 +395,14 @@ func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone)
 	localRecords := make(map[string]nbdns.SimpleRecord, 0)
 
 	for _, customZone := range customZones {
-
 		if len(customZone.Records) == 0 {
 			return nil, nil, fmt.Errorf("received an empty list of records")
 		}
 
 		muxUpdates = append(muxUpdates, muxUpdate{
-			domain:  customZone.Domain,
-			handler: s.localResolver,
+			domain:   customZone.Domain,
+			handler:  s.localResolver,
+			priority: PriorityMatchDomain,
 		})
 
 		for _, record := range customZone.Records {
@@ -413,8 +464,9 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 
 		if nsGroup.Primary {
 			muxUpdates = append(muxUpdates, muxUpdate{
-				domain:  nbdns.RootZone,
-				handler: handler,
+				domain:   nbdns.RootZone,
+				handler:  handler,
+				priority: PriorityDefault,
 			})
 			continue
 		}
@@ -430,8 +482,9 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 				return nil, fmt.Errorf("received a nameserver group with an empty domain element")
 			}
 			muxUpdates = append(muxUpdates, muxUpdate{
-				domain:  domain,
-				handler: handler,
+				domain:   domain,
+				handler:  handler,
+				priority: PriorityMatchDomain,
 			})
 		}
 	}
@@ -441,12 +494,16 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 
 func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 	muxUpdateMap := make(registeredHandlerMap)
+	handlersByPriority := make(map[string]int)
 
 	var isContainRootUpdate bool
 
+	// First register new handlers
 	for _, update := range muxUpdates {
-		s.service.RegisterMux(update.domain, update.handler)
+		s.registerHandler([]string{update.domain}, update.handler, update.priority)
 		muxUpdateMap[update.domain] = update.handler
+		handlersByPriority[update.domain] = update.priority
+
 		if existingHandler, ok := s.dnsMuxMap[update.domain]; ok {
 			existingHandler.stop()
 		}
@@ -456,6 +513,7 @@ func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 		}
 	}
 
+	// Then deregister old handlers not in the update
 	for key, existingHandler := range s.dnsMuxMap {
 		_, found := muxUpdateMap[key]
 		if !found {
@@ -464,12 +522,16 @@ func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 				existingHandler.stop()
 			} else {
 				existingHandler.stop()
-				s.service.DeregisterMux(key)
+				// Deregister with the priority that was used to register
+				if oldPriority, ok := s.handlerPriorities[key]; ok {
+					s.deregisterHandler([]string{key}, oldPriority)
+				}
 			}
 		}
 	}
 
 	s.dnsMuxMap = muxUpdateMap
+	s.handlerPriorities = handlersByPriority
 }
 
 func (s *DefaultServer) updateLocalResolver(update map[string]nbdns.SimpleRecord) {
@@ -518,13 +580,13 @@ func (s *DefaultServer) upstreamCallbacks(
 		if nsGroup.Primary {
 			removeIndex[nbdns.RootZone] = -1
 			s.currentConfig.RouteAll = false
-			s.service.DeregisterMux(nbdns.RootZone)
+			s.deregisterHandler([]string{nbdns.RootZone}, PriorityDefault)
 		}
 
 		for i, item := range s.currentConfig.Domains {
 			if _, found := removeIndex[item.Domain]; found {
 				s.currentConfig.Domains[i].Disabled = true
-				s.service.DeregisterMux(item.Domain)
+				s.deregisterHandler([]string{item.Domain}, PriorityMatchDomain)
 				removeIndex[item.Domain] = i
 			}
 		}
@@ -533,12 +595,11 @@ func (s *DefaultServer) upstreamCallbacks(
 			l.Errorf("Failed to apply nameserver deactivation on the host: %v", err)
 		}
 
-		// persist dns state right away
-		ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
-		defer cancel()
-		if err := s.stateManager.PersistState(ctx); err != nil {
-			l.Errorf("Failed to persist dns state: %v", err)
-		}
+		go func() {
+			if err := s.stateManager.PersistState(s.ctx); err != nil {
+				l.Errorf("Failed to persist dns state: %v", err)
+			}
+		}()
 
 		if runtime.GOOS == "android" && nsGroup.Primary && len(s.hostsDNSHolder.get()) > 0 {
 			s.addHostRootZone()
@@ -556,7 +617,7 @@ func (s *DefaultServer) upstreamCallbacks(
 				continue
 			}
 			s.currentConfig.Domains[i].Disabled = false
-			s.service.RegisterMux(domain, handler)
+			s.registerHandler([]string{domain}, handler, PriorityMatchDomain)
 		}
 
 		l := log.WithField("nameservers", nsGroup.NameServers)
@@ -564,7 +625,7 @@ func (s *DefaultServer) upstreamCallbacks(
 
 		if nsGroup.Primary {
 			s.currentConfig.RouteAll = true
-			s.service.RegisterMux(nbdns.RootZone, handler)
+			s.registerHandler([]string{nbdns.RootZone}, handler, PriorityDefault)
 		}
 		if err := s.hostManager.applyDNSConfig(s.currentConfig, s.stateManager); err != nil {
 			l.WithError(err).Error("reactivate temporary disabled nameserver group, DNS update apply")
@@ -595,7 +656,8 @@ func (s *DefaultServer) addHostRootZone() {
 	}
 	handler.deactivate = func(error) {}
 	handler.reactivate = func() {}
-	s.service.RegisterMux(nbdns.RootZone, handler)
+
+	s.registerHandler([]string{nbdns.RootZone}, handler, PriorityDefault)
 }
 
 func (s *DefaultServer) updateNSGroupStates(groups []*nbdns.NameServerGroup) {

client/internal/dns/server_test.go

@@ -11,7 +11,9 @@ import (
 	"time"
 
 	"github.com/golang/mock/gomock"
+	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/mock"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/firewall/uspfilter"
@@ -512,7 +514,7 @@ func TestDNSServerStartStop(t *testing.T) {
 				t.Error(err)
 			}
 
-			dnsServer.service.RegisterMux("netbird.cloud", dnsServer.localResolver)
+			dnsServer.registerHandler([]string{"netbird.cloud"}, dnsServer.localResolver, 1)
 
 			resolver := &net.Resolver{
 				PreferGo: true,
@@ -560,7 +562,9 @@ func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
-		hostManager: hostManager,
+		handlerChain:      NewHandlerChain(),
+		handlerPriorities: make(map[string]int),
+		hostManager:       hostManager,
 		currentConfig: HostDNSConfig{
 			Domains: []DomainConfig{
 				{false, "domain0", false},
@@ -782,7 +786,7 @@ func TestDNSPermanent_matchOnly(t *testing.T) {
 						Port:   53,
 					},
 				},
-				Domains: []string{"customdomain.com"},
+				Domains: []string{"google.com"},
 				Primary: false,
 			},
 		},
@@ -804,7 +808,7 @@ func TestDNSPermanent_matchOnly(t *testing.T) {
 	if ips[0] != zoneRecords[0].RData {
 		t.Fatalf("invalid zone record: %v", err)
 	}
-	_, err = resolver.LookupHost(context.Background(), "customdomain.com")
+	_, err = resolver.LookupHost(context.Background(), "google.com")
 	if err != nil {
 		t.Errorf("failed to resolve: %s", err)
 	}
@@ -872,3 +876,86 @@ func newDnsResolver(ip string, port int) *net.Resolver {
 		},
 	}
 }
+
+// MockHandler implements dns.Handler interface for testing
+type MockHandler struct {
+	mock.Mock
+}
+
+func (m *MockHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	m.Called(w, r)
+}
+
+type MockSubdomainHandler struct {
+	MockHandler
+	Subdomains bool
+}
+
+func (m *MockSubdomainHandler) MatchSubdomains() bool {
+	return m.Subdomains
+}
+
+func TestHandlerChain_DomainPriorities(t *testing.T) {
+	chain := NewHandlerChain()
+
+	dnsRouteHandler := &MockHandler{}
+	upstreamHandler := &MockSubdomainHandler{
+		Subdomains: true,
+	}
+
+	chain.AddHandler("example.com.", dnsRouteHandler, PriorityDNSRoute, nil)
+	chain.AddHandler("example.com.", upstreamHandler, PriorityMatchDomain, nil)
+
+	testCases := []struct {
+		name            string
+		query           string
+		expectedHandler dns.Handler
+	}{
+		{
+			name:            "exact domain with dns route handler",
+			query:           "example.com.",
+			expectedHandler: dnsRouteHandler,
+		},
+		{
+			name:            "subdomain should use upstream handler",
+			query:           "sub.example.com.",
+			expectedHandler: upstreamHandler,
+		},
+		{
+			name:            "deep subdomain should use upstream handler",
+			query:           "deep.sub.example.com.",
+			expectedHandler: upstreamHandler,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			r := new(dns.Msg)
+			r.SetQuestion(tc.query, dns.TypeA)
+			w := &ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			if mh, ok := tc.expectedHandler.(*MockHandler); ok {
+				mh.On("ServeDNS", mock.Anything, r).Once()
+			} else if mh, ok := tc.expectedHandler.(*MockSubdomainHandler); ok {
+				mh.On("ServeDNS", mock.Anything, r).Once()
+			}
+
+			chain.ServeDNS(w, r)
+
+			if mh, ok := tc.expectedHandler.(*MockHandler); ok {
+				mh.AssertExpectations(t)
+			} else if mh, ok := tc.expectedHandler.(*MockSubdomainHandler); ok {
+				mh.AssertExpectations(t)
+			}
+
+			// Reset mocks
+			if mh, ok := tc.expectedHandler.(*MockHandler); ok {
+				mh.ExpectedCalls = nil
+				mh.Calls = nil
+			} else if mh, ok := tc.expectedHandler.(*MockSubdomainHandler); ok {
+				mh.ExpectedCalls = nil
+				mh.Calls = nil
+			}
+		})
+	}
+}

client/internal/dns/service_listener.go

@@ -105,6 +105,7 @@ func (s *serviceViaListener) Stop() {
 }
 
 func (s *serviceViaListener) RegisterMux(pattern string, handler dns.Handler) {
+	log.Debugf("registering dns handler for pattern: %s", pattern)
 	s.dnsMux.Handle(pattern, handler)
 }
 

client/internal/dns/upstream.go

@@ -66,6 +66,15 @@ func newUpstreamResolverBase(ctx context.Context, statusRecorder *peer.Status) *
 	}
 }
 
+// String returns a string representation of the upstream resolver
+func (u *upstreamResolverBase) String() string {
+	return fmt.Sprintf("upstream %v", u.upstreamServers)
+}
+
+func (u *upstreamResolverBase) MatchSubdomains() bool {
+	return true
+}
+
 func (u *upstreamResolverBase) stop() {
 	log.Debugf("stopping serving DNS for upstreams %s", u.upstreamServers)
 	u.cancel()

client/internal/dnsfwd/forwarder.go

@@ -0,0 +1,157 @@
+package dnsfwd
+
+import (
+	"context"
+	"errors"
+	"net"
+
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+const errResolveFailed = "failed to resolve query for domain=%s: %v"
+
+type DNSForwarder struct {
+	listenAddress string
+	ttl           uint32
+	domains       []string
+
+	dnsServer *dns.Server
+	mux       *dns.ServeMux
+}
+
+func NewDNSForwarder(listenAddress string, ttl uint32) *DNSForwarder {
+	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
+	return &DNSForwarder{
+		listenAddress: listenAddress,
+		ttl:           ttl,
+	}
+}
+
+func (f *DNSForwarder) Listen(domains []string) error {
+	log.Infof("listen DNS forwarder on address=%s", f.listenAddress)
+	mux := dns.NewServeMux()
+
+	dnsServer := &dns.Server{
+		Addr:    f.listenAddress,
+		Net:     "udp",
+		Handler: mux,
+	}
+	f.dnsServer = dnsServer
+	f.mux = mux
+
+	f.UpdateDomains(domains)
+
+	return dnsServer.ListenAndServe()
+}
+
+func (f *DNSForwarder) UpdateDomains(domains []string) {
+	log.Debugf("Updating domains from %v to %v", f.domains, domains)
+
+	for _, d := range f.domains {
+		f.mux.HandleRemove(d)
+	}
+
+	newDomains := filterDomains(domains)
+	for _, d := range newDomains {
+		f.mux.HandleFunc(d, f.handleDNSQuery)
+	}
+	f.domains = newDomains
+}
+
+func (f *DNSForwarder) Close(ctx context.Context) error {
+	if f.dnsServer == nil {
+		return nil
+	}
+	return f.dnsServer.ShutdownContext(ctx)
+}
+
+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)
+
+	ips, err := net.LookupIP(domain)
+	if err != nil {
+		var dnsErr *net.DNSError
+
+		switch {
+		case errors.As(err, &dnsErr):
+			resp.Rcode = dns.RcodeServerFailure
+			if dnsErr.IsNotFound {
+				// Pass through NXDOMAIN
+				resp.Rcode = dns.RcodeNameError
+			}
+
+			if dnsErr.Server != "" {
+				log.Warnf("failed to resolve query for domain=%s server=%s: %v", domain, dnsErr.Server, err)
+			} else {
+				log.Warnf(errResolveFailed, domain, err)
+			}
+		default:
+			resp.Rcode = dns.RcodeServerFailure
+			log.Warnf(errResolveFailed, domain, err)
+		}
+
+		if err := w.WriteMsg(resp); err != nil {
+			log.Errorf("failed to write failure DNS response: %v", err)
+		}
+		return
+	}
+
+	for _, ip := range ips {
+		var respRecord dns.RR
+		if ip.To4() == nil {
+			log.Tracef("resolved domain=%s to IPv6=%s", domain, ip)
+			rr := dns.AAAA{
+				AAAA: ip,
+				Hdr: dns.RR_Header{
+					Name:   domain,
+					Rrtype: dns.TypeAAAA,
+					Class:  dns.ClassINET,
+					Ttl:    f.ttl,
+				},
+			}
+			respRecord = &rr
+		} else {
+			log.Tracef("resolved domain=%s to IPv4=%s", domain, ip)
+			rr := dns.A{
+				A: ip,
+				Hdr: dns.RR_Header{
+					Name:   domain,
+					Rrtype: dns.TypeA,
+					Class:  dns.ClassINET,
+					Ttl:    f.ttl,
+				},
+			}
+			respRecord = &rr
+		}
+		resp.Answer = append(resp.Answer, respRecord)
+	}
+
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed to write DNS response: %v", err)
+	}
+}
+
+// filterDomains returns a list of normalized domains
+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, nbdns.NormalizeZone(d))
+	}
+	return newDomains
+}

client/internal/dnsfwd/manager.go

@@ -0,0 +1,106 @@
+package dnsfwd
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"github.com/hashicorp/go-multierror"
+	log "github.com/sirupsen/logrus"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	firewall "github.com/netbirdio/netbird/client/firewall/manager"
+)
+
+const (
+	// ListenPort is the port that the DNS forwarder listens on. It has been used by the client peers also
+	ListenPort = 5353
+	dnsTTL     = 60 //seconds
+)
+
+type Manager struct {
+	firewall firewall.Manager
+
+	fwRules      []firewall.Rule
+	dnsForwarder *DNSForwarder
+}
+
+func NewManager(fw firewall.Manager) *Manager {
+	return &Manager{
+		firewall: fw,
+	}
+}
+
+func (m *Manager) Start(domains []string) error {
+	log.Infof("starting DNS forwarder")
+	if m.dnsForwarder != nil {
+		return nil
+	}
+
+	if err := m.allowDNSFirewall(); err != nil {
+		return err
+	}
+
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL)
+	go func() {
+		if err := m.dnsForwarder.Listen(domains); err != nil {
+			// todo handle close error if it is exists
+			log.Errorf("failed to start DNS forwarder, err: %v", err)
+		}
+	}()
+
+	return nil
+}
+
+func (m *Manager) UpdateDomains(domains []string) {
+	if m.dnsForwarder == nil {
+		return
+	}
+
+	m.dnsForwarder.UpdateDomains(domains)
+}
+
+func (m *Manager) Stop(ctx context.Context) error {
+	if m.dnsForwarder == nil {
+		return nil
+	}
+
+	var mErr *multierror.Error
+	if err := m.dropDNSFirewall(); err != nil {
+		mErr = multierror.Append(mErr, err)
+	}
+
+	if err := m.dnsForwarder.Close(ctx); err != nil {
+		mErr = multierror.Append(mErr, err)
+	}
+
+	m.dnsForwarder = nil
+	return nberrors.FormatErrorOrNil(mErr)
+}
+
+func (h *Manager) allowDNSFirewall() error {
+	dport := &firewall.Port{
+		IsRange: false,
+		Values:  []int{ListenPort},
+	}
+	dnsRules, err := h.firewall.AddPeerFiltering(net.ParseIP("0.0.0.0"), firewall.ProtocolUDP, nil, dport, firewall.RuleDirectionIN, firewall.ActionAccept, "", "")
+	if err != nil {
+		log.Errorf("failed to add allow DNS router rules, err: %v", err)
+		return err
+	}
+	h.fwRules = dnsRules
+
+	return nil
+}
+
+func (h *Manager) dropDNSFirewall() error {
+	var mErr *multierror.Error
+	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))
+		}
+	}
+
+	h.fwRules = nil
+	return nberrors.FormatErrorOrNil(mErr)
+}

client/internal/engine.go

@@ -4,13 +4,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"maps"
 	"math/rand"
 	"net"
 	"net/netip"
 	"reflect"
 	"runtime"
 	"slices"
+	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -20,24 +20,27 @@ import (
 	"github.com/pion/stun/v2"
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/netbirdio/netbird/client/firewall"
-	"github.com/netbirdio/netbird/client/firewall/manager"
+	firewallmanager "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/internal/acl"
 	"github.com/netbirdio/netbird/client/internal/dns"
+	"github.com/netbirdio/netbird/client/internal/dnsfwd"
 	"github.com/netbirdio/netbird/client/internal/networkmonitor"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/client/internal/relay"
 	"github.com/netbirdio/netbird/client/internal/rosenpass"
 	"github.com/netbirdio/netbird/client/internal/routemanager"
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
-
+	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 
 	nbssh "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
@@ -61,6 +64,7 @@ import (
 const (
 	PeerConnectionTimeoutMax = 45000 // ms
 	PeerConnectionTimeoutMin = 30000 // ms
+	connInitLimit            = 200
 )
 
 var ErrResetConnection = fmt.Errorf("reset connection")
@@ -114,7 +118,7 @@ type Engine struct {
 	// mgmClient is a Management Service client
 	mgmClient mgm.Client
 	// peerConns is a map that holds all the peers that are known to this peer
-	peerConns map[string]*peer.Conn
+	peerStore *peerstore.Store
 
 	beforePeerHook nbnet.AddHookFunc
 	afterPeerHook  nbnet.RemoveHookFunc
@@ -134,10 +138,6 @@ type Engine struct {
 	TURNs    []*stun.URI
 	stunTurn atomic.Value
 
-	// clientRoutes is the most recent list of clientRoutes received from the Management Service
-	clientRoutes   route.HAMap
-	clientRoutesMu sync.RWMutex
-
 	clientCtx    context.Context
 	clientCancel context.CancelFunc
 
@@ -158,9 +158,10 @@ type Engine struct {
 
 	statusRecorder *peer.Status
 
-	firewall     manager.Manager
-	routeManager routemanager.Manager
-	acl          acl.Manager
+	firewall      firewallmanager.Manager
+	routeManager  routemanager.Manager
+	acl           acl.Manager
+	dnsForwardMgr *dnsfwd.Manager
 
 	dnsServer dns.Server
 
@@ -171,7 +172,12 @@ type Engine struct {
 
 	relayManager *relayClient.Manager
 	stateManager *statemanager.Manager
-	srWatcher *guard.SRWatcher
+	srWatcher    *guard.SRWatcher
+
+	// Network map persistence
+	persistNetworkMap bool
+	latestNetworkMap  *mgmProto.NetworkMap
+	connSemaphore     *semaphoregroup.SemaphoreGroup
 }
 
 // Peer is an instance of the Connection Peer
@@ -226,7 +232,7 @@ func NewEngineWithProbes(
 		signaler:       peer.NewSignaler(signalClient, config.WgPrivateKey),
 		mgmClient:      mgmClient,
 		relayManager:   relayManager,
-		peerConns:      make(map[string]*peer.Conn),
+		peerStore:      peerstore.NewConnStore(),
 		syncMsgMux:     &sync.Mutex{},
 		config:         config,
 		mobileDep:      mobileDep,
@@ -237,6 +243,18 @@ func NewEngineWithProbes(
 		statusRecorder: statusRecorder,
 		probes:         probes,
 		checks:         checks,
+		connSemaphore:  semaphoregroup.NewSemaphoreGroup(connInitLimit),
+	}
+	if runtime.GOOS == "ios" {
+		if !fileExists(mobileDep.StateFilePath) {
+			err := createFile(mobileDep.StateFilePath)
+			if err != nil {
+				log.Errorf("failed to create state file: %v", err)
+				// we are not exiting as we can run without the state manager
+			}
+		}
+
+		engine.stateManager = statemanager.New(mobileDep.StateFilePath)
 	}
 	if path := statemanager.GetDefaultStatePath(); path != "" {
 		engine.stateManager = statemanager.New(path)
@@ -267,19 +285,26 @@ func (e *Engine) Stop() error {
 		e.routeManager.Stop(e.stateManager)
 	}
 
+	if e.dnsForwardMgr != nil {
+		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+			log.Errorf("failed to stop DNS forward: %v", err)
+		}
+		e.dnsForwardMgr = nil
+	}
+
 	if e.srWatcher != nil {
 		e.srWatcher.Close()
 	}
 
+	e.statusRecorder.ReplaceOfflinePeers([]peer.State{})
+	e.statusRecorder.UpdateDNSStates([]peer.NSGroupState{})
+	e.statusRecorder.UpdateRelayStates([]relay.ProbeResult{})
+
 	err := e.removeAllPeers()
 	if err != nil {
 		return fmt.Errorf("failed to remove all peers: %s", err)
 	}
 
-	e.clientRoutesMu.Lock()
-	e.clientRoutes = nil
-	e.clientRoutesMu.Unlock()
-
 	if e.cancel != nil {
 		e.cancel()
 	}
@@ -297,7 +322,7 @@ func (e *Engine) Stop() error {
 	if err := e.stateManager.Stop(ctx); err != nil {
 		return fmt.Errorf("failed to stop state manager: %w", err)
 	}
-	if err := e.stateManager.PersistState(ctx); err != nil {
+	if err := e.stateManager.PersistState(context.Background()); err != nil {
 		log.Errorf("failed to persist state: %v", err)
 	}
 
@@ -349,8 +374,19 @@ func (e *Engine) Start() error {
 	}
 	e.dnsServer = dnsServer
 
-	e.routeManager = routemanager.NewManager(e.ctx, e.config.WgPrivateKey.PublicKey().String(), e.config.DNSRouteInterval, e.wgInterface, e.statusRecorder, e.relayManager, initialRoutes)
-	beforePeerHook, afterPeerHook, err := e.routeManager.Init(e.stateManager)
+	e.routeManager = routemanager.NewManager(
+		e.ctx,
+		e.config.WgPrivateKey.PublicKey().String(),
+		e.config.DNSRouteInterval,
+		e.wgInterface,
+		e.statusRecorder,
+		e.relayManager,
+		initialRoutes,
+		e.stateManager,
+		dnsServer,
+		e.peerStore,
+	)
+	beforePeerHook, afterPeerHook, err := e.routeManager.Init()
 	if err != nil {
 		log.Errorf("Failed to initialize route manager: %s", err)
 	} else {
@@ -427,8 +463,8 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 	var modified []*mgmProto.RemotePeerConfig
 	for _, p := range peersUpdate {
 		peerPubKey := p.GetWgPubKey()
-		if peerConn, ok := e.peerConns[peerPubKey]; ok {
-			if peerConn.WgConfig().AllowedIps != strings.Join(p.AllowedIps, ",") {
+		if allowedIPs, ok := e.peerStore.AllowedIPs(peerPubKey); ok {
+			if allowedIPs != strings.Join(p.AllowedIps, ",") {
 				modified = append(modified, p)
 				continue
 			}
@@ -459,17 +495,12 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 // removePeers finds and removes peers that do not exist anymore in the network map received from the Management Service.
 // It also removes peers that have been modified (e.g. change of IP address). They will be added again in addPeers method.
 func (e *Engine) removePeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
-	currentPeers := make([]string, 0, len(e.peerConns))
-	for p := range e.peerConns {
-		currentPeers = append(currentPeers, p)
-	}
-
 	newPeers := make([]string, 0, len(peersUpdate))
 	for _, p := range peersUpdate {
 		newPeers = append(newPeers, p.GetWgPubKey())
 	}
 
-	toRemove := util.SliceDiff(currentPeers, newPeers)
+	toRemove := util.SliceDiff(e.peerStore.PeersPubKey(), newPeers)
 
 	for _, p := range toRemove {
 		err := e.removePeer(p)
@@ -483,7 +514,7 @@ func (e *Engine) removePeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 
 func (e *Engine) removeAllPeers() error {
 	log.Debugf("removing all peer connections")
-	for p := range e.peerConns {
+	for _, p := range e.peerStore.PeersPubKey() {
 		err := e.removePeer(p)
 		if err != nil {
 			return err
@@ -507,9 +538,8 @@ func (e *Engine) removePeer(peerKey string) error {
 		}
 	}()
 
-	conn, exists := e.peerConns[peerKey]
+	conn, exists := e.peerStore.Remove(peerKey)
 	if exists {
-		delete(e.peerConns, peerKey)
 		conn.Close()
 	}
 	return nil
@@ -538,6 +568,7 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 
 		relayMsg := wCfg.GetRelay()
 		if relayMsg != nil {
+			// when we receive token we expect valid address list too
 			c := &auth.Token{
 				Payload:   relayMsg.GetTokenPayload(),
 				Signature: relayMsg.GetTokenSignature(),
@@ -546,9 +577,16 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 				log.Errorf("failed to update relay token: %v", err)
 				return fmt.Errorf("update relay token: %w", err)
 			}
+
+			e.relayManager.UpdateServerURLs(relayMsg.Urls)
+
+			// Just in case the agent started with an MGM server where the relay was disabled but was later enabled.
+			// We can ignore all errors because the guard will manage the reconnection retries.
+			_ = e.relayManager.Serve()
+		} else {
+			e.relayManager.UpdateServerURLs(nil)
 		}
 
-		// todo update relay address in the relay manager
 		// todo update signal
 	}
 
@@ -556,13 +594,22 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 		return err
 	}
 
-	if update.GetNetworkMap() != nil {
-		// only apply new changes and ignore old ones
-		err := e.updateNetworkMap(update.GetNetworkMap())
-		if err != nil {
-			return err
-		}
+	nm := update.GetNetworkMap()
+	if nm == nil {
+		return nil
 	}
+
+	// Store network map if persistence is enabled
+	if e.persistNetworkMap {
+		e.latestNetworkMap = nm
+		log.Debugf("network map persisted with serial %d", nm.GetSerial())
+	}
+
+	// only apply new changes and ignore old ones
+	if err := e.updateNetworkMap(nm); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -641,6 +688,10 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 }
 
 func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
+	if e.wgInterface == nil {
+		return errors.New("wireguard interface is not initialized")
+	}
+
 	if e.wgInterface.Address().String() != conf.Address {
 		oldAddr := e.wgInterface.Address().String()
 		log.Debugf("updating peer address from %s to %s", oldAddr, conf.Address)
@@ -732,7 +783,6 @@ func (e *Engine) updateTURNs(turns []*mgmProto.ProtectedHostConfig) error {
 }
 
 func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
-
 	// intentionally leave it before checking serial because for now it can happen that peer IP changed but serial didn't
 	if networkMap.GetPeerConfig() != nil {
 		err := e.updateConfig(networkMap.GetPeerConfig())
@@ -752,20 +802,16 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		e.acl.ApplyFiltering(networkMap)
 	}
 
-	protoRoutes := networkMap.GetRoutes()
-	if protoRoutes == nil {
-		protoRoutes = []*mgmProto.Route{}
-	}
+	// DNS forwarder
+	dnsRouteFeatureFlag := toDNSFeatureFlag(networkMap)
+	dnsRouteDomains := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
+	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains)
 
-	_, clientRoutes, err := e.routeManager.UpdateRoutes(serial, toRoutes(protoRoutes))
-	if err != nil {
+	routes := toRoutes(networkMap.GetRoutes())
+	if err := e.routeManager.UpdateRoutes(serial, routes, dnsRouteFeatureFlag); err != nil {
 		log.Errorf("failed to update clientRoutes, err: %v", err)
 	}
 
-	e.clientRoutesMu.Lock()
-	e.clientRoutes = clientRoutes
-	e.clientRoutesMu.Unlock()
-
 	log.Debugf("got peers update from Management Service, total peers to connect to = %d", len(networkMap.GetRemotePeers()))
 
 	e.updateOfflinePeers(networkMap.GetOfflinePeers())
@@ -813,8 +859,7 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		protoDNSConfig = &mgmProto.DNSConfig{}
 	}
 
-	err = e.dnsServer.UpdateDNSServer(serial, toDNSConfig(protoDNSConfig))
-	if err != nil {
+	if err := e.dnsServer.UpdateDNSServer(serial, toDNSConfig(protoDNSConfig)); err != nil {
 		log.Errorf("failed to update dns server, err: %v", err)
 	}
 
@@ -827,7 +872,18 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 	return nil
 }
 
+func toDNSFeatureFlag(networkMap *mgmProto.NetworkMap) bool {
+	if networkMap.PeerConfig != nil {
+		return networkMap.PeerConfig.RoutingPeerDnsResolutionEnabled
+	}
+	return false
+}
+
 func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
+	if protoRoutes == nil {
+		protoRoutes = []*mgmProto.Route{}
+	}
+
 	routes := make([]*route.Route, 0)
 	for _, protoRoute := range protoRoutes {
 		var prefix netip.Prefix
@@ -838,6 +894,7 @@ func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
 				continue
 			}
 		}
+
 		convertedRoute := &route.Route{
 			ID:          route.ID(protoRoute.ID),
 			Network:     prefix,
@@ -854,6 +911,23 @@ func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
 	return routes
 }
 
+func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) []string {
+	if protoRoutes == nil {
+		protoRoutes = []*mgmProto.Route{}
+	}
+
+	var dnsRoutes []string
+	for _, protoRoute := range protoRoutes {
+		if len(protoRoute.Domains) == 0 {
+			continue
+		}
+		if protoRoute.Peer == myPubKey {
+			dnsRoutes = append(dnsRoutes, protoRoute.Domains...)
+		}
+	}
+	return dnsRoutes
+}
+
 func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig) nbdns.Config {
 	dnsUpdate := nbdns.Config{
 		ServiceEnable:    protoDNSConfig.GetServiceEnable(),
@@ -928,12 +1002,16 @@ func (e *Engine) addNewPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
 func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error {
 	peerKey := peerConfig.GetWgPubKey()
 	peerIPs := peerConfig.GetAllowedIps()
-	if _, ok := e.peerConns[peerKey]; !ok {
+	if _, ok := e.peerStore.PeerConn(peerKey); !ok {
 		conn, err := e.createPeerConn(peerKey, strings.Join(peerIPs, ","))
 		if err != nil {
 			return fmt.Errorf("create peer connection: %w", err)
 		}
-		e.peerConns[peerKey] = conn
+
+		if ok := e.peerStore.AddPeerConn(peerKey, conn); !ok {
+			conn.Close()
+			return fmt.Errorf("peer already exists: %s", peerKey)
+		}
 
 		if e.beforePeerHook != nil && e.afterPeerHook != nil {
 			conn.AddBeforeAddPeerHook(e.beforePeerHook)
@@ -1001,7 +1079,7 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs string) (*peer.Conn, e
 		},
 	}
 
-	peerConn, err := peer.NewConn(e.ctx, config, e.statusRecorder, e.signaler, e.mobileDep.IFaceDiscover, e.relayManager, e.srWatcher)
+	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
 	}
@@ -1022,8 +1100,8 @@ func (e *Engine) receiveSignalEvents() {
 			e.syncMsgMux.Lock()
 			defer e.syncMsgMux.Unlock()
 
-			conn := e.peerConns[msg.Key]
-			if conn == nil {
+			conn, ok := e.peerStore.PeerConn(msg.Key)
+			if !ok {
 				return fmt.Errorf("wrongly addressed message %s", msg.Key)
 			}
 
@@ -1081,7 +1159,7 @@ func (e *Engine) receiveSignalEvents() {
 					return err
 				}
 
-				go conn.OnRemoteCandidate(candidate, e.GetClientRoutes())
+				go conn.OnRemoteCandidate(candidate, e.routeManager.GetClientRoutes())
 			case sProto.Body_MODE:
 			}
 
@@ -1268,26 +1346,6 @@ func (e *Engine) newDnsServer() ([]*route.Route, dns.Server, error) {
 	}
 }
 
-// GetClientRoutes returns the current routes from the route map
-func (e *Engine) GetClientRoutes() route.HAMap {
-	e.clientRoutesMu.RLock()
-	defer e.clientRoutesMu.RUnlock()
-
-	return maps.Clone(e.clientRoutes)
-}
-
-// GetClientRoutesWithNetID returns the current routes from the route map, but the keys consist of the network ID only
-func (e *Engine) GetClientRoutesWithNetID() map[route.NetID][]*route.Route {
-	e.clientRoutesMu.RLock()
-	defer e.clientRoutesMu.RUnlock()
-
-	routes := make(map[route.NetID][]*route.Route, len(e.clientRoutes))
-	for id, v := range e.clientRoutes {
-		routes[id.NetID()] = v
-	}
-	return routes
-}
-
 // GetRouteManager returns the route manager
 func (e *Engine) GetRouteManager() routemanager.Manager {
 	return e.routeManager
@@ -1372,9 +1430,8 @@ func (e *Engine) receiveProbeEvents() {
 		go e.probes.WgProbe.Receive(e.ctx, func() bool {
 			log.Debug("received wg probe request")
 
-			for _, peer := range e.peerConns {
-				key := peer.GetKey()
-				wgStats, err := peer.WgConfig().WgInterface.GetStats(key)
+			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)
 				}
@@ -1451,7 +1508,7 @@ func (e *Engine) startNetworkMonitor() {
 
 func (e *Engine) addrViaRoutes(addr netip.Addr) (bool, netip.Prefix, error) {
 	var vpnRoutes []netip.Prefix
-	for _, routes := range e.GetClientRoutes() {
+	for _, routes := range e.routeManager.GetClientRoutes() {
 		if len(routes) > 0 && routes[0] != nil {
 			vpnRoutes = append(vpnRoutes, routes[0].Network)
 		}
@@ -1479,8 +1536,101 @@ func (e *Engine) stopDNSServer() {
 	e.statusRecorder.UpdateDNSStates(nsGroupStates)
 }
 
+// SetNetworkMapPersistence enables or disables network map persistence
+func (e *Engine) SetNetworkMapPersistence(enabled bool) {
+	e.syncMsgMux.Lock()
+	defer e.syncMsgMux.Unlock()
+
+	if enabled == e.persistNetworkMap {
+		return
+	}
+	e.persistNetworkMap = enabled
+	log.Debugf("Network map persistence is set to %t", enabled)
+
+	if !enabled {
+		e.latestNetworkMap = nil
+	}
+}
+
+// GetLatestNetworkMap returns the stored network map if persistence is enabled
+func (e *Engine) GetLatestNetworkMap() (*mgmProto.NetworkMap, error) {
+	e.syncMsgMux.Lock()
+	defer e.syncMsgMux.Unlock()
+
+	if !e.persistNetworkMap {
+		return nil, errors.New("network map persistence is disabled")
+	}
+
+	if e.latestNetworkMap == nil {
+		//nolint:nilnil
+		return nil, nil
+	}
+
+	// Create a deep copy to avoid external modifications
+	nm, ok := proto.Clone(e.latestNetworkMap).(*mgmProto.NetworkMap)
+	if !ok {
+
+		return nil, fmt.Errorf("failed to clone network map")
+	}
+
+	return nm, nil
+}
+
+// GetFirewallStats returns the firewall stats
+func (e *Engine) GetFirewallStats() []*firewallmanager.FlowStats {
+	if e.firewall != nil {
+		return e.firewall.CollectStats()
+	}
+	return nil
+}
+
+// updateDNSForwarder start or stop the DNS forwarder based on the domains and the feature flag
+func (e *Engine) updateDNSForwarder(enabled bool, domains []string) {
+	if !enabled {
+		if e.dnsForwardMgr == nil {
+			return
+		}
+		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+			log.Errorf("failed to stop DNS forward: %v", err)
+		}
+		return
+	}
+
+	if len(domains) > 0 {
+		log.Infof("enable domain router service for domains: %v", domains)
+		if e.dnsForwardMgr == nil {
+			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall)
+
+			if err := e.dnsForwardMgr.Start(domains); err != nil {
+				log.Errorf("failed to start DNS forward: %v", err)
+				e.dnsForwardMgr = nil
+			}
+		} else {
+			log.Infof("update domain router service for domains: %v", domains)
+			e.dnsForwardMgr.UpdateDomains(domains)
+		}
+	} else if e.dnsForwardMgr != nil {
+		log.Infof("disable domain router service")
+		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+			log.Errorf("failed to stop DNS forward: %v", err)
+		}
+		e.dnsForwardMgr = nil
+	}
+}
+
 // isChecksEqual checks if two slices of checks are equal.
 func isChecksEqual(checks []*mgmProto.Checks, oChecks []*mgmProto.Checks) bool {
+	for _, check := range checks {
+		sort.Slice(check.Files, func(i, j int) bool {
+			return check.Files[i] < check.Files[j]
+		})
+	}
+	for _, oCheck := range oChecks {
+		sort.Slice(oCheck.Files, func(i, j int) bool {
+			return oCheck.Files[i] < oCheck.Files[j]
+		})
+	}
+
 	return slices.EqualFunc(checks, oChecks, func(checks, oChecks *mgmProto.Checks) bool {
 		return slices.Equal(checks.Files, oChecks.Files)
 	})

client/internal/engine_test.go

@@ -39,6 +39,8 @@ import (
 	mgmtProto "github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/activity"
+	"github.com/netbirdio/netbird/management/server/settings"
+	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	"github.com/netbirdio/netbird/route"
@@ -245,12 +247,15 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 		nil)
 
 	wgIface := &iface.MockWGIface{
+		NameFunc: func() string { return "utun102" },
 		RemovePeerFunc: func(peerKey string) error {
 			return nil
 		},
 	}
 	engine.wgInterface = wgIface
-	engine.routeManager = routemanager.NewManager(ctx, key.PublicKey().String(), time.Minute, engine.wgInterface, engine.statusRecorder, relayMgr, nil)
+	engine.routeManager = routemanager.NewManager(ctx, key.PublicKey().String(), time.Minute, engine.wgInterface, engine.statusRecorder, relayMgr, nil, nil, nil, nil)
+	_, _, err = engine.routeManager.Init()
+	require.NoError(t, err)
 	engine.dnsServer = &dns.MockServer{
 		UpdateDNSServerFunc: func(serial uint64, update nbdns.Config) error { return nil },
 	}
@@ -388,8 +393,8 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 				return
 			}
 
-			if len(engine.peerConns) != c.expectedLen {
-				t.Errorf("expecting Engine.peerConns to be of size %d, got %d", c.expectedLen, len(engine.peerConns))
+			if len(engine.peerStore.PeersPubKey()) != c.expectedLen {
+				t.Errorf("expecting Engine.peerConns to be of size %d, got %d", c.expectedLen, len(engine.peerStore.PeersPubKey()))
 			}
 
 			if engine.networkSerial != c.expectedSerial {
@@ -397,7 +402,7 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 			}
 
 			for _, p := range c.expectedPeers {
-				conn, ok := engine.peerConns[p.GetWgPubKey()]
+				conn, ok := engine.peerStore.PeerConn(p.GetWgPubKey())
 				if !ok {
 					t.Errorf("expecting Engine.peerConns to contain peer %s", p)
 				}
@@ -622,10 +627,10 @@ func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {
 			}{}
 
 			mockRouteManager := &routemanager.MockManager{
-				UpdateRoutesFunc: func(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error) {
+				UpdateRoutesFunc: func(updateSerial uint64, newRoutes []*route.Route) error {
 					input.inputSerial = updateSerial
 					input.inputRoutes = newRoutes
-					return nil, nil, testCase.inputErr
+					return testCase.inputErr
 				},
 			}
 
@@ -798,8 +803,8 @@ func TestEngine_UpdateNetworkMapWithDNSUpdate(t *testing.T) {
 			assert.NoError(t, err, "shouldn't return error")
 
 			mockRouteManager := &routemanager.MockManager{
-				UpdateRoutesFunc: func(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error) {
-					return nil, nil, nil
+				UpdateRoutesFunc: func(updateSerial uint64, newRoutes []*route.Route) error {
+					return nil
 				},
 			}
 
@@ -1006,6 +1011,99 @@ func Test_ParseNATExternalIPMappings(t *testing.T) {
 	}
 }
 
+func Test_CheckFilesEqual(t *testing.T) {
+	testCases := []struct {
+		name         string
+		inputChecks1 []*mgmtProto.Checks
+		inputChecks2 []*mgmtProto.Checks
+		expectedBool bool
+	}{
+		{
+			name: "Equal Files In Equal Order Should Return True",
+			inputChecks1: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile2",
+					},
+				},
+			},
+			inputChecks2: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile2",
+					},
+				},
+			},
+			expectedBool: true,
+		},
+		{
+			name: "Equal Files In Reverse Order Should Return True",
+			inputChecks1: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile2",
+					},
+				},
+			},
+			inputChecks2: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile2",
+						"testfile1",
+					},
+				},
+			},
+			expectedBool: true,
+		},
+		{
+			name: "Unequal Files Should Return False",
+			inputChecks1: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile2",
+					},
+				},
+			},
+			inputChecks2: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile3",
+					},
+				},
+			},
+			expectedBool: false,
+		},
+		{
+			name: "Compared With Empty Should Return False",
+			inputChecks1: []*mgmtProto.Checks{
+				{
+					Files: []string{
+						"testfile1",
+						"testfile2",
+					},
+				},
+			},
+			inputChecks2: []*mgmtProto.Checks{
+				{
+					Files: []string{},
+				},
+			},
+			expectedBool: false,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			result := isChecksEqual(testCase.inputChecks1, testCase.inputChecks2)
+			assert.Equal(t, testCase.expectedBool, result, "result should match expected bool")
+		})
+	}
+}
+
 func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey string, i int, mgmtAddr string, signalAddr string) (*Engine, error) {
 	key, err := wgtypes.GeneratePrivateKey()
 	if err != nil {
@@ -1100,7 +1198,7 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 	}
 	s := grpc.NewServer(grpc.KeepaliveEnforcementPolicy(kaep), grpc.KeepaliveParams(kasp))
 
-	store, cleanUp, err := server.NewTestStoreFromSQL(context.Background(), testFile, config.Datadir)
+	store, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), testFile, config.Datadir)
 	if err != nil {
 		return nil, "", err
 	}
@@ -1122,7 +1220,7 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 	}
 
 	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay)
-	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, peersUpdateManager, secretsManager, nil, nil)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settings.NewManager(store), peersUpdateManager, secretsManager, nil, nil)
 	if err != nil {
 		return nil, "", err
 	}
@@ -1141,7 +1239,8 @@ func getConnectedPeers(e *Engine) int {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 	i := 0
-	for _, conn := range e.peerConns {
+	for _, id := range e.peerStore.PeersPubKey() {
+		conn, _ := e.peerStore.PeerConn(id)
 		if conn.Status() == peer.StatusConnected {
 			i++
 		}
@@ -1153,5 +1252,5 @@ func getPeers(e *Engine) int {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 
-	return len(e.peerConns)
+	return len(e.peerStore.PeersPubKey())
 }

client/internal/mobile_dependency.go

@@ -19,4 +19,5 @@ type MobileDependency struct {
 	//	iOS only
 	DnsManager     dns.IosDnsManager
 	FileDescriptor int32
+	StateFilePath  string
 }

client/internal/peer/conn.go

@@ -23,6 +23,7 @@ import (
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	"github.com/netbirdio/netbird/route"
 	nbnet "github.com/netbirdio/netbird/util/net"
+	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 )
 
 type ConnPriority int
@@ -83,7 +84,6 @@ type Conn struct {
 	signaler       *Signaler
 	relayManager   *relayClient.Manager
 	allowedIP      net.IP
-	allowedNet     string
 	handshaker     *Handshaker
 
 	onConnected    func(remoteWireGuardKey string, remoteRosenpassPubKey []byte, wireGuardIP string, remoteRosenpassAddr string)
@@ -105,13 +105,14 @@ type Conn struct {
 	wgProxyICE   wgproxy.Proxy
 	wgProxyRelay wgproxy.Proxy
 
-	guard *guard.Guard
+	guard     *guard.Guard
+	semaphore *semaphoregroup.SemaphoreGroup
 }
 
 // NewConn creates a new not opened Conn to the remote peer.
 // To establish a connection run Conn.Open
-func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Status, signaler *Signaler, iFaceDiscover stdnet.ExternalIFaceDiscover, relayManager *relayClient.Manager, srWatcher *guard.SRWatcher) (*Conn, error) {
-	allowedIP, allowedNet, err := net.ParseCIDR(config.WgConfig.AllowedIps)
+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) {
+	allowedIP, _, err := net.ParseCIDR(config.WgConfig.AllowedIps)
 	if err != nil {
 		log.Errorf("failed to parse allowedIPS: %v", err)
 		return nil, err
@@ -129,9 +130,9 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 		signaler:       signaler,
 		relayManager:   relayManager,
 		allowedIP:      allowedIP,
-		allowedNet:     allowedNet.String(),
 		statusRelay:    NewAtomicConnStatus(),
 		statusICE:      NewAtomicConnStatus(),
+		semaphore:      semaphore,
 	}
 
 	rFns := WorkerRelayCallbacks{
@@ -171,6 +172,7 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 // 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()
@@ -193,6 +195,7 @@ func (conn *Conn) Open() {
 }
 
 func (conn *Conn) startHandshakeAndReconnect(ctx context.Context) {
+	defer conn.semaphore.Done(conn.ctx)
 	conn.waitInitialRandomSleepTime(ctx)
 
 	err := conn.handshaker.sendOffer()
@@ -309,6 +312,11 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		return
 	}
 
+	if remoteConnNil(conn.log, iceConnInfo.RemoteConn) {
+		conn.log.Errorf("remote ICE connection is nil")
+		return
+	}
+
 	conn.log.Debugf("ICE connection is ready")
 
 	if conn.currentConnPriority > priority {
@@ -333,12 +341,9 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		ep = wgProxy.EndpointAddr()
 		conn.wgProxyICE = wgProxy
 	} else {
-		conn.log.Infof("direct iceConnInfo: %v", iceConnInfo.RemoteConn)
-		agentCheck(conn.log, iceConnInfo.Agent)
-		nilCheck(conn.log, iceConnInfo.RemoteConn)
 		directEp, err := net.ResolveUDPAddr("udp", iceConnInfo.RemoteConn.RemoteAddr().String())
 		if err != nil {
-			conn.log.Errorf("failed to resolveUDPaddr")
+			log.Errorf("failed to resolveUDPaddr")
 			conn.handleConfigurationFailure(err, nil)
 			return
 		}
@@ -440,7 +445,7 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 
 	if conn.iceP2PIsActive() {
 		conn.log.Debugf("do not switch to relay because current priority is: %v", conn.currentConnPriority)
-		conn.wgProxyRelay = wgProxy
+		conn.setRelayedProxy(wgProxy)
 		conn.statusRelay.Set(StatusConnected)
 		conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
 		return
@@ -463,7 +468,7 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 	wgConfigWorkaround()
 	conn.currentConnPriority = connPriorityRelay
 	conn.statusRelay.Set(StatusConnected)
-	conn.wgProxyRelay = wgProxy
+	conn.setRelayedProxy(wgProxy)
 	conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
 	conn.log.Infof("start to communicate with peer via relay")
 	conn.doOnConnected(rci.rosenpassPubKey, rci.rosenpassAddr)
@@ -592,14 +597,13 @@ func (conn *Conn) doOnConnected(remoteRosenpassPubKey []byte, remoteRosenpassAdd
 	}
 
 	if conn.onConnected != nil {
-		conn.onConnected(conn.config.Key, remoteRosenpassPubKey, conn.allowedNet, remoteRosenpassAddr)
+		conn.onConnected(conn.config.Key, remoteRosenpassPubKey, conn.allowedIP.String(), remoteRosenpassAddr)
 	}
 }
 
 func (conn *Conn) waitInitialRandomSleepTime(ctx context.Context) {
-	minWait := 100
-	maxWait := 800
-	duration := time.Duration(rand.Intn(maxWait-minWait)+minWait) * time.Millisecond
+	maxWait := 300
+	duration := time.Duration(rand.Intn(maxWait)) * time.Millisecond
 
 	timeout := time.NewTimer(duration)
 	defer timeout.Stop()
@@ -734,6 +738,20 @@ func (conn *Conn) logTraceConnState() {
 	}
 }
 
+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.wgProxyRelay = proxy
+}
+
+// AllowedIP returns the allowed IP of the remote peer
+func (conn *Conn) AllowedIP() net.IP {
+	return conn.allowedIP
+}
+
 func isController(config ConnConfig) bool {
 	return config.LocalKey > config.Key
 }

client/internal/peer/conn_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
 	"github.com/netbirdio/netbird/util"
+	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 )
 
 var connConf = ConnConfig{
@@ -46,7 +47,7 @@ func TestNewConn_interfaceFilter(t *testing.T) {
 
 func TestConn_GetKey(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-	conn, err := NewConn(context.Background(), connConf, nil, nil, nil, nil, swWatcher)
+	conn, err := NewConn(context.Background(), connConf, nil, nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -58,7 +59,7 @@ func TestConn_GetKey(t *testing.T) {
 
 func TestConn_OnRemoteOffer(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -92,7 +93,7 @@ func TestConn_OnRemoteOffer(t *testing.T) {
 
 func TestConn_OnRemoteAnswer(t *testing.T) {
 	swWatcher := guard.NewSRWatcher(nil, nil, nil, connConf.ICEConfig)
-	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}
@@ -125,7 +126,7 @@ func TestConn_OnRemoteAnswer(t *testing.T) {
 }
 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)
+	conn, err := NewConn(context.Background(), connConf, NewRecorder("https://mgm"), nil, nil, nil, swWatcher, semaphoregroup.NewSemaphoreGroup(1))
 	if err != nil {
 		return
 	}

client/internal/peer/nilcheck.go

@@ -2,53 +2,20 @@ package peer
 
 import (
 	"net"
-	"reflect"
 
-	"github.com/pion/ice/v3"
 	log "github.com/sirupsen/logrus"
 )
 
-func nilCheck(log *log.Entry, conn net.Conn) {
+func remoteConnNil(log *log.Entry, conn net.Conn) bool {
 	if conn == nil {
-		log.Infof("conn is nil")
-		return
+		log.Errorf("ice conn is nil")
+		return true
 	}
 
 	if conn.RemoteAddr() == nil {
-		log.Infof("conn.RemoteAddr() is nil")
-		return
+		log.Errorf("ICE remote address is nil")
+		return true
 	}
 
-	if reflect.ValueOf(conn.RemoteAddr()).IsNil() {
-		log.Infof("value of conn.RemoteAddr() is nil")
-		return
-	}
-}
-
-func agentCheck(log *log.Entry, agent *ice.Agent) {
-	if agent == nil {
-		log.Errorf("agent is nil")
-		return
-	}
-
-	pair, err := agent.GetSelectedCandidatePair()
-	if err != nil {
-		log.Errorf("error getting selected candidate pair: %v", err)
-		return
-	}
-
-	if pair == nil {
-		log.Errorf("pair is nil")
-		return
-	}
-
-	if pair.Remote == nil {
-		log.Errorf("pair.Remote is nil")
-		return
-	}
-
-	if pair.Remote.Address() == "" {
-		log.Errorf("address is empty")
-		return
-	}
+	return false
 }

client/internal/peer/status.go

@@ -17,6 +17,11 @@ import (
 	relayClient "github.com/netbirdio/netbird/relay/client"
 )
 
+type ResolvedDomainInfo struct {
+	Prefixes     []netip.Prefix
+	ParentDomain domain.Domain
+}
+
 // State contains the latest state of a peer
 type State struct {
 	Mux                        *sync.RWMutex
@@ -67,7 +72,7 @@ func (s *State) DeleteRoute(network string) {
 func (s *State) GetRoutes() map[string]struct{} {
 	s.Mux.RLock()
 	defer s.Mux.RUnlock()
-	return s.routes
+	return maps.Clone(s.routes)
 }
 
 // LocalPeerState contains the latest state of the local peer
@@ -138,7 +143,7 @@ type Status struct {
 	rosenpassEnabled      bool
 	rosenpassPermissive   bool
 	nsGroupStates         []NSGroupState
-	resolvedDomainsStates map[domain.Domain][]netip.Prefix
+	resolvedDomainsStates map[domain.Domain]ResolvedDomainInfo
 
 	// To reduce the number of notification invocation this bool will be true when need to call the notification
 	// Some Peer actions mostly used by in a batch when the network map has been synchronized. In these type of events
@@ -156,7 +161,7 @@ func NewRecorder(mgmAddress string) *Status {
 		offlinePeers:          make([]State, 0),
 		notifier:              newNotifier(),
 		mgmAddress:            mgmAddress,
-		resolvedDomainsStates: make(map[domain.Domain][]netip.Prefix),
+		resolvedDomainsStates: map[domain.Domain]ResolvedDomainInfo{},
 	}
 }
 
@@ -237,10 +242,6 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		peerState.IP = receivedState.IP
 	}
 
-	if receivedState.GetRoutes() != nil {
-		peerState.SetRoutes(receivedState.GetRoutes())
-	}
-
 	skipNotification := shouldSkipNotify(receivedState.ConnStatus, peerState)
 
 	if receivedState.ConnStatus != peerState.ConnStatus {
@@ -261,12 +262,40 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		return nil
 	}
 
-	ch, found := d.changeNotify[receivedState.PubKey]
-	if found && ch != nil {
-		close(ch)
-		d.changeNotify[receivedState.PubKey] = nil
+	d.notifyPeerListChanged()
+	return nil
+}
+
+func (d *Status) AddPeerStateRoute(peer string, route string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[peer]
+	if !ok {
+		return errors.New("peer doesn't exist")
 	}
 
+	peerState.AddRoute(route)
+	d.peers[peer] = peerState
+
+	// todo: consider to make sense of this notification or not
+	d.notifyPeerListChanged()
+	return nil
+}
+
+func (d *Status) RemovePeerStateRoute(peer string, route string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[peer]
+	if !ok {
+		return errors.New("peer doesn't exist")
+	}
+
+	peerState.DeleteRoute(route)
+	d.peers[peer] = peerState
+
+	// todo: consider to make sense of this notification or not
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -301,12 +330,7 @@ func (d *Status) UpdatePeerICEState(receivedState State) error {
 		return nil
 	}
 
-	ch, found := d.changeNotify[receivedState.PubKey]
-	if found && ch != nil {
-		close(ch)
-		d.changeNotify[receivedState.PubKey] = nil
-	}
-
+	d.notifyPeerStateChangeListeners(receivedState.PubKey)
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -334,12 +358,7 @@ func (d *Status) UpdatePeerRelayedState(receivedState State) error {
 		return nil
 	}
 
-	ch, found := d.changeNotify[receivedState.PubKey]
-	if found && ch != nil {
-		close(ch)
-		d.changeNotify[receivedState.PubKey] = nil
-	}
-
+	d.notifyPeerStateChangeListeners(receivedState.PubKey)
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -366,12 +385,7 @@ func (d *Status) UpdatePeerRelayedStateToDisconnected(receivedState State) error
 		return nil
 	}
 
-	ch, found := d.changeNotify[receivedState.PubKey]
-	if found && ch != nil {
-		close(ch)
-		d.changeNotify[receivedState.PubKey] = nil
-	}
-
+	d.notifyPeerStateChangeListeners(receivedState.PubKey)
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -401,12 +415,7 @@ func (d *Status) UpdatePeerICEStateToDisconnected(receivedState State) error {
 		return nil
 	}
 
-	ch, found := d.changeNotify[receivedState.PubKey]
-	if found && ch != nil {
-		close(ch)
-		d.changeNotify[receivedState.PubKey] = nil
-	}
-
+	d.notifyPeerStateChangeListeners(receivedState.PubKey)
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -477,11 +486,14 @@ func (d *Status) FinishPeerListModifications() {
 func (d *Status) GetPeerStateChangeNotifier(peer string) <-chan struct{} {
 	d.mux.Lock()
 	defer d.mux.Unlock()
+
 	ch, found := d.changeNotify[peer]
-	if !found || ch == nil {
-		ch = make(chan struct{})
-		d.changeNotify[peer] = ch
+	if found {
+		return ch
 	}
+
+	ch = make(chan struct{})
+	d.changeNotify[peer] = ch
 	return ch
 }
 
@@ -584,16 +596,27 @@ func (d *Status) UpdateDNSStates(dnsStates []NSGroupState) {
 	d.nsGroupStates = dnsStates
 }
 
-func (d *Status) UpdateResolvedDomainsStates(domain domain.Domain, prefixes []netip.Prefix) {
+func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resolvedDomain domain.Domain, prefixes []netip.Prefix) {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	d.resolvedDomainsStates[domain] = prefixes
+
+	// Store both the original domain pattern and resolved domain
+	d.resolvedDomainsStates[resolvedDomain] = ResolvedDomainInfo{
+		Prefixes:     prefixes,
+		ParentDomain: originalDomain,
+	}
 }
 
 func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	delete(d.resolvedDomainsStates, domain)
+
+	// Remove all entries that have this domain as their parent
+	for k, v := range d.resolvedDomainsStates {
+		if v.ParentDomain == domain {
+			delete(d.resolvedDomainsStates, k)
+		}
+	}
 }
 
 func (d *Status) GetRosenpassState() RosenpassState {
@@ -669,25 +692,23 @@ func (d *Status) GetRelayStates() []relay.ProbeResult {
 	// extend the list of stun, turn servers with relay address
 	relayStates := slices.Clone(d.relayStates)
 
-	var relayState relay.ProbeResult
-
 	// if the server connection is not established then we will use the general address
 	// in case of connection we will use the instance specific address
 	instanceAddr, err := d.relayMgr.RelayInstanceAddress()
 	if err != nil {
 		// TODO add their status
-		if errors.Is(err, relayClient.ErrRelayClientNotConnected) {
-			for _, r := range d.relayMgr.ServerURLs() {
-				relayStates = append(relayStates, relay.ProbeResult{
-					URI: r,
-				})
-			}
-			return relayStates
+		for _, r := range d.relayMgr.ServerURLs() {
+			relayStates = append(relayStates, relay.ProbeResult{
+				URI: r,
+				Err: err,
+			})
 		}
-		relayState.Err = err
+		return relayStates
 	}
 
-	relayState.URI = instanceAddr
+	relayState := relay.ProbeResult{
+		URI: instanceAddr,
+	}
 	return append(relayStates, relayState)
 }
 
@@ -697,7 +718,7 @@ func (d *Status) GetDNSStates() []NSGroupState {
 	return d.nsGroupStates
 }
 
-func (d *Status) GetResolvedDomainsStates() map[domain.Domain][]netip.Prefix {
+func (d *Status) GetResolvedDomainsStates() map[domain.Domain]ResolvedDomainInfo {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 	return maps.Clone(d.resolvedDomainsStates)
@@ -755,6 +776,17 @@ func (d *Status) onConnectionChanged() {
 	d.notifier.updateServerStates(d.managementState, d.signalState)
 }
 
+// notifyPeerStateChangeListeners notifies route manager about the change in peer state
+func (d *Status) notifyPeerStateChangeListeners(peerID string) {
+	ch, found := d.changeNotify[peerID]
+	if !found {
+		return
+	}
+
+	close(ch)
+	delete(d.changeNotify, peerID)
+}
+
 func (d *Status) notifyPeerListChanged() {
 	d.notifier.peerListChanged(d.numOfPeers())
 }

client/internal/peer/status_test.go

@@ -93,7 +93,7 @@ func TestGetPeerStateChangeNotifierLogic(t *testing.T) {
 
 	peerState.IP = ip
 
-	err := status.UpdatePeerState(peerState)
+	err := status.UpdatePeerRelayedStateToDisconnected(peerState)
 	assert.NoError(t, err, "shouldn't return error")
 
 	select {

client/internal/peer/worker_ice.go

@@ -29,7 +29,6 @@ type ICEConnInfo struct {
 	LocalIceCandidateEndpoint  string
 	Relayed                    bool
 	RelayedOnLocal             bool
-	Agent                      *ice.Agent
 }
 
 type WorkerICECallbacks struct {
@@ -47,8 +46,6 @@ type WorkerICE struct {
 	hasRelayOnLocally bool
 	conn              WorkerICECallbacks
 
-	selectedPriority ConnPriority
-
 	agent    *ice.Agent
 	muxAgent sync.Mutex
 
@@ -58,6 +55,9 @@ type WorkerICE struct {
 
 	localUfrag string
 	localPwd   string
+
+	// we record the last known state of the ICE agent to avoid duplicate on disconnected events
+	lastKnownState ice.ConnectionState
 }
 
 func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool, callBacks WorkerICECallbacks) (*WorkerICE, error) {
@@ -93,10 +93,8 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 
 	var preferredCandidateTypes []ice.CandidateType
 	if w.hasRelayOnLocally && remoteOfferAnswer.RelaySrvAddress != "" {
-		w.selectedPriority = connPriorityICEP2P
 		preferredCandidateTypes = icemaker.CandidateTypesP2P()
 	} else {
-		w.selectedPriority = connPriorityICETurn
 		preferredCandidateTypes = icemaker.CandidateTypes()
 	}
 
@@ -127,9 +125,6 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		w.log.Debugf("failed to dial the remote peer: %s", err)
 		return
 	}
-	w.log.Infof("check remoteConn: %v", remoteConn)
-	w.log.Infof("check remoteConn.RemoteAddr: %v", remoteConn.RemoteAddr())
-	nilCheck(w.log, remoteConn)
 	w.log.Debugf("agent dial succeeded")
 
 	pair, err := w.agent.GetSelectedCandidatePair()
@@ -158,10 +153,9 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		RemoteIceCandidateEndpoint: fmt.Sprintf("%s:%d", pair.Remote.Address(), pair.Remote.Port()),
 		Relayed:                    isRelayed(pair),
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
-		Agent:                      agent,
 	}
 	w.log.Debugf("on ICE conn read to use ready")
-	go w.conn.OnConnReady(w.selectedPriority, ci)
+	go w.conn.OnConnReady(selectedPriority(pair), ci)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -199,8 +193,7 @@ func (w *WorkerICE) Close() {
 		return
 	}
 
-	err := w.agent.Close()
-	if err != nil {
+	if err := w.agent.Close(); err != nil {
 		w.log.Warnf("failed to close ICE agent: %s", err)
 	}
 }
@@ -220,15 +213,18 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 
 	err = agent.OnConnectionStateChange(func(state ice.ConnectionState) {
 		w.log.Debugf("ICE ConnectionState has changed to %s", state.String())
-		if state == ice.ConnectionStateFailed || state == ice.ConnectionStateDisconnected {
-			w.conn.OnStatusChanged(StatusDisconnected)
-
-			w.muxAgent.Lock()
-			agentCancel()
-			_ = agent.Close()
-			w.agent = nil
-
-			w.muxAgent.Unlock()
+		switch state {
+		case ice.ConnectionStateConnected:
+			w.lastKnownState = ice.ConnectionStateConnected
+			return
+		case ice.ConnectionStateFailed, ice.ConnectionStateDisconnected:
+			if w.lastKnownState != ice.ConnectionStateDisconnected {
+				w.lastKnownState = ice.ConnectionStateDisconnected
+				w.conn.OnStatusChanged(StatusDisconnected)
+			}
+			w.closeAgent(agentCancel)
+		default:
+			return
 		}
 	})
 	if err != nil {
@@ -254,10 +250,27 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 	return agent, nil
 }
 
+func (w *WorkerICE) closeAgent(cancel context.CancelFunc) {
+	w.muxAgent.Lock()
+	defer w.muxAgent.Unlock()
+
+	cancel()
+	if err := w.agent.Close(); err != nil {
+		w.log.Warnf("failed to close ICE agent: %s", err)
+	}
+	w.agent = nil
+}
+
 func (w *WorkerICE) punchRemoteWGPort(pair *ice.CandidatePair, remoteWgPort int) {
 	// wait local endpoint configuration
 	time.Sleep(time.Second)
-	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pair.Remote.Address(), remoteWgPort))
+	addrString := pair.Remote.Address()
+	parsed, err := netip.ParseAddr(addrString)
+	if (err == nil) && (parsed.Is6()) {
+		addrString = fmt.Sprintf("[%s]", addrString)
+		//IPv6 Literals need to be wrapped in brackets for Resolve*Addr()
+	}
+	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", addrString, remoteWgPort))
 	if err != nil {
 		w.log.Warnf("got an error while resolving the udp address, err: %s", err)
 		return
@@ -327,10 +340,8 @@ func (w *WorkerICE) shouldSendExtraSrflxCandidate(candidate ice.Candidate) bool
 func (w *WorkerICE) turnAgentDial(ctx context.Context, remoteOfferAnswer *OfferAnswer) (*ice.Conn, error) {
 	isControlling := w.config.LocalKey > w.config.Key
 	if isControlling {
-		w.log.Infof("dialing remote peer %s as controlling", w.config.Key)
 		return w.agent.Dial(ctx, remoteOfferAnswer.IceCredentials.UFrag, remoteOfferAnswer.IceCredentials.Pwd)
 	} else {
-		w.log.Infof("dialing remote peer %s as controlled", w.config.Key)
 		return w.agent.Accept(ctx, remoteOfferAnswer.IceCredentials.UFrag, remoteOfferAnswer.IceCredentials.Pwd)
 	}
 }
@@ -385,3 +396,11 @@ func isRelayed(pair *ice.CandidatePair) bool {
 	}
 	return false
 }
+
+func selectedPriority(pair *ice.CandidatePair) ConnPriority {
+	if isRelayed(pair) {
+		return connPriorityICETurn
+	} else {
+		return connPriorityICEP2P
+	}
+}

client/internal/peerstore/store.go

@@ -0,0 +1,87 @@
+package peerstore
+
+import (
+	"net"
+	"sync"
+
+	"golang.org/x/exp/maps"
+
+	"github.com/netbirdio/netbird/client/internal/peer"
+)
+
+// Store is a thread-safe store for peer connections.
+type Store struct {
+	peerConns   map[string]*peer.Conn
+	peerConnsMu sync.RWMutex
+}
+
+func NewConnStore() *Store {
+	return &Store{
+		peerConns: make(map[string]*peer.Conn),
+	}
+}
+
+func (s *Store) AddPeerConn(pubKey string, conn *peer.Conn) bool {
+	s.peerConnsMu.Lock()
+	defer s.peerConnsMu.Unlock()
+
+	_, ok := s.peerConns[pubKey]
+	if ok {
+		return false
+	}
+
+	s.peerConns[pubKey] = conn
+	return true
+}
+
+func (s *Store) Remove(pubKey string) (*peer.Conn, bool) {
+	s.peerConnsMu.Lock()
+	defer s.peerConnsMu.Unlock()
+
+	p, ok := s.peerConns[pubKey]
+	if !ok {
+		return nil, false
+	}
+	delete(s.peerConns, pubKey)
+	return p, true
+}
+
+func (s *Store) AllowedIPs(pubKey string) (string, bool) {
+	s.peerConnsMu.RLock()
+	defer s.peerConnsMu.RUnlock()
+
+	p, ok := s.peerConns[pubKey]
+	if !ok {
+		return "", false
+	}
+	return p.WgConfig().AllowedIps, true
+}
+
+func (s *Store) AllowedIP(pubKey string) (net.IP, bool) {
+	s.peerConnsMu.RLock()
+	defer s.peerConnsMu.RUnlock()
+
+	p, ok := s.peerConns[pubKey]
+	if !ok {
+		return nil, false
+	}
+	return p.AllowedIP(), true
+}
+
+func (s *Store) PeerConn(pubKey string) (*peer.Conn, bool) {
+	s.peerConnsMu.RLock()
+	defer s.peerConnsMu.RUnlock()
+
+	p, ok := s.peerConns[pubKey]
+	if !ok {
+		return nil, false
+	}
+	return p, true
+}
+
+func (s *Store) PeersPubKey() []string {
+	s.peerConnsMu.RLock()
+	defer s.peerConnsMu.RUnlock()
+
+	return maps.Keys(s.peerConns)
+}

client/internal/routemanager/client.go

@@ -13,12 +13,20 @@ import (
 	"github.com/netbirdio/netbird/client/iface"
 	nbdns "github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
+	"github.com/netbirdio/netbird/client/internal/routemanager/dnsinterceptor"
 	"github.com/netbirdio/netbird/client/internal/routemanager/dynamic"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	"github.com/netbirdio/netbird/client/internal/routemanager/static"
 	"github.com/netbirdio/netbird/route"
 )
 
+const (
+	handlerTypeDynamic = iota
+	handlerTypeDomain
+	handlerTypeStatic
+)
+
 type routerPeerStatus struct {
 	connected bool
 	relayed   bool
@@ -53,7 +61,18 @@ type clientNetwork struct {
 	updateSerial        uint64
 }
 
-func newClientNetworkWatcher(ctx context.Context, dnsRouteInterval time.Duration, wgInterface iface.IWGIface, statusRecorder *peer.Status, rt *route.Route, routeRefCounter *refcounter.RouteRefCounter, allowedIPsRefCounter *refcounter.AllowedIPsRefCounter) *clientNetwork {
+func newClientNetworkWatcher(
+	ctx context.Context,
+	dnsRouteInterval time.Duration,
+	wgInterface iface.IWGIface,
+	statusRecorder *peer.Status,
+	rt *route.Route,
+	routeRefCounter *refcounter.RouteRefCounter,
+	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter,
+	dnsServer nbdns.Server,
+	peerStore *peerstore.Store,
+	useNewDNSRoute bool,
+) *clientNetwork {
 	ctx, cancel := context.WithCancel(ctx)
 
 	client := &clientNetwork{
@@ -65,7 +84,17 @@ func newClientNetworkWatcher(ctx context.Context, dnsRouteInterval time.Duration
 		routePeersNotifiers: make(map[string]chan struct{}),
 		routeUpdate:         make(chan routesUpdate),
 		peerStateUpdate:     make(chan struct{}),
-		handler:             handlerFromRoute(rt, routeRefCounter, allowedIPsRefCounter, dnsRouteInterval, statusRecorder, wgInterface),
+		handler: handlerFromRoute(
+			rt,
+			routeRefCounter,
+			allowedIPsRefCounter,
+			dnsRouteInterval,
+			statusRecorder,
+			wgInterface,
+			dnsServer,
+			peerStore,
+			useNewDNSRoute,
+		),
 	}
 	return client
 }
@@ -122,13 +151,20 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID]
 			tempScore = float64(metricDiff) * 10
 		}
 
-		// in some temporal cases, latency can be 0, so we set it to 1s to not block but try to avoid this route
-		latency := time.Second
+		// in some temporal cases, latency can be 0, so we set it to 999ms to not block but try to avoid this route
+		latency := 999 * time.Millisecond
 		if peerStatus.latency != 0 {
 			latency = peerStatus.latency
 		} else {
-			log.Warnf("peer %s has 0 latency", r.Peer)
+			log.Tracef("peer %s has 0 latency, range %s", r.Peer, c.handler)
 		}
+
+		// avoid negative tempScore on the higher latency calculation
+		if latency > 1*time.Second {
+			latency = 999 * time.Millisecond
+		}
+
+		// higher latency is worse score
 		tempScore += 1 - latency.Seconds()
 
 		if !peerStatus.relayed {
@@ -150,6 +186,8 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID]
 		}
 	}
 
+	log.Debugf("chosen route: %s, chosen score: %f, current route: %s, current score: %f", chosen, chosenScore, currID, currScore)
+
 	switch {
 	case chosen == "":
 		var peers []string
@@ -195,15 +233,20 @@ func (c *clientNetwork) watchPeerStatusChanges(ctx context.Context, peerKey stri
 func (c *clientNetwork) startPeersStatusChangeWatcher() {
 	for _, r := range c.routes {
 		_, found := c.routePeersNotifiers[r.Peer]
-		if !found {
-			c.routePeersNotifiers[r.Peer] = make(chan struct{})
-			go c.watchPeerStatusChanges(c.ctx, r.Peer, c.peerStateUpdate, c.routePeersNotifiers[r.Peer])
+		if found {
+			continue
 		}
+
+		closerChan := make(chan struct{})
+		c.routePeersNotifiers[r.Peer] = closerChan
+		go c.watchPeerStatusChanges(c.ctx, r.Peer, c.peerStateUpdate, closerChan)
 	}
 }
 
-func (c *clientNetwork) removeRouteFromWireguardPeer() error {
-	c.removeStateRoute()
+func (c *clientNetwork) removeRouteFromWireGuardPeer() error {
+	if err := c.statusRecorder.RemovePeerStateRoute(c.currentChosen.Peer, c.handler.String()); err != nil {
+		log.Warnf("Failed to update peer state: %v", err)
+	}
 
 	if err := c.handler.RemoveAllowedIPs(); err != nil {
 		return fmt.Errorf("remove allowed IPs: %w", err)
@@ -218,7 +261,7 @@ func (c *clientNetwork) removeRouteFromPeerAndSystem() error {
 
 	var merr *multierror.Error
 
-	if err := c.removeRouteFromWireguardPeer(); err != nil {
+	if err := c.removeRouteFromWireGuardPeer(); err != nil {
 		merr = multierror.Append(merr, fmt.Errorf("remove allowed IPs for peer %s: %w", c.currentChosen.Peer, err))
 	}
 	if err := c.handler.RemoveRoute(); err != nil {
@@ -257,7 +300,7 @@ func (c *clientNetwork) recalculateRouteAndUpdatePeerAndSystem() error {
 		}
 	} else {
 		// Otherwise, remove the allowed IPs from the previous peer first
-		if err := c.removeRouteFromWireguardPeer(); err != nil {
+		if err := c.removeRouteFromWireGuardPeer(); err != nil {
 			return fmt.Errorf("remove allowed IPs for peer %s: %w", c.currentChosen.Peer, err)
 		}
 	}
@@ -268,37 +311,13 @@ func (c *clientNetwork) recalculateRouteAndUpdatePeerAndSystem() error {
 		return fmt.Errorf("add allowed IPs for peer %s: %w", c.currentChosen.Peer, err)
 	}
 
-	c.addStateRoute()
-
+	err := c.statusRecorder.AddPeerStateRoute(c.currentChosen.Peer, c.handler.String())
+	if err != nil {
+		return fmt.Errorf("add peer state route: %w", err)
+	}
 	return nil
 }
 
-func (c *clientNetwork) addStateRoute() {
-	state, err := c.statusRecorder.GetPeer(c.currentChosen.Peer)
-	if err != nil {
-		log.Errorf("Failed to get peer state: %v", err)
-		return
-	}
-
-	state.AddRoute(c.handler.String())
-	if err := c.statusRecorder.UpdatePeerState(state); err != nil {
-		log.Warnf("Failed to update peer state: %v", err)
-	}
-}
-
-func (c *clientNetwork) removeStateRoute() {
-	state, err := c.statusRecorder.GetPeer(c.currentChosen.Peer)
-	if err != nil {
-		log.Errorf("Failed to get peer state: %v", err)
-		return
-	}
-
-	state.DeleteRoute(c.handler.String())
-	if err := c.statusRecorder.UpdatePeerState(state); err != nil {
-		log.Warnf("Failed to update peer state: %v", err)
-	}
-}
-
 func (c *clientNetwork) sendUpdateToClientNetworkWatcher(update routesUpdate) {
 	go func() {
 		c.routeUpdate <- update
@@ -378,10 +397,50 @@ func (c *clientNetwork) peersStateAndUpdateWatcher() {
 	}
 }
 
-func handlerFromRoute(rt *route.Route, routeRefCounter *refcounter.RouteRefCounter, allowedIPsRefCounter *refcounter.AllowedIPsRefCounter, dnsRouterInteval time.Duration, statusRecorder *peer.Status, wgInterface iface.IWGIface) RouteHandler {
-	if rt.IsDynamic() {
+func handlerFromRoute(
+	rt *route.Route,
+	routeRefCounter *refcounter.RouteRefCounter,
+	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter,
+	dnsRouterInteval time.Duration,
+	statusRecorder *peer.Status,
+	wgInterface iface.IWGIface,
+	dnsServer nbdns.Server,
+	peerStore *peerstore.Store,
+	useNewDNSRoute bool,
+) RouteHandler {
+	switch handlerType(rt, useNewDNSRoute) {
+	case handlerTypeDomain:
+		return dnsinterceptor.New(
+			rt,
+			routeRefCounter,
+			allowedIPsRefCounter,
+			statusRecorder,
+			dnsServer,
+			peerStore,
+		)
+	case handlerTypeDynamic:
 		dns := nbdns.NewServiceViaMemory(wgInterface)
-		return dynamic.NewRoute(rt, routeRefCounter, allowedIPsRefCounter, dnsRouterInteval, statusRecorder, wgInterface, fmt.Sprintf("%s:%d", dns.RuntimeIP(), dns.RuntimePort()))
+		return dynamic.NewRoute(
+			rt,
+			routeRefCounter,
+			allowedIPsRefCounter,
+			dnsRouterInteval,
+			statusRecorder,
+			wgInterface,
+			fmt.Sprintf("%s:%d", dns.RuntimeIP(), dns.RuntimePort()),
+		)
+	default:
+		return static.NewRoute(rt, routeRefCounter, allowedIPsRefCounter)
 	}
-	return static.NewRoute(rt, routeRefCounter, allowedIPsRefCounter)
+}
+
+func handlerType(rt *route.Route, useNewDNSRoute bool) int {
+	if !rt.IsDynamic() {
+		return handlerTypeStatic
+	}
+
+	if useNewDNSRoute {
+		return handlerTypeDomain
+	}
+	return handlerTypeDynamic
 }

client/internal/routemanager/client_test.go

@@ -1,6 +1,7 @@
 package routemanager
 
 import (
+	"fmt"
 	"net/netip"
 	"testing"
 	"time"
@@ -227,6 +228,64 @@ func TestGetBestrouteFromStatuses(t *testing.T) {
 			currentRoute:    "route1",
 			expectedRouteID: "route1",
 		},
+		{
+			name: "relayed routes with latency 0 should maintain previous choice",
+			statuses: map[route.ID]routerPeerStatus{
+				"route1": {
+					connected: true,
+					relayed:   true,
+					latency:   0 * time.Millisecond,
+				},
+				"route2": {
+					connected: true,
+					relayed:   true,
+					latency:   0 * time.Millisecond,
+				},
+			},
+			existingRoutes: map[route.ID]*route.Route{
+				"route1": {
+					ID:     "route1",
+					Metric: route.MaxMetric,
+					Peer:   "peer1",
+				},
+				"route2": {
+					ID:     "route2",
+					Metric: route.MaxMetric,
+					Peer:   "peer2",
+				},
+			},
+			currentRoute:    "route1",
+			expectedRouteID: "route1",
+		},
+		{
+			name: "p2p routes with latency 0 should maintain previous choice",
+			statuses: map[route.ID]routerPeerStatus{
+				"route1": {
+					connected: true,
+					relayed:   false,
+					latency:   0 * time.Millisecond,
+				},
+				"route2": {
+					connected: true,
+					relayed:   false,
+					latency:   0 * time.Millisecond,
+				},
+			},
+			existingRoutes: map[route.ID]*route.Route{
+				"route1": {
+					ID:     "route1",
+					Metric: route.MaxMetric,
+					Peer:   "peer1",
+				},
+				"route2": {
+					ID:     "route2",
+					Metric: route.MaxMetric,
+					Peer:   "peer2",
+				},
+			},
+			currentRoute:    "route1",
+			expectedRouteID: "route1",
+		},
 		{
 			name: "current route with bad score should be changed to route with better score",
 			statuses: map[route.ID]routerPeerStatus{
@@ -287,6 +346,45 @@ func TestGetBestrouteFromStatuses(t *testing.T) {
 		},
 	}
 
+	// fill the test data with random routes
+	for _, tc := range testCases {
+		for i := 0; i < 50; i++ {
+			dummyRoute := &route.Route{
+				ID:     route.ID(fmt.Sprintf("dummy_p1_%d", i)),
+				Metric: route.MinMetric,
+				Peer:   fmt.Sprintf("dummy_p1_%d", i),
+			}
+			tc.existingRoutes[dummyRoute.ID] = dummyRoute
+		}
+		for i := 0; i < 50; i++ {
+			dummyRoute := &route.Route{
+				ID:     route.ID(fmt.Sprintf("dummy_p2_%d", i)),
+				Metric: route.MinMetric,
+				Peer:   fmt.Sprintf("dummy_p1_%d", i),
+			}
+			tc.existingRoutes[dummyRoute.ID] = dummyRoute
+		}
+
+		for i := 0; i < 50; i++ {
+			id := route.ID(fmt.Sprintf("dummy_p1_%d", i))
+			dummyStatus := routerPeerStatus{
+				connected: false,
+				relayed:   true,
+				latency:   0,
+			}
+			tc.statuses[id] = dummyStatus
+		}
+		for i := 0; i < 50; i++ {
+			id := route.ID(fmt.Sprintf("dummy_p2_%d", i))
+			dummyStatus := routerPeerStatus{
+				connected: false,
+				relayed:   true,
+				latency:   0,
+			}
+			tc.statuses[id] = dummyStatus
+		}
+	}
+
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			currentRoute := &route.Route{

client/internal/routemanager/dnsinterceptor/handler.go

@@ -0,0 +1,356 @@
+package dnsinterceptor
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/netip"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/hashicorp/go-multierror"
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	nbdns "github.com/netbirdio/netbird/client/internal/dns"
+	"github.com/netbirdio/netbird/client/internal/dnsfwd"
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
+	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
+	"github.com/netbirdio/netbird/management/domain"
+	"github.com/netbirdio/netbird/route"
+)
+
+type domainMap map[domain.Domain][]netip.Prefix
+
+type DnsInterceptor struct {
+	mu                   sync.RWMutex
+	route                *route.Route
+	routeRefCounter      *refcounter.RouteRefCounter
+	allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
+	statusRecorder       *peer.Status
+	dnsServer            nbdns.Server
+	currentPeerKey       string
+	interceptedDomains   domainMap
+	peerStore            *peerstore.Store
+}
+
+func New(
+	rt *route.Route,
+	routeRefCounter *refcounter.RouteRefCounter,
+	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter,
+	statusRecorder *peer.Status,
+	dnsServer nbdns.Server,
+	peerStore *peerstore.Store,
+) *DnsInterceptor {
+	return &DnsInterceptor{
+		route:                rt,
+		routeRefCounter:      routeRefCounter,
+		allowedIPsRefcounter: allowedIPsRefCounter,
+		statusRecorder:       statusRecorder,
+		dnsServer:            dnsServer,
+		interceptedDomains:   make(domainMap),
+		peerStore:            peerStore,
+	}
+}
+
+func (d *DnsInterceptor) String() string {
+	return d.route.Domains.SafeString()
+}
+
+func (d *DnsInterceptor) AddRoute(context.Context) error {
+	d.dnsServer.RegisterHandler(d.route.Domains.ToPunycodeList(), d, nbdns.PriorityDNSRoute)
+	return nil
+}
+
+func (d *DnsInterceptor) RemoveRoute() error {
+	d.mu.Lock()
+
+	var merr *multierror.Error
+	for domain, prefixes := range d.interceptedDomains {
+		for _, prefix := range prefixes {
+			if _, err := d.routeRefCounter.Decrement(prefix); err != nil {
+				merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %v", prefix, err))
+			}
+			if d.currentPeerKey != "" {
+				if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+					merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
+				}
+			}
+		}
+		log.Debugf("removed dynamic route(s) for [%s]: %s", domain.SafeString(), strings.ReplaceAll(fmt.Sprintf("%s", prefixes), " ", ", "))
+
+	}
+	for _, domain := range d.route.Domains {
+		d.statusRecorder.DeleteResolvedDomainsStates(domain)
+	}
+
+	clear(d.interceptedDomains)
+	d.mu.Unlock()
+
+	d.dnsServer.DeregisterHandler(d.route.Domains.ToPunycodeList(), nbdns.PriorityDNSRoute)
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+func (d *DnsInterceptor) AddAllowedIPs(peerKey string) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	var merr *multierror.Error
+	for domain, prefixes := range d.interceptedDomains {
+		for _, prefix := range prefixes {
+			if ref, err := d.allowedIPsRefcounter.Increment(prefix, peerKey); err != nil {
+				merr = multierror.Append(merr, fmt.Errorf("add allowed IP %s: %v", prefix, err))
+			} else if ref.Count > 1 && ref.Out != peerKey {
+				log.Warnf("IP [%s] for domain [%s] is already routed by peer [%s]. HA routing disabled",
+					prefix.Addr(),
+					domain.SafeString(),
+					ref.Out,
+				)
+			}
+		}
+	}
+
+	d.currentPeerKey = peerKey
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+func (d *DnsInterceptor) RemoveAllowedIPs() error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	var merr *multierror.Error
+	for _, prefixes := range d.interceptedDomains {
+		for _, prefix := range prefixes {
+			if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
+			}
+		}
+	}
+
+	d.currentPeerKey = ""
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+// ServeDNS implements the dns.Handler interface
+func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	if len(r.Question) == 0 {
+		return
+	}
+	log.Tracef("received DNS request for domain=%s type=%v class=%v",
+		r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
+
+	d.mu.RLock()
+	peerKey := d.currentPeerKey
+	d.mu.RUnlock()
+
+	if peerKey == "" {
+		log.Tracef("no current peer key set, letting next handler try for domain=%s", r.Question[0].Name)
+
+		d.continueToNextHandler(w, r, "no current peer key")
+		return
+	}
+
+	upstreamIP, err := d.getUpstreamIP(peerKey)
+	if err != nil {
+		log.Errorf("failed to get upstream IP: %v", err)
+		d.continueToNextHandler(w, r, fmt.Sprintf("failed to get upstream IP: %v", err))
+		return
+	}
+
+	client := &dns.Client{
+		Timeout: 5 * time.Second,
+		Net:     "udp",
+	}
+	upstream := fmt.Sprintf("%s:%d", upstreamIP, dnsfwd.ListenPort)
+	reply, _, err := client.ExchangeContext(context.Background(), r, upstream)
+
+	var answer []dns.RR
+	if reply != nil {
+		answer = reply.Answer
+	}
+	log.Tracef("upstream %s (%s) DNS response for domain=%s answers=%v", upstreamIP, peerKey, r.Question[0].Name, answer)
+
+	if err != nil {
+		log.Errorf("failed to exchange DNS request with %s: %v", upstream, err)
+		if err := w.WriteMsg(&dns.Msg{MsgHdr: dns.MsgHdr{Rcode: dns.RcodeServerFailure, Id: r.Id}}); err != nil {
+			log.Errorf("failed writing DNS response: %v", err)
+		}
+		return
+	}
+
+	reply.Id = r.Id
+	if err := d.writeMsg(w, reply); err != nil {
+		log.Errorf("failed writing DNS response: %v", err)
+	}
+}
+
+// continueToNextHandler signals the handler chain to try the next handler
+func (d *DnsInterceptor) continueToNextHandler(w dns.ResponseWriter, r *dns.Msg, reason string) {
+	log.Tracef("continuing to next handler for domain=%s reason=%s", r.Question[0].Name, reason)
+
+	resp := new(dns.Msg)
+	resp.SetRcode(r, dns.RcodeNameError)
+	// Set Zero bit to signal handler chain to continue
+	resp.MsgHdr.Zero = true
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed writing DNS continue response: %v", err)
+	}
+}
+
+func (d *DnsInterceptor) getUpstreamIP(peerKey string) (net.IP, error) {
+	peerAllowedIP, exists := d.peerStore.AllowedIP(peerKey)
+	if !exists {
+		return nil, fmt.Errorf("peer connection not found for key: %s", peerKey)
+	}
+	return peerAllowedIP, nil
+}
+
+func (d *DnsInterceptor) writeMsg(w dns.ResponseWriter, r *dns.Msg) error {
+	if r == nil {
+		return fmt.Errorf("received nil DNS message")
+	}
+
+	if len(r.Answer) > 0 && len(r.Question) > 0 {
+		origPattern := ""
+		if writer, ok := w.(*nbdns.ResponseWriterChain); ok {
+			origPattern = writer.GetOrigPattern()
+		}
+
+		resolvedDomain := domain.Domain(r.Question[0].Name)
+
+		// already punycode via RegisterHandler()
+		originalDomain := domain.Domain(origPattern)
+		if originalDomain == "" {
+			originalDomain = resolvedDomain
+		}
+
+		var newPrefixes []netip.Prefix
+		for _, answer := range r.Answer {
+			var ip netip.Addr
+			switch rr := answer.(type) {
+			case *dns.A:
+				addr, ok := netip.AddrFromSlice(rr.A)
+				if !ok {
+					log.Tracef("failed to convert A record for domain=%s ip=%v", resolvedDomain, rr.A)
+					continue
+				}
+				ip = addr
+			case *dns.AAAA:
+				addr, ok := netip.AddrFromSlice(rr.AAAA)
+				if !ok {
+					log.Tracef("failed to convert AAAA record for domain=%s ip=%v", resolvedDomain, rr.AAAA)
+					continue
+				}
+				ip = addr
+			default:
+				continue
+			}
+
+			prefix := netip.PrefixFrom(ip, ip.BitLen())
+			newPrefixes = append(newPrefixes, prefix)
+		}
+
+		if len(newPrefixes) > 0 {
+			if err := d.updateDomainPrefixes(resolvedDomain, originalDomain, newPrefixes); err != nil {
+				log.Errorf("failed to update domain prefixes: %v", err)
+			}
+		}
+	}
+
+	if err := w.WriteMsg(r); err != nil {
+		return fmt.Errorf("failed to write DNS response: %v", err)
+	}
+
+	return nil
+}
+
+func (d *DnsInterceptor) updateDomainPrefixes(resolvedDomain, originalDomain domain.Domain, newPrefixes []netip.Prefix) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	oldPrefixes := d.interceptedDomains[resolvedDomain]
+	toAdd, toRemove := determinePrefixChanges(oldPrefixes, newPrefixes)
+
+	var merr *multierror.Error
+
+	// Add new prefixes
+	for _, prefix := range toAdd {
+		if _, err := d.routeRefCounter.Increment(prefix, struct{}{}); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("add route for IP %s: %v", prefix, err))
+			continue
+		}
+
+		if d.currentPeerKey == "" {
+			continue
+		}
+		if ref, err := d.allowedIPsRefcounter.Increment(prefix, d.currentPeerKey); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("add allowed IP %s: %v", prefix, err))
+		} else if ref.Count > 1 && ref.Out != d.currentPeerKey {
+			log.Warnf("IP [%s] for domain [%s] is already routed by peer [%s]. HA routing disabled",
+				prefix.Addr(),
+				resolvedDomain.SafeString(),
+				ref.Out,
+			)
+		}
+	}
+
+	if !d.route.KeepRoute {
+		// Remove old prefixes
+		for _, prefix := range toRemove {
+			if _, err := d.routeRefCounter.Decrement(prefix); err != nil {
+				merr = multierror.Append(merr, fmt.Errorf("remove route for IP %s: %v", prefix, err))
+			}
+			if d.currentPeerKey != "" {
+				if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+					merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
+				}
+			}
+		}
+	}
+
+	// Update domain prefixes using resolved domain as key
+	if len(toAdd) > 0 || len(toRemove) > 0 {
+		d.interceptedDomains[resolvedDomain] = newPrefixes
+		originalDomain = domain.Domain(strings.TrimSuffix(string(originalDomain), "."))
+		d.statusRecorder.UpdateResolvedDomainsStates(originalDomain, resolvedDomain, newPrefixes)
+
+		if len(toAdd) > 0 {
+			log.Debugf("added dynamic route(s) for domain=%s (pattern: domain=%s): %s",
+				resolvedDomain.SafeString(),
+				originalDomain.SafeString(),
+				toAdd)
+		}
+		if len(toRemove) > 0 {
+			log.Debugf("removed dynamic route(s) for domain=%s (pattern: domain=%s): %s",
+				resolvedDomain.SafeString(),
+				originalDomain.SafeString(),
+				toRemove)
+		}
+	}
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+func determinePrefixChanges(oldPrefixes, newPrefixes []netip.Prefix) (toAdd, toRemove []netip.Prefix) {
+	prefixSet := make(map[netip.Prefix]bool)
+	for _, prefix := range oldPrefixes {
+		prefixSet[prefix] = false
+	}
+	for _, prefix := range newPrefixes {
+		if _, exists := prefixSet[prefix]; exists {
+			prefixSet[prefix] = true
+		} else {
+			toAdd = append(toAdd, prefix)
+		}
+	}
+	for prefix, inUse := range prefixSet {
+		if !inUse {
+			toRemove = append(toRemove, prefix)
+		}
+	}
+	return
+}

client/internal/routemanager/dynamic/route.go

@@ -74,11 +74,7 @@ func NewRoute(
 }
 
 func (r *Route) String() string {
-	s, err := r.route.Domains.String()
-	if err != nil {
-		return r.route.Domains.PunycodeString()
-	}
-	return s
+	return r.route.Domains.SafeString()
 }
 
 func (r *Route) AddRoute(ctx context.Context) error {
@@ -292,7 +288,7 @@ func (r *Route) updateDynamicRoutes(ctx context.Context, newDomains domainMap) e
 		updatedPrefixes := combinePrefixes(oldPrefixes, removedPrefixes, addedPrefixes)
 		r.dynamicDomains[domain] = updatedPrefixes
 
-		r.statusRecorder.UpdateResolvedDomainsStates(domain, updatedPrefixes)
+		r.statusRecorder.UpdateResolvedDomainsStates(domain, domain, updatedPrefixes)
 	}
 
 	return nberrors.FormatErrorOrNil(merr)

client/internal/routemanager/manager.go

@@ -12,12 +12,15 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/configurer"
+	"github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/listener"
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/client/internal/routemanager/notifier"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
@@ -32,10 +35,12 @@ import (
 
 // Manager is a route manager interface
 type Manager interface {
-	Init(*statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error)
-	UpdateRoutes(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error)
+	Init() (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error)
+	UpdateRoutes(updateSerial uint64, newRoutes []*route.Route, useNewDNSRoute bool) error
 	TriggerSelection(route.HAMap)
 	GetRouteSelector() *routeselector.RouteSelector
+	GetClientRoutes() route.HAMap
+	GetClientRoutesWithNetID() map[route.NetID][]*route.Route
 	SetRouteChangeListener(listener listener.NetworkChangeListener)
 	InitialRouteRange() []string
 	EnableServerRouter(firewall firewall.Manager) error
@@ -59,6 +64,12 @@ type DefaultManager struct {
 	routeRefCounter      *refcounter.RouteRefCounter
 	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter
 	dnsRouteInterval     time.Duration
+	stateManager         *statemanager.Manager
+	// clientRoutes is the most recent list of clientRoutes received from the Management Service
+	clientRoutes   route.HAMap
+	dnsServer      dns.Server
+	peerStore      *peerstore.Store
+	useNewDNSRoute bool
 }
 
 func NewManager(
@@ -69,6 +80,9 @@ func NewManager(
 	statusRecorder *peer.Status,
 	relayMgr *relayClient.Manager,
 	initialRoutes []*route.Route,
+	stateManager *statemanager.Manager,
+	dnsServer dns.Server,
+	peerStore *peerstore.Store,
 ) *DefaultManager {
 	mCTX, cancel := context.WithCancel(ctx)
 	notifier := notifier.NewNotifier()
@@ -80,12 +94,14 @@ func NewManager(
 		dnsRouteInterval: dnsRouteInterval,
 		clientNetworks:   make(map[route.HAUniqueID]*clientNetwork),
 		relayMgr:         relayMgr,
-		routeSelector:    routeselector.NewRouteSelector(),
 		sysOps:           sysOps,
 		statusRecorder:   statusRecorder,
 		wgInterface:      wgInterface,
 		pubKey:           pubKey,
 		notifier:         notifier,
+		stateManager:     stateManager,
+		dnsServer:        dnsServer,
+		peerStore:        peerStore,
 	}
 
 	dm.routeRefCounter = refcounter.New(
@@ -114,14 +130,16 @@ func NewManager(
 	)
 
 	if runtime.GOOS == "android" {
-		cr := dm.clientRoutes(initialRoutes)
+		cr := dm.initialClientRoutes(initialRoutes)
 		dm.notifier.SetInitialClientRoutes(cr)
 	}
 	return dm
 }
 
 // Init sets up the routing
-func (m *DefaultManager) Init(stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (m *DefaultManager) Init() (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+	m.routeSelector = m.initSelector()
+
 	if nbnet.CustomRoutingDisabled() {
 		return nil, nil, nil
 	}
@@ -137,14 +155,36 @@ func (m *DefaultManager) Init(stateManager *statemanager.Manager) (nbnet.AddHook
 
 	ips := resolveURLsToIPs(initialAddresses)
 
-	beforePeerHook, afterPeerHook, err := m.sysOps.SetupRouting(ips, stateManager)
+	beforePeerHook, afterPeerHook, err := m.sysOps.SetupRouting(ips, m.stateManager)
 	if err != nil {
 		return nil, nil, fmt.Errorf("setup routing: %w", err)
 	}
+
 	log.Info("Routing setup complete")
 	return beforePeerHook, afterPeerHook, nil
 }
 
+func (m *DefaultManager) initSelector() *routeselector.RouteSelector {
+	var state *SelectorState
+	m.stateManager.RegisterState(state)
+
+	// restore selector state if it exists
+	if err := m.stateManager.LoadState(state); err != nil {
+		log.Warnf("failed to load state: %v", err)
+		return routeselector.NewRouteSelector()
+	}
+
+	if state := m.stateManager.GetState(state); state != nil {
+		if selector, ok := state.(*SelectorState); ok {
+			return (*routeselector.RouteSelector)(selector)
+		}
+
+		log.Warnf("failed to convert state with type %T to SelectorState", state)
+	}
+
+	return routeselector.NewRouteSelector()
+}
+
 func (m *DefaultManager) EnableServerRouter(firewall firewall.Manager) error {
 	var err error
 	m.serverRouter, err = newServerRouter(m.ctx, m.wgInterface, firewall, m.statusRecorder)
@@ -181,33 +221,41 @@ func (m *DefaultManager) Stop(stateManager *statemanager.Manager) {
 	}
 
 	m.ctx = nil
+
+	m.mux.Lock()
+	defer m.mux.Unlock()
+	m.clientRoutes = nil
 }
 
 // UpdateRoutes compares received routes with existing routes and removes, updates or adds them to the client and server maps
-func (m *DefaultManager) UpdateRoutes(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error) {
+func (m *DefaultManager) UpdateRoutes(updateSerial uint64, newRoutes []*route.Route, useNewDNSRoute bool) error {
 	select {
 	case <-m.ctx.Done():
 		log.Infof("not updating routes as context is closed")
-		return nil, nil, m.ctx.Err()
+		return nil
 	default:
-		m.mux.Lock()
-		defer m.mux.Unlock()
-
-		newServerRoutesMap, newClientRoutesIDMap := m.classifyRoutes(newRoutes)
-
-		filteredClientRoutes := m.routeSelector.FilterSelected(newClientRoutesIDMap)
-		m.updateClientNetworks(updateSerial, filteredClientRoutes)
-		m.notifier.OnNewRoutes(filteredClientRoutes)
-
-		if m.serverRouter != nil {
-			err := m.serverRouter.updateRoutes(newServerRoutesMap)
-			if err != nil {
-				return nil, nil, fmt.Errorf("update routes: %w", err)
-			}
-		}
-
-		return newServerRoutesMap, newClientRoutesIDMap, nil
 	}
+
+	m.mux.Lock()
+	defer m.mux.Unlock()
+	m.useNewDNSRoute = useNewDNSRoute
+
+	newServerRoutesMap, newClientRoutesIDMap := m.classifyRoutes(newRoutes)
+
+	filteredClientRoutes := m.routeSelector.FilterSelected(newClientRoutesIDMap)
+	m.updateClientNetworks(updateSerial, filteredClientRoutes)
+	m.notifier.OnNewRoutes(filteredClientRoutes)
+
+	if m.serverRouter != nil {
+		err := m.serverRouter.updateRoutes(newServerRoutesMap)
+		if err != nil {
+			return err
+		}
+	}
+
+	m.clientRoutes = newClientRoutesIDMap
+
+	return nil
 }
 
 // SetRouteChangeListener set RouteListener for route change Notifier
@@ -225,9 +273,24 @@ func (m *DefaultManager) GetRouteSelector() *routeselector.RouteSelector {
 	return m.routeSelector
 }
 
-// GetClientRoutes returns the client routes
-func (m *DefaultManager) GetClientRoutes() map[route.HAUniqueID]*clientNetwork {
-	return m.clientNetworks
+// GetClientRoutes returns most recent list of clientRoutes received from the Management Service
+func (m *DefaultManager) GetClientRoutes() route.HAMap {
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	return maps.Clone(m.clientRoutes)
+}
+
+// GetClientRoutesWithNetID returns the current routes from the route map, but the keys consist of the network ID only
+func (m *DefaultManager) GetClientRoutesWithNetID() map[route.NetID][]*route.Route {
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	routes := make(map[route.NetID][]*route.Route, len(m.clientRoutes))
+	for id, v := range m.clientRoutes {
+		routes[id.NetID()] = v
+	}
+	return routes
 }
 
 // TriggerSelection triggers the selection of routes, stopping deselected watchers and starting newly selected ones
@@ -247,11 +310,26 @@ func (m *DefaultManager) TriggerSelection(networks route.HAMap) {
 			continue
 		}
 
-		clientNetworkWatcher := newClientNetworkWatcher(m.ctx, m.dnsRouteInterval, m.wgInterface, m.statusRecorder, routes[0], m.routeRefCounter, m.allowedIPsRefCounter)
+		clientNetworkWatcher := newClientNetworkWatcher(
+			m.ctx,
+			m.dnsRouteInterval,
+			m.wgInterface,
+			m.statusRecorder,
+			routes[0],
+			m.routeRefCounter,
+			m.allowedIPsRefCounter,
+			m.dnsServer,
+			m.peerStore,
+			m.useNewDNSRoute,
+		)
 		m.clientNetworks[id] = clientNetworkWatcher
 		go clientNetworkWatcher.peersStateAndUpdateWatcher()
 		clientNetworkWatcher.sendUpdateToClientNetworkWatcher(routesUpdate{routes: routes})
 	}
+
+	if err := m.stateManager.UpdateState((*SelectorState)(m.routeSelector)); err != nil {
+		log.Errorf("failed to update state: %v", err)
+	}
 }
 
 // stopObsoleteClients stops the client network watcher for the networks that are not in the new list
@@ -272,7 +350,18 @@ func (m *DefaultManager) updateClientNetworks(updateSerial uint64, networks rout
 	for id, routes := range networks {
 		clientNetworkWatcher, found := m.clientNetworks[id]
 		if !found {
-			clientNetworkWatcher = newClientNetworkWatcher(m.ctx, m.dnsRouteInterval, m.wgInterface, m.statusRecorder, routes[0], m.routeRefCounter, m.allowedIPsRefCounter)
+			clientNetworkWatcher = newClientNetworkWatcher(
+				m.ctx,
+				m.dnsRouteInterval,
+				m.wgInterface,
+				m.statusRecorder,
+				routes[0],
+				m.routeRefCounter,
+				m.allowedIPsRefCounter,
+				m.dnsServer,
+				m.peerStore,
+				m.useNewDNSRoute,
+			)
 			m.clientNetworks[id] = clientNetworkWatcher
 			go clientNetworkWatcher.peersStateAndUpdateWatcher()
 		}
@@ -315,7 +404,7 @@ func (m *DefaultManager) classifyRoutes(newRoutes []*route.Route) (map[route.ID]
 	return newServerRoutesMap, newClientRoutesIDMap
 }
 
-func (m *DefaultManager) clientRoutes(initialRoutes []*route.Route) []*route.Route {
+func (m *DefaultManager) initialClientRoutes(initialRoutes []*route.Route) []*route.Route {
 	_, crMap := m.classifyRoutes(initialRoutes)
 	rs := make([]*route.Route, 0, len(crMap))
 	for _, routes := range crMap {

client/internal/routemanager/manager_test.go

@@ -424,9 +424,9 @@ func TestManagerUpdateRoutes(t *testing.T) {
 
 			statusRecorder := peer.NewRecorder("https://mgm")
 			ctx := context.TODO()
-			routeManager := NewManager(ctx, localPeerKey, 0, wgInterface, statusRecorder, nil, nil)
+			routeManager := NewManager(ctx, localPeerKey, 0, wgInterface, statusRecorder, nil, nil, nil, nil, nil)
 
-			_, _, err = routeManager.Init(nil)
+			_, _, err = routeManager.Init()
 
 			require.NoError(t, err, "should init route manager")
 			defer routeManager.Stop(nil)
@@ -436,11 +436,11 @@ func TestManagerUpdateRoutes(t *testing.T) {
 			}
 
 			if len(testCase.inputInitRoutes) > 0 {
-				_, _, err = routeManager.UpdateRoutes(testCase.inputSerial, testCase.inputRoutes)
+				_ = routeManager.UpdateRoutes(testCase.inputSerial, testCase.inputRoutes, false)
 				require.NoError(t, err, "should update routes with init routes")
 			}
 
-			_, _, err = routeManager.UpdateRoutes(testCase.inputSerial+uint64(len(testCase.inputInitRoutes)), testCase.inputRoutes)
+			_ = routeManager.UpdateRoutes(testCase.inputSerial+uint64(len(testCase.inputInitRoutes)), testCase.inputRoutes, false)
 			require.NoError(t, err, "should update routes")
 
 			expectedWatchers := testCase.clientNetworkWatchersExpected

client/internal/routemanager/mock.go

@@ -2,7 +2,6 @@ package routemanager
 
 import (
 	"context"
-	"fmt"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface"
@@ -15,13 +14,15 @@ import (
 
 // MockManager is the mock instance of a route manager
 type MockManager struct {
-	UpdateRoutesFunc     func(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error)
-	TriggerSelectionFunc func(haMap route.HAMap)
-	GetRouteSelectorFunc func() *routeselector.RouteSelector
-	StopFunc             func(manager *statemanager.Manager)
+	UpdateRoutesFunc             func(updateSerial uint64, newRoutes []*route.Route) error
+	TriggerSelectionFunc         func(haMap route.HAMap)
+	GetRouteSelectorFunc         func() *routeselector.RouteSelector
+	GetClientRoutesFunc          func() route.HAMap
+	GetClientRoutesWithNetIDFunc func() map[route.NetID][]*route.Route
+	StopFunc                     func(manager *statemanager.Manager)
 }
 
-func (m *MockManager) Init(*statemanager.Manager) (net.AddHookFunc, net.RemoveHookFunc, error) {
+func (m *MockManager) Init() (net.AddHookFunc, net.RemoveHookFunc, error) {
 	return nil, nil, nil
 }
 
@@ -31,11 +32,11 @@ func (m *MockManager) InitialRouteRange() []string {
 }
 
 // UpdateRoutes mock implementation of UpdateRoutes from Manager interface
-func (m *MockManager) UpdateRoutes(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error) {
+func (m *MockManager) UpdateRoutes(updateSerial uint64, newRoutes []*route.Route, b bool) error {
 	if m.UpdateRoutesFunc != nil {
 		return m.UpdateRoutesFunc(updateSerial, newRoutes)
 	}
-	return nil, nil, fmt.Errorf("method UpdateRoutes is not implemented")
+	return nil
 }
 
 func (m *MockManager) TriggerSelection(networks route.HAMap) {
@@ -52,6 +53,22 @@ func (m *MockManager) GetRouteSelector() *routeselector.RouteSelector {
 	return nil
 }
 
+// GetClientRoutes mock implementation of GetClientRoutes from Manager interface
+func (m *MockManager) GetClientRoutes() route.HAMap {
+	if m.GetClientRoutesFunc != nil {
+		return m.GetClientRoutesFunc()
+	}
+	return nil
+}
+
+// GetClientRoutesWithNetID mock implementation of GetClientRoutesWithNetID from Manager interface
+func (m *MockManager) GetClientRoutesWithNetID() map[route.NetID][]*route.Route {
+	if m.GetClientRoutesWithNetIDFunc != nil {
+		return m.GetClientRoutesWithNetIDFunc()
+	}
+	return nil
+}
+
 // Start mock implementation of Start from Manager interface
 func (m *MockManager) Start(ctx context.Context, iface *iface.WGIface) {
 }

client/internal/routemanager/refcounter/refcounter.go

@@ -47,10 +47,9 @@ type RemoveFunc[Key, O any] func(key Key, out O) error
 type Counter[Key comparable, I, O any] struct {
 	// refCountMap keeps track of the reference Ref for keys
 	refCountMap map[Key]Ref[O]
-	refCountMu  sync.Mutex
+	mu          sync.Mutex
 	// idMap keeps track of the keys associated with an ID for removal
 	idMap  map[string][]Key
-	idMu   sync.Mutex
 	add    AddFunc[Key, I, O]
 	remove RemoveFunc[Key, O]
 }
@@ -72,13 +71,14 @@ func New[Key comparable, I, O any](add AddFunc[Key, I, O], remove RemoveFunc[Key
 }
 
 // LoadData loads the data from the existing counter
+// The passed counter should not be used any longer after calling this function.
 func (rm *Counter[Key, I, O]) LoadData(
 	existingCounter *Counter[Key, I, O],
 ) {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+	existingCounter.mu.Lock()
+	defer existingCounter.mu.Unlock()
 
 	rm.refCountMap = existingCounter.refCountMap
 	rm.idMap = existingCounter.idMap
@@ -87,8 +87,8 @@ func (rm *Counter[Key, I, O]) LoadData(
 // Get retrieves the current reference count and associated data for a key.
 // If the key doesn't exist, it returns a zero value Ref and false.
 func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
 	ref, ok := rm.refCountMap[key]
 	return ref, ok
@@ -97,9 +97,13 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
 // Increment increments the reference count for the given key.
 // If this is the first reference to the key, the AddFunc is called.
 func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
+	return rm.increment(key, in)
+}
+
+func (rm *Counter[Key, I, O]) increment(key Key, in I) (Ref[O], error) {
 	ref := rm.refCountMap[key]
 	logCallerF("Increasing ref count [%d -> %d] for key %v with In [%v] Out [%v]", ref.Count, ref.Count+1, key, in, ref.Out)
 
@@ -126,10 +130,10 @@ func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {
 // IncrementWithID increments the reference count for the given key and groups it under the given ID.
 // If this is the first reference to the key, the AddFunc is called.
 func (rm *Counter[Key, I, O]) IncrementWithID(id string, key Key, in I) (Ref[O], error) {
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
-	ref, err := rm.Increment(key, in)
+	ref, err := rm.increment(key, in)
 	if err != nil {
 		return ref, fmt.Errorf("with ID: %w", err)
 	}
@@ -141,9 +145,12 @@ func (rm *Counter[Key, I, O]) IncrementWithID(id string, key Key, in I) (Ref[O],
 // Decrement decrements the reference count for the given key.
 // If the reference count reaches 0, the RemoveFunc is called.
 func (rm *Counter[Key, I, O]) Decrement(key Key) (Ref[O], error) {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+	return rm.decrement(key)
+}
 
+func (rm *Counter[Key, I, O]) decrement(key Key) (Ref[O], error) {
 	ref, ok := rm.refCountMap[key]
 	if !ok {
 		logCallerF("No reference found for key %v", key)
@@ -168,12 +175,12 @@ func (rm *Counter[Key, I, O]) Decrement(key Key) (Ref[O], error) {
 // DecrementWithID decrements the reference count for all keys associated with the given ID.
 // If the reference count reaches 0, the RemoveFunc is called.
 func (rm *Counter[Key, I, O]) DecrementWithID(id string) error {
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
 	var merr *multierror.Error
 	for _, key := range rm.idMap[id] {
-		if _, err := rm.Decrement(key); err != nil {
+		if _, err := rm.decrement(key); err != nil {
 			merr = multierror.Append(merr, err)
 		}
 	}
@@ -184,10 +191,8 @@ func (rm *Counter[Key, I, O]) DecrementWithID(id string) error {
 
 // Flush removes all references and calls RemoveFunc for each key.
 func (rm *Counter[Key, I, O]) Flush() error {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
 	var merr *multierror.Error
 	for key := range rm.refCountMap {
@@ -206,10 +211,8 @@ func (rm *Counter[Key, I, O]) Flush() error {
 
 // Clear removes all references without calling RemoveFunc.
 func (rm *Counter[Key, I, O]) Clear() {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
 
 	clear(rm.refCountMap)
 	clear(rm.idMap)
@@ -217,6 +220,9 @@ func (rm *Counter[Key, I, O]) Clear() {
 
 // MarshalJSON implements the json.Marshaler interface for Counter.
 func (rm *Counter[Key, I, O]) MarshalJSON() ([]byte, error) {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
 	return json.Marshal(struct {
 		RefCountMap map[Key]Ref[O]   `json:"refCountMap"`
 		IDMap       map[string][]Key `json:"idMap"`
@@ -228,6 +234,9 @@ func (rm *Counter[Key, I, O]) MarshalJSON() ([]byte, error) {
 
 // UnmarshalJSON implements the json.Unmarshaler interface for Counter.
 func (rm *Counter[Key, I, O]) UnmarshalJSON(data []byte) error {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
 	var temp struct {
 		RefCountMap map[Key]Ref[O]   `json:"refCountMap"`
 		IDMap       map[string][]Key `json:"idMap"`
@@ -238,6 +247,13 @@ func (rm *Counter[Key, I, O]) UnmarshalJSON(data []byte) error {
 	rm.refCountMap = temp.RefCountMap
 	rm.idMap = temp.IDMap
 
+	if temp.RefCountMap == nil {
+		temp.RefCountMap = map[Key]Ref[O]{}
+	}
+	if temp.IDMap == nil {
+		temp.IDMap = map[string][]Key{}
+	}
+
 	return nil
 }
 

client/internal/routemanager/state.go

@@ -0,0 +1,19 @@
+package routemanager
+
+import (
+	"github.com/netbirdio/netbird/client/internal/routeselector"
+)
+
+type SelectorState routeselector.RouteSelector
+
+func (s *SelectorState) Name() string {
+	return "routeselector_state"
+}
+
+func (s *SelectorState) MarshalJSON() ([]byte, error) {
+	return (*routeselector.RouteSelector)(s).MarshalJSON()
+}
+
+func (s *SelectorState) UnmarshalJSON(data []byte) error {
+	return (*routeselector.RouteSelector)(s).UnmarshalJSON(data)
+}

client/internal/routemanager/systemops/state.go

@@ -2,31 +2,28 @@ package systemops
 
 import (
 	"net/netip"
-	"sync"
 
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 )
 
-type ShutdownState struct {
-	Counter *ExclusionCounter `json:"counter,omitempty"`
-	mu      sync.RWMutex
-}
+type ShutdownState ExclusionCounter
 
 func (s *ShutdownState) Name() string {
 	return "route_state"
 }
 
 func (s *ShutdownState) Cleanup() error {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-
-	if s.Counter == nil {
-		return nil
-	}
-
 	sysops := NewSysOps(nil, nil)
 	sysops.refCounter = refcounter.New[netip.Prefix, struct{}, Nexthop](nil, sysops.removeFromRouteTable)
-	sysops.refCounter.LoadData(s.Counter)
+	sysops.refCounter.LoadData((*ExclusionCounter)(s))
 
 	return sysops.refCounter.Flush()
 }
+
+func (s *ShutdownState) MarshalJSON() ([]byte, error) {
+	return (*ExclusionCounter)(s).MarshalJSON()
+}
+
+func (s *ShutdownState) UnmarshalJSON(data []byte) error {
+	return (*ExclusionCounter)(s).UnmarshalJSON(data)
+}

client/internal/routemanager/systemops/systemops_generic.go

@@ -57,30 +57,19 @@ func (r *SysOps) setupRefCounter(initAddresses []net.IP, stateManager *statemana
 				return nexthop, refcounter.ErrIgnore
 			}
 
-			r.updateState(stateManager)
-
 			return nexthop, err
 		},
-		func(prefix netip.Prefix, nexthop Nexthop) error {
-			// remove from state even if we have trouble removing it from the route table
-			// it could be already gone
-			r.updateState(stateManager)
-
-			return r.removeFromRouteTable(prefix, nexthop)
-		},
+		r.removeFromRouteTable,
 	)
 
 	r.refCounter = refCounter
 
-	return r.setupHooks(initAddresses)
+	return r.setupHooks(initAddresses, stateManager)
 }
 
+// updateState updates state on every change so it will be persisted regularly
 func (r *SysOps) updateState(stateManager *statemanager.Manager) {
-	state := getState(stateManager)
-
-	state.Counter = r.refCounter
-
-	if err := stateManager.UpdateState(state); err != nil {
+	if err := stateManager.UpdateState((*ShutdownState)(r.refCounter)); err != nil {
 		log.Errorf("failed to update state: %v", err)
 	}
 }
@@ -336,7 +325,7 @@ func (r *SysOps) genericRemoveVPNRoute(prefix netip.Prefix, intf *net.Interface)
 	return r.removeFromRouteTable(prefix, nextHop)
 }
 
-func (r *SysOps) setupHooks(initAddresses []net.IP) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	beforeHook := func(connID nbnet.ConnectionID, ip net.IP) error {
 		prefix, err := util.GetPrefixFromIP(ip)
 		if err != nil {
@@ -347,6 +336,8 @@ func (r *SysOps) setupHooks(initAddresses []net.IP) (nbnet.AddHookFunc, nbnet.Re
 			return fmt.Errorf("adding route reference: %v", err)
 		}
 
+		r.updateState(stateManager)
+
 		return nil
 	}
 	afterHook := func(connID nbnet.ConnectionID) error {
@@ -354,6 +345,8 @@ func (r *SysOps) setupHooks(initAddresses []net.IP) (nbnet.AddHookFunc, nbnet.Re
 			return fmt.Errorf("remove route reference: %w", err)
 		}
 
+		r.updateState(stateManager)
+
 		return nil
 	}
 
@@ -532,14 +525,3 @@ func isVpnRoute(addr netip.Addr, vpnRoutes []netip.Prefix, localRoutes []netip.P
 	// Return true if the longest matching prefix is from vpnRoutes
 	return isVpn, longestPrefix
 }
-
-func getState(stateManager *statemanager.Manager) *ShutdownState {
-	var shutdownState *ShutdownState
-	if state := stateManager.GetState(shutdownState); state != nil {
-		shutdownState = state.(*ShutdownState)
-	} else {
-		shutdownState = &ShutdownState{}
-	}
-
-	return shutdownState
-}

client/internal/routemanager/systemops/systemops_linux.go

@@ -55,7 +55,7 @@ type ruleParams struct {
 
 // isLegacy determines whether to use the legacy routing setup
 func isLegacy() bool {
-	return os.Getenv("NB_USE_LEGACY_ROUTING") == "true" || nbnet.CustomRoutingDisabled()
+	return os.Getenv("NB_USE_LEGACY_ROUTING") == "true" || nbnet.CustomRoutingDisabled() || nbnet.SkipSocketMark()
 }
 
 // setIsLegacy sets the legacy routing setup
@@ -92,17 +92,6 @@ func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager
 		return r.setupRefCounter(initAddresses, stateManager)
 	}
 
-	if err = addRoutingTableName(); err != nil {
-		log.Errorf("Error adding routing table name: %v", err)
-	}
-
-	originalValues, err := sysctl.Setup(r.wgInterface)
-	if err != nil {
-		log.Errorf("Error setting up sysctl: %v", err)
-		sysctlFailed = true
-	}
-	originalSysctl = originalValues
-
 	defer func() {
 		if err != nil {
 			if cleanErr := r.CleanupRouting(stateManager); cleanErr != nil {
@@ -123,6 +112,17 @@ func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager
 		}
 	}
 
+	if err = addRoutingTableName(); err != nil {
+		log.Errorf("Error adding routing table name: %v", err)
+	}
+
+	originalValues, err := sysctl.Setup(r.wgInterface)
+	if err != nil {
+		log.Errorf("Error setting up sysctl: %v", err)
+		sysctlFailed = true
+	}
+	originalSysctl = originalValues
+
 	return nil, nil, nil
 }
 

client/internal/routemanager/systemops/systemops_windows.go

@@ -230,10 +230,13 @@ func (rm *RouteMonitor) parseUpdate(row *MIB_IPFORWARD_ROW2, notificationType MI
 	if idx != 0 {
 		intf, err := net.InterfaceByIndex(idx)
 		if err != nil {
-			return update, fmt.Errorf("get interface name: %w", err)
+			log.Warnf("failed to get interface name for index %d: %v", idx, err)
+			update.Interface = &net.Interface{
+				Index: idx,
+			}
+		} else {
+			update.Interface = intf
 		}
-
-		update.Interface = intf
 	}
 
 	log.Tracef("Received route update with destination %v, next hop %v, interface %v", row.DestinationPrefix, row.NextHop, update.Interface)

client/internal/routeselector/routeselector.go

@@ -1,8 +1,10 @@
 package routeselector
 
 import (
+	"encoding/json"
 	"fmt"
 	"slices"
+	"sync"
 
 	"github.com/hashicorp/go-multierror"
 	"golang.org/x/exp/maps"
@@ -12,6 +14,7 @@ import (
 )
 
 type RouteSelector struct {
+	mu             sync.RWMutex
 	selectedRoutes map[route.NetID]struct{}
 	selectAll      bool
 }
@@ -26,6 +29,9 @@ func NewRouteSelector() *RouteSelector {
 
 // SelectRoutes updates the selected routes based on the provided route IDs.
 func (rs *RouteSelector) SelectRoutes(routes []route.NetID, appendRoute bool, allRoutes []route.NetID) error {
+	rs.mu.Lock()
+	defer rs.mu.Unlock()
+
 	if !appendRoute {
 		rs.selectedRoutes = map[route.NetID]struct{}{}
 	}
@@ -46,6 +52,9 @@ func (rs *RouteSelector) SelectRoutes(routes []route.NetID, appendRoute bool, al
 
 // SelectAllRoutes sets the selector to select all routes.
 func (rs *RouteSelector) SelectAllRoutes() {
+	rs.mu.Lock()
+	defer rs.mu.Unlock()
+
 	rs.selectAll = true
 	rs.selectedRoutes = map[route.NetID]struct{}{}
 }
@@ -53,6 +62,9 @@ func (rs *RouteSelector) SelectAllRoutes() {
 // DeselectRoutes removes specific routes from the selection.
 // If the selector is in "select all" mode, it will transition to "select specific" mode.
 func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.NetID) error {
+	rs.mu.Lock()
+	defer rs.mu.Unlock()
+
 	if rs.selectAll {
 		rs.selectAll = false
 		rs.selectedRoutes = map[route.NetID]struct{}{}
@@ -76,12 +88,18 @@ func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.
 
 // DeselectAllRoutes deselects all routes, effectively disabling route selection.
 func (rs *RouteSelector) DeselectAllRoutes() {
+	rs.mu.Lock()
+	defer rs.mu.Unlock()
+
 	rs.selectAll = false
 	rs.selectedRoutes = map[route.NetID]struct{}{}
 }
 
 // IsSelected checks if a specific route is selected.
 func (rs *RouteSelector) IsSelected(routeID route.NetID) bool {
+	rs.mu.RLock()
+	defer rs.mu.RUnlock()
+
 	if rs.selectAll {
 		return true
 	}
@@ -91,6 +109,9 @@ func (rs *RouteSelector) IsSelected(routeID route.NetID) bool {
 
 // FilterSelected removes unselected routes from the provided map.
 func (rs *RouteSelector) FilterSelected(routes route.HAMap) route.HAMap {
+	rs.mu.RLock()
+	defer rs.mu.RUnlock()
+
 	if rs.selectAll {
 		return maps.Clone(routes)
 	}
@@ -103,3 +124,49 @@ func (rs *RouteSelector) FilterSelected(routes route.HAMap) route.HAMap {
 	}
 	return filtered
 }
+
+// MarshalJSON implements the json.Marshaler interface
+func (rs *RouteSelector) MarshalJSON() ([]byte, error) {
+	rs.mu.RLock()
+	defer rs.mu.RUnlock()
+
+	return json.Marshal(struct {
+		SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
+		SelectAll      bool                     `json:"select_all"`
+	}{
+		SelectAll:      rs.selectAll,
+		SelectedRoutes: rs.selectedRoutes,
+	})
+}
+
+// UnmarshalJSON implements the json.Unmarshaler interface
+// If the JSON is empty or null, it will initialize like a NewRouteSelector.
+func (rs *RouteSelector) UnmarshalJSON(data []byte) error {
+	rs.mu.Lock()
+	defer rs.mu.Unlock()
+
+	// Check for null or empty JSON
+	if len(data) == 0 || string(data) == "null" {
+		rs.selectedRoutes = map[route.NetID]struct{}{}
+		rs.selectAll = true
+		return nil
+	}
+
+	var temp struct {
+		SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
+		SelectAll      bool                     `json:"select_all"`
+	}
+
+	if err := json.Unmarshal(data, &temp); err != nil {
+		return err
+	}
+
+	rs.selectedRoutes = temp.SelectedRoutes
+	rs.selectAll = temp.SelectAll
+
+	if rs.selectedRoutes == nil {
+		rs.selectedRoutes = map[route.NetID]struct{}{}
+	}
+
+	return nil
+}

client/internal/routeselector/routeselector_test.go

@@ -273,3 +273,88 @@ func TestRouteSelector_FilterSelected(t *testing.T) {
 		"route2|192.168.0.0/16": {},
 	}, filtered)
 }
+
+func TestRouteSelector_NewRoutesBehavior(t *testing.T) {
+	initialRoutes := []route.NetID{"route1", "route2", "route3"}
+	newRoutes := []route.NetID{"route1", "route2", "route3", "route4", "route5"}
+
+	tests := []struct {
+		name            string
+		initialState    func(rs *routeselector.RouteSelector) error // Setup initial state
+		wantNewSelected []route.NetID                               // Expected selected routes after new routes appear
+	}{
+		{
+			name: "New routes with initial selectAll state",
+			initialState: func(rs *routeselector.RouteSelector) error {
+				rs.SelectAllRoutes()
+				return nil
+			},
+			// When selectAll is true, all routes including new ones should be selected
+			wantNewSelected: []route.NetID{"route1", "route2", "route3", "route4", "route5"},
+		},
+		{
+			name: "New routes after specific selection",
+			initialState: func(rs *routeselector.RouteSelector) error {
+				return rs.SelectRoutes([]route.NetID{"route1", "route2"}, false, initialRoutes)
+			},
+			// When specific routes were selected, new routes should remain unselected
+			wantNewSelected: []route.NetID{"route1", "route2"},
+		},
+		{
+			name: "New routes after deselect all",
+			initialState: func(rs *routeselector.RouteSelector) error {
+				rs.DeselectAllRoutes()
+				return nil
+			},
+			// After deselect all, new routes should remain unselected
+			wantNewSelected: []route.NetID{},
+		},
+		{
+			name: "New routes after deselecting specific routes",
+			initialState: func(rs *routeselector.RouteSelector) error {
+				rs.SelectAllRoutes()
+				return rs.DeselectRoutes([]route.NetID{"route1"}, initialRoutes)
+			},
+			// After deselecting specific routes, new routes should remain unselected
+			wantNewSelected: []route.NetID{"route2", "route3"},
+		},
+		{
+			name: "New routes after selecting with append",
+			initialState: func(rs *routeselector.RouteSelector) error {
+				return rs.SelectRoutes([]route.NetID{"route1"}, true, initialRoutes)
+			},
+			// When routes were appended, new routes should remain unselected
+			wantNewSelected: []route.NetID{"route1"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rs := routeselector.NewRouteSelector()
+
+			// Setup initial state
+			err := tt.initialState(rs)
+			require.NoError(t, err)
+
+			// Verify selection state with new routes
+			for _, id := range newRoutes {
+				assert.Equal(t, rs.IsSelected(id), slices.Contains(tt.wantNewSelected, id),
+					"Route %s selection state incorrect", id)
+			}
+
+			// Additional verification using FilterSelected
+			routes := route.HAMap{
+				"route1|10.0.0.0/8":     {},
+				"route2|192.168.0.0/16": {},
+				"route3|172.16.0.0/12":  {},
+				"route4|10.10.0.0/16":   {},
+				"route5|192.168.1.0/24": {},
+			}
+
+			filtered := rs.FilterSelected(routes)
+			expectedLen := len(tt.wantNewSelected)
+			assert.Equal(t, expectedLen, len(filtered),
+				"FilterSelected returned wrong number of routes, got %d want %d", len(filtered), expectedLen)
+		})
+	}
+}

client/internal/statemanager/manager.go

@@ -16,14 +16,39 @@ import (
 	"golang.org/x/exp/maps"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	errStateNotRegistered = "state %s not registered"
+	errLoadStateFile      = "load state file: %w"
 )
 
 // State interface defines the methods that all state types must implement
 type State interface {
 	Name() string
+}
+
+// CleanableState interface extends State with cleanup capability
+type CleanableState interface {
+	State
 	Cleanup() error
 }
 
+// RawState wraps raw JSON data for unregistered states
+type RawState struct {
+	data json.RawMessage
+}
+
+func (r *RawState) Name() string {
+	return "" // This is a placeholder implementation
+}
+
+// MarshalJSON implements json.Marshaler to preserve the original JSON
+func (r *RawState) MarshalJSON() ([]byte, error) {
+	return r.data, nil
+}
+
 // Manager handles the persistence and management of various states
 type Manager struct {
 	mu     sync.Mutex
@@ -73,15 +98,15 @@ func (m *Manager) Stop(ctx context.Context) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if m.cancel != nil {
-		m.cancel()
+	if m.cancel == nil {
+		return nil
+	}
+	m.cancel()
 
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-m.done:
-			return nil
-		}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-m.done:
 	}
 
 	return nil
@@ -139,7 +164,7 @@ func (m *Manager) setState(name string, state State) error {
 	defer m.mu.Unlock()
 
 	if _, exists := m.states[name]; !exists {
-		return fmt.Errorf("state %s not registered", name)
+		return fmt.Errorf(errStateNotRegistered, name)
 	}
 
 	m.states[name] = state
@@ -148,6 +173,63 @@ func (m *Manager) setState(name string, state State) error {
 	return nil
 }
 
+// DeleteStateByName handles deletion of states without cleanup.
+// It doesn't require the state to be registered.
+func (m *Manager) DeleteStateByName(stateName string) error {
+	if m == nil {
+		return nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	rawStates, err := m.loadStateFile(false)
+	if err != nil {
+		return fmt.Errorf(errLoadStateFile, err)
+	}
+	if rawStates == nil {
+		return nil
+	}
+
+	if _, exists := rawStates[stateName]; !exists {
+		return fmt.Errorf("state %s not found", stateName)
+	}
+
+	// Mark state as deleted by setting it to nil and marking it dirty
+	m.states[stateName] = nil
+	m.dirty[stateName] = struct{}{}
+
+	return nil
+}
+
+// DeleteAllStates removes all states.
+func (m *Manager) DeleteAllStates() (int, error) {
+	if m == nil {
+		return 0, nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	rawStates, err := m.loadStateFile(false)
+	if err != nil {
+		return 0, fmt.Errorf(errLoadStateFile, err)
+	}
+	if rawStates == nil {
+		return 0, nil
+	}
+
+	count := len(rawStates)
+
+	// Mark all states as deleted and dirty
+	for name := range rawStates {
+		m.states[name] = nil
+		m.dirty[name] = struct{}{}
+	}
+
+	return count, nil
+}
+
 func (m *Manager) periodicStateSave(ctx context.Context) {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
@@ -178,25 +260,18 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		return nil
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	bs, err := marshalWithPanicRecovery(m.states)
+	if err != nil {
+		return fmt.Errorf("marshal states: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
 
 	done := make(chan error, 1)
-
+	start := time.Now()
 	go func() {
-		data, err := json.MarshalIndent(m.states, "", "  ")
-		if err != nil {
-			done <- fmt.Errorf("marshal states: %w", err)
-			return
-		}
-
-		// nolint:gosec
-		if err := os.WriteFile(m.filePath, data, 0640); err != nil {
-			done <- fmt.Errorf("write state file: %w", err)
-			return
-		}
-
-		done <- nil
+		done <- util.WriteBytesWithRestrictedPermission(ctx, m.filePath, bs)
 	}()
 
 	select {
@@ -208,63 +283,175 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		}
 	}
 
-	log.Debugf("persisted shutdown states: %v", maps.Keys(m.dirty))
+	log.Debugf("persisted states: %v, took %v", maps.Keys(m.dirty), time.Since(start))
 
 	clear(m.dirty)
 
 	return nil
 }
 
-// loadState loads the existing state from the state file
-func (m *Manager) loadState() error {
+// loadStateFile reads and unmarshals the state file into a map of raw JSON messages
+func (m *Manager) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage, error) {
 	data, err := os.ReadFile(m.filePath)
 	if err != nil {
 		if errors.Is(err, fs.ErrNotExist) {
 			log.Debug("state file does not exist")
-			return nil
+			return nil, nil // nolint:nilnil
 		}
-		return fmt.Errorf("read state file: %w", err)
+		return nil, fmt.Errorf("read state file: %w", err)
 	}
 
 	var rawStates map[string]json.RawMessage
 	if err := json.Unmarshal(data, &rawStates); err != nil {
-		log.Warn("State file appears to be corrupted, attempting to delete it")
-		if err := os.Remove(m.filePath); err != nil {
-			log.Errorf("Failed to delete corrupted state file: %v", err)
-		} else {
-			log.Info("State file deleted")
+		if deleteCorrupt {
+			log.Warn("State file appears to be corrupted, attempting to delete it", err)
+			if err := os.Remove(m.filePath); err != nil {
+				log.Errorf("Failed to delete corrupted state file: %v", err)
+			} else {
+				log.Info("State file deleted")
+			}
 		}
-		return fmt.Errorf("unmarshal states: %w", err)
+		return nil, fmt.Errorf("unmarshal states: %w", err)
 	}
 
-	var merr *multierror.Error
+	return rawStates, nil
+}
 
-	for name, rawState := range rawStates {
-		stateType, ok := m.stateTypes[name]
-		if !ok {
-			merr = multierror.Append(merr, fmt.Errorf("unknown state type: %s", name))
-			continue
-		}
+// loadSingleRawState unmarshals a raw state into a concrete state object
+func (m *Manager) loadSingleRawState(name string, rawState json.RawMessage) (State, error) {
+	stateType, ok := m.stateTypes[name]
+	if !ok {
+		return nil, fmt.Errorf(errStateNotRegistered, name)
+	}
 
-		if string(rawState) == "null" {
-			continue
-		}
+	if string(rawState) == "null" {
+		return nil, nil //nolint:nilnil
+	}
 
-		statePtr := reflect.New(stateType).Interface().(State)
-		if err := json.Unmarshal(rawState, statePtr); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("unmarshal state %s: %w", name, err))
-			continue
-		}
+	statePtr := reflect.New(stateType).Interface().(State)
+	if err := json.Unmarshal(rawState, statePtr); err != nil {
+		return nil, fmt.Errorf("unmarshal state %s: %w", name, err)
+	}
 
-		m.states[name] = statePtr
+	return statePtr, nil
+}
+
+// LoadState loads a specific state from the state file
+func (m *Manager) LoadState(state State) error {
+	if m == nil {
+		return nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	rawStates, err := m.loadStateFile(false)
+	if err != nil {
+		return err
+	}
+	if rawStates == nil {
+		return nil
+	}
+
+	name := state.Name()
+	rawState, exists := rawStates[name]
+	if !exists {
+		return nil
+	}
+
+	loadedState, err := m.loadSingleRawState(name, rawState)
+	if err != nil {
+		return err
+	}
+
+	m.states[name] = loadedState
+	if loadedState != nil {
 		log.Debugf("loaded state: %s", name)
 	}
 
-	return nberrors.FormatErrorOrNil(merr)
+	return nil
 }
 
-// PerformCleanup retrieves all states from the state file for the registered states and calls Cleanup on them.
-// If the cleanup is successful, the state is marked for deletion.
+// cleanupSingleState handles the cleanup of a specific state and returns any error.
+// The caller must hold the mutex.
+func (m *Manager) cleanupSingleState(name string, rawState json.RawMessage) error {
+	// For unregistered states, preserve the raw JSON
+	if _, registered := m.stateTypes[name]; !registered {
+		m.states[name] = &RawState{data: rawState}
+		return nil
+	}
+
+	// Load the state
+	loadedState, err := m.loadSingleRawState(name, rawState)
+	if err != nil {
+		return err
+	}
+
+	if loadedState == nil {
+		return nil
+	}
+
+	// Check if state supports cleanup
+	cleanableState, isCleanable := loadedState.(CleanableState)
+	if !isCleanable {
+		// If it doesn't support cleanup, keep it as-is
+		m.states[name] = loadedState
+		return nil
+	}
+
+	// Perform cleanup
+	log.Infof("cleaning up state %s", name)
+	if err := cleanableState.Cleanup(); err != nil {
+		// On cleanup error, preserve the state
+		m.states[name] = loadedState
+		return fmt.Errorf("cleanup state: %w", err)
+	}
+
+	// Successfully cleaned up - mark for deletion
+	m.states[name] = nil
+	m.dirty[name] = struct{}{}
+	return nil
+}
+
+// CleanupStateByName loads and cleans up a specific state by name if it implements CleanableState.
+// Returns an error if the state doesn't exist, isn't registered, or cleanup fails.
+func (m *Manager) CleanupStateByName(name string) error {
+	if m == nil {
+		return nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// Check if state is registered
+	if _, registered := m.stateTypes[name]; !registered {
+		return fmt.Errorf(errStateNotRegistered, name)
+	}
+
+	// Load raw states from file
+	rawStates, err := m.loadStateFile(false)
+	if err != nil {
+		return err
+	}
+	if rawStates == nil {
+		return nil
+	}
+
+	// Check if state exists in file
+	rawState, exists := rawStates[name]
+	if !exists {
+		return nil
+	}
+
+	if err := m.cleanupSingleState(name, rawState); err != nil {
+		return fmt.Errorf("%s: %w", name, err)
+	}
+
+	return nil
+}
+
+// PerformCleanup retrieves all states from the state file and calls Cleanup on registered states that support it.
+// Unregistered states are preserved in their original state.
 func (m *Manager) PerformCleanup() error {
 	if m == nil {
 		return nil
@@ -273,26 +460,63 @@ func (m *Manager) PerformCleanup() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if err := m.loadState(); err != nil {
-		log.Warnf("Failed to load state during cleanup: %v", err)
+	// Load raw states from file
+	rawStates, err := m.loadStateFile(true)
+	if err != nil {
+		return fmt.Errorf(errLoadStateFile, err)
+	}
+	if rawStates == nil {
+		return nil
 	}
 
 	var merr *multierror.Error
-	for name, state := range m.states {
-		if state == nil {
-			// If no state was found in the state file, we don't mark the state dirty nor return an error
-			continue
-		}
 
-		log.Infof("client was not shut down properly, cleaning up %s", name)
-		if err := state.Cleanup(); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("cleanup state for %s: %w", name, err))
-		} else {
-			// mark for deletion on cleanup success
-			m.states[name] = nil
-			m.dirty[name] = struct{}{}
+	// Process each state in the file
+	for name, rawState := range rawStates {
+		if err := m.cleanupSingleState(name, rawState); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("%s: %w", name, err))
 		}
 	}
 
 	return nberrors.FormatErrorOrNil(merr)
 }
+
+// GetSavedStateNames returns all state names that are currently saved in the state file.
+func (m *Manager) GetSavedStateNames() ([]string, error) {
+	if m == nil {
+		return nil, nil
+	}
+
+	rawStates, err := m.loadStateFile(false)
+	if err != nil {
+		return nil, fmt.Errorf(errLoadStateFile, err)
+	}
+	if rawStates == nil {
+		return nil, nil
+	}
+
+	var states []string
+	for name, state := range rawStates {
+		if len(state) != 0 && string(state) != "null" {
+			states = append(states, name)
+		}
+	}
+
+	return states, nil
+}
+
+func marshalWithPanicRecovery(v any) ([]byte, error) {
+	var bs []byte
+	var err error
+
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				err = fmt.Errorf("panic during marshal: %v", r)
+			}
+		}()
+		bs, err = json.Marshal(v)
+	}()
+
+	return bs, err
+}

client/internal/statemanager/path.go

@@ -4,32 +4,20 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
-
-	log "github.com/sirupsen/logrus"
 )
 
 // GetDefaultStatePath returns the path to the state file based on the operating system
-// It returns an empty string if the path cannot be determined. It also creates the directory if it does not exist.
+// It returns an empty string if the path cannot be determined.
 func GetDefaultStatePath() string {
-	var path string
-
 	switch runtime.GOOS {
 	case "windows":
-		path = filepath.Join(os.Getenv("PROGRAMDATA"), "Netbird", "state.json")
+		return filepath.Join(os.Getenv("PROGRAMDATA"), "Netbird", "state.json")
 	case "darwin", "linux":
-		path = "/var/lib/netbird/state.json"
+		return "/var/lib/netbird/state.json"
 	case "freebsd", "openbsd", "netbsd", "dragonfly":
-		path = "/var/db/netbird/state.json"
-	// ios/android don't need state
-	default:
-		return ""
+		return "/var/db/netbird/state.json"
 	}
 
-	dir := filepath.Dir(path)
-	if err := os.MkdirAll(dir, 0755); err != nil {
-		log.Errorf("Error creating directory %s: %v. Continuing without state support.", dir, err)
-		return ""
-	}
+	return ""
 
-	return path
 }

client/ios/NetBirdSDK/client.go

@@ -59,6 +59,7 @@ func init() {
 // Client struct manage the life circle of background service
 type Client struct {
 	cfgFile               string
+	stateFile             string
 	recorder              *peer.Status
 	ctxCancel             context.CancelFunc
 	ctxCancelLock         *sync.Mutex
@@ -73,9 +74,10 @@ type Client struct {
 }
 
 // NewClient instantiate a new Client
-func NewClient(cfgFile, deviceName string, osVersion string, osName string, networkChangeListener NetworkChangeListener, dnsManager DnsManager) *Client {
+func NewClient(cfgFile, stateFile, deviceName string, osVersion string, osName string, networkChangeListener NetworkChangeListener, dnsManager DnsManager) *Client {
 	return &Client{
 		cfgFile:               cfgFile,
+		stateFile:             stateFile,
 		deviceName:            deviceName,
 		osName:                osName,
 		osVersion:             osVersion,
@@ -91,7 +93,8 @@ func (c *Client) Run(fd int32, interfaceName string) error {
 	log.Infof("Starting NetBird client")
 	log.Debugf("Tunnel uses interface: %s", interfaceName)
 	cfg, err := internal.UpdateOrCreateConfig(internal.ConfigInput{
-		ConfigPath: c.cfgFile,
+		ConfigPath:    c.cfgFile,
+		StateFilePath: c.stateFile,
 	})
 	if err != nil {
 		return err
@@ -124,7 +127,7 @@ func (c *Client) Run(fd int32, interfaceName string) error {
 	cfg.WgIface = interfaceName
 
 	c.connectClient = internal.NewConnectClient(ctx, cfg, c.recorder)
-	return c.connectClient.RunOniOS(fd, c.networkChangeListener, c.dnsManager)
+	return c.connectClient.RunOniOS(fd, c.networkChangeListener, c.dnsManager, c.stateFile)
 }
 
 // Stop the internal client and free the resources
@@ -269,8 +272,8 @@ func (c *Client) GetRoutesSelectionDetails() (*RoutesSelectionDetails, error) {
 		return nil, fmt.Errorf("not connected")
 	}
 
-	routesMap := engine.GetClientRoutesWithNetID()
 	routeManager := engine.GetRouteManager()
+	routesMap := routeManager.GetClientRoutesWithNetID()
 	if routeManager == nil {
 		return nil, fmt.Errorf("could not get route manager")
 	}
@@ -314,7 +317,7 @@ func (c *Client) GetRoutesSelectionDetails() (*RoutesSelectionDetails, error) {
 
 }
 
-func prepareRouteSelectionDetails(routes []*selectRoute, resolvedDomains map[domain.Domain][]netip.Prefix) *RoutesSelectionDetails {
+func prepareRouteSelectionDetails(routes []*selectRoute, resolvedDomains map[domain.Domain]peer.ResolvedDomainInfo) *RoutesSelectionDetails {
 	var routeSelection []RoutesSelectionInfo
 	for _, r := range routes {
 		domainList := make([]DomainInfo, 0)
@@ -322,9 +325,10 @@ func prepareRouteSelectionDetails(routes []*selectRoute, resolvedDomains map[dom
 			domainResp := DomainInfo{
 				Domain: d.SafeString(),
 			}
-			if prefixes, exists := resolvedDomains[d]; exists {
+
+			if info, exists := resolvedDomains[d]; exists {
 				var ipStrings []string
-				for _, prefix := range prefixes {
+				for _, prefix := range info.Prefixes {
 					ipStrings = append(ipStrings, prefix.Addr().String())
 				}
 				domainResp.ResolvedIPs = strings.Join(ipStrings, ", ")
@@ -362,12 +366,12 @@ func (c *Client) SelectRoute(id string) error {
 	} else {
 		log.Debugf("select route with id: %s", id)
 		routes := toNetIDs([]string{id})
-		if err := routeSelector.SelectRoutes(routes, true, maps.Keys(engine.GetClientRoutesWithNetID())); err != nil {
+		if err := routeSelector.SelectRoutes(routes, true, maps.Keys(routeManager.GetClientRoutesWithNetID())); err != nil {
 			log.Debugf("error when selecting routes: %s", err)
 			return fmt.Errorf("select routes: %w", err)
 		}
 	}
-	routeManager.TriggerSelection(engine.GetClientRoutes())
+	routeManager.TriggerSelection(routeManager.GetClientRoutes())
 	return nil
 
 }
@@ -389,12 +393,12 @@ func (c *Client) DeselectRoute(id string) error {
 	} else {
 		log.Debugf("deselect route with id: %s", id)
 		routes := toNetIDs([]string{id})
-		if err := routeSelector.DeselectRoutes(routes, maps.Keys(engine.GetClientRoutesWithNetID())); err != nil {
+		if err := routeSelector.DeselectRoutes(routes, maps.Keys(routeManager.GetClientRoutesWithNetID())); err != nil {
 			log.Debugf("error when deselecting routes: %s", err)
 			return fmt.Errorf("deselect routes: %w", err)
 		}
 	}
-	routeManager.TriggerSelection(engine.GetClientRoutes())
+	routeManager.TriggerSelection(routeManager.GetClientRoutes())
 	return nil
 }
 

client/ios/NetBirdSDK/preferences.go

@@ -10,9 +10,10 @@ type Preferences struct {
 }
 
 // NewPreferences create new Preferences instance
-func NewPreferences(configPath string) *Preferences {
+func NewPreferences(configPath string, stateFilePath string) *Preferences {
 	ci := internal.ConfigInput{
-		ConfigPath: configPath,
+		ConfigPath:    configPath,
+		StateFilePath: stateFilePath,
 	}
 	return &Preferences{ci}
 }

client/ios/NetBirdSDK/preferences_test.go

@@ -9,7 +9,8 @@ import (
 
 func TestPreferences_DefaultValues(t *testing.T) {
 	cfgFile := filepath.Join(t.TempDir(), "netbird.json")
-	p := NewPreferences(cfgFile)
+	stateFile := filepath.Join(t.TempDir(), "state.json")
+	p := NewPreferences(cfgFile, stateFile)
 	defaultVar, err := p.GetAdminURL()
 	if err != nil {
 		t.Fatalf("failed to read default value: %s", err)
@@ -42,7 +43,8 @@ func TestPreferences_DefaultValues(t *testing.T) {
 func TestPreferences_ReadUncommitedValues(t *testing.T) {
 	exampleString := "exampleString"
 	cfgFile := filepath.Join(t.TempDir(), "netbird.json")
-	p := NewPreferences(cfgFile)
+	stateFile := filepath.Join(t.TempDir(), "state.json")
+	p := NewPreferences(cfgFile, stateFile)
 
 	p.SetAdminURL(exampleString)
 	resp, err := p.GetAdminURL()
@@ -79,7 +81,8 @@ func TestPreferences_Commit(t *testing.T) {
 	exampleURL := "https://myurl.com:443"
 	examplePresharedKey := "topsecret"
 	cfgFile := filepath.Join(t.TempDir(), "netbird.json")
-	p := NewPreferences(cfgFile)
+	stateFile := filepath.Join(t.TempDir(), "state.json")
+	p := NewPreferences(cfgFile, stateFile)
 
 	p.SetAdminURL(exampleURL)
 	p.SetManagementURL(exampleURL)
@@ -90,7 +93,7 @@ func TestPreferences_Commit(t *testing.T) {
 		t.Fatalf("failed to save changes: %s", err)
 	}
 
-	p = NewPreferences(cfgFile)
+	p = NewPreferences(cfgFile, stateFile)
 	resp, err := p.GetAdminURL()
 	if err != nil {
 		t.Fatalf("failed to read admin url: %s", err)

client/proto/daemon.proto

@@ -28,14 +28,14 @@ service DaemonService {
   // GetConfig of the daemon.
   rpc GetConfig(GetConfigRequest) returns (GetConfigResponse) {}
 
-  // List available network routes
-  rpc ListRoutes(ListRoutesRequest) returns (ListRoutesResponse) {}
+  // List available networks
+  rpc ListNetworks(ListNetworksRequest) returns (ListNetworksResponse) {}
 
   // Select specific routes
-  rpc SelectRoutes(SelectRoutesRequest) returns (SelectRoutesResponse) {}
+  rpc SelectNetworks(SelectNetworksRequest) returns (SelectNetworksResponse) {}
 
   // Deselect specific routes
-  rpc DeselectRoutes(SelectRoutesRequest) returns (SelectRoutesResponse) {}
+  rpc DeselectNetworks(SelectNetworksRequest) returns (SelectNetworksResponse) {}
 
   // DebugBundle creates a debug bundle
   rpc DebugBundle(DebugBundleRequest) returns (DebugBundleResponse) {}
@@ -45,7 +45,20 @@ service DaemonService {
 
   // SetLogLevel sets the log level of the daemon
   rpc SetLogLevel(SetLogLevelRequest) returns (SetLogLevelResponse) {}
-};
+
+  // List all states
+  rpc ListStates(ListStatesRequest) returns (ListStatesResponse) {}
+
+  // Clean specific state or all states
+  rpc CleanState(CleanStateRequest) returns (CleanStateResponse) {}
+
+  // Delete specific state or all states
+  rpc DeleteState(DeleteStateRequest) returns (DeleteStateResponse) {}
+
+  // SetNetworkMapPersistence enables or disables network map persistence
+  rpc SetNetworkMapPersistence(SetNetworkMapPersistenceRequest) returns (SetNetworkMapPersistenceResponse) {}
+}
+
 
 message LoginRequest {
   // setupKey wiretrustee setup key.
@@ -177,7 +190,7 @@ message PeerState {
   int64 bytesRx = 13;
   int64 bytesTx = 14;
   bool rosenpassEnabled = 15;
-  repeated string routes = 16;
+  repeated string networks = 16;
   google.protobuf.Duration latency = 17;
   string relayAddress = 18;
 }
@@ -190,7 +203,7 @@ message LocalPeerState {
   string fqdn = 4;
   bool rosenpassEnabled = 5;
   bool rosenpassPermissive = 6;
-  repeated string routes = 7;
+  repeated string networks = 7;
 }
 
 // SignalState contains the latest state of a signal connection
@@ -231,20 +244,20 @@ message FullStatus {
   repeated NSGroupState dns_servers = 6;
 }
 
-message ListRoutesRequest {
+message ListNetworksRequest {
 }
 
-message ListRoutesResponse {
-  repeated Route routes = 1;
+message ListNetworksResponse {
+  repeated Network routes = 1;
 }
 
-message SelectRoutesRequest {
-  repeated string routeIDs = 1;
+message SelectNetworksRequest {
+  repeated string networkIDs = 1;
   bool append = 2;
   bool all = 3;
 }
 
-message SelectRoutesResponse {
+message SelectNetworksResponse {
 }
 
 message IPList {
@@ -252,9 +265,9 @@ message IPList {
 }
 
 
-message Route {
+message Network {
   string ID = 1;
-  string network = 2;
+  string range = 2;
   bool selected = 3;
   repeated string domains = 4;
   map<string, IPList> resolvedIPs = 5;
@@ -293,4 +306,46 @@ message SetLogLevelRequest {
 }
 
 message SetLogLevelResponse {
-}
+}
+
+// State represents a daemon state entry
+message State {
+  string name = 1;
+}
+
+// ListStatesRequest is empty as it requires no parameters
+message ListStatesRequest {}
+
+// ListStatesResponse contains a list of states
+message ListStatesResponse {
+  repeated State states = 1;
+}
+
+// CleanStateRequest for cleaning states
+message CleanStateRequest {
+  string state_name = 1;
+  bool all = 2;
+}
+
+// CleanStateResponse contains the result of the clean operation
+message CleanStateResponse {
+  int32 cleaned_states = 1;
+}
+
+// DeleteStateRequest for deleting states
+message DeleteStateRequest {
+  string state_name = 1;
+  bool all = 2;
+}
+
+// DeleteStateResponse contains the result of the delete operation
+message DeleteStateResponse {
+  int32 deleted_states = 1;
+}
+
+
+message SetNetworkMapPersistenceRequest {
+  bool enabled = 1;
+}
+
+message SetNetworkMapPersistenceResponse {}

client/proto/daemon_grpc.pb.go

@@ -31,18 +31,26 @@ type DaemonServiceClient interface {
 	Down(ctx context.Context, in *DownRequest, opts ...grpc.CallOption) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(ctx context.Context, in *GetConfigRequest, opts ...grpc.CallOption) (*GetConfigResponse, error)
-	// List available network routes
-	ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error)
+	// List available networks
+	ListNetworks(ctx context.Context, in *ListNetworksRequest, opts ...grpc.CallOption) (*ListNetworksResponse, error)
 	// Select specific routes
-	SelectRoutes(ctx context.Context, in *SelectRoutesRequest, opts ...grpc.CallOption) (*SelectRoutesResponse, error)
+	SelectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error)
 	// Deselect specific routes
-	DeselectRoutes(ctx context.Context, in *SelectRoutesRequest, opts ...grpc.CallOption) (*SelectRoutesResponse, error)
+	DeselectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error)
 	// DebugBundle creates a debug bundle
 	DebugBundle(ctx context.Context, in *DebugBundleRequest, opts ...grpc.CallOption) (*DebugBundleResponse, error)
 	// GetLogLevel gets the log level of the daemon
 	GetLogLevel(ctx context.Context, in *GetLogLevelRequest, opts ...grpc.CallOption) (*GetLogLevelResponse, error)
 	// SetLogLevel sets the log level of the daemon
 	SetLogLevel(ctx context.Context, in *SetLogLevelRequest, opts ...grpc.CallOption) (*SetLogLevelResponse, error)
+	// List all states
+	ListStates(ctx context.Context, in *ListStatesRequest, opts ...grpc.CallOption) (*ListStatesResponse, error)
+	// Clean specific state or all states
+	CleanState(ctx context.Context, in *CleanStateRequest, opts ...grpc.CallOption) (*CleanStateResponse, error)
+	// Delete specific state or all states
+	DeleteState(ctx context.Context, in *DeleteStateRequest, opts ...grpc.CallOption) (*DeleteStateResponse, error)
+	// SetNetworkMapPersistence enables or disables network map persistence
+	SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error)
 }
 
 type daemonServiceClient struct {
@@ -107,27 +115,27 @@ func (c *daemonServiceClient) GetConfig(ctx context.Context, in *GetConfigReques
 	return out, nil
 }
 
-func (c *daemonServiceClient) ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) {
-	out := new(ListRoutesResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListRoutes", in, out, opts...)
+func (c *daemonServiceClient) ListNetworks(ctx context.Context, in *ListNetworksRequest, opts ...grpc.CallOption) (*ListNetworksResponse, error) {
+	out := new(ListNetworksResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
 	return out, nil
 }
 
-func (c *daemonServiceClient) SelectRoutes(ctx context.Context, in *SelectRoutesRequest, opts ...grpc.CallOption) (*SelectRoutesResponse, error) {
-	out := new(SelectRoutesResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SelectRoutes", in, out, opts...)
+func (c *daemonServiceClient) SelectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error) {
+	out := new(SelectNetworksResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SelectNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
 	return out, nil
 }
 
-func (c *daemonServiceClient) DeselectRoutes(ctx context.Context, in *SelectRoutesRequest, opts ...grpc.CallOption) (*SelectRoutesResponse, error) {
-	out := new(SelectRoutesResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeselectRoutes", in, out, opts...)
+func (c *daemonServiceClient) DeselectNetworks(ctx context.Context, in *SelectNetworksRequest, opts ...grpc.CallOption) (*SelectNetworksResponse, error) {
+	out := new(SelectNetworksResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeselectNetworks", in, out, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -161,6 +169,42 @@ func (c *daemonServiceClient) SetLogLevel(ctx context.Context, in *SetLogLevelRe
 	return out, nil
 }
 
+func (c *daemonServiceClient) ListStates(ctx context.Context, in *ListStatesRequest, opts ...grpc.CallOption) (*ListStatesResponse, error) {
+	out := new(ListStatesResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListStates", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *daemonServiceClient) CleanState(ctx context.Context, in *CleanStateRequest, opts ...grpc.CallOption) (*CleanStateResponse, error) {
+	out := new(CleanStateResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/CleanState", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *daemonServiceClient) DeleteState(ctx context.Context, in *DeleteStateRequest, opts ...grpc.CallOption) (*DeleteStateResponse, error) {
+	out := new(DeleteStateResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeleteState", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *daemonServiceClient) SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error) {
+	out := new(SetNetworkMapPersistenceResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SetNetworkMapPersistence", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 // DaemonServiceServer is the server API for DaemonService service.
 // All implementations must embed UnimplementedDaemonServiceServer
 // for forward compatibility
@@ -178,18 +222,26 @@ type DaemonServiceServer interface {
 	Down(context.Context, *DownRequest) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(context.Context, *GetConfigRequest) (*GetConfigResponse, error)
-	// List available network routes
-	ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error)
+	// List available networks
+	ListNetworks(context.Context, *ListNetworksRequest) (*ListNetworksResponse, error)
 	// Select specific routes
-	SelectRoutes(context.Context, *SelectRoutesRequest) (*SelectRoutesResponse, error)
+	SelectNetworks(context.Context, *SelectNetworksRequest) (*SelectNetworksResponse, error)
 	// Deselect specific routes
-	DeselectRoutes(context.Context, *SelectRoutesRequest) (*SelectRoutesResponse, error)
+	DeselectNetworks(context.Context, *SelectNetworksRequest) (*SelectNetworksResponse, error)
 	// DebugBundle creates a debug bundle
 	DebugBundle(context.Context, *DebugBundleRequest) (*DebugBundleResponse, error)
 	// GetLogLevel gets the log level of the daemon
 	GetLogLevel(context.Context, *GetLogLevelRequest) (*GetLogLevelResponse, error)
 	// SetLogLevel sets the log level of the daemon
 	SetLogLevel(context.Context, *SetLogLevelRequest) (*SetLogLevelResponse, error)
+	// List all states
+	ListStates(context.Context, *ListStatesRequest) (*ListStatesResponse, error)
+	// Clean specific state or all states
+	CleanState(context.Context, *CleanStateRequest) (*CleanStateResponse, error)
+	// Delete specific state or all states
+	DeleteState(context.Context, *DeleteStateRequest) (*DeleteStateResponse, error)
+	// SetNetworkMapPersistence enables or disables network map persistence
+	SetNetworkMapPersistence(context.Context, *SetNetworkMapPersistenceRequest) (*SetNetworkMapPersistenceResponse, error)
 	mustEmbedUnimplementedDaemonServiceServer()
 }
 
@@ -215,14 +267,14 @@ func (UnimplementedDaemonServiceServer) Down(context.Context, *DownRequest) (*Do
 func (UnimplementedDaemonServiceServer) GetConfig(context.Context, *GetConfigRequest) (*GetConfigResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method GetConfig not implemented")
 }
-func (UnimplementedDaemonServiceServer) ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method ListRoutes not implemented")
+func (UnimplementedDaemonServiceServer) ListNetworks(context.Context, *ListNetworksRequest) (*ListNetworksResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method ListNetworks not implemented")
 }
-func (UnimplementedDaemonServiceServer) SelectRoutes(context.Context, *SelectRoutesRequest) (*SelectRoutesResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method SelectRoutes not implemented")
+func (UnimplementedDaemonServiceServer) SelectNetworks(context.Context, *SelectNetworksRequest) (*SelectNetworksResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method SelectNetworks not implemented")
 }
-func (UnimplementedDaemonServiceServer) DeselectRoutes(context.Context, *SelectRoutesRequest) (*SelectRoutesResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method DeselectRoutes not implemented")
+func (UnimplementedDaemonServiceServer) DeselectNetworks(context.Context, *SelectNetworksRequest) (*SelectNetworksResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method DeselectNetworks not implemented")
 }
 func (UnimplementedDaemonServiceServer) DebugBundle(context.Context, *DebugBundleRequest) (*DebugBundleResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method DebugBundle not implemented")
@@ -233,6 +285,18 @@ func (UnimplementedDaemonServiceServer) GetLogLevel(context.Context, *GetLogLeve
 func (UnimplementedDaemonServiceServer) SetLogLevel(context.Context, *SetLogLevelRequest) (*SetLogLevelResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method SetLogLevel not implemented")
 }
+func (UnimplementedDaemonServiceServer) ListStates(context.Context, *ListStatesRequest) (*ListStatesResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method ListStates not implemented")
+}
+func (UnimplementedDaemonServiceServer) CleanState(context.Context, *CleanStateRequest) (*CleanStateResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method CleanState not implemented")
+}
+func (UnimplementedDaemonServiceServer) DeleteState(context.Context, *DeleteStateRequest) (*DeleteStateResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method DeleteState not implemented")
+}
+func (UnimplementedDaemonServiceServer) SetNetworkMapPersistence(context.Context, *SetNetworkMapPersistenceRequest) (*SetNetworkMapPersistenceResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method SetNetworkMapPersistence not implemented")
+}
 func (UnimplementedDaemonServiceServer) mustEmbedUnimplementedDaemonServiceServer() {}
 
 // UnsafeDaemonServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -354,56 +418,56 @@ func _DaemonService_GetConfig_Handler(srv interface{}, ctx context.Context, dec
 	return interceptor(ctx, in, info, handler)
 }
 
-func _DaemonService_ListRoutes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(ListRoutesRequest)
+func _DaemonService_ListNetworks_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(ListNetworksRequest)
 	if err := dec(in); err != nil {
 		return nil, err
 	}
 	if interceptor == nil {
-		return srv.(DaemonServiceServer).ListRoutes(ctx, in)
+		return srv.(DaemonServiceServer).ListNetworks(ctx, in)
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: "/daemon.DaemonService/ListRoutes",
+		FullMethod: "/daemon.DaemonService/ListNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).ListRoutes(ctx, req.(*ListRoutesRequest))
+		return srv.(DaemonServiceServer).ListNetworks(ctx, req.(*ListNetworksRequest))
 	}
 	return interceptor(ctx, in, info, handler)
 }
 
-func _DaemonService_SelectRoutes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(SelectRoutesRequest)
+func _DaemonService_SelectNetworks_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(SelectNetworksRequest)
 	if err := dec(in); err != nil {
 		return nil, err
 	}
 	if interceptor == nil {
-		return srv.(DaemonServiceServer).SelectRoutes(ctx, in)
+		return srv.(DaemonServiceServer).SelectNetworks(ctx, in)
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: "/daemon.DaemonService/SelectRoutes",
+		FullMethod: "/daemon.DaemonService/SelectNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).SelectRoutes(ctx, req.(*SelectRoutesRequest))
+		return srv.(DaemonServiceServer).SelectNetworks(ctx, req.(*SelectNetworksRequest))
 	}
 	return interceptor(ctx, in, info, handler)
 }
 
-func _DaemonService_DeselectRoutes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(SelectRoutesRequest)
+func _DaemonService_DeselectNetworks_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(SelectNetworksRequest)
 	if err := dec(in); err != nil {
 		return nil, err
 	}
 	if interceptor == nil {
-		return srv.(DaemonServiceServer).DeselectRoutes(ctx, in)
+		return srv.(DaemonServiceServer).DeselectNetworks(ctx, in)
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: "/daemon.DaemonService/DeselectRoutes",
+		FullMethod: "/daemon.DaemonService/DeselectNetworks",
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).DeselectRoutes(ctx, req.(*SelectRoutesRequest))
+		return srv.(DaemonServiceServer).DeselectNetworks(ctx, req.(*SelectNetworksRequest))
 	}
 	return interceptor(ctx, in, info, handler)
 }
@@ -462,6 +526,78 @@ func _DaemonService_SetLogLevel_Handler(srv interface{}, ctx context.Context, de
 	return interceptor(ctx, in, info, handler)
 }
 
+func _DaemonService_ListStates_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(ListStatesRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DaemonServiceServer).ListStates(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/daemon.DaemonService/ListStates",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DaemonServiceServer).ListStates(ctx, req.(*ListStatesRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _DaemonService_CleanState_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(CleanStateRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DaemonServiceServer).CleanState(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/daemon.DaemonService/CleanState",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DaemonServiceServer).CleanState(ctx, req.(*CleanStateRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _DaemonService_DeleteState_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(DeleteStateRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DaemonServiceServer).DeleteState(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/daemon.DaemonService/DeleteState",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DaemonServiceServer).DeleteState(ctx, req.(*DeleteStateRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _DaemonService_SetNetworkMapPersistence_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(SetNetworkMapPersistenceRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DaemonServiceServer).SetNetworkMapPersistence(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/daemon.DaemonService/SetNetworkMapPersistence",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DaemonServiceServer).SetNetworkMapPersistence(ctx, req.(*SetNetworkMapPersistenceRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 // DaemonService_ServiceDesc is the grpc.ServiceDesc for DaemonService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -494,16 +630,16 @@ var DaemonService_ServiceDesc = grpc.ServiceDesc{
 			Handler:    _DaemonService_GetConfig_Handler,
 		},
 		{
-			MethodName: "ListRoutes",
-			Handler:    _DaemonService_ListRoutes_Handler,
+			MethodName: "ListNetworks",
+			Handler:    _DaemonService_ListNetworks_Handler,
 		},
 		{
-			MethodName: "SelectRoutes",
-			Handler:    _DaemonService_SelectRoutes_Handler,
+			MethodName: "SelectNetworks",
+			Handler:    _DaemonService_SelectNetworks_Handler,
 		},
 		{
-			MethodName: "DeselectRoutes",
-			Handler:    _DaemonService_DeselectRoutes_Handler,
+			MethodName: "DeselectNetworks",
+			Handler:    _DaemonService_DeselectNetworks_Handler,
 		},
 		{
 			MethodName: "DebugBundle",
@@ -517,6 +653,22 @@ var DaemonService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "SetLogLevel",
 			Handler:    _DaemonService_SetLogLevel_Handler,
 		},
+		{
+			MethodName: "ListStates",
+			Handler:    _DaemonService_ListStates_Handler,
+		},
+		{
+			MethodName: "CleanState",
+			Handler:    _DaemonService_CleanState_Handler,
+		},
+		{
+			MethodName: "DeleteState",
+			Handler:    _DaemonService_DeleteState_Handler,
+		},
+		{
+			MethodName: "SetNetworkMapPersistence",
+			Handler:    _DaemonService_SetNetworkMapPersistence_Handler,
+		},
 	},
 	Streams:  []grpc.StreamDesc{},
 	Metadata: "daemon.proto",