Merge branch 'peers-get-account-refactoring' into refactor/optimize-peer-expiration

# Conflicts:
#	management/server/peer.go

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

@@ -24,7 +24,7 @@ jobs:
           copyback: false
           release: "14.1"
           prepare: |
-            pkg install -y go pkgconf xorg
+            pkg install -y go
 
           # -x - to print all executed commands
           # -e - to faile on first error
@@ -33,7 +33,7 @@ jobs:
             time go build -o netbird client/main.go
             # check all component except management, since we do not support management server on freebsd
             time go test -timeout 1m -failfast ./base62/...
-            # NOTE: without -p1 `client/internal/dns` will fail because of `listen udp4 :33100: bind: address already in use`
+            # NOTE: without -p1 `client/internal/dns` will fail becasue of `listen udp4 :33100: bind: address already in use`
             time go test -timeout 8m -failfast -p 1 ./client/...
             time go test -timeout 1m -failfast ./dns/...
             time go test -timeout 1m -failfast ./encryption/...

.github/workflows/release.yml

@@ -9,10 +9,10 @@ on:
   pull_request:
 
 env:
-  SIGN_PIPE_VER: "v0.0.18"
+  SIGN_PIPE_VER: "v0.0.17"
   GORELEASER_VER: "v2.3.2"
   PRODUCT_NAME: "NetBird"
-  COPYRIGHT: "NetBird GmbH"
+  COPYRIGHT: "Wiretrustee UG (haftungsbeschreankt)"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || github.actor_id }}

AUTHORS

@@ -1,3 +1,3 @@
 Mikhail Bragin (https://github.com/braginini)
 Maycon Santos (https://github.com/mlsmaycon)
-NetBird GmbH
+Wiretrustee UG (haftungsbeschränkt)

CONTRIBUTOR_LICENSE_AGREEMENT.md

@@ -3,10 +3,10 @@
 We are incredibly thankful for the contributions we receive from the community.
 We require our external contributors to sign a Contributor License Agreement ("CLA") in
 order to ensure that our projects remain licensed under Free and Open Source licenses such
-as BSD-3 while allowing NetBird to build a sustainable business.
+as BSD-3 while allowing Wiretrustee to build a sustainable business.
 
-NetBird is committed to having a true Open Source Software ("OSS") license for
-our software. A CLA enables NetBird to safely commercialize our products
+Wiretrustee is committed to having a true Open Source Software ("OSS") license for
+our software. A CLA enables Wiretrustee to safely commercialize our products
 while keeping a standard OSS license with all the rights that license grants to users: the
 ability to use the project in their own projects or businesses, to republish modified
 source, or to completely fork the project.
@@ -20,11 +20,11 @@ This is a human-readable summary of (and not a substitute for) the full agreemen
 This highlights only some of key terms of the CLA. It has no legal value and you should
 carefully review all the terms of the actual CLA before agreeing.
 
-<li>Grant of copyright license. You give NetBird permission to use your copyrighted work
+<li>Grant of copyright license. You give Wiretrustee permission to use your copyrighted work
 in commercial products.
 </li>
 
-<li>Grant of patent license. If your contributed work uses a patent, you give NetBird a
+<li>Grant of patent license. If your contributed work uses a patent, you give Wiretrustee a
 license to use that patent including within commercial products. You also agree that you
 have permission to grant this license.
 </li>
@@ -45,7 +45,7 @@ more.
 # Why require a CLA?
 
 Agreeing to a CLA explicitly states that you are entitled to provide a contribution, that you cannot withdraw permission
-to use your contribution at a later date, and that NetBird has permission to use your contribution in our commercial
+to use your contribution at a later date, and that Wiretrustee has permission to use your contribution in our commercial
 products.
 
 This removes any ambiguities or uncertainties caused by not having a CLA and allows users and customers to confidently
@@ -65,25 +65,25 @@ Follow the steps given by the bot to sign the CLA. This will require you to log
 information from your account) and to fill in a few additional details such as your name and email address. We will only
 use this information for CLA tracking; none of your submitted information will be used for marketing purposes.
 
-You only have to sign the CLA once. Once you've signed the CLA, future contributions to any NetBird project will not
+You only have to sign the CLA once. Once you've signed the CLA, future contributions to any Wiretrustee project will not
 require you to sign again.
 
 # Legal Terms and Agreement
 
-In order to clarify the intellectual property license granted with Contributions from any person or entity, NetBird
-GmbH ("NetBird") must have a Contributor License Agreement ("CLA") on file that has been signed
+In order to clarify the intellectual property license granted with Contributions from any person or entity, Wiretrustee
+UG (haftungsbeschränkt) ("Wiretrustee") must have a Contributor License Agreement ("CLA") on file that has been signed
 by each Contributor, indicating agreement to the license terms below. This license does not change your rights to use
 your own Contributions for any other purpose.
 
 You accept and agree to the following terms and conditions for Your present and future Contributions submitted to
-NetBird. Except for the license granted herein to NetBird and recipients of software distributed by NetBird,
+Wiretrustee. Except for the license granted herein to Wiretrustee and recipients of software distributed by Wiretrustee,
 You reserve all right, title, and interest in and to Your Contributions.
 
 1. Definitions.
 
     ```
    "You" (or "Your") shall mean the copyright owner or legal entity authorized by the copyright owner 
-   that is making this Agreement with NetBird. For legal entities, the entity making a Contribution and all other 
+   that is making this Agreement with Wiretrustee. For legal entities, the entity making a Contribution and all other 
    entities that control, are controlled by, or are under common control with that entity are considered 
    to be a single Contributor. For the purposes of this definition, "control" means (i) the power, direct or indirect,
     to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty 
@@ -91,23 +91,23 @@ You reserve all right, title, and interest in and to Your Contributions.
    ```
    ```
     "Contribution" shall mean any original work of authorship, including any modifications or additions to 
-   an existing work, that is or previously has been intentionally submitted by You to NetBird for inclusion in, 
-   or documentation of, any of the products owned or managed by NetBird (the "Work"). 
+   an existing work, that is or previously has been intentionally submitted by You to Wiretrustee for inclusion in, 
+   or documentation of, any of the products owned or managed by Wiretrustee (the "Work"). 
    For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication 
-   sent to NetBird or its representatives, including but not limited to communication on electronic mailing lists, 
+   sent to Wiretrustee or its representatives, including but not limited to communication on electronic mailing lists, 
    source code control systems, and issue tracking systems that are managed by, or on behalf of, 
-   NetBird for the purpose of discussing and improving the Work, but excluding communication that is conspicuously 
+   Wiretrustee for the purpose of discussing and improving the Work, but excluding communication that is conspicuously 
    marked or otherwise designated in writing by You as "Not a Contribution."
    ```
 
-2. Grant of Copyright License. Subject to the terms and conditions of this Agreement, You hereby grant to NetBird
-   and to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge,
+2. Grant of Copyright License. Subject to the terms and conditions of this Agreement, You hereby grant to Wiretrustee
+   and to recipients of software distributed by Wiretrustee a perpetual, worldwide, non-exclusive, no-charge,
    royalty-free, irrevocable copyright license to reproduce, prepare derivative works of, publicly display, publicly
    perform, sublicense, and distribute Your Contributions and such derivative works.
 
 
-3. Grant of Patent License. Subject to the terms and conditions of this Agreement, You hereby grant to NetBird and
-   to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge, royalty-free,
+3. Grant of Patent License. Subject to the terms and conditions of this Agreement, You hereby grant to Wiretrustee and
+   to recipients of software distributed by Wiretrustee a perpetual, worldwide, non-exclusive, no-charge, royalty-free,
    irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import,
    and otherwise transfer the Work, where such license applies only to those patent claims licensable by You that are
    necessarily infringed by Your Contribution(s) alone or by combination of Your Contribution(s) with the Work to which
@@ -121,8 +121,8 @@ You reserve all right, title, and interest in and to Your Contributions.
    intellectual property that you create that includes your Contributions, you represent that you have received
    permission to make Contributions on behalf of that employer, that you will have received permission from your current
    and future employers for all future Contributions, that your applicable employer has waived such rights for all of
-   your current and future Contributions to NetBird, or that your employer has executed a separate Corporate CLA
-   with NetBird.
+   your current and future Contributions to Wiretrustee, or that your employer has executed a separate Corporate CLA
+   with Wiretrustee.
 
 
 5. You represent that each of Your Contributions is Your original creation (see section 7 for submissions on behalf of
@@ -138,11 +138,11 @@ You reserve all right, title, and interest in and to Your Contributions.
    MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE.
 
 
-7. Should You wish to submit work that is not Your original creation, You may submit it to NetBird separately from
+7. Should You wish to submit work that is not Your original creation, You may submit it to Wiretrustee separately from
    any Contribution, identifying the complete details of its source and of any license or other restriction (including,
    but not limited to, related patents, trademarks, and license agreements) of which you are personally aware, and
    conspicuously marking the work as "Submitted on behalf of a third-party: [named here]".
 
 
-8. You agree to notify NetBird of any facts or circumstances of which you become aware that would make these
-   representations inaccurate in any respect.
+8. You agree to notify Wiretrustee of any facts or circumstances of which you become aware that would make these
+   representations inaccurate in any respect.

LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2022 NetBird GmbH & AUTHORS
+Copyright (c) 2022 Wiretrustee UG (haftungsbeschränkt) & AUTHORS
 
 Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 
@@ -10,4 +10,4 @@ Redistribution and use in source and binary forms, with or without modification,
 
 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
 
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

client/Dockerfile-rootless

@@ -9,7 +9,6 @@ USER netbird:netbird
 
 ENV NB_FOREGROUND_MODE=true
 ENV NB_USE_NETSTACK_MODE=true
-ENV NB_ENABLE_NETSTACK_LOCAL_FORWARDING=true
 ENV NB_CONFIG=config.json
 ENV NB_DAEMON_ADDR=unix://netbird.sock
 ENV NB_DISABLE_DNS=true

client/cmd/trace.go

@@ -1,137 +0,0 @@
-package cmd
-
-import (
-	"fmt"
-	"math/rand"
-	"strings"
-
-	"github.com/spf13/cobra"
-	"google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/proto"
-)
-
-var traceCmd = &cobra.Command{
-	Use:   "trace <direction> <source-ip> <dest-ip>",
-	Short: "Trace a packet through the firewall",
-	Example: `
-  netbird debug trace in 192.168.1.10 10.10.0.2 -p tcp --sport 12345 --dport 443 --syn --ack
-  netbird debug trace out 10.10.0.1 8.8.8.8 -p udp  --dport 53
-  netbird debug trace in 10.10.0.2 10.10.0.1 -p icmp --type 8 --code 0
-  netbird debug trace in 100.64.1.1 self -p tcp --dport 80`,
-	Args: cobra.ExactArgs(3),
-	RunE: tracePacket,
-}
-
-func init() {
-	debugCmd.AddCommand(traceCmd)
-
-	traceCmd.Flags().StringP("protocol", "p", "tcp", "Protocol (tcp/udp/icmp)")
-	traceCmd.Flags().Uint16("sport", 0, "Source port")
-	traceCmd.Flags().Uint16("dport", 0, "Destination port")
-	traceCmd.Flags().Uint8("icmp-type", 0, "ICMP type")
-	traceCmd.Flags().Uint8("icmp-code", 0, "ICMP code")
-	traceCmd.Flags().Bool("syn", false, "TCP SYN flag")
-	traceCmd.Flags().Bool("ack", false, "TCP ACK flag")
-	traceCmd.Flags().Bool("fin", false, "TCP FIN flag")
-	traceCmd.Flags().Bool("rst", false, "TCP RST flag")
-	traceCmd.Flags().Bool("psh", false, "TCP PSH flag")
-	traceCmd.Flags().Bool("urg", false, "TCP URG flag")
-}
-
-func tracePacket(cmd *cobra.Command, args []string) error {
-	direction := strings.ToLower(args[0])
-	if direction != "in" && direction != "out" {
-		return fmt.Errorf("invalid direction: use 'in' or 'out'")
-	}
-
-	protocol := cmd.Flag("protocol").Value.String()
-	if protocol != "tcp" && protocol != "udp" && protocol != "icmp" {
-		return fmt.Errorf("invalid protocol: use tcp/udp/icmp")
-	}
-
-	sport, err := cmd.Flags().GetUint16("sport")
-	if err != nil {
-		return fmt.Errorf("invalid source port: %v", err)
-	}
-	dport, err := cmd.Flags().GetUint16("dport")
-	if err != nil {
-		return fmt.Errorf("invalid destination port: %v", err)
-	}
-
-	// For TCP/UDP, generate random ephemeral port (49152-65535) if not specified
-	if protocol != "icmp" {
-		if sport == 0 {
-			sport = uint16(rand.Intn(16383) + 49152)
-		}
-		if dport == 0 {
-			dport = uint16(rand.Intn(16383) + 49152)
-		}
-	}
-
-	var tcpFlags *proto.TCPFlags
-	if protocol == "tcp" {
-		syn, _ := cmd.Flags().GetBool("syn")
-		ack, _ := cmd.Flags().GetBool("ack")
-		fin, _ := cmd.Flags().GetBool("fin")
-		rst, _ := cmd.Flags().GetBool("rst")
-		psh, _ := cmd.Flags().GetBool("psh")
-		urg, _ := cmd.Flags().GetBool("urg")
-
-		tcpFlags = &proto.TCPFlags{
-			Syn: syn,
-			Ack: ack,
-			Fin: fin,
-			Rst: rst,
-			Psh: psh,
-			Urg: urg,
-		}
-	}
-
-	icmpType, _ := cmd.Flags().GetUint32("icmp-type")
-	icmpCode, _ := cmd.Flags().GetUint32("icmp-code")
-
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-
-	client := proto.NewDaemonServiceClient(conn)
-	resp, err := client.TracePacket(cmd.Context(), &proto.TracePacketRequest{
-		SourceIp:        args[1],
-		DestinationIp:   args[2],
-		Protocol:        protocol,
-		SourcePort:      uint32(sport),
-		DestinationPort: uint32(dport),
-		Direction:       direction,
-		TcpFlags:        tcpFlags,
-		IcmpType:        &icmpType,
-		IcmpCode:        &icmpCode,
-	})
-	if err != nil {
-		return fmt.Errorf("trace failed: %v", status.Convert(err).Message())
-	}
-
-	printTrace(cmd, args[1], args[2], protocol, sport, dport, resp)
-	return nil
-}
-
-func printTrace(cmd *cobra.Command, src, dst, proto string, sport, dport uint16, resp *proto.TracePacketResponse) {
-	cmd.Printf("Packet trace %s:%d -> %s:%d (%s)\n\n", src, sport, dst, dport, strings.ToUpper(proto))
-
-	for _, stage := range resp.Stages {
-		if stage.ForwardingDetails != nil {
-			cmd.Printf("%s: %s [%s]\n", stage.Name, stage.Message, *stage.ForwardingDetails)
-		} else {
-			cmd.Printf("%s: %s\n", stage.Name, stage.Message)
-		}
-	}
-
-	disposition := map[bool]string{
-		true:  "\033[32mALLOWED\033[0m", // Green
-		false: "\033[31mDENIED\033[0m",  // Red
-	}[resp.FinalDisposition]
-
-	cmd.Printf("\nFinal disposition: %s\n", disposition)
-}

client/cmd/up.go

@@ -190,7 +190,7 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command) error {
 	r.GetFullStatus()
 
 	connectClient := internal.NewConnectClient(ctx, config, r)
-	return connectClient.Run(nil)
+	return connectClient.Run()
 }
 
 func runInDaemonMode(ctx context.Context, cmd *cobra.Command) error {

client/firewall/create.go

@@ -14,13 +14,13 @@ import (
 )
 
 // NewFirewall creates a firewall manager instance
-func NewFirewall(iface IFaceMapper, _ *statemanager.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func NewFirewall(iface IFaceMapper, _ *statemanager.Manager) (firewall.Manager, error) {
 	if !iface.IsUserspaceBind() {
 		return nil, fmt.Errorf("not implemented for this OS: %s", runtime.GOOS)
 	}
 
 	// use userspace packet filtering firewall
-	fm, err := uspfilter.Create(iface, disableServerRoutes)
+	fm, err := uspfilter.Create(iface)
 	if err != nil {
 		return nil, err
 	}

client/firewall/create_linux.go

@@ -33,12 +33,12 @@ const SKIP_NFTABLES_ENV = "NB_SKIP_NFTABLES_CHECK"
 // FWType is the type for the firewall type
 type FWType int
 
-func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager) (firewall.Manager, error) {
 	// on the linux system we try to user nftables or iptables
 	// in any case, because we need to allow netbird interface traffic
 	// so we use AllowNetbird traffic from these firewall managers
 	// for the userspace packet filtering firewall
-	fm, err := createNativeFirewall(iface, stateManager, disableServerRoutes)
+	fm, err := createNativeFirewall(iface, stateManager)
 
 	if !iface.IsUserspaceBind() {
 		return fm, err
@@ -47,10 +47,10 @@ func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager, disableS
 	if err != nil {
 		log.Warnf("failed to create native firewall: %v. Proceeding with userspace", err)
 	}
-	return createUserspaceFirewall(iface, fm, disableServerRoutes)
+	return createUserspaceFirewall(iface, fm)
 }
 
-func createNativeFirewall(iface IFaceMapper, stateManager *statemanager.Manager, routes bool) (firewall.Manager, error) {
+func createNativeFirewall(iface IFaceMapper, stateManager *statemanager.Manager) (firewall.Manager, error) {
 	fm, err := createFW(iface)
 	if err != nil {
 		return nil, fmt.Errorf("create firewall: %s", err)
@@ -77,12 +77,12 @@ func createFW(iface IFaceMapper) (firewall.Manager, error) {
 	}
 }
 
-func createUserspaceFirewall(iface IFaceMapper, fm firewall.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func createUserspaceFirewall(iface IFaceMapper, fm firewall.Manager) (firewall.Manager, error) {
 	var errUsp error
 	if fm != nil {
-		fm, errUsp = uspfilter.CreateWithNativeFirewall(iface, fm, disableServerRoutes)
+		fm, errUsp = uspfilter.CreateWithNativeFirewall(iface, fm)
 	} else {
-		fm, errUsp = uspfilter.Create(iface, disableServerRoutes)
+		fm, errUsp = uspfilter.Create(iface)
 	}
 
 	if errUsp != nil {

client/firewall/iface.go

@@ -1,8 +1,6 @@
 package firewall
 
 import (
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
 	"github.com/netbirdio/netbird/client/iface/device"
 )
 
@@ -12,6 +10,4 @@ type IFaceMapper interface {
 	Address() device.WGAddress
 	IsUserspaceBind() bool
 	SetFilter(device.PacketFilter) error
-	GetDevice() *device.FilteredDevice
-	GetWGDevice() *wgdevice.Device
 }

client/firewall/iptables/acl_linux.go

@@ -3,7 +3,7 @@ package iptables
 import (
 	"fmt"
 	"net"
-	"slices"
+	"strconv"
 
 	"github.com/coreos/go-iptables/iptables"
 	"github.com/google/uuid"
@@ -86,19 +86,19 @@ func (m *aclManager) AddPeerFiltering(
 	action firewall.Action,
 	ipsetName string,
 ) ([]firewall.Rule, error) {
+	var dPortVal, sPortVal string
+	if dPort != nil && dPort.Values != nil {
+		// TODO: we support only one port per rule in current implementation of ACLs
+		dPortVal = strconv.Itoa(dPort.Values[0])
+	}
+	if sPort != nil && sPort.Values != nil {
+		sPortVal = strconv.Itoa(sPort.Values[0])
+	}
+
 	chain := chainNameInputRules
 
-	ipsetName = transformIPsetName(ipsetName, sPort, dPort)
-	specs := filterRuleSpecs(ip, string(protocol), sPort, dPort, action, ipsetName)
-
-	mangleSpecs := slices.Clone(specs)
-	mangleSpecs = append(mangleSpecs,
-		"-i", m.wgIface.Name(),
-		"-m", "addrtype", "--dst-type", "LOCAL",
-		"-j", "MARK", "--set-xmark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected),
-	)
-
-	specs = append(specs, "-j", actionToStr(action))
+	ipsetName = transformIPsetName(ipsetName, sPortVal, dPortVal)
+	specs := filterRuleSpecs(ip, string(protocol), sPortVal, dPortVal, action, ipsetName)
 	if ipsetName != "" {
 		if ipList, ipsetExists := m.ipsetStore.ipset(ipsetName); ipsetExists {
 			if err := ipset.Add(ipsetName, ip.String()); err != nil {
@@ -130,7 +130,7 @@ func (m *aclManager) AddPeerFiltering(
 		m.ipsetStore.addIpList(ipsetName, ipList)
 	}
 
-	ok, err := m.iptablesClient.Exists(tableFilter, chain, specs...)
+	ok, err := m.iptablesClient.Exists("filter", chain, specs...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to check rule: %w", err)
 	}
@@ -138,22 +138,16 @@ func (m *aclManager) AddPeerFiltering(
 		return nil, fmt.Errorf("rule already exists")
 	}
 
-	if err := m.iptablesClient.Append(tableFilter, chain, specs...); err != nil {
+	if err := m.iptablesClient.Append("filter", chain, specs...); err != nil {
 		return nil, err
 	}
 
-	if err := m.iptablesClient.Append(tableMangle, chainRTPRE, mangleSpecs...); err != nil {
-		log.Errorf("failed to add mangle rule: %v", err)
-		mangleSpecs = nil
-	}
-
 	rule := &Rule{
-		ruleID:      uuid.New().String(),
-		specs:       specs,
-		mangleSpecs: mangleSpecs,
-		ipsetName:   ipsetName,
-		ip:          ip.String(),
-		chain:       chain,
+		ruleID:    uuid.New().String(),
+		specs:     specs,
+		ipsetName: ipsetName,
+		ip:        ip.String(),
+		chain:     chain,
 	}
 
 	m.updateState()
@@ -196,12 +190,6 @@ func (m *aclManager) DeletePeerRule(rule firewall.Rule) error {
 		return fmt.Errorf("failed to delete rule: %s, %v: %w", r.chain, r.specs, err)
 	}
 
-	if r.mangleSpecs != nil {
-		if err := m.iptablesClient.Delete(tableMangle, chainRTPRE, r.mangleSpecs...); err != nil {
-			log.Errorf("failed to delete mangle rule: %v", err)
-		}
-	}
-
 	m.updateState()
 
 	return nil
@@ -322,10 +310,17 @@ func (m *aclManager) seedInitialEntries() {
 func (m *aclManager) seedInitialOptionalEntries() {
 	m.optionalEntries["FORWARD"] = []entry{
 		{
-			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected), "-j", "ACCEPT"},
+			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.PreroutingFwmarkRedirected)},
+			position: 1,
+		},
+	}
 }
 
 func (m *aclManager) appendToEntries(chainName string, spec []string) {
@@ -359,7 +354,7 @@ func (m *aclManager) updateState() {
 }
 
 // filterRuleSpecs returns the specs of a filtering rule
-func filterRuleSpecs(ip net.IP, protocol string, sPort, dPort *firewall.Port, action firewall.Action, ipsetName string) (specs []string) {
+func filterRuleSpecs(ip net.IP, protocol, sPort, dPort string, action firewall.Action, ipsetName string) (specs []string) {
 	matchByIP := true
 	// don't use IP matching if IP is ip 0.0.0.0
 	if ip.String() == "0.0.0.0" {
@@ -376,9 +371,13 @@ func filterRuleSpecs(ip net.IP, protocol string, sPort, dPort *firewall.Port, ac
 	if protocol != "all" {
 		specs = append(specs, "-p", protocol)
 	}
-	specs = append(specs, applyPort("--sport", sPort)...)
-	specs = append(specs, applyPort("--dport", dPort)...)
-	return specs
+	if sPort != "" {
+		specs = append(specs, "--sport", sPort)
+	}
+	if dPort != "" {
+		specs = append(specs, "--dport", dPort)
+	}
+	return append(specs, "-j", actionToStr(action))
 }
 
 func actionToStr(action firewall.Action) string {
@@ -388,15 +387,15 @@ func actionToStr(action firewall.Action) string {
 	return "DROP"
 }
 
-func transformIPsetName(ipsetName string, sPort, dPort *firewall.Port) string {
+func transformIPsetName(ipsetName string, sPort, dPort string) string {
 	switch {
 	case ipsetName == "":
 		return ""
-	case sPort != nil && dPort != nil:
+	case sPort != "" && dPort != "":
 		return ipsetName + "-sport-dport"
-	case sPort != nil:
+	case sPort != "":
 		return ipsetName + "-sport"
-	case dPort != nil:
+	case dPort != "":
 		return ipsetName + "-dport"
 	default:
 		return ipsetName

client/firewall/iptables/manager_linux.go

@@ -213,19 +213,6 @@ func (m *Manager) AllowNetbird() error {
 // Flush doesn't need to be implemented for this manager
 func (m *Manager) Flush() error { return nil }
 
-// SetLogLevel sets the log level for the firewall manager
-func (m *Manager) SetLogLevel(log.Level) {
-	// not supported
-}
-
-func (m *Manager) EnableRouting() error {
-	return nil
-}
-
-func (m *Manager) DisableRouting() error {
-	return nil
-}
-
 func getConntrackEstablished() []string {
 	return []string{"-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT"}
 }

client/firewall/iptables/manager_linux_test.go

@@ -72,8 +72,7 @@ func TestIptablesManager(t *testing.T) {
 	t.Run("add second rule", func(t *testing.T) {
 		ip := net.ParseIP("10.20.0.3")
 		port := &fw.Port{
-			IsRange: true,
-			Values:  []uint16{8043, 8046},
+			Values: []int{8043: 8046},
 		}
 		rule2, err = manager.AddPeerFiltering(ip, "tcp", port, nil, fw.ActionAccept, "", "accept HTTPS traffic from ports range")
 		require.NoError(t, err, "failed to add rule")
@@ -96,7 +95,7 @@ func TestIptablesManager(t *testing.T) {
 	t.Run("reset check", func(t *testing.T) {
 		// add second rule
 		ip := net.ParseIP("10.20.0.3")
-		port := &fw.Port{Values: []uint16{5353}}
+		port := &fw.Port{Values: []int{5353}}
 		_, err = manager.AddPeerFiltering(ip, "udp", nil, port, fw.ActionAccept, "", "accept Fake DNS traffic")
 		require.NoError(t, err, "failed to add rule")
 
@@ -146,7 +145,7 @@ func TestIptablesManagerIPSet(t *testing.T) {
 	t.Run("add second rule", func(t *testing.T) {
 		ip := net.ParseIP("10.20.0.3")
 		port := &fw.Port{
-			Values: []uint16{443},
+			Values: []int{443},
 		}
 		rule2, err = manager.AddPeerFiltering(ip, "tcp", port, nil, fw.ActionAccept, "default", "accept HTTPS traffic from ports range")
 		for _, r := range rule2 {
@@ -215,7 +214,7 @@ func TestIptablesCreatePerformance(t *testing.T) {
 			ip := net.ParseIP("10.20.0.100")
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
-				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
+				port := &fw.Port{Values: []int{1000 + i}}
 				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
 
 				require.NoError(t, err, "failed to add rule")

client/firewall/iptables/router_linux.go

@@ -135,16 +135,7 @@ func (r *router) AddRouteFiltering(
 	}
 
 	rule := genRouteFilteringRuleSpec(params)
-	// Insert DROP rules at the beginning, append ACCEPT rules at the end
-	var err error
-	if action == firewall.ActionDrop {
-		// after the established rule
-		err = r.iptablesClient.Insert(tableFilter, chainRTFWD, 2, rule...)
-	} else {
-		err = r.iptablesClient.Append(tableFilter, chainRTFWD, rule...)
-	}
-
-	if err != nil {
+	if err := r.iptablesClient.Append(tableFilter, chainRTFWD, rule...); err != nil {
 		return nil, fmt.Errorf("add route rule: %v", err)
 	}
 
@@ -599,10 +590,10 @@ func applyPort(flag string, port *firewall.Port) []string {
 	if len(port.Values) > 1 {
 		portList := make([]string, len(port.Values))
 		for i, p := range port.Values {
-			portList[i] = strconv.Itoa(int(p))
+			portList[i] = strconv.Itoa(p)
 		}
 		return []string{"-m", "multiport", flag, strings.Join(portList, ",")}
 	}
 
-	return []string{flag, strconv.Itoa(int(port.Values[0]))}
+	return []string{flag, strconv.Itoa(port.Values[0])}
 }

client/firewall/iptables/router_linux_test.go

@@ -239,7 +239,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			destination: netip.MustParsePrefix("10.0.0.0/24"),
 			proto:       firewall.ProtocolTCP,
 			sPort:       nil,
-			dPort:       &firewall.Port{Values: []uint16{80}},
+			dPort:       &firewall.Port{Values: []int{80}},
 			direction:   firewall.RuleDirectionIN,
 			action:      firewall.ActionAccept,
 			expectSet:   false,
@@ -252,7 +252,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			},
 			destination: netip.MustParsePrefix("10.0.0.0/8"),
 			proto:       firewall.ProtocolUDP,
-			sPort:       &firewall.Port{Values: []uint16{1024, 2048}, IsRange: true},
+			sPort:       &firewall.Port{Values: []int{1024, 2048}, IsRange: true},
 			dPort:       nil,
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionDrop,
@@ -285,7 +285,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			sources:     []netip.Prefix{netip.MustParsePrefix("172.16.0.0/12")},
 			destination: netip.MustParsePrefix("192.168.0.0/16"),
 			proto:       firewall.ProtocolTCP,
-			sPort:       &firewall.Port{Values: []uint16{80, 443, 8080}},
+			sPort:       &firewall.Port{Values: []int{80, 443, 8080}},
 			dPort:       nil,
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionAccept,
@@ -297,7 +297,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			destination: netip.MustParsePrefix("10.0.0.0/24"),
 			proto:       firewall.ProtocolUDP,
 			sPort:       nil,
-			dPort:       &firewall.Port{Values: []uint16{5000, 5100}, IsRange: true},
+			dPort:       &firewall.Port{Values: []int{5000, 5100}, IsRange: true},
 			direction:   firewall.RuleDirectionIN,
 			action:      firewall.ActionDrop,
 			expectSet:   false,
@@ -307,8 +307,8 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			sources:     []netip.Prefix{netip.MustParsePrefix("10.0.0.0/24")},
 			destination: netip.MustParsePrefix("172.16.0.0/16"),
 			proto:       firewall.ProtocolTCP,
-			sPort:       &firewall.Port{Values: []uint16{1024, 65535}, IsRange: true},
-			dPort:       &firewall.Port{Values: []uint16{22}},
+			sPort:       &firewall.Port{Values: []int{1024, 65535}, IsRange: true},
+			dPort:       &firewall.Port{Values: []int{22}},
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionAccept,
 			expectSet:   false,

client/firewall/iptables/rule.go

@@ -5,10 +5,9 @@ type Rule struct {
 	ruleID    string
 	ipsetName string
 
-	specs       []string
-	mangleSpecs []string
-	ip          string
-	chain       string
+	specs []string
+	ip    string
+	chain string
 }
 
 // GetRuleID returns the rule id

client/firewall/manager/firewall.go

@@ -99,12 +99,6 @@ type Manager interface {
 
 	// Flush the changes to firewall controller
 	Flush() error
-
-	SetLogLevel(log.Level)
-
-	EnableRouting() error
-
-	DisableRouting() error
 }
 
 func GenKey(format string, pair RouterPair) string {

client/firewall/manager/port.go

@@ -30,7 +30,7 @@ type Port struct {
 	IsRange bool
 
 	// Values contains one value for single port, multiple values for the list of ports, or two values for the range of ports
-	Values []uint16
+	Values []int
 }
 
 // String interface implementation
@@ -40,11 +40,7 @@ func (p *Port) String() string {
 		if ports != "" {
 			ports += ","
 		}
-		ports += strconv.Itoa(int(port))
+		ports += strconv.Itoa(port)
 	}
-	if p.IsRange {
-		ports = "range:" + ports
-	}
-
 	return ports
 }

client/firewall/nftables/acl_linux.go

@@ -2,9 +2,9 @@ package nftables
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"net"
-	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -46,7 +46,6 @@ type AclManager struct {
 
 	workTable       *nftables.Table
 	chainInputRules *nftables.Chain
-	chainPrerouting *nftables.Chain
 
 	ipsetStore *ipsetStore
 	rules      map[string]*Rule
@@ -119,32 +118,23 @@ func (m *AclManager) DeletePeerRule(rule firewall.Rule) error {
 	}
 
 	if r.nftSet == nil {
-		if err := m.rConn.DelRule(r.nftRule); err != nil {
+		err := m.rConn.DelRule(r.nftRule)
+		if err != nil {
 			log.Errorf("failed to delete rule: %v", err)
 		}
-		if r.mangleRule != nil {
-			if err := m.rConn.DelRule(r.mangleRule); err != nil {
-				log.Errorf("failed to delete mangle rule: %v", err)
-			}
-		}
 		delete(m.rules, r.GetRuleID())
 		return m.rConn.Flush()
 	}
 
 	ips, ok := m.ipsetStore.ips(r.nftSet.Name)
 	if !ok {
-		if err := m.rConn.DelRule(r.nftRule); err != nil {
+		err := m.rConn.DelRule(r.nftRule)
+		if err != nil {
 			log.Errorf("failed to delete rule: %v", err)
 		}
-		if r.mangleRule != nil {
-			if err := m.rConn.DelRule(r.mangleRule); err != nil {
-				log.Errorf("failed to delete mangle rule: %v", err)
-			}
-		}
 		delete(m.rules, r.GetRuleID())
 		return m.rConn.Flush()
 	}
-
 	if _, ok := ips[r.ip.String()]; ok {
 		err := m.sConn.SetDeleteElements(r.nftSet, []nftables.SetElement{{Key: r.ip.To4()}})
 		if err != nil {
@@ -163,16 +153,12 @@ func (m *AclManager) DeletePeerRule(rule firewall.Rule) error {
 		return nil
 	}
 
-	if err := m.rConn.DelRule(r.nftRule); err != nil {
+	err := m.rConn.DelRule(r.nftRule)
+	if err != nil {
 		log.Errorf("failed to delete rule: %v", err)
 	}
-	if r.mangleRule != nil {
-		if err := m.rConn.DelRule(r.mangleRule); err != nil {
-			log.Errorf("failed to delete mangle rule: %v", err)
-		}
-	}
-
-	if err := m.rConn.Flush(); err != nil {
+	err = m.rConn.Flush()
+	if err != nil {
 		return err
 	}
 
@@ -239,12 +225,9 @@ func (m *AclManager) Flush() error {
 		return err
 	}
 
-	if err := m.refreshRuleHandles(m.chainInputRules, false); err != nil {
+	if err := m.refreshRuleHandles(m.chainInputRules); err != nil {
 		log.Errorf("failed to refresh rule handles ipv4 input chain: %v", err)
 	}
-	if err := m.refreshRuleHandles(m.chainPrerouting, true); err != nil {
-		log.Errorf("failed to refresh rule handles prerouting chain: %v", err)
-	}
 
 	return nil
 }
@@ -261,11 +244,10 @@ func (m *AclManager) addIOFiltering(
 	ruleId := generatePeerRuleId(ip, sPort, dPort, action, ipset)
 	if r, ok := m.rules[ruleId]; ok {
 		return &Rule{
-			nftRule:    r.nftRule,
-			mangleRule: r.mangleRule,
-			nftSet:     r.nftSet,
-			ruleID:     r.ruleID,
-			ip:         ip,
+			r.nftRule,
+			r.nftSet,
+			r.ruleID,
+			ip,
 		}, nil
 	}
 
@@ -326,16 +308,43 @@ func (m *AclManager) addIOFiltering(
 		}
 	}
 
-	expressions = append(expressions, applyPort(sPort, true)...)
-	expressions = append(expressions, applyPort(dPort, false)...)
+	if sPort != nil && len(sPort.Values) != 0 {
+		expressions = append(expressions,
+			&expr.Payload{
+				DestRegister: 1,
+				Base:         expr.PayloadBaseTransportHeader,
+				Offset:       0,
+				Len:          2,
+			},
+			&expr.Cmp{
+				Op:       expr.CmpOpEq,
+				Register: 1,
+				Data:     encodePort(*sPort),
+			},
+		)
+	}
 
-	mainExpressions := slices.Clone(expressions)
+	if dPort != nil && len(dPort.Values) != 0 {
+		expressions = append(expressions,
+			&expr.Payload{
+				DestRegister: 1,
+				Base:         expr.PayloadBaseTransportHeader,
+				Offset:       2,
+				Len:          2,
+			},
+			&expr.Cmp{
+				Op:       expr.CmpOpEq,
+				Register: 1,
+				Data:     encodePort(*dPort),
+			},
+		)
+	}
 
 	switch action {
 	case firewall.ActionAccept:
-		mainExpressions = append(mainExpressions, &expr.Verdict{Kind: expr.VerdictAccept})
+		expressions = append(expressions, &expr.Verdict{Kind: expr.VerdictAccept})
 	case firewall.ActionDrop:
-		mainExpressions = append(mainExpressions, &expr.Verdict{Kind: expr.VerdictDrop})
+		expressions = append(expressions, &expr.Verdict{Kind: expr.VerdictDrop})
 	}
 
 	userData := []byte(strings.Join([]string{ruleId, comment}, " "))
@@ -344,82 +353,23 @@ func (m *AclManager) addIOFiltering(
 	nftRule := m.rConn.AddRule(&nftables.Rule{
 		Table:    m.workTable,
 		Chain:    chain,
-		Exprs:    mainExpressions,
+		Exprs:    expressions,
 		UserData: userData,
 	})
 
-	if err := m.rConn.Flush(); err != nil {
-		return nil, fmt.Errorf(flushError, err)
-	}
-
 	rule := &Rule{
-		nftRule:    nftRule,
-		mangleRule: m.createPreroutingRule(expressions, userData),
-		nftSet:     ipset,
-		ruleID:     ruleId,
-		ip:         ip,
+		nftRule: nftRule,
+		nftSet:  ipset,
+		ruleID:  ruleId,
+		ip:      ip,
 	}
 	m.rules[ruleId] = rule
 	if ipset != nil {
 		m.ipsetStore.AddReferenceToIpset(ipset.Name)
 	}
-
 	return rule, nil
 }
 
-func (m *AclManager) createPreroutingRule(expressions []expr.Any, userData []byte) *nftables.Rule {
-	if m.chainPrerouting == nil {
-		log.Warn("prerouting chain is not created")
-		return nil
-	}
-
-	preroutingExprs := slices.Clone(expressions)
-
-	// interface
-	preroutingExprs = append([]expr.Any{
-		&expr.Meta{
-			Key:      expr.MetaKeyIIFNAME,
-			Register: 1,
-		},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 1,
-			Data:     ifname(m.wgIface.Name()),
-		},
-	}, preroutingExprs...)
-
-	// local destination and mark
-	preroutingExprs = append(preroutingExprs,
-		&expr.Fib{
-			Register:       1,
-			ResultADDRTYPE: true,
-			FlagDADDR:      true,
-		},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 1,
-			Data:     binaryutil.NativeEndian.PutUint32(unix.RTN_LOCAL),
-		},
-
-		&expr.Immediate{
-			Register: 1,
-			Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
-		},
-		&expr.Meta{
-			Key:            expr.MetaKeyMARK,
-			Register:       1,
-			SourceRegister: true,
-		},
-	)
-
-	return m.rConn.AddRule(&nftables.Rule{
-		Table:    m.workTable,
-		Chain:    m.chainPrerouting,
-		Exprs:    preroutingExprs,
-		UserData: userData,
-	})
-}
-
 func (m *AclManager) createDefaultChains() (err error) {
 	// chainNameInputRules
 	chain := m.createChain(chainNameInputRules)
@@ -463,7 +413,7 @@ func (m *AclManager) createDefaultChains() (err error) {
 // go through the input filter as well. This will enable e.g. Docker services to keep working by accessing the
 // netbird peer IP.
 func (m *AclManager) allowRedirectedTraffic(chainFwFilter *nftables.Chain) error {
-	m.chainPrerouting = m.rConn.AddChain(&nftables.Chain{
+	preroutingChain := m.rConn.AddChain(&nftables.Chain{
 		Name:     chainNamePrerouting,
 		Table:    m.workTable,
 		Type:     nftables.ChainTypeFilter,
@@ -471,6 +421,8 @@ func (m *AclManager) allowRedirectedTraffic(chainFwFilter *nftables.Chain) error
 		Priority: nftables.ChainPriorityMangle,
 	})
 
+	m.addPreroutingRule(preroutingChain)
+
 	m.addFwmarkToForward(chainFwFilter)
 
 	if err := m.rConn.Flush(); err != nil {
@@ -480,6 +432,43 @@ func (m *AclManager) allowRedirectedTraffic(chainFwFilter *nftables.Chain) error
 	return nil
 }
 
+func (m *AclManager) addPreroutingRule(preroutingChain *nftables.Chain) {
+	m.rConn.AddRule(&nftables.Rule{
+		Table: m.workTable,
+		Chain: preroutingChain,
+		Exprs: []expr.Any{
+			&expr.Meta{
+				Key:      expr.MetaKeyIIFNAME,
+				Register: 1,
+			},
+			&expr.Cmp{
+				Op:       expr.CmpOpEq,
+				Register: 1,
+				Data:     ifname(m.wgIface.Name()),
+			},
+			&expr.Fib{
+				Register:       1,
+				ResultADDRTYPE: true,
+				FlagDADDR:      true,
+			},
+			&expr.Cmp{
+				Op:       expr.CmpOpEq,
+				Register: 1,
+				Data:     binaryutil.NativeEndian.PutUint32(unix.RTN_LOCAL),
+			},
+			&expr.Immediate{
+				Register: 1,
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
+			},
+			&expr.Meta{
+				Key:            expr.MetaKeyMARK,
+				Register:       1,
+				SourceRegister: true,
+			},
+		},
+	})
+}
+
 func (m *AclManager) addFwmarkToForward(chainFwFilter *nftables.Chain) {
 	m.rConn.InsertRule(&nftables.Rule{
 		Table: m.workTable,
@@ -495,7 +484,8 @@ func (m *AclManager) addFwmarkToForward(chainFwFilter *nftables.Chain) {
 				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
 			},
 			&expr.Verdict{
-				Kind: expr.VerdictAccept,
+				Kind:  expr.VerdictJump,
+				Chain: m.chainInputRules.Name,
 			},
 		},
 	})
@@ -642,7 +632,6 @@ func (m *AclManager) flushWithBackoff() (err error) {
 	for i := 0; ; i++ {
 		err = m.rConn.Flush()
 		if err != nil {
-			log.Debugf("failed to flush nftables: %v", err)
 			if !strings.Contains(err.Error(), "busy") {
 				return
 			}
@@ -659,7 +648,7 @@ func (m *AclManager) flushWithBackoff() (err error) {
 	return
 }
 
-func (m *AclManager) refreshRuleHandles(chain *nftables.Chain, mangle bool) error {
+func (m *AclManager) refreshRuleHandles(chain *nftables.Chain) error {
 	if m.workTable == nil || chain == nil {
 		return nil
 	}
@@ -676,11 +665,7 @@ func (m *AclManager) refreshRuleHandles(chain *nftables.Chain, mangle bool) erro
 		split := bytes.Split(rule.UserData, []byte(" "))
 		r, ok := m.rules[string(split[0])]
 		if ok {
-			if mangle {
-				*r.mangleRule = *rule
-			} else {
-				*r.nftRule = *rule
-			}
+			*r.nftRule = *rule
 		}
 	}
 
@@ -704,6 +689,12 @@ func generatePeerRuleId(ip net.IP, sPort *firewall.Port, dPort *firewall.Port, a
 	return "set:" + ipset.Name + rulesetID
 }
 
+func encodePort(port firewall.Port) []byte {
+	bs := make([]byte, 2)
+	binary.BigEndian.PutUint16(bs, uint16(port.Values[0]))
+	return bs
+}
+
 func ifname(n string) []byte {
 	b := make([]byte, 16)
 	copy(b, n+"\x00")

client/firewall/nftables/manager_linux.go

@@ -318,19 +318,6 @@ func (m *Manager) cleanupNetbirdTables() error {
 	return nil
 }
 
-// SetLogLevel sets the log level for the firewall manager
-func (m *Manager) SetLogLevel(log.Level) {
-	// not supported
-}
-
-func (m *Manager) EnableRouting() error {
-	return nil
-}
-
-func (m *Manager) DisableRouting() error {
-	return nil
-}
-
 // Flush rule/chain/set operations from the buffer
 //
 // Method also get all rules after flush and refreshes handle values in the rulesets

client/firewall/nftables/manager_linux_test.go

@@ -74,7 +74,7 @@ func TestNftablesManager(t *testing.T) {
 
 	testClient := &nftables.Conn{}
 
-	rule, err := manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{53}}, fw.ActionDrop, "", "")
+	rule, err := manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []int{53}}, fw.ActionDrop, "", "")
 	require.NoError(t, err, "failed to add rule")
 
 	err = manager.Flush()
@@ -107,7 +107,7 @@ func TestNftablesManager(t *testing.T) {
 			Kind: expr.VerdictAccept,
 		},
 	}
-	compareExprsIgnoringCounters(t, rules[0].Exprs, expectedExprs1)
+	require.ElementsMatch(t, rules[0].Exprs, expectedExprs1, "expected the same expressions")
 
 	ipToAdd, _ := netip.AddrFromSlice(ip)
 	add := ipToAdd.Unmap()
@@ -200,7 +200,7 @@ func TestNFtablesCreatePerformance(t *testing.T) {
 			ip := net.ParseIP("10.20.0.100")
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
-				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
+				port := &fw.Port{Values: []int{1000 + i}}
 				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
 				require.NoError(t, err, "failed to add rule")
 
@@ -283,7 +283,7 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 	})
 
 	ip := net.ParseIP("100.96.0.1")
-	_, err = manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept, "", "test rule")
+	_, err = manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []int{80}}, fw.ActionAccept, "", "test rule")
 	require.NoError(t, err, "failed to add peer filtering rule")
 
 	_, err = manager.AddRouteFiltering(
@@ -291,7 +291,7 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 		netip.MustParsePrefix("10.1.0.0/24"),
 		fw.ProtocolTCP,
 		nil,
-		&fw.Port{Values: []uint16{443}},
+		&fw.Port{Values: []int{443}},
 		fw.ActionAccept,
 	)
 	require.NoError(t, err, "failed to add route filtering rule")
@@ -307,18 +307,3 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 	stdout, stderr = runIptablesSave(t)
 	verifyIptablesOutput(t, stdout, stderr)
 }
-
-func compareExprsIgnoringCounters(t *testing.T, got, want []expr.Any) {
-	t.Helper()
-	require.Equal(t, len(got), len(want), "expression count mismatch")
-
-	for i := range got {
-		if _, isCounter := got[i].(*expr.Counter); isCounter {
-			_, wantIsCounter := want[i].(*expr.Counter)
-			require.True(t, wantIsCounter, "expected Counter at index %d", i)
-			continue
-		}
-
-		require.Equal(t, got[i], want[i], "expression mismatch at index %d", i)
-	}
-}

client/firewall/nftables/router_linux.go

@@ -233,13 +233,7 @@ func (r *router) AddRouteFiltering(
 		UserData: []byte(ruleKey),
 	}
 
-	// Insert DROP rules at the beginning, append ACCEPT rules at the end
-	if action == firewall.ActionDrop {
-		// TODO: Insert after the established rule
-		rule = r.conn.InsertRule(rule)
-	} else {
-		rule = r.conn.AddRule(rule)
-	}
+	rule = r.conn.AddRule(rule)
 
 	log.Tracef("Adding route rule %s", spew.Sdump(rule))
 	if err := r.conn.Flush(); err != nil {
@@ -962,12 +956,12 @@ func applyPort(port *firewall.Port, isSource bool) []expr.Any {
 			&expr.Cmp{
 				Op:       expr.CmpOpGte,
 				Register: 1,
-				Data:     binaryutil.BigEndian.PutUint16(port.Values[0]),
+				Data:     binaryutil.BigEndian.PutUint16(uint16(port.Values[0])),
 			},
 			&expr.Cmp{
 				Op:       expr.CmpOpLte,
 				Register: 1,
-				Data:     binaryutil.BigEndian.PutUint16(port.Values[1]),
+				Data:     binaryutil.BigEndian.PutUint16(uint16(port.Values[1])),
 			},
 		)
 	} else {
@@ -986,7 +980,7 @@ func applyPort(port *firewall.Port, isSource bool) []expr.Any {
 			exprs = append(exprs, &expr.Cmp{
 				Op:       expr.CmpOpEq,
 				Register: 1,
-				Data:     binaryutil.BigEndian.PutUint16(p),
+				Data:     binaryutil.BigEndian.PutUint16(uint16(p)),
 			})
 		}
 	}

client/firewall/nftables/router_linux_test.go

@@ -222,7 +222,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			destination: netip.MustParsePrefix("10.0.0.0/24"),
 			proto:       firewall.ProtocolTCP,
 			sPort:       nil,
-			dPort:       &firewall.Port{Values: []uint16{80}},
+			dPort:       &firewall.Port{Values: []int{80}},
 			direction:   firewall.RuleDirectionIN,
 			action:      firewall.ActionAccept,
 			expectSet:   false,
@@ -235,7 +235,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			},
 			destination: netip.MustParsePrefix("10.0.0.0/8"),
 			proto:       firewall.ProtocolUDP,
-			sPort:       &firewall.Port{Values: []uint16{1024, 2048}, IsRange: true},
+			sPort:       &firewall.Port{Values: []int{1024, 2048}, IsRange: true},
 			dPort:       nil,
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionDrop,
@@ -268,7 +268,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			sources:     []netip.Prefix{netip.MustParsePrefix("172.16.0.0/12")},
 			destination: netip.MustParsePrefix("192.168.0.0/16"),
 			proto:       firewall.ProtocolTCP,
-			sPort:       &firewall.Port{Values: []uint16{80, 443, 8080}},
+			sPort:       &firewall.Port{Values: []int{80, 443, 8080}},
 			dPort:       nil,
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionAccept,
@@ -280,7 +280,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			destination: netip.MustParsePrefix("10.0.0.0/24"),
 			proto:       firewall.ProtocolUDP,
 			sPort:       nil,
-			dPort:       &firewall.Port{Values: []uint16{5000, 5100}, IsRange: true},
+			dPort:       &firewall.Port{Values: []int{5000, 5100}, IsRange: true},
 			direction:   firewall.RuleDirectionIN,
 			action:      firewall.ActionDrop,
 			expectSet:   false,
@@ -290,8 +290,8 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 			sources:     []netip.Prefix{netip.MustParsePrefix("10.0.0.0/24")},
 			destination: netip.MustParsePrefix("172.16.0.0/16"),
 			proto:       firewall.ProtocolTCP,
-			sPort:       &firewall.Port{Values: []uint16{1024, 65535}, IsRange: true},
-			dPort:       &firewall.Port{Values: []uint16{22}},
+			sPort:       &firewall.Port{Values: []int{1024, 65535}, IsRange: true},
+			dPort:       &firewall.Port{Values: []int{22}},
 			direction:   firewall.RuleDirectionOUT,
 			action:      firewall.ActionAccept,
 			expectSet:   false,

client/firewall/nftables/rule_linux.go

@@ -8,11 +8,10 @@ import (
 
 // Rule to handle management of rules
 type Rule struct {
-	nftRule    *nftables.Rule
-	mangleRule *nftables.Rule
-	nftSet     *nftables.Set
-	ruleID     string
-	ip         net.IP
+	nftRule *nftables.Rule
+	nftSet  *nftables.Set
+	ruleID  string
+	ip      net.IP
 }
 
 // GetRuleID returns the rule id

client/firewall/uspfilter/allow_netbird.go

@@ -3,11 +3,6 @@
 package uspfilter
 
 import (
-	"context"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
@@ -22,29 +17,17 @@ func (m *Manager) Reset(stateManager *statemanager.Manager) error {
 
 	if m.udpTracker != nil {
 		m.udpTracker.Close()
-		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger)
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout)
 	}
 
 	if m.icmpTracker != nil {
 		m.icmpTracker.Close()
-		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger)
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout)
 	}
 
 	if m.tcpTracker != nil {
 		m.tcpTracker.Close()
-		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger)
-	}
-
-	if m.forwarder != nil {
-		m.forwarder.Stop()
-	}
-
-	if m.logger != nil {
-		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-		defer cancel()
-		if err := m.logger.Stop(ctx); err != nil {
-			log.Errorf("failed to shutdown logger: %v", err)
-		}
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout)
 	}
 
 	if m.nativeFirewall != nil {

client/firewall/uspfilter/allow_netbird_windows.go

@@ -1,11 +1,9 @@
 package uspfilter
 
 import (
-	"context"
 	"fmt"
 	"os/exec"
 	"syscall"
-	"time"
 
 	log "github.com/sirupsen/logrus"
 
@@ -31,29 +29,17 @@ func (m *Manager) Reset(*statemanager.Manager) error {
 
 	if m.udpTracker != nil {
 		m.udpTracker.Close()
-		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger)
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout)
 	}
 
 	if m.icmpTracker != nil {
 		m.icmpTracker.Close()
-		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger)
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout)
 	}
 
 	if m.tcpTracker != nil {
 		m.tcpTracker.Close()
-		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger)
-	}
-
-	if m.forwarder != nil {
-		m.forwarder.Stop()
-	}
-
-	if m.logger != nil {
-		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-		defer cancel()
-		if err := m.logger.Stop(ctx); err != nil {
-			log.Errorf("failed to shutdown logger: %v", err)
-		}
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout)
 	}
 
 	if !isWindowsFirewallReachable() {

client/firewall/uspfilter/common/iface.go

@@ -1,16 +0,0 @@
-package common
-
-import (
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
-	"github.com/netbirdio/netbird/client/iface"
-	"github.com/netbirdio/netbird/client/iface/device"
-)
-
-// IFaceMapper defines subset methods of interface required for manager
-type IFaceMapper interface {
-	SetFilter(device.PacketFilter) error
-	Address() iface.WGAddress
-	GetWGDevice() *wgdevice.Device
-	GetDevice() *device.FilteredDevice
-}

client/firewall/uspfilter/conntrack/common.go

@@ -10,11 +10,12 @@ import (
 
 // BaseConnTrack provides common fields and locking for all connection types
 type BaseConnTrack struct {
-	SourceIP   net.IP
-	DestIP     net.IP
-	SourcePort uint16
-	DestPort   uint16
-	lastSeen   atomic.Int64 // Unix nano for atomic access
+	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
@@ -24,6 +25,16 @@ 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())

client/firewall/uspfilter/conntrack/common_test.go

@@ -3,14 +3,8 @@ package conntrack
 import (
 	"net"
 	"testing"
-
-	"github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 )
 
-var logger = log.NewFromLogrus(logrus.StandardLogger())
-
 func BenchmarkIPOperations(b *testing.B) {
 	b.Run("MakeIPAddr", func(b *testing.B) {
 		ip := net.ParseIP("192.168.1.1")
@@ -40,11 +34,37 @@ func BenchmarkIPOperations(b *testing.B) {
 	})
 
 }
+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, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout)
 		defer tracker.Close()
 
 		// Generate different IPs
@@ -69,7 +89,7 @@ func BenchmarkMemoryPressure(b *testing.B) {
 	})
 
 	b.Run("UDPHighLoad", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout)
 		defer tracker.Close()
 
 		// Generate different IPs

client/firewall/uspfilter/conntrack/icmp.go

@@ -6,8 +6,6 @@ import (
 	"time"
 
 	"github.com/google/gopacket/layers"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 )
 
 const (
@@ -35,7 +33,6 @@ type ICMPConnTrack struct {
 
 // ICMPTracker manages ICMP connection states
 type ICMPTracker struct {
-	logger        *nblog.Logger
 	connections   map[ICMPConnKey]*ICMPConnTrack
 	timeout       time.Duration
 	cleanupTicker *time.Ticker
@@ -45,13 +42,12 @@ type ICMPTracker struct {
 }
 
 // NewICMPTracker creates a new ICMP connection tracker
-func NewICMPTracker(timeout time.Duration, logger *nblog.Logger) *ICMPTracker {
+func NewICMPTracker(timeout time.Duration) *ICMPTracker {
 	if timeout == 0 {
 		timeout = DefaultICMPTimeout
 	}
 
 	tracker := &ICMPTracker{
-		logger:        logger,
 		connections:   make(map[ICMPConnKey]*ICMPConnTrack),
 		timeout:       timeout,
 		cleanupTicker: time.NewTicker(ICMPCleanupInterval),
@@ -66,6 +62,7 @@ func NewICMPTracker(timeout time.Duration, logger *nblog.Logger) *ICMPTracker {
 // TrackOutbound records an outbound ICMP Echo Request
 func (t *ICMPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, id uint16, seq uint16) {
 	key := makeICMPKey(srcIP, dstIP, id, seq)
+	now := time.Now().UnixNano()
 
 	t.mutex.Lock()
 	conn, exists := t.connections[key]
@@ -83,19 +80,24 @@ func (t *ICMPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, id uint16, seq u
 			ID:       id,
 			Sequence: seq,
 		}
-		conn.UpdateLastSeen()
+		conn.lastSeen.Store(now)
+		conn.established.Store(true)
 		t.connections[key] = conn
-
-		t.logger.Trace("New ICMP connection %v", key)
 	}
 	t.mutex.Unlock()
 
-	conn.UpdateLastSeen()
+	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) bool {
-	if icmpType != uint8(layers.ICMPv4TypeEchoReply) {
+	switch icmpType {
+	case uint8(layers.ICMPv4TypeDestinationUnreachable),
+		uint8(layers.ICMPv4TypeTimeExceeded):
+		return true
+	case uint8(layers.ICMPv4TypeEchoReply):
+		// continue processing
+	default:
 		return false
 	}
 
@@ -113,7 +115,8 @@ func (t *ICMPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, id uint16, seq
 		return false
 	}
 
-	return ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
+	return conn.IsEstablished() &&
+		ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
 		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
 		conn.ID == id &&
 		conn.Sequence == seq
@@ -138,8 +141,6 @@ func (t *ICMPTracker) cleanup() {
 			t.ipPool.Put(conn.SourceIP)
 			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
-
-			t.logger.Debug("Removed ICMP connection %v (timeout)", key)
 		}
 	}
 }

client/firewall/uspfilter/conntrack/icmp_test.go

@@ -7,7 +7,7 @@ import (
 
 func BenchmarkICMPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewICMPTracker(DefaultICMPTimeout, logger)
+		tracker := NewICMPTracker(DefaultICMPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")
@@ -20,7 +20,7 @@ func BenchmarkICMPTracker(b *testing.B) {
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewICMPTracker(DefaultICMPTimeout, logger)
+		tracker := NewICMPTracker(DefaultICMPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")

client/firewall/uspfilter/conntrack/tcp.go

@@ -5,10 +5,7 @@ package conntrack
 import (
 	"net"
 	"sync"
-	"sync/atomic"
 	"time"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 )
 
 const (
@@ -64,24 +61,12 @@ type TCPConnKey struct {
 // TCPConnTrack represents a TCP connection state
 type TCPConnTrack struct {
 	BaseConnTrack
-	State       TCPState
-	established atomic.Bool
+	State TCPState
 	sync.RWMutex
 }
 
-// IsEstablished safely checks if connection is established
-func (t *TCPConnTrack) IsEstablished() bool {
-	return t.established.Load()
-}
-
-// SetEstablished safely sets the established state
-func (t *TCPConnTrack) SetEstablished(state bool) {
-	t.established.Store(state)
-}
-
 // TCPTracker manages TCP connection states
 type TCPTracker struct {
-	logger        *nblog.Logger
 	connections   map[ConnKey]*TCPConnTrack
 	mutex         sync.RWMutex
 	cleanupTicker *time.Ticker
@@ -91,9 +76,8 @@ type TCPTracker struct {
 }
 
 // NewTCPTracker creates a new TCP connection tracker
-func NewTCPTracker(timeout time.Duration, logger *nblog.Logger) *TCPTracker {
+func NewTCPTracker(timeout time.Duration) *TCPTracker {
 	tracker := &TCPTracker{
-		logger:        logger,
 		connections:   make(map[ConnKey]*TCPConnTrack),
 		cleanupTicker: time.NewTicker(TCPCleanupInterval),
 		done:          make(chan struct{}),
@@ -109,6 +93,7 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger) *TCPTracker {
 func (t *TCPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, flags uint8) {
 	// Create key before lock
 	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	now := time.Now().UnixNano()
 
 	t.mutex.Lock()
 	conn, exists := t.connections[key]
@@ -128,11 +113,9 @@ func (t *TCPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, d
 			},
 			State: TCPStateNew,
 		}
-		conn.UpdateLastSeen()
+		conn.lastSeen.Store(now)
 		conn.established.Store(false)
 		t.connections[key] = conn
-
-		t.logger.Trace("New TCP connection: %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
 	}
 	t.mutex.Unlock()
 
@@ -140,7 +123,7 @@ func (t *TCPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, d
 	conn.Lock()
 	t.updateState(conn, flags, true)
 	conn.Unlock()
-	conn.UpdateLastSeen()
+	conn.lastSeen.Store(now)
 }
 
 // IsValidInbound checks if an inbound TCP packet matches a tracked connection
@@ -188,9 +171,6 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 	if flags&TCPRst != 0 {
 		conn.State = TCPStateClosed
 		conn.SetEstablished(false)
-
-		t.logger.Trace("TCP connection reset: %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 		return
 	}
 
@@ -247,9 +227,6 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 		if flags&TCPAck != 0 {
 			conn.State = TCPStateTimeWait
 			// Keep established = false from previous state
-
-			t.logger.Trace("TCP connection closed (simultaneous) - %s:%d -> %s:%d",
-				conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 		}
 
 	case TCPStateCloseWait:
@@ -260,17 +237,11 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 	case TCPStateLastAck:
 		if flags&TCPAck != 0 {
 			conn.State = TCPStateClosed
-
-			t.logger.Trace("TCP connection gracefully closed: %s:%d -> %s:%d",
-				conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 		}
 
 	case TCPStateTimeWait:
 		// Stay in TIME-WAIT for 2MSL before transitioning to closed
 		// This is handled by the cleanup routine
-
-		t.logger.Trace("TCP connection completed - %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 	}
 }
 
@@ -347,8 +318,6 @@ func (t *TCPTracker) cleanup() {
 			t.ipPool.Put(conn.SourceIP)
 			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
-
-			t.logger.Trace("Cleaned up TCP connection: %s:%d -> %s:%d", conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 		}
 	}
 }

client/firewall/uspfilter/conntrack/tcp_test.go

@@ -9,7 +9,7 @@ import (
 )
 
 func TestTCPStateMachine(t *testing.T) {
-	tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+	tracker := NewTCPTracker(DefaultTCPTimeout)
 	defer tracker.Close()
 
 	srcIP := net.ParseIP("100.64.0.1")
@@ -154,7 +154,7 @@ func TestTCPStateMachine(t *testing.T) {
 			t.Run(tt.name, func(t *testing.T) {
 				t.Helper()
 
-				tracker = NewTCPTracker(DefaultTCPTimeout, logger)
+				tracker = NewTCPTracker(DefaultTCPTimeout)
 				tt.test(t)
 			})
 		}
@@ -162,7 +162,7 @@ func TestTCPStateMachine(t *testing.T) {
 }
 
 func TestRSTHandling(t *testing.T) {
-	tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+	tracker := NewTCPTracker(DefaultTCPTimeout)
 	defer tracker.Close()
 
 	srcIP := net.ParseIP("100.64.0.1")
@@ -233,7 +233,7 @@ func establishConnection(t *testing.T, tracker *TCPTracker, srcIP, dstIP net.IP,
 
 func BenchmarkTCPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")
@@ -246,7 +246,7 @@ func BenchmarkTCPTracker(b *testing.B) {
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")
@@ -264,7 +264,7 @@ func BenchmarkTCPTracker(b *testing.B) {
 	})
 
 	b.Run("ConcurrentAccess", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")
@@ -287,7 +287,7 @@ func BenchmarkTCPTracker(b *testing.B) {
 // Benchmark connection cleanup
 func BenchmarkCleanup(b *testing.B) {
 	b.Run("TCPCleanup", func(b *testing.B) {
-		tracker := NewTCPTracker(100*time.Millisecond, logger) // Short timeout for testing
+		tracker := NewTCPTracker(100 * time.Millisecond) // Short timeout for testing
 		defer tracker.Close()
 
 		// Pre-populate with expired connections

client/firewall/uspfilter/conntrack/udp.go

@@ -4,8 +4,6 @@ import (
 	"net"
 	"sync"
 	"time"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 )
 
 const (
@@ -22,7 +20,6 @@ type UDPConnTrack struct {
 
 // UDPTracker manages UDP connection states
 type UDPTracker struct {
-	logger        *nblog.Logger
 	connections   map[ConnKey]*UDPConnTrack
 	timeout       time.Duration
 	cleanupTicker *time.Ticker
@@ -32,13 +29,12 @@ type UDPTracker struct {
 }
 
 // NewUDPTracker creates a new UDP connection tracker
-func NewUDPTracker(timeout time.Duration, logger *nblog.Logger) *UDPTracker {
+func NewUDPTracker(timeout time.Duration) *UDPTracker {
 	if timeout == 0 {
 		timeout = DefaultUDPTimeout
 	}
 
 	tracker := &UDPTracker{
-		logger:        logger,
 		connections:   make(map[ConnKey]*UDPConnTrack),
 		timeout:       timeout,
 		cleanupTicker: time.NewTicker(UDPCleanupInterval),
@@ -53,6 +49,7 @@ func NewUDPTracker(timeout time.Duration, logger *nblog.Logger) *UDPTracker {
 // TrackOutbound records an outbound UDP connection
 func (t *UDPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16) {
 	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	now := time.Now().UnixNano()
 
 	t.mutex.Lock()
 	conn, exists := t.connections[key]
@@ -70,14 +67,13 @@ func (t *UDPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, d
 				DestPort:   dstPort,
 			},
 		}
-		conn.UpdateLastSeen()
+		conn.lastSeen.Store(now)
+		conn.established.Store(true)
 		t.connections[key] = conn
-
-		t.logger.Trace("New UDP connection: %v", conn)
 	}
 	t.mutex.Unlock()
 
-	conn.UpdateLastSeen()
+	conn.lastSeen.Store(now)
 }
 
 // IsValidInbound checks if an inbound packet matches a tracked connection
@@ -96,7 +92,8 @@ func (t *UDPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
 		return false
 	}
 
-	return ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
+	return conn.IsEstablished() &&
+		ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
 		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
 		conn.DestPort == srcPort &&
 		conn.SourcePort == dstPort
@@ -123,8 +120,6 @@ func (t *UDPTracker) cleanup() {
 			t.ipPool.Put(conn.SourceIP)
 			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
-
-			t.logger.Trace("Removed UDP connection %v (timeout)", conn)
 		}
 	}
 }

client/firewall/uspfilter/conntrack/udp_test.go

@@ -29,7 +29,7 @@ func TestNewUDPTracker(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tracker := NewUDPTracker(tt.timeout, logger)
+			tracker := NewUDPTracker(tt.timeout)
 			assert.NotNil(t, tracker)
 			assert.Equal(t, tt.wantTimeout, tracker.timeout)
 			assert.NotNil(t, tracker.connections)
@@ -40,7 +40,7 @@ func TestNewUDPTracker(t *testing.T) {
 }
 
 func TestUDPTracker_TrackOutbound(t *testing.T) {
-	tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+	tracker := NewUDPTracker(DefaultUDPTimeout)
 	defer tracker.Close()
 
 	srcIP := net.ParseIP("192.168.1.2")
@@ -58,11 +58,12 @@ func TestUDPTracker_TrackOutbound(t *testing.T) {
 	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, logger)
+	tracker := NewUDPTracker(1 * time.Second)
 	defer tracker.Close()
 
 	srcIP := net.ParseIP("192.168.1.2")
@@ -161,7 +162,6 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 		cleanupTicker: time.NewTicker(cleanupInterval),
 		done:          make(chan struct{}),
 		ipPool:        NewPreallocatedIPs(),
-		logger:        logger,
 	}
 
 	// Start cleanup routine
@@ -211,7 +211,7 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 
 func BenchmarkUDPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")
@@ -224,7 +224,7 @@ func BenchmarkUDPTracker(b *testing.B) {
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout)
 		defer tracker.Close()
 
 		srcIP := net.ParseIP("192.168.1.1")

client/firewall/uspfilter/forwarder/endpoint.go

@@ -1,81 +0,0 @@
-package forwarder
-
-import (
-	wgdevice "golang.zx2c4.com/wireguard/device"
-	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/header"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
-)
-
-// endpoint implements stack.LinkEndpoint and handles integration with the wireguard device
-type endpoint struct {
-	logger     *nblog.Logger
-	dispatcher stack.NetworkDispatcher
-	device     *wgdevice.Device
-	mtu        uint32
-}
-
-func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
-	e.dispatcher = dispatcher
-}
-
-func (e *endpoint) IsAttached() bool {
-	return e.dispatcher != nil
-}
-
-func (e *endpoint) MTU() uint32 {
-	return e.mtu
-}
-
-func (e *endpoint) Capabilities() stack.LinkEndpointCapabilities {
-	return stack.CapabilityNone
-}
-
-func (e *endpoint) MaxHeaderLength() uint16 {
-	return 0
-}
-
-func (e *endpoint) LinkAddress() tcpip.LinkAddress {
-	return ""
-}
-
-func (e *endpoint) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
-	var written int
-	for _, pkt := range pkts.AsSlice() {
-		netHeader := header.IPv4(pkt.NetworkHeader().View().AsSlice())
-
-		data := stack.PayloadSince(pkt.NetworkHeader())
-		if data == nil {
-			continue
-		}
-
-		// Send the packet through WireGuard
-		address := netHeader.DestinationAddress()
-		err := e.device.CreateOutboundPacket(data.AsSlice(), address.AsSlice())
-		if err != nil {
-			e.logger.Error("CreateOutboundPacket: %v", err)
-			continue
-		}
-		written++
-	}
-
-	return written, nil
-}
-
-func (e *endpoint) Wait() {
-	// not required
-}
-
-func (e *endpoint) ARPHardwareType() header.ARPHardwareType {
-	return header.ARPHardwareNone
-}
-
-func (e *endpoint) AddHeader(*stack.PacketBuffer) {
-	// not required
-}
-
-func (e *endpoint) ParseHeader(*stack.PacketBuffer) bool {
-	return true
-}

client/firewall/uspfilter/forwarder/forwarder.go

@@ -1,166 +0,0 @@
-package forwarder
-
-import (
-	"context"
-	"fmt"
-	"net"
-	"runtime"
-
-	log "github.com/sirupsen/logrus"
-	"gvisor.dev/gvisor/pkg/buffer"
-	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/header"
-	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/icmp"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
-
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/common"
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
-)
-
-const (
-	defaultReceiveWindow = 32768
-	defaultMaxInFlight   = 1024
-	iosReceiveWindow     = 16384
-	iosMaxInFlight       = 256
-)
-
-type Forwarder struct {
-	logger       *nblog.Logger
-	stack        *stack.Stack
-	endpoint     *endpoint
-	udpForwarder *udpForwarder
-	ctx          context.Context
-	cancel       context.CancelFunc
-	ip           net.IP
-	netstack     bool
-}
-
-func New(iface common.IFaceMapper, logger *nblog.Logger, netstack bool) (*Forwarder, error) {
-	s := stack.New(stack.Options{
-		NetworkProtocols: []stack.NetworkProtocolFactory{ipv4.NewProtocol},
-		TransportProtocols: []stack.TransportProtocolFactory{
-			tcp.NewProtocol,
-			udp.NewProtocol,
-			icmp.NewProtocol4,
-		},
-		HandleLocal: false,
-	})
-
-	mtu, err := iface.GetDevice().MTU()
-	if err != nil {
-		return nil, fmt.Errorf("get MTU: %w", err)
-	}
-	nicID := tcpip.NICID(1)
-	endpoint := &endpoint{
-		logger: logger,
-		device: iface.GetWGDevice(),
-		mtu:    uint32(mtu),
-	}
-
-	if err := s.CreateNIC(nicID, endpoint); err != nil {
-		return nil, fmt.Errorf("failed to create NIC: %v", err)
-	}
-
-	ones, _ := iface.Address().Network.Mask.Size()
-	protoAddr := tcpip.ProtocolAddress{
-		Protocol: ipv4.ProtocolNumber,
-		AddressWithPrefix: tcpip.AddressWithPrefix{
-			Address:   tcpip.AddrFromSlice(iface.Address().IP.To4()),
-			PrefixLen: ones,
-		},
-	}
-
-	if err := s.AddProtocolAddress(nicID, protoAddr, stack.AddressProperties{}); err != nil {
-		return nil, fmt.Errorf("failed to add protocol address: %s", err)
-	}
-
-	defaultSubnet, err := tcpip.NewSubnet(
-		tcpip.AddrFrom4([4]byte{0, 0, 0, 0}),
-		tcpip.MaskFromBytes([]byte{0, 0, 0, 0}),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("creating default subnet: %w", err)
-	}
-
-	if err := s.SetPromiscuousMode(nicID, true); err != nil {
-		return nil, fmt.Errorf("set promiscuous mode: %s", err)
-	}
-	if err := s.SetSpoofing(nicID, true); err != nil {
-		return nil, fmt.Errorf("set spoofing: %s", err)
-	}
-
-	s.SetRouteTable([]tcpip.Route{
-		{
-			Destination: defaultSubnet,
-			NIC:         nicID,
-		},
-	})
-
-	ctx, cancel := context.WithCancel(context.Background())
-	f := &Forwarder{
-		logger:       logger,
-		stack:        s,
-		endpoint:     endpoint,
-		udpForwarder: newUDPForwarder(mtu, logger),
-		ctx:          ctx,
-		cancel:       cancel,
-		netstack:     netstack,
-		ip:           iface.Address().IP,
-	}
-
-	receiveWindow := defaultReceiveWindow
-	maxInFlight := defaultMaxInFlight
-	if runtime.GOOS == "ios" {
-		receiveWindow = iosReceiveWindow
-		maxInFlight = iosMaxInFlight
-	}
-
-	tcpForwarder := tcp.NewForwarder(s, receiveWindow, maxInFlight, f.handleTCP)
-	s.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
-
-	udpForwarder := udp.NewForwarder(s, f.handleUDP)
-	s.SetTransportProtocolHandler(udp.ProtocolNumber, udpForwarder.HandlePacket)
-
-	s.SetTransportProtocolHandler(icmp.ProtocolNumber4, f.handleICMP)
-
-	log.Debugf("forwarder: Initialization complete with NIC %d", nicID)
-	return f, nil
-}
-
-func (f *Forwarder) InjectIncomingPacket(payload []byte) error {
-	if len(payload) < header.IPv4MinimumSize {
-		return fmt.Errorf("packet too small: %d bytes", len(payload))
-	}
-
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		Payload: buffer.MakeWithData(payload),
-	})
-	defer pkt.DecRef()
-
-	if f.endpoint.dispatcher != nil {
-		f.endpoint.dispatcher.DeliverNetworkPacket(ipv4.ProtocolNumber, pkt)
-	}
-	return nil
-}
-
-// Stop gracefully shuts down the forwarder
-func (f *Forwarder) Stop() {
-	f.cancel()
-
-	if f.udpForwarder != nil {
-		f.udpForwarder.Stop()
-	}
-
-	f.stack.Close()
-	f.stack.Wait()
-}
-
-func (f *Forwarder) determineDialAddr(addr tcpip.Address) net.IP {
-	if f.netstack && f.ip.Equal(addr.AsSlice()) {
-		return net.IPv4(127, 0, 0, 1)
-	}
-	return addr.AsSlice()
-}

client/firewall/uspfilter/forwarder/icmp.go

@@ -1,109 +0,0 @@
-package forwarder
-
-import (
-	"context"
-	"net"
-	"time"
-
-	"gvisor.dev/gvisor/pkg/tcpip/header"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-)
-
-// handleICMP handles ICMP packets from the network stack
-func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBufferPtr) bool {
-	ctx, cancel := context.WithTimeout(f.ctx, 5*time.Second)
-	defer cancel()
-
-	lc := net.ListenConfig{}
-	// TODO: support non-root
-	conn, err := lc.ListenPacket(ctx, "ip4:icmp", "0.0.0.0")
-	if err != nil {
-		f.logger.Error("Failed to create ICMP socket for %v: %v", id, err)
-
-		// This will make netstack reply on behalf of the original destination, that's ok for now
-		return false
-	}
-	defer func() {
-		if err := conn.Close(); err != nil {
-			f.logger.Debug("Failed to close ICMP socket: %v", err)
-		}
-	}()
-
-	dstIP := f.determineDialAddr(id.LocalAddress)
-	dst := &net.IPAddr{IP: dstIP}
-
-	// Get the complete ICMP message (header + data)
-	fullPacket := stack.PayloadSince(pkt.TransportHeader())
-	payload := fullPacket.AsSlice()
-
-	icmpHdr := header.ICMPv4(pkt.TransportHeader().View().AsSlice())
-
-	// For Echo Requests, send and handle response
-	switch icmpHdr.Type() {
-	case header.ICMPv4Echo:
-		return f.handleEchoResponse(icmpHdr, payload, dst, conn, id)
-	case header.ICMPv4EchoReply:
-		// dont process our own replies
-		return true
-	default:
-	}
-
-	// For other ICMP types (Time Exceeded, Destination Unreachable, etc)
-	_, err = conn.WriteTo(payload, dst)
-	if err != nil {
-		f.logger.Error("Failed to write ICMP packet for %v: %v", id, err)
-		return true
-	}
-
-	f.logger.Trace("Forwarded ICMP packet %v type=%v code=%v",
-		id, icmpHdr.Type(), icmpHdr.Code())
-
-	return true
-}
-
-func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, payload []byte, dst *net.IPAddr, conn net.PacketConn, id stack.TransportEndpointID) bool {
-	if _, err := conn.WriteTo(payload, dst); err != nil {
-		f.logger.Error("Failed to write ICMP packet for %v: %v", id, err)
-		return true
-	}
-
-	f.logger.Trace("Forwarded ICMP packet %v type=%v code=%v",
-		id, icmpHdr.Type(), icmpHdr.Code())
-
-	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
-		f.logger.Error("Failed to set read deadline for ICMP response: %v", err)
-		return true
-	}
-
-	response := make([]byte, f.endpoint.mtu)
-	n, _, err := conn.ReadFrom(response)
-	if err != nil {
-		if !isTimeout(err) {
-			f.logger.Error("Failed to read ICMP response: %v", err)
-		}
-		return true
-	}
-
-	ipHdr := make([]byte, header.IPv4MinimumSize)
-	ip := header.IPv4(ipHdr)
-	ip.Encode(&header.IPv4Fields{
-		TotalLength: uint16(header.IPv4MinimumSize + n),
-		TTL:         64,
-		Protocol:    uint8(header.ICMPv4ProtocolNumber),
-		SrcAddr:     id.LocalAddress,
-		DstAddr:     id.RemoteAddress,
-	})
-	ip.SetChecksum(^ip.CalculateChecksum())
-
-	fullPacket := make([]byte, 0, len(ipHdr)+n)
-	fullPacket = append(fullPacket, ipHdr...)
-	fullPacket = append(fullPacket, response[:n]...)
-
-	if err := f.InjectIncomingPacket(fullPacket); err != nil {
-		f.logger.Error("Failed to inject ICMP response: %v", err)
-		return true
-	}
-
-	f.logger.Trace("Forwarded ICMP echo reply for %v", id)
-	return true
-}

client/firewall/uspfilter/forwarder/tcp.go

@@ -1,90 +0,0 @@
-package forwarder
-
-import (
-	"context"
-	"fmt"
-	"io"
-	"net"
-
-	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
-	"gvisor.dev/gvisor/pkg/waiter"
-)
-
-// handleTCP is called by the TCP forwarder for new connections.
-func (f *Forwarder) handleTCP(r *tcp.ForwarderRequest) {
-	id := r.ID()
-
-	dialAddr := fmt.Sprintf("%s:%d", f.determineDialAddr(id.LocalAddress), id.LocalPort)
-
-	outConn, err := (&net.Dialer{}).DialContext(f.ctx, "tcp", dialAddr)
-	if err != nil {
-		r.Complete(true)
-		f.logger.Trace("forwarder: dial error for %v: %v", id, err)
-		return
-	}
-
-	// Create wait queue for blocking syscalls
-	wq := waiter.Queue{}
-
-	ep, epErr := r.CreateEndpoint(&wq)
-	if epErr != nil {
-		f.logger.Error("forwarder: failed to create TCP endpoint: %v", epErr)
-		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: outConn close error: %v", err)
-		}
-		r.Complete(true)
-		return
-	}
-
-	// Complete the handshake
-	r.Complete(false)
-
-	inConn := gonet.NewTCPConn(&wq, ep)
-
-	f.logger.Trace("forwarder: established TCP connection %v", id)
-
-	go f.proxyTCP(id, inConn, outConn, ep)
-}
-
-func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn, outConn net.Conn, ep tcpip.Endpoint) {
-	defer func() {
-		if err := inConn.Close(); err != nil {
-			f.logger.Debug("forwarder: inConn close error: %v", err)
-		}
-		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: outConn close error: %v", err)
-		}
-		ep.Close()
-	}()
-
-	// Create context for managing the proxy goroutines
-	ctx, cancel := context.WithCancel(f.ctx)
-	defer cancel()
-
-	errChan := make(chan error, 2)
-
-	go func() {
-		_, err := io.Copy(outConn, inConn)
-		errChan <- err
-	}()
-
-	go func() {
-		_, err := io.Copy(inConn, outConn)
-		errChan <- err
-	}()
-
-	select {
-	case <-ctx.Done():
-		f.logger.Trace("forwarder: tearing down TCP connection %v due to context done", id)
-		return
-	case err := <-errChan:
-		if err != nil && !isClosedError(err) {
-			f.logger.Error("proxyTCP: copy error: %v", err)
-		}
-		f.logger.Trace("forwarder: tearing down TCP connection %v", id)
-		return
-	}
-}

client/firewall/uspfilter/forwarder/udp.go

@@ -1,288 +0,0 @@
-package forwarder
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"net"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
-	"gvisor.dev/gvisor/pkg/waiter"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
-)
-
-const (
-	udpTimeout = 30 * time.Second
-)
-
-type udpPacketConn struct {
-	conn     *gonet.UDPConn
-	outConn  net.Conn
-	lastSeen atomic.Int64
-	cancel   context.CancelFunc
-	ep       tcpip.Endpoint
-}
-
-type udpForwarder struct {
-	sync.RWMutex
-	logger  *nblog.Logger
-	conns   map[stack.TransportEndpointID]*udpPacketConn
-	bufPool sync.Pool
-	ctx     context.Context
-	cancel  context.CancelFunc
-}
-
-type idleConn struct {
-	id   stack.TransportEndpointID
-	conn *udpPacketConn
-}
-
-func newUDPForwarder(mtu int, logger *nblog.Logger) *udpForwarder {
-	ctx, cancel := context.WithCancel(context.Background())
-	f := &udpForwarder{
-		logger: logger,
-		conns:  make(map[stack.TransportEndpointID]*udpPacketConn),
-		ctx:    ctx,
-		cancel: cancel,
-		bufPool: sync.Pool{
-			New: func() any {
-				b := make([]byte, mtu)
-				return &b
-			},
-		},
-	}
-	go f.cleanup()
-	return f
-}
-
-// Stop stops the UDP forwarder and all active connections
-func (f *udpForwarder) Stop() {
-	f.cancel()
-
-	f.Lock()
-	defer f.Unlock()
-
-	for id, conn := range f.conns {
-		conn.cancel()
-		if err := conn.conn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP conn close error for %v: %v", id, err)
-		}
-		if err := conn.outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
-		}
-
-		conn.ep.Close()
-		delete(f.conns, id)
-	}
-}
-
-// cleanup periodically removes idle UDP connections
-func (f *udpForwarder) cleanup() {
-	ticker := time.NewTicker(time.Minute)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-f.ctx.Done():
-			return
-		case <-ticker.C:
-			var idleConns []idleConn
-
-			f.RLock()
-			for id, conn := range f.conns {
-				if conn.getIdleDuration() > udpTimeout {
-					idleConns = append(idleConns, idleConn{id, conn})
-				}
-			}
-			f.RUnlock()
-
-			for _, idle := range idleConns {
-				idle.conn.cancel()
-				if err := idle.conn.conn.Close(); err != nil {
-					f.logger.Debug("forwarder: UDP conn close error for %v: %v", idle.id, err)
-				}
-				if err := idle.conn.outConn.Close(); err != nil {
-					f.logger.Debug("forwarder: UDP outConn close error for %v: %v", idle.id, err)
-				}
-
-				idle.conn.ep.Close()
-
-				f.Lock()
-				delete(f.conns, idle.id)
-				f.Unlock()
-
-				f.logger.Trace("forwarder: cleaned up idle UDP connection %v", idle.id)
-			}
-		}
-	}
-}
-
-// handleUDP is called by the UDP forwarder for new packets
-func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
-	if f.ctx.Err() != nil {
-		f.logger.Trace("forwarder: context done, dropping UDP packet")
-		return
-	}
-
-	id := r.ID()
-
-	f.udpForwarder.RLock()
-	_, exists := f.udpForwarder.conns[id]
-	f.udpForwarder.RUnlock()
-	if exists {
-		f.logger.Trace("forwarder: existing UDP connection for %v", id)
-		return
-	}
-
-	dstAddr := fmt.Sprintf("%s:%d", f.determineDialAddr(id.LocalAddress), id.LocalPort)
-	outConn, err := (&net.Dialer{}).DialContext(f.ctx, "udp", dstAddr)
-	if err != nil {
-		f.logger.Debug("forwarder: UDP dial error for %v: %v", id, err)
-		// TODO: Send ICMP error message
-		return
-	}
-
-	// Create wait queue for blocking syscalls
-	wq := waiter.Queue{}
-	ep, epErr := r.CreateEndpoint(&wq)
-	if epErr != nil {
-		f.logger.Debug("forwarder: failed to create UDP endpoint: %v", epErr)
-		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
-		}
-		return
-	}
-
-	inConn := gonet.NewUDPConn(f.stack, &wq, ep)
-	connCtx, connCancel := context.WithCancel(f.ctx)
-
-	pConn := &udpPacketConn{
-		conn:    inConn,
-		outConn: outConn,
-		cancel:  connCancel,
-		ep:      ep,
-	}
-	pConn.updateLastSeen()
-
-	f.udpForwarder.Lock()
-	// Double-check no connection was created while we were setting up
-	if _, exists := f.udpForwarder.conns[id]; exists {
-		f.udpForwarder.Unlock()
-		pConn.cancel()
-		if err := inConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", id, err)
-		}
-		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
-		}
-		return
-	}
-	f.udpForwarder.conns[id] = pConn
-	f.udpForwarder.Unlock()
-
-	f.logger.Trace("forwarder: established UDP connection to %v", id)
-	go f.proxyUDP(connCtx, pConn, id, ep)
-}
-
-func (f *Forwarder) proxyUDP(ctx context.Context, pConn *udpPacketConn, id stack.TransportEndpointID, ep tcpip.Endpoint) {
-	defer func() {
-		pConn.cancel()
-		if err := pConn.conn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", id, err)
-		}
-		if err := pConn.outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
-		}
-
-		ep.Close()
-
-		f.udpForwarder.Lock()
-		delete(f.udpForwarder.conns, id)
-		f.udpForwarder.Unlock()
-	}()
-
-	errChan := make(chan error, 2)
-
-	go func() {
-		errChan <- pConn.copy(ctx, pConn.conn, pConn.outConn, &f.udpForwarder.bufPool, "outbound->inbound")
-	}()
-
-	go func() {
-		errChan <- pConn.copy(ctx, pConn.outConn, pConn.conn, &f.udpForwarder.bufPool, "inbound->outbound")
-	}()
-
-	select {
-	case <-ctx.Done():
-		f.logger.Trace("forwarder: tearing down UDP connection %v due to context done", id)
-		return
-	case err := <-errChan:
-		if err != nil && !isClosedError(err) {
-			f.logger.Error("proxyUDP: copy error: %v", err)
-		}
-		f.logger.Trace("forwarder: tearing down UDP connection %v", id)
-		return
-	}
-}
-
-func (c *udpPacketConn) updateLastSeen() {
-	c.lastSeen.Store(time.Now().UnixNano())
-}
-
-func (c *udpPacketConn) getIdleDuration() time.Duration {
-	lastSeen := time.Unix(0, c.lastSeen.Load())
-	return time.Since(lastSeen)
-}
-
-func (c *udpPacketConn) copy(ctx context.Context, dst net.Conn, src net.Conn, bufPool *sync.Pool, direction string) error {
-	bufp := bufPool.Get().(*[]byte)
-	defer bufPool.Put(bufp)
-	buffer := *bufp
-
-	if err := src.SetReadDeadline(time.Now().Add(udpTimeout)); err != nil {
-		return fmt.Errorf("set read deadline: %w", err)
-	}
-	if err := src.SetWriteDeadline(time.Now().Add(udpTimeout)); err != nil {
-		return fmt.Errorf("set write deadline: %w", err)
-	}
-
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
-			n, err := src.Read(buffer)
-			if err != nil {
-				if isTimeout(err) {
-					continue
-				}
-				return fmt.Errorf("read from %s: %w", direction, err)
-			}
-
-			_, err = dst.Write(buffer[:n])
-			if err != nil {
-				return fmt.Errorf("write to %s: %w", direction, err)
-			}
-
-			c.updateLastSeen()
-		}
-	}
-}
-
-func isClosedError(err error) bool {
-	return errors.Is(err, net.ErrClosed) || errors.Is(err, context.Canceled)
-}
-
-func isTimeout(err error) bool {
-	var netErr net.Error
-	if errors.As(err, &netErr) {
-		return netErr.Timeout()
-	}
-	return false
-}

client/firewall/uspfilter/localip.go

@@ -1,134 +0,0 @@
-package uspfilter
-
-import (
-	"fmt"
-	"net"
-	"sync"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/common"
-)
-
-type localIPManager struct {
-	mu sync.RWMutex
-
-	// Use bitmap for IPv4 (32 bits * 2^16 = 256KB memory)
-	ipv4Bitmap [1 << 16]uint32
-}
-
-func newLocalIPManager() *localIPManager {
-	return &localIPManager{}
-}
-
-func (m *localIPManager) setBitmapBit(ip net.IP) {
-	ipv4 := ip.To4()
-	if ipv4 == nil {
-		return
-	}
-	high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-	low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-	m.ipv4Bitmap[high] |= 1 << (low % 32)
-}
-
-func (m *localIPManager) checkBitmapBit(ip net.IP) bool {
-	ipv4 := ip.To4()
-	if ipv4 == nil {
-		return false
-	}
-	high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-	low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-	return (m.ipv4Bitmap[high] & (1 << (low % 32))) != 0
-}
-
-func (m *localIPManager) processIP(ip net.IP, newIPv4Bitmap *[1 << 16]uint32, ipv4Set map[string]struct{}, ipv4Addresses *[]string) error {
-	if ipv4 := ip.To4(); ipv4 != nil {
-		high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-		low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-		if int(high) >= len(*newIPv4Bitmap) {
-			return fmt.Errorf("invalid IPv4 address: %s", ip)
-		}
-		ipStr := ip.String()
-		if _, exists := ipv4Set[ipStr]; !exists {
-			ipv4Set[ipStr] = struct{}{}
-			*ipv4Addresses = append(*ipv4Addresses, ipStr)
-			newIPv4Bitmap[high] |= 1 << (low % 32)
-		}
-	}
-	return nil
-}
-
-func (m *localIPManager) processInterface(iface net.Interface, newIPv4Bitmap *[1 << 16]uint32, ipv4Set map[string]struct{}, ipv4Addresses *[]string) {
-	addrs, err := iface.Addrs()
-	if err != nil {
-		log.Debugf("get addresses for interface %s failed: %v", iface.Name, err)
-		return
-	}
-
-	for _, addr := range addrs {
-		var ip net.IP
-		switch v := addr.(type) {
-		case *net.IPNet:
-			ip = v.IP
-		case *net.IPAddr:
-			ip = v.IP
-		default:
-			continue
-		}
-
-		if err := m.processIP(ip, newIPv4Bitmap, ipv4Set, ipv4Addresses); err != nil {
-			log.Debugf("process IP failed: %v", err)
-		}
-	}
-}
-
-func (m *localIPManager) UpdateLocalIPs(iface common.IFaceMapper) (err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			err = fmt.Errorf("panic: %v", r)
-		}
-	}()
-
-	var newIPv4Bitmap [1 << 16]uint32
-	ipv4Set := make(map[string]struct{})
-	var ipv4Addresses []string
-
-	// 127.0.0.0/8
-	high := uint16(127) << 8
-	for i := uint16(0); i < 256; i++ {
-		newIPv4Bitmap[high|i] = 0xffffffff
-	}
-
-	if iface != nil {
-		if err := m.processIP(iface.Address().IP, &newIPv4Bitmap, ipv4Set, &ipv4Addresses); err != nil {
-			return err
-		}
-	}
-
-	interfaces, err := net.Interfaces()
-	if err != nil {
-		log.Warnf("failed to get interfaces: %v", err)
-	} else {
-		for _, intf := range interfaces {
-			m.processInterface(intf, &newIPv4Bitmap, ipv4Set, &ipv4Addresses)
-		}
-	}
-
-	m.mu.Lock()
-	m.ipv4Bitmap = newIPv4Bitmap
-	m.mu.Unlock()
-
-	log.Debugf("Local IPv4 addresses: %v", ipv4Addresses)
-	return nil
-}
-
-func (m *localIPManager) IsLocalIP(ip net.IP) bool {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	if ipv4 := ip.To4(); ipv4 != nil {
-		return m.checkBitmapBit(ipv4)
-	}
-
-	return false
-}

client/firewall/uspfilter/localip_test.go

@@ -1,270 +0,0 @@
-package uspfilter
-
-import (
-	"net"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/iface"
-)
-
-func TestLocalIPManager(t *testing.T) {
-	tests := []struct {
-		name      string
-		setupAddr iface.WGAddress
-		testIP    net.IP
-		expected  bool
-	}{
-		{
-			name: "Localhost range",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("192.168.1.1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("192.168.1.0"),
-					Mask: net.CIDRMask(24, 32),
-				},
-			},
-			testIP:   net.ParseIP("127.0.0.2"),
-			expected: true,
-		},
-		{
-			name: "Localhost standard address",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("192.168.1.1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("192.168.1.0"),
-					Mask: net.CIDRMask(24, 32),
-				},
-			},
-			testIP:   net.ParseIP("127.0.0.1"),
-			expected: true,
-		},
-		{
-			name: "Localhost range edge",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("192.168.1.1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("192.168.1.0"),
-					Mask: net.CIDRMask(24, 32),
-				},
-			},
-			testIP:   net.ParseIP("127.255.255.255"),
-			expected: true,
-		},
-		{
-			name: "Local IP matches",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("192.168.1.1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("192.168.1.0"),
-					Mask: net.CIDRMask(24, 32),
-				},
-			},
-			testIP:   net.ParseIP("192.168.1.1"),
-			expected: true,
-		},
-		{
-			name: "Local IP doesn't match",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("192.168.1.1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("192.168.1.0"),
-					Mask: net.CIDRMask(24, 32),
-				},
-			},
-			testIP:   net.ParseIP("192.168.1.2"),
-			expected: false,
-		},
-		{
-			name: "IPv6 address",
-			setupAddr: iface.WGAddress{
-				IP: net.ParseIP("fe80::1"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("fe80::"),
-					Mask: net.CIDRMask(64, 128),
-				},
-			},
-			testIP:   net.ParseIP("fe80::1"),
-			expected: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			manager := newLocalIPManager()
-
-			mock := &IFaceMock{
-				AddressFunc: func() iface.WGAddress {
-					return tt.setupAddr
-				},
-			}
-
-			err := manager.UpdateLocalIPs(mock)
-			require.NoError(t, err)
-
-			result := manager.IsLocalIP(tt.testIP)
-			require.Equal(t, tt.expected, result)
-		})
-	}
-}
-
-func TestLocalIPManager_AllInterfaces(t *testing.T) {
-	manager := newLocalIPManager()
-	mock := &IFaceMock{}
-
-	// Get actual local interfaces
-	interfaces, err := net.Interfaces()
-	require.NoError(t, err)
-
-	var tests []struct {
-		ip       string
-		expected bool
-	}
-
-	// Add all local interface IPs to test cases
-	for _, iface := range interfaces {
-		addrs, err := iface.Addrs()
-		require.NoError(t, err)
-
-		for _, addr := range addrs {
-			var ip net.IP
-			switch v := addr.(type) {
-			case *net.IPNet:
-				ip = v.IP
-			case *net.IPAddr:
-				ip = v.IP
-			default:
-				continue
-			}
-
-			if ip4 := ip.To4(); ip4 != nil {
-				tests = append(tests, struct {
-					ip       string
-					expected bool
-				}{
-					ip:       ip4.String(),
-					expected: true,
-				})
-			}
-		}
-	}
-
-	// Add some external IPs as negative test cases
-	externalIPs := []string{
-		"8.8.8.8",
-		"1.1.1.1",
-		"208.67.222.222",
-	}
-	for _, ip := range externalIPs {
-		tests = append(tests, struct {
-			ip       string
-			expected bool
-		}{
-			ip:       ip,
-			expected: false,
-		})
-	}
-
-	require.NotEmpty(t, tests, "No test cases generated")
-
-	err = manager.UpdateLocalIPs(mock)
-	require.NoError(t, err)
-
-	t.Logf("Testing %d IPs", len(tests))
-	for _, tt := range tests {
-		t.Run(tt.ip, func(t *testing.T) {
-			result := manager.IsLocalIP(net.ParseIP(tt.ip))
-			require.Equal(t, tt.expected, result, "IP: %s", tt.ip)
-		})
-	}
-}
-
-// MapImplementation is a version using map[string]struct{}
-type MapImplementation struct {
-	localIPs map[string]struct{}
-}
-
-func BenchmarkIPChecks(b *testing.B) {
-	interfaces := make([]net.IP, 16)
-	for i := range interfaces {
-		interfaces[i] = net.IPv4(10, 0, byte(i>>8), byte(i))
-	}
-
-	// Setup bitmap version
-	bitmapManager := &localIPManager{
-		ipv4Bitmap: [1 << 16]uint32{},
-	}
-	for _, ip := range interfaces[:8] { // Add half of IPs
-		bitmapManager.setBitmapBit(ip)
-	}
-
-	// Setup map version
-	mapManager := &MapImplementation{
-		localIPs: make(map[string]struct{}),
-	}
-	for _, ip := range interfaces[:8] {
-		mapManager.localIPs[ip.String()] = struct{}{}
-	}
-
-	b.Run("Bitmap_Hit", func(b *testing.B) {
-		ip := interfaces[4]
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			bitmapManager.checkBitmapBit(ip)
-		}
-	})
-
-	b.Run("Bitmap_Miss", func(b *testing.B) {
-		ip := interfaces[12]
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			bitmapManager.checkBitmapBit(ip)
-		}
-	})
-
-	b.Run("Map_Hit", func(b *testing.B) {
-		ip := interfaces[4]
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			// nolint:gosimple
-			_, _ = mapManager.localIPs[ip.String()]
-		}
-	})
-
-	b.Run("Map_Miss", func(b *testing.B) {
-		ip := interfaces[12]
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			// nolint:gosimple
-			_, _ = mapManager.localIPs[ip.String()]
-		}
-	})
-}
-
-func BenchmarkWGPosition(b *testing.B) {
-	wgIP := net.ParseIP("10.10.0.1")
-
-	// Create two managers - one checks WG IP first, other checks it last
-	b.Run("WG_First", func(b *testing.B) {
-		bm := &localIPManager{ipv4Bitmap: [1 << 16]uint32{}}
-		bm.setBitmapBit(wgIP)
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			bm.checkBitmapBit(wgIP)
-		}
-	})
-
-	b.Run("WG_Last", func(b *testing.B) {
-		bm := &localIPManager{ipv4Bitmap: [1 << 16]uint32{}}
-		// Fill with other IPs first
-		for i := 0; i < 15; i++ {
-			bm.setBitmapBit(net.IPv4(10, 0, byte(i>>8), byte(i)))
-		}
-		bm.setBitmapBit(wgIP) // Add WG IP last
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			bm.checkBitmapBit(wgIP)
-		}
-	})
-}

client/firewall/uspfilter/log/log.go

@@ -1,196 +0,0 @@
-// Package logger provides a high-performance, non-blocking logger for userspace networking
-package log
-
-import (
-	"context"
-	"fmt"
-	"io"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	maxBatchSize         = 1024 * 16  // 16KB max batch size
-	maxMessageSize       = 1024 * 2   // 2KB per message
-	bufferSize           = 1024 * 256 // 256KB ring buffer
-	defaultFlushInterval = 2 * time.Second
-)
-
-// Level represents log severity
-type Level uint32
-
-const (
-	LevelPanic Level = iota
-	LevelFatal
-	LevelError
-	LevelWarn
-	LevelInfo
-	LevelDebug
-	LevelTrace
-)
-
-var levelStrings = map[Level]string{
-	LevelPanic: "PANC",
-	LevelFatal: "FATL",
-	LevelError: "ERRO",
-	LevelWarn:  "WARN",
-	LevelInfo:  "INFO",
-	LevelDebug: "DEBG",
-	LevelTrace: "TRAC",
-}
-
-// Logger is a high-performance, non-blocking logger
-type Logger struct {
-	output    io.Writer
-	level     atomic.Uint32
-	buffer    *ringBuffer
-	shutdown  chan struct{}
-	closeOnce sync.Once
-	wg        sync.WaitGroup
-
-	// Reusable buffer pool for formatting messages
-	bufPool sync.Pool
-}
-
-func NewFromLogrus(logrusLogger *log.Logger) *Logger {
-	l := &Logger{
-		output:   logrusLogger.Out,
-		buffer:   newRingBuffer(bufferSize),
-		shutdown: make(chan struct{}),
-		bufPool: sync.Pool{
-			New: func() interface{} {
-				// Pre-allocate buffer for message formatting
-				b := make([]byte, 0, maxMessageSize)
-				return &b
-			},
-		},
-	}
-	logrusLevel := logrusLogger.GetLevel()
-	l.level.Store(uint32(logrusLevel))
-	level := levelStrings[Level(logrusLevel)]
-	log.Debugf("New uspfilter logger created with loglevel %v", level)
-
-	l.wg.Add(1)
-	go l.worker()
-
-	return l
-}
-
-func (l *Logger) SetLevel(level Level) {
-	l.level.Store(uint32(level))
-
-	log.Debugf("Set uspfilter logger loglevel to %v", levelStrings[level])
-}
-
-func (l *Logger) formatMessage(buf *[]byte, level Level, format string, args ...interface{}) {
-	*buf = (*buf)[:0]
-
-	// Timestamp
-	*buf = time.Now().AppendFormat(*buf, "2006-01-02T15:04:05-07:00")
-	*buf = append(*buf, ' ')
-
-	// Level
-	*buf = append(*buf, levelStrings[level]...)
-	*buf = append(*buf, ' ')
-
-	// Message
-	if len(args) > 0 {
-		*buf = append(*buf, fmt.Sprintf(format, args...)...)
-	} else {
-		*buf = append(*buf, format...)
-	}
-
-	*buf = append(*buf, '\n')
-}
-
-func (l *Logger) log(level Level, format string, args ...interface{}) {
-	bufp := l.bufPool.Get().(*[]byte)
-	l.formatMessage(bufp, level, format, args...)
-
-	if len(*bufp) > maxMessageSize {
-		*bufp = (*bufp)[:maxMessageSize]
-	}
-	_, _ = l.buffer.Write(*bufp)
-
-	l.bufPool.Put(bufp)
-}
-
-func (l *Logger) Error(format string, args ...interface{}) {
-	if l.level.Load() >= uint32(LevelError) {
-		l.log(LevelError, format, args...)
-	}
-}
-
-func (l *Logger) Warn(format string, args ...interface{}) {
-	if l.level.Load() >= uint32(LevelWarn) {
-		l.log(LevelWarn, format, args...)
-	}
-}
-
-func (l *Logger) Info(format string, args ...interface{}) {
-	if l.level.Load() >= uint32(LevelInfo) {
-		l.log(LevelInfo, format, args...)
-	}
-}
-
-func (l *Logger) Debug(format string, args ...interface{}) {
-	if l.level.Load() >= uint32(LevelDebug) {
-		l.log(LevelDebug, format, args...)
-	}
-}
-
-func (l *Logger) Trace(format string, args ...interface{}) {
-	if l.level.Load() >= uint32(LevelTrace) {
-		l.log(LevelTrace, format, args...)
-	}
-}
-
-// worker periodically flushes the buffer
-func (l *Logger) worker() {
-	defer l.wg.Done()
-
-	ticker := time.NewTicker(defaultFlushInterval)
-	defer ticker.Stop()
-
-	buf := make([]byte, 0, maxBatchSize)
-
-	for {
-		select {
-		case <-l.shutdown:
-			return
-		case <-ticker.C:
-			// Read accumulated messages
-			n, _ := l.buffer.Read(buf[:cap(buf)])
-			if n == 0 {
-				continue
-			}
-
-			// Write batch
-			_, _ = l.output.Write(buf[:n])
-		}
-	}
-}
-
-// Stop gracefully shuts down the logger
-func (l *Logger) Stop(ctx context.Context) error {
-	done := make(chan struct{})
-
-	l.closeOnce.Do(func() {
-		close(l.shutdown)
-	})
-
-	go func() {
-		l.wg.Wait()
-		close(done)
-	}()
-
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-done:
-		return nil
-	}
-}

client/firewall/uspfilter/log/ringbuffer.go

@@ -1,85 +0,0 @@
-package log
-
-import "sync"
-
-// ringBuffer is a simple ring buffer implementation
-type ringBuffer struct {
-	buf  []byte
-	size int
-	r, w int64 // Read and write positions
-	mu   sync.Mutex
-}
-
-func newRingBuffer(size int) *ringBuffer {
-	return &ringBuffer{
-		buf:  make([]byte, size),
-		size: size,
-	}
-}
-
-func (r *ringBuffer) Write(p []byte) (n int, err error) {
-	if len(p) == 0 {
-		return 0, nil
-	}
-
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	if len(p) > r.size {
-		p = p[:r.size]
-	}
-
-	n = len(p)
-
-	// Write data, handling wrap-around
-	pos := int(r.w % int64(r.size))
-	writeLen := min(len(p), r.size-pos)
-	copy(r.buf[pos:], p[:writeLen])
-
-	// If we have more data and need to wrap around
-	if writeLen < len(p) {
-		copy(r.buf, p[writeLen:])
-	}
-
-	// Update write position
-	r.w += int64(n)
-
-	return n, nil
-}
-
-func (r *ringBuffer) Read(p []byte) (n int, err error) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	if r.w == r.r {
-		return 0, nil
-	}
-
-	// Calculate available data accounting for wraparound
-	available := int(r.w - r.r)
-	if available < 0 {
-		available += r.size
-	}
-	available = min(available, r.size)
-
-	// Limit read to buffer size
-	toRead := min(available, len(p))
-	if toRead == 0 {
-		return 0, nil
-	}
-
-	// Read data, handling wrap-around
-	pos := int(r.r % int64(r.size))
-	readLen := min(toRead, r.size-pos)
-	n = copy(p, r.buf[pos:pos+readLen])
-
-	// If we need more data and need to wrap around
-	if readLen < toRead {
-		n += copy(p[readLen:toRead], r.buf[:toRead-readLen])
-	}
-
-	// Update read position
-	r.r += int64(n)
-
-	return n, nil
-}

client/firewall/uspfilter/rule.go

@@ -2,22 +2,19 @@ package uspfilter
 
 import (
 	"net"
-	"net/netip"
 
 	"github.com/google/gopacket"
-
-	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 )
 
-// PeerRule to handle management of rules
-type PeerRule struct {
+// Rule to handle management of rules
+type Rule struct {
 	id         string
 	ip         net.IP
 	ipLayer    gopacket.LayerType
 	matchByIP  bool
 	protoLayer gopacket.LayerType
-	sPort      *firewall.Port
-	dPort      *firewall.Port
+	sPort      uint16
+	dPort      uint16
 	drop       bool
 	comment    string
 
@@ -25,21 +22,6 @@ type PeerRule struct {
 }
 
 // GetRuleID returns the rule id
-func (r *PeerRule) GetRuleID() string {
-	return r.id
-}
-
-type RouteRule struct {
-	id          string
-	sources     []netip.Prefix
-	destination netip.Prefix
-	proto       firewall.Protocol
-	srcPort     *firewall.Port
-	dstPort     *firewall.Port
-	action      firewall.Action
-}
-
-// GetRuleID returns the rule id
-func (r *RouteRule) GetRuleID() string {
+func (r *Rule) GetRuleID() string {
 	return r.id
 }

client/firewall/uspfilter/tracer.go

@@ -1,390 +0,0 @@
-package uspfilter
-
-import (
-	"fmt"
-	"net"
-	"time"
-
-	"github.com/google/gopacket"
-	"github.com/google/gopacket/layers"
-
-	fw "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
-)
-
-type PacketStage int
-
-const (
-	StageReceived PacketStage = iota
-	StageConntrack
-	StagePeerACL
-	StageRouting
-	StageRouteACL
-	StageForwarding
-	StageCompleted
-)
-
-const msgProcessingCompleted = "Processing completed"
-
-func (s PacketStage) String() string {
-	return map[PacketStage]string{
-		StageReceived:   "Received",
-		StageConntrack:  "Connection Tracking",
-		StagePeerACL:    "Peer ACL",
-		StageRouting:    "Routing",
-		StageRouteACL:   "Route ACL",
-		StageForwarding: "Forwarding",
-		StageCompleted:  "Completed",
-	}[s]
-}
-
-type ForwarderAction struct {
-	Action     string
-	RemoteAddr string
-	Error      error
-}
-
-type TraceResult struct {
-	Timestamp       time.Time
-	Stage           PacketStage
-	Message         string
-	Allowed         bool
-	ForwarderAction *ForwarderAction
-}
-
-type PacketTrace struct {
-	SourceIP        net.IP
-	DestinationIP   net.IP
-	Protocol        string
-	SourcePort      uint16
-	DestinationPort uint16
-	Direction       fw.RuleDirection
-	Results         []TraceResult
-}
-
-type TCPState struct {
-	SYN bool
-	ACK bool
-	FIN bool
-	RST bool
-	PSH bool
-	URG bool
-}
-
-type PacketBuilder struct {
-	SrcIP       net.IP
-	DstIP       net.IP
-	Protocol    fw.Protocol
-	SrcPort     uint16
-	DstPort     uint16
-	ICMPType    uint8
-	ICMPCode    uint8
-	Direction   fw.RuleDirection
-	PayloadSize int
-	TCPState    *TCPState
-}
-
-func (t *PacketTrace) AddResult(stage PacketStage, message string, allowed bool) {
-	t.Results = append(t.Results, TraceResult{
-		Timestamp: time.Now(),
-		Stage:     stage,
-		Message:   message,
-		Allowed:   allowed,
-	})
-}
-
-func (t *PacketTrace) AddResultWithForwarder(stage PacketStage, message string, allowed bool, action *ForwarderAction) {
-	t.Results = append(t.Results, TraceResult{
-		Timestamp:       time.Now(),
-		Stage:           stage,
-		Message:         message,
-		Allowed:         allowed,
-		ForwarderAction: action,
-	})
-}
-
-func (p *PacketBuilder) Build() ([]byte, error) {
-	ip := p.buildIPLayer()
-	pktLayers := []gopacket.SerializableLayer{ip}
-
-	transportLayer, err := p.buildTransportLayer(ip)
-	if err != nil {
-		return nil, err
-	}
-	pktLayers = append(pktLayers, transportLayer...)
-
-	if p.PayloadSize > 0 {
-		payload := make([]byte, p.PayloadSize)
-		pktLayers = append(pktLayers, gopacket.Payload(payload))
-	}
-
-	return serializePacket(pktLayers)
-}
-
-func (p *PacketBuilder) buildIPLayer() *layers.IPv4 {
-	return &layers.IPv4{
-		Version:  4,
-		TTL:      64,
-		Protocol: layers.IPProtocol(getIPProtocolNumber(p.Protocol)),
-		SrcIP:    p.SrcIP,
-		DstIP:    p.DstIP,
-	}
-}
-
-func (p *PacketBuilder) buildTransportLayer(ip *layers.IPv4) ([]gopacket.SerializableLayer, error) {
-	switch p.Protocol {
-	case "tcp":
-		return p.buildTCPLayer(ip)
-	case "udp":
-		return p.buildUDPLayer(ip)
-	case "icmp":
-		return p.buildICMPLayer()
-	default:
-		return nil, fmt.Errorf("unsupported protocol: %s", p.Protocol)
-	}
-}
-
-func (p *PacketBuilder) buildTCPLayer(ip *layers.IPv4) ([]gopacket.SerializableLayer, error) {
-	tcp := &layers.TCP{
-		SrcPort: layers.TCPPort(p.SrcPort),
-		DstPort: layers.TCPPort(p.DstPort),
-		Window:  65535,
-		SYN:     p.TCPState != nil && p.TCPState.SYN,
-		ACK:     p.TCPState != nil && p.TCPState.ACK,
-		FIN:     p.TCPState != nil && p.TCPState.FIN,
-		RST:     p.TCPState != nil && p.TCPState.RST,
-		PSH:     p.TCPState != nil && p.TCPState.PSH,
-		URG:     p.TCPState != nil && p.TCPState.URG,
-	}
-	if err := tcp.SetNetworkLayerForChecksum(ip); err != nil {
-		return nil, fmt.Errorf("set network layer for TCP checksum: %w", err)
-	}
-	return []gopacket.SerializableLayer{tcp}, nil
-}
-
-func (p *PacketBuilder) buildUDPLayer(ip *layers.IPv4) ([]gopacket.SerializableLayer, error) {
-	udp := &layers.UDP{
-		SrcPort: layers.UDPPort(p.SrcPort),
-		DstPort: layers.UDPPort(p.DstPort),
-	}
-	if err := udp.SetNetworkLayerForChecksum(ip); err != nil {
-		return nil, fmt.Errorf("set network layer for UDP checksum: %w", err)
-	}
-	return []gopacket.SerializableLayer{udp}, nil
-}
-
-func (p *PacketBuilder) buildICMPLayer() ([]gopacket.SerializableLayer, error) {
-	icmp := &layers.ICMPv4{
-		TypeCode: layers.CreateICMPv4TypeCode(p.ICMPType, p.ICMPCode),
-	}
-	if p.ICMPType == layers.ICMPv4TypeEchoRequest || p.ICMPType == layers.ICMPv4TypeEchoReply {
-		icmp.Id = uint16(1)
-		icmp.Seq = uint16(1)
-	}
-	return []gopacket.SerializableLayer{icmp}, nil
-}
-
-func serializePacket(layers []gopacket.SerializableLayer) ([]byte, error) {
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{
-		ComputeChecksums: true,
-		FixLengths:       true,
-	}
-	if err := gopacket.SerializeLayers(buf, opts, layers...); err != nil {
-		return nil, fmt.Errorf("serialize packet: %w", err)
-	}
-	return buf.Bytes(), nil
-}
-
-func getIPProtocolNumber(protocol fw.Protocol) int {
-	switch protocol {
-	case fw.ProtocolTCP:
-		return int(layers.IPProtocolTCP)
-	case fw.ProtocolUDP:
-		return int(layers.IPProtocolUDP)
-	case fw.ProtocolICMP:
-		return int(layers.IPProtocolICMPv4)
-	default:
-		return 0
-	}
-}
-
-func (m *Manager) TracePacketFromBuilder(builder *PacketBuilder) (*PacketTrace, error) {
-	packetData, err := builder.Build()
-	if err != nil {
-		return nil, fmt.Errorf("build packet: %w", err)
-	}
-
-	return m.TracePacket(packetData, builder.Direction), nil
-}
-
-func (m *Manager) TracePacket(packetData []byte, direction fw.RuleDirection) *PacketTrace {
-
-	d := m.decoders.Get().(*decoder)
-	defer m.decoders.Put(d)
-
-	trace := &PacketTrace{Direction: direction}
-
-	// Initial packet decoding
-	if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-		trace.AddResult(StageReceived, fmt.Sprintf("Failed to decode packet: %v", err), false)
-		return trace
-	}
-
-	// Extract base packet info
-	srcIP, dstIP := m.extractIPs(d)
-	trace.SourceIP = srcIP
-	trace.DestinationIP = dstIP
-
-	// Determine protocol and ports
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		trace.Protocol = "TCP"
-		trace.SourcePort = uint16(d.tcp.SrcPort)
-		trace.DestinationPort = uint16(d.tcp.DstPort)
-	case layers.LayerTypeUDP:
-		trace.Protocol = "UDP"
-		trace.SourcePort = uint16(d.udp.SrcPort)
-		trace.DestinationPort = uint16(d.udp.DstPort)
-	case layers.LayerTypeICMPv4:
-		trace.Protocol = "ICMP"
-	}
-
-	trace.AddResult(StageReceived, fmt.Sprintf("Received %s packet: %s:%d -> %s:%d",
-		trace.Protocol, srcIP, trace.SourcePort, dstIP, trace.DestinationPort), true)
-
-	if direction == fw.RuleDirectionOUT {
-		return m.traceOutbound(packetData, trace)
-	}
-
-	return m.traceInbound(packetData, trace, d, srcIP, dstIP)
-}
-
-func (m *Manager) traceInbound(packetData []byte, trace *PacketTrace, d *decoder, srcIP net.IP, dstIP net.IP) *PacketTrace {
-	if m.stateful && m.handleConntrackState(trace, d, srcIP, dstIP) {
-		return trace
-	}
-
-	if m.handleLocalDelivery(trace, packetData, d, srcIP, dstIP) {
-		return trace
-	}
-
-	if !m.handleRouting(trace) {
-		return trace
-	}
-
-	if m.nativeRouter {
-		return m.handleNativeRouter(trace)
-	}
-
-	return m.handleRouteACLs(trace, d, srcIP, dstIP)
-}
-
-func (m *Manager) handleConntrackState(trace *PacketTrace, d *decoder, srcIP, dstIP net.IP) bool {
-	allowed := m.isValidTrackedConnection(d, srcIP, dstIP)
-	msg := "No existing connection found"
-	if allowed {
-		msg = m.buildConntrackStateMessage(d)
-		trace.AddResult(StageConntrack, msg, true)
-		trace.AddResult(StageCompleted, "Packet allowed by connection tracking", true)
-		return true
-	}
-	trace.AddResult(StageConntrack, msg, false)
-	return false
-}
-
-func (m *Manager) buildConntrackStateMessage(d *decoder) string {
-	msg := "Matched existing connection state"
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		flags := getTCPFlags(&d.tcp)
-		msg += fmt.Sprintf(" (TCP Flags: SYN=%v ACK=%v RST=%v FIN=%v)",
-			flags&conntrack.TCPSyn != 0,
-			flags&conntrack.TCPAck != 0,
-			flags&conntrack.TCPRst != 0,
-			flags&conntrack.TCPFin != 0)
-	case layers.LayerTypeICMPv4:
-		msg += fmt.Sprintf(" (ICMP ID=%d, Seq=%d)", d.icmp4.Id, d.icmp4.Seq)
-	}
-	return msg
-}
-
-func (m *Manager) handleLocalDelivery(trace *PacketTrace, packetData []byte, d *decoder, srcIP, dstIP net.IP) bool {
-	if !m.localForwarding {
-		trace.AddResult(StageRouting, "Local forwarding disabled", false)
-		trace.AddResult(StageCompleted, "Packet dropped - local forwarding disabled", false)
-		return true
-	}
-
-	trace.AddResult(StageRouting, "Packet destined for local delivery", true)
-	blocked := m.peerACLsBlock(srcIP, packetData, m.incomingRules, d)
-
-	msg := "Allowed by peer ACL rules"
-	if blocked {
-		msg = "Blocked by peer ACL rules"
-	}
-	trace.AddResult(StagePeerACL, msg, !blocked)
-
-	if m.netstack {
-		m.addForwardingResult(trace, "proxy-local", "127.0.0.1", !blocked)
-	}
-
-	trace.AddResult(StageCompleted, msgProcessingCompleted, !blocked)
-	return true
-}
-
-func (m *Manager) handleRouting(trace *PacketTrace) bool {
-	if !m.routingEnabled {
-		trace.AddResult(StageRouting, "Routing disabled", false)
-		trace.AddResult(StageCompleted, "Packet dropped - routing disabled", false)
-		return false
-	}
-	trace.AddResult(StageRouting, "Routing enabled, checking ACLs", true)
-	return true
-}
-
-func (m *Manager) handleNativeRouter(trace *PacketTrace) *PacketTrace {
-	trace.AddResult(StageRouteACL, "Using native router, skipping ACL checks", true)
-	trace.AddResult(StageForwarding, "Forwarding via native router", true)
-	trace.AddResult(StageCompleted, msgProcessingCompleted, true)
-	return trace
-}
-
-func (m *Manager) handleRouteACLs(trace *PacketTrace, d *decoder, srcIP, dstIP net.IP) *PacketTrace {
-	proto := getProtocolFromPacket(d)
-	srcPort, dstPort := getPortsFromPacket(d)
-	allowed := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort)
-
-	msg := "Allowed by route ACLs"
-	if !allowed {
-		msg = "Blocked by route ACLs"
-	}
-	trace.AddResult(StageRouteACL, msg, allowed)
-
-	if allowed && m.forwarder != nil {
-		m.addForwardingResult(trace, "proxy-remote", fmt.Sprintf("%s:%d", dstIP, dstPort), true)
-	}
-
-	trace.AddResult(StageCompleted, msgProcessingCompleted, allowed)
-	return trace
-}
-
-func (m *Manager) addForwardingResult(trace *PacketTrace, action, remoteAddr string, allowed bool) {
-	fwdAction := &ForwarderAction{
-		Action:     action,
-		RemoteAddr: remoteAddr,
-	}
-	trace.AddResultWithForwarder(StageForwarding,
-		fmt.Sprintf("Forwarding to %s", fwdAction.Action), allowed, fwdAction)
-}
-
-func (m *Manager) traceOutbound(packetData []byte, trace *PacketTrace) *PacketTrace {
-	// will create or update the connection state
-	dropped := m.processOutgoingHooks(packetData)
-	if dropped {
-		trace.AddResult(StageCompleted, "Packet dropped by outgoing hook", false)
-	} else {
-		trace.AddResult(StageCompleted, "Packet allowed (outgoing)", true)
-	}
-	return trace
-}

client/firewall/uspfilter/uspfilter.go

@@ -1,14 +1,11 @@
 package uspfilter
 
 import (
-	"errors"
 	"fmt"
 	"net"
 	"net/netip"
 	"os"
-	"slices"
 	"strconv"
-	"strings"
 	"sync"
 
 	"github.com/google/gopacket"
@@ -17,83 +14,46 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/common"
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/forwarder"
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
-	"github.com/netbirdio/netbird/client/iface/netstack"
+	"github.com/netbirdio/netbird/client/iface"
+	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
 const layerTypeAll = 0
 
-const (
-	// EnvDisableConntrack disables the stateful filter, replies to outbound traffic won't be allowed.
-	EnvDisableConntrack = "NB_DISABLE_CONNTRACK"
+const EnvDisableConntrack = "NB_DISABLE_CONNTRACK"
 
-	// EnvDisableUserspaceRouting disables userspace routing, to-be-routed packets will be dropped.
-	EnvDisableUserspaceRouting = "NB_DISABLE_USERSPACE_ROUTING"
-
-	// EnvForceUserspaceRouter forces userspace routing even if native routing is available.
-	EnvForceUserspaceRouter = "NB_FORCE_USERSPACE_ROUTER"
-
-	// EnvEnableNetstackLocalForwarding enables forwarding of local traffic to the native stack when running netstack
-	// Leaving this on by default introduces a security risk as sockets on listening on localhost only will be accessible
-	EnvEnableNetstackLocalForwarding = "NB_ENABLE_NETSTACK_LOCAL_FORWARDING"
+var (
+	errRouteNotSupported = fmt.Errorf("route not supported with userspace firewall")
 )
 
-// RuleSet is a set of rules grouped by a string key
-type RuleSet map[string]PeerRule
-
-type RouteRules []RouteRule
-
-func (r RouteRules) Sort() {
-	slices.SortStableFunc(r, func(a, b RouteRule) int {
-		// Deny rules come first
-		if a.action == firewall.ActionDrop && b.action != firewall.ActionDrop {
-			return -1
-		}
-		if a.action != firewall.ActionDrop && b.action == firewall.ActionDrop {
-			return 1
-		}
-		return strings.Compare(a.id, b.id)
-	})
+// IFaceMapper defines subset methods of interface required for manager
+type IFaceMapper interface {
+	SetFilter(device.PacketFilter) error
+	Address() iface.WGAddress
 }
 
+// RuleSet is a set of rules grouped by a string key
+type RuleSet map[string]Rule
+
 // Manager userspace firewall manager
 type Manager struct {
 	// outgoingRules is used for hooks only
 	outgoingRules map[string]RuleSet
 	// incomingRules is used for filtering and hooks
 	incomingRules  map[string]RuleSet
-	routeRules     RouteRules
 	wgNetwork      *net.IPNet
 	decoders       sync.Pool
-	wgIface        common.IFaceMapper
+	wgIface        IFaceMapper
 	nativeFirewall firewall.Manager
 
 	mutex sync.RWMutex
 
-	// indicates whether server routes are disabled
-	disableServerRoutes bool
-	// indicates whether we forward packets not destined for ourselves
-	routingEnabled bool
-	// indicates whether we leave forwarding and filtering to the native firewall
-	nativeRouter bool
-	// indicates whether we track outbound connections
-	stateful bool
-	// indicates whether wireguards runs in netstack mode
-	netstack bool
-	// indicates whether we forward local traffic to the native stack
-	localForwarding bool
-
-	localipmanager *localIPManager
-
+	stateful    bool
 	udpTracker  *conntrack.UDPTracker
 	icmpTracker *conntrack.ICMPTracker
 	tcpTracker  *conntrack.TCPTracker
-	forwarder   *forwarder.Forwarder
-	logger      *nblog.Logger
 }
 
 // decoder for packages
@@ -110,44 +70,22 @@ type decoder struct {
 }
 
 // Create userspace firewall manager constructor
-func Create(iface common.IFaceMapper, disableServerRoutes bool) (*Manager, error) {
-	return create(iface, nil, disableServerRoutes)
+func Create(iface IFaceMapper) (*Manager, error) {
+	return create(iface)
 }
 
-func CreateWithNativeFirewall(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool) (*Manager, error) {
-	if nativeFirewall == nil {
-		return nil, errors.New("native firewall is nil")
-	}
-
-	mgr, err := create(iface, nativeFirewall, disableServerRoutes)
+func CreateWithNativeFirewall(iface IFaceMapper, nativeFirewall firewall.Manager) (*Manager, error) {
+	mgr, err := create(iface)
 	if err != nil {
 		return nil, err
 	}
 
+	mgr.nativeFirewall = nativeFirewall
 	return mgr, nil
 }
 
-func parseCreateEnv() (bool, bool) {
-	var disableConntrack, enableLocalForwarding bool
-	var err error
-	if val := os.Getenv(EnvDisableConntrack); val != "" {
-		disableConntrack, err = strconv.ParseBool(val)
-		if err != nil {
-			log.Warnf("failed to parse %s: %v", EnvDisableConntrack, err)
-		}
-	}
-	if val := os.Getenv(EnvEnableNetstackLocalForwarding); val != "" {
-		enableLocalForwarding, err = strconv.ParseBool(val)
-		if err != nil {
-			log.Warnf("failed to parse %s: %v", EnvEnableNetstackLocalForwarding, err)
-		}
-	}
-
-	return disableConntrack, enableLocalForwarding
-}
-
-func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool) (*Manager, error) {
-	disableConntrack, enableLocalForwarding := parseCreateEnv()
+func create(iface IFaceMapper) (*Manager, error) {
+	disableConntrack, _ := strconv.ParseBool(os.Getenv(EnvDisableConntrack))
 
 	m := &Manager{
 		decoders: sync.Pool{
@@ -163,183 +101,52 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 				return d
 			},
 		},
-		nativeFirewall:      nativeFirewall,
-		outgoingRules:       make(map[string]RuleSet),
-		incomingRules:       make(map[string]RuleSet),
-		wgIface:             iface,
-		localipmanager:      newLocalIPManager(),
-		disableServerRoutes: disableServerRoutes,
-		routingEnabled:      false,
-		stateful:            !disableConntrack,
-		logger:              nblog.NewFromLogrus(log.StandardLogger()),
-		netstack:            netstack.IsEnabled(),
-		// default true for non-netstack, for netstack only if explicitly enabled
-		localForwarding: !netstack.IsEnabled() || enableLocalForwarding,
-	}
-
-	if err := m.localipmanager.UpdateLocalIPs(iface); err != nil {
-		return nil, fmt.Errorf("update local IPs: %w", err)
+		outgoingRules: make(map[string]RuleSet),
+		incomingRules: make(map[string]RuleSet),
+		wgIface:       iface,
+		stateful:      !disableConntrack,
 	}
 
+	// Only initialize trackers if stateful mode is enabled
 	if disableConntrack {
 		log.Info("conntrack is disabled")
 	} else {
-		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger)
-		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger)
-		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger)
-	}
-
-	// netstack needs the forwarder for local traffic
-	if m.netstack && m.localForwarding {
-		if err := m.initForwarder(); err != nil {
-			log.Errorf("failed to initialize forwarder: %v", err)
-		}
-	}
-
-	if err := m.blockInvalidRouted(iface); err != nil {
-		log.Errorf("failed to block invalid routed traffic: %v", err)
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout)
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout)
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout)
 	}
 
 	if err := iface.SetFilter(m); err != nil {
-		return nil, fmt.Errorf("set filter: %w", err)
+		return nil, err
 	}
 	return m, nil
 }
 
-func (m *Manager) blockInvalidRouted(iface common.IFaceMapper) error {
-	if m.forwarder == nil {
-		return nil
-	}
-	wgPrefix, err := netip.ParsePrefix(iface.Address().Network.String())
-	if err != nil {
-		return fmt.Errorf("parse wireguard network: %w", err)
-	}
-	log.Debugf("blocking invalid routed traffic for %s", wgPrefix)
-
-	if _, err := m.AddRouteFiltering(
-		[]netip.Prefix{netip.PrefixFrom(netip.IPv4Unspecified(), 0)},
-		wgPrefix,
-		firewall.ProtocolALL,
-		nil,
-		nil,
-		firewall.ActionDrop,
-	); err != nil {
-		return fmt.Errorf("block wg nte : %w", err)
-	}
-
-	// TODO: Block networks that we're a client of
-
-	return nil
-}
-
-func (m *Manager) determineRouting() error {
-	var disableUspRouting, forceUserspaceRouter bool
-	var err error
-	if val := os.Getenv(EnvDisableUserspaceRouting); val != "" {
-		disableUspRouting, err = strconv.ParseBool(val)
-		if err != nil {
-			log.Warnf("failed to parse %s: %v", EnvDisableUserspaceRouting, err)
-		}
-	}
-	if val := os.Getenv(EnvForceUserspaceRouter); val != "" {
-		forceUserspaceRouter, err = strconv.ParseBool(val)
-		if err != nil {
-			log.Warnf("failed to parse %s: %v", EnvForceUserspaceRouter, err)
-		}
-	}
-
-	switch {
-	case disableUspRouting:
-		m.routingEnabled = false
-		m.nativeRouter = false
-		log.Info("userspace routing is disabled")
-
-	case m.disableServerRoutes:
-		//  if server routes are disabled we will let packets pass to the native stack
-		m.routingEnabled = true
-		m.nativeRouter = true
-
-		log.Info("server routes are disabled")
-
-	case forceUserspaceRouter:
-		m.routingEnabled = true
-		m.nativeRouter = false
-
-		log.Info("userspace routing is forced")
-
-	case !m.netstack && m.nativeFirewall != nil && m.nativeFirewall.IsServerRouteSupported():
-		// if the OS supports routing natively, then we don't need to filter/route ourselves
-		// netstack mode won't support native routing as there is no interface
-
-		m.routingEnabled = true
-		m.nativeRouter = true
-
-		log.Info("native routing is enabled")
-
-	default:
-		m.routingEnabled = true
-		m.nativeRouter = false
-
-		log.Info("userspace routing enabled by default")
-	}
-
-	if m.routingEnabled && !m.nativeRouter {
-		return m.initForwarder()
-	}
-
-	return nil
-}
-
-// initForwarder initializes the forwarder, it disables routing on errors
-func (m *Manager) initForwarder() error {
-	if m.forwarder != nil {
-		return nil
-	}
-
-	// Only supported in userspace mode as we need to inject packets back into wireguard directly
-	intf := m.wgIface.GetWGDevice()
-	if intf == nil {
-		m.routingEnabled = false
-		return errors.New("forwarding not supported")
-	}
-
-	forwarder, err := forwarder.New(m.wgIface, m.logger, m.netstack)
-	if err != nil {
-		m.routingEnabled = false
-		return fmt.Errorf("create forwarder: %w", err)
-	}
-
-	m.forwarder = forwarder
-
-	log.Debug("forwarder initialized")
-
-	return nil
-}
-
 func (m *Manager) Init(*statemanager.Manager) error {
 	return nil
 }
 
 func (m *Manager) IsServerRouteSupported() bool {
-	return true
+	if m.nativeFirewall == nil {
+		return false
+	} else {
+		return true
+	}
 }
 
 func (m *Manager) AddNatRule(pair firewall.RouterPair) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
-		return m.nativeFirewall.AddNatRule(pair)
+	if m.nativeFirewall == nil {
+		return errRouteNotSupported
 	}
-
-	// userspace routed packets are always SNATed to the inbound direction
-	// TODO: implement outbound SNAT
-	return nil
+	return m.nativeFirewall.AddNatRule(pair)
 }
 
 // RemoveNatRule removes a routing firewall rule
 func (m *Manager) RemoveNatRule(pair firewall.RouterPair) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
-		return m.nativeFirewall.RemoveNatRule(pair)
+	if m.nativeFirewall == nil {
+		return errRouteNotSupported
 	}
-	return nil
+	return m.nativeFirewall.RemoveNatRule(pair)
 }
 
 // AddPeerFiltering rule to the firewall
@@ -355,7 +162,7 @@ func (m *Manager) AddPeerFiltering(
 	_ string,
 	comment string,
 ) ([]firewall.Rule, error) {
-	r := PeerRule{
+	r := Rule{
 		id:        uuid.New().String(),
 		ip:        ip,
 		ipLayer:   layers.LayerTypeIPv6,
@@ -372,8 +179,13 @@ func (m *Manager) AddPeerFiltering(
 		r.matchByIP = false
 	}
 
-	r.sPort = sPort
-	r.dPort = dPort
+	if sPort != nil && len(sPort.Values) == 1 {
+		r.sPort = uint16(sPort.Values[0])
+	}
+
+	if dPort != nil && len(dPort.Values) == 1 {
+		r.dPort = uint16(dPort.Values[0])
+	}
 
 	switch proto {
 	case firewall.ProtocolTCP:
@@ -398,56 +210,18 @@ func (m *Manager) AddPeerFiltering(
 	return []firewall.Rule{&r}, nil
 }
 
-func (m *Manager) AddRouteFiltering(
-	sources []netip.Prefix,
-	destination netip.Prefix,
-	proto firewall.Protocol,
-	sPort *firewall.Port,
-	dPort *firewall.Port,
-	action firewall.Action,
-) (firewall.Rule, error) {
-	if m.nativeRouter && m.nativeFirewall != nil {
-		return m.nativeFirewall.AddRouteFiltering(sources, destination, proto, sPort, dPort, action)
+func (m *Manager) AddRouteFiltering(sources []netip.Prefix, destination netip.Prefix, proto firewall.Protocol, sPort *firewall.Port, dPort *firewall.Port, action firewall.Action) (firewall.Rule, error) {
+	if m.nativeFirewall == nil {
+		return nil, errRouteNotSupported
 	}
-
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	ruleID := uuid.New().String()
-	rule := RouteRule{
-		id:          ruleID,
-		sources:     sources,
-		destination: destination,
-		proto:       proto,
-		srcPort:     sPort,
-		dstPort:     dPort,
-		action:      action,
-	}
-
-	m.routeRules = append(m.routeRules, rule)
-	m.routeRules.Sort()
-
-	return &rule, nil
+	return m.nativeFirewall.AddRouteFiltering(sources, destination, proto, sPort, dPort, action)
 }
 
 func (m *Manager) DeleteRouteRule(rule firewall.Rule) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
-		return m.nativeFirewall.DeleteRouteRule(rule)
+	if m.nativeFirewall == nil {
+		return errRouteNotSupported
 	}
-
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	ruleID := rule.GetRuleID()
-	idx := slices.IndexFunc(m.routeRules, func(r RouteRule) bool {
-		return r.id == ruleID
-	})
-	if idx < 0 {
-		return fmt.Errorf("route rule not found: %s", ruleID)
-	}
-
-	m.routeRules = slices.Delete(m.routeRules, idx, idx+1)
-	return nil
+	return m.nativeFirewall.DeleteRouteRule(rule)
 }
 
 // DeletePeerRule from the firewall by rule definition
@@ -455,7 +229,7 @@ func (m *Manager) DeletePeerRule(rule firewall.Rule) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	r, ok := rule.(*PeerRule)
+	r, ok := rule.(*Rule)
 	if !ok {
 		return fmt.Errorf("delete rule: invalid rule type: %T", rule)
 	}
@@ -486,14 +260,10 @@ func (m *Manager) DropOutgoing(packetData []byte) bool {
 
 // DropIncoming filter incoming packets
 func (m *Manager) DropIncoming(packetData []byte) bool {
-	return m.dropFilter(packetData)
-}
-
-// UpdateLocalIPs updates the list of local IPs
-func (m *Manager) UpdateLocalIPs() error {
-	return m.localipmanager.UpdateLocalIPs(m.wgIface)
+	return m.dropFilter(packetData, m.incomingRules)
 }
 
+// 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()
@@ -514,11 +284,18 @@ func (m *Manager) processOutgoingHooks(packetData []byte) bool {
 		return false
 	}
 
-	// Track all protocols if stateful mode is enabled
+	// Always process UDP hooks
+	if d.decoded[1] == layers.LayerTypeUDP {
+		// Track UDP state only if enabled
+		if m.stateful {
+			m.trackUDPOutbound(d, srcIP, dstIP)
+		}
+		return m.checkUDPHooks(d, dstIP, packetData)
+	}
+
+	// Track other protocols only if stateful mode is enabled
 	if m.stateful {
 		switch d.decoded[1] {
-		case layers.LayerTypeUDP:
-			m.trackUDPOutbound(d, srcIP, dstIP)
 		case layers.LayerTypeTCP:
 			m.trackTCPOutbound(d, srcIP, dstIP)
 		case layers.LayerTypeICMPv4:
@@ -526,11 +303,6 @@ func (m *Manager) processOutgoingHooks(packetData []byte) bool {
 		}
 	}
 
-	// Process UDP hooks even if stateful mode is disabled
-	if d.decoded[1] == layers.LayerTypeUDP {
-		return m.checkUDPHooks(d, dstIP, packetData)
-	}
-
 	return false
 }
 
@@ -592,7 +364,7 @@ func (m *Manager) checkUDPHooks(d *decoder, dstIP net.IP, packetData []byte) boo
 	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 && portsMatch(rule.dPort, uint16(d.udp.DstPort)) {
+				if rule.udpHook != nil && (rule.dPort == 0 || rule.dPort == uint16(d.udp.DstPort)) {
 					return rule.udpHook(packetData)
 				}
 			}
@@ -612,9 +384,10 @@ func (m *Manager) trackICMPOutbound(d *decoder, srcIP, dstIP net.IP) {
 	}
 }
 
-// dropFilter implements filtering logic for incoming packets.
-// If it returns true, the packet should be dropped.
-func (m *Manager) dropFilter(packetData []byte) bool {
+// dropFilter implements filtering logic for incoming packets
+func (m *Manager) dropFilter(packetData []byte, rules map[string]RuleSet) bool {
+	// TODO: Disable router if --disable-server-router is set
+
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
@@ -627,127 +400,39 @@ func (m *Manager) dropFilter(packetData []byte) bool {
 
 	srcIP, dstIP := m.extractIPs(d)
 	if srcIP == nil {
-		m.logger.Error("Unknown network layer: %v", d.decoded[0])
+		log.Errorf("unknown layer: %v", d.decoded[0])
 		return true
 	}
 
-	// For all inbound traffic, first check if it matches a tracked connection.
-	// This must happen before any other filtering because the packets are statefully tracked.
+	if !m.isWireguardTraffic(srcIP, dstIP) {
+		return false
+	}
+
+	// Check connection state only if enabled
 	if m.stateful && m.isValidTrackedConnection(d, srcIP, dstIP) {
 		return false
 	}
 
-	if m.localipmanager.IsLocalIP(dstIP) {
-		return m.handleLocalTraffic(d, srcIP, dstIP, packetData)
-	}
-
-	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData)
-}
-
-// handleLocalTraffic handles local traffic.
-// If it returns true, the packet should be dropped.
-func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP net.IP, packetData []byte) bool {
-	if !m.localForwarding {
-		m.logger.Trace("Dropping local packet (local forwarding disabled): src=%s dst=%s", srcIP, dstIP)
-		return true
-	}
-
-	if m.peerACLsBlock(srcIP, packetData, m.incomingRules, d) {
-		m.logger.Trace("Dropping local packet (ACL denied): src=%s dst=%s",
-			srcIP, dstIP)
-		return true
-	}
-
-	// if running in netstack mode we need to pass this to the forwarder
-	if m.netstack {
-		m.handleNetstackLocalTraffic(packetData)
-
-		// don't process this packet further
-		return true
-	}
-
-	return false
-}
-func (m *Manager) handleNetstackLocalTraffic(packetData []byte) {
-	if m.forwarder == nil {
-		return
-	}
-
-	if err := m.forwarder.InjectIncomingPacket(packetData); err != nil {
-		m.logger.Error("Failed to inject local packet: %v", err)
-	}
-}
-
-// handleRoutedTraffic handles routed traffic.
-// If it returns true, the packet should be dropped.
-func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP net.IP, packetData []byte) bool {
-	// Drop if routing is disabled
-	if !m.routingEnabled {
-		m.logger.Trace("Dropping routed packet (routing disabled): src=%s dst=%s",
-			srcIP, dstIP)
-		return true
-	}
-
-	// Pass to native stack if native router is enabled or forced
-	if m.nativeRouter {
-		return false
-	}
-
-	proto := getProtocolFromPacket(d)
-	srcPort, dstPort := getPortsFromPacket(d)
-
-	if !m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort) {
-		m.logger.Trace("Dropping routed packet (ACL denied): src=%s:%d dst=%s:%d proto=%v",
-			srcIP, srcPort, dstIP, dstPort, proto)
-		return true
-	}
-
-	// Let forwarder handle the packet if it passed route ACLs
-	if err := m.forwarder.InjectIncomingPacket(packetData); err != nil {
-		m.logger.Error("Failed to inject incoming packet: %v", err)
-	}
-
-	// Forwarded packets shouldn't reach the native stack, hence they won't be visible in a packet capture
-	return true
-}
-
-func getProtocolFromPacket(d *decoder) firewall.Protocol {
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		return firewall.ProtocolTCP
-	case layers.LayerTypeUDP:
-		return firewall.ProtocolUDP
-	case layers.LayerTypeICMPv4, layers.LayerTypeICMPv6:
-		return firewall.ProtocolICMP
-	default:
-		return firewall.ProtocolALL
-	}
-}
-
-func getPortsFromPacket(d *decoder) (srcPort, dstPort uint16) {
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		return uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort)
-	case layers.LayerTypeUDP:
-		return uint16(d.udp.SrcPort), uint16(d.udp.DstPort)
-	default:
-		return 0, 0
-	}
+	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 {
-		m.logger.Trace("couldn't decode packet, err: %s", err)
+		log.Tracef("couldn't decode layer, err: %s", err)
 		return false
 	}
 
 	if len(d.decoded) < 2 {
-		m.logger.Trace("packet doesn't have network and transport layers")
+		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) bool {
 	switch d.decoded[1] {
 	case layers.LayerTypeTCP:
@@ -782,22 +467,7 @@ func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP) bool
 	return false
 }
 
-// isSpecialICMP returns true if the packet is a special ICMP packet that should be allowed
-func (m *Manager) isSpecialICMP(d *decoder) bool {
-	if d.decoded[1] != layers.LayerTypeICMPv4 {
-		return false
-	}
-
-	icmpType := d.icmp4.TypeCode.Type()
-	return icmpType == layers.ICMPv4TypeDestinationUnreachable ||
-		icmpType == layers.ICMPv4TypeTimeExceeded
-}
-
-func (m *Manager) peerACLsBlock(srcIP net.IP, packetData []byte, rules map[string]RuleSet, d *decoder) bool {
-	if m.isSpecialICMP(d) {
-		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
 	}
@@ -814,24 +484,7 @@ func (m *Manager) peerACLsBlock(srcIP net.IP, packetData []byte, rules map[strin
 	return true
 }
 
-func portsMatch(rulePort *firewall.Port, packetPort uint16) bool {
-	if rulePort == nil {
-		return true
-	}
-
-	if rulePort.IsRange {
-		return packetPort >= rulePort.Values[0] && packetPort <= rulePort.Values[1]
-	}
-
-	for _, p := range rulePort.Values {
-		if p == packetPort {
-			return true
-		}
-	}
-	return false
-}
-
-func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *decoder) (bool, bool) {
+func validateRule(ip net.IP, packetData []byte, rules map[string]Rule, d *decoder) (bool, bool) {
 	payloadLayer := d.decoded[1]
 	for _, rule := range rules {
 		if rule.matchByIP && !ip.Equal(rule.ip) {
@@ -848,7 +501,13 @@ func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *de
 
 		switch payloadLayer {
 		case layers.LayerTypeTCP:
-			if portsMatch(rule.sPort, uint16(d.tcp.SrcPort)) && portsMatch(rule.dPort, uint16(d.tcp.DstPort)) {
+			if rule.sPort == 0 && rule.dPort == 0 {
+				return rule.drop, true
+			}
+			if rule.sPort != 0 && rule.sPort == uint16(d.tcp.SrcPort) {
+				return rule.drop, true
+			}
+			if rule.dPort != 0 && rule.dPort == uint16(d.tcp.DstPort) {
 				return rule.drop, true
 			}
 		case layers.LayerTypeUDP:
@@ -858,7 +517,13 @@ func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *de
 				return rule.udpHook(packetData), true
 			}
 
-			if portsMatch(rule.sPort, uint16(d.udp.SrcPort)) && portsMatch(rule.dPort, uint16(d.udp.DstPort)) {
+			if rule.sPort == 0 && rule.dPort == 0 {
+				return rule.drop, true
+			}
+			if rule.sPort != 0 && rule.sPort == uint16(d.udp.SrcPort) {
+				return rule.drop, true
+			}
+			if rule.dPort != 0 && rule.dPort == uint16(d.udp.DstPort) {
 				return rule.drop, true
 			}
 		case layers.LayerTypeICMPv4, layers.LayerTypeICMPv6:
@@ -868,51 +533,6 @@ func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *de
 	return false, false
 }
 
-// routeACLsPass returns treu if the packet is allowed by the route ACLs
-func (m *Manager) routeACLsPass(srcIP, dstIP net.IP, proto firewall.Protocol, srcPort, dstPort uint16) bool {
-	m.mutex.RLock()
-	defer m.mutex.RUnlock()
-
-	srcAddr := netip.AddrFrom4([4]byte(srcIP.To4()))
-	dstAddr := netip.AddrFrom4([4]byte(dstIP.To4()))
-
-	for _, rule := range m.routeRules {
-		if m.ruleMatches(rule, srcAddr, dstAddr, proto, srcPort, dstPort) {
-			return rule.action == firewall.ActionAccept
-		}
-	}
-	return false
-}
-
-func (m *Manager) ruleMatches(rule RouteRule, srcAddr, dstAddr netip.Addr, proto firewall.Protocol, srcPort, dstPort uint16) bool {
-	if !rule.destination.Contains(dstAddr) {
-		return false
-	}
-
-	sourceMatched := false
-	for _, src := range rule.sources {
-		if src.Contains(srcAddr) {
-			sourceMatched = true
-			break
-		}
-	}
-	if !sourceMatched {
-		return false
-	}
-
-	if rule.proto != firewall.ProtocolALL && rule.proto != proto {
-		return false
-	}
-
-	if proto == firewall.ProtocolTCP || proto == firewall.ProtocolUDP {
-		if !portsMatch(rule.srcPort, srcPort) || !portsMatch(rule.dstPort, dstPort) {
-			return false
-		}
-	}
-
-	return true
-}
-
 // SetNetwork of the wireguard interface to which filtering applied
 func (m *Manager) SetNetwork(network *net.IPNet) {
 	m.wgNetwork = network
@@ -924,11 +544,11 @@ func (m *Manager) SetNetwork(network *net.IPNet) {
 func (m *Manager) AddUDPPacketHook(
 	in bool, ip net.IP, dPort uint16, hook func([]byte) bool,
 ) string {
-	r := PeerRule{
+	r := Rule{
 		id:         uuid.New().String(),
 		ip:         ip,
 		protoLayer: layers.LayerTypeUDP,
-		dPort:      &firewall.Port{Values: []uint16{dPort}},
+		dPort:      dPort,
 		ipLayer:    layers.LayerTypeIPv6,
 		comment:    fmt.Sprintf("UDP Hook direction: %v, ip:%v, dport:%d", in, ip, dPort),
 		udpHook:    hook,
@@ -941,12 +561,12 @@ func (m *Manager) AddUDPPacketHook(
 	m.mutex.Lock()
 	if in {
 		if _, ok := m.incomingRules[r.ip.String()]; !ok {
-			m.incomingRules[r.ip.String()] = make(map[string]PeerRule)
+			m.incomingRules[r.ip.String()] = make(map[string]Rule)
 		}
 		m.incomingRules[r.ip.String()][r.id] = r
 	} else {
 		if _, ok := m.outgoingRules[r.ip.String()]; !ok {
-			m.outgoingRules[r.ip.String()] = make(map[string]PeerRule)
+			m.outgoingRules[r.ip.String()] = make(map[string]Rule)
 		}
 		m.outgoingRules[r.ip.String()][r.id] = r
 	}
@@ -979,41 +599,3 @@ func (m *Manager) RemovePacketHook(hookID string) error {
 	}
 	return fmt.Errorf("hook with given id not found")
 }
-
-// SetLogLevel sets the log level for the firewall manager
-func (m *Manager) SetLogLevel(level log.Level) {
-	if m.logger != nil {
-		m.logger.SetLevel(nblog.Level(level))
-	}
-}
-
-func (m *Manager) EnableRouting() error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.determineRouting()
-}
-
-func (m *Manager) DisableRouting() error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	if m.forwarder == nil {
-		return nil
-	}
-
-	m.routingEnabled = false
-	m.nativeRouter = false
-
-	// don't stop forwarder if in use by netstack
-	if m.netstack && m.localForwarding {
-		return nil
-	}
-
-	m.forwarder.Stop()
-	m.forwarder = nil
-
-	log.Debug("forwarder stopped")
-
-	return nil
-}

client/firewall/uspfilter/uspfilter_bench_test.go

@@ -1,12 +1,9 @@
-//go:build uspbench
-
 package uspfilter
 
 import (
 	"fmt"
 	"math/rand"
 	"net"
-	"net/netip"
 	"os"
 	"strings"
 	"testing"
@@ -115,8 +112,8 @@ func BenchmarkCoreFiltering(b *testing.B) {
 				for i := 0; i < 1000; i++ { // Simulate realistic ruleset size
 					ip := generateRandomIPs(1)[0]
 					_, err := m.AddPeerFiltering(ip, fw.ProtocolTCP,
-						&fw.Port{Values: []uint16{uint16(1024 + i)}},
-						&fw.Port{Values: []uint16{80}},
+						&fw.Port{Values: []int{1024 + i}},
+						&fw.Port{Values: []int{80}},
 						fw.ActionAccept, "", "explicit return")
 					require.NoError(b, err)
 				}
@@ -158,7 +155,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 				// Create manager and basic setup
 				manager, _ := Create(&IFaceMock{
 					SetFilterFunc: func(device.PacketFilter) error { return nil },
-				}, false)
+				})
 				defer b.Cleanup(func() {
 					require.NoError(b, manager.Reset(nil))
 				})
@@ -188,7 +185,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 				// Measure inbound packet processing
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
-					manager.dropFilter(inbound)
+					manager.dropFilter(inbound, manager.incomingRules)
 				}
 			})
 		}
@@ -203,7 +200,7 @@ func BenchmarkStateScaling(b *testing.B) {
 		b.Run(fmt.Sprintf("conns_%d", count), func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -231,7 +228,7 @@ func BenchmarkStateScaling(b *testing.B) {
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(testIn)
+				manager.dropFilter(testIn, manager.incomingRules)
 			}
 		})
 	}
@@ -251,7 +248,7 @@ func BenchmarkEstablishmentOverhead(b *testing.B) {
 		b.Run(sc.name, func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -272,7 +269,7 @@ func BenchmarkEstablishmentOverhead(b *testing.B) {
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(inbound)
+				manager.dropFilter(inbound, manager.incomingRules)
 			}
 		})
 	}
@@ -450,7 +447,7 @@ func BenchmarkRoutedNetworkReturn(b *testing.B) {
 		b.Run(sc.name, func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -475,7 +472,7 @@ func BenchmarkRoutedNetworkReturn(b *testing.B) {
 					manager.processOutgoingHooks(syn)
 					// SYN-ACK
 					synack := generateTCPPacketWithFlags(b, dstIP, srcIP, 80, 1024, uint16(conntrack.TCPSyn|conntrack.TCPAck))
-					manager.dropFilter(synack)
+					manager.dropFilter(synack, manager.incomingRules)
 					// ACK
 					ack := generateTCPPacketWithFlags(b, srcIP, dstIP, 1024, 80, uint16(conntrack.TCPAck))
 					manager.processOutgoingHooks(ack)
@@ -484,7 +481,7 @@ func BenchmarkRoutedNetworkReturn(b *testing.B) {
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(inbound)
+				manager.dropFilter(inbound, manager.incomingRules)
 			}
 		})
 	}
@@ -577,7 +574,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -591,7 +588,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 			if sc.rules {
 				// Single rule to allow all return traffic from port 80
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
+					&fw.Port{Values: []int{80}},
 					nil,
 					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
@@ -621,7 +618,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 				// SYN-ACK
 				synack := generateTCPPacketWithFlags(b, dstIPs[i], srcIPs[i],
 					80, uint16(1024+i), uint16(conntrack.TCPSyn|conntrack.TCPAck))
-				manager.dropFilter(synack)
+				manager.dropFilter(synack, manager.incomingRules)
 
 				// ACK
 				ack := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
@@ -649,7 +646,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 				// First outbound data
 				manager.processOutgoingHooks(outPackets[connIdx])
 				// Then inbound response - this is what we're actually measuring
-				manager.dropFilter(inPackets[connIdx])
+				manager.dropFilter(inPackets[connIdx], manager.incomingRules)
 			}
 		})
 	}
@@ -668,7 +665,7 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -682,7 +679,7 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 			if sc.rules {
 				// Single rule to allow all return traffic from port 80
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
+					&fw.Port{Values: []int{80}},
 					nil,
 					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
@@ -757,17 +754,17 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 
 				// Connection establishment
 				manager.processOutgoingHooks(p.syn)
-				manager.dropFilter(p.synAck)
+				manager.dropFilter(p.synAck, manager.incomingRules)
 				manager.processOutgoingHooks(p.ack)
 
 				// Data transfer
 				manager.processOutgoingHooks(p.request)
-				manager.dropFilter(p.response)
+				manager.dropFilter(p.response, manager.incomingRules)
 
 				// Connection teardown
 				manager.processOutgoingHooks(p.finClient)
-				manager.dropFilter(p.ackServer)
-				manager.dropFilter(p.finServer)
+				manager.dropFilter(p.ackServer, manager.incomingRules)
+				manager.dropFilter(p.finServer, manager.incomingRules)
 				manager.processOutgoingHooks(p.ackClient)
 			}
 		})
@@ -787,7 +784,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -800,7 +797,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 			// Setup initial state based on scenario
 			if sc.rules {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
+					&fw.Port{Values: []int{80}},
 					nil,
 					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
@@ -828,7 +825,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 
 				synack := generateTCPPacketWithFlags(b, dstIPs[i], srcIPs[i],
 					80, uint16(1024+i), uint16(conntrack.TCPSyn|conntrack.TCPAck))
-				manager.dropFilter(synack)
+				manager.dropFilter(synack, manager.incomingRules)
 
 				ack := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, uint16(conntrack.TCPAck))
@@ -855,7 +852,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 
 					// Simulate bidirectional traffic
 					manager.processOutgoingHooks(outPackets[connIdx])
-					manager.dropFilter(inPackets[connIdx])
+					manager.dropFilter(inPackets[connIdx], manager.incomingRules)
 				}
 			})
 		})
@@ -875,7 +872,7 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Reset(nil))
 			})
@@ -887,7 +884,7 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 
 			if sc.rules {
 				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
+					&fw.Port{Values: []int{80}},
 					nil,
 					fw.ActionAccept, "", "return traffic")
 				require.NoError(b, err)
@@ -952,15 +949,15 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 
 					// Full connection lifecycle
 					manager.processOutgoingHooks(p.syn)
-					manager.dropFilter(p.synAck)
+					manager.dropFilter(p.synAck, manager.incomingRules)
 					manager.processOutgoingHooks(p.ack)
 
 					manager.processOutgoingHooks(p.request)
-					manager.dropFilter(p.response)
+					manager.dropFilter(p.response, manager.incomingRules)
 
 					manager.processOutgoingHooks(p.finClient)
-					manager.dropFilter(p.ackServer)
-					manager.dropFilter(p.finServer)
+					manager.dropFilter(p.ackServer, manager.incomingRules)
+					manager.dropFilter(p.finServer, manager.incomingRules)
 					manager.processOutgoingHooks(p.ackClient)
 				}
 			})
@@ -999,72 +996,3 @@ func generateTCPPacketWithFlags(b *testing.B, srcIP, dstIP net.IP, srcPort, dstP
 	require.NoError(b, gopacket.SerializeLayers(buf, opts, ipv4, tcp, gopacket.Payload("test")))
 	return buf.Bytes()
 }
-
-func BenchmarkRouteACLs(b *testing.B) {
-	manager := setupRoutedManager(b, "10.10.0.100/16")
-
-	// Add several route rules to simulate real-world scenario
-	rules := []struct {
-		sources []netip.Prefix
-		dest    netip.Prefix
-		proto   fw.Protocol
-		port    *fw.Port
-	}{
-		{
-			sources: []netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-			dest:    netip.MustParsePrefix("192.168.1.0/24"),
-			proto:   fw.ProtocolTCP,
-			port:    &fw.Port{Values: []uint16{80, 443}},
-		},
-		{
-			sources: []netip.Prefix{
-				netip.MustParsePrefix("172.16.0.0/12"),
-				netip.MustParsePrefix("10.0.0.0/8"),
-			},
-			dest:  netip.MustParsePrefix("0.0.0.0/0"),
-			proto: fw.ProtocolICMP,
-		},
-		{
-			sources: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0")},
-			dest:    netip.MustParsePrefix("192.168.0.0/16"),
-			proto:   fw.ProtocolUDP,
-			port:    &fw.Port{Values: []uint16{53}},
-		},
-	}
-
-	for _, r := range rules {
-		_, err := manager.AddRouteFiltering(
-			r.sources,
-			r.dest,
-			r.proto,
-			nil,
-			r.port,
-			fw.ActionAccept,
-		)
-		if err != nil {
-			b.Fatal(err)
-		}
-	}
-
-	// Test cases that exercise different matching scenarios
-	cases := []struct {
-		srcIP   string
-		dstIP   string
-		proto   fw.Protocol
-		dstPort uint16
-	}{
-		{"100.10.0.1", "192.168.1.100", fw.ProtocolTCP, 443}, // Match first rule
-		{"172.16.0.1", "8.8.8.8", fw.ProtocolICMP, 0},        // Match second rule
-		{"1.1.1.1", "192.168.1.53", fw.ProtocolUDP, 53},      // Match third rule
-		{"192.168.1.1", "10.0.0.1", fw.ProtocolTCP, 8080},    // No match
-	}
-
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		for _, tc := range cases {
-			srcIP := net.ParseIP(tc.srcIP)
-			dstIP := net.ParseIP(tc.dstIP)
-			manager.routeACLsPass(srcIP, dstIP, tc.proto, 0, tc.dstPort)
-		}
-	}
-}

client/firewall/uspfilter/uspfilter_test.go

@@ -9,38 +9,17 @@ import (
 
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
-	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/require"
-	wgdevice "golang.zx2c4.com/wireguard/device"
 
 	fw "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
-	"github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
 
-var logger = log.NewFromLogrus(logrus.StandardLogger())
-
 type IFaceMock struct {
-	SetFilterFunc   func(device.PacketFilter) error
-	AddressFunc     func() iface.WGAddress
-	GetWGDeviceFunc func() *wgdevice.Device
-	GetDeviceFunc   func() *device.FilteredDevice
-}
-
-func (i *IFaceMock) GetWGDevice() *wgdevice.Device {
-	if i.GetWGDeviceFunc == nil {
-		return nil
-	}
-	return i.GetWGDeviceFunc()
-}
-
-func (i *IFaceMock) GetDevice() *device.FilteredDevice {
-	if i.GetDeviceFunc == nil {
-		return nil
-	}
-	return i.GetDeviceFunc()
+	SetFilterFunc func(device.PacketFilter) error
+	AddressFunc   func() iface.WGAddress
 }
 
 func (i *IFaceMock) SetFilter(iface device.PacketFilter) error {
@@ -62,7 +41,7 @@ func TestManagerCreate(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -82,7 +61,7 @@ func TestManagerAddPeerFiltering(t *testing.T) {
 		},
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -90,7 +69,7 @@ func TestManagerAddPeerFiltering(t *testing.T) {
 
 	ip := net.ParseIP("192.168.1.1")
 	proto := fw.ProtocolTCP
-	port := &fw.Port{Values: []uint16{80}}
+	port := &fw.Port{Values: []int{80}}
 	action := fw.ActionDrop
 	comment := "Test rule"
 
@@ -116,7 +95,7 @@ func TestManagerDeleteRule(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -124,7 +103,7 @@ func TestManagerDeleteRule(t *testing.T) {
 
 	ip := net.ParseIP("192.168.1.1")
 	proto := fw.ProtocolTCP
-	port := &fw.Port{Values: []uint16{80}}
+	port := &fw.Port{Values: []int{80}}
 	action := fw.ActionDrop
 	comment := "Test rule 2"
 
@@ -187,12 +166,12 @@ func TestAddUDPPacketHook(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			manager, err := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			})
 			require.NoError(t, err)
 
 			manager.AddUDPPacketHook(tt.in, tt.ip, tt.dPort, tt.hook)
 
-			var addedRule PeerRule
+			var addedRule Rule
 			if tt.in {
 				if len(manager.incomingRules[tt.ip.String()]) != 1 {
 					t.Errorf("expected 1 incoming rule, got %d", len(manager.incomingRules))
@@ -215,8 +194,8 @@ func TestAddUDPPacketHook(t *testing.T) {
 				t.Errorf("expected ip %s, got %s", tt.ip, addedRule.ip)
 				return
 			}
-			if tt.dPort != addedRule.dPort.Values[0] {
-				t.Errorf("expected dPort %d, got %d", tt.dPort, addedRule.dPort.Values[0])
+			if tt.dPort != addedRule.dPort {
+				t.Errorf("expected dPort %d, got %d", tt.dPort, addedRule.dPort)
 				return
 			}
 			if layers.LayerTypeUDP != addedRule.protoLayer {
@@ -236,7 +215,7 @@ func TestManagerReset(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -244,7 +223,7 @@ func TestManagerReset(t *testing.T) {
 
 	ip := net.ParseIP("192.168.1.1")
 	proto := fw.ProtocolTCP
-	port := &fw.Port{Values: []uint16{80}}
+	port := &fw.Port{Values: []int{80}}
 	action := fw.ActionDrop
 	comment := "Test rule"
 
@@ -268,18 +247,9 @@ func TestManagerReset(t *testing.T) {
 func TestNotMatchByIP(t *testing.T) {
 	ifaceMock := &IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
-		AddressFunc: func() iface.WGAddress {
-			return iface.WGAddress{
-				IP: net.ParseIP("100.10.0.100"),
-				Network: &net.IPNet{
-					IP:   net.ParseIP("100.10.0.0"),
-					Mask: net.CIDRMask(16, 32),
-				},
-			}
-		},
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -328,7 +298,7 @@ func TestNotMatchByIP(t *testing.T) {
 		return
 	}
 
-	if m.dropFilter(buf.Bytes()) {
+	if m.dropFilter(buf.Bytes(), m.incomingRules) {
 		t.Errorf("expected packet to be accepted")
 		return
 	}
@@ -347,7 +317,7 @@ func TestRemovePacketHook(t *testing.T) {
 	}
 
 	// creating manager instance
-	manager, err := Create(iface, false)
+	manager, err := Create(iface)
 	if err != nil {
 		t.Fatalf("Failed to create Manager: %s", err)
 	}
@@ -393,7 +363,7 @@ func TestRemovePacketHook(t *testing.T) {
 func TestProcessOutgoingHooks(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false)
+	})
 	require.NoError(t, err)
 
 	manager.wgNetwork = &net.IPNet{
@@ -401,7 +371,7 @@ func TestProcessOutgoingHooks(t *testing.T) {
 		Mask: net.CIDRMask(16, 32),
 	}
 	manager.udpTracker.Close()
-	manager.udpTracker = conntrack.NewUDPTracker(100*time.Millisecond, logger)
+	manager.udpTracker = conntrack.NewUDPTracker(100 * time.Millisecond)
 	defer func() {
 		require.NoError(t, manager.Reset(nil))
 	}()
@@ -479,7 +449,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 			ifaceMock := &IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
 			}
-			manager, err := Create(ifaceMock, false)
+			manager, err := Create(ifaceMock)
 			require.NoError(t, err)
 			time.Sleep(time.Second)
 
@@ -493,7 +463,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 			ip := net.ParseIP("10.20.0.100")
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
-				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
+				port := &fw.Port{Values: []int{1000 + i}}
 				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
 
 				require.NoError(t, err, "failed to add rule")
@@ -506,7 +476,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false)
+	})
 	require.NoError(t, err)
 
 	manager.wgNetwork = &net.IPNet{
@@ -515,7 +485,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	}
 
 	manager.udpTracker.Close() // Close the existing tracker
-	manager.udpTracker = conntrack.NewUDPTracker(200*time.Millisecond, logger)
+	manager.udpTracker = conntrack.NewUDPTracker(200 * time.Millisecond)
 	manager.decoders = sync.Pool{
 		New: func() any {
 			d := &decoder{
@@ -636,7 +606,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	for _, cp := range checkPoints {
 		time.Sleep(cp.sleep)
 
-		drop = manager.dropFilter(inboundBuf.Bytes())
+		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
@@ -707,7 +677,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify the invalid packet is dropped
-			drop = manager.dropFilter(testBuf.Bytes())
+			drop = manager.dropFilter(testBuf.Bytes(), manager.incomingRules)
 			require.True(t, drop, tc.description)
 		})
 	}

client/iface/configurer/name.go

@@ -2,5 +2,5 @@
 
 package configurer
 
-// WgInterfaceDefault is a default interface name of Netbird
+// WgInterfaceDefault is a default interface name of Wiretrustee
 const WgInterfaceDefault = "wt0"

client/iface/configurer/name_darwin.go

@@ -2,5 +2,5 @@
 
 package configurer
 
-// WgInterfaceDefault is a default interface name of Netbird
+// WgInterfaceDefault is a default interface name of Wiretrustee
 const WgInterfaceDefault = "utun100"

client/iface/configurer/usp.go

@@ -362,7 +362,7 @@ func toWgUserspaceString(wgCfg wgtypes.Config) string {
 }
 
 func getFwmark() int {
-	if nbnet.AdvancedRouting() {
+	if runtime.GOOS == "linux" && !nbnet.CustomRoutingDisabled() {
 		return nbnet.NetbirdFwmark
 	}
 	return 0

client/iface/device.go

@@ -3,8 +3,6 @@
 package iface
 
 import (
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
@@ -17,5 +15,4 @@ type WGTunDevice interface {
 	DeviceName() string
 	Close() error
 	FilteredDevice() *device.FilteredDevice
-	Device() *wgdevice.Device
 }

client/iface/device/device_android.go

@@ -63,7 +63,7 @@ func (t *WGTunDevice) Create(routes []string, dns string, searchDomains []string
 	t.filteredDevice = newDeviceFilter(tunDevice)
 
 	log.Debugf("attaching to interface %v", name)
-	t.device = device.NewDevice(t.filteredDevice, t.iceBind, device.NewLogger(wgLogLevel(), "[netbird] "))
+	t.device = device.NewDevice(t.filteredDevice, t.iceBind, device.NewLogger(wgLogLevel(), "[wiretrustee] "))
 	// without this property mobile devices can discover remote endpoints if the configured one was wrong.
 	// this helps with support for the older NetBird clients that had a hardcoded direct mode
 	// t.device.DisableSomeRoamingForBrokenMobileSemantics()

client/iface/device/device_darwin.go

@@ -117,11 +117,6 @@ func (t *TunDevice) FilteredDevice() *FilteredDevice {
 	return t.filteredDevice
 }
 
-// Device returns the wireguard device
-func (t *TunDevice) Device() *device.Device {
-	return t.device
-}
-
 // assignAddr Adds IP address to the tunnel interface and network route based on the range provided
 func (t *TunDevice) assignAddr() error {
 	cmd := exec.Command("ifconfig", t.name, "inet", t.address.IP.String(), t.address.IP.String())

client/iface/device/device_ios.go

@@ -64,7 +64,7 @@ func (t *TunDevice) Create() (WGConfigurer, error) {
 
 	t.filteredDevice = newDeviceFilter(tunDevice)
 	log.Debug("Attaching to interface")
-	t.device = device.NewDevice(t.filteredDevice, t.iceBind, device.NewLogger(wgLogLevel(), "[netbird] "))
+	t.device = device.NewDevice(t.filteredDevice, t.iceBind, device.NewLogger(wgLogLevel(), "[wiretrustee] "))
 	// without this property mobile devices can discover remote endpoints if the configured one was wrong.
 	// this helps with support for the older NetBird clients that had a hardcoded direct mode
 	// t.device.DisableSomeRoamingForBrokenMobileSemantics()

client/iface/device/device_kernel_unix.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/pion/transport/v3"
 	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/device"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/configurer"
@@ -34,6 +33,8 @@ type TunKernelDevice struct {
 }
 
 func NewKernelDevice(name string, address WGAddress, wgPort int, key string, mtu int, transportNet transport.Net) *TunKernelDevice {
+	checkUser()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	return &TunKernelDevice{
 		ctx:          ctx,
@@ -152,11 +153,6 @@ func (t *TunKernelDevice) DeviceName() string {
 	return t.name
 }
 
-// Device returns the wireguard device, not applicable for kernel devices
-func (t *TunKernelDevice) Device() *device.Device {
-	return nil
-}
-
 func (t *TunKernelDevice) FilteredDevice() *FilteredDevice {
 	return nil
 }

client/iface/device/device_netstack.go

@@ -117,8 +117,3 @@ func (t *TunNetstackDevice) DeviceName() string {
 func (t *TunNetstackDevice) FilteredDevice() *FilteredDevice {
 	return t.filteredDevice
 }
-
-// Device returns the wireguard device
-func (t *TunNetstackDevice) Device() *device.Device {
-	return t.device
-}

client/iface/device/device_usp_unix.go

@@ -4,6 +4,8 @@ package device
 
 import (
 	"fmt"
+	"os"
+	"runtime"
 
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/device"
@@ -30,6 +32,8 @@ type USPDevice struct {
 func NewUSPDevice(name string, address WGAddress, port int, key string, mtu int, iceBind *bind.ICEBind) *USPDevice {
 	log.Infof("using userspace bind mode")
 
+	checkUser()
+
 	return &USPDevice{
 		name:    name,
 		address: address,
@@ -124,14 +128,18 @@ func (t *USPDevice) FilteredDevice() *FilteredDevice {
 	return t.filteredDevice
 }
 
-// Device returns the wireguard device
-func (t *USPDevice) Device() *device.Device {
-	return t.device
-}
-
 // assignAddr Adds IP address to the tunnel interface
 func (t *USPDevice) assignAddr() error {
 	link := newWGLink(t.name)
 
 	return link.assignAddr(t.address)
 }
+
+func checkUser() {
+	if runtime.GOOS == "freebsd" {
+		euid := os.Geteuid()
+		if euid != 0 {
+			log.Warn("newTunUSPDevice: on netbird must run as root to be able to assign address to the tun interface with ifconfig")
+		}
+	}
+}

client/iface/device/device_windows.go

@@ -150,11 +150,6 @@ func (t *TunDevice) FilteredDevice() *FilteredDevice {
 	return t.filteredDevice
 }
 
-// Device returns the wireguard device
-func (t *TunDevice) Device() *device.Device {
-	return t.device
-}
-
 func (t *TunDevice) GetInterfaceGUIDString() (string, error) {
 	if t.nativeTunDevice == nil {
 		return "", fmt.Errorf("interface has not been initialized yet")

client/iface/device_android.go

@@ -1,8 +1,6 @@
 package iface
 
 import (
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
@@ -15,5 +13,4 @@ type WGTunDevice interface {
 	DeviceName() string
 	Close() error
 	FilteredDevice() *device.FilteredDevice
-	Device() *wgdevice.Device
 }

client/iface/freebsd/link.go

@@ -203,11 +203,6 @@ func (l *Link) setAddr(ip, netmask string) error {
 		return fmt.Errorf("set interface addr: %w", err)
 	}
 
-	cmd = exec.Command("ifconfig", l.name, "inet6", "fe80::/64")
-	if out, err := cmd.CombinedOutput(); err != nil {
-		log.Debugf("adding address command '%v' failed with output: %s", cmd.String(), out)
-	}
-
 	return nil
 }
 

client/iface/iface.go

@@ -11,8 +11,6 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
 	"github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/configurer"
@@ -205,11 +203,6 @@ func (w *WGIface) GetDevice() *device.FilteredDevice {
 	return w.tun.FilteredDevice()
 }
 
-// GetWGDevice returns the WireGuard device
-func (w *WGIface) GetWGDevice() *wgdevice.Device {
-	return w.tun.Device()
-}
-
 // GetStats returns the last handshake time, rx and tx bytes for the given peer
 func (w *WGIface) GetStats(peerKey string) (configurer.WGStats, error) {
 	return w.configurer.GetStats(peerKey)

client/iface/iface_moc.go

@@ -4,7 +4,6 @@ import (
 	"net"
 	"time"
 
-	wgdevice "golang.zx2c4.com/wireguard/device"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
@@ -30,7 +29,6 @@ type MockWGIface struct {
 	SetFilterFunc              func(filter device.PacketFilter) error
 	GetFilterFunc              func() device.PacketFilter
 	GetDeviceFunc              func() *device.FilteredDevice
-	GetWGDeviceFunc            func() *wgdevice.Device
 	GetStatsFunc               func(peerKey string) (configurer.WGStats, error)
 	GetInterfaceGUIDStringFunc func() (string, error)
 	GetProxyFunc               func() wgproxy.Proxy
@@ -104,14 +102,11 @@ func (m *MockWGIface) GetDevice() *device.FilteredDevice {
 	return m.GetDeviceFunc()
 }
 
-func (m *MockWGIface) GetWGDevice() *wgdevice.Device {
-	return m.GetWGDeviceFunc()
-}
-
 func (m *MockWGIface) GetStats(peerKey string) (configurer.WGStats, error) {
 	return m.GetStatsFunc(peerKey)
 }
 
 func (m *MockWGIface) GetProxy() wgproxy.Proxy {
-	return m.GetProxyFunc()
+	//TODO implement me
+	panic("implement me")
 }

client/iface/iwginterface.go

@@ -6,7 +6,6 @@ import (
 	"net"
 	"time"
 
-	wgdevice "golang.zx2c4.com/wireguard/device"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
@@ -33,6 +32,5 @@ type IWGIface interface {
 	SetFilter(filter device.PacketFilter) error
 	GetFilter() device.PacketFilter
 	GetDevice() *device.FilteredDevice
-	GetWGDevice() *wgdevice.Device
 	GetStats(peerKey string) (configurer.WGStats, error)
 }

client/iface/iwginterface_windows.go

@@ -4,7 +4,6 @@ import (
 	"net"
 	"time"
 
-	wgdevice "golang.zx2c4.com/wireguard/device"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
 	"github.com/netbirdio/netbird/client/iface/bind"
@@ -31,7 +30,6 @@ type IWGIface interface {
 	SetFilter(filter device.PacketFilter) error
 	GetFilter() device.PacketFilter
 	GetDevice() *device.FilteredDevice
-	GetWGDevice() *wgdevice.Device
 	GetStats(peerKey string) (configurer.WGStats, error)
 	GetInterfaceGUIDString() (string, error)
 }

client/internal/acl/manager.go

@@ -268,16 +268,13 @@ func (d *DefaultManager) protoRuleToFirewallRule(
 	}
 
 	var port *firewall.Port
-	if !portInfoEmpty(r.PortInfo) {
-		port = convertPortInfo(r.PortInfo)
-	} else if r.Port != "" {
-		// old version of management, single port
+	if r.Port != "" {
 		value, err := strconv.Atoi(r.Port)
 		if err != nil {
-			return "", nil, fmt.Errorf("invalid port: %w", err)
+			return "", nil, fmt.Errorf("invalid port, skipping firewall rule")
 		}
 		port = &firewall.Port{
-			Values: []uint16{uint16(value)},
+			Values: []int{value},
 		}
 	}
 
@@ -305,22 +302,6 @@ func (d *DefaultManager) protoRuleToFirewallRule(
 	return ruleID, rules, nil
 }
 
-func portInfoEmpty(portInfo *mgmProto.PortInfo) bool {
-	if portInfo == nil {
-		return true
-	}
-
-	switch portInfo.GetPortSelection().(type) {
-	case *mgmProto.PortInfo_Port:
-		return portInfo.GetPort() == 0
-	case *mgmProto.PortInfo_Range_:
-		r := portInfo.GetRange()
-		return r == nil || r.Start == 0 || r.End == 0
-	default:
-		return true
-	}
-}
-
 func (d *DefaultManager) addInRules(
 	ip net.IP,
 	protocol firewall.Protocol,
@@ -507,7 +488,7 @@ func (d *DefaultManager) squashAcceptRules(
 
 // getRuleGroupingSelector takes all rule properties except IP address to build selector
 func (d *DefaultManager) getRuleGroupingSelector(rule *mgmProto.FirewallRule) string {
-	return fmt.Sprintf("%v:%v:%v:%s:%v", strconv.Itoa(int(rule.Direction)), rule.Action, rule.Protocol, rule.Port, rule.PortInfo)
+	return fmt.Sprintf("%v:%v:%v:%s", strconv.Itoa(int(rule.Direction)), rule.Action, rule.Protocol, rule.Port)
 }
 
 func (d *DefaultManager) rollBack(newRulePairs map[id.RuleID][]firewall.Rule) {
@@ -558,14 +539,14 @@ func convertPortInfo(portInfo *mgmProto.PortInfo) *firewall.Port {
 
 	if portInfo.GetPort() != 0 {
 		return &firewall.Port{
-			Values: []uint16{uint16(int(portInfo.GetPort()))},
+			Values: []int{int(portInfo.GetPort())},
 		}
 	}
 
 	if portInfo.GetRange() != nil {
 		return &firewall.Port{
 			IsRange: true,
-			Values:  []uint16{uint16(portInfo.GetRange().Start), uint16(portInfo.GetRange().End)},
+			Values:  []int{int(portInfo.GetRange().Start), int(portInfo.GetRange().End)},
 		}
 	}
 

client/internal/acl/manager_test.go

@@ -49,10 +49,9 @@ func TestDefaultManager(t *testing.T) {
 		IP:      ip,
 		Network: network,
 	}).AnyTimes()
-	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
 	// we receive one rule from the management so for testing purposes ignore it
-	fw, err := firewall.NewFirewall(ifaceMock, nil, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil)
 	if err != nil {
 		t.Errorf("create firewall: %v", err)
 		return
@@ -343,10 +342,9 @@ func TestDefaultManagerEnableSSHRules(t *testing.T) {
 		IP:      ip,
 		Network: network,
 	}).AnyTimes()
-	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
 	// we receive one rule from the management so for testing purposes ignore it
-	fw, err := firewall.NewFirewall(ifaceMock, nil, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil)
 	if err != nil {
 		t.Errorf("create firewall: %v", err)
 		return

client/internal/acl/mocks/iface_mapper.go

@@ -8,8 +8,6 @@ import (
 	reflect "reflect"
 
 	gomock "github.com/golang/mock/gomock"
-	wgdevice "golang.zx2c4.com/wireguard/device"
-
 	iface "github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
@@ -92,31 +90,3 @@ func (mr *MockIFaceMapperMockRecorder) SetFilter(arg0 interface{}) *gomock.Call
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetFilter", reflect.TypeOf((*MockIFaceMapper)(nil).SetFilter), arg0)
 }
-
-// GetDevice mocks base method.
-func (m *MockIFaceMapper) GetDevice() *device.FilteredDevice {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetDevice")
-	ret0, _ := ret[0].(*device.FilteredDevice)
-	return ret0
-}
-
-// GetDevice indicates an expected call of GetDevice.
-func (mr *MockIFaceMapperMockRecorder) GetDevice() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDevice", reflect.TypeOf((*MockIFaceMapper)(nil).GetDevice))
-}
-
-// GetWGDevice mocks base method.
-func (m *MockIFaceMapper) GetWGDevice() *wgdevice.Device {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetWGDevice")
-	ret0, _ := ret[0].(*wgdevice.Device)
-	return ret0
-}
-
-// GetWGDevice indicates an expected call of GetWGDevice.
-func (mr *MockIFaceMapperMockRecorder) GetWGDevice() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGDevice", reflect.TypeOf((*MockIFaceMapper)(nil).GetWGDevice))
-}

client/internal/auth/device_flow.go

@@ -2,8 +2,6 @@ package auth
 
 import (
 	"context"
-	"crypto/tls"
-	"crypto/x509"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -13,10 +11,7 @@ import (
 	"strings"
 	"time"
 
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/client/internal"
-	"github.com/netbirdio/netbird/util/embeddedroots"
 )
 
 // HostedGrantType grant type for device flow on Hosted
@@ -61,18 +56,6 @@ func NewDeviceAuthorizationFlow(config internal.DeviceAuthProviderConfig) (*Devi
 	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
 	httpTransport.MaxIdleConns = 5
 
-	certPool, err := x509.SystemCertPool()
-	if err != nil || certPool == nil {
-		log.Debugf("System cert pool not available; falling back to embedded cert, error: %v", err)
-		certPool = embeddedroots.Get()
-	} else {
-		log.Debug("Using system certificate pool.")
-	}
-
-	httpTransport.TLSClientConfig = &tls.Config{
-		RootCAs: certPool,
-	}
-
 	httpClient := &http.Client{
 		Timeout:   10 * time.Second,
 		Transport: httpTransport,

client/internal/connect.go

@@ -31,7 +31,6 @@ import (
 	relayClient "github.com/netbirdio/netbird/relay/client"
 	signal "github.com/netbirdio/netbird/signal/client"
 	"github.com/netbirdio/netbird/util"
-	nbnet "github.com/netbirdio/netbird/util/net"
 	"github.com/netbirdio/netbird/version"
 )
 
@@ -60,8 +59,13 @@ func NewConnectClient(
 }
 
 // Run with main logic.
-func (c *ConnectClient) Run(runningChan chan error) error {
-	return c.run(MobileDependency{}, runningChan)
+func (c *ConnectClient) Run() error {
+	return c.run(MobileDependency{}, nil, nil)
+}
+
+// RunWithProbes runs the client's main logic with probes attached
+func (c *ConnectClient) RunWithProbes(probes *ProbeHolder, runningChan chan error) error {
+	return c.run(MobileDependency{}, probes, runningChan)
 }
 
 // RunOnAndroid with main logic on mobile system
@@ -80,7 +84,7 @@ func (c *ConnectClient) RunOnAndroid(
 		HostDNSAddresses:      dnsAddresses,
 		DnsReadyListener:      dnsReadyListener,
 	}
-	return c.run(mobileDependency, nil)
+	return c.run(mobileDependency, nil, nil)
 }
 
 func (c *ConnectClient) RunOniOS(
@@ -98,10 +102,10 @@ func (c *ConnectClient) RunOniOS(
 		DnsManager:            dnsManager,
 		StateFilePath:         stateFilePath,
 	}
-	return c.run(mobileDependency, nil)
+	return c.run(mobileDependency, nil, nil)
 }
 
-func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan error) error {
+func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHolder, runningChan chan error) error {
 	defer func() {
 		if r := recover(); r != nil {
 			log.Panicf("Panic occurred: %v, stack trace: %s", r, string(debug.Stack()))
@@ -110,8 +114,6 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 
 	log.Infof("starting NetBird client version %s on %s/%s", version.NetbirdVersion(), runtime.GOOS, runtime.GOARCH)
 
-	nbnet.Init()
-
 	backOff := &backoff.ExponentialBackOff{
 		InitialInterval:     time.Second,
 		RandomizationFactor: 1,
@@ -180,7 +182,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 			}
 		}()
 
-		// connect (just a connection, no stream yet) and login to Management Service to get an initial global Netbird config
+		// connect (just a connection, no stream yet) and login to Management Service to get an initial global Wiretrustee config
 		loginResp, err := loginToManagement(engineCtx, mgmClient, publicSSHKey, c.config)
 		if err != nil {
 			log.Debug(err)
@@ -202,8 +204,8 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		c.statusRecorder.UpdateLocalPeerState(localPeerState)
 
 		signalURL := fmt.Sprintf("%s://%s",
-			strings.ToLower(loginResp.GetNetbirdConfig().GetSignal().GetProtocol().String()),
-			loginResp.GetNetbirdConfig().GetSignal().GetUri(),
+			strings.ToLower(loginResp.GetWiretrusteeConfig().GetSignal().GetProtocol().String()),
+			loginResp.GetWiretrusteeConfig().GetSignal().GetUri(),
 		)
 
 		c.statusRecorder.UpdateSignalAddress(signalURL)
@@ -214,8 +216,8 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 			c.statusRecorder.MarkSignalDisconnected(err)
 		}()
 
-		// with the global Netbird config in hand connect (just a connection, no stream yet) Signal
-		signalClient, err := connectToSignal(engineCtx, loginResp.GetNetbirdConfig(), myPrivateKey)
+		// with the global Wiretrustee config in hand connect (just a connection, no stream yet) Signal
+		signalClient, err := connectToSignal(engineCtx, loginResp.GetWiretrusteeConfig(), myPrivateKey)
 		if err != nil {
 			log.Error(err)
 			return wrapErr(err)
@@ -259,7 +261,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		checks := loginResp.GetChecks()
 
 		c.engineMutex.Lock()
-		c.engine = NewEngine(engineCtx, cancel, signalClient, mgmClient, relayManager, engineConfig, mobileDependency, c.statusRecorder, checks)
+		c.engine = NewEngineWithProbes(engineCtx, cancel, signalClient, mgmClient, relayManager, engineConfig, mobileDependency, c.statusRecorder, probes, checks)
 		c.engine.SetNetworkMapPersistence(c.persistNetworkMap)
 		c.engineMutex.Unlock()
 
@@ -314,7 +316,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 }
 
 func parseRelayInfo(loginResp *mgmProto.LoginResponse) ([]string, *hmac.Token) {
-	relayCfg := loginResp.GetNetbirdConfig().GetRelay()
+	relayCfg := loginResp.GetWiretrusteeConfig().GetRelay()
 	if relayCfg == nil {
 		return nil, nil
 	}
@@ -443,7 +445,7 @@ func createEngineConfig(key wgtypes.Key, config *Config, peerConfig *mgmProto.Pe
 }
 
 // connectToSignal creates Signal Service client and established a connection
-func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourPrivateKey wgtypes.Key) (*signal.GrpcClient, error) {
+func connectToSignal(ctx context.Context, wtConfig *mgmProto.WiretrusteeConfig, ourPrivateKey wgtypes.Key) (*signal.GrpcClient, error) {
 	var sigTLSEnabled bool
 	if wtConfig.Signal.Protocol == mgmProto.HostConfig_HTTPS {
 		sigTLSEnabled = true
@@ -460,7 +462,7 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourP
 	return signalClient, nil
 }
 
-// loginToManagement creates Management Services client, establishes a connection, logs-in and gets a global Netbird config (signal, turn, stun hosts, etc)
+// loginToManagement creates Management Services client, establishes a connection, logs-in and gets a global Wiretrustee config (signal, turn, stun hosts, etc)
 func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte, config *Config) (*mgmProto.LoginResponse, error) {
 
 	serverPublicKey, err := client.GetServerPublicKey()

client/internal/dns/handler_chain.go

@@ -12,7 +12,7 @@ import (
 const (
 	PriorityDNSRoute    = 100
 	PriorityMatchDomain = 50
-	PriorityDefault     = 1
+	PriorityDefault     = 0
 )
 
 type SubdomainMatcher interface {
@@ -26,6 +26,7 @@ type HandlerEntry struct {
 	Pattern         string
 	OrigPattern     string
 	IsWildcard      bool
+	StopHandler     handlerWithStop
 	MatchSubdomains bool
 }
 
@@ -63,7 +64,7 @@ func (w *ResponseWriterChain) GetOrigPattern() string {
 }
 
 // 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) {
+func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority int, stopHandler handlerWithStop) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -77,6 +78,9 @@ func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority
 	// First remove any existing handler with same pattern (case-insensitive) and priority
 	for i := len(c.handlers) - 1; i >= 0; i-- {
 		if strings.EqualFold(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
 		}
@@ -97,33 +101,21 @@ func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority
 		Pattern:         pattern,
 		OrigPattern:     origPattern,
 		IsWildcard:      isWildcard,
+		StopHandler:     stopHandler,
 		MatchSubdomains: matchSubdomains,
 	}
 
-	pos := c.findHandlerPosition(entry)
-	c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...)
-}
-
-// findHandlerPosition determines where to insert a new handler based on priority and specificity
-func (c *HandlerChain) findHandlerPosition(newEntry HandlerEntry) int {
+	// Insert handler in priority order
+	pos := 0
 	for i, h := range c.handlers {
-		// prio first
-		if h.Priority < newEntry.Priority {
-			return i
-		}
-
-		// domain specificity next
-		if h.Priority == newEntry.Priority {
-			newDots := strings.Count(newEntry.Pattern, ".")
-			existingDots := strings.Count(h.Pattern, ".")
-			if newDots > existingDots {
-				return i
-			}
+		if h.Priority < priority {
+			pos = i
+			break
 		}
+		pos = i + 1
 	}
 
-	// add at end
-	return len(c.handlers)
+	c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...)
 }
 
 // RemoveHandler removes a handler for the given pattern and priority
@@ -137,6 +129,9 @@ func (c *HandlerChain) RemoveHandler(pattern string, priority int) {
 	for i := len(c.handlers) - 1; i >= 0; i-- {
 		entry := c.handlers[i]
 		if strings.EqualFold(entry.OrigPattern, pattern) && entry.Priority == priority {
+			if entry.StopHandler != nil {
+				entry.StopHandler.stop()
+			}
 			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
 			return
 		}
@@ -172,8 +167,8 @@ func (c *HandlerChain) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	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 match_subdomain=%v priority=%d",
-				h.Pattern, h.OrigPattern, h.IsWildcard, h.MatchSubdomains, h.Priority)
+			log.Tracef("  - pattern: domain=%s original: domain=%s wildcard=%v priority=%d",
+				h.Pattern, h.OrigPattern, h.IsWildcard, h.Priority)
 		}
 	}
 
@@ -198,13 +193,13 @@ func (c *HandlerChain) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 		}
 
 		if !matched {
-			log.Tracef("trying domain match: request: domain=%s pattern: domain=%s wildcard=%v match_subdomain=%v priority=%d matched=false",
-				qname, entry.OrigPattern, entry.MatchSubdomains, entry.IsWildcard, entry.Priority)
+			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 priority=%d",
-			qname, entry.OrigPattern, entry.IsWildcard, entry.MatchSubdomains, entry.Priority)
+		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,

client/internal/dns/handler_chain_test.go

@@ -21,9 +21,9 @@ func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
 	dnsRouteHandler := &nbdns.MockHandler{}
 
 	// Setup handlers with different priorities
-	chain.AddHandler("example.com.", defaultHandler, nbdns.PriorityDefault)
-	chain.AddHandler("example.com.", matchDomainHandler, nbdns.PriorityMatchDomain)
-	chain.AddHandler("example.com.", dnsRouteHandler, nbdns.PriorityDNSRoute)
+	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)
@@ -138,7 +138,7 @@ func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
 				pattern = "*." + tt.handlerDomain[2:]
 			}
 
-			chain.AddHandler(pattern, handler, nbdns.PriorityDefault)
+			chain.AddHandler(pattern, handler, nbdns.PriorityDefault, nil)
 
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryDomain, dns.TypeA)
@@ -253,7 +253,7 @@ func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
 					handler.On("ServeDNS", mock.Anything, mock.Anything).Maybe()
 				}
 
-				chain.AddHandler(tt.handlers[i].pattern, handler, tt.handlers[i].priority)
+				chain.AddHandler(tt.handlers[i].pattern, handler, tt.handlers[i].priority, nil)
 			}
 
 			// Create and execute request
@@ -280,9 +280,9 @@ func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
 	handler3 := &nbdns.MockHandler{}
 
 	// Add handlers in priority order
-	chain.AddHandler("example.com.", handler1, nbdns.PriorityDNSRoute)
-	chain.AddHandler("example.com.", handler2, nbdns.PriorityMatchDomain)
-	chain.AddHandler("example.com.", handler3, nbdns.PriorityDefault)
+	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)
@@ -416,7 +416,7 @@ func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 				if op.action == "add" {
 					handler := &nbdns.MockHandler{}
 					handlers[op.priority] = handler
-					chain.AddHandler(op.pattern, handler, op.priority)
+					chain.AddHandler(op.pattern, handler, op.priority, nil)
 				} else {
 					chain.RemoveHandler(op.pattern, op.priority)
 				}
@@ -471,9 +471,9 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	r.SetQuestion(testQuery, dns.TypeA)
 
 	// Add handlers in mixed order
-	chain.AddHandler(testDomain, defaultHandler, nbdns.PriorityDefault)
-	chain.AddHandler(testDomain, routeHandler, nbdns.PriorityDNSRoute)
-	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain)
+	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{}}
@@ -653,7 +653,7 @@ func TestHandlerChain_CaseSensitivity(t *testing.T) {
 					handler = mockHandler
 				}
 
-				chain.AddHandler(pattern, handler, h.priority)
+				chain.AddHandler(pattern, handler, h.priority, nil)
 			}
 
 			// Execute request
@@ -677,156 +677,3 @@ func TestHandlerChain_CaseSensitivity(t *testing.T) {
 		})
 	}
 }
-
-func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) {
-	tests := []struct {
-		name     string
-		scenario string
-		ops      []struct {
-			action    string
-			pattern   string
-			priority  int
-			subdomain bool
-		}
-		query         string
-		expectedMatch string
-	}{
-		{
-			name:     "more specific domain matches first",
-			scenario: "sub.example.com should match before example.com",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "sub.example.com.", nbdns.PriorityMatchDomain, false},
-			},
-			query:         "sub.example.com.",
-			expectedMatch: "sub.example.com.",
-		},
-		{
-			name:     "more specific domain matches first, both match subdomains",
-			scenario: "sub.example.com should match before example.com",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "sub.example.com.", nbdns.PriorityMatchDomain, true},
-			},
-			query:         "sub.example.com.",
-			expectedMatch: "sub.example.com.",
-		},
-		{
-			name:     "maintain specificity order after removal",
-			scenario: "after removing most specific, should fall back to less specific",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "sub.example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "test.sub.example.com.", nbdns.PriorityMatchDomain, false},
-				{"remove", "test.sub.example.com.", nbdns.PriorityMatchDomain, false},
-			},
-			query:         "test.sub.example.com.",
-			expectedMatch: "sub.example.com.",
-		},
-		{
-			name:     "priority overrides specificity",
-			scenario: "less specific domain with higher priority should match first",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "sub.example.com.", nbdns.PriorityMatchDomain, false},
-				{"add", "example.com.", nbdns.PriorityDNSRoute, true},
-			},
-			query:         "sub.example.com.",
-			expectedMatch: "example.com.",
-		},
-		{
-			name:     "equal priority respects specificity",
-			scenario: "with equal priority, more specific domain should match",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "other.example.com.", nbdns.PriorityMatchDomain, true},
-				{"add", "sub.example.com.", nbdns.PriorityMatchDomain, false},
-			},
-			query:         "sub.example.com.",
-			expectedMatch: "sub.example.com.",
-		},
-		{
-			name:     "specific matches before wildcard",
-			scenario: "specific domain should match before wildcard at same priority",
-			ops: []struct {
-				action    string
-				pattern   string
-				priority  int
-				subdomain bool
-			}{
-				{"add", "*.example.com.", nbdns.PriorityDNSRoute, false},
-				{"add", "sub.example.com.", nbdns.PriorityDNSRoute, false},
-			},
-			query:         "sub.example.com.",
-			expectedMatch: "sub.example.com.",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			chain := nbdns.NewHandlerChain()
-			handlers := make(map[string]*nbdns.MockSubdomainHandler)
-
-			for _, op := range tt.ops {
-				if op.action == "add" {
-					handler := &nbdns.MockSubdomainHandler{Subdomains: op.subdomain}
-					handlers[op.pattern] = handler
-					chain.AddHandler(op.pattern, handler, op.priority)
-				} else {
-					chain.RemoveHandler(op.pattern, op.priority)
-				}
-			}
-
-			r := new(dns.Msg)
-			r.SetQuestion(tt.query, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
-
-			// Setup handler expectations
-			for pattern, handler := range handlers {
-				if pattern == tt.expectedMatch {
-					handler.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
-						w := args.Get(0).(dns.ResponseWriter)
-						r := args.Get(1).(*dns.Msg)
-						resp := new(dns.Msg)
-						resp.SetReply(r)
-						assert.NoError(t, w.WriteMsg(resp))
-					}).Once()
-				}
-			}
-
-			chain.ServeDNS(w, r)
-
-			for pattern, handler := range handlers {
-				if pattern == tt.expectedMatch {
-					handler.AssertNumberOfCalls(t, "ServeDNS", 1)
-				} else {
-					handler.AssertNumberOfCalls(t, "ServeDNS", 0)
-				}
-			}
-		})
-	}
-}

client/internal/dns/host_windows.go

@@ -1,51 +1,35 @@
 package dns
 
 import (
-	"errors"
 	"fmt"
 	"io"
 	"strings"
-	"syscall"
 
-	"github.com/hashicorp/go-multierror"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/sys/windows/registry"
 
-	nberrors "github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
-var (
-	userenv = syscall.NewLazyDLL("userenv.dll")
-
-	// https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-refreshpolicyex
-	refreshPolicyExFn = userenv.NewProc("RefreshPolicyEx")
-)
-
 const (
-	dnsPolicyConfigMatchPath    = `SYSTEM\CurrentControlSet\Services\Dnscache\Parameters\DnsPolicyConfig\NetBird-Match`
-	gpoDnsPolicyRoot            = `SOFTWARE\Policies\Microsoft\Windows NT\DNSClient`
-	gpoDnsPolicyConfigMatchPath = gpoDnsPolicyRoot + `\DnsPolicyConfig\NetBird-Match`
-
+	dnsPolicyConfigMatchPath            = `SYSTEM\CurrentControlSet\Services\Dnscache\Parameters\DnsPolicyConfig\NetBird-Match`
 	dnsPolicyConfigVersionKey           = "Version"
 	dnsPolicyConfigVersionValue         = 2
 	dnsPolicyConfigNameKey              = "Name"
 	dnsPolicyConfigGenericDNSServersKey = "GenericDNSServers"
 	dnsPolicyConfigConfigOptionsKey     = "ConfigOptions"
 	dnsPolicyConfigConfigOptionsValue   = 0x8
+)
 
+const (
 	interfaceConfigPath          = `SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces`
 	interfaceConfigNameServerKey = "NameServer"
 	interfaceConfigSearchListKey = "SearchList"
-
-	// RP_FORCE: Reapply all policies even if no policy change was detected
-	rpForce = 0x1
 )
 
 type registryConfigurator struct {
 	guid       string
 	routingAll bool
-	gpo        bool
 }
 
 func newHostManager(wgInterface WGIface) (*registryConfigurator, error) {
@@ -53,20 +37,12 @@ func newHostManager(wgInterface WGIface) (*registryConfigurator, error) {
 	if err != nil {
 		return nil, err
 	}
+	return newHostManagerWithGuid(guid)
+}
 
-	var useGPO bool
-	k, err := registry.OpenKey(registry.LOCAL_MACHINE, gpoDnsPolicyRoot, registry.QUERY_VALUE)
-	if err != nil {
-		log.Debugf("failed to open GPO DNS policy root: %v", err)
-	} else {
-		closer(k)
-		useGPO = true
-		log.Infof("detected GPO DNS policy configuration, using policy store")
-	}
-
+func newHostManagerWithGuid(guid string) (*registryConfigurator, error) {
 	return &registryConfigurator{
 		guid: guid,
-		gpo:  useGPO,
 	}, nil
 }
 
@@ -75,23 +51,30 @@ func (r *registryConfigurator) supportCustomPort() bool {
 }
 
 func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+	var err error
 	if config.RouteAll {
-		if err := r.addDNSSetupForAll(config.ServerIP); err != nil {
+		err = r.addDNSSetupForAll(config.ServerIP)
+		if err != nil {
 			return fmt.Errorf("add dns setup: %w", err)
 		}
 	} else if r.routingAll {
-		if err := r.deleteInterfaceRegistryKeyProperty(interfaceConfigNameServerKey); err != nil {
+		err = r.deleteInterfaceRegistryKeyProperty(interfaceConfigNameServerKey)
+		if err != nil {
 			return fmt.Errorf("delete interface registry key property: %w", err)
 		}
 		r.routingAll = false
 		log.Infof("removed %s as main DNS forwarder for this peer", config.ServerIP)
 	}
 
-	if err := stateManager.UpdateState(&ShutdownState{Guid: r.guid, GPO: r.gpo}); err != nil {
+	if err := stateManager.UpdateState(&ShutdownState{Guid: r.guid}); err != nil {
 		log.Errorf("failed to update shutdown state: %s", err)
 	}
 
-	var searchDomains, matchDomains []string
+	var (
+		searchDomains []string
+		matchDomains  []string
+	)
+
 	for _, dConf := range config.Domains {
 		if dConf.Disabled {
 			continue
@@ -103,16 +86,16 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 	}
 
 	if len(matchDomains) != 0 {
-		if err := r.addDNSMatchPolicy(matchDomains, config.ServerIP); err != nil {
-			return fmt.Errorf("add dns match policy: %w", err)
-		}
+		err = r.addDNSMatchPolicy(matchDomains, config.ServerIP)
 	} else {
-		if err := r.removeDNSMatchPolicies(); err != nil {
-			return fmt.Errorf("remove dns match policies: %w", err)
-		}
+		err = removeRegistryKeyFromDNSPolicyConfig(dnsPolicyConfigMatchPath)
+	}
+	if err != nil {
+		return fmt.Errorf("add dns match policy: %w", err)
 	}
 
-	if err := r.updateSearchDomains(searchDomains); err != nil {
+	err = r.updateSearchDomains(searchDomains)
+	if err != nil {
 		return fmt.Errorf("update search domains: %w", err)
 	}
 
@@ -120,8 +103,9 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 }
 
 func (r *registryConfigurator) addDNSSetupForAll(ip string) error {
-	if err := r.setInterfaceRegistryKeyStringValue(interfaceConfigNameServerKey, ip); err != nil {
-		return fmt.Errorf("adding dns setup for all failed: %w", err)
+	err := r.setInterfaceRegistryKeyStringValue(interfaceConfigNameServerKey, ip)
+	if err != nil {
+		return fmt.Errorf("adding dns setup for all failed with error: %w", err)
 	}
 	r.routingAll = true
 	log.Infof("configured %s:53 as main DNS forwarder for this peer", ip)
@@ -129,54 +113,64 @@ func (r *registryConfigurator) addDNSSetupForAll(ip string) error {
 }
 
 func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip string) error {
-	// if the gpo key is present, we need to put our DNS settings there, otherwise our config might be ignored
-	// see https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-gpnrpt/8cc31cb9-20cb-4140-9e85-3e08703b4745
-	policyPath := dnsPolicyConfigMatchPath
-	if r.gpo {
-		policyPath = gpoDnsPolicyConfigMatchPath
-	}
-
-	if err := removeRegistryKeyFromDNSPolicyConfig(policyPath); err != nil {
-		return fmt.Errorf("remove existing dns policy: %w", err)
-	}
-
-	regKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, policyPath, registry.SET_VALUE)
-	if err != nil {
-		return fmt.Errorf("create registry key HKEY_LOCAL_MACHINE\\%s: %w", policyPath, err)
-	}
-	defer closer(regKey)
-
-	if err := regKey.SetDWordValue(dnsPolicyConfigVersionKey, dnsPolicyConfigVersionValue); err != nil {
-		return fmt.Errorf("set %s: %w", dnsPolicyConfigVersionKey, err)
-	}
-
-	if err := regKey.SetStringsValue(dnsPolicyConfigNameKey, domains); err != nil {
-		return fmt.Errorf("set %s: %w", dnsPolicyConfigNameKey, err)
-	}
-
-	if err := regKey.SetStringValue(dnsPolicyConfigGenericDNSServersKey, ip); err != nil {
-		return fmt.Errorf("set %s: %w", dnsPolicyConfigGenericDNSServersKey, err)
-	}
-
-	if err := regKey.SetDWordValue(dnsPolicyConfigConfigOptionsKey, dnsPolicyConfigConfigOptionsValue); err != nil {
-		return fmt.Errorf("set %s: %w", dnsPolicyConfigConfigOptionsKey, err)
-	}
-
-	if r.gpo {
-		if err := refreshGroupPolicy(); err != nil {
-			log.Warnf("failed to refresh group policy: %v", err)
+	_, err := registry.OpenKey(registry.LOCAL_MACHINE, dnsPolicyConfigMatchPath, registry.QUERY_VALUE)
+	if err == nil {
+		err = registry.DeleteKey(registry.LOCAL_MACHINE, dnsPolicyConfigMatchPath)
+		if err != nil {
+			return fmt.Errorf("unable to remove existing key from registry, key: HKEY_LOCAL_MACHINE\\%s, error: %w", dnsPolicyConfigMatchPath, err)
 		}
 	}
 
-	log.Infof("added %d match domains. Domain list: %s", len(domains), domains)
+	regKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, dnsPolicyConfigMatchPath, registry.SET_VALUE)
+	if err != nil {
+		return fmt.Errorf("unable to create registry key, key: HKEY_LOCAL_MACHINE\\%s, error: %w", dnsPolicyConfigMatchPath, err)
+	}
+
+	err = regKey.SetDWordValue(dnsPolicyConfigVersionKey, dnsPolicyConfigVersionValue)
+	if err != nil {
+		return fmt.Errorf("unable to set registry value for %s, error: %w", dnsPolicyConfigVersionKey, err)
+	}
+
+	err = regKey.SetStringsValue(dnsPolicyConfigNameKey, domains)
+	if err != nil {
+		return fmt.Errorf("unable to set registry value for %s, error: %w", dnsPolicyConfigNameKey, err)
+	}
+
+	err = regKey.SetStringValue(dnsPolicyConfigGenericDNSServersKey, ip)
+	if err != nil {
+		return fmt.Errorf("unable to set registry value for %s, error: %w", dnsPolicyConfigGenericDNSServersKey, err)
+	}
+
+	err = regKey.SetDWordValue(dnsPolicyConfigConfigOptionsKey, dnsPolicyConfigConfigOptionsValue)
+	if err != nil {
+		return fmt.Errorf("unable to set registry value for %s, error: %w", dnsPolicyConfigConfigOptionsKey, err)
+	}
+
+	log.Infof("added %d match domains to the state. Domain list: %s", len(domains), domains)
+
+	return nil
+}
+
+func (r *registryConfigurator) restoreHostDNS() error {
+	if err := removeRegistryKeyFromDNSPolicyConfig(dnsPolicyConfigMatchPath); err != nil {
+		log.Errorf("remove registry key from dns policy config: %s", err)
+	}
+
+	if err := r.deleteInterfaceRegistryKeyProperty(interfaceConfigSearchListKey); err != nil {
+		return fmt.Errorf("remove interface registry key: %w", err)
+	}
+
 	return nil
 }
 
 func (r *registryConfigurator) updateSearchDomains(domains []string) error {
-	if err := r.setInterfaceRegistryKeyStringValue(interfaceConfigSearchListKey, strings.Join(domains, ",")); err != nil {
-		return fmt.Errorf("update search domains: %w", err)
+	err := r.setInterfaceRegistryKeyStringValue(interfaceConfigSearchListKey, strings.Join(domains, ","))
+	if err != nil {
+		return fmt.Errorf("adding search domain failed with error: %w", err)
 	}
-	log.Infof("updated search domains: %s", domains)
+
+	log.Infof("updated the search domains in the registry with %d domains. Domain list: %s", len(domains), domains)
+
 	return nil
 }
 
@@ -187,9 +181,11 @@ func (r *registryConfigurator) setInterfaceRegistryKeyStringValue(key, value str
 	}
 	defer closer(regKey)
 
-	if err := regKey.SetStringValue(key, value); err != nil {
-		return fmt.Errorf("set key %s=%s: %w", key, value, err)
+	err = regKey.SetStringValue(key, value)
+	if err != nil {
+		return fmt.Errorf("applying key %s with value \"%s\" for interface failed with error: %w", key, value, err)
 	}
+
 	return nil
 }
 
@@ -200,91 +196,43 @@ func (r *registryConfigurator) deleteInterfaceRegistryKeyProperty(propertyKey st
 	}
 	defer closer(regKey)
 
-	if err := regKey.DeleteValue(propertyKey); err != nil {
-		return fmt.Errorf("delete registry key %s: %w", propertyKey, err)
+	err = regKey.DeleteValue(propertyKey)
+	if err != nil {
+		return fmt.Errorf("deleting registry key %s for interface failed with error: %w", propertyKey, err)
 	}
+
 	return nil
 }
 
 func (r *registryConfigurator) getInterfaceRegistryKey() (registry.Key, error) {
+	var regKey registry.Key
+
 	regKeyPath := interfaceConfigPath + "\\" + r.guid
+
 	regKey, err := registry.OpenKey(registry.LOCAL_MACHINE, regKeyPath, registry.SET_VALUE)
 	if err != nil {
-		return regKey, fmt.Errorf("open HKEY_LOCAL_MACHINE\\%s: %w", regKeyPath, err)
+		return regKey, fmt.Errorf("unable to open the interface registry key, key: HKEY_LOCAL_MACHINE\\%s, error: %w", regKeyPath, err)
 	}
+
 	return regKey, nil
 }
 
-func (r *registryConfigurator) restoreHostDNS() error {
-	if err := r.removeDNSMatchPolicies(); err != nil {
-		log.Errorf("remove dns match policies: %s", err)
-	}
-
-	if err := r.deleteInterfaceRegistryKeyProperty(interfaceConfigSearchListKey); err != nil {
-		return fmt.Errorf("remove interface registry key: %w", err)
-	}
-
-	return nil
-}
-
-func (r *registryConfigurator) removeDNSMatchPolicies() error {
-	var merr *multierror.Error
-	if err := removeRegistryKeyFromDNSPolicyConfig(dnsPolicyConfigMatchPath); err != nil {
-		merr = multierror.Append(merr, fmt.Errorf("remove local registry key: %w", err))
-	}
-
-	if err := removeRegistryKeyFromDNSPolicyConfig(gpoDnsPolicyConfigMatchPath); err != nil {
-		merr = multierror.Append(merr, fmt.Errorf("remove GPO registry key: %w", err))
-	}
-
-	if err := refreshGroupPolicy(); err != nil {
-		merr = multierror.Append(merr, fmt.Errorf("refresh group policy: %w", err))
-	}
-
-	return nberrors.FormatErrorOrNil(merr)
-}
-
 func (r *registryConfigurator) restoreUncleanShutdownDNS() error {
-	return r.restoreHostDNS()
+	if err := r.restoreHostDNS(); err != nil {
+		return fmt.Errorf("restoring dns via registry: %w", err)
+	}
+	return nil
 }
 
 func removeRegistryKeyFromDNSPolicyConfig(regKeyPath string) error {
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, regKeyPath, registry.QUERY_VALUE)
-	if err != nil {
-		log.Debugf("failed to open HKEY_LOCAL_MACHINE\\%s: %v", regKeyPath, err)
-		return nil
-	}
-
-	closer(k)
-	if err := registry.DeleteKey(registry.LOCAL_MACHINE, regKeyPath); err != nil {
-		return fmt.Errorf("delete HKEY_LOCAL_MACHINE\\%s: %w", regKeyPath, err)
-	}
-
-	return nil
-}
-
-func refreshGroupPolicy() error {
-	// refreshPolicyExFn.Call() panics if the func is not found
-	defer func() {
-		if r := recover(); r != nil {
-			log.Errorf("Recovered from panic: %v", r)
+	if err == nil {
+		defer closer(k)
+		err = registry.DeleteKey(registry.LOCAL_MACHINE, regKeyPath)
+		if err != nil {
+			return fmt.Errorf("unable to remove existing key from registry, key: HKEY_LOCAL_MACHINE\\%s, error: %w", regKeyPath, err)
 		}
-	}()
-
-	ret, _, err := refreshPolicyExFn.Call(
-		// bMachine = TRUE (computer policy)
-		uintptr(1),
-		// dwOptions = RP_FORCE
-		uintptr(rpForce),
-	)
-
-	if ret == 0 {
-		if err != nil && !errors.Is(err, syscall.Errno(0)) {
-			return fmt.Errorf("RefreshPolicyEx failed: %w", err)
-		}
-		return fmt.Errorf("RefreshPolicyEx failed")
 	}
-
 	return nil
 }
 

client/internal/dns/local.go

@@ -29,15 +29,10 @@ func (d *localResolver) String() string {
 	return fmt.Sprintf("local resolver [%d records]", len(d.registeredMap))
 }
 
-// ID returns the unique handler ID
-func (d *localResolver) id() handlerID {
-	return "local-resolver"
-}
-
 // ServeDNS handles a DNS request
 func (d *localResolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	if len(r.Question) > 0 {
-		log.Tracef("received local question: domain=%s type=%v class=%v", r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
+		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{}

client/internal/dns/server.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/netip"
 	"runtime"
+	"strings"
 	"sync"
 
 	"github.com/miekg/dns"
@@ -41,12 +42,7 @@ type Server interface {
 	ProbeAvailability()
 }
 
-type handlerID string
-
-type nsGroupsByDomain struct {
-	domain string
-	groups []*nbdns.NameServerGroup
-}
+type registeredHandlerMap map[string]handlerWithStop
 
 // DefaultServer dns server object
 type DefaultServer struct {
@@ -56,6 +52,7 @@ type DefaultServer struct {
 	mux                sync.Mutex
 	service            service
 	dnsMuxMap          registeredHandlerMap
+	handlerPriorities  map[string]int
 	localResolver      *localResolver
 	wgInterface        WGIface
 	hostManager        hostManager
@@ -80,17 +77,14 @@ type handlerWithStop interface {
 	dns.Handler
 	stop()
 	probeAvailability()
-	id() handlerID
 }
 
-type handlerWrapper struct {
+type muxUpdate struct {
 	domain   string
 	handler  handlerWithStop
 	priority int
 }
 
-type registeredHandlerMap map[handlerID]handlerWrapper
-
 // NewDefaultServer returns a new dns server
 func NewDefaultServer(
 	ctx context.Context,
@@ -164,12 +158,13 @@ func newDefaultServer(
 ) *DefaultServer {
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
-		ctx:          ctx,
-		ctxCancel:    stop,
-		disableSys:   disableSys,
-		service:      dnsService,
-		handlerChain: NewHandlerChain(),
-		dnsMuxMap:    make(registeredHandlerMap),
+		ctx:               ctx,
+		ctxCancel:         stop,
+		disableSys:        disableSys,
+		service:           dnsService,
+		handlerChain:      NewHandlerChain(),
+		dnsMuxMap:         make(registeredHandlerMap),
+		handlerPriorities: make(map[string]int),
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
@@ -197,7 +192,8 @@ func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, p
 			log.Warn("skipping empty domain")
 			continue
 		}
-		s.handlerChain.AddHandler(domain, handler, priority)
+		s.handlerChain.AddHandler(domain, handler, priority, nil)
+		s.handlerPriorities[domain] = priority
 		s.service.RegisterMux(nbdns.NormalizeZone(domain), s.handlerChain)
 	}
 }
@@ -213,15 +209,14 @@ func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
 	log.Debugf("deregistering handler %v with priority %d", domains, priority)
 
 	for _, domain := range domains {
-		if domain == "" {
-			log.Warn("skipping empty domain")
-			continue
-		}
-
 		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))
 		}
 	}
@@ -288,24 +283,14 @@ func (s *DefaultServer) Stop() {
 
 // OnUpdatedHostDNSServer update the DNS servers addresses for root zones
 // It will be applied if the mgm server do not enforce DNS settings for root zone
-
 func (s *DefaultServer) OnUpdatedHostDNSServer(hostsDnsList []string) {
 	s.hostsDNSHolder.set(hostsDnsList)
 
-	// Check if there's any root handler
-	var hasRootHandler bool
-	for _, handler := range s.dnsMuxMap {
-		if handler.domain == nbdns.RootZone {
-			hasRootHandler = true
-			break
-		}
-	}
-
-	if hasRootHandler {
+	_, ok := s.dnsMuxMap[nbdns.RootZone]
+	if ok {
 		log.Debugf("on new host DNS config but skip to apply it")
 		return
 	}
-
 	log.Debugf("update host DNS settings: %+v", hostsDnsList)
 	s.addHostRootZone()
 }
@@ -379,7 +364,7 @@ func (s *DefaultServer) ProbeAvailability() {
 		go func(mux handlerWithStop) {
 			defer wg.Done()
 			mux.probeAvailability()
-		}(mux.handler)
+		}(mux)
 	}
 	wg.Wait()
 }
@@ -434,8 +419,8 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	return nil
 }
 
-func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone) ([]handlerWrapper, map[string]nbdns.SimpleRecord, error) {
-	var muxUpdates []handlerWrapper
+func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone) ([]muxUpdate, map[string]nbdns.SimpleRecord, error) {
+	var muxUpdates []muxUpdate
 	localRecords := make(map[string]nbdns.SimpleRecord, 0)
 
 	for _, customZone := range customZones {
@@ -443,7 +428,7 @@ func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone)
 			return nil, nil, fmt.Errorf("received an empty list of records")
 		}
 
-		muxUpdates = append(muxUpdates, handlerWrapper{
+		muxUpdates = append(muxUpdates, muxUpdate{
 			domain:   customZone.Domain,
 			handler:  s.localResolver,
 			priority: PriorityMatchDomain,
@@ -461,59 +446,15 @@ func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone)
 	return muxUpdates, localRecords, nil
 }
 
-func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.NameServerGroup) ([]handlerWrapper, error) {
-	var muxUpdates []handlerWrapper
+func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.NameServerGroup) ([]muxUpdate, error) {
 
+	var muxUpdates []muxUpdate
 	for _, nsGroup := range nameServerGroups {
 		if len(nsGroup.NameServers) == 0 {
 			log.Warn("received a nameserver group with empty nameserver list")
 			continue
 		}
 
-		if !nsGroup.Primary && len(nsGroup.Domains) == 0 {
-			return nil, fmt.Errorf("received a non primary nameserver group with an empty domain list")
-		}
-
-		for _, domain := range nsGroup.Domains {
-			if domain == "" {
-				return nil, fmt.Errorf("received a nameserver group with an empty domain element")
-			}
-		}
-	}
-
-	groupedNS := groupNSGroupsByDomain(nameServerGroups)
-
-	for _, domainGroup := range groupedNS {
-		basePriority := PriorityMatchDomain
-		if domainGroup.domain == nbdns.RootZone {
-			basePriority = PriorityDefault
-		}
-
-		updates, err := s.createHandlersForDomainGroup(domainGroup, basePriority)
-		if err != nil {
-			return nil, err
-		}
-		muxUpdates = append(muxUpdates, updates...)
-	}
-
-	return muxUpdates, nil
-}
-
-func (s *DefaultServer) createHandlersForDomainGroup(domainGroup nsGroupsByDomain, basePriority int) ([]handlerWrapper, error) {
-	var muxUpdates []handlerWrapper
-
-	for i, nsGroup := range domainGroup.groups {
-		// Decrement priority by handler index (0, 1, 2, ...) to avoid conflicts
-		priority := basePriority - i
-
-		// Check if we're about to overlap with the next priority tier
-		if basePriority == PriorityMatchDomain && priority <= PriorityDefault {
-			log.Warnf("too many handlers for domain=%s, would overlap with default priority tier (diff=%d). Skipping remaining handlers",
-				domainGroup.domain, PriorityMatchDomain-PriorityDefault)
-			break
-		}
-
-		log.Debugf("creating handler for domain=%s with priority=%d", domainGroup.domain, priority)
 		handler, err := newUpstreamResolver(
 			s.ctx,
 			s.wgInterface.Name(),
@@ -521,12 +462,10 @@ func (s *DefaultServer) createHandlersForDomainGroup(domainGroup nsGroupsByDomai
 			s.wgInterface.Address().Network,
 			s.statusRecorder,
 			s.hostsDNSHolder,
-			domainGroup.domain,
 		)
 		if err != nil {
-			return nil, fmt.Errorf("create upstream resolver: %v", err)
+			return nil, fmt.Errorf("unable to create a new upstream resolver, error: %v", err)
 		}
-
 		for _, ns := range nsGroup.NameServers {
 			if ns.NSType != nbdns.UDPNameServerType {
 				log.Warnf("skipping nameserver %s with type %s, this peer supports only %s",
@@ -550,47 +489,78 @@ func (s *DefaultServer) createHandlersForDomainGroup(domainGroup nsGroupsByDomai
 		// after some period defined by upstream it tries to reactivate self by calling this hook
 		// everything we need here is just to re-apply current configuration because it already
 		// contains this upstream settings (temporal deactivation not removed it)
-		handler.deactivate, handler.reactivate = s.upstreamCallbacks(nsGroup, handler, priority)
+		handler.deactivate, handler.reactivate = s.upstreamCallbacks(nsGroup, handler)
 
-		muxUpdates = append(muxUpdates, handlerWrapper{
-			domain:   domainGroup.domain,
-			handler:  handler,
-			priority: priority,
-		})
+		if nsGroup.Primary {
+			muxUpdates = append(muxUpdates, muxUpdate{
+				domain:   nbdns.RootZone,
+				handler:  handler,
+				priority: PriorityDefault,
+			})
+			continue
+		}
+
+		if len(nsGroup.Domains) == 0 {
+			handler.stop()
+			return nil, fmt.Errorf("received a non primary nameserver group with an empty domain list")
+		}
+
+		for _, domain := range nsGroup.Domains {
+			if domain == "" {
+				handler.stop()
+				return nil, fmt.Errorf("received a nameserver group with an empty domain element")
+			}
+			muxUpdates = append(muxUpdates, muxUpdate{
+				domain:   domain,
+				handler:  handler,
+				priority: PriorityMatchDomain,
+			})
+		}
 	}
 
 	return muxUpdates, nil
 }
 
-func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
-	// this will introduce a short period of time when the server is not able to handle DNS requests
-	for _, existing := range s.dnsMuxMap {
-		s.deregisterHandler([]string{existing.domain}, existing.priority)
-		existing.handler.stop()
-	}
-
+func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 	muxUpdateMap := make(registeredHandlerMap)
-	var containsRootUpdate bool
+	handlersByPriority := make(map[string]int)
 
+	var isContainRootUpdate bool
+
+	// First register new handlers
 	for _, update := range muxUpdates {
-		if update.domain == nbdns.RootZone {
-			containsRootUpdate = true
-		}
 		s.registerHandler([]string{update.domain}, update.handler, update.priority)
-		muxUpdateMap[update.handler.id()] = update
+		muxUpdateMap[update.domain] = update.handler
+		handlersByPriority[update.domain] = update.priority
+
+		if existingHandler, ok := s.dnsMuxMap[update.domain]; ok {
+			existingHandler.stop()
+		}
+
+		if update.domain == nbdns.RootZone {
+			isContainRootUpdate = true
+		}
 	}
 
-	// If there's no root update and we had a root handler, restore it
-	if !containsRootUpdate {
-		for _, existing := range s.dnsMuxMap {
-			if existing.domain == nbdns.RootZone {
+	// Then deregister old handlers not in the update
+	for key, existingHandler := range s.dnsMuxMap {
+		_, found := muxUpdateMap[key]
+		if !found {
+			if !isContainRootUpdate && key == nbdns.RootZone {
 				s.addHostRootZone()
-				break
+				existingHandler.stop()
+			} else {
+				existingHandler.stop()
+				// 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) {
@@ -623,7 +593,6 @@ func getNSHostPort(ns nbdns.NameServer) string {
 func (s *DefaultServer) upstreamCallbacks(
 	nsGroup *nbdns.NameServerGroup,
 	handler dns.Handler,
-	priority int,
 ) (deactivate func(error), reactivate func()) {
 	var removeIndex map[string]int
 	deactivate = func(err error) {
@@ -640,13 +609,13 @@ func (s *DefaultServer) upstreamCallbacks(
 		if nsGroup.Primary {
 			removeIndex[nbdns.RootZone] = -1
 			s.currentConfig.RouteAll = false
-			s.deregisterHandler([]string{nbdns.RootZone}, priority)
+			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.deregisterHandler([]string{item.Domain}, priority)
+				s.deregisterHandler([]string{item.Domain}, PriorityMatchDomain)
 				removeIndex[item.Domain] = i
 			}
 		}
@@ -666,8 +635,8 @@ func (s *DefaultServer) upstreamCallbacks(
 		}
 
 		s.updateNSState(nsGroup, err, false)
-	}
 
+	}
 	reactivate = func() {
 		s.mux.Lock()
 		defer s.mux.Unlock()
@@ -677,7 +646,7 @@ func (s *DefaultServer) upstreamCallbacks(
 				continue
 			}
 			s.currentConfig.Domains[i].Disabled = false
-			s.registerHandler([]string{domain}, handler, priority)
+			s.registerHandler([]string{domain}, handler, PriorityMatchDomain)
 		}
 
 		l := log.WithField("nameservers", nsGroup.NameServers)
@@ -685,7 +654,7 @@ func (s *DefaultServer) upstreamCallbacks(
 
 		if nsGroup.Primary {
 			s.currentConfig.RouteAll = true
-			s.registerHandler([]string{nbdns.RootZone}, handler, priority)
+			s.registerHandler([]string{nbdns.RootZone}, handler, PriorityDefault)
 		}
 
 		if s.hostManager != nil {
@@ -707,7 +676,6 @@ func (s *DefaultServer) addHostRootZone() {
 		s.wgInterface.Address().Network,
 		s.statusRecorder,
 		s.hostsDNSHolder,
-		nbdns.RootZone,
 	)
 	if err != nil {
 		log.Errorf("unable to create a new upstream resolver, error: %v", err)
@@ -764,34 +732,5 @@ func generateGroupKey(nsGroup *nbdns.NameServerGroup) string {
 	for _, ns := range nsGroup.NameServers {
 		servers = append(servers, fmt.Sprintf("%s:%d", ns.IP, ns.Port))
 	}
-	return fmt.Sprintf("%v_%v", servers, nsGroup.Domains)
-}
-
-// groupNSGroupsByDomain groups nameserver groups by their match domains
-func groupNSGroupsByDomain(nsGroups []*nbdns.NameServerGroup) []nsGroupsByDomain {
-	domainMap := make(map[string][]*nbdns.NameServerGroup)
-
-	for _, group := range nsGroups {
-		if group.Primary {
-			domainMap[nbdns.RootZone] = append(domainMap[nbdns.RootZone], group)
-			continue
-		}
-
-		for _, domain := range group.Domains {
-			if domain == "" {
-				continue
-			}
-			domainMap[domain] = append(domainMap[domain], group)
-		}
-	}
-
-	var result []nsGroupsByDomain
-	for domain, groups := range domainMap {
-		result = append(result, nsGroupsByDomain{
-			domain: domain,
-			groups: groups,
-		})
-	}
-
-	return result
+	return fmt.Sprintf("%s_%s_%s", nsGroup.ID, nsGroup.Name, strings.Join(servers, ","))
 }

client/internal/dns/server_test.go

@@ -13,7 +13,6 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
@@ -89,18 +88,6 @@ func init() {
 	formatter.SetTextFormatter(log.StandardLogger())
 }
 
-func generateDummyHandler(domain string, servers []nbdns.NameServer) *upstreamResolverBase {
-	var srvs []string
-	for _, srv := range servers {
-		srvs = append(srvs, getNSHostPort(srv))
-	}
-	return &upstreamResolverBase{
-		domain:          domain,
-		upstreamServers: srvs,
-		cancel:          func() {},
-	}
-}
-
 func TestUpdateDNSServer(t *testing.T) {
 	nameServers := []nbdns.NameServer{
 		{
@@ -153,37 +140,15 @@ func TestUpdateDNSServer(t *testing.T) {
 					},
 				},
 			},
-			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
-					domain:   "netbird.io",
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-				dummyHandler.id(): handlerWrapper{
-					domain:   "netbird.cloud",
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-				generateDummyHandler(".", nameServers).id(): handlerWrapper{
-					domain:   nbdns.RootZone,
-					handler:  dummyHandler,
-					priority: PriorityDefault,
-				},
-			},
-			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
+			expectedUpstreamMap: registeredHandlerMap{"netbird.io": dummyHandler, "netbird.cloud": dummyHandler, nbdns.RootZone: dummyHandler},
+			expectedLocalMap:    registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
 		},
 		{
-			name:         "New Config Should Succeed",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
-			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
-					domain:   buildRecordKey(zoneRecords[0].Name, 1, 1),
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-			},
-			initSerial:  0,
-			inputSerial: 1,
+			name:            "New Config Should Succeed",
+			initLocalMap:    registrationMap{"netbird.cloud": struct{}{}},
+			initUpstreamMap: registeredHandlerMap{buildRecordKey(zoneRecords[0].Name, 1, 1): dummyHandler},
+			initSerial:      0,
+			inputSerial:     1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -199,19 +164,8 @@ func TestUpdateDNSServer(t *testing.T) {
 					},
 				},
 			},
-			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
-					domain:   "netbird.io",
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-				"local-resolver": handlerWrapper{
-					domain:   "netbird.cloud",
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-			},
-			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
+			expectedUpstreamMap: registeredHandlerMap{"netbird.io": dummyHandler, "netbird.cloud": dummyHandler},
+			expectedLocalMap:    registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
 		},
 		{
 			name:            "Smaller Config Serial Should Be Skipped",
@@ -288,15 +242,9 @@ func TestUpdateDNSServer(t *testing.T) {
 			shouldFail: true,
 		},
 		{
-			name:         "Empty Config Should Succeed and Clean Maps",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
-			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
-					domain:   zoneRecords[0].Name,
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-			},
+			name:                "Empty Config Should Succeed and Clean Maps",
+			initLocalMap:        registrationMap{"netbird.cloud": struct{}{}},
+			initUpstreamMap:     registeredHandlerMap{zoneRecords[0].Name: dummyHandler},
 			initSerial:          0,
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: true},
@@ -304,15 +252,9 @@ func TestUpdateDNSServer(t *testing.T) {
 			expectedLocalMap:    make(registrationMap),
 		},
 		{
-			name:         "Disabled Service Should clean map",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
-			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
-					domain:   zoneRecords[0].Name,
-					handler:  dummyHandler,
-					priority: PriorityMatchDomain,
-				},
-			},
+			name:                "Disabled Service Should clean map",
+			initLocalMap:        registrationMap{"netbird.cloud": struct{}{}},
+			initUpstreamMap:     registeredHandlerMap{zoneRecords[0].Name: dummyHandler},
 			initSerial:          0,
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: false},
@@ -479,13 +421,7 @@ func TestDNSFakeResolverHandleUpdates(t *testing.T) {
 		}
 	}()
 
-	dnsServer.dnsMuxMap = registeredHandlerMap{
-		"id1": handlerWrapper{
-			domain:   zoneRecords[0].Name,
-			handler:  &localResolver{},
-			priority: PriorityMatchDomain,
-		},
-	}
+	dnsServer.dnsMuxMap = registeredHandlerMap{zoneRecords[0].Name: &localResolver{}}
 	dnsServer.localResolver.registeredMap = registrationMap{"netbird.cloud": struct{}{}}
 	dnsServer.updateSerial = 0
 
@@ -626,8 +562,9 @@ func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
-		handlerChain: NewHandlerChain(),
-		hostManager:  hostManager,
+		handlerChain:      NewHandlerChain(),
+		handlerPriorities: make(map[string]int),
+		hostManager:       hostManager,
 		currentConfig: HostDNSConfig{
 			Domains: []DomainConfig{
 				{false, "domain0", false},
@@ -656,7 +593,7 @@ func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 		NameServers: []nbdns.NameServer{
 			{IP: netip.MustParseAddr("8.8.0.0"), NSType: nbdns.UDPNameServerType, Port: 53},
 		},
-	}, nil, 0)
+	}, nil)
 
 	deactivate(nil)
 	expected := "domain0,domain2"
@@ -912,7 +849,7 @@ func createWgInterfaceWithBind(t *testing.T) (*iface.WGIface, error) {
 		return nil, err
 	}
 
-	pf, err := uspfilter.Create(wgIface, false)
+	pf, err := uspfilter.Create(wgIface)
 	if err != nil {
 		t.Fatalf("failed to create uspfilter: %v", err)
 		return nil, err
@@ -966,8 +903,8 @@ func TestHandlerChain_DomainPriorities(t *testing.T) {
 		Subdomains: true,
 	}
 
-	chain.AddHandler("example.com.", dnsRouteHandler, PriorityDNSRoute)
-	chain.AddHandler("example.com.", upstreamHandler, PriorityMatchDomain)
+	chain.AddHandler("example.com.", dnsRouteHandler, PriorityDNSRoute, nil)
+	chain.AddHandler("example.com.", upstreamHandler, PriorityMatchDomain, nil)
 
 	testCases := []struct {
 		name            string
@@ -1022,421 +959,3 @@ func TestHandlerChain_DomainPriorities(t *testing.T) {
 		})
 	}
 }
-
-type mockHandler struct {
-	Id string
-}
-
-func (m *mockHandler) ServeDNS(dns.ResponseWriter, *dns.Msg) {}
-func (m *mockHandler) stop()                                 {}
-func (m *mockHandler) probeAvailability()                    {}
-func (m *mockHandler) id() handlerID                         { return handlerID(m.Id) }
-
-type mockService struct{}
-
-func (m *mockService) Listen() error                   { return nil }
-func (m *mockService) Stop()                           {}
-func (m *mockService) RuntimeIP() string               { return "127.0.0.1" }
-func (m *mockService) RuntimePort() int                { return 53 }
-func (m *mockService) RegisterMux(string, dns.Handler) {}
-func (m *mockService) DeregisterMux(string)            {}
-
-func TestDefaultServer_UpdateMux(t *testing.T) {
-	baseMatchHandlers := registeredHandlerMap{
-		"upstream-group1": {
-			domain: "example.com",
-			handler: &mockHandler{
-				Id: "upstream-group1",
-			},
-			priority: PriorityMatchDomain,
-		},
-		"upstream-group2": {
-			domain: "example.com",
-			handler: &mockHandler{
-				Id: "upstream-group2",
-			},
-			priority: PriorityMatchDomain - 1,
-		},
-	}
-
-	baseRootHandlers := registeredHandlerMap{
-		"upstream-root1": {
-			domain: ".",
-			handler: &mockHandler{
-				Id: "upstream-root1",
-			},
-			priority: PriorityDefault,
-		},
-		"upstream-root2": {
-			domain: ".",
-			handler: &mockHandler{
-				Id: "upstream-root2",
-			},
-			priority: PriorityDefault - 1,
-		},
-	}
-
-	baseMixedHandlers := registeredHandlerMap{
-		"upstream-group1": {
-			domain: "example.com",
-			handler: &mockHandler{
-				Id: "upstream-group1",
-			},
-			priority: PriorityMatchDomain,
-		},
-		"upstream-group2": {
-			domain: "example.com",
-			handler: &mockHandler{
-				Id: "upstream-group2",
-			},
-			priority: PriorityMatchDomain - 1,
-		},
-		"upstream-other": {
-			domain: "other.com",
-			handler: &mockHandler{
-				Id: "upstream-other",
-			},
-			priority: PriorityMatchDomain,
-		},
-	}
-
-	tests := []struct {
-		name             string
-		initialHandlers  registeredHandlerMap
-		updates          []handlerWrapper
-		expectedHandlers map[string]string // map[handlerID]domain
-		description      string
-	}{
-		{
-			name:            "Remove group1 from update",
-			initialHandlers: baseMatchHandlers,
-			updates: []handlerWrapper{
-				// Only group2 remains
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group2",
-					},
-					priority: PriorityMatchDomain - 1,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group2": "example.com",
-			},
-			description: "When group1 is not included in the update, it should be removed while group2 remains",
-		},
-		{
-			name:            "Remove group2 from update",
-			initialHandlers: baseMatchHandlers,
-			updates: []handlerWrapper{
-				// Only group1 remains
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group1",
-					},
-					priority: PriorityMatchDomain,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group1": "example.com",
-			},
-			description: "When group2 is not included in the update, it should be removed while group1 remains",
-		},
-		{
-			name:            "Add group3 in first position",
-			initialHandlers: baseMatchHandlers,
-			updates: []handlerWrapper{
-				// Add group3 with highest priority
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group3",
-					},
-					priority: PriorityMatchDomain + 1,
-				},
-				// Keep existing groups with their original priorities
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group1",
-					},
-					priority: PriorityMatchDomain,
-				},
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group2",
-					},
-					priority: PriorityMatchDomain - 1,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group1": "example.com",
-				"upstream-group2": "example.com",
-				"upstream-group3": "example.com",
-			},
-			description: "When adding group3 with highest priority, it should be first in chain while maintaining existing groups",
-		},
-		{
-			name:            "Add group3 in last position",
-			initialHandlers: baseMatchHandlers,
-			updates: []handlerWrapper{
-				// Keep existing groups with their original priorities
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group1",
-					},
-					priority: PriorityMatchDomain,
-				},
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group2",
-					},
-					priority: PriorityMatchDomain - 1,
-				},
-				// Add group3 with lowest priority
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group3",
-					},
-					priority: PriorityMatchDomain - 2,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group1": "example.com",
-				"upstream-group2": "example.com",
-				"upstream-group3": "example.com",
-			},
-			description: "When adding group3 with lowest priority, it should be last in chain while maintaining existing groups",
-		},
-		// Root zone tests
-		{
-			name:            "Remove root1 from update",
-			initialHandlers: baseRootHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root2",
-					},
-					priority: PriorityDefault - 1,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-root2": ".",
-			},
-			description: "When root1 is not included in the update, it should be removed while root2 remains",
-		},
-		{
-			name:            "Remove root2 from update",
-			initialHandlers: baseRootHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root1",
-					},
-					priority: PriorityDefault,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-root1": ".",
-			},
-			description: "When root2 is not included in the update, it should be removed while root1 remains",
-		},
-		{
-			name:            "Add root3 in first position",
-			initialHandlers: baseRootHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root3",
-					},
-					priority: PriorityDefault + 1,
-				},
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root1",
-					},
-					priority: PriorityDefault,
-				},
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root2",
-					},
-					priority: PriorityDefault - 1,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-root1": ".",
-				"upstream-root2": ".",
-				"upstream-root3": ".",
-			},
-			description: "When adding root3 with highest priority, it should be first in chain while maintaining existing root handlers",
-		},
-		{
-			name:            "Add root3 in last position",
-			initialHandlers: baseRootHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root1",
-					},
-					priority: PriorityDefault,
-				},
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root2",
-					},
-					priority: PriorityDefault - 1,
-				},
-				{
-					domain: ".",
-					handler: &mockHandler{
-						Id: "upstream-root3",
-					},
-					priority: PriorityDefault - 2,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-root1": ".",
-				"upstream-root2": ".",
-				"upstream-root3": ".",
-			},
-			description: "When adding root3 with lowest priority, it should be last in chain while maintaining existing root handlers",
-		},
-		// Mixed domain tests
-		{
-			name:            "Update with mixed domains - remove one of duplicate domain",
-			initialHandlers: baseMixedHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group1",
-					},
-					priority: PriorityMatchDomain,
-				},
-				{
-					domain: "other.com",
-					handler: &mockHandler{
-						Id: "upstream-other",
-					},
-					priority: PriorityMatchDomain,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group1": "example.com",
-				"upstream-other":  "other.com",
-			},
-			description: "When updating mixed domains, should correctly handle removal of one duplicate while maintaining other domains",
-		},
-		{
-			name:            "Update with mixed domains - add new domain",
-			initialHandlers: baseMixedHandlers,
-			updates: []handlerWrapper{
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group1",
-					},
-					priority: PriorityMatchDomain,
-				},
-				{
-					domain: "example.com",
-					handler: &mockHandler{
-						Id: "upstream-group2",
-					},
-					priority: PriorityMatchDomain - 1,
-				},
-				{
-					domain: "other.com",
-					handler: &mockHandler{
-						Id: "upstream-other",
-					},
-					priority: PriorityMatchDomain,
-				},
-				{
-					domain: "new.com",
-					handler: &mockHandler{
-						Id: "upstream-new",
-					},
-					priority: PriorityMatchDomain,
-				},
-			},
-			expectedHandlers: map[string]string{
-				"upstream-group1": "example.com",
-				"upstream-group2": "example.com",
-				"upstream-other":  "other.com",
-				"upstream-new":    "new.com",
-			},
-			description: "When updating mixed domains, should maintain existing duplicates and add new domain",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			server := &DefaultServer{
-				dnsMuxMap:    tt.initialHandlers,
-				handlerChain: NewHandlerChain(),
-				service:      &mockService{},
-			}
-
-			// Perform the update
-			server.updateMux(tt.updates)
-
-			// Verify the results
-			assert.Equal(t, len(tt.expectedHandlers), len(server.dnsMuxMap),
-				"Number of handlers after update doesn't match expected")
-
-			// Check each expected handler
-			for id, expectedDomain := range tt.expectedHandlers {
-				handler, exists := server.dnsMuxMap[handlerID(id)]
-				assert.True(t, exists, "Expected handler %s not found", id)
-				if exists {
-					assert.Equal(t, expectedDomain, handler.domain,
-						"Domain mismatch for handler %s", id)
-				}
-			}
-
-			// Verify no unexpected handlers exist
-			for handlerID := range server.dnsMuxMap {
-				_, expected := tt.expectedHandlers[string(handlerID)]
-				assert.True(t, expected, "Unexpected handler found: %s", handlerID)
-			}
-
-			// Verify the handlerChain state and order
-			previousPriority := 0
-			for _, chainEntry := range server.handlerChain.handlers {
-				// Verify priority order
-				if previousPriority > 0 {
-					assert.True(t, chainEntry.Priority <= previousPriority,
-						"Handlers in chain not properly ordered by priority")
-				}
-				previousPriority = chainEntry.Priority
-
-				// Verify handler exists in mux
-				foundInMux := false
-				for _, muxEntry := range server.dnsMuxMap {
-					if chainEntry.Handler == muxEntry.handler &&
-						chainEntry.Priority == muxEntry.priority &&
-						chainEntry.Pattern == dns.Fqdn(muxEntry.domain) {
-						foundInMux = true
-						break
-					}
-				}
-				assert.True(t, foundInMux,
-					"Handler in chain not found in dnsMuxMap")
-			}
-		})
-	}
-}

client/internal/dns/unclean_shutdown_windows.go

@@ -6,7 +6,6 @@ import (
 
 type ShutdownState struct {
 	Guid string
-	GPO  bool
 }
 
 func (s *ShutdownState) Name() string {
@@ -14,9 +13,9 @@ func (s *ShutdownState) Name() string {
 }
 
 func (s *ShutdownState) Cleanup() error {
-	manager := &registryConfigurator{
-		guid: s.Guid,
-		gpo:  s.GPO,
+	manager, err := newHostManagerWithGuid(s.Guid)
+	if err != nil {
+		return fmt.Errorf("create host manager: %w", err)
 	}
 
 	if err := manager.restoreUncleanShutdownDNS(); err != nil {

client/internal/dns/upstream.go

@@ -2,13 +2,9 @@ package dns
 
 import (
 	"context"
-	"crypto/sha256"
-	"encoding/hex"
 	"errors"
 	"fmt"
 	"net"
-	"slices"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -44,7 +40,6 @@ type upstreamResolverBase struct {
 	cancel           context.CancelFunc
 	upstreamClient   upstreamClient
 	upstreamServers  []string
-	domain           string
 	disabled         bool
 	failsCount       atomic.Int32
 	successCount     atomic.Int32
@@ -58,13 +53,12 @@ type upstreamResolverBase struct {
 	statusRecorder *peer.Status
 }
 
-func newUpstreamResolverBase(ctx context.Context, statusRecorder *peer.Status, domain string) *upstreamResolverBase {
+func newUpstreamResolverBase(ctx context.Context, statusRecorder *peer.Status) *upstreamResolverBase {
 	ctx, cancel := context.WithCancel(ctx)
 
 	return &upstreamResolverBase{
 		ctx:              ctx,
 		cancel:           cancel,
-		domain:           domain,
 		upstreamTimeout:  upstreamTimeout,
 		reactivatePeriod: reactivatePeriod,
 		failsTillDeact:   failsTillDeact,
@@ -77,17 +71,6 @@ func (u *upstreamResolverBase) String() string {
 	return fmt.Sprintf("upstream %v", u.upstreamServers)
 }
 
-// ID returns the unique handler ID
-func (u *upstreamResolverBase) id() handlerID {
-	servers := slices.Clone(u.upstreamServers)
-	slices.Sort(servers)
-
-	hash := sha256.New()
-	hash.Write([]byte(u.domain + ":"))
-	hash.Write([]byte(strings.Join(servers, ",")))
-	return handlerID("upstream-" + hex.EncodeToString(hash.Sum(nil)[:8]))
-}
-
 func (u *upstreamResolverBase) MatchSubdomains() bool {
 	return true
 }
@@ -104,7 +87,7 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 		u.checkUpstreamFails(err)
 	}()
 
-	log.Tracef("received upstream question: domain=%s type=%v class=%v", r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
+	log.WithField("question", r.Question[0]).Trace("received an upstream question")
 	// set the AuthenticatedData flag and the EDNS0 buffer size to 4096 bytes to support larger dns records
 	if r.Extra == nil {
 		r.SetEdns0(4096, false)
@@ -113,7 +96,6 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 
 	select {
 	case <-u.ctx.Done():
-		log.Tracef("%s has been stopped", u)
 		return
 	default:
 	}
@@ -130,36 +112,41 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 
 		if err != nil {
 			if errors.Is(err, context.DeadlineExceeded) || isTimeout(err) {
-				log.Warnf("upstream %s timed out for question domain=%s", upstream, r.Question[0].Name)
+				log.WithError(err).WithField("upstream", upstream).
+					Warn("got an error while connecting to upstream")
 				continue
 			}
-			log.Warnf("failed to query upstream %s for question domain=%s: %s", upstream, r.Question[0].Name, err)
-			continue
+			u.failsCount.Add(1)
+			log.WithError(err).WithField("upstream", upstream).
+				Error("got other error while querying the upstream")
+			return
 		}
 
-		if rm == nil || !rm.Response {
-			log.Warnf("no response from upstream %s for question domain=%s", upstream, r.Question[0].Name)
-			continue
+		if rm == nil {
+			log.WithError(err).WithField("upstream", upstream).
+				Warn("no response from upstream")
+			return
+		}
+		// those checks need to be independent of each other due to memory address issues
+		if !rm.Response {
+			log.WithError(err).WithField("upstream", upstream).
+				Warn("no response from upstream")
+			return
 		}
 
 		u.successCount.Add(1)
-		log.Tracef("took %s to query the upstream %s for question domain=%s", t, upstream, r.Question[0].Name)
+		log.Tracef("took %s to query the upstream %s", t, upstream)
 
-		if err = w.WriteMsg(rm); err != nil {
-			log.Errorf("failed to write DNS response for question domain=%s: %s", r.Question[0].Name, err)
+		err = w.WriteMsg(rm)
+		if err != nil {
+			log.WithError(err).Error("got an error while writing the upstream resolver response")
 		}
 		// count the fails only if they happen sequentially
 		u.failsCount.Store(0)
 		return
 	}
 	u.failsCount.Add(1)
-	log.Errorf("all queries to the %s failed for question domain=%s", u, r.Question[0].Name)
-
-	m := new(dns.Msg)
-	m.SetRcode(r, dns.RcodeServerFailure)
-	if err := w.WriteMsg(m); err != nil {
-		log.Errorf("failed to write error response for %s for question domain=%s: %s", u, r.Question[0].Name, err)
-	}
+	log.Error("all queries to the upstream nameservers failed with timeout")
 }
 
 // checkUpstreamFails counts fails and disables or enables upstream resolving

client/internal/dns/upstream_android.go

@@ -27,9 +27,8 @@ func newUpstreamResolver(
 	_ *net.IPNet,
 	statusRecorder *peer.Status,
 	hostsDNSHolder *hostsDNSHolder,
-	domain string,
 ) (*upstreamResolver, error) {
-	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder, domain)
+	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder)
 	c := &upstreamResolver{
 		upstreamResolverBase: upstreamResolverBase,
 		hostsDNSHolder:       hostsDNSHolder,

client/internal/dns/upstream_general.go

@@ -23,9 +23,8 @@ func newUpstreamResolver(
 	_ *net.IPNet,
 	statusRecorder *peer.Status,
 	_ *hostsDNSHolder,
-	domain string,
 ) (*upstreamResolver, error) {
-	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder, domain)
+	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder)
 	nonIOS := &upstreamResolver{
 		upstreamResolverBase: upstreamResolverBase,
 	}

client/internal/dns/upstream_ios.go

@@ -30,9 +30,8 @@ func newUpstreamResolver(
 	net *net.IPNet,
 	statusRecorder *peer.Status,
 	_ *hostsDNSHolder,
-	domain string,
 ) (*upstreamResolverIOS, error) {
-	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder, domain)
+	upstreamResolverBase := newUpstreamResolverBase(ctx, statusRecorder)
 
 	ios := &upstreamResolverIOS{
 		upstreamResolverBase: upstreamResolverBase,

client/internal/dns/upstream_test.go

@@ -20,7 +20,6 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 		timeout             time.Duration
 		cancelCTX           bool
 		expectedAnswer      string
-		acceptNXDomain      bool
 	}{
 		{
 			name:           "Should Resolve A Record",
@@ -37,11 +36,11 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			expectedAnswer: "1.1.1.1",
 		},
 		{
-			name:           "Should Not Resolve If Can't Connect To Both Servers",
-			inputMSG:       new(dns.Msg).SetQuestion("one.one.one.one.", dns.TypeA),
-			InputServers:   []string{"8.0.0.0:53", "8.0.0.1:53"},
-			timeout:        200 * time.Millisecond,
-			acceptNXDomain: true,
+			name:                "Should Not Resolve If Can't Connect To Both Servers",
+			inputMSG:            new(dns.Msg).SetQuestion("one.one.one.one.", dns.TypeA),
+			InputServers:        []string{"8.0.0.0:53", "8.0.0.1:53"},
+			timeout:             200 * time.Millisecond,
+			responseShouldBeNil: true,
 		},
 		{
 			name:                "Should Not Resolve If Parent Context Is Canceled",
@@ -52,11 +51,14 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			responseShouldBeNil: true,
 		},
 	}
+	// should resolve if first upstream times out
+	// should not write when both fails
+	// should not resolve if parent context is canceled
 
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.TODO())
-			resolver, _ := newUpstreamResolver(ctx, "", net.IP{}, &net.IPNet{}, nil, nil, ".")
+			resolver, _ := newUpstreamResolver(ctx, "", net.IP{}, &net.IPNet{}, nil, nil)
 			resolver.upstreamServers = testCase.InputServers
 			resolver.upstreamTimeout = testCase.timeout
 			if testCase.cancelCTX {
@@ -82,22 +84,16 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 				t.Fatalf("should write a response message")
 			}
 
-			if testCase.acceptNXDomain && responseMSG.Rcode == dns.RcodeNameError {
-				return
+			foundAnswer := false
+			for _, answer := range responseMSG.Answer {
+				if strings.Contains(answer.String(), testCase.expectedAnswer) {
+					foundAnswer = true
+					break
+				}
 			}
 
-			if testCase.expectedAnswer != "" {
-				foundAnswer := false
-				for _, answer := range responseMSG.Answer {
-					if strings.Contains(answer.String(), testCase.expectedAnswer) {
-						foundAnswer = true
-						break
-					}
-				}
-
-				if !foundAnswer {
-					t.Errorf("couldn't find the required answer, %s, in the dns response", testCase.expectedAnswer)
-				}
+			if !foundAnswer {
+				t.Errorf("couldn't find the required answer, %s, in the dns response", testCase.expectedAnswer)
 			}
 		})
 	}

client/internal/dnsfwd/manager.go

@@ -81,7 +81,7 @@ func (m *Manager) Stop(ctx context.Context) error {
 func (h *Manager) allowDNSFirewall() error {
 	dport := &firewall.Port{
 		IsRange: false,
-		Values:  []uint16{ListenPort},
+		Values:  []int{ListenPort},
 	}
 
 	if h.firewall == nil {

client/internal/engine.go

@@ -13,6 +13,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
@@ -42,13 +43,13 @@ import (
 	"github.com/netbirdio/netbird/client/internal/routemanager"
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
-	"github.com/netbirdio/netbird/management/domain"
 	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 
 	nbssh "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
 	nbdns "github.com/netbirdio/netbird/dns"
 	mgm "github.com/netbirdio/netbird/management/client"
+	"github.com/netbirdio/netbird/management/domain"
 	mgmProto "github.com/netbirdio/netbird/management/proto"
 	auth "github.com/netbirdio/netbird/relay/auth/hmac"
 	relayClient "github.com/netbirdio/netbird/relay/client"
@@ -145,7 +146,7 @@ type Engine struct {
 	STUNs []*stun.URI
 	// TURNs is a list of STUN servers used by ICE
 	TURNs    []*stun.URI
-	stunTurn icemaker.StunTurn
+	stunTurn atomic.Value
 
 	clientCtx    context.Context
 	clientCancel context.CancelFunc
@@ -174,6 +175,8 @@ type Engine struct {
 
 	dnsServer dns.Server
 
+	probes *ProbeHolder
+
 	// checks are the client-applied posture checks that need to be evaluated on the client
 	checks []*mgmProto.Checks
 
@@ -193,11 +196,7 @@ type Peer struct {
 	WgAllowedIps string
 }
 
-type localIpUpdater interface {
-	UpdateLocalIPs() error
-}
-
-// NewEngine creates a new Connection Engine with probes attached
+// NewEngine creates a new Connection Engine
 func NewEngine(
 	clientCtx context.Context,
 	clientCancel context.CancelFunc,
@@ -208,6 +207,33 @@ func NewEngine(
 	mobileDep MobileDependency,
 	statusRecorder *peer.Status,
 	checks []*mgmProto.Checks,
+) *Engine {
+	return NewEngineWithProbes(
+		clientCtx,
+		clientCancel,
+		signalClient,
+		mgmClient,
+		relayManager,
+		config,
+		mobileDep,
+		statusRecorder,
+		nil,
+		checks,
+	)
+}
+
+// NewEngineWithProbes creates a new Connection Engine with probes attached
+func NewEngineWithProbes(
+	clientCtx context.Context,
+	clientCancel context.CancelFunc,
+	signalClient signal.Client,
+	mgmClient mgm.Client,
+	relayManager *relayClient.Manager,
+	config *EngineConfig,
+	mobileDep MobileDependency,
+	statusRecorder *peer.Status,
+	probes *ProbeHolder,
+	checks []*mgmProto.Checks,
 ) *Engine {
 	engine := &Engine{
 		clientCtx:      clientCtx,
@@ -225,6 +251,7 @@ func NewEngine(
 		networkSerial:  0,
 		sshServerFunc:  nbssh.DefaultSSHServer,
 		statusRecorder: statusRecorder,
+		probes:         probes,
 		checks:         checks,
 		connSemaphore:  semaphoregroup.NewSemaphoreGroup(connInitLimit),
 	}
@@ -423,6 +450,7 @@ func (e *Engine) Start() error {
 
 	e.receiveSignalEvents()
 	e.receiveManagementEvents()
+	e.receiveProbeEvents()
 
 	// starting network monitor at the very last to avoid disruptions
 	e.startNetworkMonitor()
@@ -437,7 +465,7 @@ func (e *Engine) createFirewall() error {
 	}
 
 	var err error
-	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager, e.config.DisableServerRoutes)
+	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager)
 	if err != nil || e.firewall == nil {
 		log.Errorf("failed creating firewall manager: %s", err)
 		return nil
@@ -467,7 +495,7 @@ func (e *Engine) initFirewall() error {
 	}
 
 	rosenpassPort := e.rpManager.GetAddress().Port
-	port := manager.Port{Values: []uint16{uint16(rosenpassPort)}}
+	port := manager.Port{Values: []int{rosenpassPort}}
 
 	// this rule is static and will be torn down on engine down by the firewall manager
 	if _, err := e.firewall.AddPeerFiltering(
@@ -611,8 +639,8 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 
-	if update.GetNetbirdConfig() != nil {
-		wCfg := update.GetNetbirdConfig()
+	if update.GetWiretrusteeConfig() != nil {
+		wCfg := update.GetWiretrusteeConfig()
 		err := e.updateTURNs(wCfg.GetTurns())
 		if err != nil {
 			return fmt.Errorf("update TURNs: %w", err)
@@ -887,14 +915,6 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		e.acl.ApplyFiltering(networkMap)
 	}
 
-	if e.firewall != nil {
-		if localipfw, ok := e.firewall.(localIpUpdater); ok {
-			if err := localipfw.UpdateLocalIPs(); err != nil {
-				log.Errorf("failed to update local IPs: %v", err)
-			}
-		}
-	}
-
 	// DNS forwarder
 	dnsRouteFeatureFlag := toDNSFeatureFlag(networkMap)
 	dnsRouteDomains := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
@@ -1458,11 +1478,6 @@ func (e *Engine) GetRouteManager() routemanager.Manager {
 	return e.routeManager
 }
 
-// GetFirewallManager returns the firewall manager
-func (e *Engine) GetFirewallManager() manager.Manager {
-	return e.firewall
-}
-
 func findIPFromInterfaceName(ifaceName string) (net.IP, error) {
 	iface, err := net.InterfaceByName(ifaceName)
 	if err != nil {
@@ -1498,58 +1513,72 @@ func (e *Engine) getRosenpassAddr() string {
 	return ""
 }
 
-// RunHealthProbes executes health checks for Signal, Management, Relay and WireGuard services
-// and updates the status recorder with the latest states.
-func (e *Engine) RunHealthProbes() bool {
-	signalHealthy := e.signal.IsHealthy()
-	log.Debugf("signal health check: healthy=%t", signalHealthy)
-
-	managementHealthy := e.mgmClient.IsHealthy()
-	log.Debugf("management health check: healthy=%t", managementHealthy)
-
-	results := append(e.probeSTUNs(), e.probeTURNs()...)
-	e.statusRecorder.UpdateRelayStates(results)
-
-	relayHealthy := true
-	for _, res := range results {
-		if res.Err != nil {
-			relayHealthy = false
-			break
-		}
+func (e *Engine) receiveProbeEvents() {
+	if e.probes == nil {
+		return
 	}
-	log.Debugf("relay health check: healthy=%t", relayHealthy)
-
-	for _, key := range e.peerStore.PeersPubKey() {
-		wgStats, err := e.wgInterface.GetStats(key)
-		if err != nil {
-			log.Debugf("failed to get wg stats for peer %s: %s", key, err)
-			continue
-		}
-		// wgStats could be zero value, in which case we just reset the stats
-		if err := e.statusRecorder.UpdateWireGuardPeerState(key, wgStats); err != nil {
-			log.Debugf("failed to update wg stats for peer %s: %s", key, err)
-		}
+	if e.probes.SignalProbe != nil {
+		go e.probes.SignalProbe.Receive(e.ctx, func() bool {
+			healthy := e.signal.IsHealthy()
+			log.Debugf("received signal probe request, healthy: %t", healthy)
+			return healthy
+		})
 	}
 
-	allHealthy := signalHealthy && managementHealthy && relayHealthy
-	log.Debugf("all health checks completed: healthy=%t", allHealthy)
-	return allHealthy
+	if e.probes.MgmProbe != nil {
+		go e.probes.MgmProbe.Receive(e.ctx, func() bool {
+			healthy := e.mgmClient.IsHealthy()
+			log.Debugf("received management probe request, healthy: %t", healthy)
+			return healthy
+		})
+	}
+
+	if e.probes.RelayProbe != nil {
+		go e.probes.RelayProbe.Receive(e.ctx, func() bool {
+			healthy := true
+
+			results := append(e.probeSTUNs(), e.probeTURNs()...)
+			e.statusRecorder.UpdateRelayStates(results)
+
+			// A single failed server will result in a "failed" probe
+			for _, res := range results {
+				if res.Err != nil {
+					healthy = false
+					break
+				}
+			}
+
+			log.Debugf("received relay probe request, healthy: %t", healthy)
+			return healthy
+		})
+	}
+
+	if e.probes.WgProbe != nil {
+		go e.probes.WgProbe.Receive(e.ctx, func() bool {
+			log.Debug("received wg probe request")
+
+			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)
+				}
+				// wgStats could be zero value, in which case we just reset the stats
+				if err := e.statusRecorder.UpdateWireGuardPeerState(key, wgStats); err != nil {
+					log.Debugf("failed to update wg stats for peer %s: %s", key, err)
+				}
+			}
+
+			return true
+		})
+	}
 }
 
 func (e *Engine) probeSTUNs() []relay.ProbeResult {
-	e.syncMsgMux.Lock()
-	stuns := slices.Clone(e.STUNs)
-	e.syncMsgMux.Unlock()
-
-	return relay.ProbeAll(e.ctx, relay.ProbeSTUN, stuns)
+	return relay.ProbeAll(e.ctx, relay.ProbeSTUN, e.STUNs)
 }
 
 func (e *Engine) probeTURNs() []relay.ProbeResult {
-	e.syncMsgMux.Lock()
-	turns := slices.Clone(e.TURNs)
-	e.syncMsgMux.Unlock()
-
-	return relay.ProbeAll(e.ctx, relay.ProbeTURN, turns)
+	return relay.ProbeAll(e.ctx, relay.ProbeTURN, e.TURNs)
 }
 
 func (e *Engine) restartEngine() {
@@ -1674,14 +1703,6 @@ func (e *Engine) GetLatestNetworkMap() (*mgmProto.NetworkMap, error) {
 	return nm, nil
 }
 
-// GetWgAddr returns the wireguard address
-func (e *Engine) GetWgAddr() net.IP {
-	if e.wgInterface == nil {
-		return nil
-	}
-	return e.wgInterface.Address().IP
-}
-
 // 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 {

client/internal/peer/conn.go

@@ -2,7 +2,6 @@ package peer
 
 import (
 	"context"
-	"fmt"
 	"math/rand"
 	"net"
 	"os"
@@ -29,28 +28,12 @@ import (
 
 type ConnPriority int
 
-func (cp ConnPriority) String() string {
-	switch cp {
-	case connPriorityNone:
-		return "None"
-	case connPriorityRelay:
-		return "PriorityRelay"
-	case connPriorityICETurn:
-		return "PriorityICETurn"
-	case connPriorityICEP2P:
-		return "PriorityICEP2P"
-	default:
-		return fmt.Sprintf("ConnPriority(%d)", cp)
-	}
-}
-
 const (
 	defaultWgKeepAlive = 25 * time.Second
 
-	connPriorityNone    ConnPriority = 0
 	connPriorityRelay   ConnPriority = 1
-	connPriorityICETurn ConnPriority = 2
-	connPriorityICEP2P  ConnPriority = 3
+	connPriorityICETurn ConnPriority = 1
+	connPriorityICEP2P  ConnPriority = 2
 )
 
 type WgConfig struct {
@@ -83,6 +66,14 @@ type ConnConfig struct {
 	ICEConfig icemaker.Config
 }
 
+type WorkerCallbacks struct {
+	OnRelayReadyCallback func(info RelayConnInfo)
+	OnRelayStatusChanged func(ConnStatus)
+
+	OnICEConnReadyCallback func(ConnPriority, ICEConnInfo)
+	OnICEStatusChanged     func(ConnStatus)
+}
+
 type Conn struct {
 	log            *log.Entry
 	mu             sync.Mutex
@@ -144,11 +135,21 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 		semaphore:      semaphore,
 	}
 
+	rFns := WorkerRelayCallbacks{
+		OnConnReady:    conn.relayConnectionIsReady,
+		OnDisconnected: conn.onWorkerRelayStateDisconnected,
+	}
+
+	wFns := WorkerICECallbacks{
+		OnConnReady:     conn.iCEConnectionIsReady,
+		OnStatusChanged: conn.onWorkerICEStateDisconnected,
+	}
+
 	ctrl := isController(config)
-	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, conn, relayManager)
+	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, relayManager, rFns)
 
 	relayIsSupportedLocally := conn.workerRelay.RelayIsSupportedLocally()
-	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, conn, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally)
+	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally, wFns)
 	if err != nil {
 		return nil, err
 	}
@@ -303,7 +304,7 @@ func (conn *Conn) GetKey() string {
 }
 
 // configureConnection starts proxying traffic from/to local Wireguard and sets connection status to StatusConnected
-func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEConnInfo) {
+func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICEConnInfo) {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -316,10 +317,9 @@ func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEC
 		return
 	}
 
-	// this never should happen, because Relay is the lower priority and ICE always close the deprecated connection before upgrade
-	// todo consider to remove this check
+	conn.log.Debugf("ICE connection is ready")
+
 	if conn.currentConnPriority > priority {
-		conn.log.Infof("current connection priority (%s) is higher than the new one (%s), do not upgrade connection", conn.currentConnPriority, priority)
 		conn.statusICE.Set(StatusConnected)
 		conn.updateIceState(iceConnInfo)
 		return
@@ -375,7 +375,8 @@ func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEC
 	conn.doOnConnected(iceConnInfo.RosenpassPubKey, iceConnInfo.RosenpassAddr)
 }
 
-func (conn *Conn) onICEStateDisconnected() {
+// todo review to make sense to handle connecting and disconnected status also?
+func (conn *Conn) onWorkerICEStateDisconnected(newState ConnStatus) {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -383,7 +384,7 @@ func (conn *Conn) onICEStateDisconnected() {
 		return
 	}
 
-	conn.log.Tracef("ICE connection state changed to disconnected")
+	conn.log.Tracef("ICE connection state changed to %s", newState)
 
 	if conn.wgProxyICE != nil {
 		if err := conn.wgProxyICE.CloseConn(); err != nil {
@@ -393,7 +394,7 @@ func (conn *Conn) onICEStateDisconnected() {
 
 	// switch back to relay connection
 	if conn.isReadyToUpgrade() {
-		conn.log.Infof("ICE disconnected, set Relay to active connection")
+		conn.log.Debugf("ICE disconnected, set Relay to active connection")
 		conn.wgProxyRelay.Work()
 
 		if err := conn.configureWGEndpoint(conn.wgProxyRelay.EndpointAddr()); err != nil {
@@ -401,16 +402,12 @@ func (conn *Conn) onICEStateDisconnected() {
 		}
 		conn.workerRelay.EnableWgWatcher(conn.ctx)
 		conn.currentConnPriority = connPriorityRelay
-	} else {
-		conn.log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", connPriorityNone.String())
-		conn.currentConnPriority = connPriorityNone
 	}
 
-	changed := conn.statusICE.Get() != StatusDisconnected
-	if changed {
-		conn.guard.SetICEConnDisconnected()
-	}
-	conn.statusICE.Set(StatusDisconnected)
+	changed := conn.statusICE.Get() != newState && newState != StatusConnecting
+	conn.statusICE.Set(newState)
+
+	conn.guard.SetICEConnDisconnected(changed)
 
 	peerState := State{
 		PubKey:           conn.config.Key,
@@ -425,7 +422,7 @@ func (conn *Conn) onICEStateDisconnected() {
 	}
 }
 
-func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
+func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -447,7 +444,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 	conn.log.Infof("created new wgProxy for relay connection: %s", wgProxy.EndpointAddr().String())
 
 	if conn.iceP2PIsActive() {
-		conn.log.Debugf("do not switch to relay because current priority is: %s", conn.currentConnPriority.String())
+		conn.log.Debugf("do not switch to relay because current priority is: %v", conn.currentConnPriority)
 		conn.setRelayedProxy(wgProxy)
 		conn.statusRelay.Set(StatusConnected)
 		conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
@@ -477,7 +474,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 	conn.doOnConnected(rci.rosenpassPubKey, rci.rosenpassAddr)
 }
 
-func (conn *Conn) onRelayDisconnected() {
+func (conn *Conn) onWorkerRelayStateDisconnected() {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -500,10 +497,8 @@ func (conn *Conn) onRelayDisconnected() {
 	}
 
 	changed := conn.statusRelay.Get() != StatusDisconnected
-	if changed {
-		conn.guard.SetRelayedConnDisconnected()
-	}
 	conn.statusRelay.Set(StatusDisconnected)
+	conn.guard.SetRelayedConnDisconnected(changed)
 
 	peerState := State{
 		PubKey:           conn.config.Key,

client/internal/peer/conn_status_test.go

@@ -1,9 +1,8 @@
 package peer
 
 import (
+	"github.com/magiconair/properties/assert"
 	"testing"
-
-	"github.com/stretchr/testify/assert"
 )
 
 func TestConnStatus_String(t *testing.T) {

client/internal/peer/conn_test.go

@@ -7,7 +7,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/magiconair/properties/assert"
 
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"

client/internal/peer/guard/guard.go

@@ -29,8 +29,8 @@ type Guard struct {
 	isConnectedOnAllWay     isConnectedFunc
 	timeout                 time.Duration
 	srWatcher               *SRWatcher
-	relayedConnDisconnected chan struct{}
-	iCEConnDisconnected     chan struct{}
+	relayedConnDisconnected chan bool
+	iCEConnDisconnected     chan bool
 }
 
 func NewGuard(log *log.Entry, isController bool, isConnectedFn isConnectedFunc, timeout time.Duration, srWatcher *SRWatcher) *Guard {
@@ -41,8 +41,8 @@ func NewGuard(log *log.Entry, isController bool, isConnectedFn isConnectedFunc,
 		isConnectedOnAllWay:     isConnectedFn,
 		timeout:                 timeout,
 		srWatcher:               srWatcher,
-		relayedConnDisconnected: make(chan struct{}, 1),
-		iCEConnDisconnected:     make(chan struct{}, 1),
+		relayedConnDisconnected: make(chan bool, 1),
+		iCEConnDisconnected:     make(chan bool, 1),
 	}
 }
 
@@ -54,16 +54,16 @@ func (g *Guard) Start(ctx context.Context) {
 	}
 }
 
-func (g *Guard) SetRelayedConnDisconnected() {
+func (g *Guard) SetRelayedConnDisconnected(changed bool) {
 	select {
-	case g.relayedConnDisconnected <- struct{}{}:
+	case g.relayedConnDisconnected <- changed:
 	default:
 	}
 }
 
-func (g *Guard) SetICEConnDisconnected() {
+func (g *Guard) SetICEConnDisconnected(changed bool) {
 	select {
-	case g.iCEConnDisconnected <- struct{}{}:
+	case g.iCEConnDisconnected <- changed:
 	default:
 	}
 }
@@ -96,13 +96,19 @@ func (g *Guard) reconnectLoopWithRetry(ctx context.Context) {
 				g.triggerOfferSending()
 			}
 
-		case <-g.relayedConnDisconnected:
+		case changed := <-g.relayedConnDisconnected:
+			if !changed {
+				continue
+			}
 			g.log.Debugf("Relay connection changed, reset reconnection ticker")
 			ticker.Stop()
 			ticker = g.prepareExponentTicker(ctx)
 			tickerChannel = ticker.C
 
-		case <-g.iCEConnDisconnected:
+		case changed := <-g.iCEConnDisconnected:
+			if !changed {
+				continue
+			}
 			g.log.Debugf("ICE connection changed, reset reconnection ticker")
 			ticker.Stop()
 			ticker = g.prepareExponentTicker(ctx)
@@ -132,10 +138,16 @@ func (g *Guard) listenForDisconnectEvents(ctx context.Context) {
 	g.log.Infof("start listen for reconnect events...")
 	for {
 		select {
-		case <-g.relayedConnDisconnected:
+		case changed := <-g.relayedConnDisconnected:
+			if !changed {
+				continue
+			}
 			g.log.Debugf("Relay connection changed, triggering reconnect")
 			g.triggerOfferSending()
-		case <-g.iCEConnDisconnected:
+		case changed := <-g.iCEConnDisconnected:
+			if !changed {
+				continue
+			}
 			g.log.Debugf("ICE state changed, try to send new offer")
 			g.triggerOfferSending()
 		case <-srReconnectedChan:

client/internal/peer/ice/StunTurn.go

@@ -1,22 +0,0 @@
-package ice
-
-import (
-	"sync/atomic"
-
-	"github.com/pion/stun/v2"
-)
-
-type StunTurn atomic.Value
-
-func (s *StunTurn) Load() []*stun.URI {
-	v := (*atomic.Value)(s).Load()
-	if v == nil {
-		return nil
-	}
-
-	return v.([]*stun.URI)
-}
-
-func (s *StunTurn) Store(value []*stun.URI) {
-	(*atomic.Value)(s).Store(value)
-}

client/internal/peer/ice/StunTurn_test.go

@@ -1,13 +0,0 @@
-package ice
-
-import (
-	"testing"
-)
-
-func TestStunTurn_LoadEmpty(t *testing.T) {
-	var stStunTurn StunTurn
-	got := stStunTurn.Load()
-	if len(got) != 0 {
-		t.Errorf("StunTurn.Load() = %v, want %v", got, nil)
-	}
-}

client/internal/peer/ice/agent.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/pion/ice/v3"
 	"github.com/pion/randutil"
+	"github.com/pion/stun/v2"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/internal/stdnet"
@@ -38,7 +39,7 @@ func NewAgent(iFaceDiscover stdnet.ExternalIFaceDiscover, config Config, candida
 	agentConfig := &ice.AgentConfig{
 		MulticastDNSMode:       ice.MulticastDNSModeDisabled,
 		NetworkTypes:           []ice.NetworkType{ice.NetworkTypeUDP4, ice.NetworkTypeUDP6},
-		Urls:                   config.StunTurn.Load(),
+		Urls:                   config.StunTurn.Load().([]*stun.URI),
 		CandidateTypes:         candidateTypes,
 		InterfaceFilter:        stdnet.InterfaceFilter(config.InterfaceBlackList),
 		UDPMux:                 config.UDPMux,

client/internal/peer/ice/config.go

@@ -1,12 +1,14 @@
 package ice
 
 import (
+	"sync/atomic"
+
 	"github.com/pion/ice/v3"
 )
 
 type Config struct {
 	// StunTurn is a list of STUN and TURN URLs
-	StunTurn *StunTurn // []*stun.URI
+	StunTurn *atomic.Value // []*stun.URI
 
 	// InterfaceBlackList is a list of machine interfaces that should be filtered out by ICE Candidate gathering
 	// (e.g. if eth0 is in the list, host candidate of this interface won't be used)

client/internal/peer/status.go

@@ -721,9 +721,7 @@ func (d *Status) GetRelayStates() []relay.ProbeResult {
 func (d *Status) GetDNSStates() []NSGroupState {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-
-	// shallow copy is good enough, as slices fields are currently not updated
-	return slices.Clone(d.nsGroupStates)
+	return d.nsGroupStates
 }
 
 func (d *Status) GetResolvedDomainsStates() map[domain.Domain]ResolvedDomainInfo {

client/internal/peer/wg_watcher.go

@@ -1,154 +0,0 @@
-package peer
-
-import (
-	"context"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/iface/configurer"
-)
-
-const (
-	wgHandshakePeriod = 3 * time.Minute
-)
-
-var (
-	wgHandshakeOvertime = 30 * time.Second // allowed delay in network
-	checkPeriod         = wgHandshakePeriod + wgHandshakeOvertime
-)
-
-type WGInterfaceStater interface {
-	GetStats(key string) (configurer.WGStats, error)
-}
-
-type WGWatcher struct {
-	log           *log.Entry
-	wgIfaceStater WGInterfaceStater
-	peerKey       string
-
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-	ctxLock   sync.Mutex
-	waitGroup sync.WaitGroup
-}
-
-func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey string) *WGWatcher {
-	return &WGWatcher{
-		log:           log,
-		wgIfaceStater: wgIfaceStater,
-		peerKey:       peerKey,
-	}
-}
-
-// EnableWgWatcher starts the WireGuard watcher. If it is already enabled, it will return immediately and do nothing.
-func (w *WGWatcher) EnableWgWatcher(parentCtx context.Context, onDisconnectedFn func()) {
-	w.log.Debugf("enable WireGuard watcher")
-	w.ctxLock.Lock()
-	defer w.ctxLock.Unlock()
-
-	if w.ctx != nil && w.ctx.Err() == nil {
-		w.log.Errorf("WireGuard watcher already enabled")
-		return
-	}
-
-	ctx, ctxCancel := context.WithCancel(parentCtx)
-	w.ctx = ctx
-	w.ctxCancel = ctxCancel
-
-	initialHandshake, err := w.wgState()
-	if err != nil {
-		w.log.Warnf("failed to read initial wg stats: %v", err)
-	}
-
-	w.waitGroup.Add(1)
-	go w.periodicHandshakeCheck(ctx, ctxCancel, onDisconnectedFn, initialHandshake)
-}
-
-// DisableWgWatcher stops the WireGuard watcher and wait for the watcher to exit
-func (w *WGWatcher) DisableWgWatcher() {
-	w.ctxLock.Lock()
-	defer w.ctxLock.Unlock()
-
-	if w.ctxCancel == nil {
-		return
-	}
-
-	w.log.Debugf("disable WireGuard watcher")
-
-	w.ctxCancel()
-	w.ctxCancel = nil
-	w.waitGroup.Wait()
-}
-
-// wgStateCheck help to check the state of the WireGuard handshake and relay connection
-func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, ctxCancel context.CancelFunc, onDisconnectedFn func(), initialHandshake time.Time) {
-	w.log.Infof("WireGuard watcher started")
-	defer w.waitGroup.Done()
-
-	timer := time.NewTimer(wgHandshakeOvertime)
-	defer timer.Stop()
-	defer ctxCancel()
-
-	lastHandshake := initialHandshake
-
-	for {
-		select {
-		case <-timer.C:
-			handshake, ok := w.handshakeCheck(lastHandshake)
-			if !ok {
-				onDisconnectedFn()
-				return
-			}
-			lastHandshake = *handshake
-
-			resetTime := time.Until(handshake.Add(checkPeriod))
-			timer.Reset(resetTime)
-
-			w.log.Debugf("WireGuard watcher reset timer: %v", resetTime)
-		case <-ctx.Done():
-			w.log.Infof("WireGuard watcher stopped")
-			return
-		}
-	}
-}
-
-// handshakeCheck checks the WireGuard handshake and return the new handshake time if it is different from the previous one
-func (w *WGWatcher) handshakeCheck(lastHandshake time.Time) (*time.Time, bool) {
-	handshake, err := w.wgState()
-	if err != nil {
-		w.log.Errorf("failed to read wg stats: %v", err)
-		return nil, false
-	}
-
-	w.log.Tracef("previous handshake, handshake: %v, %v", lastHandshake, handshake)
-
-	// the current know handshake did not change
-	if handshake.Equal(lastHandshake) {
-		w.log.Warnf("WireGuard handshake timed out, closing relay connection: %v", handshake)
-		return nil, false
-	}
-
-	// in case if the machine is suspended, the handshake time will be in the past
-	if handshake.Add(checkPeriod).Before(time.Now()) {
-		w.log.Warnf("WireGuard handshake timed out, closing relay connection: %v", handshake)
-		return nil, false
-	}
-
-	// error handling for handshake time in the future
-	if handshake.After(time.Now()) {
-		w.log.Warnf("WireGuard handshake is in the future, closing relay connection: %v", handshake)
-		return nil, false
-	}
-
-	return &handshake, true
-}
-
-func (w *WGWatcher) wgState() (time.Time, error) {
-	wgState, err := w.wgIfaceStater.GetStats(w.peerKey)
-	if err != nil {
-		return time.Time{}, err
-	}
-	return wgState.LastHandshake, nil
-}

client/internal/peer/wg_watcher_test.go

@@ -1,98 +0,0 @@
-package peer
-
-import (
-	"context"
-	"testing"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-
-	"github.com/netbirdio/netbird/client/iface/configurer"
-)
-
-type MocWgIface struct {
-	initial       bool
-	lastHandshake time.Time
-	stop          bool
-}
-
-func (m *MocWgIface) GetStats(key string) (configurer.WGStats, error) {
-	if !m.initial {
-		m.initial = true
-		return configurer.WGStats{}, nil
-	}
-
-	if !m.stop {
-		m.lastHandshake = time.Now()
-	}
-
-	stats := configurer.WGStats{
-		LastHandshake: m.lastHandshake,
-	}
-
-	return stats, nil
-}
-
-func (m *MocWgIface) disconnect() {
-	m.stop = true
-}
-
-func TestWGWatcher_EnableWgWatcher(t *testing.T) {
-	checkPeriod = 5 * time.Second
-	wgHandshakeOvertime = 1 * time.Second
-
-	mlog := log.WithField("peer", "tet")
-	mocWgIface := &MocWgIface{}
-	watcher := NewWGWatcher(mlog, mocWgIface, "")
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	onDisconnected := make(chan struct{}, 1)
-	watcher.EnableWgWatcher(ctx, func() {
-		mlog.Infof("onDisconnectedFn")
-		onDisconnected <- struct{}{}
-	})
-
-	// wait for initial reading
-	time.Sleep(2 * time.Second)
-	mocWgIface.disconnect()
-
-	select {
-	case <-onDisconnected:
-	case <-time.After(10 * time.Second):
-		t.Errorf("timeout")
-	}
-	watcher.DisableWgWatcher()
-}
-
-func TestWGWatcher_ReEnable(t *testing.T) {
-	checkPeriod = 5 * time.Second
-	wgHandshakeOvertime = 1 * time.Second
-
-	mlog := log.WithField("peer", "tet")
-	mocWgIface := &MocWgIface{}
-	watcher := NewWGWatcher(mlog, mocWgIface, "")
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	onDisconnected := make(chan struct{}, 1)
-
-	watcher.EnableWgWatcher(ctx, func() {})
-	watcher.DisableWgWatcher()
-
-	watcher.EnableWgWatcher(ctx, func() {
-		onDisconnected <- struct{}{}
-	})
-
-	time.Sleep(2 * time.Second)
-	mocWgIface.disconnect()
-
-	select {
-	case <-onDisconnected:
-	case <-time.After(10 * time.Second):
-		t.Errorf("timeout")
-	}
-	watcher.DisableWgWatcher()
-}

client/internal/peer/worker_ice.go

@@ -31,15 +31,20 @@ type ICEConnInfo struct {
 	RelayedOnLocal             bool
 }
 
+type WorkerICECallbacks struct {
+	OnConnReady     func(ConnPriority, ICEConnInfo)
+	OnStatusChanged func(ConnStatus)
+}
+
 type WorkerICE struct {
 	ctx               context.Context
 	log               *log.Entry
 	config            ConnConfig
-	conn              *Conn
 	signaler          *Signaler
 	iFaceDiscover     stdnet.ExternalIFaceDiscover
 	statusRecorder    *Status
 	hasRelayOnLocally bool
+	conn              WorkerICECallbacks
 
 	agent    *ice.Agent
 	muxAgent sync.Mutex
@@ -55,16 +60,16 @@ type WorkerICE struct {
 	lastKnownState ice.ConnectionState
 }
 
-func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, conn *Conn, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool) (*WorkerICE, error) {
+func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool, callBacks WorkerICECallbacks) (*WorkerICE, error) {
 	w := &WorkerICE{
 		ctx:               ctx,
 		log:               log,
 		config:            config,
-		conn:              conn,
 		signaler:          signaler,
 		iFaceDiscover:     ifaceDiscover,
 		statusRecorder:    statusRecorder,
 		hasRelayOnLocally: hasRelayOnLocally,
+		conn:              callBacks,
 	}
 
 	localUfrag, localPwd, err := icemaker.GenerateICECredentials()
@@ -149,8 +154,8 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		Relayed:                    isRelayed(pair),
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
 	}
-	w.log.Debugf("on ICE conn is ready to use")
-	go w.conn.onICEConnectionIsReady(selectedPriority(pair), ci)
+	w.log.Debugf("on ICE conn read to use ready")
+	go w.conn.OnConnReady(selectedPriority(pair), ci)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -215,7 +220,7 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 		case ice.ConnectionStateFailed, ice.ConnectionStateDisconnected:
 			if w.lastKnownState != ice.ConnectionStateDisconnected {
 				w.lastKnownState = ice.ConnectionStateDisconnected
-				w.conn.onICEStateDisconnected()
+				w.conn.OnStatusChanged(StatusDisconnected)
 			}
 			w.closeAgent(agentCancel)
 		default:
@@ -250,10 +255,6 @@ func (w *WorkerICE) closeAgent(cancel context.CancelFunc) {
 	defer w.muxAgent.Unlock()
 
 	cancel()
-	if w.agent == nil {
-		return
-	}
-
 	if err := w.agent.Close(); err != nil {
 		w.log.Warnf("failed to close ICE agent: %s", err)
 	}

client/internal/peer/worker_relay.go

@@ -6,41 +6,52 @@ import (
 	"net"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 
 	relayClient "github.com/netbirdio/netbird/relay/client"
 )
 
+var (
+	wgHandshakePeriod   = 3 * time.Minute
+	wgHandshakeOvertime = 30 * time.Second
+)
+
 type RelayConnInfo struct {
 	relayedConn     net.Conn
 	rosenpassPubKey []byte
 	rosenpassAddr   string
 }
 
+type WorkerRelayCallbacks struct {
+	OnConnReady    func(RelayConnInfo)
+	OnDisconnected func()
+}
+
 type WorkerRelay struct {
 	log          *log.Entry
 	isController bool
 	config       ConnConfig
-	conn         *Conn
 	relayManager relayClient.ManagerService
+	callBacks    WorkerRelayCallbacks
 
-	relayedConn net.Conn
-	relayLock   sync.Mutex
+	relayedConn      net.Conn
+	relayLock        sync.Mutex
+	ctxWgWatch       context.Context
+	ctxCancelWgWatch context.CancelFunc
+	ctxLock          sync.Mutex
 
 	relaySupportedOnRemotePeer atomic.Bool
-
-	wgWatcher *WGWatcher
 }
 
-func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, conn *Conn, relayManager relayClient.ManagerService) *WorkerRelay {
+func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, relayManager relayClient.ManagerService, callbacks WorkerRelayCallbacks) *WorkerRelay {
 	r := &WorkerRelay{
 		log:          log,
 		isController: ctrl,
 		config:       config,
-		conn:         conn,
 		relayManager: relayManager,
-		wgWatcher:    NewWGWatcher(log, config.WgConfig.WgInterface, config.Key),
+		callBacks:    callbacks,
 	}
 	return r
 }
@@ -76,7 +87,7 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relayedConn = relayedConn
 	w.relayLock.Unlock()
 
-	err = w.relayManager.AddCloseListener(srv, w.onRelayClientDisconnected)
+	err = w.relayManager.AddCloseListener(srv, w.onRelayMGDisconnected)
 	if err != nil {
 		log.Errorf("failed to add close listener: %s", err)
 		_ = relayedConn.Close()
@@ -84,7 +95,7 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	}
 
 	w.log.Debugf("peer conn opened via Relay: %s", srv)
-	go w.conn.onRelayConnectionIsReady(RelayConnInfo{
+	go w.callBacks.OnConnReady(RelayConnInfo{
 		relayedConn:     relayedConn,
 		rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
 		rosenpassAddr:   remoteOfferAnswer.RosenpassAddr,
@@ -92,11 +103,32 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 }
 
 func (w *WorkerRelay) EnableWgWatcher(ctx context.Context) {
-	w.wgWatcher.EnableWgWatcher(ctx, w.onWGDisconnected)
+	w.log.Debugf("enable WireGuard watcher")
+	w.ctxLock.Lock()
+	defer w.ctxLock.Unlock()
+
+	if w.ctxWgWatch != nil && w.ctxWgWatch.Err() == nil {
+		return
+	}
+
+	ctx, ctxCancel := context.WithCancel(ctx)
+	w.ctxWgWatch = ctx
+	w.ctxCancelWgWatch = ctxCancel
+
+	w.wgStateCheck(ctx, ctxCancel)
 }
 
 func (w *WorkerRelay) DisableWgWatcher() {
-	w.wgWatcher.DisableWgWatcher()
+	w.ctxLock.Lock()
+	defer w.ctxLock.Unlock()
+
+	if w.ctxCancelWgWatch == nil {
+		return
+	}
+
+	w.log.Debugf("disable WireGuard watcher")
+
+	w.ctxCancelWgWatch()
 }
 
 func (w *WorkerRelay) RelayInstanceAddress() (string, error) {
@@ -118,17 +150,57 @@ func (w *WorkerRelay) CloseConn() {
 		return
 	}
 
-	if err := w.relayedConn.Close(); err != nil {
+	err := w.relayedConn.Close()
+	if err != nil {
 		w.log.Warnf("failed to close relay connection: %v", err)
 	}
 }
 
-func (w *WorkerRelay) onWGDisconnected() {
-	w.relayLock.Lock()
-	_ = w.relayedConn.Close()
-	w.relayLock.Unlock()
+// wgStateCheck help to check the state of the WireGuard handshake and relay connection
+func (w *WorkerRelay) wgStateCheck(ctx context.Context, ctxCancel context.CancelFunc) {
+	w.log.Debugf("WireGuard watcher started")
+	lastHandshake, err := w.wgState()
+	if err != nil {
+		w.log.Warnf("failed to read wg stats: %v", err)
+		lastHandshake = time.Time{}
+	}
+
+	go func(lastHandshake time.Time) {
+		timer := time.NewTimer(wgHandshakeOvertime)
+		defer timer.Stop()
+		defer ctxCancel()
+
+		for {
+			select {
+			case <-timer.C:
+				handshake, err := w.wgState()
+				if err != nil {
+					w.log.Errorf("failed to read wg stats: %v", err)
+					timer.Reset(wgHandshakeOvertime)
+					continue
+				}
+
+				w.log.Tracef("previous handshake, handshake: %v, %v", lastHandshake, handshake)
+
+				if handshake.Equal(lastHandshake) {
+					w.log.Infof("WireGuard handshake timed out, closing relay connection: %v", handshake)
+					w.relayLock.Lock()
+					_ = w.relayedConn.Close()
+					w.relayLock.Unlock()
+					w.callBacks.OnDisconnected()
+					return
+				}
+
+				resetTime := time.Until(handshake.Add(wgHandshakePeriod + wgHandshakeOvertime))
+				lastHandshake = handshake
+				timer.Reset(resetTime)
+			case <-ctx.Done():
+				w.log.Debugf("WireGuard watcher stopped")
+				return
+			}
+		}
+	}(lastHandshake)
 
-	w.conn.onRelayDisconnected()
 }
 
 func (w *WorkerRelay) isRelaySupported(answer *OfferAnswer) bool {
@@ -145,7 +217,20 @@ func (w *WorkerRelay) preferredRelayServer(myRelayAddress, remoteRelayAddress st
 	return remoteRelayAddress
 }
 
-func (w *WorkerRelay) onRelayClientDisconnected() {
-	w.wgWatcher.DisableWgWatcher()
-	go w.conn.onRelayDisconnected()
+func (w *WorkerRelay) wgState() (time.Time, error) {
+	wgState, err := w.config.WgConfig.WgInterface.GetStats(w.config.Key)
+	if err != nil {
+		return time.Time{}, err
+	}
+	return wgState.LastHandshake, nil
+}
+
+func (w *WorkerRelay) onRelayMGDisconnected() {
+	w.ctxLock.Lock()
+	defer w.ctxLock.Unlock()
+
+	if w.ctxCancelWgWatch != nil {
+		w.ctxCancelWgWatch()
+	}
+	go w.callBacks.OnDisconnected()
 }