[management] Remove context from database calls (#2863)

Improve the performance on Linux and Android in case of P2P connections

.github/FUNDING.yml

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

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

@@ -13,6 +13,7 @@ concurrency:
 jobs:
   test:
     strategy:
+      fail-fast: false
       matrix:
         arch: [ '386','amd64' ]
         store: [ 'sqlite', 'postgres']

client/anonymize/anonymize.go

@@ -201,6 +201,8 @@ func isWellKnown(addr netip.Addr) bool {
 		"2606:4700:4700::1111", "2606:4700:4700::1001", // Cloudflare DNS IPv6
 		"9.9.9.9", "149.112.112.112", // Quad9 DNS IPv4
 		"2620:fe::fe", "2620:fe::9", // Quad9 DNS IPv6
+
+		"128.0.0.0", "8000::", // 2nd split subnet for default routes
 	}
 
 	if slices.Contains(wellKnown, addr.String()) {

client/cmd/status.go

@@ -680,7 +680,7 @@ func parsePeers(peers peersStateOutput, rosenpassEnabled, rosenpassPermissive bo
 func skipDetailByFilters(peerState *proto.PeerState, isConnected bool) bool {
 	statusEval := false
 	ipEval := false
-	nameEval := false
+	nameEval := true
 
 	if statusFilter != "" {
 		lowerStatusFilter := strings.ToLower(statusFilter)
@@ -700,11 +700,13 @@ func skipDetailByFilters(peerState *proto.PeerState, isConnected bool) bool {
 
 	if len(prefixNamesFilter) > 0 {
 		for prefixNameFilter := range prefixNamesFilterMap {
-			if !strings.HasPrefix(peerState.Fqdn, prefixNameFilter) {
-				nameEval = true
+			if strings.HasPrefix(peerState.Fqdn, prefixNameFilter) {
+				nameEval = false
 				break
 			}
 		}
+	} else {
+		nameEval = false
 	}
 
 	return statusEval || ipEval || nameEval

client/firewall/iptables/acl_linux.go

@@ -352,14 +352,14 @@ func (m *aclManager) seedInitialEntries() {
 func (m *aclManager) seedInitialOptionalEntries() {
 	m.optionalEntries["FORWARD"] = []entry{
 		{
-			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark), "-j", chainNameInputRules},
+			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected), "-j", chainNameInputRules},
 			position: 2,
 		},
 	}
 
 	m.optionalEntries["PREROUTING"] = []entry{
 		{
-			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark)},
+			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected)},
 			position: 1,
 		},
 	}

client/firewall/iptables/router_linux.go

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

client/firewall/iptables/router_linux_test.go

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

client/firewall/manager/firewall.go

@@ -17,6 +17,7 @@ import (
 const (
 	ForwardingFormatPrefix = "netbird-fwd-"
 	ForwardingFormat       = "netbird-fwd-%s-%t"
+	PreroutingFormat       = "netbird-prerouting-%s-%t"
 	NatFormat              = "netbird-nat-%s-%t"
 )
 

client/firewall/nftables/acl_linux.go

@@ -520,7 +520,7 @@ func (m *AclManager) addPreroutingRule(preroutingChain *nftables.Chain) {
 			},
 			&expr.Immediate{
 				Register: 1,
-				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
 			},
 			&expr.Meta{
 				Key:            expr.MetaKeyMARK,
@@ -543,7 +543,7 @@ func (m *AclManager) addFwmarkToForward(chainFwFilter *nftables.Chain) {
 			&expr.Cmp{
 				Op:       expr.CmpOpEq,
 				Register: 1,
-				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmarkRedirected),
 			},
 			&expr.Verdict{
 				Kind:  expr.VerdictJump,

client/firewall/nftables/router_linux.go

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

client/firewall/nftables/router_linux_test.go

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

client/iface/bind/ice_bind.go

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

client/iface/wgproxy/bind/proxy.go

@@ -104,8 +104,8 @@ func (p *ProxyBind) proxyToLocal(ctx context.Context) {
 		}
 	}()
 
-	buf := make([]byte, 1500)
 	for {
+		buf := make([]byte, 1500)
 		n, err := p.remoteConn.Read(buf)
 		if err != nil {
 			if ctx.Err() != nil {

client/internal/peer/conn.go

@@ -309,6 +309,11 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		return
 	}
 
+	if remoteConnNil(conn.log, iceConnInfo.RemoteConn) {
+		conn.log.Errorf("remote ICE connection is nil")
+		return
+	}
+
 	conn.log.Debugf("ICE connection is ready")
 
 	if conn.currentConnPriority > priority {

client/internal/peer/nilcheck.go

@@ -0,0 +1,21 @@
+package peer
+
+import (
+	"net"
+
+	log "github.com/sirupsen/logrus"
+)
+
+func remoteConnNil(log *log.Entry, conn net.Conn) bool {
+	if conn == nil {
+		log.Errorf("ice conn is nil")
+		return true
+	}
+
+	if conn.RemoteAddr() == nil {
+		log.Errorf("ICE remote address is nil")
+		return true
+	}
+
+	return false
+}

client/server/panic_generic.go

@@ -0,0 +1,7 @@
+//go:build !windows
+
+package server
+
+func handlePanicLog() error {
+	return nil
+}

client/server/panic_windows.go

@@ -0,0 +1,83 @@
+package server
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"syscall"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	windowsPanicLogEnvVar = "NB_WINDOWS_PANIC_LOG"
+	// STD_ERROR_HANDLE ((DWORD)-12) = 4294967284
+	stdErrorHandle = ^uintptr(11)
+)
+
+var (
+	kernel32 = syscall.NewLazyDLL("kernel32.dll")
+
+	// https://learn.microsoft.com/en-us/windows/console/setstdhandle
+	setStdHandleFn = kernel32.NewProc("SetStdHandle")
+)
+
+func handlePanicLog() error {
+	logPath := os.Getenv(windowsPanicLogEnvVar)
+	if logPath == "" {
+		return nil
+	}
+
+	// Ensure the directory exists
+	logDir := filepath.Dir(logPath)
+	if err := os.MkdirAll(logDir, 0750); err != nil {
+		return fmt.Errorf("create panic log directory: %w", err)
+	}
+	if err := util.EnforcePermission(logPath); err != nil {
+		return fmt.Errorf("enforce permission on panic log file: %w", err)
+	}
+
+	// Open log file with append mode
+	f, err := os.OpenFile(logPath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
+	if err != nil {
+		return fmt.Errorf("open panic log file: %w", err)
+	}
+
+	// Redirect stderr to the file
+	if err = redirectStderr(f); err != nil {
+		if closeErr := f.Close(); closeErr != nil {
+			log.Warnf("failed to close file after redirect error: %v", closeErr)
+		}
+		return fmt.Errorf("redirect stderr: %w", err)
+	}
+
+	log.Infof("successfully configured panic logging to: %s", logPath)
+	return nil
+}
+
+// redirectStderr redirects stderr to the provided file
+func redirectStderr(f *os.File) error {
+	// Get the current process's stderr handle
+	if err := setStdHandle(f); err != nil {
+		return fmt.Errorf("failed to set stderr handle: %w", err)
+	}
+
+	// Also set os.Stderr for Go's standard library
+	os.Stderr = f
+
+	return nil
+}
+
+func setStdHandle(f *os.File) error {
+	handle := f.Fd()
+	r0, _, e1 := setStdHandleFn.Call(stdErrorHandle, handle)
+	if r0 == 0 {
+		if e1 != nil {
+			return e1
+		}
+		return syscall.EINVAL
+	}
+	return nil
+}

client/server/server.go

@@ -97,6 +97,10 @@ func (s *Server) Start() error {
 	defer s.mutex.Unlock()
 	state := internal.CtxGetState(s.rootCtx)
 
+	if err := handlePanicLog(); err != nil {
+		log.Warnf("failed to redirect stderr: %v", err)
+	}
+
 	if err := restoreResidualState(s.rootCtx); err != nil {
 		log.Warnf(errRestoreResidualState, err)
 	}
@@ -622,6 +626,8 @@ func (s *Server) Down(ctx context.Context, _ *proto.DownRequest) (*proto.DownRes
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
+	s.oauthAuthFlow = oauthAuthFlow{}
+
 	if s.actCancel == nil {
 		return nil, fmt.Errorf("service is not up")
 	}

funding.json

@@ -0,0 +1,126 @@
+{
+  "version": "v1.0.0",
+  "entity": {
+    "type": "organisation",
+    "role": "owner",
+    "name": "NetBird GmbH",
+    "email": "hello@netbird.io",
+    "phone": "",
+    "description": "NetBird GmbH is a Berlin-based software company specializing in the development of open-source network security solutions. Network security is utterly complex and expensive, accessible only to companies with multi-million dollar IT budgets. In contrast, there are millions of companies left behind. Our mission is to create an advanced network and cybersecurity platform that is both easy-to-use and affordable for teams of all sizes and budgets. By leveraging the open-source strategy and technological advancements, NetBird aims to set the industry standard for connecting and securing IT infrastructure.",
+    "webpageUrl": {
+      "url": "https://github.com/netbirdio"
+    }
+  },
+  "projects": [
+    {
+      "guid": "netbird",
+      "name": "NetBird",
+      "description": "NetBird is a configuration-free peer-to-peer private network and a centralized access control system combined in a single open-source platform. It makes it easy to create secure WireGuard-based private networks for your organization or home.",
+      "webpageUrl": {
+        "url": "https://github.com/netbirdio/netbird"
+      },
+      "repositoryUrl": {
+        "url": "https://github.com/netbirdio/netbird"
+      },
+      "licenses": [
+        "BSD-3"
+      ],
+      "tags": [
+        "network-security",
+        "vpn",
+        "developer-tools",
+        "ztna",
+        "zero-trust",
+        "remote-access",
+        "wireguard",
+        "peer-to-peer",
+        "private-networking",
+        "software-defined-networking"
+      ]
+    }
+  ],
+  "funding": {
+    "channels": [
+      {
+        "guid": "github-sponsors",
+        "type": "payment-provider",
+        "address": "https://github.com/sponsors/netbirdio",
+        "description": ""
+      },
+      {
+        "guid": "bank-transfer",
+        "type": "bank",
+        "address": "",
+        "description": "Contact us at hello@netbird.io for bank transfer details."
+      }
+    ],
+    "plans": [
+      {
+        "guid": "support-yearly",
+        "status": "active",
+        "name": "Support Open Source Development and Maintenance - Yearly",
+        "description": "This will help us partially cover the yearly cost of maintaining the open-source NetBird project.",
+        "amount": 100000,
+        "currency": "USD",
+        "frequency": "yearly",
+        "channels": [
+          "github-sponsors",
+          "bank-transfer"
+        ]
+      },
+      {
+        "guid": "support-one-time-year",
+        "status": "active",
+        "name": "Support Open Source Development and Maintenance - One Year",
+        "description": "This will help us partially cover the yearly cost of maintaining the open-source NetBird project.",
+        "amount": 100000,
+        "currency": "USD",
+        "frequency": "one-time",
+        "channels": [
+          "github-sponsors",
+          "bank-transfer"
+        ]
+      },
+      {
+        "guid": "support-one-time-monthly",
+        "status": "active",
+        "name": "Support Open Source Development and Maintenance - Monthly",
+        "description": "This will help us partially cover the monthly cost of maintaining the open-source NetBird project.",
+        "amount": 10000,
+        "currency": "USD",
+        "frequency": "monthly",
+        "channels": [
+          "github-sponsors",
+          "bank-transfer"
+        ]
+      },
+      {
+        "guid": "support-monthly",
+        "status": "active",
+        "name": "Support Open Source Development and Maintenance - One Month",
+        "description": "This will help us partially cover the monthly cost of maintaining the open-source NetBird project.",
+        "amount": 10000,
+        "currency": "USD",
+        "frequency": "monthly",
+        "channels": [
+          "github-sponsors",
+          "bank-transfer"
+        ]
+      },
+      {
+        "guid": "goodwill",
+        "status": "active",
+        "name": "Goodwill Plan",
+        "description": "Pay anything you wish to show your goodwill for the project.",
+        "amount": 0,
+        "currency": "USD",
+        "frequency": "monthly",
+        "channels": [
+          "github-sponsors",
+          "bank-transfer"
+        ]
+      }
+    ],
+    "history": null
+  }
+}

go.mod

@@ -60,7 +60,7 @@ require (
 	github.com/miekg/dns v1.1.59
 	github.com/mitchellh/hashstructure/v2 v2.0.2
 	github.com/nadoo/ipset v0.5.0
-	github.com/netbirdio/management-integrations/integrations v0.0.0-20240929132811-9af486d346fd
+	github.com/netbirdio/management-integrations/integrations v0.0.0-20241106153857-de8e2beb5254
 	github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20241010133937-e0df50df217d
 	github.com/okta/okta-sdk-golang/v2 v2.18.0
 	github.com/oschwald/maxminddb-golang v1.12.0
@@ -71,7 +71,6 @@ require (
 	github.com/pion/transport/v3 v3.0.1
 	github.com/pion/turn/v3 v3.0.1
 	github.com/prometheus/client_golang v1.19.1
-	github.com/r3labs/diff/v3 v3.0.1
 	github.com/rs/xid v1.3.0
 	github.com/shirou/gopsutil/v3 v3.24.4
 	github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
@@ -156,7 +155,7 @@ require (
 	github.com/go-text/typesetting v0.1.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/google/btree v1.0.1 // indirect
+	github.com/google/btree v1.1.2 // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.3 // indirect
@@ -211,8 +210,6 @@ require (
 	github.com/tklauser/go-sysconf v0.3.14 // indirect
 	github.com/tklauser/numcpus v0.8.0 // indirect
 	github.com/vishvananda/netns v0.0.4 // indirect
-	github.com/vmihailenco/msgpack/v5 v5.3.5 // indirect
-	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	github.com/yuin/goldmark v1.7.1 // indirect
 	github.com/zeebo/blake3 v0.2.3 // indirect
 	go.opencensus.io v0.24.0 // indirect
@@ -231,7 +228,7 @@ require (
 	gopkg.in/square/go-jose.v2 v2.6.0 // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/tomb.v2 v2.0.0-20161208151619-d5d1b5820637 // indirect
-	gvisor.dev/gvisor v0.0.0-20230927004350-cbd86285d259 // indirect
+	gvisor.dev/gvisor v0.0.0-20231020174304-b8a429915ff1 // indirect
 	k8s.io/apimachinery v0.26.2 // indirect
 )
 
@@ -239,7 +236,7 @@ replace github.com/kardianos/service => github.com/netbirdio/service v0.0.0-2024
 
 replace github.com/getlantern/systray => github.com/netbirdio/systray v0.0.0-20231030152038-ef1ed2a27949
 
-replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-20240105182236-6c340dd55aed
+replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73
 
 replace github.com/cloudflare/circl => github.com/cunicu/circl v0.0.0-20230801113412-fec58fc7b5f6
 

go.sum

@@ -297,8 +297,8 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
-github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
-github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
+github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU=
+github.com/google/btree v1.1.2/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
@@ -521,14 +521,14 @@ github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944 h1:TDtJKmM6S
 github.com/netbirdio/go-netroute v0.0.0-20240611143515-f59b0e1d3944/go.mod h1:sHA6TRxjQ6RLbnI+3R4DZo2Eseg/iKiPRfNmcuNySVQ=
 github.com/netbirdio/ice/v3 v3.0.0-20240315174635-e72a50fcb64e h1:PURA50S8u4mF6RrkYYCAvvPCixhqqEiEy3Ej6avh04c=
 github.com/netbirdio/ice/v3 v3.0.0-20240315174635-e72a50fcb64e/go.mod h1:YMLU7qbKfVjmEv7EoZPIVEI+kNYxWCdPK3VS0BU+U4Q=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20240929132811-9af486d346fd h1:phKq1S1Y/lnqEhP5Qknta733+rPX16dRDHM7hKkot9c=
-github.com/netbirdio/management-integrations/integrations v0.0.0-20240929132811-9af486d346fd/go.mod h1:nykwWZnxb+sJz2Z//CEq45CMRWSHllH8pODKRB8eY7Y=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20241106153857-de8e2beb5254 h1:L8mNd3tBxMdnQNxMNJ+/EiwHwizNOMy8/nHLVGNfjpg=
+github.com/netbirdio/management-integrations/integrations v0.0.0-20241106153857-de8e2beb5254/go.mod h1:nykwWZnxb+sJz2Z//CEq45CMRWSHllH8pODKRB8eY7Y=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502 h1:3tHlFmhTdX9axERMVN63dqyFqnvuD+EMJHzM7mNGON8=
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20241010133937-e0df50df217d h1:bRq5TKgC7Iq20pDiuC54yXaWnAVeS5PdGpSokFTlR28=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20241010133937-e0df50df217d/go.mod h1:5/sjFmLb8O96B5737VCqhHyGRzNFIaN/Bu7ZodXc3qQ=
-github.com/netbirdio/wireguard-go v0.0.0-20240105182236-6c340dd55aed h1:t0UADZUJDaaZgfKrt8JUPrOLL9Mg/ryjP85RAH53qgs=
-github.com/netbirdio/wireguard-go v0.0.0-20240105182236-6c340dd55aed/go.mod h1:tkCQ4FQXmpAgYVh++1cq16/dH4QJtmvpRv19DWGAHSA=
+github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73 h1:jayg97LH/jJlvpIHVxueTfa+tfQ+FY8fy2sIhCwkz0g=
+github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73/go.mod h1:tkCQ4FQXmpAgYVh++1cq16/dH4QJtmvpRv19DWGAHSA=
 github.com/nicksnyder/go-i18n/v2 v2.4.0 h1:3IcvPOAvnCKwNm0TB0dLDTuawWEj+ax/RERNC+diLMM=
 github.com/nicksnyder/go-i18n/v2 v2.4.0/go.mod h1:nxYSZE9M0bf3Y70gPQjN9ha7XNHX7gMc814+6wVyEI4=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
@@ -605,8 +605,6 @@ github.com/prometheus/common v0.53.0 h1:U2pL9w9nmJwJDa4qqLQ3ZaePJ6ZTwt7cMD3AG3+a
 github.com/prometheus/common v0.53.0/go.mod h1:BrxBKv3FWBIGXw89Mg1AeBq7FSyRzXWI3l3e7W3RN5U=
 github.com/prometheus/procfs v0.15.0 h1:A82kmvXJq2jTu5YUhSGNlYoxh85zLnKgPz4bMZgI5Ek=
 github.com/prometheus/procfs v0.15.0/go.mod h1:Y0RJ/Y5g5wJpkTisOtqwDSo4HwhGmLB4VQSw2sQJLHk=
-github.com/r3labs/diff/v3 v3.0.1 h1:CBKqf3XmNRHXKmdU7mZP1w7TV0pDyVCis1AUHtA4Xtg=
-github.com/r3labs/diff/v3 v3.0.1/go.mod h1:f1S9bourRbiM66NskseyUdo0fTmEE0qKrikYJX63dgo=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
@@ -699,10 +697,6 @@ github.com/vishvananda/netlink v1.2.1-beta.2/go.mod h1:twkDnbuQxJYemMlGd4JFIcuhg
 github.com/vishvananda/netns v0.0.0-20200728191858-db3c7e526aae/go.mod h1:DD4vA1DwXk04H54A1oHXtwZmA0grkVMdPxx/VGLCah0=
 github.com/vishvananda/netns v0.0.4 h1:Oeaw1EM2JMxD51g9uhtC0D7erkIjgmj8+JZc26m1YX8=
 github.com/vishvananda/netns v0.0.4/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
-github.com/vmihailenco/msgpack/v5 v5.3.5 h1:5gO0H1iULLWGhs2H5tbAHIZTV8/cYafcFOr9znI5mJU=
-github.com/vmihailenco/msgpack/v5 v5.3.5/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc=
-github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
-github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
@@ -1238,8 +1232,8 @@ gorm.io/gorm v1.25.7-0.20240204074919-46816ad31dde h1:9DShaph9qhkIYw7QF91I/ynrr4
 gorm.io/gorm v1.25.7-0.20240204074919-46816ad31dde/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
 gotest.tools/v3 v3.5.0 h1:Ljk6PdHdOhAb5aDMWXjDLMMhph+BpztA4v1QdqEW2eY=
 gotest.tools/v3 v3.5.0/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=
-gvisor.dev/gvisor v0.0.0-20230927004350-cbd86285d259 h1:TbRPT0HtzFP3Cno1zZo7yPzEEnfu8EjLfl6IU9VfqkQ=
-gvisor.dev/gvisor v0.0.0-20230927004350-cbd86285d259/go.mod h1:AVgIgHMwK63XvmAzWG9vLQ41YnVHN0du0tEC46fI7yY=
+gvisor.dev/gvisor v0.0.0-20231020174304-b8a429915ff1 h1:qDCwdCWECGnwQSQC01Dpnp09fRHxJs9PbktotUqG+hs=
+gvisor.dev/gvisor v0.0.0-20231020174304-b8a429915ff1/go.mod h1:8hmigyCdYtw5xJGfQDJzSH5Ju8XEIDBnpyi8+O6GRt8=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

management/server/account.go

@@ -1249,7 +1249,7 @@ func (am *DefaultAccountManager) peerInactivityExpirationJob(ctx context.Context
 
 		account, err := am.Store.GetAccount(ctx, accountID)
 		if err != nil {
-			log.Errorf("failed getting account %s expiring peers", account.Id)
+			log.Errorf("failed getting account %s expiring peers", accountID)
 			return account.GetNextInactivePeerExpiration()
 		}
 

management/server/account_test.go

@@ -29,14 +29,18 @@ import (
 )
 
 type MocIntegratedValidator struct {
+	ValidatePeerFunc func(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, bool, error)
 }
 
 func (a MocIntegratedValidator) ValidateExtraSettings(_ context.Context, newExtraSettings *account.ExtraSettings, oldExtraSettings *account.ExtraSettings, peers map[string]*nbpeer.Peer, userID string, accountID string) error {
 	return nil
 }
 
-func (a MocIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, error) {
-	return update, nil
+func (a MocIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, bool, error) {
+	if a.ValidatePeerFunc != nil {
+		return a.ValidatePeerFunc(context.Background(), update, peer, userID, accountID, dnsDomain, peersGroup, extraSettings)
+	}
+	return update, false, nil
 }
 func (a MocIntegratedValidator) GetValidatedPeers(accountID string, groups map[string]*group.Group, peers map[string]*nbpeer.Peer, extraSettings *account.ExtraSettings) (map[string]struct{}, error) {
 	validatedPeers := make(map[string]struct{})
@@ -978,6 +982,110 @@ func TestAccountManager_DeleteAccount(t *testing.T) {
 	}
 }
 
+func BenchmarkTest_GetAccountWithclaims(b *testing.B) {
+	claims := jwtclaims.AuthorizationClaims{
+		Domain:         "example.com",
+		UserId:         "pvt-domain-user",
+		DomainCategory: PrivateCategory,
+	}
+
+	publicClaims := jwtclaims.AuthorizationClaims{
+		Domain:         "test.com",
+		UserId:         "public-domain-user",
+		DomainCategory: PublicCategory,
+	}
+
+	am, err := createManager(b)
+	if err != nil {
+		b.Fatal(err)
+		return
+	}
+	id, err := am.getAccountIDWithAuthorizationClaims(context.Background(), claims)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	pid, err := am.getAccountIDWithAuthorizationClaims(context.Background(), publicClaims)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	users := genUsers("priv", 100)
+
+	acc, err := am.Store.GetAccount(context.Background(), id)
+	if err != nil {
+		b.Fatal(err)
+	}
+	acc.Users = users
+
+	err = am.Store.SaveAccount(context.Background(), acc)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	userP := genUsers("pub", 100)
+
+	pacc, err := am.Store.GetAccount(context.Background(), pid)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	pacc.Users = userP
+
+	err = am.Store.SaveAccount(context.Background(), pacc)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.Run("public without account ID", func(b *testing.B) {
+		//b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), publicClaims)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	b.Run("private without account ID", func(b *testing.B) {
+		//b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), claims)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	b.Run("private with account ID", func(b *testing.B) {
+		claims.AccountId = id
+		//b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), claims)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+}
+
+func genUsers(p string, n int) map[string]*User {
+	users := map[string]*User{}
+	now := time.Now()
+	for i := 0; i < n; i++ {
+		users[fmt.Sprintf("%s-%d", p, i)] = &User{
+			Id:         fmt.Sprintf("%s-%d", p, i),
+			Role:       UserRoleAdmin,
+			LastLogin:  now,
+			CreatedAt:  now,
+			Issued:     "api",
+			AutoGroups: []string{"one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"},
+		}
+	}
+	return users
+}
+
 func TestAccountManager_AddPeer(t *testing.T) {
 	manager, err := createManager(t)
 	if err != nil {

management/server/differs/netip.go

@@ -1,82 +0,0 @@
-package differs
-
-import (
-	"fmt"
-	"net/netip"
-	"reflect"
-
-	"github.com/r3labs/diff/v3"
-)
-
-// NetIPAddr is a custom differ for netip.Addr
-type NetIPAddr struct {
-	DiffFunc func(path []string, a, b reflect.Value, p interface{}) error
-}
-
-func (differ NetIPAddr) Match(a, b reflect.Value) bool {
-	return diff.AreType(a, b, reflect.TypeOf(netip.Addr{}))
-}
-
-func (differ NetIPAddr) Diff(_ diff.DiffType, _ diff.DiffFunc, cl *diff.Changelog, path []string, a, b reflect.Value, _ interface{}) error {
-	if a.Kind() == reflect.Invalid {
-		cl.Add(diff.CREATE, path, nil, b.Interface())
-		return nil
-	}
-
-	if b.Kind() == reflect.Invalid {
-		cl.Add(diff.DELETE, path, a.Interface(), nil)
-		return nil
-	}
-
-	fromAddr, ok1 := a.Interface().(netip.Addr)
-	toAddr, ok2 := b.Interface().(netip.Addr)
-	if !ok1 || !ok2 {
-		return fmt.Errorf("invalid type for netip.Addr")
-	}
-
-	if fromAddr.String() != toAddr.String() {
-		cl.Add(diff.UPDATE, path, fromAddr.String(), toAddr.String())
-	}
-
-	return nil
-}
-
-func (differ NetIPAddr) InsertParentDiffer(dfunc func(path []string, a, b reflect.Value, p interface{}) error) {
-	differ.DiffFunc = dfunc //nolint
-}
-
-// NetIPPrefix is a custom differ for netip.Prefix
-type NetIPPrefix struct {
-	DiffFunc func(path []string, a, b reflect.Value, p interface{}) error
-}
-
-func (differ NetIPPrefix) Match(a, b reflect.Value) bool {
-	return diff.AreType(a, b, reflect.TypeOf(netip.Prefix{}))
-}
-
-func (differ NetIPPrefix) Diff(_ diff.DiffType, _ diff.DiffFunc, cl *diff.Changelog, path []string, a, b reflect.Value, _ interface{}) error {
-	if a.Kind() == reflect.Invalid {
-		cl.Add(diff.CREATE, path, nil, b.Interface())
-		return nil
-	}
-	if b.Kind() == reflect.Invalid {
-		cl.Add(diff.DELETE, path, a.Interface(), nil)
-		return nil
-	}
-
-	fromPrefix, ok1 := a.Interface().(netip.Prefix)
-	toPrefix, ok2 := b.Interface().(netip.Prefix)
-	if !ok1 || !ok2 {
-		return fmt.Errorf("invalid type for netip.Addr")
-	}
-
-	if fromPrefix.String() != toPrefix.String() {
-		cl.Add(diff.UPDATE, path, fromPrefix.String(), toPrefix.String())
-	}
-
-	return nil
-}
-
-func (differ NetIPPrefix) InsertParentDiffer(dfunc func(path []string, a, b reflect.Value, p interface{}) error) {
-	differ.DiffFunc = dfunc //nolint
-}

management/server/dns_test.go

@@ -8,9 +8,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/telemetry"
-	"github.com/stretchr/testify/assert"
 
 	"github.com/stretchr/testify/require"
 
@@ -521,23 +522,64 @@ func TestDNSAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	err = manager.SaveGroup(context.Background(), account.Id, userID, &group.Group{
-		ID:    "groupA",
-		Name:  "GroupA",
-		Peers: []string{peer1.ID, peer2.ID, peer3.ID},
-	})
-	assert.NoError(t, err)
+	// Creating DNS settings with groups that have no peers should not update account peers or send peer update
+	t.Run("creating dns setting with unused groups", func(t *testing.T) {
+		done := make(chan struct{})
+		go func() {
+			peerShouldNotReceiveUpdate(t, updMsg)
+			close(done)
+		}()
 
-	_, err = manager.CreateNameServerGroup(
-		context.Background(), account.Id, "ns-group-1", "ns-group-1", []dns.NameServer{{
-			IP:     netip.MustParseAddr(peer1.IP.String()),
-			NSType: dns.UDPNameServerType,
-			Port:   dns.DefaultDNSPort,
-		}},
-		[]string{"groupA"},
-		true, []string{}, true, userID, false,
-	)
-	assert.NoError(t, err)
+		_, err = manager.CreateNameServerGroup(
+			context.Background(), account.Id, "ns-group", "ns-group", []dns.NameServer{{
+				IP:     netip.MustParseAddr(peer1.IP.String()),
+				NSType: dns.UDPNameServerType,
+				Port:   dns.DefaultDNSPort,
+			}},
+			[]string{"groupB"},
+			true, []string{}, true, userID, false,
+		)
+		assert.NoError(t, err)
+
+		select {
+		case <-done:
+		case <-time.After(time.Second):
+			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+		}
+	})
+
+	// Creating DNS settings with groups that have peers should update account peers and send peer update
+	t.Run("creating dns setting with used groups", func(t *testing.T) {
+		err = manager.SaveGroup(context.Background(), account.Id, userID, &group.Group{
+			ID:    "groupA",
+			Name:  "GroupA",
+			Peers: []string{peer1.ID, peer2.ID, peer3.ID},
+		})
+		assert.NoError(t, err)
+
+		done := make(chan struct{})
+		go func() {
+			peerShouldReceiveUpdate(t, updMsg)
+			close(done)
+		}()
+
+		_, err = manager.CreateNameServerGroup(
+			context.Background(), account.Id, "ns-group-1", "ns-group-1", []dns.NameServer{{
+				IP:     netip.MustParseAddr(peer1.IP.String()),
+				NSType: dns.UDPNameServerType,
+				Port:   dns.DefaultDNSPort,
+			}},
+			[]string{"groupA"},
+			true, []string{}, true, userID, false,
+		)
+		assert.NoError(t, err)
+
+		select {
+		case <-done:
+		case <-time.After(time.Second):
+			t.Error("timeout waiting for peerShouldReceiveUpdate")
+		}
+	})
 
 	// Saving DNS settings with groups that have peers should update account peers and send peer update
 	t.Run("saving dns setting with used groups", func(t *testing.T) {
@@ -559,27 +601,6 @@ func TestDNSAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Saving unchanged DNS settings with used groups should update account peers and not send peer update
-	// since there is no change in the network map
-	t.Run("saving unchanged dns setting with used groups", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveDNSSettings(context.Background(), account.Id, userID, &DNSSettings{
-			DisabledManagementGroups: []string{"groupA", "groupB"},
-		})
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Removing group with no peers from DNS settings  should not trigger updates to account peers or send peer updates
 	t.Run("removing group with no peers from dns settings", func(t *testing.T) {
 		done := make(chan struct{})

management/server/group_test.go

@@ -8,12 +8,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	nbdns "github.com/netbirdio/netbird/dns"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/route"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -536,29 +537,6 @@ func TestGroupAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Saving an unchanged group should trigger account peers update and not send peer update
-	// since there is no change in the network map
-	t.Run("saving unchanged group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
-			ID:    "groupA",
-			Name:  "GroupA",
-			Peers: []string{peer1.ID, peer2.ID},
-		})
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// adding peer to a used group should update account peers and send peer update
 	t.Run("adding peer to linked group", func(t *testing.T) {
 		done := make(chan struct{})

management/server/integrated_validator/interface.go

@@ -11,7 +11,7 @@ import (
 // IntegratedValidator interface exists to avoid the circle dependencies
 type IntegratedValidator interface {
 	ValidateExtraSettings(ctx context.Context, newExtraSettings *account.ExtraSettings, oldExtraSettings *account.ExtraSettings, peers map[string]*nbpeer.Peer, userID string, accountID string) error
-	ValidatePeer(ctx context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, error)
+	ValidatePeer(ctx context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, bool, error)
 	PreparePeer(ctx context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *account.ExtraSettings) *nbpeer.Peer
 	IsNotValidPeer(ctx context.Context, accountID string, peer *nbpeer.Peer, peersGroup []string, extraSettings *account.ExtraSettings) (bool, bool, error)
 	GetValidatedPeers(accountID string, groups map[string]*nbgroup.Group, peers map[string]*nbpeer.Peer, extraSettings *account.ExtraSettings) (map[string]struct{}, error)

management/server/management_test.go

@@ -453,8 +453,8 @@ func (a MocIntegratedValidator) ValidateExtraSettings(_ context.Context, newExtr
 	return nil
 }
 
-func (a MocIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, error) {
-	return update, nil
+func (a MocIntegratedValidator) ValidatePeer(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *account.ExtraSettings) (*nbpeer.Peer, bool, error) {
+	return update, false, nil
 }
 
 func (a MocIntegratedValidator) GetValidatedPeers(accountID string, groups map[string]*group.Group, peers map[string]*nbpeer.Peer, extraSettings *account.ExtraSettings) (map[string]struct{}, error) {

management/server/nameserver_test.go

@@ -1065,36 +1065,6 @@ func TestNameServerAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// saving unchanged nameserver group should update account peers and not send peer update
-	t.Run("saving unchanged nameserver group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		newNameServerGroupB.NameServers = []nbdns.NameServer{
-			{
-				IP:     netip.MustParseAddr("1.1.1.2"),
-				NSType: nbdns.UDPNameServerType,
-				Port:   nbdns.DefaultDNSPort,
-			},
-			{
-				IP:     netip.MustParseAddr("8.8.8.8"),
-				NSType: nbdns.UDPNameServerType,
-				Port:   nbdns.DefaultDNSPort,
-			},
-		}
-		err = manager.SaveNameServerGroup(context.Background(), account.Id, userID, newNameServerGroupB)
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Deleting a nameserver group should update account peers and send peer update
 	t.Run("deleting nameserver group", func(t *testing.T) {
 		done := make(chan struct{})

management/server/network.go

@@ -41,9 +41,9 @@ type Network struct {
 	Dns        string
 	// Serial is an ID that increments by 1 when any change to the network happened (e.g. new peer has been added).
 	// Used to synchronize state to the client apps.
-	Serial uint64 `diff:"-"`
+	Serial uint64
 
-	mu sync.Mutex `json:"-" gorm:"-" diff:"-"`
+	mu sync.Mutex `json:"-" gorm:"-"`
 }
 
 // NewNetwork creates a new Network initializing it with a Serial=0

management/server/peer.go

@@ -189,7 +189,8 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		return nil, status.Errorf(status.NotFound, "peer %s not found", update.ID)
 	}
 
-	update, err = am.integratedPeerValidator.ValidatePeer(ctx, update, peer, userID, accountID, am.GetDNSDomain(), account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
+	var requiresPeerUpdates bool
+	update, requiresPeerUpdates, err = am.integratedPeerValidator.ValidatePeer(ctx, update, peer, userID, accountID, am.GetDNSDomain(), account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
 		return nil, err
 	}
@@ -265,7 +266,7 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		return nil, err
 	}
 
-	if peerLabelUpdated {
+	if peerLabelUpdated || requiresPeerUpdates {
 		am.updateAccountPeers(ctx, account)
 	}
 
@@ -589,6 +590,12 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		return nil, nil, nil, fmt.Errorf("error getting account: %w", err)
 	}
 
+	allGroup, err := account.GetGroupAll()
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("error getting all group ID: %w", err)
+	}
+
+	groupsToAdd = append(groupsToAdd, allGroup.ID)
 	if areGroupChangesAffectPeers(account, groupsToAdd) {
 		am.updateAccountPeers(ctx, account)
 	}

management/server/peer/peer.go

@@ -4,6 +4,7 @@ import (
 	"net"
 	"net/netip"
 	"slices"
+	"sort"
 	"time"
 )
 
@@ -19,33 +20,33 @@ type Peer struct {
 	// IP address of the Peer
 	IP net.IP `gorm:"serializer:json"`
 	// Meta is a Peer system meta data
-	Meta PeerSystemMeta `gorm:"embedded;embeddedPrefix:meta_" diff:"-"`
+	Meta PeerSystemMeta `gorm:"embedded;embeddedPrefix:meta_"`
 	// Name is peer's name (machine name)
 	Name string
 	// DNSLabel is the parsed peer name for domain resolution. It is used to form an FQDN by appending the account's
 	// domain to the peer label. e.g. peer-dns-label.netbird.cloud
 	DNSLabel string
 	// Status peer's management connection status
-	Status *PeerStatus `gorm:"embedded;embeddedPrefix:peer_status_" diff:"-"`
+	Status *PeerStatus `gorm:"embedded;embeddedPrefix:peer_status_"`
 	// The user ID that registered the peer
-	UserID string `diff:"-"`
+	UserID string
 	// SSHKey is a public SSH key of the peer
 	SSHKey string
 	// SSHEnabled indicates whether SSH server is enabled on the peer
 	SSHEnabled bool
 	// LoginExpirationEnabled indicates whether peer's login expiration is enabled and once expired the peer has to re-login.
 	// Works with LastLogin
-	LoginExpirationEnabled bool `diff:"-"`
+	LoginExpirationEnabled bool
 
-	InactivityExpirationEnabled bool `diff:"-"`
+	InactivityExpirationEnabled bool
 	// LastLogin the time when peer performed last login operation
-	LastLogin time.Time `diff:"-"`
+	LastLogin time.Time
 	// CreatedAt records the time the peer was created
-	CreatedAt time.Time `diff:"-"`
+	CreatedAt time.Time
 	// Indicate ephemeral peer attribute
-	Ephemeral bool `diff:"-"`
+	Ephemeral bool
 	// Geo location based on connection IP
-	Location Location `gorm:"embedded;embeddedPrefix:location_" diff:"-"`
+	Location Location `gorm:"embedded;embeddedPrefix:location_"`
 }
 
 type PeerStatus struct { //nolint:revive
@@ -107,6 +108,12 @@ type PeerSystemMeta struct { //nolint:revive
 }
 
 func (p PeerSystemMeta) isEqual(other PeerSystemMeta) bool {
+	sort.Slice(p.NetworkAddresses, func(i, j int) bool {
+		return p.NetworkAddresses[i].Mac < p.NetworkAddresses[j].Mac
+	})
+	sort.Slice(other.NetworkAddresses, func(i, j int) bool {
+		return other.NetworkAddresses[i].Mac < other.NetworkAddresses[j].Mac
+	})
 	equalNetworkAddresses := slices.EqualFunc(p.NetworkAddresses, other.NetworkAddresses, func(addr NetworkAddress, oAddr NetworkAddress) bool {
 		return addr.Mac == oAddr.Mac && addr.NetIP == oAddr.NetIP
 	})
@@ -114,6 +121,12 @@ func (p PeerSystemMeta) isEqual(other PeerSystemMeta) bool {
 		return false
 	}
 
+	sort.Slice(p.Files, func(i, j int) bool {
+		return p.Files[i].Path < p.Files[j].Path
+	})
+	sort.Slice(other.Files, func(i, j int) bool {
+		return other.Files[i].Path < other.Files[j].Path
+	})
 	equalFiles := slices.EqualFunc(p.Files, other.Files, func(file File, oFile File) bool {
 		return file.Path == oFile.Path && file.Exist == oFile.Exist && file.ProcessIsRunning == oFile.ProcessIsRunning
 	})

management/server/peer/peer_test.go

@@ -2,6 +2,7 @@ package peer
 
 import (
 	"fmt"
+	"net/netip"
 	"testing"
 )
 
@@ -29,3 +30,56 @@ func BenchmarkFQDN(b *testing.B) {
 		}
 	})
 }
+
+func TestIsEqual(t *testing.T) {
+	meta1 := PeerSystemMeta{
+		NetworkAddresses: []NetworkAddress{{
+			NetIP: netip.MustParsePrefix("192.168.1.2/24"),
+			Mac:   "2",
+		},
+			{
+				NetIP: netip.MustParsePrefix("192.168.1.0/24"),
+				Mac:   "1",
+			},
+		},
+		Files: []File{
+			{
+				Path:             "/etc/hosts1",
+				Exist:            true,
+				ProcessIsRunning: true,
+			},
+			{
+				Path:             "/etc/hosts2",
+				Exist:            false,
+				ProcessIsRunning: false,
+			},
+		},
+	}
+	meta2 := PeerSystemMeta{
+		NetworkAddresses: []NetworkAddress{
+			{
+				NetIP: netip.MustParsePrefix("192.168.1.0/24"),
+				Mac:   "1",
+			},
+			{
+				NetIP: netip.MustParsePrefix("192.168.1.2/24"),
+				Mac:   "2",
+			},
+		},
+		Files: []File{
+			{
+				Path:             "/etc/hosts2",
+				Exist:            false,
+				ProcessIsRunning: false,
+			},
+			{
+				Path:             "/etc/hosts1",
+				Exist:            true,
+				ProcessIsRunning: true,
+			},
+		},
+	}
+	if !meta1.isEqual(meta2) {
+		t.Error("meta1 should be equal to meta2")
+	}
+}

management/server/peer_test.go

@@ -22,6 +22,7 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/proto"
+	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
@@ -1398,6 +1399,50 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
+	t.Run("validator requires update", func(t *testing.T) {
+		requireUpdateFunc := func(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *nbAccount.ExtraSettings) (*nbpeer.Peer, bool, error) {
+			return update, true, nil
+		}
+
+		manager.integratedPeerValidator = MocIntegratedValidator{ValidatePeerFunc: requireUpdateFunc}
+		done := make(chan struct{})
+		go func() {
+			peerShouldReceiveUpdate(t, updMsg)
+			close(done)
+		}()
+
+		_, err = manager.UpdatePeer(context.Background(), account.Id, userID, peer1)
+		require.NoError(t, err)
+
+		select {
+		case <-done:
+		case <-time.After(time.Second):
+			t.Error("timeout waiting for peerShouldReceiveUpdate")
+		}
+	})
+
+	t.Run("validator requires no update", func(t *testing.T) {
+		requireNoUpdateFunc := func(_ context.Context, update *nbpeer.Peer, peer *nbpeer.Peer, userID string, accountID string, dnsDomain string, peersGroup []string, extraSettings *nbAccount.ExtraSettings) (*nbpeer.Peer, bool, error) {
+			return update, false, nil
+		}
+
+		manager.integratedPeerValidator = MocIntegratedValidator{ValidatePeerFunc: requireNoUpdateFunc}
+		done := make(chan struct{})
+		go func() {
+			peerShouldNotReceiveUpdate(t, updMsg)
+			close(done)
+		}()
+
+		_, err = manager.UpdatePeer(context.Background(), account.Id, userID, peer1)
+		require.NoError(t, err)
+
+		select {
+		case <-done:
+		case <-time.After(time.Second):
+			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+		}
+	})
+
 	// Adding peer to group linked with policy should update account peers and send peer update
 	t.Run("adding peer to group linked with policy", func(t *testing.T) {
 		err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{

management/server/policy.go

@@ -405,7 +405,9 @@ func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, po
 
 	am.StoreEvent(ctx, userID, policy.ID, accountID, activity.PolicyRemoved, policy.EventMeta())
 
-	am.updateAccountPeers(ctx, account)
+	if anyGroupHasPeers(account, policy.ruleGroups()) {
+		am.updateAccountPeers(ctx, account)
+	}
 
 	return nil
 }

management/server/policy_test.go

@@ -854,16 +854,11 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	})
 	assert.NoError(t, err)
 
-	updMsg1 := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
+	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
 	t.Cleanup(func() {
 		manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
 	})
 
-	updMsg2 := manager.peersUpdateManager.CreateChannel(context.Background(), peer2.ID)
-	t.Cleanup(func() {
-		manager.peersUpdateManager.CloseChannel(context.Background(), peer2.ID)
-	})
-
 	// Saving policy with rule groups with no peers should not update account's peers and not send peer update
 	t.Run("saving policy with rule groups with no peers", func(t *testing.T) {
 		policy := Policy{
@@ -883,7 +878,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg1)
+			peerShouldNotReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -918,7 +913,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg1)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -953,7 +948,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg2)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -987,7 +982,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg1)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1021,7 +1016,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg1)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1056,7 +1051,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg1)
+			peerShouldNotReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1090,7 +1085,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg1)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1104,46 +1099,13 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Saving unchanged policy should trigger account peers update but not send peer update
-	t.Run("saving unchanged policy", func(t *testing.T) {
-		policy := Policy{
-			ID:      "policy-source-destination-peers",
-			Enabled: true,
-			Rules: []*PolicyRule{
-				{
-					ID:            xid.New().String(),
-					Enabled:       true,
-					Sources:       []string{"groupA"},
-					Destinations:  []string{"groupD"},
-					Bidirectional: true,
-					Action:        PolicyTrafficActionAccept,
-				},
-			},
-		}
-
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg1)
-			close(done)
-		}()
-
-		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Deleting policy should trigger account peers update and send peer update
 	t.Run("deleting policy with source and destination groups with peers", func(t *testing.T) {
 		policyID := "policy-source-destination-peers"
 
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg1)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1164,7 +1126,7 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 		policyID := "policy-destination-has-peers-source-none"
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg2)
+			peerShouldReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 
@@ -1180,10 +1142,10 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 
 	// Deleting policy with no peers in groups should not update account's peers and not send peer update
 	t.Run("deleting policy with no peers in groups", func(t *testing.T) {
-		policyID := "policy-rule-groups-no-peers" // Deleting the policy created in Case 2
+		policyID := "policy-rule-groups-no-peers"
 		done := make(chan struct{})
 		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg1)
+			peerShouldNotReceiveUpdate(t, updMsg)
 			close(done)
 		}()
 

management/server/posture_checks_test.go

@@ -5,10 +5,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/netbirdio/netbird/management/server/group"
 	"github.com/rs/xid"
 	"github.com/stretchr/testify/assert"
 
+	"github.com/netbirdio/netbird/management/server/group"
+
 	"github.com/netbirdio/netbird/management/server/posture"
 )
 
@@ -264,25 +265,6 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Saving unchanged posture check should not trigger account peers update and not send peer update
-	// since there is no change in the network map
-	t.Run("saving unchanged posture check", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Removing posture check from policy should trigger account peers update and send peer update
 	t.Run("removing posture check from policy", func(t *testing.T) {
 		done := make(chan struct{})
@@ -412,50 +394,9 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Updating linked posture check to policy where source has peers but destination does not,
-	// should not trigger account peers update or send peer update
-	t.Run("updating linked posture check to policy where source has peers but destination does not", func(t *testing.T) {
-		policy = Policy{
-			ID:      "policyB",
-			Enabled: true,
-			Rules: []*PolicyRule{
-				{
-					Enabled:       true,
-					Sources:       []string{"groupA"},
-					Destinations:  []string{"groupB"},
-					Bidirectional: true,
-					Action:        PolicyTrafficActionAccept,
-				},
-			},
-			SourcePostureChecks: []string{postureCheck.ID},
-		}
-		err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
-		assert.NoError(t, err)
-
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		postureCheck.Checks = posture.ChecksDefinition{
-			NBVersionCheck: &posture.NBVersionCheck{
-				MinVersion: "0.29.0",
-			},
-		}
-		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Updating linked client posture check to policy where source has peers but destination does not,
 	// should trigger account peers update and send peer update
-	t.Run("updating linked client posture check to policy where source has peers but destination does not", func(t *testing.T) {
+	t.Run("updating linked posture check to policy where source has peers but destination does not", func(t *testing.T) {
 		policy = Policy{
 			ID:      "policyB",
 			Enabled: true,

management/server/route_test.go

@@ -1938,26 +1938,6 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Updating unchanged route should update account peers and not send peer update
-	t.Run("updating unchanged route", func(t *testing.T) {
-		baseRoute.Groups = []string{routeGroup1, routeGroup2}
-
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveRoute(context.Background(), account.Id, userID, &baseRoute)
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(time.Second):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
-
 	// Deleting the route should update account peers and send peer update
 	t.Run("deleting route", func(t *testing.T) {
 		done := make(chan struct{})

management/server/sql_store.go

@@ -69,9 +69,17 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine StoreEngine, metr
 	if err != nil {
 		conns = runtime.NumCPU()
 	}
+
+	if storeEngine == SqliteStoreEngine {
+		if err == nil {
+			log.WithContext(ctx).Warnf("setting NB_SQL_MAX_OPEN_CONNS is not supported for sqlite, using default value 1")
+		}
+		conns = 1
+	}
+
 	sql.SetMaxOpenConns(conns)
 
-	log.Infof("Set max open db connections to %d", conns)
+	log.WithContext(ctx).Infof("Set max open db connections to %d", conns)
 
 	if err := migrate(ctx, db); err != nil {
 		return nil, fmt.Errorf("migrate: %w", err)
@@ -296,7 +304,7 @@ func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.
 	peerCopy := peer.Copy()
 	peerCopy.AccountID = accountID
 
-	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+	err := s.db.Transaction(func(tx *gorm.DB) error {
 		// check if peer exists before saving
 		var peerID string
 		result := tx.Model(&nbpeer.Peer{}).Select("id").Find(&peerID, accountAndIDQueryCondition, accountID, peer.ID)
@@ -331,7 +339,7 @@ func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID
 	}
 
 	fieldsToUpdate := []string{"domain", "domain_category", "is_domain_primary_account"}
-	result := s.db.WithContext(ctx).Model(&Account{}).
+	result := s.db.Model(&Account{}).
 		Select(fieldsToUpdate).
 		Where(idQueryCondition, accountID).
 		Updates(&accountCopy)
@@ -403,14 +411,19 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 		}
 		usersToSave = append(usersToSave, *user)
 	}
-	return s.db.Session(&gorm.Session{FullSaveAssociations: true}).
+	err := s.db.Session(&gorm.Session{FullSaveAssociations: true}).
 		Clauses(clause.OnConflict{UpdateAll: true}).
 		Create(&usersToSave).Error
+	if err != nil {
+		return status.Errorf(status.Internal, "failed to save users to store: %v", err)
+	}
+
+	return nil
 }
 
 // SaveUser saves the given user to the database.
 func (s *SqlStore) SaveUser(ctx context.Context, lockStrength LockingStrength, user *User) error {
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
 	if result.Error != nil {
 		return status.Errorf(status.Internal, "failed to save user to store: %v", result.Error)
 	}
@@ -423,7 +436,7 @@ func (s *SqlStore) SaveGroups(ctx context.Context, lockStrength LockingStrength,
 		return nil
 	}
 
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&groups)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&groups)
 	if result.Error != nil {
 		return status.Errorf(status.Internal, "failed to save groups to store: %v", result.Error)
 	}
@@ -452,7 +465,7 @@ func (s *SqlStore) GetAccountByPrivateDomain(ctx context.Context, domain string)
 
 func (s *SqlStore) GetAccountIDByPrivateDomain(ctx context.Context, lockStrength LockingStrength, domain string) (string, error) {
 	var accountID string
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("id").
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("id").
 		Where("domain = ? and is_domain_primary_account = ? and domain_category = ?",
 			strings.ToLower(domain), true, PrivateCategory,
 		).First(&accountID)
@@ -469,7 +482,7 @@ func (s *SqlStore) GetAccountIDByPrivateDomain(ctx context.Context, lockStrength
 
 func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*Account, error) {
 	var key SetupKey
-	result := s.db.WithContext(ctx).Select("account_id").First(&key, keyQueryCondition, setupKey)
+	result := s.db.Select("account_id").First(&key, keyQueryCondition, setupKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -529,7 +542,7 @@ func (s *SqlStore) GetUserByTokenID(ctx context.Context, tokenID string) (*User,
 
 func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*User, error) {
 	var user User
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		Preload(clause.Associations).First(&user, idQueryCondition, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -662,7 +675,7 @@ func (s *SqlStore) GetAccount(ctx context.Context, accountID string) (*Account,
 
 func (s *SqlStore) GetAccountByUser(ctx context.Context, userID string) (*Account, error) {
 	var user User
-	result := s.db.WithContext(ctx).Select("account_id").First(&user, idQueryCondition, userID)
+	result := s.db.Select("account_id").First(&user, idQueryCondition, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -679,7 +692,7 @@ func (s *SqlStore) GetAccountByUser(ctx context.Context, userID string) (*Accoun
 
 func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*Account, error) {
 	var peer nbpeer.Peer
-	result := s.db.WithContext(ctx).Select("account_id").First(&peer, idQueryCondition, peerID)
+	result := s.db.Select("account_id").First(&peer, idQueryCondition, peerID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -696,8 +709,7 @@ func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*Acco
 
 func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (*Account, error) {
 	var peer nbpeer.Peer
-
-	result := s.db.WithContext(ctx).Select("account_id").First(&peer, keyQueryCondition, peerKey)
+	result := s.db.Select("account_id").First(&peer, keyQueryCondition, peerKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -715,7 +727,7 @@ func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (
 func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string) (string, error) {
 	var peer nbpeer.Peer
 	var accountID string
-	result := s.db.WithContext(ctx).Model(&peer).Select("account_id").Where(keyQueryCondition, peerKey).First(&accountID)
+	result := s.db.Model(&peer).Select("account_id").Where(keyQueryCondition, peerKey).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -741,7 +753,7 @@ func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
 
 func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string) (string, error) {
 	var accountID string
-	result := s.db.WithContext(ctx).Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
+	result := s.db.Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
@@ -760,7 +772,7 @@ func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength
 	var ipJSONStrings []string
 
 	// Fetch the IP addresses as JSON strings
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Where("account_id = ?", accountID).
 		Pluck("ip", &ipJSONStrings)
 	if result.Error != nil {
@@ -785,8 +797,7 @@ func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength
 
 func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string) ([]string, error) {
 	var labels []string
-
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Where("account_id = ?", accountID).
 		Pluck("dns_label", &labels)
 
@@ -803,8 +814,7 @@ func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength Lock
 
 func (s *SqlStore) GetAccountNetwork(ctx context.Context, lockStrength LockingStrength, accountID string) (*Network, error) {
 	var accountNetwork AccountNetwork
-
-	if err := s.db.WithContext(ctx).Model(&Account{}).Where(idQueryCondition, accountID).First(&accountNetwork).Error; err != nil {
+	if err := s.db.Model(&Account{}).Where(idQueryCondition, accountID).First(&accountNetwork).Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, status.NewAccountNotFoundError(accountID)
 		}
@@ -815,7 +825,7 @@ func (s *SqlStore) GetAccountNetwork(ctx context.Context, lockStrength LockingSt
 
 func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error) {
 	var peer nbpeer.Peer
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).First(&peer, keyQueryCondition, peerKey)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&peer, keyQueryCondition, peerKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "peer not found")
@@ -828,7 +838,7 @@ func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength Locking
 
 func (s *SqlStore) GetAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*Settings, error) {
 	var accountSettings AccountSettings
-	if err := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Where(idQueryCondition, accountID).First(&accountSettings).Error; err != nil {
+	if err := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Where(idQueryCondition, accountID).First(&accountSettings).Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "settings not found")
 		}
@@ -840,12 +850,12 @@ func (s *SqlStore) GetAccountSettings(ctx context.Context, lockStrength LockingS
 // SaveUserLastLogin stores the last login time for a user in DB.
 func (s *SqlStore) SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error {
 	var user User
-
-	result := s.db.WithContext(ctx).First(&user, accountAndIDQueryCondition, accountID, userID)
+	result := s.db.First(&user, accountAndIDQueryCondition, accountID, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.NewUserNotFoundError(userID)
 		}
+
 		return status.NewGetUserFromStoreError()
 	}
 	user.LastLogin = lastLogin
@@ -972,7 +982,7 @@ func NewPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore,
 
 func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error) {
 	var setupKey SetupKey
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		First(&setupKey, keyQueryCondition, key)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -984,7 +994,7 @@ func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength Locking
 }
 
 func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error {
-	result := s.db.WithContext(ctx).Model(&SetupKey{}).
+	result := s.db.Model(&SetupKey{}).
 		Where(idQueryCondition, setupKeyID).
 		Updates(map[string]interface{}{
 			"used_times": gorm.Expr("used_times + 1"),
@@ -1004,8 +1014,7 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string
 
 func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error {
 	var group nbgroup.Group
-
-	result := s.db.WithContext(ctx).Where("account_id = ? AND name = ?", accountID, "All").First(&group)
+	result := s.db.Where("account_id = ? AND name = ?", accountID, "All").First(&group)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.Errorf(status.NotFound, "group 'All' not found for account")
@@ -1030,12 +1039,12 @@ func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peer
 
 func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountId string, peerId string, groupID string) error {
 	var group nbgroup.Group
-
-	result := s.db.WithContext(ctx).Where(accountAndIDQueryCondition, accountId, groupID).First(&group)
+	result := s.db.Where(accountAndIDQueryCondition, accountId, groupID).First(&group)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.Errorf(status.NotFound, "group not found for account")
 		}
+
 		return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
 	}
 
@@ -1056,11 +1065,11 @@ func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountId string, peerId
 
 // GetUserPeers retrieves peers for a user.
 func (s *SqlStore) GetUserPeers(ctx context.Context, lockStrength LockingStrength, accountID, userID string) ([]*nbpeer.Peer, error) {
-	return getRecords[*nbpeer.Peer](s.db.WithContext(ctx).Where("user_id = ?", userID), lockStrength, accountID)
+	return getRecords[*nbpeer.Peer](s.db.Where("user_id = ?", userID), lockStrength, accountID)
 }
 
 func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error {
-	if err := s.db.WithContext(ctx).Create(peer).Error; err != nil {
+	if err := s.db.Create(peer).Error; err != nil {
 		return status.Errorf(status.Internal, "issue adding peer to account: %s", err)
 	}
 
@@ -1068,7 +1077,7 @@ func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) erro
 }
 
 func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, accountId string) error {
-	result := s.db.WithContext(ctx).Model(&Account{}).Where(idQueryCondition, accountId).Update("network_serial", gorm.Expr("network_serial + 1"))
+	result := s.db.Model(&Account{}).Where(idQueryCondition, accountId).Update("network_serial", gorm.Expr("network_serial + 1"))
 	if result.Error != nil {
 		return status.Errorf(status.Internal, "issue incrementing network serial count: %s", result.Error)
 	}
@@ -1076,7 +1085,7 @@ func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, accountId string)
 }
 
 func (s *SqlStore) ExecuteInTransaction(ctx context.Context, operation func(store Store) error) error {
-	tx := s.db.WithContext(ctx).Begin()
+	tx := s.db.Begin()
 	if tx.Error != nil {
 		return tx.Error
 	}
@@ -1101,8 +1110,7 @@ func (s *SqlStore) GetDB() *gorm.DB {
 
 func (s *SqlStore) GetAccountDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*DNSSettings, error) {
 	var accountDNSSettings AccountDNSSettings
-
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
 		First(&accountDNSSettings, idQueryCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -1116,8 +1124,7 @@ func (s *SqlStore) GetAccountDNSSettings(ctx context.Context, lockStrength Locki
 // AccountExists checks whether an account exists by the given ID.
 func (s *SqlStore) AccountExists(ctx context.Context, lockStrength LockingStrength, id string) (bool, error) {
 	var accountID string
-
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
 		Select("id").First(&accountID, idQueryCondition, id)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -1132,8 +1139,7 @@ func (s *SqlStore) AccountExists(ctx context.Context, lockStrength LockingStreng
 // GetAccountDomainAndCategory retrieves the Domain and DomainCategory fields for an account based on the given accountID.
 func (s *SqlStore) GetAccountDomainAndCategory(ctx context.Context, lockStrength LockingStrength, accountID string) (string, string, error) {
 	var account Account
-
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("domain", "domain_category").
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("domain", "domain_category").
 		Where(idQueryCondition, accountID).First(&account)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -1147,7 +1153,7 @@ func (s *SqlStore) GetAccountDomainAndCategory(ctx context.Context, lockStrength
 
 // GetGroupByID retrieves a group by ID and account ID.
 func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error) {
-	return getRecordByID[nbgroup.Group](s.db.WithContext(ctx).Preload(clause.Associations), lockStrength, groupID, accountID)
+	return getRecordByID[nbgroup.Group](s.db.Preload(clause.Associations), lockStrength, groupID, accountID)
 }
 
 // GetGroupByName retrieves a group by name and account ID.
@@ -1156,7 +1162,7 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 
 	// TODO: This fix is accepted for now, but if we need to handle this more frequently
 	// we may need to reconsider changing the types.
-	query := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations)
+	query := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations)
 	if s.storeEngine == PostgresStoreEngine {
 		query = query.Order("json_array_length(peers::json) DESC")
 	} else {
@@ -1175,7 +1181,7 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 
 // SaveGroup saves a group to the store.
 func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error {
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
 	if result.Error != nil {
 		return status.Errorf(status.Internal, "failed to save group to store: %v", result.Error)
 	}
@@ -1184,56 +1190,56 @@ func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength,
 
 // GetAccountPolicies retrieves policies for an account.
 func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) {
-	return getRecords[*Policy](s.db.WithContext(ctx).Preload(clause.Associations), lockStrength, accountID)
+	return getRecords[*Policy](s.db.Preload(clause.Associations), lockStrength, accountID)
 }
 
 // GetPolicyByID retrieves a policy by its ID and account ID.
 func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, policyID string, accountID string) (*Policy, error) {
-	return getRecordByID[Policy](s.db.WithContext(ctx).Preload(clause.Associations), lockStrength, policyID, accountID)
+	return getRecordByID[Policy](s.db.Preload(clause.Associations), lockStrength, policyID, accountID)
 }
 
 // GetAccountPostureChecks retrieves posture checks for an account.
 func (s *SqlStore) GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error) {
-	return getRecords[*posture.Checks](s.db.WithContext(ctx), lockStrength, accountID)
+	return getRecords[*posture.Checks](s.db, lockStrength, accountID)
 }
 
 // GetPostureChecksByID retrieves posture checks by their ID and account ID.
 func (s *SqlStore) GetPostureChecksByID(ctx context.Context, lockStrength LockingStrength, postureCheckID string, accountID string) (*posture.Checks, error) {
-	return getRecordByID[posture.Checks](s.db.WithContext(ctx), lockStrength, postureCheckID, accountID)
+	return getRecordByID[posture.Checks](s.db, lockStrength, postureCheckID, accountID)
 }
 
 // GetAccountRoutes retrieves network routes for an account.
 func (s *SqlStore) GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error) {
-	return getRecords[*route.Route](s.db.WithContext(ctx), lockStrength, accountID)
+	return getRecords[*route.Route](s.db, lockStrength, accountID)
 }
 
 // GetRouteByID retrieves a route by its ID and account ID.
 func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrength, routeID string, accountID string) (*route.Route, error) {
-	return getRecordByID[route.Route](s.db.WithContext(ctx), lockStrength, routeID, accountID)
+	return getRecordByID[route.Route](s.db, lockStrength, routeID, accountID)
 }
 
 // GetAccountSetupKeys retrieves setup keys for an account.
 func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error) {
-	return getRecords[*SetupKey](s.db.WithContext(ctx), lockStrength, accountID)
+	return getRecords[*SetupKey](s.db, lockStrength, accountID)
 }
 
 // GetSetupKeyByID retrieves a setup key by its ID and account ID.
 func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error) {
-	return getRecordByID[SetupKey](s.db.WithContext(ctx), lockStrength, setupKeyID, accountID)
+	return getRecordByID[SetupKey](s.db, lockStrength, setupKeyID, accountID)
 }
 
 // GetAccountNameServerGroups retrieves name server groups for an account.
 func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
-	return getRecords[*nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, accountID)
+	return getRecords[*nbdns.NameServerGroup](s.db, lockStrength, accountID)
 }
 
 // GetNameServerGroupByID retrieves a name server group by its ID and account ID.
 func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, nsGroupID string, accountID string) (*nbdns.NameServerGroup, error) {
-	return getRecordByID[nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, nsGroupID, accountID)
+	return getRecordByID[nbdns.NameServerGroup](s.db, lockStrength, nsGroupID, accountID)
 }
 
 func (s *SqlStore) DeleteSetupKey(ctx context.Context, accountID, keyID string) error {
-	return deleteRecordByID[SetupKey](s.db.WithContext(ctx), LockingStrengthUpdate, keyID, accountID)
+	return deleteRecordByID[SetupKey](s.db, LockingStrengthUpdate, keyID, accountID)
 }
 
 // getRecords retrieves records from the database based on the account ID.

management/server/status/error.go

@@ -3,6 +3,7 @@ package status
 import (
 	"errors"
 	"fmt"
+	"time"
 )
 
 const (
@@ -115,6 +116,11 @@ func NewGetUserFromStoreError() error {
 	return Errorf(Internal, "issue getting user from store")
 }
 
+// NewStoreContextCanceledError creates a new Error with Internal type for a canceled store context
+func NewStoreContextCanceledError(duration time.Duration) error {
+	return Errorf(Internal, "store access: context canceled after %v", duration)
+}
+
 // NewInvalidKeyIDError creates a new Error with InvalidArgument type for an issue getting a setup key
 func NewInvalidKeyIDError() error {
 	return Errorf(InvalidArgument, "invalid key ID")

management/server/updatechannel.go

@@ -2,13 +2,9 @@ package server
 
 import (
 	"context"
-	"fmt"
-	"runtime/debug"
 	"sync"
 	"time"
 
-	"github.com/netbirdio/netbird/management/server/differs"
-	"github.com/r3labs/diff/v3"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/proto"
@@ -25,8 +21,6 @@ type UpdateMessage struct {
 type PeersUpdateManager struct {
 	// peerChannels is an update channel indexed by Peer.ID
 	peerChannels map[string]chan *UpdateMessage
-	// peerNetworkMaps is the UpdateMessage indexed by Peer.ID.
-	peerUpdateMessage map[string]*UpdateMessage
 	// channelsMux keeps the mutex to access peerChannels
 	channelsMux *sync.RWMutex
 	// metrics provides method to collect application metrics
@@ -36,10 +30,9 @@ type PeersUpdateManager struct {
 // NewPeersUpdateManager returns a new instance of PeersUpdateManager
 func NewPeersUpdateManager(metrics telemetry.AppMetrics) *PeersUpdateManager {
 	return &PeersUpdateManager{
-		peerChannels:      make(map[string]chan *UpdateMessage),
-		peerUpdateMessage: make(map[string]*UpdateMessage),
-		channelsMux:       &sync.RWMutex{},
-		metrics:           metrics,
+		peerChannels: make(map[string]chan *UpdateMessage),
+		channelsMux:  &sync.RWMutex{},
+		metrics:      metrics,
 	}
 }
 
@@ -48,15 +41,6 @@ func (p *PeersUpdateManager) SendUpdate(ctx context.Context, peerID string, upda
 	start := time.Now()
 	var found, dropped bool
 
-	// skip sending sync update to the peer if there is no change in update message,
-	// it will not check on turn credential refresh as we do not send network map or client posture checks
-	if update.NetworkMap != nil {
-		updated := p.handlePeerMessageUpdate(ctx, peerID, update)
-		if !updated {
-			return
-		}
-	}
-
 	p.channelsMux.Lock()
 
 	defer func() {
@@ -66,16 +50,6 @@ func (p *PeersUpdateManager) SendUpdate(ctx context.Context, peerID string, upda
 		}
 	}()
 
-	if update.NetworkMap != nil {
-		lastSentUpdate := p.peerUpdateMessage[peerID]
-		if lastSentUpdate != nil && lastSentUpdate.Update.NetworkMap.GetSerial() > update.Update.NetworkMap.GetSerial() {
-			log.WithContext(ctx).Debugf("peer %s new network map serial: %d not greater than last sent: %d, skip sending update",
-				peerID, update.Update.NetworkMap.GetSerial(), lastSentUpdate.Update.NetworkMap.GetSerial())
-			return
-		}
-		p.peerUpdateMessage[peerID] = update
-	}
-
 	if channel, ok := p.peerChannels[peerID]; ok {
 		found = true
 		select {
@@ -108,7 +82,6 @@ func (p *PeersUpdateManager) CreateChannel(ctx context.Context, peerID string) c
 		closed = true
 		delete(p.peerChannels, peerID)
 		close(channel)
-		delete(p.peerUpdateMessage, peerID)
 	}
 	// mbragin: todo shouldn't it be more? or configurable?
 	channel := make(chan *UpdateMessage, channelBufferSize)
@@ -123,7 +96,6 @@ func (p *PeersUpdateManager) closeChannel(ctx context.Context, peerID string) {
 	if channel, ok := p.peerChannels[peerID]; ok {
 		delete(p.peerChannels, peerID)
 		close(channel)
-		delete(p.peerUpdateMessage, peerID)
 	}
 
 	log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
@@ -200,72 +172,3 @@ func (p *PeersUpdateManager) HasChannel(peerID string) bool {
 
 	return ok
 }
-
-// handlePeerMessageUpdate checks if the update message for a peer is new and should be sent.
-func (p *PeersUpdateManager) handlePeerMessageUpdate(ctx context.Context, peerID string, update *UpdateMessage) bool {
-	p.channelsMux.RLock()
-	lastSentUpdate := p.peerUpdateMessage[peerID]
-	p.channelsMux.RUnlock()
-
-	if lastSentUpdate != nil {
-		updated, err := isNewPeerUpdateMessage(ctx, lastSentUpdate, update)
-		if err != nil {
-			log.WithContext(ctx).Errorf("error checking for SyncResponse updates: %v", err)
-			return false
-		}
-		if !updated {
-			log.WithContext(ctx).Debugf("peer %s network map is not updated, skip sending update", peerID)
-			return false
-		}
-	}
-
-	return true
-}
-
-// isNewPeerUpdateMessage checks if the given current update message is a new update that should be sent.
-func isNewPeerUpdateMessage(ctx context.Context, lastSentUpdate, currUpdateToSend *UpdateMessage) (isNew bool, err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			log.WithContext(ctx).Panicf("comparing peer update messages. Trace: %s", debug.Stack())
-			isNew, err = true, nil
-		}
-	}()
-
-	if lastSentUpdate.Update.NetworkMap.GetSerial() > currUpdateToSend.Update.NetworkMap.GetSerial() {
-		return false, nil
-	}
-
-	differ, err := diff.NewDiffer(
-		diff.CustomValueDiffers(&differs.NetIPAddr{}),
-		diff.CustomValueDiffers(&differs.NetIPPrefix{}),
-	)
-	if err != nil {
-		return false, fmt.Errorf("failed to create differ: %v", err)
-	}
-
-	lastSentFiles := getChecksFiles(lastSentUpdate.Update.Checks)
-	currFiles := getChecksFiles(currUpdateToSend.Update.Checks)
-
-	changelog, err := differ.Diff(lastSentFiles, currFiles)
-	if err != nil {
-		return false, fmt.Errorf("failed to diff checks: %v", err)
-	}
-	if len(changelog) > 0 {
-		return true, nil
-	}
-
-	changelog, err = differ.Diff(lastSentUpdate.NetworkMap, currUpdateToSend.NetworkMap)
-	if err != nil {
-		return false, fmt.Errorf("failed to diff network map: %v", err)
-	}
-	return len(changelog) > 0, nil
-}
-
-// getChecksFiles returns a list of files from the given checks.
-func getChecksFiles(checks []*proto.Checks) []string {
-	files := make([]string, 0, len(checks))
-	for _, check := range checks {
-		files = append(files, check.GetFiles()...)
-	}
-	return files
-}

management/server/updatechannel_test.go

@@ -2,19 +2,10 @@ package server
 
 import (
 	"context"
-	"net"
-	"net/netip"
 	"testing"
 	"time"
 
-	nbdns "github.com/netbirdio/netbird/dns"
-	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/proto"
-	nbpeer "github.com/netbirdio/netbird/management/server/peer"
-	"github.com/netbirdio/netbird/management/server/posture"
-	nbroute "github.com/netbirdio/netbird/route"
-	"github.com/netbirdio/netbird/util"
-	"github.com/stretchr/testify/assert"
 )
 
 // var peersUpdater *PeersUpdateManager
@@ -86,470 +77,3 @@ func TestCloseChannel(t *testing.T) {
 		t.Error("Error closing the channel")
 	}
 }
-
-func TestHandlePeerMessageUpdate(t *testing.T) {
-	tests := []struct {
-		name           string
-		peerID         string
-		existingUpdate *UpdateMessage
-		newUpdate      *UpdateMessage
-		expectedResult bool
-	}{
-		{
-			name:   "update message with turn credentials update",
-			peerID: "peer",
-			newUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					WiretrusteeConfig: &proto.WiretrusteeConfig{},
-				},
-			},
-			expectedResult: true,
-		},
-		{
-			name:   "update message for peer without existing update",
-			peerID: "peer1",
-			newUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 1},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 2}},
-			},
-			expectedResult: true,
-		},
-		{
-			name:   "update message with no changes in update",
-			peerID: "peer2",
-			existingUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 1},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
-			},
-			newUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 1},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
-			},
-			expectedResult: false,
-		},
-		{
-			name:   "update message with changes in checks",
-			peerID: "peer3",
-			existingUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 1},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
-			},
-			newUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 2},
-					Checks: []*proto.Checks{
-						{
-							Files: []string{"/usr/bin/netbird"},
-						},
-					},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 2}},
-			},
-			expectedResult: true,
-		},
-		{
-			name:   "update message with lower serial number",
-			peerID: "peer4",
-			existingUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 2},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 2}},
-			},
-			newUpdate: &UpdateMessage{
-				Update: &proto.SyncResponse{
-					NetworkMap: &proto.NetworkMap{Serial: 1},
-				},
-				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
-			},
-			expectedResult: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			p := NewPeersUpdateManager(nil)
-			ctx := context.Background()
-
-			if tt.existingUpdate != nil {
-				p.peerUpdateMessage[tt.peerID] = tt.existingUpdate
-			}
-
-			result := p.handlePeerMessageUpdate(ctx, tt.peerID, tt.newUpdate)
-			assert.Equal(t, tt.expectedResult, result)
-		})
-	}
-}
-
-func TestIsNewPeerUpdateMessage(t *testing.T) {
-	t.Run("Unchanged value", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.False(t, message)
-	})
-
-	t.Run("Unchanged value with serial incremented", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.False(t, message)
-	})
-
-	t.Run("Updating routes network", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.Routes[0].Network = netip.MustParsePrefix("1.1.1.1/32")
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-
-	})
-
-	t.Run("Updating routes groups", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.Routes[0].Groups = []string{"randomGroup1"}
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating network map peers", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newPeer := &nbpeer.Peer{
-			IP:         net.ParseIP("192.168.1.4"),
-			SSHEnabled: true,
-			Key:        "peer4-key",
-			DNSLabel:   "peer4",
-			SSHKey:     "peer4-ssh-key",
-		}
-		newUpdateMessage2.NetworkMap.Peers = append(newUpdateMessage2.NetworkMap.Peers, newPeer)
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating process check", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-
-		newUpdateMessage2 := createMockUpdateMessage(t)
-		newUpdateMessage2.Update.NetworkMap.Serial++
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.False(t, message)
-
-		newUpdateMessage3 := createMockUpdateMessage(t)
-		newUpdateMessage3.Update.Checks = []*proto.Checks{}
-		newUpdateMessage3.Update.NetworkMap.Serial++
-		message, err = isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage3)
-		assert.NoError(t, err)
-		assert.True(t, message)
-
-		newUpdateMessage4 := createMockUpdateMessage(t)
-		check := &posture.Checks{
-			Checks: posture.ChecksDefinition{
-				ProcessCheck: &posture.ProcessCheck{
-					Processes: []posture.Process{
-						{
-							LinuxPath: "/usr/local/netbird",
-							MacPath:   "/usr/bin/netbird",
-						},
-					},
-				},
-			},
-		}
-		newUpdateMessage4.Update.Checks = []*proto.Checks{toProtocolCheck(check)}
-		newUpdateMessage4.Update.NetworkMap.Serial++
-		message, err = isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage4)
-		assert.NoError(t, err)
-		assert.True(t, message)
-
-		newUpdateMessage5 := createMockUpdateMessage(t)
-		check = &posture.Checks{
-			Checks: posture.ChecksDefinition{
-				ProcessCheck: &posture.ProcessCheck{
-					Processes: []posture.Process{
-						{
-							LinuxPath:   "/usr/bin/netbird",
-							WindowsPath: "C:\\Program Files\\netbird\\netbird.exe",
-							MacPath:     "/usr/local/netbird",
-						},
-					},
-				},
-			},
-		}
-		newUpdateMessage5.Update.Checks = []*proto.Checks{toProtocolCheck(check)}
-		newUpdateMessage5.Update.NetworkMap.Serial++
-		message, err = isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage5)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating DNS configuration", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newDomain := "newexample.com"
-		newUpdateMessage2.NetworkMap.DNSConfig.NameServerGroups[0].Domains = append(
-			newUpdateMessage2.NetworkMap.DNSConfig.NameServerGroups[0].Domains,
-			newDomain,
-		)
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating peer IP", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.Peers[0].IP = net.ParseIP("192.168.1.10")
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating firewall rule", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.FirewallRules[0].Port = "443"
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Add new firewall rule", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newRule := &FirewallRule{
-			PeerIP:    "192.168.1.3",
-			Direction: firewallRuleDirectionOUT,
-			Action:    string(PolicyTrafficActionDrop),
-			Protocol:  string(PolicyRuleProtocolUDP),
-			Port:      "53",
-		}
-		newUpdateMessage2.NetworkMap.FirewallRules = append(newUpdateMessage2.NetworkMap.FirewallRules, newRule)
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Removing nameserver", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.DNSConfig.NameServerGroups[0].NameServers = make([]nbdns.NameServer, 0)
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating name server IP", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.DNSConfig.NameServerGroups[0].NameServers[0].IP = netip.MustParseAddr("8.8.4.4")
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-	t.Run("Updating custom DNS zone", func(t *testing.T) {
-		newUpdateMessage1 := createMockUpdateMessage(t)
-		newUpdateMessage2 := createMockUpdateMessage(t)
-
-		newUpdateMessage2.NetworkMap.DNSConfig.CustomZones[0].Records[0].RData = "100.64.0.2"
-		newUpdateMessage2.Update.NetworkMap.Serial++
-
-		message, err := isNewPeerUpdateMessage(context.Background(), newUpdateMessage1, newUpdateMessage2)
-		assert.NoError(t, err)
-		assert.True(t, message)
-	})
-
-}
-
-func createMockUpdateMessage(t *testing.T) *UpdateMessage {
-	t.Helper()
-
-	_, ipNet, err := net.ParseCIDR("192.168.1.0/24")
-	if err != nil {
-		t.Fatal(err)
-	}
-	domainList, err := domain.FromStringList([]string{"example.com"})
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	config := &Config{
-		Signal: &Host{
-			Proto:    "https",
-			URI:      "signal.uri",
-			Username: "",
-			Password: "",
-		},
-		Stuns: []*Host{{URI: "stun.uri", Proto: UDP}},
-		TURNConfig: &TURNConfig{
-			Turns: []*Host{{URI: "turn.uri", Proto: UDP, Username: "turn-user", Password: "turn-pass"}},
-		},
-	}
-	peer := &nbpeer.Peer{
-		IP:         net.ParseIP("192.168.1.1"),
-		SSHEnabled: true,
-		Key:        "peer-key",
-		DNSLabel:   "peer1",
-		SSHKey:     "peer1-ssh-key",
-	}
-
-	secretManager := NewTimeBasedAuthSecretsManager(
-		NewPeersUpdateManager(nil),
-		&TURNConfig{
-			TimeBasedCredentials: false,
-			CredentialsTTL: util.Duration{
-				Duration: defaultDuration,
-			},
-			Secret: "secret",
-			Turns:  []*Host{TurnTestHost},
-		},
-		&Relay{
-			Addresses:      []string{"localhost:0"},
-			CredentialsTTL: util.Duration{Duration: time.Hour},
-			Secret:         "secret",
-		},
-	)
-
-	networkMap := &NetworkMap{
-		Network:      &Network{Net: *ipNet, Serial: 1000},
-		Peers:        []*nbpeer.Peer{{IP: net.ParseIP("192.168.1.2"), Key: "peer2-key", DNSLabel: "peer2", SSHEnabled: true, SSHKey: "peer2-ssh-key"}},
-		OfflinePeers: []*nbpeer.Peer{{IP: net.ParseIP("192.168.1.3"), Key: "peer3-key", DNSLabel: "peer3", SSHEnabled: true, SSHKey: "peer3-ssh-key"}},
-		Routes: []*nbroute.Route{
-			{
-				ID:          "route1",
-				Network:     netip.MustParsePrefix("10.0.0.0/24"),
-				KeepRoute:   true,
-				NetID:       "route1",
-				Peer:        "peer1",
-				NetworkType: 1,
-				Masquerade:  true,
-				Metric:      9999,
-				Enabled:     true,
-				Groups:      []string{"test1", "test2"},
-			},
-			{
-				ID:          "route2",
-				Domains:     domainList,
-				KeepRoute:   true,
-				NetID:       "route2",
-				Peer:        "peer1",
-				NetworkType: 1,
-				Masquerade:  true,
-				Metric:      9999,
-				Enabled:     true,
-				Groups:      []string{"test1", "test2"},
-			},
-		},
-		DNSConfig: nbdns.Config{
-			ServiceEnable: true,
-			NameServerGroups: []*nbdns.NameServerGroup{
-				{
-					NameServers: []nbdns.NameServer{{
-						IP:     netip.MustParseAddr("8.8.8.8"),
-						NSType: nbdns.UDPNameServerType,
-						Port:   nbdns.DefaultDNSPort,
-					}},
-					Primary:              true,
-					Domains:              []string{"example.com"},
-					Enabled:              true,
-					SearchDomainsEnabled: true,
-				},
-				{
-					ID: "ns1",
-					NameServers: []nbdns.NameServer{{
-						IP:     netip.MustParseAddr("1.1.1.1"),
-						NSType: nbdns.UDPNameServerType,
-						Port:   nbdns.DefaultDNSPort,
-					}},
-					Groups:               []string{"group1"},
-					Primary:              true,
-					Domains:              []string{"example.com"},
-					Enabled:              true,
-					SearchDomainsEnabled: true,
-				},
-			},
-			CustomZones: []nbdns.CustomZone{{Domain: "example.com", Records: []nbdns.SimpleRecord{{Name: "example.com", Type: 1, Class: "IN", TTL: 60, RData: "100.64.0.1"}}}},
-		},
-		FirewallRules: []*FirewallRule{
-			{PeerIP: "192.168.1.2", Direction: firewallRuleDirectionIN, Action: string(PolicyTrafficActionAccept), Protocol: string(PolicyRuleProtocolTCP), Port: "80"},
-		},
-	}
-	dnsName := "example.com"
-	checks := []*posture.Checks{
-		{
-			Checks: posture.ChecksDefinition{
-				ProcessCheck: &posture.ProcessCheck{
-					Processes: []posture.Process{
-						{
-							LinuxPath:   "/usr/bin/netbird",
-							WindowsPath: "C:\\Program Files\\netbird\\netbird.exe",
-							MacPath:     "/usr/bin/netbird",
-						},
-					},
-				},
-			},
-		},
-	}
-	dnsCache := &DNSConfigCache{}
-
-	turnToken, err := secretManager.GenerateTurnToken()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	relayToken, err := secretManager.GenerateRelayToken()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	return &UpdateMessage{
-		Update:     toSyncResponse(context.Background(), config, peer, turnToken, relayToken, networkMap, dnsName, checks, dnsCache),
-		NetworkMap: networkMap,
-	}
-}

relay/client/manager.go

@@ -16,6 +16,7 @@ import (
 
 var (
 	relayCleanupInterval = 60 * time.Second
+	keepUnusedServerTime = 5 * time.Second
 
 	ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
 )
@@ -27,10 +28,13 @@ type RelayTrack struct {
 	sync.RWMutex
 	relayClient *Client
 	err         error
+	created     time.Time
 }
 
 func NewRelayTrack() *RelayTrack {
-	return &RelayTrack{}
+	return &RelayTrack{
+		created: time.Now(),
+	}
 }
 
 type OnServerCloseListener func()
@@ -302,6 +306,18 @@ func (m *Manager) cleanUpUnusedRelays() {
 
 	for addr, rt := range m.relayClients {
 		rt.Lock()
+		// if the connection failed to the server the relay client will be nil
+		// but the instance will be kept in the relayClients until the next locking
+		if rt.err != nil {
+			rt.Unlock()
+			continue
+		}
+
+		if time.Since(rt.created) <= keepUnusedServerTime {
+			rt.Unlock()
+			continue
+		}
+
 		if rt.relayClient.HasConns() {
 			rt.Unlock()
 			continue

relay/client/manager_test.go

@@ -288,8 +288,9 @@ func TestForeginAutoClose(t *testing.T) {
 		t.Fatalf("failed to close connection: %s", err)
 	}
 
-	t.Logf("waiting for relay cleanup: %s", relayCleanupInterval+1*time.Second)
-	time.Sleep(relayCleanupInterval + 1*time.Second)
+	timeout := relayCleanupInterval + keepUnusedServerTime + 1*time.Second
+	t.Logf("waiting for relay cleanup: %s", timeout)
+	time.Sleep(timeout)
 	if len(mgr.relayClients) != 0 {
 		t.Errorf("expected 0, got %d", len(mgr.relayClients))
 	}

relay/client/picker_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"testing"
-	"time"
 )
 
 func TestServerPicker_UnavailableServers(t *testing.T) {
@@ -13,7 +12,7 @@ func TestServerPicker_UnavailableServers(t *testing.T) {
 		PeerID:     "test",
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
 	defer cancel()
 
 	go func() {

relay/server/listener/ws/listener.go

@@ -63,13 +63,14 @@ func (l *Listener) Shutdown(ctx context.Context) error {
 }
 
 func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
+	connRemoteAddr := remoteAddr(r)
 	wsConn, err := websocket.Accept(w, r, nil)
 	if err != nil {
-		log.Errorf("failed to accept ws connection from %s: %s", r.RemoteAddr, err)
+		log.Errorf("failed to accept ws connection from %s: %s", connRemoteAddr, err)
 		return
 	}
 
-	rAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
+	rAddr, err := net.ResolveTCPAddr("tcp", connRemoteAddr)
 	if err != nil {
 		err = wsConn.Close(websocket.StatusInternalError, "internal error")
 		if err != nil {
@@ -90,3 +91,10 @@ func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
 	conn := NewConn(wsConn, lAddr, rAddr)
 	l.acceptFn(conn)
 }
+
+func remoteAddr(r *http.Request) string {
+	if r.Header.Get("X-Real-Ip") == "" || r.Header.Get("X-Real-Port") == "" {
+		return r.RemoteAddr
+	}
+	return fmt.Sprintf("%s:%s", r.Header.Get("X-Real-Ip"), r.Header.Get("X-Real-Port"))
+}

util/net/net.go

@@ -11,8 +11,11 @@ import (
 
 const (
 	// NetbirdFwmark is the fwmark value used by Netbird via wireguard
-	NetbirdFwmark    = 0x1BD00
-	PreroutingFwmark = 0x1BD01
+	NetbirdFwmark = 0x1BD00
+
+	PreroutingFwmarkRedirected       = 0x1BD01
+	PreroutingFwmarkMasquerade       = 0x1BD11
+	PreroutingFwmarkMasqueradeReturn = 0x1BD12
 
 	envDisableCustomRouting = "NB_DISABLE_CUSTOM_ROUTING"
 )