Add logging around routes

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

client/cmd/debug.go

@@ -168,7 +168,7 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 
 	client := proto.NewDaemonServiceClient(conn)
 
-	stat, err := client.Status(cmd.Context(), &proto.StatusRequest{ShouldRunProbes: true})
+	stat, err := client.Status(cmd.Context(), &proto.StatusRequest{})
 	if err != nil {
 		return fmt.Errorf("failed to get status: %v", status.Convert(err).Message())
 	}
@@ -303,7 +303,7 @@ func setSyncResponsePersistence(cmd *cobra.Command, args []string) error {
 
 func getStatusOutput(cmd *cobra.Command, anon bool) string {
 	var statusOutputString string
-	statusResp, err := getStatus(cmd.Context(), true)
+	statusResp, err := getStatus(cmd.Context())
 	if err != nil {
 		cmd.PrintErrf("Failed to get status: %v\n", err)
 	} else {

client/cmd/status.go

@@ -68,7 +68,7 @@ func statusFunc(cmd *cobra.Command, args []string) error {
 
 	ctx := internal.CtxInitState(cmd.Context())
 
-	resp, err := getStatus(ctx, false)
+	resp, err := getStatus(ctx)
 	if err != nil {
 		return err
 	}
@@ -121,7 +121,7 @@ func statusFunc(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-func getStatus(ctx context.Context, shouldRunProbes bool) (*proto.StatusResponse, error) {
+func getStatus(ctx context.Context) (*proto.StatusResponse, error) {
 	conn, err := DialClientGRPCServer(ctx, daemonAddr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to daemon error: %v\n"+
@@ -130,7 +130,7 @@ func getStatus(ctx context.Context, shouldRunProbes bool) (*proto.StatusResponse
 	}
 	defer conn.Close()
 
-	resp, err := proto.NewDaemonServiceClient(conn).Status(ctx, &proto.StatusRequest{GetFullPeerStatus: true, ShouldRunProbes: shouldRunProbes})
+	resp, err := proto.NewDaemonServiceClient(conn).Status(ctx, &proto.StatusRequest{GetFullPeerStatus: true, ShouldRunProbes: true})
 	if err != nil {
 		return nil, fmt.Errorf("status failed: %v", status.Convert(err).Message())
 	}

client/firewall/iptables/manager_linux.go

@@ -260,22 +260,6 @@ func (m *Manager) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
 	return m.router.UpdateSet(set, prefixes)
 }
 
-// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services.
-func (m *Manager) AddInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.AddInboundDNAT(localAddr, protocol, sourcePort, targetPort)
-}
-
-// RemoveInboundDNAT removes an inbound DNAT rule.
-func (m *Manager) RemoveInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.RemoveInboundDNAT(localAddr, protocol, sourcePort, targetPort)
-}
-
 func getConntrackEstablished() []string {
 	return []string{"-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT"}
 }

client/firewall/iptables/router_linux.go

@@ -880,54 +880,6 @@ func (r *router) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services.
-func (r *router) AddInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	ruleID := fmt.Sprintf("inbound-dnat-%s-%s-%d-%d", localAddr.String(), protocol, sourcePort, targetPort)
-
-	if _, exists := r.rules[ruleID]; exists {
-		return nil
-	}
-
-	dnatRule := []string{
-		"-i", r.wgIface.Name(),
-		"-p", strings.ToLower(string(protocol)),
-		"--dport", strconv.Itoa(int(sourcePort)),
-		"-d", localAddr.String(),
-		"-m", "addrtype", "--dst-type", "LOCAL",
-		"-j", "DNAT",
-		"--to-destination", ":" + strconv.Itoa(int(targetPort)),
-	}
-
-	ruleInfo := ruleInfo{
-		table: tableNat,
-		chain: chainRTRDR,
-		rule:  dnatRule,
-	}
-
-	if err := r.iptablesClient.Append(ruleInfo.table, ruleInfo.chain, ruleInfo.rule...); err != nil {
-		return fmt.Errorf("add inbound DNAT rule: %w", err)
-	}
-	r.rules[ruleID] = ruleInfo.rule
-
-	r.updateState()
-	return nil
-}
-
-// RemoveInboundDNAT removes an inbound DNAT rule.
-func (r *router) RemoveInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	ruleID := fmt.Sprintf("inbound-dnat-%s-%s-%d-%d", localAddr.String(), protocol, sourcePort, targetPort)
-
-	if dnatRule, exists := r.rules[ruleID]; exists {
-		if err := r.iptablesClient.Delete(tableNat, chainRTRDR, dnatRule...); err != nil {
-			return fmt.Errorf("delete inbound DNAT rule: %w", err)
-		}
-		delete(r.rules, ruleID)
-	}
-
-	r.updateState()
-	return nil
-}
-
 func applyPort(flag string, port *firewall.Port) []string {
 	if port == nil {
 		return nil

client/firewall/manager/firewall.go

@@ -151,20 +151,14 @@ type Manager interface {
 
 	DisableRouting() error
 
-	// AddDNATRule adds outbound DNAT rule for forwarding external traffic to the NetBird network.
+	// AddDNATRule adds a DNAT rule
 	AddDNATRule(ForwardRule) (Rule, error)
 
-	// DeleteDNATRule deletes the outbound DNAT rule.
+	// DeleteDNATRule deletes a DNAT rule
 	DeleteDNATRule(Rule) error
 
 	// UpdateSet updates the set with the given prefixes
 	UpdateSet(hash Set, prefixes []netip.Prefix) error
-
-	// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services
-	AddInboundDNAT(localAddr netip.Addr, protocol Protocol, sourcePort, targetPort uint16) error
-
-	// RemoveInboundDNAT removes inbound DNAT rule
-	RemoveInboundDNAT(localAddr netip.Addr, protocol Protocol, sourcePort, targetPort uint16) error
 }
 
 func GenKey(format string, pair RouterPair) string {

client/firewall/nftables/manager_linux.go

@@ -376,22 +376,6 @@ func (m *Manager) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
 	return m.router.UpdateSet(set, prefixes)
 }
 
-// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services.
-func (m *Manager) AddInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.AddInboundDNAT(localAddr, protocol, sourcePort, targetPort)
-}
-
-// RemoveInboundDNAT removes an inbound DNAT rule.
-func (m *Manager) RemoveInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	return m.router.RemoveInboundDNAT(localAddr, protocol, sourcePort, targetPort)
-}
-
 func (m *Manager) createWorkTable() (*nftables.Table, error) {
 	tables, err := m.rConn.ListTablesOfFamily(nftables.TableFamilyIPv4)
 	if err != nil {

client/firewall/nftables/router_linux.go

@@ -1350,103 +1350,6 @@ func (r *router) UpdateSet(set firewall.Set, prefixes []netip.Prefix) error {
 	return nil
 }
 
-// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services.
-func (r *router) AddInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	ruleID := fmt.Sprintf("inbound-dnat-%s-%s-%d-%d", localAddr.String(), protocol, sourcePort, targetPort)
-
-	if _, exists := r.rules[ruleID]; exists {
-		return nil
-	}
-
-	protoNum, err := protoToInt(protocol)
-	if err != nil {
-		return fmt.Errorf("convert protocol to number: %w", err)
-	}
-
-	exprs := []expr.Any{
-		&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 1,
-			Data:     ifname(r.wgIface.Name()),
-		},
-		&expr.Meta{Key: expr.MetaKeyL4PROTO, Register: 2},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 2,
-			Data:     []byte{protoNum},
-		},
-		&expr.Payload{
-			DestRegister: 3,
-			Base:         expr.PayloadBaseTransportHeader,
-			Offset:       2,
-			Len:          2,
-		},
-		&expr.Cmp{
-			Op:       expr.CmpOpEq,
-			Register: 3,
-			Data:     binaryutil.BigEndian.PutUint16(sourcePort),
-		},
-	}
-
-	exprs = append(exprs, applyPrefix(netip.PrefixFrom(localAddr, 32), false)...)
-
-	exprs = append(exprs,
-		&expr.Immediate{
-			Register: 1,
-			Data:     localAddr.AsSlice(),
-		},
-		&expr.Immediate{
-			Register: 2,
-			Data:     binaryutil.BigEndian.PutUint16(targetPort),
-		},
-		&expr.NAT{
-			Type:        expr.NATTypeDestNAT,
-			Family:      uint32(nftables.TableFamilyIPv4),
-			RegAddrMin:  1,
-			RegProtoMin: 2,
-			RegProtoMax: 0,
-		},
-	)
-
-	dnatRule := &nftables.Rule{
-		Table:    r.workTable,
-		Chain:    r.chains[chainNameRoutingRdr],
-		Exprs:    exprs,
-		UserData: []byte(ruleID),
-	}
-	r.conn.AddRule(dnatRule)
-
-	if err := r.conn.Flush(); err != nil {
-		return fmt.Errorf("add inbound DNAT rule: %w", err)
-	}
-
-	r.rules[ruleID] = dnatRule
-
-	return nil
-}
-
-// RemoveInboundDNAT removes an inbound DNAT rule.
-func (r *router) RemoveInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	if err := r.refreshRulesMap(); err != nil {
-		return fmt.Errorf(refreshRulesMapError, err)
-	}
-
-	ruleID := fmt.Sprintf("inbound-dnat-%s-%s-%d-%d", localAddr.String(), protocol, sourcePort, targetPort)
-
-	if rule, exists := r.rules[ruleID]; exists {
-		if err := r.conn.DelRule(rule); err != nil {
-			return fmt.Errorf("delete inbound DNAT rule %s: %w", ruleID, err)
-		}
-		if err := r.conn.Flush(); err != nil {
-			return fmt.Errorf("flush delete inbound DNAT rule: %w", err)
-		}
-		delete(r.rules, ruleID)
-	}
-
-	return nil
-}
-
 // applyNetwork generates nftables expressions for networks (CIDR) or sets
 func (r *router) applyNetwork(
 	network firewall.Network,

client/firewall/uspfilter/conntrack/common.go

@@ -22,8 +22,6 @@ type BaseConnTrack struct {
 	PacketsRx atomic.Uint64
 	BytesTx   atomic.Uint64
 	BytesRx   atomic.Uint64
-
-	DNATOrigPort atomic.Uint32
 }
 
 // these small methods will be inlined by the compiler

client/firewall/uspfilter/conntrack/tcp.go

@@ -157,7 +157,7 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftyp
 	return tracker
 }
 
-func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, uint16, bool) {
+func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, bool) {
 	key := ConnKey{
 		SrcIP:   srcIP,
 		DstIP:   dstIP,
@@ -171,30 +171,28 @@ func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort ui
 
 	if exists {
 		t.updateState(key, conn, flags, direction, size)
-		return key, uint16(conn.DNATOrigPort.Load()), true
+		return key, true
 	}
 
-	return key, 0, false
+	return key, false
 }
 
-// TrackOutbound records an outbound TCP connection and returns the original port if DNAT reversal is needed
-func (t *TCPTracker) TrackOutbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, size int) uint16 {
-	if _, origPort, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, flags, nftypes.Egress, size); exists {
-		return origPort
+// TrackOutbound records an outbound TCP connection
+func (t *TCPTracker) TrackOutbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, flags, nftypes.Egress, size); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Egress, nil, size)
 	}
-	// if (inverted direction) conn is not tracked, track this direction
-	t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Egress, nil, size, 0)
-	return 0
 }
 
 // TrackInbound processes an inbound TCP packet and updates connection state
-func (t *TCPTracker) TrackInbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, ruleID []byte, size int, dnatOrigPort uint16) {
-	t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Ingress, ruleID, size, dnatOrigPort)
+func (t *TCPTracker) TrackInbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, ruleID []byte, size int) {
+	t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Ingress, ruleID, size)
 }
 
 // track is the common implementation for tracking both inbound and outbound connections
-func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int, origPort uint16) {
-	key, _, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, flags, direction, size)
+func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, flags, direction, size)
 	if exists || flags&TCPSyn == 0 {
 		return
 	}
@@ -212,13 +210,8 @@ func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, fla
 
 	conn.tombstone.Store(false)
 	conn.state.Store(int32(TCPStateNew))
-	conn.DNATOrigPort.Store(uint32(origPort))
 
-	if origPort != 0 {
-		t.logger.Trace4("New %s TCP connection: %s (port DNAT %d -> %d)", direction, key, origPort, dstPort)
-	} else {
-		t.logger.Trace2("New %s TCP connection: %s", direction, key)
-	}
+	t.logger.Trace2("New %s TCP connection: %s", direction, key)
 	t.updateState(key, conn, flags, direction, size)
 
 	t.mutex.Lock()
@@ -456,21 +449,6 @@ func (t *TCPTracker) cleanup() {
 	}
 }
 
-// GetConnection safely retrieves a connection state
-func (t *TCPTracker) GetConnection(srcIP netip.Addr, srcPort uint16, dstIP netip.Addr, dstPort uint16) (*TCPConnTrack, bool) {
-	t.mutex.RLock()
-	defer t.mutex.RUnlock()
-
-	key := ConnKey{
-		SrcIP:   srcIP,
-		DstIP:   dstIP,
-		SrcPort: srcPort,
-		DstPort: dstPort,
-	}
-	conn, exists := t.connections[key]
-	return conn, exists
-}
-
 // Close stops the cleanup routine and releases resources
 func (t *TCPTracker) Close() {
 	t.tickerCancel()

client/firewall/uspfilter/conntrack/tcp_test.go

@@ -603,7 +603,7 @@ func TestTCPInboundInitiatedConnection(t *testing.T) {
 	serverPort := uint16(80)
 
 	// 1. Client sends SYN (we receive it as inbound)
-	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPSyn, nil, 100, 0)
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPSyn, nil, 100)
 
 	key := ConnKey{
 		SrcIP:   clientIP,
@@ -623,12 +623,12 @@ func TestTCPInboundInitiatedConnection(t *testing.T) {
 	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPSyn|TCPAck, 100)
 
 	// 3. Client sends ACK to complete handshake
-	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100, 0)
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100)
 	require.Equal(t, TCPStateEstablished, conn.GetState(), "Connection should be ESTABLISHED after handshake completion")
 
 	// 4. Test data transfer
 	// Client sends data
-	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPPush|TCPAck, nil, 1000, 0)
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPPush|TCPAck, nil, 1000)
 
 	// Server sends ACK for data
 	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPAck, 100)
@@ -637,7 +637,7 @@ func TestTCPInboundInitiatedConnection(t *testing.T) {
 	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPPush|TCPAck, 1500)
 
 	// Client sends ACK for data
-	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100, 0)
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100)
 
 	// Verify state and counters
 	require.Equal(t, TCPStateEstablished, conn.GetState())

client/firewall/uspfilter/conntrack/udp.go

@@ -58,23 +58,20 @@ func NewUDPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftyp
 	return tracker
 }
 
-// TrackOutbound records an outbound UDP connection and returns the original port if DNAT reversal is needed
-func (t *UDPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, size int) uint16 {
-	_, origPort, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, nftypes.Egress, size)
-	if exists {
-		return origPort
+// TrackOutbound records an outbound UDP connection
+func (t *UDPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, nftypes.Egress, size); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Egress, nil, size)
 	}
-	// if (inverted direction) conn is not tracked, track this direction
-	t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Egress, nil, size, 0)
-	return 0
 }
 
 // TrackInbound records an inbound UDP connection
-func (t *UDPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, ruleID []byte, size int, dnatOrigPort uint16) {
-	t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Ingress, ruleID, size, dnatOrigPort)
+func (t *UDPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, ruleID []byte, size int) {
+	t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Ingress, ruleID, size)
 }
 
-func (t *UDPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, size int) (ConnKey, uint16, bool) {
+func (t *UDPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, size int) (ConnKey, bool) {
 	key := ConnKey{
 		SrcIP:   srcIP,
 		DstIP:   dstIP,
@@ -89,15 +86,15 @@ func (t *UDPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort
 	if exists {
 		conn.UpdateLastSeen()
 		conn.UpdateCounters(direction, size)
-		return key, uint16(conn.DNATOrigPort.Load()), true
+		return key, true
 	}
 
-	return key, 0, false
+	return key, false
 }
 
 // track is the common implementation for tracking both inbound and outbound connections
-func (t *UDPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, ruleID []byte, size int, origPort uint16) {
-	key, _, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, direction, size)
+func (t *UDPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, ruleID []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, direction, size)
 	if exists {
 		return
 	}
@@ -112,7 +109,6 @@ func (t *UDPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, d
 		SourcePort: srcPort,
 		DestPort:   dstPort,
 	}
-	conn.DNATOrigPort.Store(uint32(origPort))
 	conn.UpdateLastSeen()
 	conn.UpdateCounters(direction, size)
 
@@ -120,11 +116,7 @@ func (t *UDPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, d
 	t.connections[key] = conn
 	t.mutex.Unlock()
 
-	if origPort != 0 {
-		t.logger.Trace4("New %s UDP connection: %s (port DNAT %d -> %d)", direction, key, origPort, dstPort)
-	} else {
-		t.logger.Trace2("New %s UDP connection: %s", direction, key)
-	}
+	t.logger.Trace2("New %s UDP connection: %s", direction, key)
 	t.sendEvent(nftypes.TypeStart, conn, ruleID)
 }
 

client/firewall/uspfilter/filter.go

@@ -50,12 +50,6 @@ const (
 
 var errNatNotSupported = errors.New("nat not supported with userspace firewall")
 
-// serviceKey represents a protocol/port combination for netstack service registry
-type serviceKey struct {
-	protocol gopacket.LayerType
-	port     uint16
-}
-
 // RuleSet is a set of rules grouped by a string key
 type RuleSet map[string]PeerRule
 
@@ -115,13 +109,6 @@ type Manager struct {
 	dnatMappings map[netip.Addr]netip.Addr
 	dnatMutex    sync.RWMutex
 	dnatBiMap    *biDNATMap
-
-	portDNATEnabled atomic.Bool
-	portDNATRules   []portDNATRule
-	portDNATMutex   sync.RWMutex
-
-	netstackServices     map[serviceKey]struct{}
-	netstackServiceMutex sync.RWMutex
 }
 
 // decoder for packages
@@ -135,8 +122,6 @@ type decoder struct {
 	icmp6   layers.ICMPv6
 	decoded []gopacket.LayerType
 	parser  *gopacket.DecodingLayerParser
-
-	dnatOrigPort uint16
 }
 
 // Create userspace firewall manager constructor
@@ -211,8 +196,6 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 		netstack:            netstack.IsEnabled(),
 		localForwarding:     enableLocalForwarding,
 		dnatMappings:        make(map[netip.Addr]netip.Addr),
-		portDNATRules:       []portDNATRule{},
-		netstackServices:    make(map[serviceKey]struct{}),
 	}
 	m.routingEnabled.Store(false)
 
@@ -647,7 +630,7 @@ func (m *Manager) filterOutbound(packetData []byte, size int) bool {
 		return true
 	}
 
-	m.trackOutbound(d, srcIP, dstIP, packetData, size)
+	m.trackOutbound(d, srcIP, dstIP, size)
 	m.translateOutboundDNAT(packetData, d)
 
 	return false
@@ -691,26 +674,14 @@ func getTCPFlags(tcp *layers.TCP) uint8 {
 	return flags
 }
 
-func (m *Manager) trackOutbound(d *decoder, srcIP, dstIP netip.Addr, packetData []byte, size int) {
+func (m *Manager) trackOutbound(d *decoder, srcIP, dstIP netip.Addr, size int) {
 	transport := d.decoded[1]
 	switch transport {
 	case layers.LayerTypeUDP:
-		origPort := m.udpTracker.TrackOutbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), size)
-		if origPort == 0 {
-			break
-		}
-		if err := m.rewriteUDPPort(packetData, d, origPort, sourcePortOffset); err != nil {
-			m.logger.Error1("failed to rewrite UDP port: %v", err)
-		}
+		m.udpTracker.TrackOutbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), size)
 	case layers.LayerTypeTCP:
 		flags := getTCPFlags(&d.tcp)
-		origPort := m.tcpTracker.TrackOutbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, size)
-		if origPort == 0 {
-			break
-		}
-		if err := m.rewriteTCPPort(packetData, d, origPort, sourcePortOffset); err != nil {
-			m.logger.Error1("failed to rewrite TCP port: %v", err)
-		}
+		m.tcpTracker.TrackOutbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, size)
 	case layers.LayerTypeICMPv4:
 		m.icmpTracker.TrackOutbound(srcIP, dstIP, d.icmp4.Id, d.icmp4.TypeCode, d.icmp4.Payload, size)
 	}
@@ -720,15 +691,13 @@ func (m *Manager) trackInbound(d *decoder, srcIP, dstIP netip.Addr, ruleID []byt
 	transport := d.decoded[1]
 	switch transport {
 	case layers.LayerTypeUDP:
-		m.udpTracker.TrackInbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), ruleID, size, d.dnatOrigPort)
+		m.udpTracker.TrackInbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), ruleID, size)
 	case layers.LayerTypeTCP:
 		flags := getTCPFlags(&d.tcp)
-		m.tcpTracker.TrackInbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, ruleID, size, d.dnatOrigPort)
+		m.tcpTracker.TrackInbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, ruleID, size)
 	case layers.LayerTypeICMPv4:
 		m.icmpTracker.TrackInbound(srcIP, dstIP, d.icmp4.Id, d.icmp4.TypeCode, ruleID, d.icmp4.Payload, size)
 	}
-
-	d.dnatOrigPort = 0
 }
 
 // udpHooksDrop checks if any UDP hooks should drop the packet
@@ -790,20 +759,10 @@ func (m *Manager) filterInbound(packetData []byte, size int) bool {
 		return false
 	}
 
-	// TODO: optimize port DNAT by caching matched rules in conntrack
-	if translated := m.translateInboundPortDNAT(packetData, d, srcIP, dstIP); translated {
-		// Re-decode after port DNAT translation to update port information
-		if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-			m.logger.Error1("failed to re-decode packet after port DNAT: %v", err)
-			return true
-		}
-		srcIP, dstIP = m.extractIPs(d)
-	}
-
 	if translated := m.translateInboundReverse(packetData, d); translated {
 		// Re-decode after translation to get original addresses
 		if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-			m.logger.Error1("failed to re-decode packet after reverse DNAT: %v", err)
+			m.logger.Error1("Failed to re-decode packet after reverse DNAT: %v", err)
 			return true
 		}
 		srcIP, dstIP = m.extractIPs(d)
@@ -848,7 +807,9 @@ func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP netip.Addr, packet
 		return true
 	}
 
-	if m.shouldForward(d, dstIP) {
+	// If requested we pass local traffic to internal interfaces to the forwarder.
+	// netstack doesn't have an interface to forward packets to the native stack so we always need to use the forwarder.
+	if m.localForwarding && (m.netstack || dstIP != m.wgIface.Address().IP) {
 		return m.handleForwardedLocalTraffic(packetData)
 	}
 
@@ -1282,86 +1243,3 @@ func (m *Manager) DisableRouting() error {
 
 	return nil
 }
-
-// RegisterNetstackService registers a service as listening on the netstack for the given protocol and port
-func (m *Manager) RegisterNetstackService(protocol nftypes.Protocol, port uint16) {
-	m.netstackServiceMutex.Lock()
-	defer m.netstackServiceMutex.Unlock()
-	layerType := m.protocolToLayerType(protocol)
-	key := serviceKey{protocol: layerType, port: port}
-	m.netstackServices[key] = struct{}{}
-	m.logger.Debug3("RegisterNetstackService: registered %s:%d (layerType=%s)", protocol, port, layerType)
-	m.logger.Debug1("RegisterNetstackService: current registry size: %d", len(m.netstackServices))
-}
-
-// UnregisterNetstackService removes a service from the netstack registry
-func (m *Manager) UnregisterNetstackService(protocol nftypes.Protocol, port uint16) {
-	m.netstackServiceMutex.Lock()
-	defer m.netstackServiceMutex.Unlock()
-	layerType := m.protocolToLayerType(protocol)
-	key := serviceKey{protocol: layerType, port: port}
-	delete(m.netstackServices, key)
-	m.logger.Debug2("Unregistered netstack service on protocol %s port %d", protocol, port)
-}
-
-// protocolToLayerType converts nftypes.Protocol to gopacket.LayerType for internal use
-func (m *Manager) protocolToLayerType(protocol nftypes.Protocol) gopacket.LayerType {
-	switch protocol {
-	case nftypes.TCP:
-		return layers.LayerTypeTCP
-	case nftypes.UDP:
-		return layers.LayerTypeUDP
-	case nftypes.ICMP:
-		return layers.LayerTypeICMPv4
-	default:
-		return gopacket.LayerType(0) // Invalid/unknown
-	}
-}
-
-// shouldForward determines if a packet should be forwarded to the forwarder.
-// The forwarder handles routing packets to the native OS network stack.
-// Returns true if packet should go to the forwarder, false if it should go to netstack listeners or the native stack directly.
-func (m *Manager) shouldForward(d *decoder, dstIP netip.Addr) bool {
-	// not enabled, never forward
-	if !m.localForwarding {
-		return false
-	}
-
-	// netstack always needs to forward because it's lacking a native interface
-	// exception for registered netstack services, those should go to netstack listeners
-	if m.netstack {
-		return !m.hasMatchingNetstackService(d)
-	}
-
-	// traffic to our other local interfaces (not NetBird IP) - always forward
-	if dstIP != m.wgIface.Address().IP {
-		return true
-	}
-
-	// traffic to our NetBird IP, not netstack mode - send to netstack listeners
-	return false
-}
-
-// hasMatchingNetstackService checks if there's a registered netstack service for this packet
-func (m *Manager) hasMatchingNetstackService(d *decoder) bool {
-	if len(d.decoded) < 2 {
-		return false
-	}
-
-	var dstPort uint16
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		dstPort = uint16(d.tcp.DstPort)
-	case layers.LayerTypeUDP:
-		dstPort = uint16(d.udp.DstPort)
-	default:
-		return false
-	}
-
-	key := serviceKey{protocol: d.decoded[1], port: dstPort}
-	m.netstackServiceMutex.RLock()
-	_, exists := m.netstackServices[key]
-	m.netstackServiceMutex.RUnlock()
-
-	return exists
-}

client/firewall/uspfilter/log/log.go

@@ -50,8 +50,6 @@ type logMessage struct {
 	arg4   any
 	arg5   any
 	arg6   any
-	arg7   any
-	arg8   any
 }
 
 // Logger is a high-performance, non-blocking logger
@@ -96,6 +94,7 @@ func (l *Logger) SetLevel(level Level) {
 	log.Debugf("Set uspfilter logger loglevel to %v", levelStrings[level])
 }
 
+
 func (l *Logger) Error(format string) {
 	if l.level.Load() >= uint32(LevelError) {
 		select {
@@ -186,15 +185,6 @@ func (l *Logger) Debug2(format string, arg1, arg2 any) {
 	}
 }
 
-func (l *Logger) Debug3(format string, arg1, arg2, arg3 any) {
-	if l.level.Load() >= uint32(LevelDebug) {
-		select {
-		case l.msgChannel <- logMessage{level: LevelDebug, format: format, arg1: arg1, arg2: arg2, arg3: arg3}:
-		default:
-		}
-	}
-}
-
 func (l *Logger) Trace1(format string, arg1 any) {
 	if l.level.Load() >= uint32(LevelTrace) {
 		select {
@@ -249,16 +239,6 @@ func (l *Logger) Trace6(format string, arg1, arg2, arg3, arg4, arg5, arg6 any) {
 	}
 }
 
-// Trace8 logs a trace message with 8 arguments (8 placeholder in format string)
-func (l *Logger) Trace8(format string, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8 any) {
-	if l.level.Load() >= uint32(LevelTrace) {
-		select {
-		case l.msgChannel <- logMessage{level: LevelTrace, format: format, arg1: arg1, arg2: arg2, arg3: arg3, arg4: arg4, arg5: arg5, arg6: arg6, arg7: arg7, arg8: arg8}:
-		default:
-		}
-	}
-}
-
 func (l *Logger) formatMessage(buf *[]byte, msg logMessage) {
 	*buf = (*buf)[:0]
 	*buf = time.Now().AppendFormat(*buf, "2006-01-02T15:04:05-07:00")
@@ -280,12 +260,6 @@ func (l *Logger) formatMessage(buf *[]byte, msg logMessage) {
 						argCount++
 						if msg.arg6 != nil {
 							argCount++
-							if msg.arg7 != nil {
-								argCount++
-								if msg.arg8 != nil {
-									argCount++
-								}
-							}
 						}
 					}
 				}
@@ -309,10 +283,6 @@ func (l *Logger) formatMessage(buf *[]byte, msg logMessage) {
 		formatted = fmt.Sprintf(msg.format, msg.arg1, msg.arg2, msg.arg3, msg.arg4, msg.arg5)
 	case 6:
 		formatted = fmt.Sprintf(msg.format, msg.arg1, msg.arg2, msg.arg3, msg.arg4, msg.arg5, msg.arg6)
-	case 7:
-		formatted = fmt.Sprintf(msg.format, msg.arg1, msg.arg2, msg.arg3, msg.arg4, msg.arg5, msg.arg6, msg.arg7)
-	case 8:
-		formatted = fmt.Sprintf(msg.format, msg.arg1, msg.arg2, msg.arg3, msg.arg4, msg.arg5, msg.arg6, msg.arg7, msg.arg8)
 	}
 
 	*buf = append(*buf, formatted...)
@@ -420,4 +390,4 @@ func (l *Logger) Stop(ctx context.Context) error {
 	case <-done:
 		return nil
 	}
-}
+}

client/firewall/uspfilter/nat.go

@@ -5,9 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net/netip"
-	"slices"
 
-	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
@@ -15,21 +13,6 @@ import (
 
 var ErrIPv4Only = errors.New("only IPv4 is supported for DNAT")
 
-var (
-	errInvalidIPHeaderLength = errors.New("invalid IP header length")
-)
-
-const (
-	// Port offsets in TCP/UDP headers
-	sourcePortOffset      = 0
-	destinationPortOffset = 2
-
-	// IP address offsets in IPv4 header
-	sourceIPOffset      = 12
-	destinationIPOffset = 16
-)
-
-// ipv4Checksum calculates IPv4 header checksum.
 func ipv4Checksum(header []byte) uint16 {
 	if len(header) < 20 {
 		return 0
@@ -69,7 +52,6 @@ func ipv4Checksum(header []byte) uint16 {
 	return ^uint16(sum)
 }
 
-// icmpChecksum calculates ICMP checksum.
 func icmpChecksum(data []byte) uint16 {
 	var sum1, sum2, sum3, sum4 uint32
 	i := 0
@@ -107,21 +89,11 @@ func icmpChecksum(data []byte) uint16 {
 	return ^uint16(sum)
 }
 
-// biDNATMap maintains bidirectional DNAT mappings.
 type biDNATMap struct {
 	forward map[netip.Addr]netip.Addr
 	reverse map[netip.Addr]netip.Addr
 }
 
-// portDNATRule represents a port-specific DNAT rule.
-type portDNATRule struct {
-	protocol   gopacket.LayerType
-	origPort   uint16
-	targetPort uint16
-	targetIP   netip.Addr
-}
-
-// newBiDNATMap creates a new bidirectional DNAT mapping structure.
 func newBiDNATMap() *biDNATMap {
 	return &biDNATMap{
 		forward: make(map[netip.Addr]netip.Addr),
@@ -129,13 +101,11 @@ func newBiDNATMap() *biDNATMap {
 	}
 }
 
-// set adds a bidirectional DNAT mapping between original and translated addresses.
 func (b *biDNATMap) set(original, translated netip.Addr) {
 	b.forward[original] = translated
 	b.reverse[translated] = original
 }
 
-// delete removes a bidirectional DNAT mapping for the given original address.
 func (b *biDNATMap) delete(original netip.Addr) {
 	if translated, exists := b.forward[original]; exists {
 		delete(b.forward, original)
@@ -143,25 +113,19 @@ func (b *biDNATMap) delete(original netip.Addr) {
 	}
 }
 
-// getTranslated returns the translated address for a given original address.
 func (b *biDNATMap) getTranslated(original netip.Addr) (netip.Addr, bool) {
 	translated, exists := b.forward[original]
 	return translated, exists
 }
 
-// getOriginal returns the original address for a given translated address.
 func (b *biDNATMap) getOriginal(translated netip.Addr) (netip.Addr, bool) {
 	original, exists := b.reverse[translated]
 	return original, exists
 }
 
-// AddInternalDNATMapping adds a 1:1 IP address mapping for internal DNAT translation.
 func (m *Manager) AddInternalDNATMapping(originalAddr, translatedAddr netip.Addr) error {
-	if !originalAddr.IsValid() {
-		return fmt.Errorf("invalid original IP address")
-	}
-	if !translatedAddr.IsValid() {
-		return fmt.Errorf("invalid translated IP address")
+	if !originalAddr.IsValid() || !translatedAddr.IsValid() {
+		return fmt.Errorf("invalid IP addresses")
 	}
 
 	if m.localipmanager.IsLocalIP(translatedAddr) {
@@ -171,6 +135,7 @@ func (m *Manager) AddInternalDNATMapping(originalAddr, translatedAddr netip.Addr
 	m.dnatMutex.Lock()
 	defer m.dnatMutex.Unlock()
 
+	// Initialize both maps together if either is nil
 	if m.dnatMappings == nil || m.dnatBiMap == nil {
 		m.dnatMappings = make(map[netip.Addr]netip.Addr)
 		m.dnatBiMap = newBiDNATMap()
@@ -186,7 +151,7 @@ func (m *Manager) AddInternalDNATMapping(originalAddr, translatedAddr netip.Addr
 	return nil
 }
 
-// RemoveInternalDNATMapping removes a 1:1 IP address mapping.
+// RemoveInternalDNATMapping removes a 1:1 IP address mapping
 func (m *Manager) RemoveInternalDNATMapping(originalAddr netip.Addr) error {
 	m.dnatMutex.Lock()
 	defer m.dnatMutex.Unlock()
@@ -204,7 +169,7 @@ func (m *Manager) RemoveInternalDNATMapping(originalAddr netip.Addr) error {
 	return nil
 }
 
-// getDNATTranslation returns the translated address if a mapping exists.
+// getDNATTranslation returns the translated address if a mapping exists
 func (m *Manager) getDNATTranslation(addr netip.Addr) (netip.Addr, bool) {
 	if !m.dnatEnabled.Load() {
 		return addr, false
@@ -216,7 +181,7 @@ func (m *Manager) getDNATTranslation(addr netip.Addr) (netip.Addr, bool) {
 	return translated, exists
 }
 
-// findReverseDNATMapping finds original address for return traffic.
+// findReverseDNATMapping finds original address for return traffic
 func (m *Manager) findReverseDNATMapping(translatedAddr netip.Addr) (netip.Addr, bool) {
 	if !m.dnatEnabled.Load() {
 		return translatedAddr, false
@@ -228,12 +193,16 @@ func (m *Manager) findReverseDNATMapping(translatedAddr netip.Addr) (netip.Addr,
 	return original, exists
 }
 
-// translateOutboundDNAT applies DNAT translation to outbound packets.
+// translateOutboundDNAT applies DNAT translation to outbound packets
 func (m *Manager) translateOutboundDNAT(packetData []byte, d *decoder) bool {
 	if !m.dnatEnabled.Load() {
 		return false
 	}
 
+	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 {
+		return false
+	}
+
 	dstIP := netip.AddrFrom4([4]byte{packetData[16], packetData[17], packetData[18], packetData[19]})
 
 	translatedIP, exists := m.getDNATTranslation(dstIP)
@@ -241,8 +210,8 @@ func (m *Manager) translateOutboundDNAT(packetData []byte, d *decoder) bool {
 		return false
 	}
 
-	if err := m.rewritePacketIP(packetData, d, translatedIP, destinationIPOffset); err != nil {
-		m.logger.Error1("failed to rewrite packet destination: %v", err)
+	if err := m.rewritePacketDestination(packetData, d, translatedIP); err != nil {
+		m.logger.Error1("Failed to rewrite packet destination: %v", err)
 		return false
 	}
 
@@ -250,12 +219,16 @@ func (m *Manager) translateOutboundDNAT(packetData []byte, d *decoder) bool {
 	return true
 }
 
-// translateInboundReverse applies reverse DNAT to inbound return traffic.
+// translateInboundReverse applies reverse DNAT to inbound return traffic
 func (m *Manager) translateInboundReverse(packetData []byte, d *decoder) bool {
 	if !m.dnatEnabled.Load() {
 		return false
 	}
 
+	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 {
+		return false
+	}
+
 	srcIP := netip.AddrFrom4([4]byte{packetData[12], packetData[13], packetData[14], packetData[15]})
 
 	originalIP, exists := m.findReverseDNATMapping(srcIP)
@@ -263,8 +236,8 @@ func (m *Manager) translateInboundReverse(packetData []byte, d *decoder) bool {
 		return false
 	}
 
-	if err := m.rewritePacketIP(packetData, d, originalIP, sourceIPOffset); err != nil {
-		m.logger.Error1("failed to rewrite packet source: %v", err)
+	if err := m.rewritePacketSource(packetData, d, originalIP); err != nil {
+		m.logger.Error1("Failed to rewrite packet source: %v", err)
 		return false
 	}
 
@@ -272,21 +245,21 @@ func (m *Manager) translateInboundReverse(packetData []byte, d *decoder) bool {
 	return true
 }
 
-// rewritePacketIP replaces an IP address (source or destination) in the packet and updates checksums.
-func (m *Manager) rewritePacketIP(packetData []byte, d *decoder, newIP netip.Addr, ipOffset int) error {
-	if !newIP.Is4() {
+// rewritePacketDestination replaces destination IP in the packet
+func (m *Manager) rewritePacketDestination(packetData []byte, d *decoder, newIP netip.Addr) error {
+	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 || !newIP.Is4() {
 		return ErrIPv4Only
 	}
 
-	var oldIP [4]byte
-	copy(oldIP[:], packetData[ipOffset:ipOffset+4])
-	newIPBytes := newIP.As4()
+	var oldDst [4]byte
+	copy(oldDst[:], packetData[16:20])
+	newDst := newIP.As4()
 
-	copy(packetData[ipOffset:ipOffset+4], newIPBytes[:])
+	copy(packetData[16:20], newDst[:])
 
 	ipHeaderLen := int(d.ip4.IHL) * 4
 	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
-		return errInvalidIPHeaderLength
+		return fmt.Errorf("invalid IP header length")
 	}
 
 	binary.BigEndian.PutUint16(packetData[10:12], 0)
@@ -296,9 +269,44 @@ func (m *Manager) rewritePacketIP(packetData []byte, d *decoder, newIP netip.Add
 	if len(d.decoded) > 1 {
 		switch d.decoded[1] {
 		case layers.LayerTypeTCP:
-			m.updateTCPChecksum(packetData, ipHeaderLen, oldIP[:], newIPBytes[:])
+			m.updateTCPChecksum(packetData, ipHeaderLen, oldDst[:], newDst[:])
 		case layers.LayerTypeUDP:
-			m.updateUDPChecksum(packetData, ipHeaderLen, oldIP[:], newIPBytes[:])
+			m.updateUDPChecksum(packetData, ipHeaderLen, oldDst[:], newDst[:])
+		case layers.LayerTypeICMPv4:
+			m.updateICMPChecksum(packetData, ipHeaderLen)
+		}
+	}
+
+	return nil
+}
+
+// rewritePacketSource replaces the source IP address in the packet
+func (m *Manager) rewritePacketSource(packetData []byte, d *decoder, newIP netip.Addr) error {
+	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 || !newIP.Is4() {
+		return ErrIPv4Only
+	}
+
+	var oldSrc [4]byte
+	copy(oldSrc[:], packetData[12:16])
+	newSrc := newIP.As4()
+
+	copy(packetData[12:16], newSrc[:])
+
+	ipHeaderLen := int(d.ip4.IHL) * 4
+	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
+		return fmt.Errorf("invalid IP header length")
+	}
+
+	binary.BigEndian.PutUint16(packetData[10:12], 0)
+	ipChecksum := ipv4Checksum(packetData[:ipHeaderLen])
+	binary.BigEndian.PutUint16(packetData[10:12], ipChecksum)
+
+	if len(d.decoded) > 1 {
+		switch d.decoded[1] {
+		case layers.LayerTypeTCP:
+			m.updateTCPChecksum(packetData, ipHeaderLen, oldSrc[:], newSrc[:])
+		case layers.LayerTypeUDP:
+			m.updateUDPChecksum(packetData, ipHeaderLen, oldSrc[:], newSrc[:])
 		case layers.LayerTypeICMPv4:
 			m.updateICMPChecksum(packetData, ipHeaderLen)
 		}
@@ -307,7 +315,6 @@ func (m *Manager) rewritePacketIP(packetData []byte, d *decoder, newIP netip.Add
 	return nil
 }
 
-// updateTCPChecksum updates TCP checksum after IP address change per RFC 1624.
 func (m *Manager) updateTCPChecksum(packetData []byte, ipHeaderLen int, oldIP, newIP []byte) {
 	tcpStart := ipHeaderLen
 	if len(packetData) < tcpStart+18 {
@@ -320,7 +327,6 @@ func (m *Manager) updateTCPChecksum(packetData []byte, ipHeaderLen int, oldIP, n
 	binary.BigEndian.PutUint16(packetData[checksumOffset:checksumOffset+2], newChecksum)
 }
 
-// updateUDPChecksum updates UDP checksum after IP address change per RFC 1624.
 func (m *Manager) updateUDPChecksum(packetData []byte, ipHeaderLen int, oldIP, newIP []byte) {
 	udpStart := ipHeaderLen
 	if len(packetData) < udpStart+8 {
@@ -338,7 +344,6 @@ func (m *Manager) updateUDPChecksum(packetData []byte, ipHeaderLen int, oldIP, n
 	binary.BigEndian.PutUint16(packetData[checksumOffset:checksumOffset+2], newChecksum)
 }
 
-// updateICMPChecksum recalculates ICMP checksum after packet modification.
 func (m *Manager) updateICMPChecksum(packetData []byte, ipHeaderLen int) {
 	icmpStart := ipHeaderLen
 	if len(packetData) < icmpStart+8 {
@@ -351,7 +356,7 @@ func (m *Manager) updateICMPChecksum(packetData []byte, ipHeaderLen int) {
 	binary.BigEndian.PutUint16(icmpData[2:4], checksum)
 }
 
-// incrementalUpdate performs incremental checksum update per RFC 1624.
+// incrementalUpdate performs incremental checksum update per RFC 1624
 func incrementalUpdate(oldChecksum uint16, oldBytes, newBytes []byte) uint16 {
 	sum := uint32(^oldChecksum)
 
@@ -386,7 +391,7 @@ func incrementalUpdate(oldChecksum uint16, oldBytes, newBytes []byte) uint16 {
 	return ^uint16(sum)
 }
 
-// AddDNATRule adds outbound DNAT rule for forwarding external traffic to NetBird network.
+// AddDNATRule adds a DNAT rule (delegates to native firewall for port forwarding)
 func (m *Manager) AddDNATRule(rule firewall.ForwardRule) (firewall.Rule, error) {
 	if m.nativeFirewall == nil {
 		return nil, errNatNotSupported
@@ -394,184 +399,10 @@ func (m *Manager) AddDNATRule(rule firewall.ForwardRule) (firewall.Rule, error)
 	return m.nativeFirewall.AddDNATRule(rule)
 }
 
-// DeleteDNATRule deletes outbound DNAT rule.
+// DeleteDNATRule deletes a DNAT rule (delegates to native firewall)
 func (m *Manager) DeleteDNATRule(rule firewall.Rule) error {
 	if m.nativeFirewall == nil {
 		return errNatNotSupported
 	}
 	return m.nativeFirewall.DeleteDNATRule(rule)
 }
-
-// addPortRedirection adds a port redirection rule.
-func (m *Manager) addPortRedirection(targetIP netip.Addr, protocol gopacket.LayerType, sourcePort, targetPort uint16) error {
-	m.portDNATMutex.Lock()
-	defer m.portDNATMutex.Unlock()
-
-	rule := portDNATRule{
-		protocol:   protocol,
-		origPort:   sourcePort,
-		targetPort: targetPort,
-		targetIP:   targetIP,
-	}
-
-	m.portDNATRules = append(m.portDNATRules, rule)
-	m.portDNATEnabled.Store(true)
-
-	return nil
-}
-
-// AddInboundDNAT adds an inbound DNAT rule redirecting traffic from NetBird peers to local services.
-func (m *Manager) AddInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	var layerType gopacket.LayerType
-	switch protocol {
-	case firewall.ProtocolTCP:
-		layerType = layers.LayerTypeTCP
-	case firewall.ProtocolUDP:
-		layerType = layers.LayerTypeUDP
-	default:
-		return fmt.Errorf("unsupported protocol: %s", protocol)
-	}
-
-	return m.addPortRedirection(localAddr, layerType, sourcePort, targetPort)
-}
-
-// removePortRedirection removes a port redirection rule.
-func (m *Manager) removePortRedirection(targetIP netip.Addr, protocol gopacket.LayerType, sourcePort, targetPort uint16) error {
-	m.portDNATMutex.Lock()
-	defer m.portDNATMutex.Unlock()
-
-	m.portDNATRules = slices.DeleteFunc(m.portDNATRules, func(rule portDNATRule) bool {
-		return rule.protocol == protocol && rule.origPort == sourcePort && rule.targetPort == targetPort && rule.targetIP.Compare(targetIP) == 0
-	})
-
-	if len(m.portDNATRules) == 0 {
-		m.portDNATEnabled.Store(false)
-	}
-
-	return nil
-}
-
-// RemoveInboundDNAT removes an inbound DNAT rule.
-func (m *Manager) RemoveInboundDNAT(localAddr netip.Addr, protocol firewall.Protocol, sourcePort, targetPort uint16) error {
-	var layerType gopacket.LayerType
-	switch protocol {
-	case firewall.ProtocolTCP:
-		layerType = layers.LayerTypeTCP
-	case firewall.ProtocolUDP:
-		layerType = layers.LayerTypeUDP
-	default:
-		return fmt.Errorf("unsupported protocol: %s", protocol)
-	}
-
-	return m.removePortRedirection(localAddr, layerType, sourcePort, targetPort)
-}
-
-// translateInboundPortDNAT applies port-specific DNAT translation to inbound packets.
-func (m *Manager) translateInboundPortDNAT(packetData []byte, d *decoder, srcIP, dstIP netip.Addr) bool {
-	if !m.portDNATEnabled.Load() {
-		return false
-	}
-
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		dstPort := uint16(d.tcp.DstPort)
-		return m.applyPortRule(packetData, d, srcIP, dstIP, dstPort, layers.LayerTypeTCP, m.rewriteTCPPort)
-	case layers.LayerTypeUDP:
-		dstPort := uint16(d.udp.DstPort)
-		return m.applyPortRule(packetData, d, netip.Addr{}, dstIP, dstPort, layers.LayerTypeUDP, m.rewriteUDPPort)
-	default:
-		return false
-	}
-}
-
-type portRewriteFunc func(packetData []byte, d *decoder, newPort uint16, portOffset int) error
-
-func (m *Manager) applyPortRule(packetData []byte, d *decoder, srcIP, dstIP netip.Addr, port uint16, protocol gopacket.LayerType, rewriteFn portRewriteFunc) bool {
-	m.portDNATMutex.RLock()
-	defer m.portDNATMutex.RUnlock()
-
-	for _, rule := range m.portDNATRules {
-		if rule.protocol != protocol || rule.targetIP.Compare(dstIP) != 0 {
-			continue
-		}
-
-		if rule.targetPort == port && rule.targetIP.Compare(srcIP) == 0 {
-			return false
-		}
-
-		if rule.origPort != port {
-			continue
-		}
-
-		if err := rewriteFn(packetData, d, rule.targetPort, destinationPortOffset); err != nil {
-			m.logger.Error1("failed to rewrite port: %v", err)
-			return false
-		}
-		d.dnatOrigPort = rule.origPort
-		return true
-	}
-	return false
-}
-
-// rewriteTCPPort rewrites a TCP port (source or destination) and updates checksum.
-func (m *Manager) rewriteTCPPort(packetData []byte, d *decoder, newPort uint16, portOffset int) error {
-	ipHeaderLen := int(d.ip4.IHL) * 4
-	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
-		return errInvalidIPHeaderLength
-	}
-
-	tcpStart := ipHeaderLen
-	if len(packetData) < tcpStart+4 {
-		return fmt.Errorf("packet too short for TCP header")
-	}
-
-	portStart := tcpStart + portOffset
-	oldPort := binary.BigEndian.Uint16(packetData[portStart : portStart+2])
-	binary.BigEndian.PutUint16(packetData[portStart:portStart+2], newPort)
-
-	if len(packetData) >= tcpStart+18 {
-		checksumOffset := tcpStart + 16
-		oldChecksum := binary.BigEndian.Uint16(packetData[checksumOffset : checksumOffset+2])
-
-		var oldPortBytes, newPortBytes [2]byte
-		binary.BigEndian.PutUint16(oldPortBytes[:], oldPort)
-		binary.BigEndian.PutUint16(newPortBytes[:], newPort)
-
-		newChecksum := incrementalUpdate(oldChecksum, oldPortBytes[:], newPortBytes[:])
-		binary.BigEndian.PutUint16(packetData[checksumOffset:checksumOffset+2], newChecksum)
-	}
-
-	return nil
-}
-
-// rewriteUDPPort rewrites a UDP port (source or destination) and updates checksum.
-func (m *Manager) rewriteUDPPort(packetData []byte, d *decoder, newPort uint16, portOffset int) error {
-	ipHeaderLen := int(d.ip4.IHL) * 4
-	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
-		return errInvalidIPHeaderLength
-	}
-
-	udpStart := ipHeaderLen
-	if len(packetData) < udpStart+8 {
-		return fmt.Errorf("packet too short for UDP header")
-	}
-
-	portStart := udpStart + portOffset
-	oldPort := binary.BigEndian.Uint16(packetData[portStart : portStart+2])
-	binary.BigEndian.PutUint16(packetData[portStart:portStart+2], newPort)
-
-	checksumOffset := udpStart + 6
-	if len(packetData) >= udpStart+8 {
-		oldChecksum := binary.BigEndian.Uint16(packetData[checksumOffset : checksumOffset+2])
-		if oldChecksum != 0 {
-			var oldPortBytes, newPortBytes [2]byte
-			binary.BigEndian.PutUint16(oldPortBytes[:], oldPort)
-			binary.BigEndian.PutUint16(newPortBytes[:], newPort)
-
-			newChecksum := incrementalUpdate(oldChecksum, oldPortBytes[:], newPortBytes[:])
-			binary.BigEndian.PutUint16(packetData[checksumOffset:checksumOffset+2], newChecksum)
-		}
-	}
-
-	return nil
-}

client/firewall/uspfilter/nat_bench_test.go

@@ -414,127 +414,3 @@ func BenchmarkChecksumOptimizations(b *testing.B) {
 		}
 	})
 }
-
-// BenchmarkPortDNAT measures the performance of port DNAT operations
-func BenchmarkPortDNAT(b *testing.B) {
-	scenarios := []struct {
-		name         string
-		proto        layers.IPProtocol
-		setupDNAT    bool
-		useMatchPort bool
-		description  string
-	}{
-		{
-			name:         "tcp_inbound_dnat_match",
-			proto:        layers.IPProtocolTCP,
-			setupDNAT:    true,
-			useMatchPort: true,
-			description:  "TCP inbound port DNAT translation (22 → 22022)",
-		},
-		{
-			name:         "tcp_inbound_dnat_nomatch",
-			proto:        layers.IPProtocolTCP,
-			setupDNAT:    true,
-			useMatchPort: false,
-			description:  "TCP inbound with DNAT configured but no port match",
-		},
-		{
-			name:         "tcp_inbound_no_dnat",
-			proto:        layers.IPProtocolTCP,
-			setupDNAT:    false,
-			useMatchPort: false,
-			description:  "TCP inbound without DNAT (baseline)",
-		},
-		{
-			name:         "udp_inbound_dnat_match",
-			proto:        layers.IPProtocolUDP,
-			setupDNAT:    true,
-			useMatchPort: true,
-			description:  "UDP inbound port DNAT translation (5353 → 22054)",
-		},
-		{
-			name:         "udp_inbound_dnat_nomatch",
-			proto:        layers.IPProtocolUDP,
-			setupDNAT:    true,
-			useMatchPort: false,
-			description:  "UDP inbound with DNAT configured but no port match",
-		},
-		{
-			name:         "udp_inbound_no_dnat",
-			proto:        layers.IPProtocolUDP,
-			setupDNAT:    false,
-			useMatchPort: false,
-			description:  "UDP inbound without DNAT (baseline)",
-		},
-	}
-
-	for _, sc := range scenarios {
-		b.Run(sc.name, func(b *testing.B) {
-			manager, err := Create(&IFaceMock{
-				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false, flowLogger)
-			require.NoError(b, err)
-			defer func() {
-				require.NoError(b, manager.Close(nil))
-			}()
-
-			// Set logger to error level to reduce noise during benchmarking
-			manager.SetLogLevel(log.ErrorLevel)
-			defer func() {
-				// Restore to info level after benchmark
-				manager.SetLogLevel(log.InfoLevel)
-			}()
-
-			localAddr := netip.MustParseAddr("100.0.2.175")
-			clientIP := netip.MustParseAddr("100.0.169.249")
-
-			var origPort, targetPort, testPort uint16
-			if sc.proto == layers.IPProtocolTCP {
-				origPort, targetPort = 22, 22022
-			} else {
-				origPort, targetPort = 5353, 22054
-			}
-
-			if sc.useMatchPort {
-				testPort = origPort
-			} else {
-				testPort = 443 // Different port
-			}
-
-			// Setup port DNAT mapping if needed
-			if sc.setupDNAT {
-				err := manager.AddInboundDNAT(localAddr, protocolToFirewall(sc.proto), origPort, targetPort)
-				require.NoError(b, err)
-			}
-
-			// Pre-establish inbound connection for outbound reverse test
-			if sc.setupDNAT && sc.useMatchPort {
-				inboundPacket := generateDNATTestPacket(b, clientIP, localAddr, sc.proto, 54321, origPort)
-				manager.filterInbound(inboundPacket, 0)
-			}
-
-			b.ResetTimer()
-			b.ReportAllocs()
-
-			// Benchmark inbound DNAT translation
-			b.Run("inbound", func(b *testing.B) {
-				for i := 0; i < b.N; i++ {
-					// Create fresh packet each time
-					packet := generateDNATTestPacket(b, clientIP, localAddr, sc.proto, 54321, testPort)
-					manager.filterInbound(packet, 0)
-				}
-			})
-
-			// Benchmark outbound reverse DNAT translation (only if DNAT is set up and port matches)
-			if sc.setupDNAT && sc.useMatchPort {
-				b.Run("outbound_reverse", func(b *testing.B) {
-					for i := 0; i < b.N; i++ {
-						// Create fresh return packet (from target port)
-						packet := generateDNATTestPacket(b, localAddr, clientIP, sc.proto, targetPort, 54321)
-						manager.filterOutbound(packet, 0)
-					}
-				})
-			}
-		})
-	}
-}

client/firewall/uspfilter/nat_test.go

@@ -8,7 +8,6 @@ import (
 	"github.com/google/gopacket/layers"
 	"github.com/stretchr/testify/require"
 
-	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface/device"
 )
 
@@ -144,111 +143,3 @@ func TestDNATMappingManagement(t *testing.T) {
 	err = manager.RemoveInternalDNATMapping(originalIP)
 	require.Error(t, err, "Should error when removing non-existent mapping")
 }
-
-func TestInboundPortDNAT(t *testing.T) {
-	manager, err := Create(&IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false, flowLogger)
-	require.NoError(t, err)
-	defer func() {
-		require.NoError(t, manager.Close(nil))
-	}()
-
-	localAddr := netip.MustParseAddr("100.0.2.175")
-	clientIP := netip.MustParseAddr("100.0.169.249")
-
-	testCases := []struct {
-		name       string
-		protocol   layers.IPProtocol
-		sourcePort uint16
-		targetPort uint16
-	}{
-		{"TCP SSH", layers.IPProtocolTCP, 22, 22022},
-		{"UDP DNS", layers.IPProtocolUDP, 5353, 22054},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			err := manager.AddInboundDNAT(localAddr, protocolToFirewall(tc.protocol), tc.sourcePort, tc.targetPort)
-			require.NoError(t, err)
-
-			inboundPacket := generateDNATTestPacket(t, clientIP, localAddr, tc.protocol, 54321, tc.sourcePort)
-			d := parsePacket(t, inboundPacket)
-
-			translated := manager.translateInboundPortDNAT(inboundPacket, d, clientIP, localAddr)
-			require.True(t, translated, "Inbound packet should be translated")
-
-			d = parsePacket(t, inboundPacket)
-			var dstPort uint16
-			switch tc.protocol {
-			case layers.IPProtocolTCP:
-				dstPort = uint16(d.tcp.DstPort)
-			case layers.IPProtocolUDP:
-				dstPort = uint16(d.udp.DstPort)
-			}
-
-			require.Equal(t, tc.targetPort, dstPort, "Destination port should be rewritten to target port")
-
-			err = manager.RemoveInboundDNAT(localAddr, protocolToFirewall(tc.protocol), tc.sourcePort, tc.targetPort)
-			require.NoError(t, err)
-		})
-	}
-}
-
-func TestInboundPortDNATNegative(t *testing.T) {
-	manager, err := Create(&IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false, flowLogger)
-	require.NoError(t, err)
-	defer func() {
-		require.NoError(t, manager.Close(nil))
-	}()
-
-	localAddr := netip.MustParseAddr("100.0.2.175")
-	clientIP := netip.MustParseAddr("100.0.169.249")
-
-	err = manager.AddInboundDNAT(localAddr, firewall.ProtocolTCP, 22, 22022)
-	require.NoError(t, err)
-
-	testCases := []struct {
-		name     string
-		protocol layers.IPProtocol
-		srcIP    netip.Addr
-		dstIP    netip.Addr
-		srcPort  uint16
-		dstPort  uint16
-	}{
-		{"Wrong port", layers.IPProtocolTCP, clientIP, localAddr, 54321, 80},
-		{"Wrong IP", layers.IPProtocolTCP, clientIP, netip.MustParseAddr("100.64.0.99"), 54321, 22},
-		{"Wrong protocol", layers.IPProtocolUDP, clientIP, localAddr, 54321, 22},
-		{"ICMP", layers.IPProtocolICMPv4, clientIP, localAddr, 0, 0},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			packet := generateDNATTestPacket(t, tc.srcIP, tc.dstIP, tc.protocol, tc.srcPort, tc.dstPort)
-			d := parsePacket(t, packet)
-
-			translated := manager.translateInboundPortDNAT(packet, d, tc.srcIP, tc.dstIP)
-			require.False(t, translated, "Packet should NOT be translated for %s", tc.name)
-
-			d = parsePacket(t, packet)
-			if tc.protocol == layers.IPProtocolTCP {
-				require.Equal(t, tc.dstPort, uint16(d.tcp.DstPort), "Port should remain unchanged")
-			} else if tc.protocol == layers.IPProtocolUDP {
-				require.Equal(t, tc.dstPort, uint16(d.udp.DstPort), "Port should remain unchanged")
-			}
-		})
-	}
-}
-
-func protocolToFirewall(proto layers.IPProtocol) firewall.Protocol {
-	switch proto {
-	case layers.IPProtocolTCP:
-		return firewall.ProtocolTCP
-	case layers.IPProtocolUDP:
-		return firewall.ProtocolUDP
-	default:
-		return firewall.ProtocolALL
-	}
-}

client/firewall/uspfilter/tracer.go

@@ -16,33 +16,25 @@ type PacketStage int
 
 const (
 	StageReceived PacketStage = iota
-	StageInboundPortDNAT
-	StageInbound1to1NAT
 	StageConntrack
 	StagePeerACL
 	StageRouting
 	StageRouteACL
 	StageForwarding
 	StageCompleted
-	StageOutbound1to1NAT
-	StageOutboundPortReverse
 )
 
 const msgProcessingCompleted = "Processing completed"
 
 func (s PacketStage) String() string {
 	return map[PacketStage]string{
-		StageReceived:            "Received",
-		StageInboundPortDNAT:     "Inbound Port DNAT",
-		StageInbound1to1NAT:      "Inbound 1:1 NAT",
-		StageConntrack:           "Connection Tracking",
-		StagePeerACL:             "Peer ACL",
-		StageRouting:             "Routing",
-		StageRouteACL:            "Route ACL",
-		StageForwarding:          "Forwarding",
-		StageCompleted:           "Completed",
-		StageOutbound1to1NAT:     "Outbound 1:1 NAT",
-		StageOutboundPortReverse: "Outbound DNAT Reverse",
+		StageReceived:   "Received",
+		StageConntrack:  "Connection Tracking",
+		StagePeerACL:    "Peer ACL",
+		StageRouting:    "Routing",
+		StageRouteACL:   "Route ACL",
+		StageForwarding: "Forwarding",
+		StageCompleted:  "Completed",
 	}[s]
 }
 
@@ -269,10 +261,6 @@ func (m *Manager) TracePacket(packetData []byte, direction fw.RuleDirection) *Pa
 }
 
 func (m *Manager) traceInbound(packetData []byte, trace *PacketTrace, d *decoder, srcIP netip.Addr, dstIP netip.Addr) *PacketTrace {
-	if m.handleInboundDNAT(trace, packetData, d, &srcIP, &dstIP) {
-		return trace
-	}
-
 	if m.stateful && m.handleConntrackState(trace, d, srcIP, dstIP) {
 		return trace
 	}
@@ -412,16 +400,7 @@ func (m *Manager) addForwardingResult(trace *PacketTrace, action, remoteAddr str
 }
 
 func (m *Manager) traceOutbound(packetData []byte, trace *PacketTrace) *PacketTrace {
-	d := m.decoders.Get().(*decoder)
-	defer m.decoders.Put(d)
-
-	if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-		trace.AddResult(StageCompleted, "Packet dropped - decode error", false)
-		return trace
-	}
-
-	m.handleOutboundDNAT(trace, packetData, d)
-
+	// will create or update the connection state
 	dropped := m.filterOutbound(packetData, 0)
 	if dropped {
 		trace.AddResult(StageCompleted, "Packet dropped by outgoing hook", false)
@@ -430,199 +409,3 @@ func (m *Manager) traceOutbound(packetData []byte, trace *PacketTrace) *PacketTr
 	}
 	return trace
 }
-
-func (m *Manager) handleInboundDNAT(trace *PacketTrace, packetData []byte, d *decoder, srcIP, dstIP *netip.Addr) bool {
-	portDNATApplied := m.traceInboundPortDNAT(trace, packetData, d)
-	if portDNATApplied {
-		if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-			trace.AddResult(StageInboundPortDNAT, "Failed to re-decode after port DNAT", false)
-			return true
-		}
-		*srcIP, *dstIP = m.extractIPs(d)
-		trace.DestinationPort = m.getDestPort(d)
-	}
-
-	nat1to1Applied := m.traceInbound1to1NAT(trace, packetData, d)
-	if nat1to1Applied {
-		if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
-			trace.AddResult(StageInbound1to1NAT, "Failed to re-decode after 1:1 NAT", false)
-			return true
-		}
-		*srcIP, *dstIP = m.extractIPs(d)
-	}
-
-	return false
-}
-
-func (m *Manager) traceInboundPortDNAT(trace *PacketTrace, packetData []byte, d *decoder) bool {
-	if !m.portDNATEnabled.Load() {
-		trace.AddResult(StageInboundPortDNAT, "Port DNAT not enabled", true)
-		return false
-	}
-
-	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 {
-		trace.AddResult(StageInboundPortDNAT, "Not IPv4, skipping port DNAT", true)
-		return false
-	}
-
-	if len(d.decoded) < 2 {
-		trace.AddResult(StageInboundPortDNAT, "No transport layer, skipping port DNAT", true)
-		return false
-	}
-
-	protocol := d.decoded[1]
-	if protocol != layers.LayerTypeTCP && protocol != layers.LayerTypeUDP {
-		trace.AddResult(StageInboundPortDNAT, "Not TCP/UDP, skipping port DNAT", true)
-		return false
-	}
-
-	srcIP := netip.AddrFrom4([4]byte{packetData[12], packetData[13], packetData[14], packetData[15]})
-	dstIP := netip.AddrFrom4([4]byte{packetData[16], packetData[17], packetData[18], packetData[19]})
-	var originalPort uint16
-	if protocol == layers.LayerTypeTCP {
-		originalPort = uint16(d.tcp.DstPort)
-	} else {
-		originalPort = uint16(d.udp.DstPort)
-	}
-
-	translated := m.translateInboundPortDNAT(packetData, d, srcIP, dstIP)
-	if translated {
-		ipHeaderLen := int((packetData[0] & 0x0F) * 4)
-		translatedPort := uint16(packetData[ipHeaderLen+2])<<8 | uint16(packetData[ipHeaderLen+3])
-
-		protoStr := "TCP"
-		if protocol == layers.LayerTypeUDP {
-			protoStr = "UDP"
-		}
-		msg := fmt.Sprintf("%s port DNAT applied: %s:%d -> %s:%d", protoStr, dstIP, originalPort, dstIP, translatedPort)
-		trace.AddResult(StageInboundPortDNAT, msg, true)
-		return true
-	}
-
-	trace.AddResult(StageInboundPortDNAT, "No matching port DNAT rule", true)
-	return false
-}
-
-func (m *Manager) traceInbound1to1NAT(trace *PacketTrace, packetData []byte, d *decoder) bool {
-	if !m.dnatEnabled.Load() {
-		trace.AddResult(StageInbound1to1NAT, "1:1 NAT not enabled", true)
-		return false
-	}
-
-	srcIP := netip.AddrFrom4([4]byte{packetData[12], packetData[13], packetData[14], packetData[15]})
-
-	translated := m.translateInboundReverse(packetData, d)
-	if translated {
-		m.dnatMutex.RLock()
-		translatedIP, exists := m.dnatBiMap.getOriginal(srcIP)
-		m.dnatMutex.RUnlock()
-
-		if exists {
-			msg := fmt.Sprintf("1:1 NAT reverse applied: %s -> %s", srcIP, translatedIP)
-			trace.AddResult(StageInbound1to1NAT, msg, true)
-			return true
-		}
-	}
-
-	trace.AddResult(StageInbound1to1NAT, "No matching 1:1 NAT rule", true)
-	return false
-}
-
-func (m *Manager) handleOutboundDNAT(trace *PacketTrace, packetData []byte, d *decoder) {
-	m.traceOutbound1to1NAT(trace, packetData, d)
-	m.traceOutboundPortReverse(trace, packetData, d)
-}
-
-func (m *Manager) traceOutbound1to1NAT(trace *PacketTrace, packetData []byte, d *decoder) bool {
-	if !m.dnatEnabled.Load() {
-		trace.AddResult(StageOutbound1to1NAT, "1:1 NAT not enabled", true)
-		return false
-	}
-
-	dstIP := netip.AddrFrom4([4]byte{packetData[16], packetData[17], packetData[18], packetData[19]})
-
-	translated := m.translateOutboundDNAT(packetData, d)
-	if translated {
-		m.dnatMutex.RLock()
-		translatedIP, exists := m.dnatMappings[dstIP]
-		m.dnatMutex.RUnlock()
-
-		if exists {
-			msg := fmt.Sprintf("1:1 NAT applied: %s -> %s", dstIP, translatedIP)
-			trace.AddResult(StageOutbound1to1NAT, msg, true)
-			return true
-		}
-	}
-
-	trace.AddResult(StageOutbound1to1NAT, "No matching 1:1 NAT rule", true)
-	return false
-}
-
-func (m *Manager) traceOutboundPortReverse(trace *PacketTrace, packetData []byte, d *decoder) bool {
-	if !m.portDNATEnabled.Load() {
-		trace.AddResult(StageOutboundPortReverse, "Port DNAT not enabled", true)
-		return false
-	}
-
-	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 {
-		trace.AddResult(StageOutboundPortReverse, "Not IPv4, skipping port reverse", true)
-		return false
-	}
-
-	if len(d.decoded) < 2 {
-		trace.AddResult(StageOutboundPortReverse, "No transport layer, skipping port reverse", true)
-		return false
-	}
-
-	srcIP := netip.AddrFrom4([4]byte{packetData[12], packetData[13], packetData[14], packetData[15]})
-	dstIP := netip.AddrFrom4([4]byte{packetData[16], packetData[17], packetData[18], packetData[19]})
-
-	var origPort uint16
-	transport := d.decoded[1]
-	switch transport {
-	case layers.LayerTypeTCP:
-		srcPort := uint16(d.tcp.SrcPort)
-		dstPort := uint16(d.tcp.DstPort)
-		conn, exists := m.tcpTracker.GetConnection(dstIP, dstPort, srcIP, srcPort)
-		if exists {
-			origPort = uint16(conn.DNATOrigPort.Load())
-		}
-		if origPort != 0 {
-			msg := fmt.Sprintf("TCP DNAT reverse (tracked connection): %s:%d -> %s:%d", srcIP, srcPort, srcIP, origPort)
-			trace.AddResult(StageOutboundPortReverse, msg, true)
-			return true
-		}
-	case layers.LayerTypeUDP:
-		srcPort := uint16(d.udp.SrcPort)
-		dstPort := uint16(d.udp.DstPort)
-		conn, exists := m.udpTracker.GetConnection(dstIP, dstPort, srcIP, srcPort)
-		if exists {
-			origPort = uint16(conn.DNATOrigPort.Load())
-		}
-		if origPort != 0 {
-			msg := fmt.Sprintf("UDP DNAT reverse (tracked connection): %s:%d -> %s:%d", srcIP, srcPort, srcIP, origPort)
-			trace.AddResult(StageOutboundPortReverse, msg, true)
-			return true
-		}
-	default:
-		trace.AddResult(StageOutboundPortReverse, "Not TCP/UDP, skipping port reverse", true)
-		return false
-	}
-
-	trace.AddResult(StageOutboundPortReverse, "No tracked connection for DNAT reverse", true)
-	return false
-}
-
-func (m *Manager) getDestPort(d *decoder) uint16 {
-	if len(d.decoded) < 2 {
-		return 0
-	}
-	switch d.decoded[1] {
-	case layers.LayerTypeTCP:
-		return uint16(d.tcp.DstPort)
-	case layers.LayerTypeUDP:
-		return uint16(d.udp.DstPort)
-	default:
-		return 0
-	}
-}

client/firewall/uspfilter/tracer_test.go

@@ -104,8 +104,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -128,8 +126,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -157,8 +153,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -185,8 +179,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -212,8 +204,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StageRouteACL,
@@ -238,8 +228,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StageRouteACL,
@@ -258,8 +246,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StageRouteACL,
@@ -278,8 +264,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StageCompleted,
@@ -303,8 +287,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageCompleted,
 			},
@@ -319,8 +301,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageOutbound1to1NAT,
-				StageOutboundPortReverse,
 				StageCompleted,
 			},
 			expectedAllow: true,
@@ -339,8 +319,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -362,8 +340,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -386,8 +362,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -408,8 +382,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageConntrack,
 				StageRouting,
 				StagePeerACL,
@@ -434,8 +406,6 @@ func TestTracePacket(t *testing.T) {
 			},
 			expectedStages: []PacketStage{
 				StageReceived,
-				StageInboundPortDNAT,
-				StageInbound1to1NAT,
 				StageRouting,
 				StagePeerACL,
 				StageCompleted,

client/internal/connect.go

@@ -25,6 +25,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
+	nbnet "github.com/netbirdio/netbird/client/net"
 	cProto "github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
@@ -34,7 +35,6 @@ import (
 	relayClient "github.com/netbirdio/netbird/shared/relay/client"
 	signal "github.com/netbirdio/netbird/shared/signal/client"
 	"github.com/netbirdio/netbird/util"
-	nbnet "github.com/netbirdio/netbird/client/net"
 	"github.com/netbirdio/netbird/version"
 )
 
@@ -289,15 +289,18 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		}
 
 		<-engineCtx.Done()
+
 		c.engineMutex.Lock()
-		if c.engine != nil && c.engine.wgInterface != nil {
-			log.Infof("ensuring %s is removed, Netbird engine context cancelled", c.engine.wgInterface.Name())
-			if err := c.engine.Stop(); err != nil {
+		engine := c.engine
+		c.engine = nil
+		c.engineMutex.Unlock()
+
+		if engine != nil && engine.wgInterface != nil {
+			log.Infof("ensuring %s is removed, Netbird engine context cancelled", engine.wgInterface.Name())
+			if err := engine.Stop(); err != nil {
 				log.Errorf("Failed to stop engine: %v", err)
 			}
-			c.engine = nil
 		}
-		c.engineMutex.Unlock()
 		c.statusRecorder.ClientTeardown()
 
 		backOff.Reset()
@@ -382,19 +385,12 @@ func (c *ConnectClient) Status() StatusType {
 }
 
 func (c *ConnectClient) Stop() error {
-	if c == nil {
-		return nil
+	engine := c.Engine()
+	if engine != nil {
+		if err := engine.Stop(); err != nil {
+			return fmt.Errorf("stop engine: %w", err)
+		}
 	}
-	c.engineMutex.Lock()
-	defer c.engineMutex.Unlock()
-
-	if c.engine == nil {
-		return nil
-	}
-	if err := c.engine.Stop(); err != nil {
-		return fmt.Errorf("stop engine: %w", err)
-	}
-
 	return nil
 }
 

client/internal/dns/server.go

@@ -65,8 +65,9 @@ type hostManagerWithOriginalNS interface {
 
 // DefaultServer dns server object
 type DefaultServer struct {
-	ctx       context.Context
-	ctxCancel context.CancelFunc
+	ctx        context.Context
+	ctxCancel  context.CancelFunc
+	shutdownWg sync.WaitGroup
 	// disableSys disables system DNS management (e.g., /etc/resolv.conf updates) while keeping the DNS service running.
 	// This is different from ServiceEnable=false from management which completely disables the DNS service.
 	disableSys         bool
@@ -318,6 +319,7 @@ func (s *DefaultServer) DnsIP() netip.Addr {
 // Stop stops the server
 func (s *DefaultServer) Stop() {
 	s.ctxCancel()
+	s.shutdownWg.Wait()
 
 	s.mux.Lock()
 	defer s.mux.Unlock()
@@ -507,8 +509,9 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 
 	s.applyHostConfig()
 
+	s.shutdownWg.Add(1)
 	go func() {
-		// persist dns state right away
+		defer s.shutdownWg.Done()
 		if err := s.stateManager.PersistState(s.ctx); err != nil {
 			log.Errorf("Failed to persist dns state: %v", err)
 		}

client/internal/dnsfwd/cache_test.go

@@ -83,3 +83,4 @@ func TestCacheMiss(t *testing.T) {
 		t.Fatalf("expected cache miss, got=%v ok=%v", got, ok)
 	}
 }
+

client/internal/dnsfwd/forwarder.go

@@ -14,7 +14,6 @@ import (
 	"github.com/hashicorp/go-multierror"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/tun/netstack"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
@@ -34,7 +33,7 @@ type firewaller interface {
 }
 
 type DNSForwarder struct {
-	listenAddress  netip.AddrPort
+	listenAddress  string
 	ttl            uint32
 	statusRecorder *peer.Status
 
@@ -48,11 +47,9 @@ type DNSForwarder struct {
 	firewall   firewaller
 	resolver   resolver
 	cache      *cache
-
-	wgIface wgIface
 }
 
-func NewDNSForwarder(listenAddress netip.AddrPort, ttl uint32, firewall firewaller, statusRecorder *peer.Status, wgIface wgIface) *DNSForwarder {
+func NewDNSForwarder(listenAddress string, ttl uint32, firewall firewaller, statusRecorder *peer.Status) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
 		listenAddress:  listenAddress,
@@ -61,46 +58,30 @@ func NewDNSForwarder(listenAddress netip.AddrPort, ttl uint32, firewall firewall
 		statusRecorder: statusRecorder,
 		resolver:       net.DefaultResolver,
 		cache:          newCache(),
-		wgIface:        wgIface,
 	}
 }
 
 func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
-	var netstackNet *netstack.Net
-	if f.wgIface != nil {
-		netstackNet = f.wgIface.GetNet()
-	}
-
-	addrDesc := f.listenAddress.String()
-	if netstackNet != nil {
-		addrDesc = fmt.Sprintf("netstack %s", f.listenAddress)
-	}
-	log.Infof("starting DNS forwarder on address=%s", addrDesc)
-
-	udpLn, err := f.createUDPListener(netstackNet)
-	if err != nil {
-		return fmt.Errorf("create UDP listener: %w", err)
-	}
-
-	tcpLn, err := f.createTCPListener(netstackNet)
-	if err != nil {
-		return fmt.Errorf("create TCP listener: %w", err)
-	}
+	log.Infof("starting DNS forwarder on address=%s", f.listenAddress)
 
+	// UDP server
 	mux := dns.NewServeMux()
 	f.mux = mux
 	mux.HandleFunc(".", f.handleDNSQueryUDP)
 	f.dnsServer = &dns.Server{
-		PacketConn: udpLn,
-		Handler:    mux,
+		Addr:    f.listenAddress,
+		Net:     "udp",
+		Handler: mux,
 	}
 
+	// TCP server
 	tcpMux := dns.NewServeMux()
 	f.tcpMux = tcpMux
 	tcpMux.HandleFunc(".", f.handleDNSQueryTCP)
 	f.tcpServer = &dns.Server{
-		Listener: tcpLn,
-		Handler:  tcpMux,
+		Addr:    f.listenAddress,
+		Net:     "tcp",
+		Handler: tcpMux,
 	}
 
 	f.UpdateDomains(entries)
@@ -108,33 +89,18 @@ func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
 	errCh := make(chan error, 2)
 
 	go func() {
-		log.Infof("DNS UDP listener running on %s", addrDesc)
-		errCh <- f.dnsServer.ActivateAndServe()
+		log.Infof("DNS UDP listener running on %s", f.listenAddress)
+		errCh <- f.dnsServer.ListenAndServe()
 	}()
 	go func() {
-		log.Infof("DNS TCP listener running on %s", addrDesc)
-		errCh <- f.tcpServer.ActivateAndServe()
+		log.Infof("DNS TCP listener running on %s", f.listenAddress)
+		errCh <- f.tcpServer.ListenAndServe()
 	}()
 
+	// return the first error we get (e.g. bind failure or shutdown)
 	return <-errCh
 }
 
-func (f *DNSForwarder) createUDPListener(netstackNet *netstack.Net) (net.PacketConn, error) {
-	if netstackNet != nil {
-		return netstackNet.ListenUDPAddrPort(f.listenAddress)
-	}
-
-	return net.ListenUDP("udp", net.UDPAddrFromAddrPort(f.listenAddress))
-}
-
-func (f *DNSForwarder) createTCPListener(netstackNet *netstack.Net) (net.Listener, error) {
-	if netstackNet != nil {
-		return netstackNet.ListenTCPAddrPort(f.listenAddress)
-	}
-
-	return net.ListenTCP("tcp", net.TCPAddrFromAddrPort(f.listenAddress))
-}
-
 func (f *DNSForwarder) UpdateDomains(entries []*ForwarderEntry) {
 	f.mutex.Lock()
 	defer f.mutex.Unlock()

client/internal/dnsfwd/forwarder_test.go

@@ -297,7 +297,7 @@ func TestDNSForwarder_UnauthorizedDomainAccess(t *testing.T) {
 				mockResolver.On("LookupNetIP", mock.Anything, "ip4", dns.Fqdn(tt.queryDomain)).Return([]netip.Addr{fakeIP}, nil)
 			}
 
-			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
+			forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
 			forwarder.resolver = mockResolver
 
 			d, err := domain.FromString(tt.configuredDomain)
@@ -402,7 +402,7 @@ func TestDNSForwarder_FirewallSetUpdates(t *testing.T) {
 			mockResolver := &MockResolver{}
 
 			// Set up forwarder
-			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
+			forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
 			forwarder.resolver = mockResolver
 
 			// Create entries and track sets
@@ -489,7 +489,7 @@ func TestDNSForwarder_MultipleIPsInSingleUpdate(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	// Configure a single domain
@@ -584,7 +584,7 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
+			forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
 
 			d, err := domain.FromString(tt.configured)
 			require.NoError(t, err)
@@ -616,7 +616,7 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
 func TestDNSForwarder_TCPTruncation(t *testing.T) {
 	// Test that large UDP responses are truncated with TC bit set
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	d, _ := domain.FromString("example.com")
@@ -652,7 +652,7 @@ func TestDNSForwarder_TCPTruncation(t *testing.T) {
 // a subsequent upstream failure still returns a successful response from cache.
 func TestDNSForwarder_ServeFromCacheOnUpstreamFailure(t *testing.T) {
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("example.com")
@@ -696,7 +696,7 @@ func TestDNSForwarder_ServeFromCacheOnUpstreamFailure(t *testing.T) {
 // Verifies that cache normalization works across casing and trailing dot variations.
 func TestDNSForwarder_CacheNormalizationCasingAndDot(t *testing.T) {
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("ExAmPlE.CoM")
@@ -742,7 +742,7 @@ func TestDNSForwarder_MultipleOverlappingPatterns(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	// Set up complex overlapping patterns
@@ -804,7 +804,7 @@ func TestDNSForwarder_NodataVsNxdomain(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("example.com")
@@ -925,7 +925,7 @@ func TestDNSForwarder_NodataVsNxdomain(t *testing.T) {
 
 func TestDNSForwarder_EmptyQuery(t *testing.T) {
 	// Test handling of malformed query with no questions
-	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
+	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
 
 	query := &dns.Msg{}
 	// Don't set any question

client/internal/dnsfwd/manager.go

@@ -4,33 +4,27 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"net/netip"
-	"os"
-	"strconv"
+	"sync"
 
 	"github.com/hashicorp/go-multierror"
 	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/tun/netstack"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
 	"github.com/netbirdio/netbird/shared/management/domain"
 )
 
-const (
-	dnsTTL        = 60
-	envServerPort = "NB_DNS_FORWARDER_PORT"
+var (
+	// ListenPort is the port that the DNS forwarder listens on. It has been used by the client peers also
+	listenPort   uint16 = 5353
+	listenPortMu sync.RWMutex
 )
 
-// wgIface defines the interface for WireGuard interface operations needed by the DNS forwarder.
-type wgIface interface {
-	GetNet() *netstack.Net
-	Address() wgaddr.Address
-}
+const (
+	dnsTTL = 60 //seconds
+)
 
 // ForwarderEntry is a mapping from a domain to a resource ID and a hash of the parent domain list.
 type ForwarderEntry struct {
@@ -42,30 +36,28 @@ type ForwarderEntry struct {
 type Manager struct {
 	firewall       firewall.Manager
 	statusRecorder *peer.Status
-	wgIface        wgIface
-	serverPort     uint16
 
 	fwRules      []firewall.Rule
 	tcpRules     []firewall.Rule
 	dnsForwarder *DNSForwarder
 }
 
-func NewManager(fw firewall.Manager, statusRecorder *peer.Status, wgIface wgIface) *Manager {
-	serverPort := nbdns.ForwarderServerPort
-	if envPort := os.Getenv(envServerPort); envPort != "" {
-		if port, err := strconv.ParseUint(envPort, 10, 16); err == nil && port > 0 {
-			serverPort = uint16(port)
-			log.Infof("using custom DNS forwarder port from %s: %d", envServerPort, serverPort)
-		} else {
-			log.Warnf("invalid %s value %q, using default %d", envServerPort, envPort, nbdns.ForwarderServerPort)
-		}
-	}
+func ListenPort() uint16 {
+	listenPortMu.RLock()
+	defer listenPortMu.RUnlock()
+	return listenPort
+}
 
+func SetListenPort(port uint16) {
+	listenPortMu.Lock()
+	listenPort = port
+	listenPortMu.Unlock()
+}
+
+func NewManager(fw firewall.Manager, statusRecorder *peer.Status) *Manager {
 	return &Manager{
 		firewall:       fw,
 		statusRecorder: statusRecorder,
-		wgIface:        wgIface,
-		serverPort:     serverPort,
 	}
 }
 
@@ -79,25 +71,7 @@ func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
 		return err
 	}
 
-	localAddr := m.wgIface.Address().IP
-
-	if localAddr.IsValid() && m.firewall != nil {
-		if err := m.firewall.AddInboundDNAT(localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
-			log.Warnf("failed to add DNS UDP DNAT rule: %v", err)
-		} else {
-			log.Infof("added DNS UDP DNAT rule: %s:%d -> %s:%d", localAddr, nbdns.ForwarderClientPort, localAddr, m.serverPort)
-		}
-
-		if err := m.firewall.AddInboundDNAT(localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
-			log.Warnf("failed to add DNS TCP DNAT rule: %v", err)
-		} else {
-			log.Infof("added DNS TCP DNAT rule: %s:%d -> %s:%d", localAddr, nbdns.ForwarderClientPort, localAddr, m.serverPort)
-		}
-	}
-
-	listenAddress := netip.AddrPortFrom(localAddr, m.serverPort)
-	m.dnsForwarder = NewDNSForwarder(listenAddress, dnsTTL, m.firewall, m.statusRecorder, m.wgIface)
-
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort()), dnsTTL, m.firewall, m.statusRecorder)
 	go func() {
 		if err := m.dnsForwarder.Listen(fwdEntries); err != nil {
 			// todo handle close error if it is exists
@@ -122,18 +96,6 @@ func (m *Manager) Stop(ctx context.Context) error {
 	}
 
 	var mErr *multierror.Error
-
-	localAddr := m.wgIface.Address().IP
-	if localAddr.IsValid() && m.firewall != nil {
-		if err := m.firewall.RemoveInboundDNAT(localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
-			mErr = multierror.Append(mErr, fmt.Errorf("remove DNS UDP DNAT rule: %w", err))
-		}
-
-		if err := m.firewall.RemoveInboundDNAT(localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
-			mErr = multierror.Append(mErr, fmt.Errorf("remove DNS TCP DNAT rule: %w", err))
-		}
-	}
-
 	if err := m.dropDNSFirewall(); err != nil {
 		mErr = multierror.Append(mErr, err)
 	}
@@ -149,7 +111,7 @@ func (m *Manager) Stop(ctx context.Context) error {
 func (m *Manager) allowDNSFirewall() error {
 	dport := &firewall.Port{
 		IsRange: false,
-		Values:  []uint16{m.serverPort},
+		Values:  []uint16{ListenPort()},
 	}
 
 	if m.firewall == nil {

client/internal/engine.go

@@ -200,10 +200,13 @@ type Engine struct {
 	flowManager         nftypes.FlowManager
 
 	// WireGuard interface monitor
-	wgIfaceMonitor   *WGIfaceMonitor
-	wgIfaceMonitorWg sync.WaitGroup
+	wgIfaceMonitor *WGIfaceMonitor
 
-	probeStunTurn *relay.StunTurnProbe
+	// shutdownWg tracks all long-running goroutines to ensure clean shutdown
+	shutdownWg sync.WaitGroup
+
+	// dns forwarder port
+	dnsFwdPort uint16
 }
 
 // Peer is an instance of the Connection Peer
@@ -246,7 +249,7 @@ func NewEngine(
 		statusRecorder: statusRecorder,
 		checks:         checks,
 		connSemaphore:  semaphoregroup.NewSemaphoreGroup(connInitLimit),
-		probeStunTurn:  relay.NewStunTurnProbe(relay.DefaultCacheTTL),
+		dnsFwdPort:     dnsfwd.ListenPort(),
 	}
 
 	sm := profilemanager.NewServiceManager("")
@@ -325,10 +328,6 @@ func (e *Engine) Stop() error {
 		e.cancel()
 	}
 
-	// very ugly but we want to remove peers from the WireGuard interface first before removing interface.
-	// Removing peers happens in the conn.Close() asynchronously
-	time.Sleep(500 * time.Millisecond)
-
 	e.close()
 
 	// stop flow manager after wg interface is gone
@@ -336,8 +335,6 @@ func (e *Engine) Stop() error {
 		e.flowManager.Close()
 	}
 
-	log.Infof("stopped Netbird Engine")
-
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 	defer cancel()
 
@@ -348,12 +345,52 @@ func (e *Engine) Stop() error {
 		log.Errorf("failed to persist state: %v", err)
 	}
 
-	// Stop WireGuard interface monitor and wait for it to exit
-	e.wgIfaceMonitorWg.Wait()
+	timeout := e.calculateShutdownTimeout()
+	log.Debugf("waiting for goroutines to finish with timeout: %v", timeout)
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	if err := waitWithContext(shutdownCtx, &e.shutdownWg); err != nil {
+		log.Warnf("shutdown timeout exceeded after %v, some goroutines may still be running", timeout)
+	}
+
+	log.Infof("stopped Netbird Engine")
 
 	return nil
 }
 
+// calculateShutdownTimeout returns shutdown timeout: 10s base + 100ms per peer, capped at 30s.
+func (e *Engine) calculateShutdownTimeout() time.Duration {
+	peerCount := len(e.peerStore.PeersPubKey())
+
+	baseTimeout := 10 * time.Second
+	perPeerTimeout := time.Duration(peerCount) * 100 * time.Millisecond
+	timeout := baseTimeout + perPeerTimeout
+
+	maxTimeout := 30 * time.Second
+	if timeout > maxTimeout {
+		timeout = maxTimeout
+	}
+
+	return timeout
+}
+
+// waitWithContext waits for WaitGroup with timeout, returns ctx.Err() on timeout.
+func waitWithContext(ctx context.Context, wg *sync.WaitGroup) error {
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // Start creates a new WireGuard tunnel interface and listens to events from Signal and Management services
 // Connections to remote peers are not established here.
 // However, they will be established once an event with a list of peers to connect to will be received from Management Service
@@ -483,14 +520,14 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 
 	// monitor WireGuard interface lifecycle and restart engine on changes
 	e.wgIfaceMonitor = NewWGIfaceMonitor()
-	e.wgIfaceMonitorWg.Add(1)
+	e.shutdownWg.Add(1)
 
 	go func() {
-		defer e.wgIfaceMonitorWg.Done()
+		defer e.shutdownWg.Done()
 
 		if shouldRestart, err := e.wgIfaceMonitor.Start(e.ctx, e.wgInterface.Name()); shouldRestart {
 			log.Infof("WireGuard interface monitor: %s, restarting engine", err)
-			e.restartEngine()
+			e.triggerClientRestart()
 		} else if err != nil {
 			log.Warnf("WireGuard interface monitor: %s", err)
 		}
@@ -891,7 +928,9 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 			if err != nil {
 				return fmt.Errorf("create ssh server: %w", err)
 			}
+			e.shutdownWg.Add(1)
 			go func() {
+				defer e.shutdownWg.Done()
 				// blocking
 				err = e.sshServer.Start()
 				if err != nil {
@@ -949,7 +988,9 @@ func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
 // receiveManagementEvents connects to the Management Service event stream to receive updates from the management service
 // E.g. when a new peer has been registered and we are allowed to connect to it.
 func (e *Engine) receiveManagementEvents() {
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		info, err := system.GetInfoWithChecks(e.ctx, e.checks)
 		if err != nil {
 			log.Warnf("failed to get system info with checks: %v", err)
@@ -1059,14 +1100,10 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		protoDNSConfig = &mgmProto.DNSConfig{}
 	}
 
-	dnsConfig := toDNSConfig(protoDNSConfig, e.wgInterface.Address().Network)
-
-	if err := e.dnsServer.UpdateDNSServer(serial, dnsConfig); err != nil {
+	if err := e.dnsServer.UpdateDNSServer(serial, toDNSConfig(protoDNSConfig, e.wgInterface.Address().Network)); err != nil {
 		log.Errorf("failed to update dns server, err: %v", err)
 	}
 
-	e.routeManager.SetDNSForwarderPort(dnsConfig.ForwarderPort)
-
 	// apply routes first, route related actions might depend on routing being enabled
 	routes := toRoutes(networkMap.GetRoutes())
 	serverRoutes, clientRoutes := e.routeManager.ClassifyRoutes(routes)
@@ -1087,7 +1124,7 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 	}
 
 	fwdEntries := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), routes)
-	e.updateDNSForwarder(dnsRouteFeatureFlag, fwdEntries)
+	e.updateDNSForwarder(dnsRouteFeatureFlag, fwdEntries, uint16(protoDNSConfig.ForwarderPort))
 
 	// Ingress forward rules
 	forwardingRules, err := e.updateForwardRules(networkMap.GetForwardingRules())
@@ -1211,16 +1248,10 @@ func toRouteDomains(myPubKey string, routes []*route.Route) []*dnsfwd.ForwarderE
 }
 
 func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig, network netip.Prefix) nbdns.Config {
-	forwarderPort := uint16(protoDNSConfig.GetForwarderPort())
-	if forwarderPort == 0 {
-		forwarderPort = nbdns.ForwarderClientPort
-	}
-
 	dnsUpdate := nbdns.Config{
 		ServiceEnable:    protoDNSConfig.GetServiceEnable(),
 		CustomZones:      make([]nbdns.CustomZone, 0),
 		NameServerGroups: make([]*nbdns.NameServerGroup, 0),
-		ForwarderPort:    forwarderPort,
 	}
 
 	for _, zone := range protoDNSConfig.GetCustomZones() {
@@ -1377,7 +1408,9 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 
 // receiveSignalEvents connects to the Signal Service event stream to negotiate connection with remote peers
 func (e *Engine) receiveSignalEvents() {
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		// connect to a stream of messages coming from the signal server
 		err := e.signal.Receive(e.ctx, func(msg *sProto.Message) error {
 			e.syncMsgMux.Lock()
@@ -1676,7 +1709,7 @@ func (e *Engine) getRosenpassAddr() string {
 
 // RunHealthProbes executes health checks for Signal, Management, Relay and WireGuard services
 // and updates the status recorder with the latest states.
-func (e *Engine) RunHealthProbes(waitForResult bool) bool {
+func (e *Engine) RunHealthProbes() bool {
 	e.syncMsgMux.Lock()
 
 	signalHealthy := e.signal.IsHealthy()
@@ -1708,12 +1741,8 @@ func (e *Engine) RunHealthProbes(waitForResult bool) bool {
 	}
 
 	e.syncMsgMux.Unlock()
-	var results []relay.ProbeResult
-	if waitForResult {
-		results = e.probeStunTurn.ProbeAllWaitResult(e.ctx, stuns, turns)
-	} else {
-		results = e.probeStunTurn.ProbeAll(e.ctx, stuns, turns)
-	}
+
+	results := e.probeICE(stuns, turns)
 	e.statusRecorder.UpdateRelayStates(results)
 
 	relayHealthy := true
@@ -1730,8 +1759,17 @@ func (e *Engine) RunHealthProbes(waitForResult bool) bool {
 	return allHealthy
 }
 
-// restartEngine restarts the engine by cancelling the client context
-func (e *Engine) restartEngine() {
+func (e *Engine) probeICE(stuns, turns []*stun.URI) []relay.ProbeResult {
+	return append(
+		relay.ProbeAll(e.ctx, relay.ProbeSTUN, stuns),
+		relay.ProbeAll(e.ctx, relay.ProbeTURN, turns)...,
+	)
+}
+
+// triggerClientRestart triggers a full client restart by cancelling the client context.
+// Note: This does NOT just restart the engine - it cancels the entire client context,
+// which causes the connect client's retry loop to create a completely new engine.
+func (e *Engine) triggerClientRestart() {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 
@@ -1753,7 +1791,9 @@ func (e *Engine) startNetworkMonitor() {
 	}
 
 	e.networkMonitor = networkmonitor.New()
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		if err := e.networkMonitor.Listen(e.ctx); err != nil {
 			if errors.Is(err, context.Canceled) {
 				log.Infof("network monitor stopped")
@@ -1763,8 +1803,8 @@ func (e *Engine) startNetworkMonitor() {
 			return
 		}
 
-		log.Infof("Network monitor: detected network change, restarting engine")
-		e.restartEngine()
+		log.Infof("Network monitor: detected network change, triggering client restart")
+		e.triggerClientRestart()
 	}()
 }
 
@@ -1849,75 +1889,63 @@ func (e *Engine) GetWgAddr() netip.Addr {
 func (e *Engine) updateDNSForwarder(
 	enabled bool,
 	fwdEntries []*dnsfwd.ForwarderEntry,
+	forwarderPort uint16,
 ) {
 	if e.config.DisableServerRoutes {
 		return
 	}
 
+	if forwarderPort > 0 {
+		dnsfwd.SetListenPort(forwarderPort)
+	}
+
 	if !enabled {
-		e.stopDNSForwarder()
+		if e.dnsForwardMgr == nil {
+			return
+		}
+		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+			log.Errorf("failed to stop DNS forward: %v", err)
+		}
 		return
 	}
 
 	if len(fwdEntries) > 0 {
-		if e.dnsForwardMgr == nil {
-			e.startDNSForwarder(fwdEntries)
-		} else {
+		switch {
+		case e.dnsForwardMgr == nil:
+			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder)
+			if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
+				log.Errorf("failed to start DNS forward: %v", err)
+				e.dnsForwardMgr = nil
+			}
+			log.Infof("started domain router service with %d entries", len(fwdEntries))
+		case e.dnsFwdPort != forwarderPort:
+			log.Infof("updating domain router service port from %d to %d", e.dnsFwdPort, forwarderPort)
+			e.restartDnsFwd(fwdEntries, forwarderPort)
+			e.dnsFwdPort = forwarderPort
+
+		default:
 			e.dnsForwardMgr.UpdateDomains(fwdEntries)
 		}
 	} else if e.dnsForwardMgr != nil {
 		log.Infof("disable domain router service")
-		e.stopDNSForwarder()
-	}
-}
-
-func (e *Engine) startDNSForwarder(fwdEntries []*dnsfwd.ForwarderEntry) {
-	e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder, e.wgInterface)
-	e.registerDNSServices()
-
-	if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
-		log.Errorf("failed to start DNS forward: %v", err)
+		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+			log.Errorf("failed to stop DNS forward: %v", err)
+		}
 		e.dnsForwardMgr = nil
-		return
 	}
 
-	log.Infof("started domain router service with %d entries", len(fwdEntries))
 }
 
-func (e *Engine) stopDNSForwarder() {
-	if e.dnsForwardMgr == nil {
-		return
-	}
-
+func (e *Engine) restartDnsFwd(fwdEntries []*dnsfwd.ForwarderEntry, forwarderPort uint16) {
+	log.Infof("updating domain router service port from %d to %d", e.dnsFwdPort, forwarderPort)
+	// stop and start the forwarder to apply the new port
 	if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
 		log.Errorf("failed to stop DNS forward: %v", err)
 	}
-
-	e.unregisterDNSServices()
-	e.dnsForwardMgr = nil
-}
-
-func (e *Engine) registerDNSServices() {
-	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
-		if registrar, ok := e.firewall.(interface {
-			RegisterNetstackService(protocol nftypes.Protocol, port uint16)
-		}); ok {
-			registrar.RegisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
-			registrar.RegisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
-			log.Debugf("registered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
-		}
-	}
-}
-
-func (e *Engine) unregisterDNSServices() {
-	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
-		if registrar, ok := e.firewall.(interface {
-			UnregisterNetstackService(protocol nftypes.Protocol, port uint16)
-		}); ok {
-			registrar.UnregisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
-			registrar.UnregisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
-			log.Debugf("unregistered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
-		}
+	e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder)
+	if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
+		log.Errorf("failed to start DNS forward: %v", err)
+		e.dnsForwardMgr = nil
 	}
 }
 

client/internal/netflow/logger/logger.go

@@ -10,10 +10,10 @@ import (
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/internal/dnsfwd"
 	"github.com/netbirdio/netbird/client/internal/netflow/store"
 	"github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	"github.com/netbirdio/netbird/dns"
 )
 
 type rcvChan chan *types.EventFields
@@ -138,8 +138,7 @@ func (l *Logger) UpdateConfig(dnsCollection, exitNodeCollection bool) {
 
 func (l *Logger) shouldStore(event *types.EventFields, isExitNode bool) bool {
 	// check dns collection
-	if !l.dnsCollection.Load() && event.Protocol == types.UDP &&
-		(event.DestPort == 53 || event.DestPort == dns.ForwarderClientPort || event.DestPort == dns.ForwarderServerPort) {
+	if !l.dnsCollection.Load() && event.Protocol == types.UDP && (event.DestPort == 53 || event.DestPort == uint16(dnsfwd.ListenPort())) {
 		return false
 	}
 

client/internal/netflow/manager.go

@@ -24,6 +24,7 @@ import (
 // Manager handles netflow tracking and logging
 type Manager struct {
 	mux            sync.Mutex
+	shutdownWg     sync.WaitGroup
 	logger         nftypes.FlowLogger
 	flowConfig     *nftypes.FlowConfig
 	conntrack      nftypes.ConnTracker
@@ -105,8 +106,15 @@ func (m *Manager) resetClient() error {
 	ctx, cancel := context.WithCancel(context.Background())
 	m.cancel = cancel
 
-	go m.receiveACKs(ctx, flowClient)
-	go m.startSender(ctx)
+	m.shutdownWg.Add(2)
+	go func() {
+		defer m.shutdownWg.Done()
+		m.receiveACKs(ctx, flowClient)
+	}()
+	go func() {
+		defer m.shutdownWg.Done()
+		m.startSender(ctx)
+	}()
 
 	return nil
 }
@@ -176,11 +184,12 @@ func (m *Manager) Update(update *nftypes.FlowConfig) error {
 // Close cleans up all resources
 func (m *Manager) Close() {
 	m.mux.Lock()
-	defer m.mux.Unlock()
-
 	if err := m.disableFlow(); err != nil {
 		log.Warnf("failed to disable flow manager: %v", err)
 	}
+	m.mux.Unlock()
+
+	m.shutdownWg.Wait()
 }
 
 // GetLogger returns the flow logger

client/internal/networkmonitor/check_change_bsd.go

@@ -1,4 +1,4 @@
-//go:build (darwin && !ios) || dragonfly || freebsd || netbsd || openbsd
+//go:build dragonfly || freebsd || netbsd || openbsd
 
 package networkmonitor
 

client/internal/networkmonitor/check_change_darwin.go

@@ -0,0 +1,344 @@
+//go:build darwin && !ios
+
+package networkmonitor
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"hash/fnv"
+	"net/netip"
+	"os/exec"
+	"syscall"
+	"time"
+	"unsafe"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/net/route"
+	"golang.org/x/sys/unix"
+
+	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
+)
+
+// todo: refactor to not use static functions
+
+func checkChange(ctx context.Context, nexthopv4, nexthopv6 systemops.Nexthop) error {
+	fd, err := unix.Socket(syscall.AF_ROUTE, syscall.SOCK_RAW, syscall.AF_UNSPEC)
+	if err != nil {
+		return fmt.Errorf("open routing socket: %v", err)
+	}
+
+	defer func() {
+		err := unix.Close(fd)
+		if err != nil && !errors.Is(err, unix.EBADF) {
+			log.Warnf("Network monitor: failed to close routing socket: %v", err)
+		}
+	}()
+
+	routeChanged := make(chan struct{})
+	go func() {
+		routeCheck(ctx, fd, nexthopv4, nexthopv6)
+		close(routeChanged)
+	}()
+
+	wakeUp := make(chan struct{})
+	go func() {
+		wakeUpListen(ctx)
+		close(wakeUp)
+	}()
+
+	gatewayChanged := make(chan string)
+	go func() {
+		gatewayPoll(ctx, nexthopv4, nexthopv6, gatewayChanged)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-routeChanged:
+		log.Infof("route change detected via routing socket")
+		return nil
+	case <-wakeUp:
+		log.Infof("wakeup detected via sleep hash change")
+		return nil
+	case reason := <-gatewayChanged:
+		log.Infof("gateway change detected via polling: %s", reason)
+		return nil
+	}
+}
+
+func routeCheck(ctx context.Context, fd int, nexthopv4 systemops.Nexthop, nexthopv6 systemops.Nexthop) {
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+
+		buf := make([]byte, 2048)
+		n, err := unix.Read(fd, buf)
+		if err != nil {
+			if !errors.Is(err, unix.EBADF) && !errors.Is(err, unix.EINVAL) {
+				log.Warnf("Network monitor: failed to read from routing socket: %v", err)
+			}
+			continue
+		}
+		if n < unix.SizeofRtMsghdr {
+			log.Debugf("Network monitor: read from routing socket returned less than expected: %d bytes", n)
+			continue
+		}
+
+		msg := (*unix.RtMsghdr)(unsafe.Pointer(&buf[0]))
+
+		switch msg.Type {
+		// handle route changes
+		case unix.RTM_ADD, syscall.RTM_DELETE:
+			route, err := parseRouteMessage(buf[:n])
+			if err != nil {
+				log.Debugf("Network monitor: error parsing routing message: %v", err)
+				continue
+			}
+
+			if route.Dst.Bits() != 0 {
+				continue
+			}
+
+			intf := "<nil>"
+			if route.Interface != nil {
+				intf = route.Interface.Name
+			}
+			switch msg.Type {
+			case unix.RTM_ADD:
+				log.Infof("Network monitor: default route changed: via %s, interface %s", route.Gw, intf)
+				return
+			case unix.RTM_DELETE:
+				if nexthopv4.Intf != nil && route.Gw.Compare(nexthopv4.IP) == 0 || nexthopv6.Intf != nil && route.Gw.Compare(nexthopv6.IP) == 0 {
+					log.Infof("Network monitor: default route removed: via %s, interface %s", route.Gw, intf)
+					return
+				}
+			}
+		}
+	}
+}
+
+func parseRouteMessage(buf []byte) (*systemops.Route, error) {
+	msgs, err := route.ParseRIB(route.RIBTypeRoute, buf)
+	if err != nil {
+		return nil, fmt.Errorf("parse RIB: %v", err)
+	}
+
+	if len(msgs) != 1 {
+		return nil, fmt.Errorf("unexpected RIB message msgs: %v", msgs)
+	}
+
+	msg, ok := msgs[0].(*route.RouteMessage)
+	if !ok {
+		return nil, fmt.Errorf("unexpected RIB message type: %T", msgs[0])
+	}
+
+	return systemops.MsgToRoute(msg)
+}
+
+func wakeUpListen(ctx context.Context) {
+	log.Infof("start to watch for system wakeups")
+	var (
+		initialHash uint32
+		err         error
+	)
+
+	// Keep retrying until initial sysctl succeeds or context is canceled
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("exit from wakeUpListen initial hash detection due to context cancellation")
+			return
+		default:
+			initialHash, err = readSleepTimeHash()
+			if err != nil {
+				log.Errorf("failed to detect initial sleep time: %v", err)
+				select {
+				case <-ctx.Done():
+					log.Info("exit from wakeUpListen initial hash detection due to context cancellation")
+					return
+				case <-time.After(3 * time.Second):
+					continue
+				}
+			}
+			log.Infof("initial wakeup hash: %d", initialHash)
+			break
+		}
+		break
+	}
+
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	lastCheck := time.Now()
+	const maxTickerDrift = 1 * time.Minute
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("context canceled, stopping wakeUpListen")
+			return
+
+		case <-ticker.C:
+			now := time.Now()
+			elapsed := now.Sub(lastCheck)
+
+			// If more time passed than expected, system likely slept (informational only)
+			if elapsed > maxTickerDrift {
+				upOut, err := exec.Command("uptime").Output()
+				if err != nil {
+					log.Errorf("failed to run uptime command: %v", err)
+					upOut = []byte("unknown")
+				}
+				log.Infof("Time drift detected (potential wakeup): expected ~5s, actual %s, uptime: %s", elapsed, upOut)
+
+				currentV4, errV4 := systemops.GetNextHop(netip.IPv4Unspecified())
+				currentV6, errV6 := systemops.GetNextHop(netip.IPv6Unspecified())
+				if errV4 == nil {
+					log.Infof("Current IPv4 default gateway: %s via %s", currentV4.IP, currentV4.Intf.Name)
+				} else {
+					log.Debugf("No IPv4 default gateway: %v", errV4)
+				}
+				if errV6 == nil {
+					log.Infof("Current IPv6 default gateway: %s via %s", currentV6.IP, currentV6.Intf.Name)
+				} else {
+					log.Debugf("No IPv6 default gateway: %v", errV6)
+				}
+			}
+
+			newHash, err := readSleepTimeHash()
+			if err != nil {
+				log.Errorf("failed to read sleep time hash: %v", err)
+				lastCheck = now
+				continue
+			}
+
+			if newHash == initialHash {
+				log.Debugf("no wakeup detected (hash unchanged: %d, time drift: %s)", initialHash, elapsed)
+				lastCheck = now
+				continue
+			}
+
+			upOut, err := exec.Command("uptime").Output()
+			if err != nil {
+				log.Errorf("failed to run uptime command: %v", err)
+				upOut = []byte("unknown")
+			}
+			log.Infof("Wakeup detected via hash change: %d -> %d, uptime: %s", initialHash, newHash, upOut)
+
+			currentV4, errV4 := systemops.GetNextHop(netip.IPv4Unspecified())
+			currentV6, errV6 := systemops.GetNextHop(netip.IPv6Unspecified())
+			if errV4 == nil {
+				log.Infof("Current IPv4 default gateway after wakeup: %s via %s", currentV4.IP, currentV4.Intf.Name)
+			} else {
+				log.Debugf("No IPv4 default gateway after wakeup: %v", errV4)
+			}
+			if errV6 == nil {
+				log.Infof("Current IPv6 default gateway after wakeup: %s via %s", currentV6.IP, currentV6.Intf.Name)
+			} else {
+				log.Debugf("No IPv6 default gateway after wakeup: %v", errV6)
+			}
+
+			return
+		}
+	}
+}
+
+func readSleepTimeHash() (uint32, error) {
+	cmd := exec.Command("sysctl", "kern.sleeptime")
+	out, err := cmd.Output()
+	if err != nil {
+		return 0, fmt.Errorf("failed to run sysctl: %w", err)
+	}
+
+	h, err := hash(out)
+	if err != nil {
+		return 0, fmt.Errorf("failed to compute hash: %w", err)
+	}
+
+	return h, nil
+}
+
+func hash(data []byte) (uint32, error) {
+	hasher := fnv.New32a()
+	if _, err := hasher.Write(data); err != nil {
+		return 0, err
+	}
+	return hasher.Sum32(), nil
+}
+
+// gatewayPoll polls the default gateway every 5 seconds to detect changes that might be missed by routing socket or wake-up detection.
+func gatewayPoll(ctx context.Context, initialV4, initialV6 systemops.Nexthop, changed chan<- string) {
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	log.Infof("Gateway polling started - initial v4: %s via %v, v6: %s via %v",
+		initialV4.IP, initialV4.Intf, initialV6.IP, initialV6.Intf)
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Debug("context canceled, stopping gateway polling")
+			return
+
+		case <-ticker.C:
+			currentV4, errV4 := systemops.GetNextHop(netip.IPv4Unspecified())
+			currentV6, errV6 := systemops.GetNextHop(netip.IPv6Unspecified())
+
+			var reason string
+
+			if errV4 == nil && initialV4.IP.IsValid() {
+				if currentV4.IP.Compare(initialV4.IP) != 0 {
+					reason = fmt.Sprintf("IPv4 gateway changed from %s to %s", initialV4.IP, currentV4.IP)
+					log.Infof("Gateway poll detected change: %s", reason)
+					changed <- reason
+					return
+				}
+				if initialV4.Intf != nil && currentV4.Intf != nil && currentV4.Intf.Name != initialV4.Intf.Name {
+					reason = fmt.Sprintf("IPv4 interface changed from %s to %s", initialV4.Intf.Name, currentV4.Intf.Name)
+					log.Infof("Gateway poll detected change: %s", reason)
+					changed <- reason
+					return
+				}
+			} else if errV4 == nil && !initialV4.IP.IsValid() {
+				reason = "IPv4 gateway appeared"
+				log.Infof("Gateway poll detected change: %s (new: %s)", reason, currentV4.IP)
+				changed <- reason
+				return
+			} else if errV4 != nil && initialV4.IP.IsValid() {
+				reason = "IPv4 gateway disappeared"
+				log.Infof("Gateway poll detected change: %s", reason)
+				changed <- reason
+				return
+			}
+
+			if errV6 == nil && initialV6.IP.IsValid() {
+				if currentV6.IP.Compare(initialV6.IP) != 0 {
+					reason = fmt.Sprintf("IPv6 gateway changed from %s to %s", initialV6.IP, currentV6.IP)
+					log.Infof("Gateway poll detected change: %s", reason)
+					changed <- reason
+					return
+				}
+				if initialV6.Intf != nil && currentV6.Intf != nil && currentV6.Intf.Name != initialV6.Intf.Name {
+					reason = fmt.Sprintf("IPv6 interface changed from %s to %s", initialV6.Intf.Name, currentV6.Intf.Name)
+					log.Infof("Gateway poll detected change: %s", reason)
+					changed <- reason
+					return
+				}
+			} else if errV6 == nil && !initialV6.IP.IsValid() {
+				reason = "IPv6 gateway appeared"
+				log.Infof("Gateway poll detected change: %s (new: %s)", reason, currentV6.IP)
+				changed <- reason
+				return
+			} else if errV6 != nil && initialV6.IP.IsValid() {
+				reason = "IPv6 gateway disappeared"
+				log.Infof("Gateway poll detected change: %s", reason)
+				changed <- reason
+				return
+			}
+
+			log.Debugf("Gateway poll: no change detected")
+		}
+	}
+}

client/internal/networkmonitor/monitor.go

@@ -88,6 +88,7 @@ func (nw *NetworkMonitor) Listen(ctx context.Context) (err error) {
 	event := make(chan struct{}, 1)
 	go nw.checkChanges(ctx, event, nexthop4, nexthop6)
 
+	log.Infof("start watching for network changes")
 	// debounce changes
 	timer := time.NewTimer(0)
 	timer.Stop()

client/internal/peer/guard/sr_watcher.go

@@ -19,11 +19,11 @@ type SRWatcher struct {
 	signalClient chNotifier
 	relayManager chNotifier
 
-	listeners     map[chan struct{}]struct{}
-	mu            sync.Mutex
-	iFaceDiscover stdnet.ExternalIFaceDiscover
-	iceConfig     ice.Config
-
+	listeners        map[chan struct{}]struct{}
+	mu               sync.Mutex
+	shutdownWg       sync.WaitGroup
+	iFaceDiscover    stdnet.ExternalIFaceDiscover
+	iceConfig        ice.Config
 	cancelIceMonitor context.CancelFunc
 }
 
@@ -52,7 +52,11 @@ func (w *SRWatcher) Start() {
 	w.cancelIceMonitor = cancel
 
 	iceMonitor := NewICEMonitor(w.iFaceDiscover, w.iceConfig, GetICEMonitorPeriod())
-	go iceMonitor.Start(ctx, w.onICEChanged)
+	w.shutdownWg.Add(1)
+	go func() {
+		defer w.shutdownWg.Done()
+		iceMonitor.Start(ctx, w.onICEChanged)
+	}()
 	w.signalClient.SetOnReconnectedListener(w.onReconnected)
 	w.relayManager.SetOnReconnectedListener(w.onReconnected)
 
@@ -60,14 +64,16 @@ func (w *SRWatcher) Start() {
 
 func (w *SRWatcher) Close() {
 	w.mu.Lock()
-	defer w.mu.Unlock()
-
 	if w.cancelIceMonitor == nil {
+		w.mu.Unlock()
 		return
 	}
 	w.cancelIceMonitor()
 	w.signalClient.SetOnReconnectedListener(nil)
 	w.relayManager.SetOnReconnectedListener(nil)
+	w.mu.Unlock()
+
+	w.shutdownWg.Wait()
 }
 
 func (w *SRWatcher) NewListener() chan struct{} {

client/internal/relay/relay.go

@@ -2,8 +2,6 @@ package relay
 
 import (
 	"context"
-	"crypto/sha256"
-	"errors"
 	"fmt"
 	"net"
 	"sync"
@@ -17,15 +15,6 @@ import (
 	nbnet "github.com/netbirdio/netbird/client/net"
 )
 
-const (
-	DefaultCacheTTL = 20 * time.Second
-	probeTimeout    = 6 * time.Second
-)
-
-var (
-	ErrCheckInProgress = errors.New("probe check is already in progress")
-)
-
 // ProbeResult holds the info about the result of a relay probe request
 type ProbeResult struct {
 	URI  string
@@ -33,164 +22,8 @@ type ProbeResult struct {
 	Addr string
 }
 
-type StunTurnProbe struct {
-	cacheResults    []ProbeResult
-	cacheTimestamp  time.Time
-	cacheKey        string
-	cacheTTL        time.Duration
-	probeInProgress bool
-	probeDone       chan struct{}
-	mu              sync.Mutex
-}
-
-func NewStunTurnProbe(cacheTTL time.Duration) *StunTurnProbe {
-	return &StunTurnProbe{
-		cacheTTL: cacheTTL,
-	}
-}
-
-func (p *StunTurnProbe) ProbeAllWaitResult(ctx context.Context, stuns []*stun.URI, turns []*stun.URI) []ProbeResult {
-	cacheKey := generateCacheKey(stuns, turns)
-
-	p.mu.Lock()
-	if p.probeInProgress {
-		doneChan := p.probeDone
-		p.mu.Unlock()
-
-		select {
-		case <-ctx.Done():
-			log.Debugf("Context cancelled while waiting for probe results")
-			return createErrorResults(stuns, turns)
-		case <-doneChan:
-			return p.getCachedResults(cacheKey, stuns, turns)
-		}
-	}
-
-	p.probeInProgress = true
-	probeDone := make(chan struct{})
-	p.probeDone = probeDone
-	p.mu.Unlock()
-
-	p.doProbe(ctx, stuns, turns, cacheKey)
-	close(probeDone)
-
-	return p.getCachedResults(cacheKey, stuns, turns)
-}
-
-// ProbeAll probes all given servers asynchronously and returns the results
-func (p *StunTurnProbe) ProbeAll(ctx context.Context, stuns []*stun.URI, turns []*stun.URI) []ProbeResult {
-	cacheKey := generateCacheKey(stuns, turns)
-
-	p.mu.Lock()
-
-	if results := p.checkCache(cacheKey); results != nil {
-		p.mu.Unlock()
-		return results
-	}
-
-	if p.probeInProgress {
-		p.mu.Unlock()
-		return createErrorResults(stuns, turns)
-	}
-
-	p.probeInProgress = true
-	probeDone := make(chan struct{})
-	p.probeDone = probeDone
-	log.Infof("started new probe for STUN, TURN servers")
-	go func() {
-		p.doProbe(ctx, stuns, turns, cacheKey)
-		close(probeDone)
-	}()
-
-	p.mu.Unlock()
-
-	timer := time.NewTimer(1300 * time.Millisecond)
-	defer timer.Stop()
-
-	select {
-	case <-ctx.Done():
-		log.Debugf("Context cancelled while waiting for probe results")
-		return createErrorResults(stuns, turns)
-	case <-probeDone:
-		// when the probe is return fast, return the results right away
-		return p.getCachedResults(cacheKey, stuns, turns)
-	case <-timer.C:
-		// if the probe takes longer than 1.3s, return error results to avoid blocking
-		return createErrorResults(stuns, turns)
-	}
-}
-
-func (p *StunTurnProbe) checkCache(cacheKey string) []ProbeResult {
-	if p.cacheKey == cacheKey && len(p.cacheResults) > 0 {
-		age := time.Since(p.cacheTimestamp)
-		if age < p.cacheTTL {
-			results := append([]ProbeResult(nil), p.cacheResults...)
-			log.Debugf("returning cached probe results (age: %v)", age)
-			return results
-		}
-	}
-	return nil
-}
-
-func (p *StunTurnProbe) getCachedResults(cacheKey string, stuns []*stun.URI, turns []*stun.URI) []ProbeResult {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-
-	if p.cacheKey == cacheKey && len(p.cacheResults) > 0 {
-		return append([]ProbeResult(nil), p.cacheResults...)
-	}
-	return createErrorResults(stuns, turns)
-}
-
-func (p *StunTurnProbe) doProbe(ctx context.Context, stuns []*stun.URI, turns []*stun.URI, cacheKey string) {
-	defer func() {
-		p.mu.Lock()
-		p.probeInProgress = false
-		p.mu.Unlock()
-	}()
-	results := make([]ProbeResult, len(stuns)+len(turns))
-
-	var wg sync.WaitGroup
-	for i, uri := range stuns {
-		wg.Add(1)
-		go func(idx int, stunURI *stun.URI) {
-			defer wg.Done()
-
-			probeCtx, cancel := context.WithTimeout(ctx, probeTimeout)
-			defer cancel()
-
-			results[idx].URI = stunURI.String()
-			results[idx].Addr, results[idx].Err = p.probeSTUN(probeCtx, stunURI)
-		}(i, uri)
-	}
-
-	stunOffset := len(stuns)
-	for i, uri := range turns {
-		wg.Add(1)
-		go func(idx int, turnURI *stun.URI) {
-			defer wg.Done()
-
-			probeCtx, cancel := context.WithTimeout(ctx, probeTimeout)
-			defer cancel()
-
-			results[idx].URI = turnURI.String()
-			results[idx].Addr, results[idx].Err = p.probeTURN(probeCtx, turnURI)
-		}(stunOffset+i, uri)
-	}
-
-	wg.Wait()
-
-	p.mu.Lock()
-	p.cacheResults = results
-	p.cacheTimestamp = time.Now()
-	p.cacheKey = cacheKey
-	p.mu.Unlock()
-
-	log.Debug("Stored new probe results in cache")
-}
-
 // ProbeSTUN tries binding to the given STUN uri and acquiring an address
-func (p *StunTurnProbe) probeSTUN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
+func ProbeSTUN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
 	defer func() {
 		if probeErr != nil {
 			log.Debugf("stun probe error from %s: %s", uri, probeErr)
@@ -250,7 +83,7 @@ func (p *StunTurnProbe) probeSTUN(ctx context.Context, uri *stun.URI) (addr stri
 }
 
 // ProbeTURN tries allocating a session from the given TURN URI
-func (p *StunTurnProbe) probeTURN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
+func ProbeTURN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
 	defer func() {
 		if probeErr != nil {
 			log.Debugf("turn probe error from %s: %s", uri, probeErr)
@@ -327,28 +160,28 @@ func (p *StunTurnProbe) probeTURN(ctx context.Context, uri *stun.URI) (addr stri
 	return relayConn.LocalAddr().String(), nil
 }
 
-func createErrorResults(stuns []*stun.URI, turns []*stun.URI) []ProbeResult {
-	total := len(stuns) + len(turns)
-	results := make([]ProbeResult, total)
+// ProbeAll probes all given servers asynchronously and returns the results
+func ProbeAll(
+	ctx context.Context,
+	fn func(ctx context.Context, uri *stun.URI) (addr string, probeErr error),
+	relays []*stun.URI,
+) []ProbeResult {
+	results := make([]ProbeResult, len(relays))
 
-	allURIs := append(append([]*stun.URI{}, stuns...), turns...)
-	for i, uri := range allURIs {
-		results[i] = ProbeResult{
-			URI: uri.String(),
-			Err: ErrCheckInProgress,
-		}
+	var wg sync.WaitGroup
+	for i, uri := range relays {
+		ctx, cancel := context.WithTimeout(ctx, 6*time.Second)
+		defer cancel()
+
+		wg.Add(1)
+		go func(res *ProbeResult, stunURI *stun.URI) {
+			defer wg.Done()
+			res.URI = stunURI.String()
+			res.Addr, res.Err = fn(ctx, stunURI)
+		}(&results[i], uri)
 	}
 
+	wg.Wait()
+
 	return results
 }
-
-func generateCacheKey(stuns []*stun.URI, turns []*stun.URI) string {
-	h := sha256.New()
-	for _, uri := range stuns {
-		h.Write([]byte(uri.String()))
-	}
-	for _, uri := range turns {
-		h.Write([]byte(uri.String()))
-	}
-	return fmt.Sprintf("%x", h.Sum(nil))
-}

client/internal/routemanager/common/params.go

@@ -1,7 +1,6 @@
 package common
 
 import (
-	"sync/atomic"
 	"time"
 
 	"github.com/netbirdio/netbird/client/firewall/manager"
@@ -26,5 +25,4 @@ type HandlerParams struct {
 	UseNewDNSRoute       bool
 	Firewall             manager.Manager
 	FakeIPManager        *fakeip.Manager
-	ForwarderPort        *atomic.Uint32
 }

client/internal/routemanager/dnsinterceptor/handler.go

@@ -8,7 +8,6 @@ import (
 	"runtime"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
@@ -19,6 +18,7 @@ import (
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	nbdns "github.com/netbirdio/netbird/client/internal/dns"
+	"github.com/netbirdio/netbird/client/internal/dnsfwd"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/client/internal/routemanager/common"
@@ -55,7 +55,6 @@ type DnsInterceptor struct {
 	peerStore            *peerstore.Store
 	firewall             firewall.Manager
 	fakeIPManager        *fakeip.Manager
-	forwarderPort        *atomic.Uint32
 }
 
 func New(params common.HandlerParams) *DnsInterceptor {
@@ -70,7 +69,6 @@ func New(params common.HandlerParams) *DnsInterceptor {
 		firewall:             params.Firewall,
 		fakeIPManager:        params.FakeIPManager,
 		interceptedDomains:   make(domainMap),
-		forwarderPort:        params.ForwarderPort,
 	}
 }
 
@@ -259,7 +257,7 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 		r.MsgHdr.AuthenticatedData = true
 	}
 
-	upstream := fmt.Sprintf("%s:%d", upstreamIP.String(), uint16(d.forwarderPort.Load()))
+	upstream := fmt.Sprintf("%s:%d", upstreamIP.String(), dnsfwd.ListenPort())
 	ctx, cancel := context.WithTimeout(context.Background(), dnsTimeout)
 	defer cancel()
 

client/internal/routemanager/manager.go

@@ -10,7 +10,6 @@ import (
 	"runtime"
 	"slices"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -24,7 +23,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/netstack"
 	"github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/listener"
-	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/client/internal/routemanager/client"
@@ -56,7 +54,6 @@ type Manager interface {
 	SetRouteChangeListener(listener listener.NetworkChangeListener)
 	InitialRouteRange() []string
 	SetFirewall(firewall.Manager) error
-	SetDNSForwarderPort(port uint16)
 	Stop(stateManager *statemanager.Manager)
 }
 
@@ -81,6 +78,7 @@ type DefaultManager struct {
 	ctx                  context.Context
 	stop                 context.CancelFunc
 	mux                  sync.Mutex
+	shutdownWg           sync.WaitGroup
 	clientNetworks       map[route.HAUniqueID]*client.Watcher
 	routeSelector        *routeselector.RouteSelector
 	serverRouter         *server.Router
@@ -104,7 +102,6 @@ type DefaultManager struct {
 	disableServerRoutes bool
 	activeRoutes        map[route.HAUniqueID]client.RouteHandler
 	fakeIPManager       *fakeip.Manager
-	dnsForwarderPort    atomic.Uint32
 }
 
 func NewManager(config ManagerConfig) *DefaultManager {
@@ -134,7 +131,6 @@ func NewManager(config ManagerConfig) *DefaultManager {
 		disableServerRoutes: config.DisableServerRoutes,
 		activeRoutes:        make(map[route.HAUniqueID]client.RouteHandler),
 	}
-	dm.dnsForwarderPort.Store(uint32(nbdns.ForwarderClientPort))
 
 	useNoop := netstack.IsEnabled() || config.DisableClientRoutes
 	dm.setupRefCounters(useNoop)
@@ -275,14 +271,10 @@ func (m *DefaultManager) SetFirewall(firewall firewall.Manager) error {
 	return nil
 }
 
-// SetDNSForwarderPort sets the DNS forwarder port for route handlers
-func (m *DefaultManager) SetDNSForwarderPort(port uint16) {
-	m.dnsForwarderPort.Store(uint32(port))
-}
-
 // Stop stops the manager watchers and clean firewall rules
 func (m *DefaultManager) Stop(stateManager *statemanager.Manager) {
 	m.stop()
+	m.shutdownWg.Wait()
 	if m.serverRouter != nil {
 		m.serverRouter.CleanUp()
 	}
@@ -355,7 +347,6 @@ func (m *DefaultManager) updateSystemRoutes(newRoutes route.HAMap) error {
 			UseNewDNSRoute:       m.useNewDNSRoute,
 			Firewall:             m.firewall,
 			FakeIPManager:        m.fakeIPManager,
-			ForwarderPort:        &m.dnsForwarderPort,
 		}
 		handler := client.HandlerFromRoute(params)
 		if err := handler.AddRoute(m.ctx); err != nil {
@@ -485,7 +476,11 @@ func (m *DefaultManager) TriggerSelection(networks route.HAMap) {
 		}
 		clientNetworkWatcher := client.NewWatcher(config)
 		m.clientNetworks[id] = clientNetworkWatcher
-		go clientNetworkWatcher.Start()
+		m.shutdownWg.Add(1)
+		go func() {
+			defer m.shutdownWg.Done()
+			clientNetworkWatcher.Start()
+		}()
 		clientNetworkWatcher.SendUpdate(client.RoutesUpdate{Routes: routes})
 	}
 
@@ -527,7 +522,11 @@ func (m *DefaultManager) updateClientNetworks(updateSerial uint64, networks rout
 			}
 			clientNetworkWatcher = client.NewWatcher(config)
 			m.clientNetworks[id] = clientNetworkWatcher
-			go clientNetworkWatcher.Start()
+			m.shutdownWg.Add(1)
+			go func() {
+				defer m.shutdownWg.Done()
+				clientNetworkWatcher.Start()
+			}()
 		}
 		update := client.RoutesUpdate{
 			UpdateSerial: updateSerial,

client/internal/routemanager/mock.go

@@ -90,10 +90,6 @@ func (m *MockManager) SetFirewall(firewall.Manager) error {
 	panic("implement me")
 }
 
-// SetDNSForwarderPort mock implementation of SetDNSForwarderPort from Manager interface
-func (m *MockManager) SetDNSForwarderPort(port uint16) {
-}
-
 // Stop mock implementation of Stop from Manager interface
 func (m *MockManager) Stop(stateManager *statemanager.Manager) {
 	if m.StopFunc != nil {

client/internal/routemanager/systemops/systemops_unix.go

@@ -105,11 +105,31 @@ func (r *SysOps) FlushMarkedRoutes() error {
 }
 
 func (r *SysOps) addToRouteTable(prefix netip.Prefix, nexthop Nexthop) error {
+	if prefix.IsSingleIP() {
+		log.Debugf("Adding single IP route: %s via %s", prefix, formatNexthop(nexthop))
+	}
 	return r.routeSocket(unix.RTM_ADD, prefix, nexthop)
 }
 
 func (r *SysOps) removeFromRouteTable(prefix netip.Prefix, nexthop Nexthop) error {
-	return r.routeSocket(unix.RTM_DELETE, prefix, nexthop)
+	if prefix.IsSingleIP() {
+		log.Debugf("Removing single IP route: %s via %s", prefix, formatNexthop(nexthop))
+	}
+
+	if err := r.routeSocket(unix.RTM_DELETE, prefix, nexthop); err != nil {
+		return err
+	}
+
+	if prefix.IsSingleIP() {
+		log.Debugf("Route removal completed for %s, verifying...", prefix)
+		if exists := r.verifyRouteRemoved(prefix); exists {
+			log.Warnf("Route %s still exists in routing table after removal", prefix)
+		} else {
+			log.Debugf("Verified route %s successfully removed", prefix)
+		}
+	}
+
+	return nil
 }
 
 func (r *SysOps) routeSocket(action int, prefix netip.Prefix, nexthop Nexthop) error {
@@ -276,3 +296,51 @@ func prefixToRouteNetmask(prefix netip.Prefix) (route.Addr, error) {
 
 	return nil, fmt.Errorf("unknown IP version in prefix: %s", prefix.Addr().String())
 }
+
+// formatNexthop returns a string representation of the nexthop for logging.
+func formatNexthop(nexthop Nexthop) string {
+	if nexthop.IP.IsValid() {
+		return nexthop.IP.String()
+	}
+	if nexthop.Intf != nil {
+		return nexthop.Intf.Name
+	}
+	return "direct"
+}
+
+// verifyRouteRemoved checks if a route still exists in the routing table.
+func (r *SysOps) verifyRouteRemoved(prefix netip.Prefix) bool {
+	rib, err := retryFetchRIB()
+	if err != nil {
+		log.Debugf("Failed to fetch RIB for route verification: %v", err)
+		return false
+	}
+
+	msgs, err := route.ParseRIB(route.RIBTypeRoute, rib)
+	if err != nil {
+		log.Debugf("Failed to parse RIB for route verification: %v", err)
+		return false
+	}
+
+	for _, msg := range msgs {
+		rtMsg, ok := msg.(*route.RouteMessage)
+		if !ok {
+			continue
+		}
+
+		if rtMsg.Flags&routeProtoFlag == 0 {
+			continue
+		}
+
+		routeInfo, err := MsgToRoute(rtMsg)
+		if err != nil {
+			continue
+		}
+
+		if routeInfo.Dst == prefix {
+			return true
+		}
+	}
+
+	return false
+}

client/server/server.go

@@ -1057,7 +1057,10 @@ func (s *Server) Status(
 	s.statusRecorder.UpdateRosenpass(s.config.RosenpassEnabled, s.config.RosenpassPermissive)
 
 	if msg.GetFullPeerStatus {
-		s.runProbes(msg.ShouldRunProbes)
+		if msg.ShouldRunProbes {
+			s.runProbes()
+		}
+
 		fullStatus := s.statusRecorder.GetFullStatus()
 		pbFullStatus := toProtoFullStatus(fullStatus)
 		pbFullStatus.Events = s.statusRecorder.GetEventHistory()
@@ -1067,7 +1070,7 @@ func (s *Server) Status(
 	return &statusResponse, nil
 }
 
-func (s *Server) runProbes(waitForProbeResult bool) {
+func (s *Server) runProbes() {
 	if s.connectClient == nil {
 		return
 	}
@@ -1078,7 +1081,7 @@ func (s *Server) runProbes(waitForProbeResult bool) {
 	}
 
 	if time.Since(s.lastProbe) > probeThreshold {
-		if engine.RunHealthProbes(waitForProbeResult) {
+		if engine.RunHealthProbes() {
 			s.lastProbe = time.Now()
 		}
 	}

client/status/status.go

@@ -15,7 +15,6 @@ import (
 
 	"github.com/netbirdio/netbird/client/anonymize"
 	"github.com/netbirdio/netbird/client/internal/peer"
-	probeRelay "github.com/netbirdio/netbird/client/internal/relay"
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/version"
@@ -341,16 +340,10 @@ func ParseGeneralSummary(overview OutputOverview, showURL bool, showRelays bool,
 		for _, relay := range overview.Relays.Details {
 			available := "Available"
 			reason := ""
-
 			if !relay.Available {
-				if relay.Error == probeRelay.ErrCheckInProgress.Error() {
-					available = "Checking..."
-				} else {
-					available = "Unavailable"
-					reason = fmt.Sprintf(", reason: %s", relay.Error)
-				}
+				available = "Unavailable"
+				reason = fmt.Sprintf(", reason: %s", relay.Error)
 			}
-
 			relaysString += fmt.Sprintf("\n  [%s] is %s%s", relay.URI, available, reason)
 		}
 	} else {

client/ui/client_ui.go

@@ -296,8 +296,6 @@ type serviceClient struct {
 	mExitNodeDeselectAll *systray.MenuItem
 	logFile              string
 	wLoginURL            fyne.Window
-
-	connectCancel context.CancelFunc
 }
 
 type menuHandler struct {
@@ -594,15 +592,17 @@ func (s *serviceClient) getSettingsForm() *widget.Form {
 	}
 }
 
-func (s *serviceClient) login(ctx context.Context, openURL bool) (*proto.LoginResponse, error) {
+func (s *serviceClient) login(openURL bool) (*proto.LoginResponse, error) {
 	conn, err := s.getSrvClient(defaultFailTimeout)
 	if err != nil {
-		return nil, fmt.Errorf("get daemon client: %w", err)
+		log.Errorf("get client: %v", err)
+		return nil, err
 	}
 
 	activeProf, err := s.profileManager.GetActiveProfile()
 	if err != nil {
-		return nil, fmt.Errorf("get active profile: %w", err)
+		log.Errorf("get active profile: %v", err)
+		return nil, err
 	}
 
 	currUser, err := user.Current()
@@ -610,71 +610,84 @@ func (s *serviceClient) login(ctx context.Context, openURL bool) (*proto.LoginRe
 		return nil, fmt.Errorf("get current user: %w", err)
 	}
 
-	loginResp, err := conn.Login(ctx, &proto.LoginRequest{
+	loginResp, err := conn.Login(s.ctx, &proto.LoginRequest{
 		IsUnixDesktopClient: runtime.GOOS == "linux" || runtime.GOOS == "freebsd",
 		ProfileName:         &activeProf.Name,
 		Username:            &currUser.Username,
 	})
 	if err != nil {
-		return nil, fmt.Errorf("login to management: %w", err)
+		log.Errorf("login to management URL with: %v", err)
+		return nil, err
 	}
 
 	if loginResp.NeedsSSOLogin && openURL {
-		if err = s.handleSSOLogin(ctx, loginResp, conn); err != nil {
-			return nil, fmt.Errorf("SSO login: %w", err)
+		err = s.handleSSOLogin(loginResp, conn)
+		if err != nil {
+			log.Errorf("handle SSO login failed: %v", err)
+			return nil, err
 		}
 	}
 
 	return loginResp, nil
 }
 
-func (s *serviceClient) handleSSOLogin(ctx context.Context, loginResp *proto.LoginResponse, conn proto.DaemonServiceClient) error {
-	if err := openURL(loginResp.VerificationURIComplete); err != nil {
-		return fmt.Errorf("open browser: %w", err)
+func (s *serviceClient) handleSSOLogin(loginResp *proto.LoginResponse, conn proto.DaemonServiceClient) error {
+	err := openURL(loginResp.VerificationURIComplete)
+	if err != nil {
+		log.Errorf("opening the verification uri in the browser failed: %v", err)
+		return err
 	}
 
-	resp, err := conn.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode})
+	resp, err := conn.WaitSSOLogin(s.ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode})
 	if err != nil {
-		return fmt.Errorf("wait for SSO login: %w", err)
+		log.Errorf("waiting sso login failed with: %v", err)
+		return err
 	}
 
 	if resp.Email != "" {
-		if err := s.profileManager.SetActiveProfileState(&profilemanager.ProfileState{
+		err := s.profileManager.SetActiveProfileState(&profilemanager.ProfileState{
 			Email: resp.Email,
-		}); err != nil {
-			log.Debugf("failed to set profile state: %v", err)
+		})
+		if err != nil {
+			log.Warnf("failed to set profile state: %v", err)
 		} else {
 			s.mProfile.refresh()
 		}
+
 	}
 
 	return nil
 }
 
-func (s *serviceClient) menuUpClick(ctx context.Context) error {
+func (s *serviceClient) menuUpClick() error {
 	systray.SetTemplateIcon(iconConnectingMacOS, s.icConnecting)
 	conn, err := s.getSrvClient(defaultFailTimeout)
 	if err != nil {
 		systray.SetTemplateIcon(iconErrorMacOS, s.icError)
-		return fmt.Errorf("get daemon client: %w", err)
+		log.Errorf("get client: %v", err)
+		return err
 	}
 
-	_, err = s.login(ctx, true)
+	_, err = s.login(true)
 	if err != nil {
-		return fmt.Errorf("login: %w", err)
+		log.Errorf("login failed with: %v", err)
+		return err
 	}
 
-	status, err := conn.Status(ctx, &proto.StatusRequest{})
+	status, err := conn.Status(s.ctx, &proto.StatusRequest{})
 	if err != nil {
-		return fmt.Errorf("get status: %w", err)
+		log.Errorf("get service status: %v", err)
+		return err
 	}
 
 	if status.Status == string(internal.StatusConnected) {
+		log.Warnf("already connected")
 		return nil
 	}
 
-	if _, err := conn.Up(ctx, &proto.UpRequest{}); err != nil {
-		return fmt.Errorf("start connection: %w", err)
+	if _, err := s.conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
+		log.Errorf("up service: %v", err)
+		return err
 	}
 
 	return nil
@@ -684,20 +697,24 @@ func (s *serviceClient) menuDownClick() error {
 	systray.SetTemplateIcon(iconConnectingMacOS, s.icConnecting)
 	conn, err := s.getSrvClient(defaultFailTimeout)
 	if err != nil {
-		return fmt.Errorf("get daemon client: %w", err)
+		log.Errorf("get client: %v", err)
+		return err
 	}
 
 	status, err := conn.Status(s.ctx, &proto.StatusRequest{})
 	if err != nil {
-		return fmt.Errorf("get status: %w", err)
+		log.Errorf("get service status: %v", err)
+		return err
 	}
 
 	if status.Status != string(internal.StatusConnected) && status.Status != string(internal.StatusConnecting) {
+		log.Warnf("already down")
 		return nil
 	}
 
-	if _, err := conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
-		return fmt.Errorf("stop connection: %w", err)
+	if _, err := s.conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
+		log.Errorf("down service: %v", err)
+		return err
 	}
 
 	return nil
@@ -833,7 +850,6 @@ func (s *serviceClient) onTrayReady() {
 
 	newProfileMenuArgs := &newProfileMenuArgs{
 		ctx:                  s.ctx,
-		serviceClient:        s,
 		profileManager:       s.profileManager,
 		eventHandler:         s.eventHandler,
 		profileMenuItem:      profileMenuItem,
@@ -1365,7 +1381,7 @@ func (s *serviceClient) showLoginURL() context.CancelFunc {
 			return
 		}
 
-		resp, err := s.login(ctx, false)
+		resp, err := s.login(false)
 		if err != nil {
 			log.Errorf("failed to fetch login URL: %v", err)
 			return
@@ -1385,7 +1401,7 @@ func (s *serviceClient) showLoginURL() context.CancelFunc {
 			return
 		}
 
-		_, err = conn.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: resp.UserCode})
+		_, err = conn.WaitSSOLogin(s.ctx, &proto.WaitSSOLoginRequest{UserCode: resp.UserCode})
 		if err != nil {
 			log.Errorf("Waiting sso login failed with: %v", err)
 			label.SetText("Waiting login failed, please create \na debug bundle in the settings and contact support.")
@@ -1393,7 +1409,7 @@ func (s *serviceClient) showLoginURL() context.CancelFunc {
 		}
 
 		label.SetText("Re-authentication successful.\nReconnecting")
-		status, err := conn.Status(ctx, &proto.StatusRequest{})
+		status, err := conn.Status(s.ctx, &proto.StatusRequest{})
 		if err != nil {
 			log.Errorf("get service status: %v", err)
 			return
@@ -1406,7 +1422,7 @@ func (s *serviceClient) showLoginURL() context.CancelFunc {
 			return
 		}
 
-		_, err = conn.Up(ctx, &proto.UpRequest{})
+		_, err = conn.Up(s.ctx, &proto.UpRequest{})
 		if err != nil {
 			label.SetText("Reconnecting failed, please create \na debug bundle in the settings and contact support.")
 			log.Errorf("Reconnecting failed with: %v", err)

client/ui/event_handler.go

@@ -12,8 +12,6 @@ import (
 	"fyne.io/fyne/v2"
 	"fyne.io/systray"
 	log "github.com/sirupsen/logrus"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/version"
@@ -69,55 +67,20 @@ func (h *eventHandler) listen(ctx context.Context) {
 
 func (h *eventHandler) handleConnectClick() {
 	h.client.mUp.Disable()
-
-	if h.client.connectCancel != nil {
-		h.client.connectCancel()
-	}
-
-	connectCtx, connectCancel := context.WithCancel(h.client.ctx)
-	h.client.connectCancel = connectCancel
-
 	go func() {
-		defer connectCancel()
-
-		if err := h.client.menuUpClick(connectCtx); err != nil {
-			st, ok := status.FromError(err)
-			if errors.Is(err, context.Canceled) || (ok && st.Code() == codes.Canceled) {
-				log.Debugf("connect operation cancelled by user")
-			} else {
-				h.client.app.SendNotification(fyne.NewNotification("Error", "Failed to connect"))
-				log.Errorf("connect failed: %v", err)
-			}
-		}
-
-		if err := h.client.updateStatus(); err != nil {
-			log.Debugf("failed to update status after connect: %v", err)
+		defer h.client.mUp.Enable()
+		if err := h.client.menuUpClick(); err != nil {
+			h.client.app.SendNotification(fyne.NewNotification("Error", "Failed to connect to NetBird service"))
 		}
 	}()
 }
 
 func (h *eventHandler) handleDisconnectClick() {
 	h.client.mDown.Disable()
-
-	if h.client.connectCancel != nil {
-		log.Debugf("cancelling ongoing connect operation")
-		h.client.connectCancel()
-		h.client.connectCancel = nil
-	}
-
 	go func() {
+		defer h.client.mDown.Enable()
 		if err := h.client.menuDownClick(); err != nil {
-			st, ok := status.FromError(err)
-			if !errors.Is(err, context.Canceled) && !(ok && st.Code() == codes.Canceled) {
-				h.client.app.SendNotification(fyne.NewNotification("Error", "Failed to disconnect"))
-				log.Errorf("disconnect failed: %v", err)
-			} else {
-				log.Debugf("disconnect cancelled or already disconnecting")
-			}
-		}
-
-		if err := h.client.updateStatus(); err != nil {
-			log.Debugf("failed to update status after disconnect: %v", err)
+			h.client.app.SendNotification(fyne.NewNotification("Error", "Failed to connect to NetBird daemon"))
 		}
 	}()
 }
@@ -282,6 +245,6 @@ func (h *eventHandler) logout(ctx context.Context) error {
 	}
 
 	h.client.getSrvConfig()
-
+	
 	return nil
 }

client/ui/profile.go

@@ -387,7 +387,6 @@ type subItem struct {
 type profileMenu struct {
 	mu                    sync.Mutex
 	ctx                   context.Context
-	serviceClient         *serviceClient
 	profileManager        *profilemanager.ProfileManager
 	eventHandler          *eventHandler
 	profileMenuItem       *systray.MenuItem
@@ -397,7 +396,7 @@ type profileMenu struct {
 	logoutSubItem         *subItem
 	profilesState         []Profile
 	downClickCallback     func() error
-	upClickCallback       func(context.Context) error
+	upClickCallback       func() error
 	getSrvClientCallback  func(timeout time.Duration) (proto.DaemonServiceClient, error)
 	loadSettingsCallback  func()
 	app                   fyne.App
@@ -405,13 +404,12 @@ type profileMenu struct {
 
 type newProfileMenuArgs struct {
 	ctx                  context.Context
-	serviceClient        *serviceClient
 	profileManager       *profilemanager.ProfileManager
 	eventHandler         *eventHandler
 	profileMenuItem      *systray.MenuItem
 	emailMenuItem        *systray.MenuItem
 	downClickCallback    func() error
-	upClickCallback      func(context.Context) error
+	upClickCallback      func() error
 	getSrvClientCallback func(timeout time.Duration) (proto.DaemonServiceClient, error)
 	loadSettingsCallback func()
 	app                  fyne.App
@@ -420,7 +418,6 @@ type newProfileMenuArgs struct {
 func newProfileMenu(args newProfileMenuArgs) *profileMenu {
 	p := profileMenu{
 		ctx:                  args.ctx,
-		serviceClient:        args.serviceClient,
 		profileManager:       args.profileManager,
 		eventHandler:         args.eventHandler,
 		profileMenuItem:      args.profileMenuItem,
@@ -572,19 +569,10 @@ func (p *profileMenu) refresh() {
 						}
 					}
 
-					if p.serviceClient.connectCancel != nil {
-						p.serviceClient.connectCancel()
-					}
-
-					connectCtx, connectCancel := context.WithCancel(p.ctx)
-					p.serviceClient.connectCancel = connectCancel
-
-					if err := p.upClickCallback(connectCtx); err != nil {
+					if err := p.upClickCallback(); err != nil {
 						log.Errorf("failed to handle up click after switching profile: %v", err)
 					}
 
-					connectCancel()
-
 					p.refresh()
 					p.loadSettingsCallback()
 				}

dns/dns.go

@@ -19,10 +19,6 @@ const (
 	RootZone = "."
 	// DefaultClass is the class supported by the system
 	DefaultClass = "IN"
-	// ForwarderClientPort is the port clients connect to. DNAT rewrites packets from ForwarderClientPort to ForwarderServerPort.
-	ForwarderClientPort uint16 = 5353
-	// ForwarderServerPort is the port the DNS forwarder actually listens on. Packets to ForwarderClientPort are DNATed here.
-	ForwarderServerPort uint16 = 22054
 )
 
 const invalidHostLabel = "[^a-zA-Z0-9-]+"
@@ -35,8 +31,6 @@ type Config struct {
 	NameServerGroups []*NameServerGroup
 	// CustomZones contains a list of custom zone
 	CustomZones []CustomZone
-	// ForwarderPort is the port clients should connect to on routing peers for DNS forwarding
-	ForwarderPort uint16
 }
 
 // CustomZone represents a custom zone to be resolved by the dns server

go.mod

@@ -56,7 +56,6 @@ require (
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/hashicorp/go-secure-stdlib/base62 v0.1.2
 	github.com/hashicorp/go-version v1.6.0
-	github.com/jackc/pgx/v5 v5.5.5
 	github.com/libdns/route53 v1.5.0
 	github.com/libp2p/go-netroute v0.2.1
 	github.com/mdlayher/socket v0.5.1
@@ -184,6 +183,7 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/pgx/v5 v5.5.5 // indirect
 	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/jeandeaual/go-locale v0.0.0-20240223122105-ce5225dcaa49 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect

management/server/dns.go

@@ -21,8 +21,8 @@ import (
 )
 
 const (
-	dnsForwarderPort = nbdns.ForwarderServerPort
-	oldForwarderPort = nbdns.ForwarderClientPort
+	dnsForwarderPort = 22054
+	oldForwarderPort = 5353
 )
 
 const dnsForwarderPortMinVersion = "v0.59.0"
@@ -196,7 +196,7 @@ func validateDNSSettings(ctx context.Context, transaction store.Store, accountID
 // If all peers have the required version, it returns the new well-known port (22054), otherwise returns 0.
 func computeForwarderPort(peers []*nbpeer.Peer, requiredVersion string) int64 {
 	if len(peers) == 0 {
-		return int64(oldForwarderPort)
+		return oldForwarderPort
 	}
 
 	reqVer := semver.Canonical(requiredVersion)
@@ -211,17 +211,17 @@ func computeForwarderPort(peers []*nbpeer.Peer, requiredVersion string) int64 {
 		peerVersion := semver.Canonical("v" + peer.Meta.WtVersion)
 		if peerVersion == "" {
 			// If any peer doesn't have version info, return 0
-			return int64(oldForwarderPort)
+			return oldForwarderPort
 		}
 
 		// Compare versions
 		if semver.Compare(peerVersion, reqVer) < 0 {
-			return int64(oldForwarderPort)
+			return oldForwarderPort
 		}
 	}
 
 	// All peers have the required version or newer
-	return int64(dnsForwarderPort)
+	return dnsForwarderPort
 }
 
 // toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache

management/server/dns_test.go

@@ -394,7 +394,7 @@ func BenchmarkToProtocolDNSConfig(b *testing.B) {
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				toProtocolDNSConfig(testData, cache, int64(dnsForwarderPort))
+				toProtocolDNSConfig(testData, cache, dnsForwarderPort)
 			}
 		})
 
@@ -402,7 +402,7 @@ func BenchmarkToProtocolDNSConfig(b *testing.B) {
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				cache := &DNSConfigCache{}
-				toProtocolDNSConfig(testData, cache, int64(dnsForwarderPort))
+				toProtocolDNSConfig(testData, cache, dnsForwarderPort)
 			}
 		})
 	}
@@ -455,13 +455,13 @@ func TestToProtocolDNSConfigWithCache(t *testing.T) {
 	}
 
 	// First run with config1
-	result1 := toProtocolDNSConfig(config1, &cache, int64(dnsForwarderPort))
+	result1 := toProtocolDNSConfig(config1, &cache, dnsForwarderPort)
 
 	// Second run with config2
-	result2 := toProtocolDNSConfig(config2, &cache, int64(dnsForwarderPort))
+	result2 := toProtocolDNSConfig(config2, &cache, dnsForwarderPort)
 
 	// Third run with config1 again
-	result3 := toProtocolDNSConfig(config1, &cache, int64(dnsForwarderPort))
+	result3 := toProtocolDNSConfig(config1, &cache, dnsForwarderPort)
 
 	// Verify that result1 and result3 are identical
 	if !reflect.DeepEqual(result1, result3) {
@@ -486,7 +486,7 @@ func TestComputeForwarderPort(t *testing.T) {
 	// Test with empty peers list
 	peers := []*nbpeer.Peer{}
 	result := computeForwarderPort(peers, "v0.59.0")
-	if result != int64(oldForwarderPort) {
+	if result != oldForwarderPort {
 		t.Errorf("Expected %d for empty peers list, got %d", oldForwarderPort, result)
 	}
 
@@ -504,7 +504,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result != int64(oldForwarderPort) {
+	if result != oldForwarderPort {
 		t.Errorf("Expected %d for peers with old versions, got %d", oldForwarderPort, result)
 	}
 
@@ -522,7 +522,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result != int64(dnsForwarderPort) {
+	if result != dnsForwarderPort {
 		t.Errorf("Expected %d for peers with new versions, got %d", dnsForwarderPort, result)
 	}
 
@@ -540,7 +540,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result != int64(oldForwarderPort) {
+	if result != oldForwarderPort {
 		t.Errorf("Expected %d for peers with mixed versions, got %d", oldForwarderPort, result)
 	}
 
@@ -553,7 +553,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result != int64(oldForwarderPort) {
+	if result != oldForwarderPort {
 		t.Errorf("Expected %d for peers with empty version, got %d", oldForwarderPort, result)
 	}
 
@@ -565,7 +565,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result == int64(oldForwarderPort) {
+	if result == oldForwarderPort {
 		t.Errorf("Expected %d for peers with dev version, got %d", dnsForwarderPort, result)
 	}
 
@@ -578,7 +578,7 @@ func TestComputeForwarderPort(t *testing.T) {
 		},
 	}
 	result = computeForwarderPort(peers, "v0.59.0")
-	if result != int64(oldForwarderPort) {
+	if result != oldForwarderPort {
 		t.Errorf("Expected %d for peers with unknown version, got %d", oldForwarderPort, result)
 	}
 }

management/server/peer_test.go

@@ -1161,7 +1161,7 @@ func TestToSyncResponse(t *testing.T) {
 	}
 	dnsCache := &DNSConfigCache{}
 	accountSettings := &types.Settings{RoutingPeerDNSResolutionEnabled: true}
-	response := toSyncResponse(context.Background(), config, peer, turnRelayToken, turnRelayToken, networkMap, dnsName, checks, dnsCache, accountSettings, nil, []string{}, int64(dnsForwarderPort))
+	response := toSyncResponse(context.Background(), config, peer, turnRelayToken, turnRelayToken, networkMap, dnsName, checks, dnsCache, accountSettings, nil, []string{}, dnsForwarderPort)
 
 	assert.NotNil(t, response)
 	// assert peer config

management/server/store/sqlstore_bench_test.go

@@ -1,951 +0,0 @@
-package store
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"net"
-	"net/netip"
-	"sort"
-	"sync"
-	"testing"
-	"time"
-
-	"gorm.io/driver/postgres"
-	"gorm.io/gorm"
-	"gorm.io/gorm/clause"
-
-	"github.com/jackc/pgx/v5/pgxpool"
-	log "github.com/sirupsen/logrus"
-	"github.com/stretchr/testify/assert"
-
-	nbdns "github.com/netbirdio/netbird/dns"
-	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
-	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
-	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
-	nbpeer "github.com/netbirdio/netbird/management/server/peer"
-	"github.com/netbirdio/netbird/management/server/posture"
-	"github.com/netbirdio/netbird/management/server/testutil"
-	"github.com/netbirdio/netbird/management/server/types"
-	"github.com/netbirdio/netbird/route"
-	"github.com/netbirdio/netbird/shared/management/status"
-)
-
-func (s *SqlStore) GetAccountSlow(ctx context.Context, accountID string) (*types.Account, error) {
-	start := time.Now()
-	defer func() {
-		elapsed := time.Since(start)
-		if elapsed > 1*time.Second {
-			log.WithContext(ctx).Tracef("GetAccount for account %s exceeded 1s, took: %v", accountID, elapsed)
-		}
-	}()
-
-	var account types.Account
-	result := s.db.Model(&account).
-		Omit("GroupsG").
-		Preload("UsersG.PATsG"). // have to be specified as this is nested reference
-		Preload(clause.Associations).
-		Take(&account, idQueryCondition, accountID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("error when getting account %s from the store: %s", accountID, result.Error)
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewAccountNotFoundError(accountID)
-		}
-		return nil, status.NewGetAccountFromStoreError(result.Error)
-	}
-
-	// we have to manually preload policy rules as it seems that gorm preloading doesn't do it for us
-	for i, policy := range account.Policies {
-		var rules []*types.PolicyRule
-		err := s.db.Model(&types.PolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
-		if err != nil {
-			return nil, status.Errorf(status.NotFound, "rule not found")
-		}
-		account.Policies[i].Rules = rules
-	}
-
-	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
-	for _, key := range account.SetupKeysG {
-		account.SetupKeys[key.Key] = key.Copy()
-	}
-	account.SetupKeysG = nil
-
-	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
-	for _, peer := range account.PeersG {
-		account.Peers[peer.ID] = peer.Copy()
-	}
-	account.PeersG = nil
-
-	account.Users = make(map[string]*types.User, len(account.UsersG))
-	for _, user := range account.UsersG {
-		user.PATs = make(map[string]*types.PersonalAccessToken, len(user.PATs))
-		for _, pat := range user.PATsG {
-			user.PATs[pat.ID] = pat.Copy()
-		}
-		account.Users[user.Id] = user.Copy()
-	}
-	account.UsersG = nil
-
-	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
-	for _, group := range account.GroupsG {
-		account.Groups[group.ID] = group.Copy()
-	}
-	account.GroupsG = nil
-
-	var groupPeers []types.GroupPeer
-	s.db.Model(&types.GroupPeer{}).Where("account_id = ?", accountID).
-		Find(&groupPeers)
-	for _, groupPeer := range groupPeers {
-		if group, ok := account.Groups[groupPeer.GroupID]; ok {
-			group.Peers = append(group.Peers, groupPeer.PeerID)
-		} else {
-			log.WithContext(ctx).Warnf("group %s not found for group peer %s in account %s", groupPeer.GroupID, groupPeer.PeerID, accountID)
-		}
-	}
-
-	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
-	for _, route := range account.RoutesG {
-		account.Routes[route.ID] = route.Copy()
-	}
-	account.RoutesG = nil
-
-	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
-	for _, ns := range account.NameServerGroupsG {
-		account.NameServerGroups[ns.ID] = ns.Copy()
-	}
-	account.NameServerGroupsG = nil
-
-	return &account, nil
-}
-
-func (s *SqlStore) GetAccountGormOpt(ctx context.Context, accountID string) (*types.Account, error) {
-	start := time.Now()
-	defer func() {
-		elapsed := time.Since(start)
-		if elapsed > 1*time.Second {
-			log.WithContext(ctx).Tracef("GetAccount for account %s exceeded 1s, took: %v", accountID, elapsed)
-		}
-	}()
-
-	var account types.Account
-	result := s.db.Model(&account).
-		Preload("UsersG.PATsG"). // have to be specified as this is nested reference
-		Preload("Policies.Rules").
-		Preload("SetupKeysG").
-		Preload("PeersG").
-		Preload("UsersG").
-		Preload("GroupsG.GroupPeers").
-		Preload("RoutesG").
-		Preload("NameServerGroupsG").
-		Preload("PostureChecks").
-		Preload("Networks").
-		Preload("NetworkRouters").
-		Preload("NetworkResources").
-		Preload("Onboarding").
-		Take(&account, idQueryCondition, accountID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("error when getting account %s from the store: %s", accountID, result.Error)
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewAccountNotFoundError(accountID)
-		}
-		return nil, status.NewGetAccountFromStoreError(result.Error)
-	}
-
-	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
-	for _, key := range account.SetupKeysG {
-		if key.UpdatedAt.IsZero() {
-			key.UpdatedAt = key.CreatedAt
-		}
-		if key.AutoGroups == nil {
-			key.AutoGroups = []string{}
-		}
-		account.SetupKeys[key.Key] = &key
-	}
-	account.SetupKeysG = nil
-
-	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
-	for _, peer := range account.PeersG {
-		account.Peers[peer.ID] = &peer
-	}
-	account.PeersG = nil
-	account.Users = make(map[string]*types.User, len(account.UsersG))
-	for _, user := range account.UsersG {
-		user.PATs = make(map[string]*types.PersonalAccessToken, len(user.PATs))
-		for _, pat := range user.PATsG {
-			pat.UserID = ""
-			user.PATs[pat.ID] = &pat
-		}
-		if user.AutoGroups == nil {
-			user.AutoGroups = []string{}
-		}
-		account.Users[user.Id] = &user
-		user.PATsG = nil
-	}
-	account.UsersG = nil
-	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
-	for _, group := range account.GroupsG {
-		group.Peers = make([]string, len(group.GroupPeers))
-		for i, gp := range group.GroupPeers {
-			group.Peers[i] = gp.PeerID
-		}
-		if group.Resources == nil {
-			group.Resources = []types.Resource{}
-		}
-		account.Groups[group.ID] = group
-	}
-	account.GroupsG = nil
-
-	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
-	for _, route := range account.RoutesG {
-		account.Routes[route.ID] = &route
-	}
-	account.RoutesG = nil
-	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
-	for _, ns := range account.NameServerGroupsG {
-		ns.AccountID = ""
-		if ns.NameServers == nil {
-			ns.NameServers = []nbdns.NameServer{}
-		}
-		if ns.Groups == nil {
-			ns.Groups = []string{}
-		}
-		if ns.Domains == nil {
-			ns.Domains = []string{}
-		}
-		account.NameServerGroups[ns.ID] = &ns
-	}
-	account.NameServerGroupsG = nil
-	return &account, nil
-}
-
-func connectDBforTest(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
-	config, err := pgxpool.ParseConfig(dsn)
-	if err != nil {
-		return nil, fmt.Errorf("unable to parse database config: %w", err)
-	}
-
-	config.MaxConns = 12
-	config.MinConns = 2
-	config.MaxConnLifetime = time.Hour
-	config.HealthCheckPeriod = time.Minute
-
-	pool, err := pgxpool.NewWithConfig(ctx, config)
-	if err != nil {
-		return nil, fmt.Errorf("unable to create connection pool: %w", err)
-	}
-
-	if err := pool.Ping(ctx); err != nil {
-		pool.Close()
-		return nil, fmt.Errorf("unable to ping database: %w", err)
-	}
-	return pool, nil
-}
-
-func setupBenchmarkDB(b testing.TB) (*SqlStore, func(), string) {
-	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
-	if err != nil {
-		b.Fatalf("failed to create test container: %v", err)
-	}
-
-	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
-	if err != nil {
-		b.Fatalf("failed to connect database: %v", err)
-	}
-
-	pool, err := connectDBforTest(context.Background(), dsn)
-	if err != nil {
-		b.Fatalf("failed to connect database: %v", err)
-	}
-
-	models := []interface{}{
-		&types.Account{}, &types.SetupKey{}, &nbpeer.Peer{}, &types.User{},
-		&types.PersonalAccessToken{}, &types.Group{}, &types.GroupPeer{},
-		&types.Policy{}, &types.PolicyRule{}, &route.Route{},
-		&nbdns.NameServerGroup{}, &posture.Checks{}, &networkTypes.Network{},
-		&routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{},
-		&types.AccountOnboarding{},
-	}
-
-	for i := len(models) - 1; i >= 0; i-- {
-		err := db.Migrator().DropTable(models[i])
-		if err != nil {
-			b.Fatalf("failed to drop table: %v", err)
-		}
-	}
-
-	err = db.AutoMigrate(models...)
-	if err != nil {
-		b.Fatalf("failed to migrate database: %v", err)
-	}
-
-	store := &SqlStore{
-		db:   db,
-		pool: pool,
-	}
-
-	const (
-		accountID           = "benchmark-account-id"
-		numUsers            = 20
-		numPatsPerUser      = 3
-		numSetupKeys        = 25
-		numPeers            = 200
-		numGroups           = 30
-		numPolicies         = 50
-		numRulesPerPolicy   = 10
-		numRoutes           = 40
-		numNSGroups         = 10
-		numPostureChecks    = 15
-		numNetworks         = 5
-		numNetworkRouters   = 5
-		numNetworkResources = 10
-	)
-
-	_, ipNet, _ := net.ParseCIDR("100.64.0.0/10")
-	acc := types.Account{
-		Id:                     accountID,
-		CreatedBy:              "benchmark-user",
-		CreatedAt:              time.Now(),
-		Domain:                 "benchmark.com",
-		IsDomainPrimaryAccount: true,
-		Network: &types.Network{
-			Identifier: "benchmark-net",
-			Net:        *ipNet,
-			Serial:     1,
-		},
-		DNSSettings: types.DNSSettings{
-			DisabledManagementGroups: []string{"group-disabled-1"},
-		},
-		Settings: &types.Settings{},
-	}
-	if err := db.Create(&acc).Error; err != nil {
-		b.Fatalf("create account: %v", err)
-	}
-
-	var setupKeys []types.SetupKey
-	for i := 0; i < numSetupKeys; i++ {
-		setupKeys = append(setupKeys, types.SetupKey{
-			Id:        fmt.Sprintf("keyid-%d", i),
-			AccountID: accountID,
-			Key:       fmt.Sprintf("key-%d", i),
-			Name:      fmt.Sprintf("Benchmark Key %d", i),
-			ExpiresAt: &time.Time{},
-		})
-	}
-	if err := db.Create(&setupKeys).Error; err != nil {
-		b.Fatalf("create setup keys: %v", err)
-	}
-
-	var peers []nbpeer.Peer
-	for i := 0; i < numPeers; i++ {
-		peers = append(peers, nbpeer.Peer{
-			ID:        fmt.Sprintf("peer-%d", i),
-			AccountID: accountID,
-			Key:       fmt.Sprintf("peerkey-%d", i),
-			IP:        net.ParseIP(fmt.Sprintf("100.64.0.%d", i+1)),
-			Name:      fmt.Sprintf("peer-name-%d", i),
-			Status:    &nbpeer.PeerStatus{Connected: i%2 == 0, LastSeen: time.Now()},
-		})
-	}
-	if err := db.Create(&peers).Error; err != nil {
-		b.Fatalf("create peers: %v", err)
-	}
-
-	for i := 0; i < numUsers; i++ {
-		userID := fmt.Sprintf("user-%d", i)
-		user := types.User{Id: userID, AccountID: accountID}
-		if err := db.Create(&user).Error; err != nil {
-			b.Fatalf("create user %s: %v", userID, err)
-		}
-
-		var pats []types.PersonalAccessToken
-		for j := 0; j < numPatsPerUser; j++ {
-			pats = append(pats, types.PersonalAccessToken{
-				ID:     fmt.Sprintf("pat-%d-%d", i, j),
-				UserID: userID,
-				Name:   fmt.Sprintf("PAT %d for User %d", j, i),
-			})
-		}
-		if err := db.Create(&pats).Error; err != nil {
-			b.Fatalf("create pats for user %s: %v", userID, err)
-		}
-	}
-
-	var groups []*types.Group
-	for i := 0; i < numGroups; i++ {
-		groups = append(groups, &types.Group{
-			ID:        fmt.Sprintf("group-%d", i),
-			AccountID: accountID,
-			Name:      fmt.Sprintf("Group %d", i),
-		})
-	}
-	if err := db.Create(&groups).Error; err != nil {
-		b.Fatalf("create groups: %v", err)
-	}
-
-	for i := 0; i < numPolicies; i++ {
-		policyID := fmt.Sprintf("policy-%d", i)
-		policy := types.Policy{ID: policyID, AccountID: accountID, Name: fmt.Sprintf("Policy %d", i), Enabled: true}
-		if err := db.Create(&policy).Error; err != nil {
-			b.Fatalf("create policy %s: %v", policyID, err)
-		}
-
-		var rules []*types.PolicyRule
-		for j := 0; j < numRulesPerPolicy; j++ {
-			rules = append(rules, &types.PolicyRule{
-				ID:       fmt.Sprintf("rule-%d-%d", i, j),
-				PolicyID: policyID,
-				Name:     fmt.Sprintf("Rule %d for Policy %d", j, i),
-				Enabled:  true,
-				Protocol: "all",
-			})
-		}
-		if err := db.Create(&rules).Error; err != nil {
-			b.Fatalf("create rules for policy %s: %v", policyID, err)
-		}
-	}
-
-	var routes []route.Route
-	for i := 0; i < numRoutes; i++ {
-		routes = append(routes, route.Route{
-			ID:          route.ID(fmt.Sprintf("route-%d", i)),
-			AccountID:   accountID,
-			Description: fmt.Sprintf("Route %d", i),
-			Network:     netip.MustParsePrefix(fmt.Sprintf("192.168.%d.0/24", i)),
-			Enabled:     true,
-		})
-	}
-	if err := db.Create(&routes).Error; err != nil {
-		b.Fatalf("create routes: %v", err)
-	}
-
-	var nsGroups []nbdns.NameServerGroup
-	for i := 0; i < numNSGroups; i++ {
-		nsGroups = append(nsGroups, nbdns.NameServerGroup{
-			ID:          fmt.Sprintf("nsg-%d", i),
-			AccountID:   accountID,
-			Name:        fmt.Sprintf("NS Group %d", i),
-			Description: "Benchmark NS Group",
-			Enabled:     true,
-		})
-	}
-	if err := db.Create(&nsGroups).Error; err != nil {
-		b.Fatalf("create nsgroups: %v", err)
-	}
-
-	var postureChecks []*posture.Checks
-	for i := 0; i < numPostureChecks; i++ {
-		postureChecks = append(postureChecks, &posture.Checks{
-			ID:        fmt.Sprintf("pc-%d", i),
-			AccountID: accountID,
-			Name:      fmt.Sprintf("Posture Check %d", i),
-		})
-	}
-	if err := db.Create(&postureChecks).Error; err != nil {
-		b.Fatalf("create posture checks: %v", err)
-	}
-
-	var networks []*networkTypes.Network
-	for i := 0; i < numNetworks; i++ {
-		networks = append(networks, &networkTypes.Network{
-			ID:        fmt.Sprintf("nettype-%d", i),
-			AccountID: accountID,
-			Name:      fmt.Sprintf("Network Type %d", i),
-		})
-	}
-	if err := db.Create(&networks).Error; err != nil {
-		b.Fatalf("create networks: %v", err)
-	}
-
-	var networkRouters []*routerTypes.NetworkRouter
-	for i := 0; i < numNetworkRouters; i++ {
-		networkRouters = append(networkRouters, &routerTypes.NetworkRouter{
-			ID:        fmt.Sprintf("router-%d", i),
-			AccountID: accountID,
-			NetworkID: networks[i%numNetworks].ID,
-			Peer:      peers[i%numPeers].ID,
-		})
-	}
-	if err := db.Create(&networkRouters).Error; err != nil {
-		b.Fatalf("create network routers: %v", err)
-	}
-
-	var networkResources []*resourceTypes.NetworkResource
-	for i := 0; i < numNetworkResources; i++ {
-		networkResources = append(networkResources, &resourceTypes.NetworkResource{
-			ID:        fmt.Sprintf("resource-%d", i),
-			AccountID: accountID,
-			NetworkID: networks[i%numNetworks].ID,
-			Name:      fmt.Sprintf("Resource %d", i),
-		})
-	}
-	if err := db.Create(&networkResources).Error; err != nil {
-		b.Fatalf("create network resources: %v", err)
-	}
-
-	onboarding := types.AccountOnboarding{
-		AccountID:             accountID,
-		OnboardingFlowPending: true,
-	}
-	if err := db.Create(&onboarding).Error; err != nil {
-		b.Fatalf("create onboarding: %v", err)
-	}
-
-	return store, cleanup, accountID
-}
-
-func BenchmarkGetAccount(b *testing.B) {
-	store, cleanup, accountID := setupBenchmarkDB(b)
-	defer cleanup()
-	ctx := context.Background()
-	b.ResetTimer()
-	b.ReportAllocs()
-	b.Run("old", func(b *testing.B) {
-		for range b.N {
-			_, err := store.GetAccountSlow(ctx, accountID)
-			if err != nil {
-				b.Fatalf("GetAccountSlow failed: %v", err)
-			}
-		}
-	})
-	b.Run("gorm opt", func(b *testing.B) {
-		for range b.N {
-			_, err := store.GetAccountGormOpt(ctx, accountID)
-			if err != nil {
-				b.Fatalf("GetAccountFast failed: %v", err)
-			}
-		}
-	})
-	b.Run("raw", func(b *testing.B) {
-		for range b.N {
-			_, err := store.GetAccount(ctx, accountID)
-			if err != nil {
-				b.Fatalf("GetAccountPureSQL failed: %v", err)
-			}
-		}
-	})
-	store.pool.Close()
-}
-
-func TestAccountEquivalence(t *testing.T) {
-	store, cleanup, accountID := setupBenchmarkDB(t)
-	defer cleanup()
-	ctx := context.Background()
-
-	type getAccountFunc func(context.Context, string) (*types.Account, error)
-
-	tests := []struct {
-		name      string
-		expectedF getAccountFunc
-		actualF   getAccountFunc
-	}{
-		{"old vs new", store.GetAccountSlow, store.GetAccountGormOpt},
-		{"old vs raw", store.GetAccountSlow, store.GetAccount},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			expected, errOld := tt.expectedF(ctx, accountID)
-			assert.NoError(t, errOld, "expected function should not return an error")
-			assert.NotNil(t, expected, "expected should not be nil")
-
-			actual, errNew := tt.actualF(ctx, accountID)
-			assert.NoError(t, errNew, "actual function should not return an error")
-			assert.NotNil(t, actual, "actual should not be nil")
-			testAccountEquivalence(t, expected, actual)
-		})
-	}
-
-	expected, errOld := store.GetAccountSlow(ctx, accountID)
-	assert.NoError(t, errOld, "GetAccountSlow should not return an error")
-	assert.NotNil(t, expected, "expected should not be nil")
-
-	actual, errNew := store.GetAccount(ctx, accountID)
-	assert.NoError(t, errNew, "GetAccount (new) should not return an error")
-	assert.NotNil(t, actual, "actual should not be nil")
-}
-
-func testAccountEquivalence(t *testing.T, expected, actual *types.Account) {
-	assert.Equal(t, expected.Id, actual.Id, "Account IDs should be equal")
-	assert.Equal(t, expected.CreatedBy, actual.CreatedBy, "Account CreatedBy fields should be equal")
-	assert.WithinDuration(t, expected.CreatedAt, actual.CreatedAt, time.Second, "Account CreatedAt timestamps should be within a second")
-	assert.Equal(t, expected.Domain, actual.Domain, "Account Domains should be equal")
-	assert.Equal(t, expected.DomainCategory, actual.DomainCategory, "Account DomainCategories should be equal")
-	assert.Equal(t, expected.IsDomainPrimaryAccount, actual.IsDomainPrimaryAccount, "Account IsDomainPrimaryAccount flags should be equal")
-	assert.Equal(t, expected.Network, actual.Network, "Embedded Account Network structs should be equal")
-	assert.Equal(t, expected.DNSSettings, actual.DNSSettings, "Embedded Account DNSSettings structs should be equal")
-	assert.Equal(t, expected.Onboarding, actual.Onboarding, "Embedded Account Onboarding structs should be equal")
-
-	assert.Len(t, actual.SetupKeys, len(expected.SetupKeys), "SetupKeys maps should have the same number of elements")
-	for key, oldVal := range expected.SetupKeys {
-		newVal, ok := actual.SetupKeys[key]
-		assert.True(t, ok, "SetupKey with key '%s' should exist in new account", key)
-		assert.Equal(t, *oldVal, *newVal, "SetupKey with key '%s' should be equal", key)
-	}
-
-	assert.Len(t, actual.Peers, len(expected.Peers), "Peers maps should have the same number of elements")
-	for key, oldVal := range expected.Peers {
-		newVal, ok := actual.Peers[key]
-		assert.True(t, ok, "Peer with ID '%s' should exist in new account", key)
-		assert.Equal(t, *oldVal, *newVal, "Peer with ID '%s' should be equal", key)
-	}
-
-	assert.Len(t, actual.Users, len(expected.Users), "Users maps should have the same number of elements")
-	for key, oldUser := range expected.Users {
-		newUser, ok := actual.Users[key]
-		assert.True(t, ok, "User with ID '%s' should exist in new account", key)
-
-		assert.Len(t, newUser.PATs, len(oldUser.PATs), "PATs map for user '%s' should have the same size", key)
-		for patKey, oldPAT := range oldUser.PATs {
-			newPAT, patOk := newUser.PATs[patKey]
-			assert.True(t, patOk, "PAT with ID '%s' for user '%s' should exist in new user object", patKey, key)
-			assert.Equal(t, *oldPAT, *newPAT, "PAT with ID '%s' for user '%s' should be equal", patKey, key)
-		}
-
-		oldUser.PATs = nil
-		newUser.PATs = nil
-		assert.Equal(t, *oldUser, *newUser, "User struct for ID '%s' (without PATs) should be equal", key)
-	}
-
-	assert.Len(t, actual.Groups, len(expected.Groups), "Groups maps should have the same number of elements")
-	for key, oldVal := range expected.Groups {
-		newVal, ok := actual.Groups[key]
-		assert.True(t, ok, "Group with ID '%s' should exist in new account", key)
-		sort.Strings(oldVal.Peers)
-		sort.Strings(newVal.Peers)
-		assert.Equal(t, *oldVal, *newVal, "Group with ID '%s' should be equal", key)
-	}
-
-	assert.Len(t, actual.Routes, len(expected.Routes), "Routes maps should have the same number of elements")
-	for key, oldVal := range expected.Routes {
-		newVal, ok := actual.Routes[key]
-		assert.True(t, ok, "Route with ID '%s' should exist in new account", key)
-		assert.Equal(t, *oldVal, *newVal, "Route with ID '%s' should be equal", key)
-	}
-
-	assert.Len(t, actual.NameServerGroups, len(expected.NameServerGroups), "NameServerGroups maps should have the same number of elements")
-	for key, oldVal := range expected.NameServerGroups {
-		newVal, ok := actual.NameServerGroups[key]
-		assert.True(t, ok, "NameServerGroup with ID '%s' should exist in new account", key)
-		assert.Equal(t, *oldVal, *newVal, "NameServerGroup with ID '%s' should be equal", key)
-	}
-
-	assert.Len(t, actual.Policies, len(expected.Policies), "Policies slices should have the same number of elements")
-	sort.Slice(expected.Policies, func(i, j int) bool { return expected.Policies[i].ID < expected.Policies[j].ID })
-	sort.Slice(actual.Policies, func(i, j int) bool { return actual.Policies[i].ID < actual.Policies[j].ID })
-	for i := range expected.Policies {
-		sort.Slice(expected.Policies[i].Rules, func(j, k int) bool { return expected.Policies[i].Rules[j].ID < expected.Policies[i].Rules[k].ID })
-		sort.Slice(actual.Policies[i].Rules, func(j, k int) bool { return actual.Policies[i].Rules[j].ID < actual.Policies[i].Rules[k].ID })
-		assert.Equal(t, *expected.Policies[i], *actual.Policies[i], "Policy with ID '%s' should be equal", expected.Policies[i].ID)
-	}
-
-	assert.Len(t, actual.PostureChecks, len(expected.PostureChecks), "PostureChecks slices should have the same number of elements")
-	sort.Slice(expected.PostureChecks, func(i, j int) bool { return expected.PostureChecks[i].ID < expected.PostureChecks[j].ID })
-	sort.Slice(actual.PostureChecks, func(i, j int) bool { return actual.PostureChecks[i].ID < actual.PostureChecks[j].ID })
-	for i := range expected.PostureChecks {
-		assert.Equal(t, *expected.PostureChecks[i], *actual.PostureChecks[i], "PostureCheck with ID '%s' should be equal", expected.PostureChecks[i].ID)
-	}
-
-	assert.Len(t, actual.Networks, len(expected.Networks), "Networks slices should have the same number of elements")
-	sort.Slice(expected.Networks, func(i, j int) bool { return expected.Networks[i].ID < expected.Networks[j].ID })
-	sort.Slice(actual.Networks, func(i, j int) bool { return actual.Networks[i].ID < actual.Networks[j].ID })
-	for i := range expected.Networks {
-		assert.Equal(t, *expected.Networks[i], *actual.Networks[i], "Network with ID '%s' should be equal", expected.Networks[i].ID)
-	}
-
-	assert.Len(t, actual.NetworkRouters, len(expected.NetworkRouters), "NetworkRouters slices should have the same number of elements")
-	sort.Slice(expected.NetworkRouters, func(i, j int) bool { return expected.NetworkRouters[i].ID < expected.NetworkRouters[j].ID })
-	sort.Slice(actual.NetworkRouters, func(i, j int) bool { return actual.NetworkRouters[i].ID < actual.NetworkRouters[j].ID })
-	for i := range expected.NetworkRouters {
-		assert.Equal(t, *expected.NetworkRouters[i], *actual.NetworkRouters[i], "NetworkRouter with ID '%s' should be equal", expected.NetworkRouters[i].ID)
-	}
-
-	assert.Len(t, actual.NetworkResources, len(expected.NetworkResources), "NetworkResources slices should have the same number of elements")
-	sort.Slice(expected.NetworkResources, func(i, j int) bool { return expected.NetworkResources[i].ID < expected.NetworkResources[j].ID })
-	sort.Slice(actual.NetworkResources, func(i, j int) bool { return actual.NetworkResources[i].ID < actual.NetworkResources[j].ID })
-	for i := range expected.NetworkResources {
-		assert.Equal(t, *expected.NetworkResources[i], *actual.NetworkResources[i], "NetworkResource with ID '%s' should be equal", expected.NetworkResources[i].ID)
-	}
-}
-
-func (s *SqlStore) GetAccountPureSQL(ctx context.Context, accountID string) (*types.Account, error) {
-	account, err := s.getAccount(ctx, accountID)
-	if err != nil {
-		return nil, err
-	}
-
-	var wg sync.WaitGroup
-	errChan := make(chan error, 12)
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		keys, err := s.getSetupKeys(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.SetupKeysG = keys
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		peers, err := s.getPeers(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.PeersG = peers
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		users, err := s.getUsers(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.UsersG = users
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		groups, err := s.getGroups(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.GroupsG = groups
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		policies, err := s.getPolicies(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.Policies = policies
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		routes, err := s.getRoutes(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.RoutesG = routes
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		nsgs, err := s.getNameServerGroups(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.NameServerGroupsG = nsgs
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		checks, err := s.getPostureChecks(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.PostureChecks = checks
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		networks, err := s.getNetworks(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.Networks = networks
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		routers, err := s.getNetworkRouters(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.NetworkRouters = routers
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		resources, err := s.getNetworkResources(ctx, accountID)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		account.NetworkResources = resources
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		err := s.getAccountOnboarding(ctx, accountID, account)
-		if err != nil {
-			errChan <- err
-			return
-		}
-	}()
-
-	wg.Wait()
-	close(errChan)
-	for e := range errChan {
-		if e != nil {
-			return nil, e
-		}
-	}
-
-	var userIDs []string
-	for _, u := range account.UsersG {
-		userIDs = append(userIDs, u.Id)
-	}
-	var policyIDs []string
-	for _, p := range account.Policies {
-		policyIDs = append(policyIDs, p.ID)
-	}
-	var groupIDs []string
-	for _, g := range account.GroupsG {
-		groupIDs = append(groupIDs, g.ID)
-	}
-
-	wg.Add(3)
-	errChan = make(chan error, 3)
-
-	var pats []types.PersonalAccessToken
-	go func() {
-		defer wg.Done()
-		var err error
-		pats, err = s.getPersonalAccessTokens(ctx, userIDs)
-		if err != nil {
-			errChan <- err
-		}
-	}()
-
-	var rules []*types.PolicyRule
-	go func() {
-		defer wg.Done()
-		var err error
-		rules, err = s.getPolicyRules(ctx, policyIDs)
-		if err != nil {
-			errChan <- err
-		}
-	}()
-
-	var groupPeers []types.GroupPeer
-	go func() {
-		defer wg.Done()
-		var err error
-		groupPeers, err = s.getGroupPeers(ctx, groupIDs)
-		if err != nil {
-			errChan <- err
-		}
-	}()
-
-	wg.Wait()
-	close(errChan)
-	for e := range errChan {
-		if e != nil {
-			return nil, e
-		}
-	}
-
-	patsByUserID := make(map[string][]*types.PersonalAccessToken)
-	for i := range pats {
-		pat := &pats[i]
-		patsByUserID[pat.UserID] = append(patsByUserID[pat.UserID], pat)
-		pat.UserID = ""
-	}
-
-	rulesByPolicyID := make(map[string][]*types.PolicyRule)
-	for _, rule := range rules {
-		rulesByPolicyID[rule.PolicyID] = append(rulesByPolicyID[rule.PolicyID], rule)
-	}
-
-	peersByGroupID := make(map[string][]string)
-	for _, gp := range groupPeers {
-		peersByGroupID[gp.GroupID] = append(peersByGroupID[gp.GroupID], gp.PeerID)
-	}
-
-	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
-	for i := range account.SetupKeysG {
-		key := &account.SetupKeysG[i]
-		account.SetupKeys[key.Key] = key
-	}
-
-	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
-	for i := range account.PeersG {
-		peer := &account.PeersG[i]
-		account.Peers[peer.ID] = peer
-	}
-
-	account.Users = make(map[string]*types.User, len(account.UsersG))
-	for i := range account.UsersG {
-		user := &account.UsersG[i]
-		user.PATs = make(map[string]*types.PersonalAccessToken)
-		if userPats, ok := patsByUserID[user.Id]; ok {
-			for j := range userPats {
-				pat := userPats[j]
-				user.PATs[pat.ID] = pat
-			}
-		}
-		account.Users[user.Id] = user
-	}
-
-	for i := range account.Policies {
-		policy := account.Policies[i]
-		if policyRules, ok := rulesByPolicyID[policy.ID]; ok {
-			policy.Rules = policyRules
-		}
-	}
-
-	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
-	for i := range account.GroupsG {
-		group := account.GroupsG[i]
-		if peerIDs, ok := peersByGroupID[group.ID]; ok {
-			group.Peers = peerIDs
-		}
-		account.Groups[group.ID] = group
-	}
-
-	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
-	for i := range account.RoutesG {
-		route := &account.RoutesG[i]
-		account.Routes[route.ID] = route
-	}
-
-	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
-	for i := range account.NameServerGroupsG {
-		nsg := &account.NameServerGroupsG[i]
-		nsg.AccountID = ""
-		account.NameServerGroups[nsg.ID] = nsg
-	}
-
-	account.SetupKeysG = nil
-	account.PeersG = nil
-	account.UsersG = nil
-	account.GroupsG = nil
-	account.RoutesG = nil
-	account.NameServerGroupsG = nil
-
-	return account, nil
-}

management/server/store/store.go

@@ -468,9 +468,6 @@ func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine)
 	closeConnection := func() {
 		cleanup()
 		store.Close(ctx)
-		if store.pool != nil {
-			store.pool.Close()
-		}
 	}
 
 	return store, closeConnection, nil

route/route.go

@@ -124,7 +124,6 @@ func (r *Route) EventMeta() map[string]any {
 func (r *Route) Copy() *Route {
 	route := &Route{
 		ID:                  r.ID,
-		AccountID:           r.AccountID,
 		Description:         r.Description,
 		NetID:               r.NetID,
 		Network:             r.Network,

shared/management/proto/management.proto

@@ -410,7 +410,7 @@ message DNSConfig {
   bool ServiceEnable = 1;
   repeated NameServerGroup NameServerGroups = 2;
   repeated CustomZone CustomZones = 3;
-  int64 ForwarderPort = 4 [deprecated = true];
+  int64 ForwarderPort = 4;
 }
 
 // CustomZone represents a dns.CustomZone