Refactor the worker callbacks to channel

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

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

client/anonymize/anonymize.go

@@ -201,8 +201,6 @@ 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/service.go

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

client/cmd/service_controller.go

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

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.PreroutingFwmarkRedirected), "-j", chainNameInputRules},
+			spec:     []string{"-m", "mark", "--mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark), "-j", chainNameInputRules},
 			position: 2,
 		},
 	}
 
 	m.optionalEntries["PREROUTING"] = []entry{
 		{
-			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmarkRedirected)},
+			spec:     []string{"-t", "mangle", "-i", m.wgIface.Name(), "-m", "addrtype", "--dst-type", "LOCAL", "-j", "MARK", "--set-mark", fmt.Sprintf("%#x", nbnet.PreroutingFwmark)},
 			position: 1,
 		},
 	}

client/firewall/iptables/router_linux.go

@@ -18,24 +18,22 @@ import (
 	"github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
-	nbnet "github.com/netbirdio/netbird/util/net"
+)
+
+const (
+	ipv4Nat = "netbird-rt-nat"
 )
 
 // 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"
 )
 
@@ -325,25 +323,24 @@ func (r *router) Reset() error {
 }
 
 func (r *router) cleanUpDefaultForwardRules() error {
-	if err := r.cleanJumpRules(); err != nil {
-		return fmt.Errorf("clean jump rules: %w", err)
+	err := r.cleanJumpRules()
+	if err != nil {
+		return err
 	}
 
 	log.Debug("flushing routing related tables")
-	for _, chainInfo := range []struct {
-		chain string
-		table string
-	}{
-		{chainRTFWD, tableFilter},
-		{chainRTNAT, tableNat},
-		{chainRTPRE, tableMangle},
-	} {
-		ok, err := r.iptablesClient.ChainExists(chainInfo.table, chainInfo.chain)
+	for _, chain := range []string{chainRTFWD, chainRTNAT} {
+		table := r.getTableForChain(chain)
+
+		ok, err := r.iptablesClient.ChainExists(table, chain)
 		if err != nil {
-			return fmt.Errorf("check chain %s in table %s: %w", chainInfo.chain, chainInfo.table, err)
+			log.Errorf("failed check chain %s, error: %v", chain, err)
+			return err
 		} else if ok {
-			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)
+			err = r.iptablesClient.ClearAndDeleteChain(table, chain)
+			if err != nil {
+				log.Errorf("failed cleaning chain %s, error: %v", chain, err)
+				return err
 			}
 		}
 	}
@@ -352,16 +349,9 @@ func (r *router) cleanUpDefaultForwardRules() error {
 }
 
 func (r *router) createContainers() error {
-	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)
+	for _, chain := range []string{chainRTFWD, chainRTNAT} {
+		if err := r.createAndSetupChain(chain); err != nil {
+			return fmt.Errorf("create chain %s: %w", chain, err)
 		}
 	}
 
@@ -369,10 +359,6 @@ 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)
 	}
@@ -380,32 +366,6 @@ 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)
 
@@ -417,14 +377,10 @@ func (r *router) createAndSetupChain(chain string) error {
 }
 
 func (r *router) getTableForChain(chain string) string {
-	switch chain {
-	case chainRTNAT:
+	if chain == chainRTNAT {
 		return tableNat
-	case chainRTPRE:
-		return tableMangle
-	default:
-		return tableFilter
 	}
+	return tableFilter
 }
 
 func (r *router) insertEstablishedRule(chain string) error {
@@ -442,39 +398,25 @@ func (r *router) insertEstablishedRule(chain string) error {
 }
 
 func (r *router) addJumpRules() error {
-	// 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)
+	rule := []string{"-j", chainRTNAT}
+	err := r.iptablesClient.Insert(tableNat, chainPOSTROUTING, 1, rule...)
+	if err != nil {
+		return err
 	}
-	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
+	r.rules[ipv4Nat] = rule
 
 	return nil
 }
 
 func (r *router) cleanJumpRules() error {
-	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)
+	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)
 		}
 	}
+
 	return nil
 }
 
@@ -482,35 +424,19 @@ 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(tableMangle, chainRTPRE, rule...); err != nil {
-			return fmt.Errorf("error while removing existing marking rule for %s: %v", pair.Destination, err)
+		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)
 		}
 		delete(r.rules, ruleKey)
 	}
 
-	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)
+	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)
 	}
 
 	r.rules[ruleKey] = rule
+
 	return nil
 }
 
@@ -518,12 +444,13 @@ 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(tableMangle, chainRTPRE, rule...); err != nil {
-			return fmt.Errorf("error while removing marking rule for %s: %v", pair.Destination, err)
+		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)
 		}
+
 		delete(r.rules, ruleKey)
 	} else {
-		log.Debugf("marking rule %s not found", ruleKey)
+		log.Debugf("nat rule %s not found", ruleKey)
 	}
 
 	return nil
@@ -555,6 +482,16 @@ 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,18 +3,17 @@
 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 {
@@ -35,24 +34,14 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 	require.NoError(t, manager.init(nil))
 
 	defer func() {
-		assert.NoError(t, manager.Reset(), "shouldn't return error")
+		_ = manager.Reset()
 	}()
 
-	// 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")
+	require.Len(t, manager.rules, 2, "should have created rules map")
 
-	exists, err := manager.iptablesClient.Exists(tableNat, chainPOSTROUTING, "-j", chainRTNAT)
+	exists, err := manager.iptablesClient.Exists(tableNat, chainPOSTROUTING, manager.rules[ipv4Nat]...)
 	require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainPOSTROUTING)
-	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")
+	require.True(t, exists, "postrouting rule should exist")
 
 	pair := firewall.RouterPair{
 		ID:          "abc",
@@ -60,15 +49,22 @@ 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.AddNatRule(pair)
-	require.NoError(t, err, "adding NAT rule should not return error")
+	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.Reset()
 	require.NoError(t, err, "shouldn't return error")
 }
 
 func TestIptablesManager_AddNatRule(t *testing.T) {
+
 	if !isIptablesSupported() {
 		t.SkipNow()
 	}
@@ -83,66 +79,52 @@ func TestIptablesManager_AddNatRule(t *testing.T) {
 			require.NoError(t, manager.init(nil))
 
 			defer func() {
-				assert.NoError(t, manager.Reset(), "shouldn't return error")
+				err := manager.Reset()
+				if err != nil {
+					log.Errorf("failed to reset iptables manager: %s", err)
+				}
 			}()
 
 			err = manager.AddNatRule(testCase.InputPair)
-			require.NoError(t, err, "marking rule should be inserted")
+			require.NoError(t, err, "forwarding pair should be inserted")
 
 			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),
-			}
+			natRule := genRuleSpec(routingFinalNatJump, testCase.InputPair.Source, testCase.InputPair.Destination, ifaceMock.Name(), false)
 
-			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)
+			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, "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")
+				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, "marking rule should not be created")
-				_, found := manager.rules[natRuleKey]
-				require.False(t, found, "marking rule should not exist in the map")
+				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")
 			}
 
-			// 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),
-			}
+			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(tableMangle, chainRTPRE, inverseMarkingRule...)
-			require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableMangle, chainRTPRE)
+			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)
 			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")
+				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")
 			} 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")
+				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")
 			}
 		})
 	}
 }
 
 func TestIptablesManager_RemoveNatRule(t *testing.T) {
+
 	if !isIptablesSupported() {
 		t.SkipNow()
 	}
@@ -155,52 +137,42 @@ func TestIptablesManager_RemoveNatRule(t *testing.T) {
 			require.NoError(t, err, "shouldn't return error")
 			require.NoError(t, manager.init(nil))
 			defer func() {
-				assert.NoError(t, manager.Reset(), "shouldn't return error")
+				_ = manager.Reset()
 			}()
 
-			err = manager.AddNatRule(testCase.InputPair)
-			require.NoError(t, err, "should add NAT rule without 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.RemoveNatRule(testCase.InputPair)
 			require.NoError(t, err, "shouldn't return error")
 
-			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")
+			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")
 
 			_, found := manager.rules[natRuleKey]
-			require.False(t, found, "marking rule should not exist in the manager map")
+			require.False(t, found, "nat rule should exist in the manager map")
 
-			// 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),
-			}
+			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")
 
-			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")
+			_, found = manager.rules[inNatRuleKey]
+			require.False(t, found, "income nat rule should exist in the manager map")
 		})
 	}
 }

client/firewall/manager/firewall.go

@@ -17,7 +17,6 @@ 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.PreroutingFwmarkRedirected),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
 			},
 			&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.PreroutingFwmarkRedirected),
+				Data:     binaryutil.NativeEndian.PutUint32(nbnet.PreroutingFwmark),
 			},
 			&expr.Verdict{
 				Kind:  expr.VerdictJump,

client/firewall/nftables/router_linux.go

@@ -21,7 +21,6 @@ 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 (
@@ -125,6 +124,7 @@ 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,21 +133,6 @@ 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)
 	}
@@ -437,149 +422,59 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 	sourceExp := generateCIDRMatcherExpressions(true, pair.Source)
 	destExp := generateCIDRMatcherExpressions(false, pair.Destination)
 
-	op := expr.CmpOpEq
+	dir := expr.MetaKeyIIFNAME
+	notDir := expr.MetaKeyOIFNAME
 	if pair.Inverse {
-		op = expr.CmpOpNeq
+		dir = expr.MetaKeyOIFNAME
+		notDir = expr.MetaKeyIIFNAME
 	}
 
+	lo := ifname("lo")
+	intf := ifname(r.wgIface.Name())
+
 	exprs := []expr.Any{
-		// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
-		// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
-		&expr.Ct{
-			Key:      expr.CtKeySTATE,
+		&expr.Meta{
+			Key:      dir,
 			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.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.Cmp{
 			Op:       expr.CmpOpNeq,
 			Register: 1,
-			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()),
+			Data:     lo,
 		},
 	}
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
-
-	var markValue uint32 = nbnet.PreroutingFwmarkMasquerade
-	if pair.Inverse {
-		markValue = nbnet.PreroutingFwmarkMasqueradeReturn
-	}
-
 	exprs = append(exprs,
-		&expr.Immediate{
-			Register: 1,
-			Data:     binaryutil.NativeEndian.PutUint32(markValue),
-		},
-		&expr.Meta{
-			Key:            expr.MetaKeyMARK,
-			SourceRegister: true,
-			Register:       1,
-		},
+		&expr.Counter{}, &expr.Masq{},
 	)
 
-	ruleKey := firewall.GenKey(firewall.PreroutingFormat, pair)
+	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
 
 	if _, exists := r.rules[ruleKey]; exists {
 		if err := r.removeNatRule(pair); err != nil {
-			return fmt.Errorf("remove prerouting rule: %w", err)
+			return fmt.Errorf("remove routing rule: %w", err)
 		}
 	}
 
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
-		Chain:    r.chains[chainNamePrerouting],
+		Chain:    r.chains[chainNameRoutingNat],
 		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
 }
 
@@ -828,18 +723,18 @@ func (r *router) removeAcceptForwardRulesIptables(ipt *iptables.IPTables) error
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-// RemoveNatRule removes the prerouting mark rule
+// RemoveNatRule removes a nftables rule pair from nat chains
 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 prerouting rule: %w", err)
+		return fmt.Errorf("remove nat rule: %w", err)
 	}
 
 	if err := r.removeNatRule(firewall.GetInversePair(pair)); err != nil {
-		return fmt.Errorf("remove inverse prerouting rule: %w", err)
+		return fmt.Errorf("remove inverse nat rule: %w", err)
 	}
 
 	if err := r.removeLegacyRouteRule(pair); err != nil {
@@ -854,20 +749,21 @@ 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.PreroutingFormat, pair)
+	ruleKey := firewall.GenKey(firewall.NatFormat, pair)
 
 	if rule, exists := r.rules[ruleKey]; exists {
 		err := r.conn.DelRule(rule)
 		if err != nil {
-			return fmt.Errorf("remove prerouting rule %s -> %s: %v", pair.Source, pair.Destination, err)
+			return fmt.Errorf("remove nat rule %s -> %s: %v", pair.Source, pair.Destination, err)
 		}
 
-		log.Debugf("nftables: removed prerouting rule %s -> %s", pair.Source, pair.Destination)
+		log.Debugf("nftables: removed nat rule %s -> %s", pair.Source, pair.Destination)
 
 		delete(r.rules, ruleKey)
 	} else {
-		log.Debugf("nftables: prerouting rule %s not found", ruleKey)
+		log.Debugf("nftables: nat rule %s not found", ruleKey)
 	}
 
 	return nil

client/firewall/nftables/router_linux_test.go

@@ -10,7 +10,6 @@ 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"
@@ -33,87 +32,100 @@ 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) {
-			// 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))
+			manager, err := newRouter(table, ifaceMock)
+			require.NoError(t, err, "failed to create router")
+			require.NoError(t, manager.init(table))
 
 			nftablesTestingClient := &nftables.Conn{}
 
-			rtr := manager.router
-			err = rtr.AddNatRule(testCase.InputPair)
+			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)
 			require.NoError(t, err, "pair should be inserted")
 
-			t.Cleanup(func() {
-				require.NoError(t, rtr.RemoveNatRule(testCase.InputPair), "failed to remove rule")
-			})
+			defer func(manager *router, pair firewall.RouterPair) {
+				require.NoError(t, manager.RemoveNatRule(pair), "failed to remove rule")
+			}(manager, testCase.InputPair)
 
 			if testCase.InputPair.Masquerade {
-				// 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,
-					},
+				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},
 					&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"),
+					},
+				)
 
-				// Build CIDR matching expressions
-				sourceExp := generateCIDRMatcherExpressions(true, testCase.InputPair.Source)
-				destExp := generateCIDRMatcherExpressions(false, testCase.InputPair.Destination)
-
-				// 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)
+				natRuleKey := firewall.GenKey(firewall.NatFormat, testCase.InputPair)
 				found := 0
-				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
-							}
+				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 in prerouting chain")
+				require.Equal(t, 1, found, "should find at least 1 rule to test")
 			}
+
+			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.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))
+				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
+						}
+					}
+				}
+				require.Equal(t, 1, found, "should find at least 1 rule to test")
+			}
+
 		})
 	}
 }
@@ -123,66 +135,68 @@ 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 := Create(ifaceMock)
-			t.Cleanup(func() {
-				require.NoError(t, manager.Reset(nil))
+			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),
 			})
-			require.NoError(t, err)
-			require.NoError(t, manager.Init(nil))
 
-			rtr := manager.router
+			sourceExp = generateCIDRMatcherExpressions(true, firewall.GetInversePair(testCase.InputPair).Source)
+			destExp = generateCIDRMatcherExpressions(false, firewall.GetInversePair(testCase.InputPair).Destination)
 
-			// First add the NAT rule using the router's method
-			err = rtr.AddNatRule(testCase.InputPair)
-			require.NoError(t, err, "should add NAT rule")
+			natExp = append(sourceExp, append(destExp, &expr.Counter{}, &expr.Masq{})...) //nolint:gocritic
+			inNatRuleKey := firewall.GenKey(firewall.NatFormat, firewall.GetInversePair(testCase.InputPair))
 
-			// 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")
+			insertedInNat := nftablesTestingClient.InsertRule(&nftables.Rule{
+				Table:    manager.workTable,
+				Chain:    manager.chains[chainNameRoutingNat],
+				Exprs:    natExp,
+				UserData: []byte(inNatRuleKey),
+			})
 
-			// Now remove the rule
-			err = rtr.RemoveNatRule(testCase.InputPair)
-			require.NoError(t, err, "shouldn't return error when removing rule")
+			err = nftablesTestingClient.Flush()
+			require.NoError(t, err, "shouldn't return error")
 
-			// 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")
+			err = manager.Reset()
+			require.NoError(t, err, "shouldn't return error")
 
-			// 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
+			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")
 					}
 				}
-				if foundCounter {
-					break
-				}
 			}
-			require.True(t, foundCounter, "static postrouting rule should remain")
 		})
 	}
 }

client/iface/bind/ice_bind.go

@@ -12,7 +12,6 @@ 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"
 )
 
@@ -25,8 +24,8 @@ type receiverCreator struct {
 	iceBind *ICEBind
 }
 
-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)
+func (rc receiverCreator) CreateIPv4ReceiverFn(msgPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) wgConn.ReceiveFunc {
+	return rc.iceBind.createIPv4ReceiverFn(msgPool, pc, conn)
 }
 
 // ICEBind is a bind implementation with two main features:
@@ -155,7 +154,7 @@ func (b *ICEBind) Send(bufs [][]byte, ep wgConn.Endpoint) error {
 	return nil
 }
 
-func (s *ICEBind) createIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool, msgsPool *sync.Pool) wgConn.ReceiveFunc {
+func (s *ICEBind) createIPv4ReceiverFn(ipv4MsgsPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) wgConn.ReceiveFunc {
 	s.muUDPMux.Lock()
 	defer s.muUDPMux.Unlock()
 
@@ -167,30 +166,16 @@ func (s *ICEBind) createIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, r
 		},
 	)
 	return func(bufs [][]byte, sizes []int, eps []wgConn.Endpoint) (n int, err error) {
-		msgs := getMessages(msgsPool)
+		msgs := ipv4MsgsPool.Get().(*[]ipv4.Message)
+		defer ipv4MsgsPool.Put(msgs)
 		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" || 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
-				}
+		if runtime.GOOS == "linux" {
+			numMsgs, err = pc.ReadBatch(*msgs, 0)
+			if err != nil {
+				return 0, err
 			}
 		} else {
 			msg := &(*msgs)[0]
@@ -206,12 +191,11 @@ func (s *ICEBind) createIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, r
 			// todo: handle err
 			ok, _ := s.filterOutStunMessages(msg.Buffers, msg.N, msg.Addr)
 			if ok {
-				continue
-			}
-			sizes[i] = msg.N
-			if sizes[i] == 0 {
-				continue
+				sizes[i] = 0
+			} else {
+				sizes[i] = msg.N
 			}
+
 			addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
 			ep := &wgConn.StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation
 			wgConn.GetSrcFromControl(msg.OOB[:msg.NN], ep)
@@ -289,15 +273,3 @@ 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

@@ -2,7 +2,6 @@ package bind
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"net"
 	"net/netip"
@@ -95,10 +94,7 @@ func (p *ProxyBind) close() error {
 
 	p.Bind.RemoveEndpoint(p.wgAddr)
 
-	if rErr := p.remoteConn.Close(); rErr != nil && !errors.Is(rErr, net.ErrClosed) {
-		return rErr
-	}
-	return nil
+	return p.remoteConn.Close()
 }
 
 func (p *ProxyBind) proxyToLocal(ctx context.Context) {
@@ -108,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/iface/wgproxy/ebpf/wrapper.go

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

client/iface/wgproxy/udp/proxy.go

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

client/internal/connect.go

@@ -207,8 +207,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		c.statusRecorder.MarkSignalDisconnected(nil)
 		defer func() {
-			_, err := state.Status()
-			c.statusRecorder.MarkSignalDisconnected(err)
+			c.statusRecorder.MarkSignalDisconnected(state.err)
 		}()
 
 		// with the global Wiretrustee config in hand connect (just a connection, no stream yet) Signal

client/internal/peer/conn.go

@@ -134,36 +134,29 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 		statusICE:      NewAtomicConnStatus(),
 	}
 
-	rFns := WorkerRelayCallbacks{
-		OnConnReady:    conn.relayConnectionIsReady,
-		OnDisconnected: conn.onWorkerRelayStateDisconnected,
-	}
-
-	wFns := WorkerICECallbacks{
-		OnConnReady:     conn.iCEConnectionIsReady,
-		OnStatusChanged: conn.onWorkerICEStateDisconnected,
-	}
-
 	ctrl := isController(config)
-	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, relayManager, rFns)
+	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, relayManager)
 
 	relayIsSupportedLocally := conn.workerRelay.RelayIsSupportedLocally()
-	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally, wFns)
+	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally)
 	if err != nil {
 		return nil, err
 	}
 
 	conn.handshaker = NewHandshaker(ctx, connLog, config, signaler, conn.workerICE, conn.workerRelay)
 
-	conn.handshaker.AddOnNewOfferListener(conn.workerRelay.OnNewOffer)
+	conn.handshaker.AddOnNewOfferListener(func(remoteOfferAnswer *OfferAnswer) {
+		conn.workerRelay.OnNewOffer(ctx, remoteOfferAnswer)
+	})
 	if os.Getenv("NB_FORCE_RELAY") != "true" {
-		conn.handshaker.AddOnNewOfferListener(conn.workerICE.OnNewOffer)
+		conn.handshaker.AddOnNewOfferListener(func(remoteOfferAnswer *OfferAnswer) {
+			conn.workerICE.OnNewOffer(ctx, remoteOfferAnswer)
+		})
 	}
 
 	conn.guard = guard.NewGuard(connLog, ctrl, conn.isConnectedOnAllWay, config.Timeout, srWatcher)
 
 	go conn.handshaker.Listen()
-
 	return conn, nil
 }
 
@@ -190,6 +183,7 @@ func (conn *Conn) Open() {
 	}
 
 	go conn.startHandshakeAndReconnect(conn.ctx)
+	go conn.listenWorkersEvents()
 }
 
 func (conn *Conn) startHandshakeAndReconnect(ctx context.Context) {
@@ -301,7 +295,7 @@ func (conn *Conn) GetKey() string {
 }
 
 // configureConnection starts proxying traffic from/to local Wireguard and sets connection status to StatusConnected
-func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICEConnInfo) {
+func (conn *Conn) iCEConnectionIsReady(iceConnInfo ICEConnInfo) {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -309,14 +303,9 @@ 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 {
+	if conn.currentConnPriority > iceConnInfo.ConnPriority {
 		conn.statusICE.Set(StatusConnected)
 		conn.updateIceState(iceConnInfo)
 		return
@@ -338,7 +327,7 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		ep = wgProxy.EndpointAddr()
 		conn.wgProxyICE = wgProxy
 	} else {
-		directEp, err := net.ResolveUDPAddr("udp", iceConnInfo.RemoteConn.RemoteAddr().String())
+		directEp, err := net.ResolveUDPAddr("udp", iceConnInfo.RemoteIceCandidateEndpoint)
 		if err != nil {
 			log.Errorf("failed to resolveUDPaddr")
 			conn.handleConfigurationFailure(err, nil)
@@ -366,14 +355,13 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		return
 	}
 	wgConfigWorkaround()
-	conn.currentConnPriority = priority
+	conn.currentConnPriority = iceConnInfo.ConnPriority
 	conn.statusICE.Set(StatusConnected)
 	conn.updateIceState(iceConnInfo)
 	conn.doOnConnected(iceConnInfo.RosenpassPubKey, iceConnInfo.RosenpassAddr)
 }
 
-// todo review to make sense to handle connecting and disconnected status also?
-func (conn *Conn) onWorkerICEStateDisconnected(newState ConnStatus) {
+func (conn *Conn) onWorkerICEStateDisconnected() {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 
@@ -381,7 +369,7 @@ func (conn *Conn) onWorkerICEStateDisconnected(newState ConnStatus) {
 		return
 	}
 
-	conn.log.Tracef("ICE connection state changed to %s", newState)
+	conn.log.Tracef("ICE connection state changed to disconnected")
 
 	if conn.wgProxyICE != nil {
 		if err := conn.wgProxyICE.CloseConn(); err != nil {
@@ -401,8 +389,8 @@ func (conn *Conn) onWorkerICEStateDisconnected(newState ConnStatus) {
 		conn.currentConnPriority = connPriorityRelay
 	}
 
-	changed := conn.statusICE.Get() != newState && newState != StatusConnecting
-	conn.statusICE.Set(newState)
+	changed := conn.statusICE.Get() != stateDisconnected
+	conn.statusICE.Set(stateDisconnected)
 
 	conn.guard.SetICEConnDisconnected(changed)
 
@@ -442,7 +430,7 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 
 	if conn.iceP2PIsActive() {
 		conn.log.Debugf("do not switch to relay because current priority is: %v", conn.currentConnPriority)
-		conn.setRelayedProxy(wgProxy)
+		conn.wgProxyRelay = wgProxy
 		conn.statusRelay.Set(StatusConnected)
 		conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
 		return
@@ -465,7 +453,7 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 	wgConfigWorkaround()
 	conn.currentConnPriority = connPriorityRelay
 	conn.statusRelay.Set(StatusConnected)
-	conn.setRelayedProxy(wgProxy)
+	conn.wgProxyRelay = wgProxy
 	conn.updateRelayStatus(rci.relayedConn.RemoteAddr().String(), rci.rosenpassPubKey)
 	conn.log.Infof("start to communicate with peer via relay")
 	conn.doOnConnected(rci.rosenpassPubKey, rci.rosenpassAddr)
@@ -736,13 +724,31 @@ func (conn *Conn) logTraceConnState() {
 	}
 }
 
-func (conn *Conn) setRelayedProxy(proxy wgproxy.Proxy) {
-	if conn.wgProxyRelay != nil {
-		if err := conn.wgProxyRelay.CloseConn(); err != nil {
-			conn.log.Warnf("failed to close deprecated wg proxy conn: %v", err)
+func (conn *Conn) listenWorkersEvents() {
+	for {
+		select {
+		case e := <-conn.workerRelay.EventChan:
+			switch e.ConnStatus {
+			case StatusConnected:
+				conn.relayConnectionIsReady(e.RelayConnInfo)
+			case StatusDisconnected:
+				conn.onWorkerRelayStateDisconnected()
+			default:
+				log.Errorf("unexpected relay connection status: %v", e.ConnStatus)
+			}
+		case e := <-conn.workerICE.EventChan:
+			switch e.ConnStatus {
+			case StatusConnected:
+				conn.iCEConnectionIsReady(e.ICEConnInfo)
+			case StatusDisconnected:
+				conn.onWorkerICEStateDisconnected()
+			default:
+				log.Errorf("unexpected ICE connection status: %v", e.ConnStatus)
+			}
+		case <-conn.ctx.Done():
+			return
 		}
 	}
-	conn.wgProxyRelay = proxy
 }
 
 func isController(config ConnConfig) bool {

client/internal/peer/nilcheck.go

@@ -1,21 +0,0 @@
-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/internal/peer/status.go

@@ -67,7 +67,7 @@ func (s *State) DeleteRoute(network string) {
 func (s *State) GetRoutes() map[string]struct{} {
 	s.Mux.RLock()
 	defer s.Mux.RUnlock()
-	return maps.Clone(s.routes)
+	return s.routes
 }
 
 // LocalPeerState contains the latest state of the local peer
@@ -237,6 +237,10 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		peerState.IP = receivedState.IP
 	}
 
+	if receivedState.GetRoutes() != nil {
+		peerState.SetRoutes(receivedState.GetRoutes())
+	}
+
 	skipNotification := shouldSkipNotify(receivedState.ConnStatus, peerState)
 
 	if receivedState.ConnStatus != peerState.ConnStatus {
@@ -257,40 +261,12 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		return nil
 	}
 
-	d.notifyPeerListChanged()
-	return nil
-}
-
-func (d *Status) AddPeerStateRoute(peer string, route string) error {
-	d.mux.Lock()
-	defer d.mux.Unlock()
-
-	peerState, ok := d.peers[peer]
-	if !ok {
-		return errors.New("peer doesn't exist")
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
 	}
 
-	peerState.AddRoute(route)
-	d.peers[peer] = peerState
-
-	// todo: consider to make sense of this notification or not
-	d.notifyPeerListChanged()
-	return nil
-}
-
-func (d *Status) RemovePeerStateRoute(peer string, route string) error {
-	d.mux.Lock()
-	defer d.mux.Unlock()
-
-	peerState, ok := d.peers[peer]
-	if !ok {
-		return errors.New("peer doesn't exist")
-	}
-
-	peerState.DeleteRoute(route)
-	d.peers[peer] = peerState
-
-	// todo: consider to make sense of this notification or not
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -325,7 +301,12 @@ func (d *Status) UpdatePeerICEState(receivedState State) error {
 		return nil
 	}
 
-	d.notifyPeerStateChangeListeners(receivedState.PubKey)
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
+	}
+
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -353,7 +334,12 @@ func (d *Status) UpdatePeerRelayedState(receivedState State) error {
 		return nil
 	}
 
-	d.notifyPeerStateChangeListeners(receivedState.PubKey)
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
+	}
+
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -380,7 +366,12 @@ func (d *Status) UpdatePeerRelayedStateToDisconnected(receivedState State) error
 		return nil
 	}
 
-	d.notifyPeerStateChangeListeners(receivedState.PubKey)
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
+	}
+
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -410,7 +401,12 @@ func (d *Status) UpdatePeerICEStateToDisconnected(receivedState State) error {
 		return nil
 	}
 
-	d.notifyPeerStateChangeListeners(receivedState.PubKey)
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
+	}
+
 	d.notifyPeerListChanged()
 	return nil
 }
@@ -481,14 +477,11 @@ func (d *Status) FinishPeerListModifications() {
 func (d *Status) GetPeerStateChangeNotifier(peer string) <-chan struct{} {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-
 	ch, found := d.changeNotify[peer]
-	if found {
-		return ch
+	if !found || ch == nil {
+		ch = make(chan struct{})
+		d.changeNotify[peer] = ch
 	}
-
-	ch = make(chan struct{})
-	d.changeNotify[peer] = ch
 	return ch
 }
 
@@ -762,17 +755,6 @@ func (d *Status) onConnectionChanged() {
 	d.notifier.updateServerStates(d.managementState, d.signalState)
 }
 
-// notifyPeerStateChangeListeners notifies route manager about the change in peer state
-func (d *Status) notifyPeerStateChangeListeners(peerID string) {
-	ch, found := d.changeNotify[peerID]
-	if !found {
-		return
-	}
-
-	close(ch)
-	delete(d.changeNotify, peerID)
-}
-
 func (d *Status) notifyPeerListChanged() {
 	d.notifier.peerListChanged(d.numOfPeers())
 }

client/internal/peer/status_test.go

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

client/internal/peer/worker_ice.go

@@ -19,8 +19,14 @@ import (
 	"github.com/netbirdio/netbird/route"
 )
 
+type ICEEvent struct {
+	ConnStatus  ConnStatus
+	ICEConnInfo ICEConnInfo
+}
+
 type ICEConnInfo struct {
 	RemoteConn                 net.Conn
+	RemoteAddr                 net.Addr
 	RosenpassPubKey            []byte
 	RosenpassAddr              string
 	LocalIceCandidateType      string
@@ -29,14 +35,11 @@ type ICEConnInfo struct {
 	LocalIceCandidateEndpoint  string
 	Relayed                    bool
 	RelayedOnLocal             bool
-}
-
-type WorkerICECallbacks struct {
-	OnConnReady     func(ConnPriority, ICEConnInfo)
-	OnStatusChanged func(ConnStatus)
+	ConnPriority               ConnPriority
 }
 
 type WorkerICE struct {
+	EventChan         chan ICEEvent
 	ctx               context.Context
 	log               *log.Entry
 	config            ConnConfig
@@ -44,7 +47,6 @@ type WorkerICE struct {
 	iFaceDiscover     stdnet.ExternalIFaceDiscover
 	statusRecorder    *Status
 	hasRelayOnLocally bool
-	conn              WorkerICECallbacks
 
 	selectedPriority ConnPriority
 
@@ -57,13 +59,11 @@ type WorkerICE struct {
 
 	localUfrag string
 	localPwd   string
-
-	// we record the last known state of the ICE agent to avoid duplicate on disconnected events
-	lastKnownState ice.ConnectionState
 }
 
-func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool, callBacks WorkerICECallbacks) (*WorkerICE, error) {
+func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool) (*WorkerICE, error) {
 	w := &WorkerICE{
+		EventChan:         make(chan ICEEvent, 2),
 		ctx:               ctx,
 		log:               log,
 		config:            config,
@@ -71,7 +71,6 @@ func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signal
 		iFaceDiscover:     ifaceDiscover,
 		statusRecorder:    statusRecorder,
 		hasRelayOnLocally: hasRelayOnLocally,
-		conn:              callBacks,
 	}
 
 	localUfrag, localPwd, err := icemaker.GenerateICECredentials()
@@ -83,7 +82,7 @@ func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, signal
 	return w, nil
 }
 
-func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
+func (w *WorkerICE) OnNewOffer(_ context.Context, remoteOfferAnswer *OfferAnswer) {
 	w.log.Debugf("OnNewOffer for ICE")
 	w.muxAgent.Lock()
 
@@ -136,6 +135,11 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		return
 	}
 
+	if pair == nil {
+		w.log.Errorf("remote address is nil, ICE conn already closed")
+		return
+	}
+
 	if !isRelayCandidate(pair.Local) {
 		// dynamically set remote WireGuard port if other side specified a different one from the default one
 		remoteWgPort := iface.DefaultWgPort
@@ -157,9 +161,13 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		RemoteIceCandidateEndpoint: fmt.Sprintf("%s:%d", pair.Remote.Address(), pair.Remote.Port()),
 		Relayed:                    isRelayed(pair),
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
+		ConnPriority:               w.selectedPriority,
 	}
 	w.log.Debugf("on ICE conn read to use ready")
-	go w.conn.OnConnReady(w.selectedPriority, ci)
+	select {
+	case w.EventChan <- ICEEvent{ConnStatus: StatusConnected, ICEConnInfo: ci}:
+	case <-w.ctx.Done():
+	}
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -197,7 +205,8 @@ func (w *WorkerICE) Close() {
 		return
 	}
 
-	if err := w.agent.Close(); err != nil {
+	err := w.agent.Close()
+	if err != nil {
 		w.log.Warnf("failed to close ICE agent: %s", err)
 	}
 }
@@ -217,18 +226,18 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 
 	err = agent.OnConnectionStateChange(func(state ice.ConnectionState) {
 		w.log.Debugf("ICE ConnectionState has changed to %s", state.String())
-		switch state {
-		case ice.ConnectionStateConnected:
-			w.lastKnownState = ice.ConnectionStateConnected
-			return
-		case ice.ConnectionStateFailed, ice.ConnectionStateDisconnected:
-			if w.lastKnownState != ice.ConnectionStateDisconnected {
-				w.lastKnownState = ice.ConnectionStateDisconnected
-				w.conn.OnStatusChanged(StatusDisconnected)
+		if state == ice.ConnectionStateFailed || state == ice.ConnectionStateDisconnected {
+			select {
+			case w.EventChan <- ICEEvent{ConnStatus: StatusDisconnected}:
+			case <-w.ctx.Done():
 			}
-			w.closeAgent(agentCancel)
-		default:
-			return
+
+			w.muxAgent.Lock()
+			agentCancel()
+			_ = agent.Close()
+			w.agent = nil
+
+			w.muxAgent.Unlock()
 		}
 	})
 	if err != nil {
@@ -254,17 +263,6 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 	return agent, nil
 }
 
-func (w *WorkerICE) closeAgent(cancel context.CancelFunc) {
-	w.muxAgent.Lock()
-	defer w.muxAgent.Unlock()
-
-	cancel()
-	if err := w.agent.Close(); err != nil {
-		w.log.Warnf("failed to close ICE agent: %s", err)
-	}
-	w.agent = nil
-}
-
 func (w *WorkerICE) punchRemoteWGPort(pair *ice.CandidatePair, remoteWgPort int) {
 	// wait local endpoint configuration
 	time.Sleep(time.Second)

client/internal/peer/worker_relay.go

@@ -18,23 +18,23 @@ var (
 	wgHandshakeOvertime = 30 * time.Second
 )
 
+type RelayEvent struct {
+	ConnStatus    ConnStatus
+	RelayConnInfo RelayConnInfo
+}
+
 type RelayConnInfo struct {
 	relayedConn     net.Conn
 	rosenpassPubKey []byte
 	rosenpassAddr   string
 }
 
-type WorkerRelayCallbacks struct {
-	OnConnReady    func(RelayConnInfo)
-	OnDisconnected func()
-}
-
 type WorkerRelay struct {
+	EventChan    chan RelayEvent
 	log          *log.Entry
 	isController bool
 	config       ConnConfig
 	relayManager relayClient.ManagerService
-	callBacks    WorkerRelayCallbacks
 
 	relayedConn      net.Conn
 	relayLock        sync.Mutex
@@ -45,18 +45,18 @@ type WorkerRelay struct {
 	relaySupportedOnRemotePeer atomic.Bool
 }
 
-func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, relayManager relayClient.ManagerService, callbacks WorkerRelayCallbacks) *WorkerRelay {
+func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, relayManager relayClient.ManagerService) *WorkerRelay {
 	r := &WorkerRelay{
+		EventChan:    make(chan RelayEvent, 2),
 		log:          log,
 		isController: ctrl,
 		config:       config,
 		relayManager: relayManager,
-		callBacks:    callbacks,
 	}
 	return r
 }
 
-func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
+func (w *WorkerRelay) OnNewOffer(ctx context.Context, remoteOfferAnswer *OfferAnswer) {
 	if !w.isRelaySupported(remoteOfferAnswer) {
 		w.log.Infof("Relay is not supported by remote peer")
 		w.relaySupportedOnRemotePeer.Store(false)
@@ -87,7 +87,9 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relayedConn = relayedConn
 	w.relayLock.Unlock()
 
-	err = w.relayManager.AddCloseListener(srv, w.onRelayMGDisconnected)
+	err = w.relayManager.AddCloseListener(srv, func() {
+		w.onRelayMGDisconnected(ctx)
+	})
 	if err != nil {
 		log.Errorf("failed to add close listener: %s", err)
 		_ = relayedConn.Close()
@@ -95,11 +97,17 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	}
 
 	w.log.Debugf("peer conn opened via Relay: %s", srv)
-	go w.callBacks.OnConnReady(RelayConnInfo{
-		relayedConn:     relayedConn,
-		rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
-		rosenpassAddr:   remoteOfferAnswer.RosenpassAddr,
-	})
+	select {
+	case w.EventChan <- RelayEvent{
+		ConnStatus: StatusConnected,
+		RelayConnInfo: RelayConnInfo{
+			relayedConn:     relayedConn,
+			rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
+			rosenpassAddr:   remoteOfferAnswer.RosenpassAddr,
+		},
+	}:
+	case <-ctx.Done():
+	}
 }
 
 func (w *WorkerRelay) EnableWgWatcher(ctx context.Context) {
@@ -187,7 +195,11 @@ func (w *WorkerRelay) wgStateCheck(ctx context.Context, ctxCancel context.Cancel
 					w.relayLock.Lock()
 					_ = w.relayedConn.Close()
 					w.relayLock.Unlock()
-					w.callBacks.OnDisconnected()
+
+					select {
+					case w.EventChan <- RelayEvent{ConnStatus: StatusDisconnected}:
+					case <-ctx.Done():
+					}
 					return
 				}
 
@@ -225,12 +237,16 @@ func (w *WorkerRelay) wgState() (time.Time, error) {
 	return wgState.LastHandshake, nil
 }
 
-func (w *WorkerRelay) onRelayMGDisconnected() {
+func (w *WorkerRelay) onRelayMGDisconnected(ctx context.Context) {
 	w.ctxLock.Lock()
 	defer w.ctxLock.Unlock()
 
 	if w.ctxCancelWgWatch != nil {
 		w.ctxCancelWgWatch()
 	}
-	go w.callBacks.OnDisconnected()
+
+	select {
+	case w.EventChan <- RelayEvent{ConnStatus: StatusDisconnected}:
+	case <-ctx.Done():
+	}
 }

client/internal/routemanager/client.go

@@ -122,20 +122,13 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID]
 			tempScore = float64(metricDiff) * 10
 		}
 
-		// in some temporal cases, latency can be 0, so we set it to 999ms to not block but try to avoid this route
-		latency := 999 * time.Millisecond
+		// in some temporal cases, latency can be 0, so we set it to 1s to not block but try to avoid this route
+		latency := time.Second
 		if peerStatus.latency != 0 {
 			latency = peerStatus.latency
 		} else {
-			log.Tracef("peer %s has 0 latency, range %s", r.Peer, c.handler)
+			log.Warnf("peer %s has 0 latency", r.Peer)
 		}
-
-		// avoid negative tempScore on the higher latency calculation
-		if latency > 1*time.Second {
-			latency = 999 * time.Millisecond
-		}
-
-		// higher latency is worse score
 		tempScore += 1 - latency.Seconds()
 
 		if !peerStatus.relayed {
@@ -157,8 +150,6 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID]
 		}
 	}
 
-	log.Debugf("chosen route: %s, chosen score: %f, current route: %s, current score: %f", chosen, chosenScore, currID, currScore)
-
 	switch {
 	case chosen == "":
 		var peers []string
@@ -204,20 +195,15 @@ func (c *clientNetwork) watchPeerStatusChanges(ctx context.Context, peerKey stri
 func (c *clientNetwork) startPeersStatusChangeWatcher() {
 	for _, r := range c.routes {
 		_, found := c.routePeersNotifiers[r.Peer]
-		if found {
-			continue
+		if !found {
+			c.routePeersNotifiers[r.Peer] = make(chan struct{})
+			go c.watchPeerStatusChanges(c.ctx, r.Peer, c.peerStateUpdate, c.routePeersNotifiers[r.Peer])
 		}
-
-		closerChan := make(chan struct{})
-		c.routePeersNotifiers[r.Peer] = closerChan
-		go c.watchPeerStatusChanges(c.ctx, r.Peer, c.peerStateUpdate, closerChan)
 	}
 }
 
-func (c *clientNetwork) removeRouteFromWireGuardPeer() error {
-	if err := c.statusRecorder.RemovePeerStateRoute(c.currentChosen.Peer, c.handler.String()); err != nil {
-		log.Warnf("Failed to update peer state: %v", err)
-	}
+func (c *clientNetwork) removeRouteFromWireguardPeer() error {
+	c.removeStateRoute()
 
 	if err := c.handler.RemoveAllowedIPs(); err != nil {
 		return fmt.Errorf("remove allowed IPs: %w", err)
@@ -232,7 +218,7 @@ func (c *clientNetwork) removeRouteFromPeerAndSystem() error {
 
 	var merr *multierror.Error
 
-	if err := c.removeRouteFromWireGuardPeer(); err != nil {
+	if err := c.removeRouteFromWireguardPeer(); err != nil {
 		merr = multierror.Append(merr, fmt.Errorf("remove allowed IPs for peer %s: %w", c.currentChosen.Peer, err))
 	}
 	if err := c.handler.RemoveRoute(); err != nil {
@@ -271,7 +257,7 @@ func (c *clientNetwork) recalculateRouteAndUpdatePeerAndSystem() error {
 		}
 	} else {
 		// Otherwise, remove the allowed IPs from the previous peer first
-		if err := c.removeRouteFromWireGuardPeer(); err != nil {
+		if err := c.removeRouteFromWireguardPeer(); err != nil {
 			return fmt.Errorf("remove allowed IPs for peer %s: %w", c.currentChosen.Peer, err)
 		}
 	}
@@ -282,13 +268,37 @@ func (c *clientNetwork) recalculateRouteAndUpdatePeerAndSystem() error {
 		return fmt.Errorf("add allowed IPs for peer %s: %w", c.currentChosen.Peer, err)
 	}
 
-	err := c.statusRecorder.AddPeerStateRoute(c.currentChosen.Peer, c.handler.String())
-	if err != nil {
-		return fmt.Errorf("add peer state route: %w", err)
-	}
+	c.addStateRoute()
+
 	return nil
 }
 
+func (c *clientNetwork) addStateRoute() {
+	state, err := c.statusRecorder.GetPeer(c.currentChosen.Peer)
+	if err != nil {
+		log.Errorf("Failed to get peer state: %v", err)
+		return
+	}
+
+	state.AddRoute(c.handler.String())
+	if err := c.statusRecorder.UpdatePeerState(state); err != nil {
+		log.Warnf("Failed to update peer state: %v", err)
+	}
+}
+
+func (c *clientNetwork) removeStateRoute() {
+	state, err := c.statusRecorder.GetPeer(c.currentChosen.Peer)
+	if err != nil {
+		log.Errorf("Failed to get peer state: %v", err)
+		return
+	}
+
+	state.DeleteRoute(c.handler.String())
+	if err := c.statusRecorder.UpdatePeerState(state); err != nil {
+		log.Warnf("Failed to update peer state: %v", err)
+	}
+}
+
 func (c *clientNetwork) sendUpdateToClientNetworkWatcher(update routesUpdate) {
 	go func() {
 		c.routeUpdate <- update

client/internal/routemanager/refcounter/refcounter.go

@@ -217,11 +217,6 @@ func (rm *Counter[Key, I, O]) Clear() {
 
 // MarshalJSON implements the json.Marshaler interface for Counter.
 func (rm *Counter[Key, I, O]) MarshalJSON() ([]byte, error) {
-	rm.refCountMu.Lock()
-	defer rm.refCountMu.Unlock()
-	rm.idMu.Lock()
-	defer rm.idMu.Unlock()
-
 	return json.Marshal(struct {
 		RefCountMap map[Key]Ref[O]   `json:"refCountMap"`
 		IDMap       map[string][]Key `json:"idMap"`

client/server/panic_generic.go

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

client/server/panic_windows.go

@@ -1,83 +0,0 @@
-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,10 +97,6 @@ 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)
 	}
@@ -626,8 +622,6 @@ 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")
 	}

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-20241106153857-de8e2beb5254
+	github.com/netbirdio/management-integrations/integrations v0.0.0-20240929132811-9af486d346fd
 	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,6 +71,7 @@ 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
@@ -210,6 +211,8 @@ 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
@@ -236,7 +239,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-20241107152827-57d8513b5f73
+replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-20240105182236-6c340dd55aed
 
 replace github.com/cloudflare/circl => github.com/cunicu/circl v0.0.0-20230801113412-fec58fc7b5f6
 

go.sum

@@ -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-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/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/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-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/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/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,6 +605,8 @@ 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=
@@ -697,6 +699,10 @@ 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=

management/server/account.go

@@ -110,10 +110,11 @@ type AccountManager interface {
 	SaveGroups(ctx context.Context, accountID, userID string, newGroups []*nbgroup.Group) error
 	DeleteGroup(ctx context.Context, accountId, userId, groupID string) error
 	DeleteGroups(ctx context.Context, accountId, userId string, groupIDs []string) error
+	ListGroups(ctx context.Context, accountId string) ([]*nbgroup.Group, error)
 	GroupAddPeer(ctx context.Context, accountId, groupID, peerID string) error
 	GroupDeletePeer(ctx context.Context, accountId, groupID, peerID string) error
 	GetPolicy(ctx context.Context, accountID, policyID, userID string) (*Policy, error)
-	SavePolicy(ctx context.Context, accountID, userID string, policy *Policy) (*Policy, error)
+	SavePolicy(ctx context.Context, accountID, userID string, policy *Policy, isUpdate bool) error
 	DeletePolicy(ctx context.Context, accountID, policyID, userID string) error
 	ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error)
 	GetRoute(ctx context.Context, accountID string, routeID route.ID, userID string) (*route.Route, error)
@@ -139,7 +140,7 @@ type AccountManager interface {
 	HasConnectedChannel(peerID string) bool
 	GetExternalCacheManager() ExternalCacheManager
 	GetPostureChecks(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error)
-	SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error)
+	SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) error
 	DeletePostureChecks(ctx context.Context, accountID, postureChecksID, userID string) error
 	ListPostureChecks(ctx context.Context, accountID, userID string) ([]*posture.Checks, error)
 	GetIdpManager() idp.Manager
@@ -1248,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", accountID)
+			log.Errorf("failed getting account %s expiring peers", account.Id)
 			return account.GetNextInactivePeerExpiration()
 		}
 
@@ -1434,7 +1435,7 @@ func isNil(i idp.Manager) bool {
 // addAccountIDToIDPAppMeta update user's  app metadata in idp manager
 func (am *DefaultAccountManager) addAccountIDToIDPAppMeta(ctx context.Context, userID string, accountID string) error {
 	if !isNil(am.idpManager) {
-		accountUsers, err := am.Store.GetAccountUsers(ctx, LockingStrengthShare, accountID)
+		accountUsers, err := am.Store.GetAccountUsers(ctx, accountID)
 		if err != nil {
 			return err
 		}
@@ -2028,7 +2029,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 			return fmt.Errorf("error getting user: %w", err)
 		}
 
-		groups, err := transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		groups, err := transaction.GetAccountGroups(ctx, accountID)
 		if err != nil {
 			return fmt.Errorf("error getting account groups: %w", err)
 		}
@@ -2058,7 +2059,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 
 		// Propagate changes to peers if group propagation is enabled
 		if settings.GroupsPropagationEnabled {
-			groups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+			groups, err = transaction.GetAccountGroups(ctx, accountID)
 			if err != nil {
 				return fmt.Errorf("error getting account groups: %w", err)
 			}
@@ -2082,7 +2083,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 				return fmt.Errorf("error saving groups: %w", err)
 			}
 
-			if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
+			if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
 				return fmt.Errorf("error incrementing network serial: %w", err)
 			}
 		}
@@ -2100,7 +2101,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 	}
 
 	for _, g := range addNewGroups {
-		group, err := am.Store.GetGroupByID(ctx, LockingStrengthShare, accountID, g)
+		group, err := am.Store.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
 		if err != nil {
 			log.WithContext(ctx).Debugf("group %s not found while saving user activity event of account %s", g, accountID)
 		} else {
@@ -2113,7 +2114,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 	}
 
 	for _, g := range removeOldGroups {
-		group, err := am.Store.GetGroupByID(ctx, LockingStrengthShare, accountID, g)
+		group, err := am.Store.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
 		if err != nil {
 			log.WithContext(ctx).Debugf("group %s not found while saving user activity event of account %s", g, accountID)
 		} else {
@@ -2126,19 +2127,14 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 	}
 
 	if settings.GroupsPropagationEnabled {
-		removedGroupAffectsPeers, err := areGroupChangesAffectPeers(ctx, am.Store, accountID, removeOldGroups)
+		account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 		if err != nil {
-			return err
+			return fmt.Errorf("error getting account: %w", err)
 		}
 
-		newGroupsAffectsPeers, err := areGroupChangesAffectPeers(ctx, am.Store, accountID, addNewGroups)
-		if err != nil {
-			return err
-		}
-
-		if removedGroupAffectsPeers || newGroupsAffectsPeers {
+		if areGroupChangesAffectPeers(account, addNewGroups) || areGroupChangesAffectPeers(account, removeOldGroups) {
 			log.WithContext(ctx).Tracef("user %s: JWT group membership changed, updating account peers", claims.UserId)
-			am.updateAccountPeers(ctx, accountID)
+			am.updateAccountPeers(ctx, account)
 		}
 	}
 
@@ -2294,12 +2290,12 @@ func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return nil, nil, nil, status.NewGetAccountError(err)
+		return nil, nil, nil, err
 	}
 
 	peer, netMap, postureChecks, err := am.SyncPeer(ctx, PeerSync{WireGuardPubKey: peerPubKey, Meta: meta}, account)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("error syncing peer: %w", err)
+		return nil, nil, nil, err
 	}
 
 	err = am.MarkPeerConnected(ctx, peerPubKey, true, realIP, account)
@@ -2318,7 +2314,7 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return status.NewGetAccountError(err)
+		return err
 	}
 
 	err = am.MarkPeerConnected(ctx, peerPubKey, false, nil, account)
@@ -2402,7 +2398,12 @@ func (am *DefaultAccountManager) CheckUserAccessByJWTGroups(ctx context.Context,
 
 func (am *DefaultAccountManager) onPeersInvalidated(ctx context.Context, accountID string) {
 	log.WithContext(ctx).Debugf("validated peers has been invalidated for account %s", accountID)
-	am.updateAccountPeers(ctx, accountID)
+	updatedAccount, err := am.Store.GetAccount(ctx, accountID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to get account %s: %v", accountID, err)
+		return
+	}
+	am.updateAccountPeers(ctx, updatedAccount)
 }
 
 func (am *DefaultAccountManager) FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {

management/server/account_test.go

@@ -29,18 +29,14 @@ 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, 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) 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) 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{})
@@ -982,110 +978,6 @@ 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 {
@@ -1238,7 +1130,8 @@ func TestAccountManager_NetworkUpdates_SaveGroup(t *testing.T) {
 		return
 	}
 
-	_, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+	policy := Policy{
+		ID:      "policy",
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -1249,7 +1142,8 @@ func TestAccountManager_NetworkUpdates_SaveGroup(t *testing.T) {
 				Action:        PolicyTrafficActionAccept,
 			},
 		},
-	})
+	}
+	err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 	require.NoError(t, err)
 
 	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
@@ -1318,6 +1212,19 @@ func TestAccountManager_NetworkUpdates_SavePolicy(t *testing.T) {
 	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
 	defer manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
 
+	policy := Policy{
+		Enabled: true,
+		Rules: []*PolicyRule{
+			{
+				Enabled:       true,
+				Sources:       []string{"groupA"},
+				Destinations:  []string{"groupA"},
+				Bidirectional: true,
+				Action:        PolicyTrafficActionAccept,
+			},
+		},
+	}
+
 	wg := sync.WaitGroup{}
 	wg.Add(1)
 	go func() {
@@ -1330,19 +1237,7 @@ func TestAccountManager_NetworkUpdates_SavePolicy(t *testing.T) {
 		}
 	}()
 
-	_, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-		Enabled: true,
-		Rules: []*PolicyRule{
-			{
-				Enabled:       true,
-				Sources:       []string{"groupA"},
-				Destinations:  []string{"groupA"},
-				Bidirectional: true,
-				Action:        PolicyTrafficActionAccept,
-			},
-		},
-	})
-	if err != nil {
+	if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil {
 		t.Errorf("delete default rule: %v", err)
 		return
 	}
@@ -1363,7 +1258,7 @@ func TestAccountManager_NetworkUpdates_DeletePeer(t *testing.T) {
 		return
 	}
 
-	_, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+	policy := Policy{
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -1374,8 +1269,9 @@ func TestAccountManager_NetworkUpdates_DeletePeer(t *testing.T) {
 				Action:        PolicyTrafficActionAccept,
 			},
 		},
-	})
-	if err != nil {
+	}
+
+	if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil {
 		t.Errorf("save policy: %v", err)
 		return
 	}
@@ -1409,20 +1305,13 @@ func TestAccountManager_NetworkUpdates_DeleteGroup(t *testing.T) {
 	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
 	defer manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
 
-	err := manager.SaveGroup(context.Background(), account.Id, userID, &group.Group{
+	group := group.Group{
 		ID:    "groupA",
 		Name:  "GroupA",
 		Peers: []string{peer1.ID, peer2.ID, peer3.ID},
-	})
-
-	require.NoError(t, err, "failed to save group")
-
-	if err := manager.DeletePolicy(context.Background(), account.Id, account.Policies[0].ID, userID); err != nil {
-		t.Errorf("delete default rule: %v", err)
-		return
 	}
 
-	policy, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+	policy := Policy{
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -1433,8 +1322,14 @@ func TestAccountManager_NetworkUpdates_DeleteGroup(t *testing.T) {
 				Action:        PolicyTrafficActionAccept,
 			},
 		},
-	})
-	if err != nil {
+	}
+
+	if err := manager.DeletePolicy(context.Background(), account.Id, account.Policies[0].ID, userID); err != nil {
+		t.Errorf("delete default rule: %v", err)
+		return
+	}
+
+	if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil {
 		t.Errorf("save policy: %v", err)
 		return
 	}
@@ -1457,7 +1352,7 @@ func TestAccountManager_NetworkUpdates_DeleteGroup(t *testing.T) {
 		return
 	}
 
-	if err := manager.DeleteGroup(context.Background(), account.Id, userID, "groupA"); err != nil {
+	if err := manager.DeleteGroup(context.Background(), account.Id, "", group.ID); err != nil {
 		t.Errorf("delete group: %v", err)
 		return
 	}
@@ -2711,7 +2606,7 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		assert.NoError(t, err, "unable to get user")
 		assert.Len(t, user.AutoGroups, 0)
 
-		group1, err := manager.Store.GetGroupByID(context.Background(), LockingStrengthShare, "accountID", "group1")
+		group1, err := manager.Store.GetGroupByID(context.Background(), LockingStrengthShare, "group1", "accountID")
 		assert.NoError(t, err, "unable to get group")
 		assert.Equal(t, group1.Issued, group.GroupIssuedAPI, "group should be api issued")
 	})
@@ -2731,7 +2626,7 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		assert.NoError(t, err, "unable to get user")
 		assert.Len(t, user.AutoGroups, 1)
 
-		group1, err := manager.Store.GetGroupByID(context.Background(), LockingStrengthShare, "accountID", "group1")
+		group1, err := manager.Store.GetGroupByID(context.Background(), LockingStrengthShare, "group1", "accountID")
 		assert.NoError(t, err, "unable to get group")
 		assert.Equal(t, group1.Issued, group.GroupIssuedAPI, "group should be api issued")
 	})
@@ -2770,7 +2665,7 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		err = manager.syncJWTGroups(context.Background(), "accountID", claims)
 		assert.NoError(t, err, "unable to sync jwt groups")
 
-		groups, err := manager.Store.GetAccountGroups(context.Background(), LockingStrengthShare, "accountID")
+		groups, err := manager.Store.GetAccountGroups(context.Background(), "accountID")
 		assert.NoError(t, err)
 		assert.Len(t, groups, 3, "new group3 should be added")
 

management/server/differs/netip.go

@@ -0,0 +1,82 @@
+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.go

@@ -3,7 +3,6 @@ package server
 import (
 	"context"
 	"fmt"
-	"slices"
 	"strconv"
 	"sync"
 
@@ -86,12 +85,8 @@ func (am *DefaultAccountManager) GetDNSSettings(ctx context.Context, accountID s
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power are allowed to view DNS settings")
 	}
 
 	return am.Store.GetAccountDNSSettings(ctx, LockingStrengthShare, accountID)
@@ -99,137 +94,64 @@ func (am *DefaultAccountManager) GetDNSSettings(ctx context.Context, accountID s
 
 // SaveDNSSettings validates a user role and updates the account's DNS settings
 func (am *DefaultAccountManager) SaveDNSSettings(ctx context.Context, accountID string, userID string, dnsSettingsToSave *DNSSettings) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
+	account, err := am.Store.GetAccount(ctx, accountID)
+	if err != nil {
+		return err
+	}
+
+	user, err := account.FindUser(userID)
+	if err != nil {
+		return err
+	}
+
+	if !user.HasAdminPower() {
+		return status.Errorf(status.PermissionDenied, "only users with admin power are allowed to update DNS settings")
+	}
+
 	if dnsSettingsToSave == nil {
 		return status.Errorf(status.InvalidArgument, "the dns settings provided are nil")
 	}
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
-	if err != nil {
-		return err
-	}
-
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	if !user.HasAdminPower() {
-		return status.NewAdminPermissionError()
-	}
-
-	var updateAccountPeers bool
-	var eventsToStore []func()
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateDNSSettings(ctx, transaction, accountID, dnsSettingsToSave); err != nil {
-			return err
-		}
-
-		oldSettings, err := transaction.GetAccountDNSSettings(ctx, LockingStrengthUpdate, accountID)
+	if len(dnsSettingsToSave.DisabledManagementGroups) != 0 {
+		err = validateGroups(dnsSettingsToSave.DisabledManagementGroups, account.Groups)
 		if err != nil {
 			return err
 		}
+	}
 
-		addedGroups := difference(dnsSettingsToSave.DisabledManagementGroups, oldSettings.DisabledManagementGroups)
-		removedGroups := difference(oldSettings.DisabledManagementGroups, dnsSettingsToSave.DisabledManagementGroups)
+	oldSettings := account.DNSSettings.Copy()
+	account.DNSSettings = dnsSettingsToSave.Copy()
 
-		updateAccountPeers, err = areDNSSettingChangesAffectPeers(ctx, transaction, accountID, addedGroups, removedGroups)
-		if err != nil {
-			return err
-		}
+	addedGroups := difference(dnsSettingsToSave.DisabledManagementGroups, oldSettings.DisabledManagementGroups)
+	removedGroups := difference(oldSettings.DisabledManagementGroups, dnsSettingsToSave.DisabledManagementGroups)
 
-		events := am.prepareDNSSettingsEvents(ctx, transaction, accountID, userID, addedGroups, removedGroups)
-		eventsToStore = append(eventsToStore, events...)
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.SaveDNSSettings(ctx, LockingStrengthUpdate, accountID, dnsSettingsToSave)
-	})
-	if err != nil {
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
 
-	for _, storeEvent := range eventsToStore {
-		storeEvent()
+	for _, id := range addedGroups {
+		group := account.GetGroup(id)
+		meta := map[string]any{"group": group.Name, "group_id": group.ID}
+		am.StoreEvent(ctx, userID, accountID, accountID, activity.GroupAddedToDisabledManagementGroups, meta)
 	}
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+	for _, id := range removedGroups {
+		group := account.GetGroup(id)
+		meta := map[string]any{"group": group.Name, "group_id": group.ID}
+		am.StoreEvent(ctx, userID, accountID, accountID, activity.GroupRemovedFromDisabledManagementGroups, meta)
+	}
+
+	if anyGroupHasPeers(account, addedGroups) || anyGroupHasPeers(account, removedGroups) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
 }
 
-// prepareGroupEvents prepares a list of event functions to be stored.
-func (am *DefaultAccountManager) prepareDNSSettingsEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string) []func() {
-	var eventsToStore []func()
-
-	modifiedGroups := slices.Concat(addedGroups, removedGroups)
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, modifiedGroups)
-	if err != nil {
-		log.WithContext(ctx).Debugf("failed to get groups for dns settings events: %v", err)
-		return nil
-	}
-
-	for _, groupID := range addedGroups {
-		group, ok := groups[groupID]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToDisabledManagementGroups activity", groupID)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID}
-			am.StoreEvent(ctx, userID, accountID, accountID, activity.GroupAddedToDisabledManagementGroups, meta)
-		})
-
-	}
-
-	for _, groupID := range removedGroups {
-		group, ok := groups[groupID]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromDisabledManagementGroups activity", groupID)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID}
-			am.StoreEvent(ctx, userID, accountID, accountID, activity.GroupRemovedFromDisabledManagementGroups, meta)
-		})
-	}
-
-	return eventsToStore
-}
-
-// areDNSSettingChangesAffectPeers checks if the DNS settings changes affect any peers.
-func areDNSSettingChangesAffectPeers(ctx context.Context, transaction Store, accountID string, addedGroups, removedGroups []string) (bool, error) {
-	hasPeers, err := anyGroupHasPeers(ctx, transaction, accountID, addedGroups)
-	if err != nil {
-		return false, err
-	}
-
-	if hasPeers {
-		return true, nil
-	}
-
-	return anyGroupHasPeers(ctx, transaction, accountID, removedGroups)
-}
-
-// validateDNSSettings validates the DNS settings.
-func validateDNSSettings(ctx context.Context, transaction Store, accountID string, settings *DNSSettings) error {
-	if len(settings.DisabledManagementGroups) == 0 {
-		return nil
-	}
-
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, settings.DisabledManagementGroups)
-	if err != nil {
-		return err
-	}
-
-	return validateGroups(settings.DisabledManagementGroups, groups)
-}
-
 // toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache
 func toProtocolDNSConfig(update nbdns.Config, cache *DNSConfigCache) *proto.DNSConfig {
 	protoUpdate := &proto.DNSConfig{

management/server/dns_test.go

@@ -8,10 +8,9 @@ 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"
 
@@ -522,64 +521,23 @@ func TestDNSAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// 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", "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")
-		}
+	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 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")
-		}
-	})
+	_, 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)
 
 	// 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) {
@@ -601,6 +559,27 @@ 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.go

@@ -37,12 +37,8 @@ func (am *DefaultAccountManager) CheckGroupPermissions(ctx context.Context, acco
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() && settings.RegularUsersViewBlocked {
-		return status.NewAdminPermissionError()
+	if (!user.IsAdminOrServiceUser() && settings.RegularUsersViewBlocked) || user.AccountID != accountID {
+		return status.Errorf(status.PermissionDenied, "groups are blocked for users")
 	}
 
 	return nil
@@ -53,7 +49,8 @@ func (am *DefaultAccountManager) GetGroup(ctx context.Context, accountID, groupI
 	if err := am.CheckGroupPermissions(ctx, accountID, userID); err != nil {
 		return nil, err
 	}
-	return am.Store.GetGroupByID(ctx, LockingStrengthShare, accountID, groupID)
+
+	return am.Store.GetGroupByID(ctx, LockingStrengthShare, groupID, accountID)
 }
 
 // GetAllGroups returns all groups in an account
@@ -61,12 +58,13 @@ func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID, us
 	if err := am.CheckGroupPermissions(ctx, accountID, userID); err != nil {
 		return nil, err
 	}
-	return am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+
+	return am.Store.GetAccountGroups(ctx, accountID)
 }
 
 // GetGroupByName filters all groups in an account by name and returns the one with the most peers
 func (am *DefaultAccountManager) GetGroupByName(ctx context.Context, groupName, accountID string) (*nbgroup.Group, error) {
-	return am.Store.GetGroupByName(ctx, LockingStrengthShare, accountID, groupName)
+	return am.Store.GetGroupByName(ctx, LockingStrengthShare, groupName, accountID)
 }
 
 // SaveGroup object of the peers
@@ -79,74 +77,79 @@ func (am *DefaultAccountManager) SaveGroup(ctx context.Context, accountID, userI
 // SaveGroups adds new groups to the account.
 // Note: This function does not acquire the global lock.
 // It is the caller's responsibility to ensure proper locking is in place before invoking this method.
-func (am *DefaultAccountManager) SaveGroups(ctx context.Context, accountID, userID string, groups []*nbgroup.Group) error {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+func (am *DefaultAccountManager) SaveGroups(ctx context.Context, accountID, userID string, newGroups []*nbgroup.Group) error {
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
 	var eventsToStore []func()
-	var groupsToSave []*nbgroup.Group
-	var updateAccountPeers bool
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		groupIDs := make([]string, 0, len(groups))
-		for _, newGroup := range groups {
-			if err = validateNewGroup(ctx, transaction, accountID, newGroup); err != nil {
-				return err
+	for _, newGroup := range newGroups {
+		if newGroup.ID == "" && newGroup.Issued != nbgroup.GroupIssuedAPI {
+			return status.Errorf(status.InvalidArgument, "%s group without ID set", newGroup.Issued)
+		}
+
+		if newGroup.ID == "" && newGroup.Issued == nbgroup.GroupIssuedAPI {
+			existingGroup, err := account.FindGroupByName(newGroup.Name)
+			if err != nil {
+				s, ok := status.FromError(err)
+				if !ok || s.ErrorType != status.NotFound {
+					return err
+				}
 			}
 
-			newGroup.AccountID = accountID
-			groupsToSave = append(groupsToSave, newGroup)
-			groupIDs = append(groupIDs, newGroup.ID)
+			// Avoid duplicate groups only for the API issued groups.
+			// Integration or JWT groups can be duplicated as they are coming from the IdP that we don't have control of.
+			if existingGroup != nil {
+				return status.Errorf(status.AlreadyExists, "group with name %s already exists", newGroup.Name)
+			}
 
-			events := am.prepareGroupEvents(ctx, transaction, accountID, userID, newGroup)
-			eventsToStore = append(eventsToStore, events...)
+			newGroup.ID = xid.New().String()
 		}
 
-		updateAccountPeers, err = areGroupChangesAffectPeers(ctx, transaction, accountID, groupIDs)
-		if err != nil {
-			return err
+		for _, peerID := range newGroup.Peers {
+			if account.Peers[peerID] == nil {
+				return status.Errorf(status.InvalidArgument, "peer with ID \"%s\" not found", peerID)
+			}
 		}
 
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
+		oldGroup := account.Groups[newGroup.ID]
+		account.Groups[newGroup.ID] = newGroup
 
-		return transaction.SaveGroups(ctx, LockingStrengthUpdate, groupsToSave)
-	})
-	if err != nil {
+		events := am.prepareGroupEvents(ctx, userID, accountID, newGroup, oldGroup, account)
+		eventsToStore = append(eventsToStore, events...)
+	}
+
+	newGroupIDs := make([]string, 0, len(newGroups))
+	for _, newGroup := range newGroups {
+		newGroupIDs = append(newGroupIDs, newGroup.ID)
+	}
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
 
+	if areGroupChangesAffectPeers(account, newGroupIDs) {
+		am.updateAccountPeers(ctx, account)
+	}
+
 	for _, storeEvent := range eventsToStore {
 		storeEvent()
 	}
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
 	return nil
 }
 
 // prepareGroupEvents prepares a list of event functions to be stored.
-func (am *DefaultAccountManager) prepareGroupEvents(ctx context.Context, transaction Store, accountID, userID string, newGroup *nbgroup.Group) []func() {
+func (am *DefaultAccountManager) prepareGroupEvents(ctx context.Context, userID string, accountID string, newGroup, oldGroup *nbgroup.Group, account *Account) []func() {
 	var eventsToStore []func()
 
 	addedPeers := make([]string, 0)
 	removedPeers := make([]string, 0)
 
-	oldGroup, err := transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, newGroup.ID)
-	if err == nil && oldGroup != nil {
+	if oldGroup != nil {
 		addedPeers = difference(newGroup.Peers, oldGroup.Peers)
 		removedPeers = difference(oldGroup.Peers, newGroup.Peers)
 	} else {
@@ -156,42 +159,35 @@ func (am *DefaultAccountManager) prepareGroupEvents(ctx context.Context, transac
 		})
 	}
 
-	modifiedPeers := slices.Concat(addedPeers, removedPeers)
-	peers, err := transaction.GetPeersByIDs(ctx, LockingStrengthShare, accountID, modifiedPeers)
-	if err != nil {
-		log.WithContext(ctx).Debugf("failed to get peers for group events: %v", err)
-		return nil
-	}
-
-	for _, peerID := range addedPeers {
-		peer, ok := peers[peerID]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding peer: %s GroupAddedToPeer activity: peer not found in store", peerID)
+	for _, p := range addedPeers {
+		peer := account.Peers[p]
+		if peer == nil {
+			log.WithContext(ctx).Errorf("peer %s not found under account %s while saving group", p, accountID)
 			continue
 		}
-
+		peerCopy := peer // copy to avoid closure issues
 		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{
-				"group": newGroup.Name, "group_id": newGroup.ID,
-				"peer_ip": peer.IP.String(), "peer_fqdn": peer.FQDN(am.GetDNSDomain()),
-			}
-			am.StoreEvent(ctx, userID, peer.ID, accountID, activity.GroupAddedToPeer, meta)
+			am.StoreEvent(ctx, userID, peerCopy.ID, accountID, activity.GroupAddedToPeer,
+				map[string]any{
+					"group": newGroup.Name, "group_id": newGroup.ID, "peer_ip": peerCopy.IP.String(),
+					"peer_fqdn": peerCopy.FQDN(am.GetDNSDomain()),
+				})
 		})
 	}
 
-	for _, peerID := range removedPeers {
-		peer, ok := peers[peerID]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding peer: %s GroupRemovedFromPeer activity: peer not found in store", peerID)
+	for _, p := range removedPeers {
+		peer := account.Peers[p]
+		if peer == nil {
+			log.WithContext(ctx).Errorf("peer %s not found under account %s while saving group", p, accountID)
 			continue
 		}
-
+		peerCopy := peer // copy to avoid closure issues
 		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{
-				"group": newGroup.Name, "group_id": newGroup.ID,
-				"peer_ip": peer.IP.String(), "peer_fqdn": peer.FQDN(am.GetDNSDomain()),
-			}
-			am.StoreEvent(ctx, userID, peer.ID, accountID, activity.GroupRemovedFromPeer, meta)
+			am.StoreEvent(ctx, userID, peerCopy.ID, accountID, activity.GroupRemovedFromPeer,
+				map[string]any{
+					"group": newGroup.Name, "group_id": newGroup.ID, "peer_ip": peerCopy.IP.String(),
+					"peer_fqdn": peerCopy.FQDN(am.GetDNSDomain()),
+				})
 		})
 	}
 
@@ -214,47 +210,40 @@ func difference(a, b []string) []string {
 }
 
 // DeleteGroup object of the peers.
-func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, userID, groupID string) error {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountId, userId, groupID string) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountId)
+	defer unlock()
+
+	account, err := am.Store.GetAccount(ctx, accountId)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
+	group, ok := account.Groups[groupID]
+	if !ok {
+		return nil
 	}
 
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
-	var group *nbgroup.Group
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		group, err = transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, groupID)
-		if err != nil {
-			return err
-		}
-
-		if group.IsGroupAll() {
-			return status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed")
-		}
-
-		if err = validateDeleteGroup(ctx, transaction, group, userID); err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.DeleteGroup(ctx, LockingStrengthUpdate, accountID, groupID)
-	})
+	allGroup, err := account.GetGroupAll()
 	if err != nil {
 		return err
 	}
 
-	am.StoreEvent(ctx, userID, groupID, accountID, activity.GroupDeleted, group.EventMeta())
+	if allGroup.ID == groupID {
+		return status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed")
+	}
+
+	if err = validateDeleteGroup(account, group, userId); err != nil {
+		return err
+	}
+	delete(account.Groups, groupID)
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, userId, groupID, accountId, activity.GroupDeleted, group.EventMeta())
 
 	return nil
 }
@@ -265,90 +254,93 @@ func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, use
 //
 // If an error occurs while deleting a group, the function skips it and continues deleting other groups.
 // Errors are collected and returned at the end.
-func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountID, userID string, groupIDs []string) error {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountId, userId string, groupIDs []string) error {
+	account, err := am.Store.GetAccount(ctx, accountId)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
 	var allErrors error
-	var groupIDsToDelete []string
-	var deletedGroups []*nbgroup.Group
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		for _, groupID := range groupIDs {
-			group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, groupID)
-			if err != nil {
-				continue
-			}
-
-			if err := validateDeleteGroup(ctx, transaction, group, userID); err != nil {
-				allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete group %s: %w", groupID, err))
-				continue
-			}
-
-			groupIDsToDelete = append(groupIDsToDelete, groupID)
-			deletedGroups = append(deletedGroups, group)
+	deletedGroups := make([]*nbgroup.Group, 0, len(groupIDs))
+	for _, groupID := range groupIDs {
+		group, ok := account.Groups[groupID]
+		if !ok {
+			continue
 		}
 
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
+		if err := validateDeleteGroup(account, group, userId); err != nil {
+			allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete group %s: %w", groupID, err))
+			continue
 		}
 
-		return transaction.DeleteGroups(ctx, LockingStrengthUpdate, accountID, groupIDsToDelete)
-	})
-	if err != nil {
+		delete(account.Groups, groupID)
+		deletedGroups = append(deletedGroups, group)
+	}
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
 
-	for _, group := range deletedGroups {
-		am.StoreEvent(ctx, userID, group.ID, accountID, activity.GroupDeleted, group.EventMeta())
+	for _, g := range deletedGroups {
+		am.StoreEvent(ctx, userId, g.ID, accountId, activity.GroupDeleted, g.EventMeta())
 	}
 
 	return allErrors
 }
 
+// ListGroups objects of the peers
+func (am *DefaultAccountManager) ListGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error) {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
+	account, err := am.Store.GetAccount(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	groups := make([]*nbgroup.Group, 0, len(account.Groups))
+	for _, item := range account.Groups {
+		groups = append(groups, item)
+	}
+
+	return groups, nil
+}
+
 // GroupAddPeer appends peer to the group
 func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, groupID, peerID string) error {
-	var group *nbgroup.Group
-	var updateAccountPeers bool
-	var err error
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
-		if err != nil {
-			return err
-		}
-
-		if updated := group.AddPeer(peerID); !updated {
-			return nil
-		}
-
-		updateAccountPeers, err = areGroupChangesAffectPeers(ctx, transaction, accountID, []string{groupID})
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.SaveGroup(ctx, LockingStrengthUpdate, group)
-	})
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+	group, ok := account.Groups[groupID]
+	if !ok {
+		return status.Errorf(status.NotFound, "group with ID %s not found", groupID)
+	}
+
+	add := true
+	for _, itemID := range group.Peers {
+		if itemID == peerID {
+			add = false
+			break
+		}
+	}
+	if add {
+		group.Peers = append(group.Peers, peerID)
+	}
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	if areGroupChangesAffectPeers(account, []string{group.ID}) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
@@ -356,80 +348,41 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr
 
 // GroupDeletePeer removes peer from the group
 func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, groupID, peerID string) error {
-	var group *nbgroup.Group
-	var updateAccountPeers bool
-	var err error
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
-		if err != nil {
-			return err
-		}
-
-		if updated := group.RemovePeer(peerID); !updated {
-			return nil
-		}
-
-		updateAccountPeers, err = areGroupChangesAffectPeers(ctx, transaction, accountID, []string{groupID})
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.SaveGroup(ctx, LockingStrengthUpdate, group)
-	})
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+	group, ok := account.Groups[groupID]
+	if !ok {
+		return status.Errorf(status.NotFound, "group with ID %s not found", groupID)
 	}
 
-	return nil
-}
-
-// validateNewGroup validates the new group for existence and required fields.
-func validateNewGroup(ctx context.Context, transaction Store, accountID string, newGroup *nbgroup.Group) error {
-	if newGroup.ID == "" && newGroup.Issued != nbgroup.GroupIssuedAPI {
-		return status.Errorf(status.InvalidArgument, "%s group without ID set", newGroup.Issued)
-	}
-
-	if newGroup.ID == "" && newGroup.Issued == nbgroup.GroupIssuedAPI {
-		existingGroup, err := transaction.GetGroupByName(ctx, LockingStrengthShare, accountID, newGroup.Name)
-		if err != nil {
-			if s, ok := status.FromError(err); !ok || s.Type() != status.NotFound {
+	account.Network.IncSerial()
+	for i, itemID := range group.Peers {
+		if itemID == peerID {
+			group.Peers = append(group.Peers[:i], group.Peers[i+1:]...)
+			if err := am.Store.SaveAccount(ctx, account); err != nil {
 				return err
 			}
 		}
-
-		// Prevent duplicate groups for API-issued groups.
-		// Integration or JWT groups can be duplicated as they are coming from the IdP that we don't have control of.
-		if existingGroup != nil {
-			return status.Errorf(status.AlreadyExists, "group with name %s already exists", newGroup.Name)
-		}
-
-		newGroup.ID = xid.New().String()
 	}
 
-	for _, peerID := range newGroup.Peers {
-		_, err := transaction.GetPeerByID(ctx, LockingStrengthShare, accountID, peerID)
-		if err != nil {
-			return status.Errorf(status.InvalidArgument, "peer with ID \"%s\" not found", peerID)
-		}
+	if areGroupChangesAffectPeers(account, []string{group.ID}) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
 }
 
-func validateDeleteGroup(ctx context.Context, transaction Store, group *nbgroup.Group, userID string) error {
+func validateDeleteGroup(account *Account, group *nbgroup.Group, userID string) error {
 	// disable a deleting integration group if the initiator is not an admin service user
 	if group.Issued == nbgroup.GroupIssuedIntegration {
-		executingUser, err := transaction.GetUserByUserID(ctx, LockingStrengthShare, userID)
-		if err != nil {
+		executingUser := account.Users[userID]
+		if executingUser == nil {
 			return status.Errorf(status.NotFound, "user not found")
 		}
 		if executingUser.Role != UserRoleAdmin || !executingUser.IsServiceUser {
@@ -437,77 +390,51 @@ func validateDeleteGroup(ctx context.Context, transaction Store, group *nbgroup.
 		}
 	}
 
-	if isLinked, linkedRoute := isGroupLinkedToRoute(ctx, transaction, group.AccountID, group.ID); isLinked {
+	if isLinked, linkedRoute := isGroupLinkedToRoute(account.Routes, group.ID); isLinked {
 		return &GroupLinkError{"route", string(linkedRoute.NetID)}
 	}
 
-	if isLinked, linkedDns := isGroupLinkedToDns(ctx, transaction, group.AccountID, group.ID); isLinked {
+	if isLinked, linkedDns := isGroupLinkedToDns(account.NameServerGroups, group.ID); isLinked {
 		return &GroupLinkError{"name server groups", linkedDns.Name}
 	}
 
-	if isLinked, linkedPolicy := isGroupLinkedToPolicy(ctx, transaction, group.AccountID, group.ID); isLinked {
+	if isLinked, linkedPolicy := isGroupLinkedToPolicy(account.Policies, group.ID); isLinked {
 		return &GroupLinkError{"policy", linkedPolicy.Name}
 	}
 
-	if isLinked, linkedSetupKey := isGroupLinkedToSetupKey(ctx, transaction, group.AccountID, group.ID); isLinked {
+	if isLinked, linkedSetupKey := isGroupLinkedToSetupKey(account.SetupKeys, group.ID); isLinked {
 		return &GroupLinkError{"setup key", linkedSetupKey.Name}
 	}
 
-	if isLinked, linkedUser := isGroupLinkedToUser(ctx, transaction, group.AccountID, group.ID); isLinked {
+	if isLinked, linkedUser := isGroupLinkedToUser(account.Users, group.ID); isLinked {
 		return &GroupLinkError{"user", linkedUser.Id}
 	}
 
-	return checkGroupLinkedToSettings(ctx, transaction, group)
-}
-
-// checkGroupLinkedToSettings verifies if a group is linked to any settings in the account.
-func checkGroupLinkedToSettings(ctx context.Context, transaction Store, group *nbgroup.Group) error {
-	dnsSettings, err := transaction.GetAccountDNSSettings(ctx, LockingStrengthShare, group.AccountID)
-	if err != nil {
-		return err
-	}
-
-	if slices.Contains(dnsSettings.DisabledManagementGroups, group.ID) {
+	if slices.Contains(account.DNSSettings.DisabledManagementGroups, group.ID) {
 		return &GroupLinkError{"disabled DNS management groups", group.Name}
 	}
 
-	settings, err := transaction.GetAccountSettings(ctx, LockingStrengthShare, group.AccountID)
-	if err != nil {
-		return err
-	}
-
-	if settings.Extra != nil && slices.Contains(settings.Extra.IntegratedValidatorGroups, group.ID) {
-		return &GroupLinkError{"integrated validator", group.Name}
+	if account.Settings.Extra != nil {
+		if slices.Contains(account.Settings.Extra.IntegratedValidatorGroups, group.ID) {
+			return &GroupLinkError{"integrated validator", group.Name}
+		}
 	}
 
 	return nil
 }
 
 // isGroupLinkedToRoute checks if a group is linked to any route in the account.
-func isGroupLinkedToRoute(ctx context.Context, transaction Store, accountID string, groupID string) (bool, *route.Route) {
-	routes, err := transaction.GetAccountRoutes(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("error retrieving routes while checking group linkage: %v", err)
-		return false, nil
-	}
-
+func isGroupLinkedToRoute(routes map[route.ID]*route.Route, groupID string) (bool, *route.Route) {
 	for _, r := range routes {
 		if slices.Contains(r.Groups, groupID) || slices.Contains(r.PeerGroups, groupID) {
 			return true, r
 		}
 	}
-
 	return false, nil
 }
 
 // isGroupLinkedToPolicy checks if a group is linked to any policy in the account.
-func isGroupLinkedToPolicy(ctx context.Context, transaction Store, accountID string, groupID string) (bool, *Policy) {
-	policies, err := transaction.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("error retrieving policies while checking group linkage: %v", err)
-		return false, nil
-	}
-
+func isGroupLinkedToPolicy(policies []*Policy, groupID string) (bool, *Policy) {
 	for _, policy := range policies {
 		for _, rule := range policy.Rules {
 			if slices.Contains(rule.Sources, groupID) || slices.Contains(rule.Destinations, groupID) {
@@ -519,13 +446,7 @@ func isGroupLinkedToPolicy(ctx context.Context, transaction Store, accountID str
 }
 
 // isGroupLinkedToDns checks if a group is linked to any nameserver group in the account.
-func isGroupLinkedToDns(ctx context.Context, transaction Store, accountID string, groupID string) (bool, *nbdns.NameServerGroup) {
-	nameServerGroups, err := transaction.GetAccountNameServerGroups(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("error retrieving name server groups while checking group linkage: %v", err)
-		return false, nil
-	}
-
+func isGroupLinkedToDns(nameServerGroups map[string]*nbdns.NameServerGroup, groupID string) (bool, *nbdns.NameServerGroup) {
 	for _, dns := range nameServerGroups {
 		for _, g := range dns.Groups {
 			if g == groupID {
@@ -533,18 +454,11 @@ func isGroupLinkedToDns(ctx context.Context, transaction Store, accountID string
 			}
 		}
 	}
-
 	return false, nil
 }
 
 // isGroupLinkedToSetupKey checks if a group is linked to any setup key in the account.
-func isGroupLinkedToSetupKey(ctx context.Context, transaction Store, accountID string, groupID string) (bool, *SetupKey) {
-	setupKeys, err := transaction.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("error retrieving setup keys while checking group linkage: %v", err)
-		return false, nil
-	}
-
+func isGroupLinkedToSetupKey(setupKeys map[string]*SetupKey, groupID string) (bool, *SetupKey) {
 	for _, setupKey := range setupKeys {
 		if slices.Contains(setupKey.AutoGroups, groupID) {
 			return true, setupKey
@@ -554,13 +468,7 @@ func isGroupLinkedToSetupKey(ctx context.Context, transaction Store, accountID s
 }
 
 // isGroupLinkedToUser checks if a group is linked to any user in the account.
-func isGroupLinkedToUser(ctx context.Context, transaction Store, accountID string, groupID string) (bool, *User) {
-	users, err := transaction.GetAccountUsers(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("error retrieving users while checking group linkage: %v", err)
-		return false, nil
-	}
-
+func isGroupLinkedToUser(users map[string]*User, groupID string) (bool, *User) {
 	for _, user := range users {
 		if slices.Contains(user.AutoGroups, groupID) {
 			return true, user
@@ -569,36 +477,8 @@ func isGroupLinkedToUser(ctx context.Context, transaction Store, accountID strin
 	return false, nil
 }
 
-// areGroupChangesAffectPeers checks if any changes to the specified groups will affect peers.
-func areGroupChangesAffectPeers(ctx context.Context, transaction Store, accountID string, groupIDs []string) (bool, error) {
-	if len(groupIDs) == 0 {
-		return false, nil
-	}
-
-	dnsSettings, err := transaction.GetAccountDNSSettings(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		return false, err
-	}
-
-	for _, groupID := range groupIDs {
-		if slices.Contains(dnsSettings.DisabledManagementGroups, groupID) {
-			return true, nil
-		}
-		if linked, _ := isGroupLinkedToDns(ctx, transaction, accountID, groupID); linked {
-			return true, nil
-		}
-		if linked, _ := isGroupLinkedToPolicy(ctx, transaction, accountID, groupID); linked {
-			return true, nil
-		}
-		if linked, _ := isGroupLinkedToRoute(ctx, transaction, accountID, groupID); linked {
-			return true, nil
-		}
-	}
-
-	return false, nil
-}
-
-func (am *DefaultAccountManager) anyGroupHasPeers(account *Account, groupIDs []string) bool {
+// anyGroupHasPeers checks if any of the given groups in the account have peers.
+func anyGroupHasPeers(account *Account, groupIDs []string) bool {
 	for _, groupID := range groupIDs {
 		if group, exists := account.Groups[groupID]; exists && group.HasPeers() {
 			return true
@@ -607,18 +487,21 @@ func (am *DefaultAccountManager) anyGroupHasPeers(account *Account, groupIDs []s
 	return false
 }
 
-// anyGroupHasPeers checks if any of the given groups in the account have peers.
-func anyGroupHasPeers(ctx context.Context, transaction Store, accountID string, groupIDs []string) (bool, error) {
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, groupIDs)
-	if err != nil {
-		return false, err
-	}
-
-	for _, group := range groups {
-		if group.HasPeers() {
-			return true, nil
+func areGroupChangesAffectPeers(account *Account, groupIDs []string) bool {
+	for _, groupID := range groupIDs {
+		if slices.Contains(account.DNSSettings.DisabledManagementGroups, groupID) {
+			return true
+		}
+		if linked, _ := isGroupLinkedToDns(account.NameServerGroups, groupID); linked {
+			return true
+		}
+		if linked, _ := isGroupLinkedToPolicy(account.Policies, groupID); linked {
+			return true
+		}
+		if linked, _ := isGroupLinkedToRoute(account.Routes, groupID); linked {
+			return true
 		}
 	}
 
-	return false, nil
+	return false
 }

management/server/group/group.go

@@ -49,35 +49,3 @@ func (g *Group) Copy() *Group {
 func (g *Group) HasPeers() bool {
 	return len(g.Peers) > 0
 }
-
-// IsGroupAll checks if the group is a default "All" group.
-func (g *Group) IsGroupAll() bool {
-	return g.Name == "All"
-}
-
-// AddPeer adds peerID to Peers if not present, returning true if added.
-func (g *Group) AddPeer(peerID string) bool {
-	if peerID == "" {
-		return false
-	}
-
-	for _, itemID := range g.Peers {
-		if itemID == peerID {
-			return false
-		}
-	}
-
-	g.Peers = append(g.Peers, peerID)
-	return true
-}
-
-// RemovePeer removes peerID from Peers if present, returning true if removed.
-func (g *Group) RemovePeer(peerID string) bool {
-	for i, itemID := range g.Peers {
-		if itemID == peerID {
-			g.Peers = append(g.Peers[:i], g.Peers[i+1:]...)
-			return true
-		}
-	}
-	return false
-}

management/server/group/group_test.go

@@ -1,90 +0,0 @@
-package group
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestAddPeer(t *testing.T) {
-	t.Run("add new peer to empty slice", func(t *testing.T) {
-		group := &Group{Peers: []string{}}
-		peerID := "peer1"
-		assert.True(t, group.AddPeer(peerID))
-		assert.Contains(t, group.Peers, peerID)
-	})
-
-	t.Run("add new peer to nil slice", func(t *testing.T) {
-		group := &Group{Peers: nil}
-		peerID := "peer1"
-		assert.True(t, group.AddPeer(peerID))
-		assert.Contains(t, group.Peers, peerID)
-	})
-
-	t.Run("add new peer to non-empty slice", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2"}}
-		peerID := "peer3"
-		assert.True(t, group.AddPeer(peerID))
-		assert.Contains(t, group.Peers, peerID)
-	})
-
-	t.Run("add duplicate peer", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2"}}
-		peerID := "peer1"
-		assert.False(t, group.AddPeer(peerID))
-		assert.Equal(t, 2, len(group.Peers))
-	})
-
-	t.Run("add empty peer", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2"}}
-		peerID := ""
-		assert.False(t, group.AddPeer(peerID))
-		assert.Equal(t, 2, len(group.Peers))
-	})
-}
-
-func TestRemovePeer(t *testing.T) {
-	t.Run("remove existing peer from slice", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2", "peer3"}}
-		peerID := "peer2"
-		assert.True(t, group.RemovePeer(peerID))
-		assert.NotContains(t, group.Peers, peerID)
-		assert.Equal(t, 2, len(group.Peers))
-	})
-
-	t.Run("remove peer from empty slice", func(t *testing.T) {
-		group := &Group{Peers: []string{}}
-		peerID := "peer1"
-		assert.False(t, group.RemovePeer(peerID))
-		assert.Equal(t, 0, len(group.Peers))
-	})
-
-	t.Run("remove peer from nil slice", func(t *testing.T) {
-		group := &Group{Peers: nil}
-		peerID := "peer1"
-		assert.False(t, group.RemovePeer(peerID))
-		assert.Nil(t, group.Peers)
-	})
-
-	t.Run("remove non-existent peer", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2"}}
-		peerID := "peer3"
-		assert.False(t, group.RemovePeer(peerID))
-		assert.Equal(t, 2, len(group.Peers))
-	})
-
-	t.Run("remove peer from single-item slice", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1"}}
-		peerID := "peer1"
-		assert.True(t, group.RemovePeer(peerID))
-		assert.Equal(t, 0, len(group.Peers))
-		assert.NotContains(t, group.Peers, peerID)
-	})
-
-	t.Run("remove empty peer", func(t *testing.T) {
-		group := &Group{Peers: []string{"peer1", "peer2"}}
-		peerID := ""
-		assert.False(t, group.RemovePeer(peerID))
-		assert.Equal(t, 2, len(group.Peers))
-	})
-}

management/server/group_test.go

@@ -8,13 +8,12 @@ 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 (
@@ -500,7 +499,8 @@ func TestGroupAccountPeersUpdate(t *testing.T) {
 	})
 
 	// adding a group to policy
-	_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+	err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+		ID:      "policy",
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -511,7 +511,7 @@ func TestGroupAccountPeersUpdate(t *testing.T) {
 				Action:        PolicyTrafficActionAccept,
 			},
 		},
-	})
+	}, false)
 	assert.NoError(t, err)
 
 	// Saving a group linked to policy should update account peers and send peer update
@@ -536,6 +536,29 @@ 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/grpcserver.go

@@ -6,7 +6,6 @@ import (
 	"net"
 	"net/netip"
 	"strings"
-	"sync"
 	"time"
 
 	pb "github.com/golang/protobuf/proto" // nolint
@@ -39,7 +38,6 @@ type GRPCServer struct {
 	jwtClaimsExtractor *jwtclaims.ClaimsExtractor
 	appMetrics         telemetry.AppMetrics
 	ephemeralManager   *EphemeralManager
-	peerLocks          sync.Map
 }
 
 // NewServer creates a new Management server
@@ -150,13 +148,6 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 	// nolint:staticcheck
 	ctx = context.WithValue(ctx, nbContext.PeerIDKey, peerKey.String())
 
-	unlock := s.acquirePeerLockByUID(ctx, peerKey.String())
-	defer func() {
-		if unlock != nil {
-			unlock()
-		}
-	}()
-
 	accountID, err := s.accountManager.GetAccountIDForPeerKey(ctx, peerKey.String())
 	if err != nil {
 		// nolint:staticcheck
@@ -180,7 +171,6 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 
 	peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), extractPeerMeta(ctx, syncReq.GetMeta()), realIP)
 	if err != nil {
-		log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err)
 		return mapError(ctx, err)
 	}
 
@@ -200,15 +190,11 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 		s.appMetrics.GRPCMetrics().CountSyncRequestDuration(time.Since(reqStart))
 	}
 
-	unlock()
-	unlock = nil
-
 	return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv)
 }
 
 // handleUpdates sends updates to the connected peer until the updates channel is closed.
 func (s *GRPCServer) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *UpdateMessage, srv proto.ManagementService_SyncServer) error {
-	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
 	for {
 		select {
 		// condition when there are some updates
@@ -259,18 +245,10 @@ func (s *GRPCServer) sendUpdate(ctx context.Context, accountID string, peerKey w
 }
 
 func (s *GRPCServer) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer) {
-	unlock := s.acquirePeerLockByUID(ctx, peer.Key)
-	defer unlock()
-
-	err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
-	}
 	s.peersUpdateManager.CloseChannel(ctx, peer.ID)
 	s.secretsManager.CancelRefresh(peer.ID)
+	_ = s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key)
 	s.ephemeralManager.OnPeerDisconnected(ctx, peer)
-
-	log.WithContext(ctx).Tracef("peer %s has been disconnected", peer.Key)
 }
 
 func (s *GRPCServer) validateToken(ctx context.Context, jwtToken string) (string, error) {
@@ -296,24 +274,6 @@ func (s *GRPCServer) validateToken(ctx context.Context, jwtToken string) (string
 	return claims.UserId, nil
 }
 
-func (s *GRPCServer) acquirePeerLockByUID(ctx context.Context, uniqueID string) (unlock func()) {
-	log.WithContext(ctx).Tracef("acquiring peer lock for ID %s", uniqueID)
-
-	start := time.Now()
-	value, _ := s.peerLocks.LoadOrStore(uniqueID, &sync.RWMutex{})
-	mtx := value.(*sync.RWMutex)
-	mtx.Lock()
-	log.WithContext(ctx).Tracef("acquired peer lock for ID %s in %v", uniqueID, time.Since(start))
-	start = time.Now()
-
-	unlock = func() {
-		mtx.Unlock()
-		log.WithContext(ctx).Tracef("released peer lock for ID %s in %v", uniqueID, time.Since(start))
-	}
-
-	return unlock
-}
-
 // maps internal internalStatus.Error to gRPC status.Error
 func mapError(ctx context.Context, err error) error {
 	if e, ok := internalStatus.FromError(err); ok {

management/server/http/policies_handler.go

@@ -6,8 +6,10 @@ import (
 	"strconv"
 
 	"github.com/gorilla/mux"
-	"github.com/netbirdio/netbird/management/server"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
+	"github.com/rs/xid"
+
+	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/http/api"
 	"github.com/netbirdio/netbird/management/server/http/util"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
@@ -120,9 +122,14 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
 		return
 	}
 
-	policy := &server.Policy{
+	isUpdate := policyID != ""
+
+	if policyID == "" {
+		policyID = xid.New().String()
+	}
+
+	policy := server.Policy{
 		ID:          policyID,
-		AccountID:   accountID,
 		Name:        req.Name,
 		Enabled:     req.Enabled,
 		Description: req.Description,
@@ -130,7 +137,6 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
 	for _, rule := range req.Rules {
 		pr := server.PolicyRule{
 			ID:            policyID, // TODO: when policy can contain multiple rules, need refactor
-			PolicyID:      policyID,
 			Name:          rule.Name,
 			Destinations:  rule.Destinations,
 			Sources:       rule.Sources,
@@ -219,8 +225,7 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
 		policy.SourcePostureChecks = *req.SourcePostureChecks
 	}
 
-	policy, err := h.accountManager.SavePolicy(r.Context(), accountID, userID, policy)
-	if err != nil {
+	if err := h.accountManager.SavePolicy(r.Context(), accountID, userID, &policy, isUpdate); err != nil {
 		util.WriteError(r.Context(), err, w)
 		return
 	}
@@ -231,7 +236,7 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
 		return
 	}
 
-	resp := toPolicyResponse(allGroups, policy)
+	resp := toPolicyResponse(allGroups, &policy)
 	if len(resp.Rules) == 0 {
 		util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w)
 		return

management/server/http/policies_handler_test.go

@@ -38,12 +38,12 @@ func initPoliciesTestData(policies ...*server.Policy) *Policies {
 				}
 				return policy, nil
 			},
-			SavePolicyFunc: func(_ context.Context, _, _ string, policy *server.Policy) (*server.Policy, error) {
+			SavePolicyFunc: func(_ context.Context, _, _ string, policy *server.Policy, _ bool) error {
 				if !strings.HasPrefix(policy.ID, "id-") {
 					policy.ID = "id-was-set"
 					policy.Rules[0].ID = "id-was-set"
 				}
-				return policy, nil
+				return nil
 			},
 			GetAllGroupsFunc: func(ctx context.Context, accountID, userID string) ([]*nbgroup.Group, error) {
 				return []*nbgroup.Group{{ID: "F"}, {ID: "G"}}, nil

management/server/http/posture_checks_handler.go

@@ -169,8 +169,7 @@ func (p *PostureChecksHandler) savePostureChecks(w http.ResponseWriter, r *http.
 		return
 	}
 
-	postureChecks, err = p.accountManager.SavePostureChecks(r.Context(), accountID, userID, postureChecks)
-	if err != nil {
+	if err := p.accountManager.SavePostureChecks(r.Context(), accountID, userID, postureChecks); err != nil {
 		util.WriteError(r.Context(), err, w)
 		return
 	}

management/server/http/posture_checks_handler_test.go

@@ -40,15 +40,15 @@ func initPostureChecksTestData(postureChecks ...*posture.Checks) *PostureChecksH
 				}
 				return p, nil
 			},
-			SavePostureChecksFunc: func(_ context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error) {
+			SavePostureChecksFunc: func(_ context.Context, accountID, userID string, postureChecks *posture.Checks) error {
 				postureChecks.ID = "postureCheck"
 				testPostureChecks[postureChecks.ID] = postureChecks
 
 				if err := postureChecks.Validate(); err != nil {
-					return nil, status.Errorf(status.InvalidArgument, err.Error()) //nolint
+					return status.Errorf(status.InvalidArgument, err.Error()) //nolint
 				}
 
-				return postureChecks, nil
+				return nil
 			},
 			DeletePostureChecksFunc: func(_ context.Context, accountID, postureChecksID, userID string) error {
 				_, ok := testPostureChecks[postureChecksID]

management/server/http/routes_handler.go

@@ -149,7 +149,7 @@ func (h *RoutesHandler) validateRoute(req api.PostApiRoutesJSONRequestBody) erro
 	}
 
 	if req.Peer == nil && req.PeerGroups == nil {
-		return status.Errorf(status.InvalidArgument, "either 'peer' or 'peer_groups' should be provided")
+		return status.Errorf(status.InvalidArgument, "either 'peer' or 'peers_group' should be provided")
 	}
 
 	if req.Peer != nil && req.PeerGroups != nil {

management/server/integrated_validator.go

@@ -52,23 +52,26 @@ func (am *DefaultAccountManager) UpdateIntegratedValidatorGroups(ctx context.Con
 	return am.Store.SaveAccount(ctx, a)
 }
 
-func (am *DefaultAccountManager) GroupValidation(ctx context.Context, accountID string, groupIDs []string) (bool, error) {
-	if len(groupIDs) == 0 {
+func (am *DefaultAccountManager) GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error) {
+	if len(groups) == 0 {
 		return true, nil
 	}
-
-	err := am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		for _, groupID := range groupIDs {
-			_, err := transaction.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
-			if err != nil {
-				return err
-			}
-		}
-		return nil
-	})
+	accountsGroups, err := am.ListGroups(ctx, accountId)
 	if err != nil {
 		return false, err
 	}
+	for _, group := range groups {
+		var found bool
+		for _, accountGroup := range accountsGroups {
+			if accountGroup.ID == group {
+				found = true
+				break
+			}
+		}
+		if !found {
+			return false, nil
+		}
+	}
 
 	return true, nil
 }

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, bool, 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)
 	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, bool, error) {
-	return update, false, 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) GetValidatedPeers(accountID string, groups map[string]*group.Group, peers map[string]*nbpeer.Peer, extraSettings *account.ExtraSettings) (map[string]struct{}, error) {

management/server/mock_server/account_mock.go

@@ -45,11 +45,12 @@ type MockAccountManager struct {
 	SaveGroupsFunc                      func(ctx context.Context, accountID, userID string, groups []*group.Group) error
 	DeleteGroupFunc                     func(ctx context.Context, accountID, userId, groupID string) error
 	DeleteGroupsFunc                    func(ctx context.Context, accountId, userId string, groupIDs []string) error
+	ListGroupsFunc                      func(ctx context.Context, accountID string) ([]*group.Group, error)
 	GroupAddPeerFunc                    func(ctx context.Context, accountID, groupID, peerID string) error
 	GroupDeletePeerFunc                 func(ctx context.Context, accountID, groupID, peerID string) error
 	DeleteRuleFunc                      func(ctx context.Context, accountID, ruleID, userID string) error
 	GetPolicyFunc                       func(ctx context.Context, accountID, policyID, userID string) (*server.Policy, error)
-	SavePolicyFunc                      func(ctx context.Context, accountID, userID string, policy *server.Policy) (*server.Policy, error)
+	SavePolicyFunc                      func(ctx context.Context, accountID, userID string, policy *server.Policy, isUpdate bool) error
 	DeletePolicyFunc                    func(ctx context.Context, accountID, policyID, userID string) error
 	ListPoliciesFunc                    func(ctx context.Context, accountID, userID string) ([]*server.Policy, error)
 	GetUsersFromAccountFunc             func(ctx context.Context, accountID, userID string) ([]*server.UserInfo, error)
@@ -96,7 +97,7 @@ type MockAccountManager struct {
 	HasConnectedChannelFunc             func(peerID string) bool
 	GetExternalCacheManagerFunc         func() server.ExternalCacheManager
 	GetPostureChecksFunc                func(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error)
-	SavePostureChecksFunc               func(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error)
+	SavePostureChecksFunc               func(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) error
 	DeletePostureChecksFunc             func(ctx context.Context, accountID, postureChecksID, userID string) error
 	ListPostureChecksFunc               func(ctx context.Context, accountID, userID string) ([]*posture.Checks, error)
 	GetIdpManagerFunc                   func() idp.Manager
@@ -353,6 +354,14 @@ func (am *MockAccountManager) DeleteGroups(ctx context.Context, accountId, userI
 	return status.Errorf(codes.Unimplemented, "method DeleteGroups is not implemented")
 }
 
+// ListGroups mock implementation of ListGroups from server.AccountManager interface
+func (am *MockAccountManager) ListGroups(ctx context.Context, accountID string) ([]*group.Group, error) {
+	if am.ListGroupsFunc != nil {
+		return am.ListGroupsFunc(ctx, accountID)
+	}
+	return nil, status.Errorf(codes.Unimplemented, "method ListGroups is not implemented")
+}
+
 // GroupAddPeer mock implementation of GroupAddPeer from server.AccountManager interface
 func (am *MockAccountManager) GroupAddPeer(ctx context.Context, accountID, groupID, peerID string) error {
 	if am.GroupAddPeerFunc != nil {
@@ -386,11 +395,11 @@ func (am *MockAccountManager) GetPolicy(ctx context.Context, accountID, policyID
 }
 
 // SavePolicy mock implementation of SavePolicy from server.AccountManager interface
-func (am *MockAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *server.Policy) (*server.Policy, error) {
+func (am *MockAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *server.Policy, isUpdate bool) error {
 	if am.SavePolicyFunc != nil {
-		return am.SavePolicyFunc(ctx, accountID, userID, policy)
+		return am.SavePolicyFunc(ctx, accountID, userID, policy, isUpdate)
 	}
-	return nil, status.Errorf(codes.Unimplemented, "method SavePolicy is not implemented")
+	return status.Errorf(codes.Unimplemented, "method SavePolicy is not implemented")
 }
 
 // DeletePolicy mock implementation of DeletePolicy from server.AccountManager interface
@@ -730,11 +739,11 @@ func (am *MockAccountManager) GetPostureChecks(ctx context.Context, accountID, p
 }
 
 // SavePostureChecks mocks SavePostureChecks of the AccountManager interface
-func (am *MockAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error) {
+func (am *MockAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) error {
 	if am.SavePostureChecksFunc != nil {
 		return am.SavePostureChecksFunc(ctx, accountID, userID, postureChecks)
 	}
-	return nil, status.Errorf(codes.Unimplemented, "method SavePostureChecks is not implemented")
+	return status.Errorf(codes.Unimplemented, "method SavePostureChecks is not implemented")
 }
 
 // DeletePostureChecks mocks DeletePostureChecks of the AccountManager interface

management/server/nameserver.go

@@ -24,34 +24,26 @@ func (am *DefaultAccountManager) GetNameServerGroup(ctx context.Context, account
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can view name server groups")
 	}
 
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	return am.Store.GetNameServerGroupByID(ctx, LockingStrengthShare, accountID, nsGroupID)
+	return am.Store.GetNameServerGroupByID(ctx, LockingStrengthShare, nsGroupID, accountID)
 }
 
 // CreateNameServerGroup creates and saves a new nameserver group
 func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, accountID string, name, description string, nameServerList []nbdns.NameServer, groups []string, primary bool, domains []string, enabled bool, userID string, searchDomainEnabled bool) (*nbdns.NameServerGroup, error) {
+
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
 	newNSGroup := &nbdns.NameServerGroup{
 		ID:                   xid.New().String(),
-		AccountID:            accountID,
 		Name:                 name,
 		Description:          description,
 		NameServers:          nameServerList,
@@ -62,34 +54,27 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
 		SearchDomainsEnabled: searchDomainEnabled,
 	}
 
-	var updateAccountPeers bool
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateNameServerGroup(ctx, transaction, accountID, newNSGroup); err != nil {
-			return err
-		}
-
-		updateAccountPeers, err = anyGroupHasPeers(ctx, transaction, accountID, newNSGroup.Groups)
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.SaveNameServerGroup(ctx, LockingStrengthUpdate, newNSGroup)
-	})
+	err = validateNameServerGroup(false, newNSGroup, account)
 	if err != nil {
 		return nil, err
 	}
 
-	am.StoreEvent(ctx, userID, newNSGroup.ID, accountID, activity.NameserverGroupCreated, newNSGroup.EventMeta())
-
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+	if account.NameServerGroups == nil {
+		account.NameServerGroups = make(map[string]*nbdns.NameServerGroup)
 	}
 
+	account.NameServerGroups[newNSGroup.ID] = newNSGroup
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return nil, err
+	}
+
+	if anyGroupHasPeers(account, newNSGroup.Groups) {
+		am.updateAccountPeers(ctx, account)
+	}
+	am.StoreEvent(ctx, userID, newNSGroup.ID, accountID, activity.NameserverGroupCreated, newNSGroup.EventMeta())
+
 	return newNSGroup.Copy(), nil
 }
 
@@ -102,96 +87,59 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
 		return status.Errorf(status.InvalidArgument, "nameserver group provided is nil")
 	}
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	var updateAccountPeers bool
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		oldNSGroup, err := transaction.GetNameServerGroupByID(ctx, LockingStrengthShare, accountID, nsGroupToSave.ID)
-		if err != nil {
-			return err
-		}
-		nsGroupToSave.AccountID = accountID
-
-		if err = validateNameServerGroup(ctx, transaction, accountID, nsGroupToSave); err != nil {
-			return err
-		}
-
-		updateAccountPeers, err = areNameServerGroupChangesAffectPeers(ctx, transaction, nsGroupToSave, oldNSGroup)
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.SaveNameServerGroup(ctx, LockingStrengthUpdate, nsGroupToSave)
-	})
+	err = validateNameServerGroup(true, nsGroupToSave, account)
 	if err != nil {
 		return err
 	}
 
+	oldNSGroup := account.NameServerGroups[nsGroupToSave.ID]
+	account.NameServerGroups[nsGroupToSave.ID] = nsGroupToSave
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	if areNameServerGroupChangesAffectPeers(account, nsGroupToSave, oldNSGroup) {
+		am.updateAccountPeers(ctx, account)
+	}
 	am.StoreEvent(ctx, userID, nsGroupToSave.ID, accountID, activity.NameserverGroupUpdated, nsGroupToSave.EventMeta())
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
 	return nil
 }
 
 // DeleteNameServerGroup deletes nameserver group with nsGroupID
 func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, accountID, nsGroupID, userID string) error {
+
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
+	nsGroup := account.NameServerGroups[nsGroupID]
+	if nsGroup == nil {
+		return status.Errorf(status.NotFound, "nameserver group %s wasn't found", nsGroupID)
 	}
+	delete(account.NameServerGroups, nsGroupID)
 
-	var nsGroup *nbdns.NameServerGroup
-	var updateAccountPeers bool
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		nsGroup, err = transaction.GetNameServerGroupByID(ctx, LockingStrengthUpdate, accountID, nsGroupID)
-		if err != nil {
-			return err
-		}
-
-		updateAccountPeers, err = anyGroupHasPeers(ctx, transaction, accountID, nsGroup.Groups)
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.DeleteNameServerGroup(ctx, LockingStrengthUpdate, accountID, nsGroupID)
-	})
-	if err != nil {
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
 
+	if anyGroupHasPeers(account, nsGroup.Groups) {
+		am.updateAccountPeers(ctx, account)
+	}
 	am.StoreEvent(ctx, userID, nsGroup.ID, accountID, activity.NameserverGroupDeleted, nsGroup.EventMeta())
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
 	return nil
 }
 
@@ -202,62 +150,44 @@ func (am *DefaultAccountManager) ListNameServerGroups(ctx context.Context, accou
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can view name server groups")
 	}
 
 	return am.Store.GetAccountNameServerGroups(ctx, LockingStrengthShare, accountID)
 }
 
-func validateNameServerGroup(ctx context.Context, transaction Store, accountID string, nameserverGroup *nbdns.NameServerGroup) error {
+func validateNameServerGroup(existingGroup bool, nameserverGroup *nbdns.NameServerGroup, account *Account) error {
+	nsGroupID := ""
+	if existingGroup {
+		nsGroupID = nameserverGroup.ID
+		_, found := account.NameServerGroups[nsGroupID]
+		if !found {
+			return status.Errorf(status.NotFound, "nameserver group with ID %s was not found", nsGroupID)
+		}
+	}
+
 	err := validateDomainInput(nameserverGroup.Primary, nameserverGroup.Domains, nameserverGroup.SearchDomainsEnabled)
 	if err != nil {
 		return err
 	}
 
+	err = validateNSGroupName(nameserverGroup.Name, nsGroupID, account.NameServerGroups)
+	if err != nil {
+		return err
+	}
+
 	err = validateNSList(nameserverGroup.NameServers)
 	if err != nil {
 		return err
 	}
 
-	nsServerGroups, err := transaction.GetAccountNameServerGroups(ctx, LockingStrengthShare, accountID)
+	err = validateGroups(nameserverGroup.Groups, account.Groups)
 	if err != nil {
 		return err
 	}
 
-	err = validateNSGroupName(nameserverGroup.Name, nameserverGroup.ID, nsServerGroups)
-	if err != nil {
-		return err
-	}
-
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, nameserverGroup.Groups)
-	if err != nil {
-		return err
-	}
-
-	return validateGroups(nameserverGroup.Groups, groups)
-}
-
-// areNameServerGroupChangesAffectPeers checks if the changes in the nameserver group affect the peers.
-func areNameServerGroupChangesAffectPeers(ctx context.Context, transaction Store, newNSGroup, oldNSGroup *nbdns.NameServerGroup) (bool, error) {
-	if !newNSGroup.Enabled && !oldNSGroup.Enabled {
-		return false, nil
-	}
-
-	hasPeers, err := anyGroupHasPeers(ctx, transaction, newNSGroup.AccountID, newNSGroup.Groups)
-	if err != nil {
-		return false, err
-	}
-
-	if hasPeers {
-		return true, nil
-	}
-
-	return anyGroupHasPeers(ctx, transaction, oldNSGroup.AccountID, oldNSGroup.Groups)
+	return nil
 }
 
 func validateDomainInput(primary bool, domains []string, searchDomainsEnabled bool) error {
@@ -283,14 +213,14 @@ func validateDomainInput(primary bool, domains []string, searchDomainsEnabled bo
 	return nil
 }
 
-func validateNSGroupName(name, nsGroupID string, groups []*nbdns.NameServerGroup) error {
+func validateNSGroupName(name, nsGroupID string, nsGroupMap map[string]*nbdns.NameServerGroup) error {
 	if utf8.RuneCountInString(name) > nbdns.MaxGroupNameChar || name == "" {
 		return status.Errorf(status.InvalidArgument, "nameserver group name should be between 1 and %d", nbdns.MaxGroupNameChar)
 	}
 
-	for _, nsGroup := range groups {
+	for _, nsGroup := range nsGroupMap {
 		if name == nsGroup.Name && nsGroup.ID != nsGroupID {
-			return status.Errorf(status.InvalidArgument, "nameserver group with name %s already exist", name)
+			return status.Errorf(status.InvalidArgument, "a nameserver group with name %s already exist", name)
 		}
 	}
 
@@ -298,8 +228,8 @@ func validateNSGroupName(name, nsGroupID string, groups []*nbdns.NameServerGroup
 }
 
 func validateNSList(list []nbdns.NameServer) error {
-	nsListLength := len(list)
-	if nsListLength == 0 || nsListLength > 3 {
+	nsListLenght := len(list)
+	if nsListLenght == 0 || nsListLenght > 3 {
 		return status.Errorf(status.InvalidArgument, "the list of nameservers should be 1 or 3, got %d", len(list))
 	}
 	return nil
@@ -314,7 +244,14 @@ func validateGroups(list []string, groups map[string]*nbgroup.Group) error {
 		if id == "" {
 			return status.Errorf(status.InvalidArgument, "group ID should not be empty string")
 		}
-		if _, found := groups[id]; !found {
+		found := false
+		for groupID := range groups {
+			if id == groupID {
+				found = true
+				break
+			}
+		}
+		if !found {
 			return status.Errorf(status.InvalidArgument, "group id %s not found", id)
 		}
 	}
@@ -340,3 +277,11 @@ func validateDomain(domain string) error {
 
 	return nil
 }
+
+// areNameServerGroupChangesAffectPeers checks if the changes in the nameserver group affect the peers.
+func areNameServerGroupChangesAffectPeers(account *Account, newNSGroup, oldNSGroup *nbdns.NameServerGroup) bool {
+	if !newNSGroup.Enabled && !oldNSGroup.Enabled {
+		return false
+	}
+	return anyGroupHasPeers(account, newNSGroup.Groups) || anyGroupHasPeers(account, oldNSGroup.Groups)
+}

management/server/nameserver_test.go

@@ -1065,6 +1065,36 @@ 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
+	Serial uint64 `diff:"-"`
 
-	mu sync.Mutex `json:"-" gorm:"-"`
+	mu sync.Mutex `json:"-" gorm:"-" diff:"-"`
 }
 
 // NewNetwork creates a new Network initializing it with a Serial=0

management/server/peer.go

@@ -110,16 +110,14 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
 func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, connected bool, realIP net.IP, account *Account) error {
 	peer, err := account.FindPeerByPubKey(peerPubKey)
 	if err != nil {
-		return fmt.Errorf("failed to find peer by pub key: %w", err)
+		return err
 	}
 
 	expired, err := am.updatePeerStatusAndLocation(ctx, peer, connected, realIP, account)
 	if err != nil {
-		return fmt.Errorf("failed to update peer status and location: %w", err)
+		return err
 	}
 
-	log.WithContext(ctx).Debugf("mark peer %s connected: %t", peer.ID, connected)
-
 	if peer.AddedWithSSOLogin() {
 		if peer.LoginExpirationEnabled && account.Settings.PeerLoginExpirationEnabled {
 			am.checkAndSchedulePeerLoginExpiration(ctx, account)
@@ -133,7 +131,7 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
 	if expired {
 		// we need to update other peers because when peer login expires all other peers are notified to disconnect from
 		// the expired one. Here we notify them that connection is now allowed again.
-		am.updateAccountPeers(ctx, account.Id)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
@@ -170,7 +168,7 @@ func (am *DefaultAccountManager) updatePeerStatusAndLocation(ctx context.Context
 
 	err := am.Store.SavePeerStatus(account.Id, peer.ID, *newStatus)
 	if err != nil {
-		return false, fmt.Errorf("failed to save peer status: %w", err)
+		return false, err
 	}
 
 	return oldStatus.LoginExpired, nil
@@ -191,8 +189,7 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		return nil, status.Errorf(status.NotFound, "peer %s not found", update.ID)
 	}
 
-	var requiresPeerUpdates bool
-	update, requiresPeerUpdates, err = am.integratedPeerValidator.ValidatePeer(ctx, update, peer, userID, accountID, am.GetDNSDomain(), account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
+	update, err = am.integratedPeerValidator.ValidatePeer(ctx, update, peer, userID, accountID, am.GetDNSDomain(), account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
 		return nil, err
 	}
@@ -268,8 +265,8 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		return nil, err
 	}
 
-	if peerLabelUpdated || requiresPeerUpdates {
-		am.updateAccountPeers(ctx, accountID)
+	if peerLabelUpdated {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return peer, nil
@@ -333,10 +330,7 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 		return err
 	}
 
-	updateAccountPeers, err := am.isPeerInActiveGroup(ctx, account, peerID)
-	if err != nil {
-		return err
-	}
+	updateAccountPeers := isPeerInActiveGroup(account, peerID)
 
 	err = am.deletePeers(ctx, account, []string{peerID}, userID)
 	if err != nil {
@@ -349,7 +343,7 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 	}
 
 	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
@@ -556,7 +550,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 			return fmt.Errorf("failed to add peer to account: %w", err)
 		}
 
-		err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID)
+		err = transaction.IncrementNetworkSerial(ctx, accountID)
 		if err != nil {
 			return fmt.Errorf("failed to increment network serial: %w", err)
 		}
@@ -592,22 +586,11 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 
 	account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 	if err != nil {
-		return nil, nil, nil, status.NewGetAccountError(err)
+		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)
-
-	newGroupsAffectsPeers, err := areGroupChangesAffectPeers(ctx, am.Store, accountID, groupsToAdd)
-	if err != nil {
-		return nil, nil, nil, err
-	}
-
-	if newGroupsAffectsPeers {
-		am.updateAccountPeers(ctx, accountID)
+	if areGroupChangesAffectPeers(account, groupsToAdd) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
@@ -615,11 +598,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		return nil, nil, nil, err
 	}
 
-	postureChecks, err := am.getPeerPostureChecks(ctx, account.Id, newPeer.ID)
-	if err != nil {
-		return nil, nil, nil, err
-	}
-
+	postureChecks := am.getPeerPostureChecks(account, newPeer)
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	networkMap := account.GetPeerNetworkMap(ctx, newPeer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
 	return newPeer, networkMap, postureChecks, nil
@@ -654,7 +633,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if peer.UserID != "" {
 		user, err := account.FindUser(peer.UserID)
 		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to get user: %w", err)
+			return nil, nil, nil, err
 		}
 
 		err = checkIfPeerOwnerIsBlocked(peer, user)
@@ -671,17 +650,17 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if updated {
 		err = am.Store.SavePeer(ctx, account.Id, peer)
 		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to save peer: %w", err)
+			return nil, nil, nil, err
 		}
 
 		if sync.UpdateAccountPeers {
-			am.updateAccountPeers(ctx, account.Id)
+			am.updateAccountPeers(ctx, account)
 		}
 	}
 
 	peerNotValid, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to validate peer: %w", err)
+		return nil, nil, nil, err
 	}
 
 	var postureChecks []*posture.Checks
@@ -694,18 +673,14 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	}
 
 	if isStatusChanged {
-		am.updateAccountPeers(ctx, account.Id)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	validPeersMap, err := am.GetValidatedPeers(account)
-	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to get validated peers: %w", err)
-	}
-
-	postureChecks, err = am.getPeerPostureChecks(ctx, account.Id, peer.ID)
 	if err != nil {
 		return nil, nil, nil, err
 	}
+	postureChecks = am.getPeerPostureChecks(account, peer)
 
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, validPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
@@ -783,7 +758,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		}
 	}
 
-	groups, err := am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+	groups, err := am.Store.GetAccountGroups(ctx, accountID)
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -829,7 +804,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 	}
 
 	if updateRemotePeers || isStatusChanged {
-		am.updateAccountPeers(ctx, accountID)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return am.getValidatedPeerWithMap(ctx, isRequiresApproval, account, peer)
@@ -878,11 +853,7 @@ func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, is
 	if err != nil {
 		return nil, nil, nil, err
 	}
-
-	postureChecks, err = am.getPeerPostureChecks(ctx, account.Id, peer.ID)
-	if err != nil {
-		return nil, nil, nil, err
-	}
+	postureChecks = am.getPeerPostureChecks(account, peer)
 
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
@@ -996,7 +967,7 @@ func (am *DefaultAccountManager) GetPeer(ctx context.Context, accountID, peerID,
 
 // updateAccountPeers updates all peers that belong to an account.
 // Should be called when changes have to be synced to peers.
-func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, accountID string) {
+func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, account *Account) {
 	start := time.Now()
 	defer func() {
 		if am.metrics != nil {
@@ -1004,11 +975,6 @@ func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, account
 		}
 	}()
 
-	account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to send out updates to peers: %v", err)
-		return
-	}
 	peers := account.GetPeers()
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
@@ -1035,12 +1001,7 @@ func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, account
 			defer wg.Done()
 			defer func() { <-semaphore }()
 
-			postureChecks, err := am.getPeerPostureChecks(ctx, account.Id, p.ID)
-			if err != nil {
-				log.WithContext(ctx).Errorf("failed to send out updates to peers, failed to get peer: %s posture checks: %v", p.ID, err)
-				return
-			}
-
+			postureChecks := am.getPeerPostureChecks(account, p)
 			remotePeerNetworkMap := account.GetPeerNetworkMap(ctx, p.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
 			update := toSyncResponse(ctx, nil, p, nil, nil, remotePeerNetworkMap, am.GetDNSDomain(), postureChecks, dnsCache)
 			am.peersUpdateManager.SendUpdate(ctx, p.ID, &UpdateMessage{Update: update, NetworkMap: remotePeerNetworkMap})
@@ -1060,12 +1021,12 @@ func ConvertSliceToMap(existingLabels []string) map[string]struct{} {
 
 // IsPeerInActiveGroup checks if the given peer is part of a group that is used
 // in an active DNS, route, or ACL configuration.
-func (am *DefaultAccountManager) isPeerInActiveGroup(ctx context.Context, account *Account, peerID string) (bool, error) {
+func isPeerInActiveGroup(account *Account, peerID string) bool {
 	peerGroupIDs := make([]string, 0)
 	for _, group := range account.Groups {
 		if slices.Contains(group.Peers, peerID) {
 			peerGroupIDs = append(peerGroupIDs, group.ID)
 		}
 	}
-	return areGroupChangesAffectPeers(ctx, am.Store, account.Id, peerGroupIDs)
+	return areGroupChangesAffectPeers(account, peerGroupIDs)
 }

management/server/peer/peer.go

@@ -20,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_"`
+	Meta PeerSystemMeta `gorm:"embedded;embeddedPrefix:meta_" diff:"-"`
 	// 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_"`
+	Status *PeerStatus `gorm:"embedded;embeddedPrefix:peer_status_" diff:"-"`
 	// The user ID that registered the peer
-	UserID string
+	UserID string `diff:"-"`
 	// 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
+	LoginExpirationEnabled bool `diff:"-"`
 
-	InactivityExpirationEnabled bool
+	InactivityExpirationEnabled bool `diff:"-"`
 	// LastLogin the time when peer performed last login operation
-	LastLogin time.Time
+	LastLogin time.Time `diff:"-"`
 	// CreatedAt records the time the peer was created
-	CreatedAt time.Time
+	CreatedAt time.Time `diff:"-"`
 	// Indicate ephemeral peer attribute
-	Ephemeral bool
+	Ephemeral bool `diff:"-"`
 	// Geo location based on connection IP
-	Location Location `gorm:"embedded;embeddedPrefix:location_"`
+	Location Location `gorm:"embedded;embeddedPrefix:location_" diff:"-"`
 }
 
 type PeerStatus struct { //nolint:revive

management/server/peer_test.go

@@ -22,7 +22,6 @@ 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"
@@ -283,12 +282,14 @@ func TestAccountManager_GetNetworkMapWithPolicy(t *testing.T) {
 	var (
 		group1 nbgroup.Group
 		group2 nbgroup.Group
+		policy Policy
 	)
 
 	group1.ID = xid.New().String()
 	group2.ID = xid.New().String()
 	group1.Name = "src"
 	group2.Name = "dst"
+	policy.ID = xid.New().String()
 	group1.Peers = append(group1.Peers, peer1.ID)
 	group2.Peers = append(group2.Peers, peer2.ID)
 
@@ -303,20 +304,18 @@ func TestAccountManager_GetNetworkMapWithPolicy(t *testing.T) {
 		return
 	}
 
-	policy := &Policy{
-		Name:    "test",
-		Enabled: true,
-		Rules: []*PolicyRule{
-			{
-				Enabled:       true,
-				Sources:       []string{group1.ID},
-				Destinations:  []string{group2.ID},
-				Bidirectional: true,
-				Action:        PolicyTrafficActionAccept,
-			},
+	policy.Name = "test"
+	policy.Enabled = true
+	policy.Rules = []*PolicyRule{
+		{
+			Enabled:       true,
+			Sources:       []string{group1.ID},
+			Destinations:  []string{group2.ID},
+			Bidirectional: true,
+			Action:        PolicyTrafficActionAccept,
 		},
 	}
-	policy, err = manager.SavePolicy(context.Background(), account.Id, userID, policy)
+	err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 	if err != nil {
 		t.Errorf("expecting rule to be added, got failure %v", err)
 		return
@@ -364,7 +363,7 @@ func TestAccountManager_GetNetworkMapWithPolicy(t *testing.T) {
 	}
 
 	policy.Enabled = false
-	_, err = manager.SavePolicy(context.Background(), account.Id, userID, policy)
+	err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 	if err != nil {
 		t.Errorf("expecting rule to be added, got failure %v", err)
 		return
@@ -877,7 +876,7 @@ func BenchmarkUpdateAccountPeers(b *testing.B) {
 			start := time.Now()
 
 			for i := 0; i < b.N; i++ {
-				manager.updateAccountPeers(ctx, account.Id)
+				manager.updateAccountPeers(ctx, account)
 			}
 
 			duration := time.Since(start)
@@ -1399,53 +1398,10 @@ 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{
+		err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+			ID:      "policy",
 			Enabled: true,
 			Rules: []*PolicyRule{
 				{
@@ -1456,7 +1412,7 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-		})
+		}, false)
 		require.NoError(t, err)
 
 		done := make(chan struct{})

management/server/policy.go

@@ -3,13 +3,13 @@ package server
 import (
 	"context"
 	_ "embed"
+	"slices"
 	"strconv"
 	"strings"
 
-	"github.com/netbirdio/netbird/management/proto"
-	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
@@ -125,7 +125,6 @@ type PolicyRule struct {
 func (pm *PolicyRule) Copy() *PolicyRule {
 	rule := &PolicyRule{
 		ID:            pm.ID,
-		PolicyID:      pm.PolicyID,
 		Name:          pm.Name,
 		Description:   pm.Description,
 		Enabled:       pm.Enabled,
@@ -172,7 +171,6 @@ type Policy struct {
 func (p *Policy) Copy() *Policy {
 	c := &Policy{
 		ID:                  p.ID,
-		AccountID:           p.AccountID,
 		Name:                p.Name,
 		Description:         p.Description,
 		Enabled:             p.Enabled,
@@ -345,72 +343,44 @@ func (am *DefaultAccountManager) GetPolicy(ctx context.Context, accountID, polic
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power are allowed to view policies")
 	}
 
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	return am.Store.GetPolicyByID(ctx, LockingStrengthShare, accountID, policyID)
+	return am.Store.GetPolicyByID(ctx, LockingStrengthShare, policyID, accountID)
 }
 
 // SavePolicy in the store
-func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *Policy) (*Policy, error) {
+func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *Policy, isUpdate bool) error {
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	var isUpdate = policy.ID != ""
-	var updateAccountPeers bool
-	var action = activity.PolicyAdded
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validatePolicy(ctx, transaction, accountID, policy); err != nil {
-			return err
-		}
-
-		updateAccountPeers, err = arePolicyChangesAffectPeers(ctx, transaction, accountID, policy, isUpdate)
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		saveFunc := transaction.CreatePolicy
-		if isUpdate {
-			action = activity.PolicyUpdated
-			saveFunc = transaction.SavePolicy
-		}
-
-		return saveFunc(ctx, LockingStrengthUpdate, policy)
-	})
+	updateAccountPeers, err := am.savePolicy(account, policy, isUpdate)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	action := activity.PolicyAdded
+	if isUpdate {
+		action = activity.PolicyUpdated
+	}
 	am.StoreEvent(ctx, userID, policy.ID, accountID, action, policy.EventMeta())
 
 	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+		am.updateAccountPeers(ctx, account)
 	}
 
-	return policy, nil
+	return nil
 }
 
 // DeletePolicy from the store
@@ -418,136 +388,110 @@ func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, po
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
-	var policy *Policy
-	var updateAccountPeers bool
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		policy, err = transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policyID)
-		if err != nil {
-			return err
-		}
-
-		updateAccountPeers, err = arePolicyChangesAffectPeers(ctx, transaction, accountID, policy, false)
-		if err != nil {
-			return err
-		}
-
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
-
-		return transaction.DeletePolicy(ctx, LockingStrengthUpdate, accountID, policyID)
-	})
+	policy, err := am.deletePolicy(account, policyID)
 	if err != nil {
 		return err
 	}
 
-	am.StoreEvent(ctx, userID, policyID, accountID, activity.PolicyRemoved, policy.EventMeta())
-
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
 	}
 
+	am.StoreEvent(ctx, userID, policy.ID, accountID, activity.PolicyRemoved, policy.EventMeta())
+
+	am.updateAccountPeers(ctx, account)
+
 	return nil
 }
 
-// ListPolicies from the store.
+// ListPolicies from the store
 func (am *DefaultAccountManager) ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error) {
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power are allowed to view policies")
 	}
 
 	return am.Store.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
 }
 
-// arePolicyChangesAffectPeers checks if changes to a policy will affect any associated peers.
-func arePolicyChangesAffectPeers(ctx context.Context, transaction Store, accountID string, policy *Policy, isUpdate bool) (bool, error) {
-	if isUpdate {
-		existingPolicy, err := transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policy.ID)
-		if err != nil {
-			return false, err
+func (am *DefaultAccountManager) deletePolicy(account *Account, policyID string) (*Policy, error) {
+	policyIdx := -1
+	for i, policy := range account.Policies {
+		if policy.ID == policyID {
+			policyIdx = i
+			break
 		}
-
-		if !policy.Enabled && !existingPolicy.Enabled {
-			return false, nil
-		}
-
-		hasPeers, err := anyGroupHasPeers(ctx, transaction, policy.AccountID, existingPolicy.ruleGroups())
-		if err != nil {
-			return false, err
-		}
-
-		if hasPeers {
-			return true, nil
-		}
-
-		return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
+	}
+	if policyIdx < 0 {
+		return nil, status.Errorf(status.NotFound, "rule with ID %s doesn't exist", policyID)
 	}
 
-	return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
+	policy := account.Policies[policyIdx]
+	account.Policies = append(account.Policies[:policyIdx], account.Policies[policyIdx+1:]...)
+	return policy, nil
 }
 
-// validatePolicy validates the policy and its rules.
-func validatePolicy(ctx context.Context, transaction Store, accountID string, policy *Policy) error {
-	if policy.ID != "" {
-		_, err := transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policy.ID)
-		if err != nil {
-			return err
-		}
-	} else {
-		policy.ID = xid.New().String()
-		policy.AccountID = accountID
+// savePolicy saves or updates a policy in the given account.
+// If isUpdate is true, the function updates the existing policy; otherwise, it adds a new policy.
+func (am *DefaultAccountManager) savePolicy(account *Account, policyToSave *Policy, isUpdate bool) (bool, error) {
+	for index, rule := range policyToSave.Rules {
+		rule.Sources = filterValidGroupIDs(account, rule.Sources)
+		rule.Destinations = filterValidGroupIDs(account, rule.Destinations)
+		policyToSave.Rules[index] = rule
 	}
 
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, policy.ruleGroups())
-	if err != nil {
-		return err
+	if policyToSave.SourcePostureChecks != nil {
+		policyToSave.SourcePostureChecks = filterValidPostureChecks(account, policyToSave.SourcePostureChecks)
 	}
 
-	postureChecks, err := transaction.GetPostureChecksByIDs(ctx, LockingStrengthShare, accountID, policy.SourcePostureChecks)
-	if err != nil {
-		return err
-	}
-
-	for i, rule := range policy.Rules {
-		ruleCopy := rule.Copy()
-		if ruleCopy.ID == "" {
-			ruleCopy.ID = xid.New().String()
-			ruleCopy.PolicyID = policy.ID
+	if isUpdate {
+		policyIdx := slices.IndexFunc(account.Policies, func(policy *Policy) bool { return policy.ID == policyToSave.ID })
+		if policyIdx < 0 {
+			return false, status.Errorf(status.NotFound, "couldn't find policy id %s", policyToSave.ID)
 		}
 
-		ruleCopy.Sources = getValidGroupIDs(groups, ruleCopy.Sources)
-		ruleCopy.Destinations = getValidGroupIDs(groups, ruleCopy.Destinations)
-		policy.Rules[i] = ruleCopy
+		oldPolicy := account.Policies[policyIdx]
+		// Update the existing policy
+		account.Policies[policyIdx] = policyToSave
+
+		if !policyToSave.Enabled && !oldPolicy.Enabled {
+			return false, nil
+		}
+		updateAccountPeers := anyGroupHasPeers(account, oldPolicy.ruleGroups()) || anyGroupHasPeers(account, policyToSave.ruleGroups())
+
+		return updateAccountPeers, nil
 	}
 
-	if policy.SourcePostureChecks != nil {
-		policy.SourcePostureChecks = getValidPostureCheckIDs(postureChecks, policy.SourcePostureChecks)
-	}
+	// Add the new policy to the account
+	account.Policies = append(account.Policies, policyToSave)
 
-	return nil
+	return anyGroupHasPeers(account, policyToSave.ruleGroups()), nil
+}
+
+func toProtocolFirewallRules(rules []*FirewallRule) []*proto.FirewallRule {
+	result := make([]*proto.FirewallRule, len(rules))
+	for i := range rules {
+		rule := rules[i]
+
+		result[i] = &proto.FirewallRule{
+			PeerIP:    rule.PeerIP,
+			Direction: getProtoDirection(rule.Direction),
+			Action:    getProtoAction(rule.Action),
+			Protocol:  getProtoProtocol(rule.Protocol),
+			Port:      rule.Port,
+		}
+	}
+	return result
 }
 
 // getAllPeersFromGroups for given peer ID and list of groups
@@ -628,42 +572,27 @@ func (a *Account) getPostureChecks(postureChecksID string) *posture.Checks {
 	return nil
 }
 
-// getValidPostureCheckIDs filters and returns only the valid posture check IDs from the provided list.
-func getValidPostureCheckIDs(postureChecks map[string]*posture.Checks, postureChecksIds []string) []string {
-	validIDs := make([]string, 0, len(postureChecksIds))
+// filterValidPostureChecks filters and returns the posture check IDs from the given list
+// that are valid within the provided account.
+func filterValidPostureChecks(account *Account, postureChecksIds []string) []string {
+	result := make([]string, 0, len(postureChecksIds))
 	for _, id := range postureChecksIds {
-		if _, exists := postureChecks[id]; exists {
-			validIDs = append(validIDs, id)
-		}
-	}
-
-	return validIDs
-}
-
-// getValidGroupIDs filters and returns only the valid group IDs from the provided list.
-func getValidGroupIDs(groups map[string]*nbgroup.Group, groupIDs []string) []string {
-	validIDs := make([]string, 0, len(groupIDs))
-	for _, id := range groupIDs {
-		if _, exists := groups[id]; exists {
-			validIDs = append(validIDs, id)
-		}
-	}
-
-	return validIDs
-}
-
-// toProtocolFirewallRules converts the firewall rules to the protocol firewall rules.
-func toProtocolFirewallRules(rules []*FirewallRule) []*proto.FirewallRule {
-	result := make([]*proto.FirewallRule, len(rules))
-	for i := range rules {
-		rule := rules[i]
-
-		result[i] = &proto.FirewallRule{
-			PeerIP:    rule.PeerIP,
-			Direction: getProtoDirection(rule.Direction),
-			Action:    getProtoAction(rule.Action),
-			Protocol:  getProtoProtocol(rule.Protocol),
-			Port:      rule.Port,
+		for _, postureCheck := range account.PostureChecks {
+			if id == postureCheck.ID {
+				result = append(result, id)
+				continue
+			}
+		}
+	}
+	return result
+}
+
+// filterValidGroupIDs filters a list of group IDs and returns only the ones present in the account's group map.
+func filterValidGroupIDs(account *Account, groupIDs []string) []string {
+	result := make([]string, 0, len(groupIDs))
+	for _, groupID := range groupIDs {
+		if _, exists := account.Groups[groupID]; exists {
+			result = append(result, groupID)
 		}
 	}
 	return result

management/server/policy_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/rs/xid"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/exp/slices"
 
@@ -853,28 +854,24 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	})
 	assert.NoError(t, err)
 
-	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
+	updMsg1 := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
 	t.Cleanup(func() {
 		manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
 	})
 
-	var policyWithGroupRulesNoPeers *Policy
-	var policyWithDestinationPeersOnly *Policy
-	var policyWithSourceAndDestinationPeers *Policy
+	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) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		policyWithGroupRulesNoPeers, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-			AccountID: account.Id,
-			Enabled:   true,
+		policy := Policy{
+			ID:      "policy-rule-groups-no-peers",
+			Enabled: true,
 			Rules: []*PolicyRule{
 				{
+					ID:            xid.New().String(),
 					Enabled:       true,
 					Sources:       []string{"groupB"},
 					Destinations:  []string{"groupC"},
@@ -882,7 +879,15 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-		})
+		}
+
+		done := make(chan struct{})
+		go func() {
+			peerShouldNotReceiveUpdate(t, updMsg1)
+			close(done)
+		}()
+
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		select {
@@ -895,17 +900,12 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Saving policy with source group containing peers, but destination group without peers should
 	// update account's peers and send peer update
 	t.Run("saving policy where source has peers but destination does not", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-			AccountID: account.Id,
-			Enabled:   true,
+		policy := Policy{
+			ID:      "policy-source-has-peers-destination-none",
+			Enabled: true,
 			Rules: []*PolicyRule{
 				{
+					ID:            xid.New().String(),
 					Enabled:       true,
 					Sources:       []string{"groupA"},
 					Destinations:  []string{"groupB"},
@@ -914,7 +914,15 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-		})
+		}
+
+		done := make(chan struct{})
+		go func() {
+			peerShouldReceiveUpdate(t, updMsg1)
+			close(done)
+		}()
+
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		select {
@@ -927,18 +935,13 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Saving policy with destination group containing peers, but source group without peers should
 	// update account's peers and send peer update
 	t.Run("saving policy where destination has peers but source does not", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		policyWithDestinationPeersOnly, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-			AccountID: account.Id,
-			Enabled:   true,
+		policy := Policy{
+			ID:      "policy-destination-has-peers-source-none",
+			Enabled: true,
 			Rules: []*PolicyRule{
 				{
-					Enabled:       true,
+					ID:            xid.New().String(),
+					Enabled:       false,
 					Sources:       []string{"groupC"},
 					Destinations:  []string{"groupD"},
 					Bidirectional: true,
@@ -946,7 +949,15 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-		})
+		}
+
+		done := make(chan struct{})
+		go func() {
+			peerShouldReceiveUpdate(t, updMsg2)
+			close(done)
+		}()
+
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		select {
@@ -959,17 +970,12 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Saving policy with destination and source groups containing peers should update account's peers
 	// and send peer update
 	t.Run("saving policy with source and destination groups with peers", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		policyWithSourceAndDestinationPeers, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-			AccountID: account.Id,
-			Enabled:   true,
+		policy := Policy{
+			ID:      "policy-source-destination-peers",
+			Enabled: true,
 			Rules: []*PolicyRule{
 				{
+					ID:            xid.New().String(),
 					Enabled:       true,
 					Sources:       []string{"groupA"},
 					Destinations:  []string{"groupD"},
@@ -977,7 +983,15 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-		})
+		}
+
+		done := make(chan struct{})
+		go func() {
+			peerShouldReceiveUpdate(t, updMsg1)
+			close(done)
+		}()
+
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		select {
@@ -990,14 +1004,28 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Disabling policy with destination and source groups containing peers should update account's peers
 	// and send peer update
 	t.Run("disabling policy with source and destination groups with peers", func(t *testing.T) {
+		policy := Policy{
+			ID:      "policy-source-destination-peers",
+			Enabled: false,
+			Rules: []*PolicyRule{
+				{
+					ID:            xid.New().String(),
+					Enabled:       true,
+					Sources:       []string{"groupA"},
+					Destinations:  []string{"groupD"},
+					Bidirectional: true,
+					Action:        PolicyTrafficActionAccept,
+				},
+			},
+		}
+
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
+			peerShouldReceiveUpdate(t, updMsg1)
 			close(done)
 		}()
 
-		policyWithSourceAndDestinationPeers.Enabled = false
-		policyWithSourceAndDestinationPeers, err = manager.SavePolicy(context.Background(), account.Id, userID, policyWithSourceAndDestinationPeers)
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 		assert.NoError(t, err)
 
 		select {
@@ -1010,15 +1038,29 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Updating disabled policy with destination and source groups containing peers should not update account's peers
 	// or send peer update
 	t.Run("updating disabled policy with source and destination groups with peers", func(t *testing.T) {
+		policy := Policy{
+			ID:          "policy-source-destination-peers",
+			Description: "updated description",
+			Enabled:     false,
+			Rules: []*PolicyRule{
+				{
+					ID:            xid.New().String(),
+					Enabled:       true,
+					Sources:       []string{"groupA"},
+					Destinations:  []string{"groupA"},
+					Bidirectional: true,
+					Action:        PolicyTrafficActionAccept,
+				},
+			},
+		}
+
 		done := make(chan struct{})
 		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
+			peerShouldNotReceiveUpdate(t, updMsg1)
 			close(done)
 		}()
 
-		policyWithSourceAndDestinationPeers.Description = "updated description"
-		policyWithSourceAndDestinationPeers.Rules[0].Destinations = []string{"groupA"}
-		policyWithSourceAndDestinationPeers, err = manager.SavePolicy(context.Background(), account.Id, userID, policyWithSourceAndDestinationPeers)
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 		assert.NoError(t, err)
 
 		select {
@@ -1031,14 +1073,28 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Enabling policy with destination and source groups containing peers should update account's peers
 	// and send peer update
 	t.Run("enabling policy with source and destination groups with peers", 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() {
-			peerShouldReceiveUpdate(t, updMsg)
+			peerShouldReceiveUpdate(t, updMsg1)
 			close(done)
 		}()
 
-		policyWithSourceAndDestinationPeers.Enabled = true
-		policyWithSourceAndDestinationPeers, err = manager.SavePolicy(context.Background(), account.Id, userID, policyWithSourceAndDestinationPeers)
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 		assert.NoError(t, err)
 
 		select {
@@ -1048,15 +1104,50 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// 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) {
+	// 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() {
-			peerShouldReceiveUpdate(t, updMsg)
+			peerShouldNotReceiveUpdate(t, updMsg1)
 			close(done)
 		}()
 
-		err := manager.DeletePolicy(context.Background(), account.Id, policyWithSourceAndDestinationPeers.ID, userID)
+		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)
+			close(done)
+		}()
+
+		err := manager.DeletePolicy(context.Background(), account.Id, policyID, userID)
 		assert.NoError(t, err)
 
 		select {
@@ -1070,13 +1161,14 @@ func TestPolicyAccountPeersUpdate(t *testing.T) {
 	// Deleting policy with destination group containing peers, but source group without peers should
 	// update account's peers and send peer update
 	t.Run("deleting policy where destination has peers but source does not", func(t *testing.T) {
+		policyID := "policy-destination-has-peers-source-none"
 		done := make(chan struct{})
 		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
+			peerShouldReceiveUpdate(t, updMsg2)
 			close(done)
 		}()
 
-		err := manager.DeletePolicy(context.Background(), account.Id, policyWithDestinationPeersOnly.ID, userID)
+		err := manager.DeletePolicy(context.Background(), account.Id, policyID, userID)
 		assert.NoError(t, err)
 
 		select {
@@ -1088,13 +1180,14 @@ 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
 		done := make(chan struct{})
 		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
+			peerShouldNotReceiveUpdate(t, updMsg1)
 			close(done)
 		}()
 
-		err := manager.DeletePolicy(context.Background(), account.Id, policyWithGroupRulesNoPeers.ID, userID)
+		err := manager.DeletePolicy(context.Background(), account.Id, policyID, userID)
 		assert.NoError(t, err)
 
 		select {

management/server/posture/checks.go

@@ -7,6 +7,8 @@ import (
 	"regexp"
 
 	"github.com/hashicorp/go-version"
+	"github.com/rs/xid"
+
 	"github.com/netbirdio/netbird/management/server/http/api"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/status"
@@ -170,6 +172,10 @@ func NewChecksFromAPIPostureCheckUpdate(source api.PostureCheckUpdate, postureCh
 }
 
 func buildPostureCheck(postureChecksID string, name string, description string, checks api.Checks) (*Checks, error) {
+	if postureChecksID == "" {
+		postureChecksID = xid.New().String()
+	}
+
 	postureChecks := Checks{
 		ID:          postureChecksID,
 		Name:        name,

management/server/posture_checks.go

@@ -2,15 +2,16 @@ package server
 
 import (
 	"context"
-	"fmt"
 	"slices"
 
 	"github.com/netbirdio/netbird/management/server/activity"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/status"
-	"github.com/rs/xid"
-	log "github.com/sirupsen/logrus"
-	"golang.org/x/exp/maps"
+)
+
+const (
+	errMsgPostureAdminOnly = "only users with admin power are allowed to view posture checks"
 )
 
 func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error) {
@@ -19,279 +20,219 @@ func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
+	if !user.HasAdminPower() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
 	}
 
-	if !user.HasAdminPower() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	return am.Store.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, postureChecksID)
+	return am.Store.GetPostureChecksByID(ctx, LockingStrengthShare, postureChecksID, accountID)
 }
 
-// SavePostureChecks saves a posture check.
-func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error) {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
-	if err != nil {
-		return nil, err
-	}
+func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if !user.HasAdminPower() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	var updateAccountPeers bool
-	var isUpdate = postureChecks.ID != ""
-	var action = activity.PostureCheckCreated
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validatePostureChecks(ctx, transaction, accountID, postureChecks); err != nil {
-			return err
-		}
-
-		if isUpdate {
-			updateAccountPeers, err = arePostureCheckChangesAffectPeers(ctx, transaction, accountID, postureChecks.ID)
-			if err != nil {
-				return err
-			}
-
-			if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-				return err
-			}
-
-			action = activity.PostureCheckUpdated
-		}
-
-		postureChecks.AccountID = accountID
-		return transaction.SavePostureChecks(ctx, LockingStrengthUpdate, postureChecks)
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	am.StoreEvent(ctx, userID, postureChecks.ID, accountID, action, postureChecks.EventMeta())
-
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
-	return postureChecks, nil
-}
-
-// DeletePostureChecks deletes a posture check by ID.
-func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accountID, postureChecksID, userID string) error {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
+	user, err := account.FindUser(userID)
+	if err != nil {
+		return err
 	}
 
 	if !user.HasAdminPower() {
-		return status.NewAdminPermissionError()
+		return status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
 	}
 
-	var postureChecks *posture.Checks
+	if err := postureChecks.Validate(); err != nil {
+		return status.Errorf(status.InvalidArgument, err.Error()) //nolint
+	}
 
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		postureChecks, err = transaction.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, postureChecksID)
-		if err != nil {
-			return err
-		}
+	exists, uniqName := am.savePostureChecks(account, postureChecks)
 
-		if err = isPostureCheckLinkedToPolicy(ctx, transaction, postureChecksID, accountID); err != nil {
-			return err
-		}
+	// we do not allow create new posture checks with non uniq name
+	if !exists && !uniqName {
+		return status.Errorf(status.PreconditionFailed, "Posture check name should be unique")
+	}
 
-		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return err
-		}
+	action := activity.PostureCheckCreated
+	if exists {
+		action = activity.PostureCheckUpdated
+		account.Network.IncSerial()
+	}
 
-		return transaction.DeletePostureChecks(ctx, LockingStrengthUpdate, accountID, postureChecksID)
-	})
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, userID, postureChecks.ID, accountID, action, postureChecks.EventMeta())
+
+	if arePostureCheckChangesAffectingPeers(account, postureChecks.ID, exists) {
+		am.updateAccountPeers(ctx, account)
+	}
+
+	return nil
+}
+
+func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accountID, postureChecksID, userID string) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
+	user, err := account.FindUser(userID)
+	if err != nil {
+		return err
+	}
+
+	if !user.HasAdminPower() {
+		return status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
+	}
+
+	postureChecks, err := am.deletePostureChecks(account, postureChecksID)
+	if err != nil {
+		return err
+	}
+
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
 	am.StoreEvent(ctx, userID, postureChecks.ID, accountID, activity.PostureCheckDeleted, postureChecks.EventMeta())
 
 	return nil
 }
 
-// ListPostureChecks returns a list of posture checks.
 func (am *DefaultAccountManager) ListPostureChecks(ctx context.Context, accountID, userID string) ([]*posture.Checks, error) {
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if !user.HasAdminPower() {
-		return nil, status.NewAdminPermissionError()
+	if !user.HasAdminPower() || user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
 	}
 
 	return am.Store.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
 }
 
+func (am *DefaultAccountManager) savePostureChecks(account *Account, postureChecks *posture.Checks) (exists, uniqName bool) {
+	uniqName = true
+	for i, p := range account.PostureChecks {
+		if !exists && p.ID == postureChecks.ID {
+			account.PostureChecks[i] = postureChecks
+			exists = true
+		}
+		if p.Name == postureChecks.Name {
+			uniqName = false
+		}
+	}
+	if !exists {
+		account.PostureChecks = append(account.PostureChecks, postureChecks)
+	}
+	return
+}
+
+func (am *DefaultAccountManager) deletePostureChecks(account *Account, postureChecksID string) (*posture.Checks, error) {
+	postureChecksIdx := -1
+	for i, postureChecks := range account.PostureChecks {
+		if postureChecks.ID == postureChecksID {
+			postureChecksIdx = i
+			break
+		}
+	}
+	if postureChecksIdx < 0 {
+		return nil, status.Errorf(status.NotFound, "posture checks with ID %s doesn't exist", postureChecksID)
+	}
+
+	// Check if posture check is linked to any policy
+	if isLinked, linkedPolicy := isPostureCheckLinkedToPolicy(account, postureChecksID); isLinked {
+		return nil, status.Errorf(status.PreconditionFailed, "posture checks have been linked to policy: %s", linkedPolicy.Name)
+	}
+
+	postureChecks := account.PostureChecks[postureChecksIdx]
+	account.PostureChecks = append(account.PostureChecks[:postureChecksIdx], account.PostureChecks[postureChecksIdx+1:]...)
+
+	return postureChecks, nil
+}
+
 // getPeerPostureChecks returns the posture checks applied for a given peer.
-func (am *DefaultAccountManager) getPeerPostureChecks(ctx context.Context, accountID string, peerID string) ([]*posture.Checks, error) {
-	peerPostureChecks := make(map[string]*posture.Checks)
+func (am *DefaultAccountManager) getPeerPostureChecks(account *Account, peer *nbpeer.Peer) []*posture.Checks {
+	peerPostureChecks := make(map[string]posture.Checks)
 
-	err := am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		postureChecks, err := transaction.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
-		if err != nil {
-			return err
-		}
-
-		if len(postureChecks) == 0 {
-			return nil
-		}
-
-		policies, err := transaction.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
-		if err != nil {
-			return err
-		}
-
-		for _, policy := range policies {
-			if !policy.Enabled {
-				continue
-			}
-
-			if err = addPolicyPostureChecks(ctx, transaction, accountID, peerID, policy, peerPostureChecks); err != nil {
-				return err
-			}
-		}
-
-		return nil
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	return maps.Values(peerPostureChecks), nil
-}
-
-// arePostureCheckChangesAffectPeers checks if the changes in posture checks are affecting peers.
-func arePostureCheckChangesAffectPeers(ctx context.Context, transaction Store, accountID, postureCheckID string) (bool, error) {
-	policies, err := transaction.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		return false, err
-	}
-
-	for _, policy := range policies {
-		if slices.Contains(policy.SourcePostureChecks, postureCheckID) {
-			hasPeers, err := anyGroupHasPeers(ctx, transaction, accountID, policy.ruleGroups())
-			if err != nil {
-				return false, err
-			}
-
-			if hasPeers {
-				return true, nil
-			}
-		}
-	}
-
-	return false, nil
-}
-
-// validatePostureChecks validates the posture checks.
-func validatePostureChecks(ctx context.Context, transaction Store, accountID string, postureChecks *posture.Checks) error {
-	if err := postureChecks.Validate(); err != nil {
-		return status.Errorf(status.InvalidArgument, err.Error()) //nolint
-	}
-
-	// If the posture check already has an ID, verify its existence in the store.
-	if postureChecks.ID != "" {
-		if _, err := transaction.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, postureChecks.ID); err != nil {
-			return err
-		}
+	if len(account.PostureChecks) == 0 {
 		return nil
 	}
 
-	// For new posture checks, ensure no duplicates by name.
-	checks, err := transaction.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		return err
-	}
+	for _, policy := range account.Policies {
+		if !policy.Enabled {
+			continue
+		}
 
-	for _, check := range checks {
-		if check.Name == postureChecks.Name && check.ID != postureChecks.ID {
-			return status.Errorf(status.InvalidArgument, "posture checks with name %s already exists", postureChecks.Name)
+		if isPeerInPolicySourceGroups(peer.ID, account, policy) {
+			addPolicyPostureChecks(account, policy, peerPostureChecks)
 		}
 	}
 
-	postureChecks.ID = xid.New().String()
-
-	return nil
-}
-
-// addPolicyPostureChecks adds posture checks from a policy to the peer posture checks map if the peer is in the policy's source groups.
-func addPolicyPostureChecks(ctx context.Context, transaction Store, accountID, peerID string, policy *Policy, peerPostureChecks map[string]*posture.Checks) error {
-	isInGroup, err := isPeerInPolicySourceGroups(ctx, transaction, accountID, peerID, policy)
-	if err != nil {
-		return err
+	postureChecksList := make([]*posture.Checks, 0, len(peerPostureChecks))
+	for _, check := range peerPostureChecks {
+		checkCopy := check
+		postureChecksList = append(postureChecksList, &checkCopy)
 	}
 
-	if !isInGroup {
-		return nil
-	}
-
-	for _, sourcePostureCheckID := range policy.SourcePostureChecks {
-		postureCheck, err := transaction.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, sourcePostureCheckID)
-		if err != nil {
-			return err
-		}
-		peerPostureChecks[sourcePostureCheckID] = postureCheck
-	}
-
-	return nil
+	return postureChecksList
 }
 
 // isPeerInPolicySourceGroups checks if a peer is present in any of the policy rule source groups.
-func isPeerInPolicySourceGroups(ctx context.Context, transaction Store, accountID, peerID string, policy *Policy) (bool, error) {
+func isPeerInPolicySourceGroups(peerID string, account *Account, policy *Policy) bool {
 	for _, rule := range policy.Rules {
 		if !rule.Enabled {
 			continue
 		}
 
 		for _, sourceGroup := range rule.Sources {
-			group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, sourceGroup)
-			if err != nil {
-				log.WithContext(ctx).Debugf("failed to check peer in policy source group: %v", err)
-				return false, fmt.Errorf("failed to check peer in policy source group: %w", err)
-			}
-
-			if slices.Contains(group.Peers, peerID) {
-				return true, nil
+			group, ok := account.Groups[sourceGroup]
+			if ok && slices.Contains(group.Peers, peerID) {
+				return true
 			}
 		}
 	}
 
+	return false
+}
+
+func addPolicyPostureChecks(account *Account, policy *Policy, peerPostureChecks map[string]posture.Checks) {
+	for _, sourcePostureCheckID := range policy.SourcePostureChecks {
+		for _, postureCheck := range account.PostureChecks {
+			if postureCheck.ID == sourcePostureCheckID {
+				peerPostureChecks[sourcePostureCheckID] = *postureCheck
+			}
+		}
+	}
+}
+
+func isPostureCheckLinkedToPolicy(account *Account, postureChecksID string) (bool, *Policy) {
+	for _, policy := range account.Policies {
+		if slices.Contains(policy.SourcePostureChecks, postureChecksID) {
+			return true, policy
+		}
+	}
 	return false, nil
 }
 
-// isPostureCheckLinkedToPolicy checks whether the posture check is linked to any account policy.
-func isPostureCheckLinkedToPolicy(ctx context.Context, transaction Store, postureChecksID, accountID string) error {
-	policies, err := transaction.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		return err
+// arePostureCheckChangesAffectingPeers checks if the changes in posture checks are affecting peers.
+func arePostureCheckChangesAffectingPeers(account *Account, postureCheckID string, exists bool) bool {
+	if !exists {
+		return false
 	}
 
-	for _, policy := range policies {
-		if slices.Contains(policy.SourcePostureChecks, postureChecksID) {
-			return status.Errorf(status.PreconditionFailed, "posture checks have been linked to policy: %s", policy.Name)
-		}
+	isLinked, linkedPolicy := isPostureCheckLinkedToPolicy(account, postureCheckID)
+	if !isLinked {
+		return false
 	}
-
-	return nil
+	return anyGroupHasPeers(account, linkedPolicy.ruleGroups())
 }

management/server/posture_checks_test.go

@@ -5,10 +5,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	"github.com/netbirdio/netbird/management/server/group"
+	"github.com/rs/xid"
+	"github.com/stretchr/testify/assert"
 
 	"github.com/netbirdio/netbird/management/server/posture"
 )
@@ -16,6 +15,7 @@ import (
 const (
 	adminUserID      = "adminUserID"
 	regularUserID    = "regularUserID"
+	postureCheckID   = "existing-id"
 	postureCheckName = "Existing check"
 )
 
@@ -32,7 +32,7 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 
 	t.Run("Generic posture check flow", func(t *testing.T) {
 		// regular users can not create checks
-		_, err = am.SavePostureChecks(context.Background(), account.Id, regularUserID, &posture.Checks{})
+		err := am.SavePostureChecks(context.Background(), account.Id, regularUserID, &posture.Checks{})
 		assert.Error(t, err)
 
 		// regular users cannot list check
@@ -40,7 +40,8 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 		assert.Error(t, err)
 
 		// should be possible to create posture check with uniq name
-		postureCheck, err := am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+		err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+			ID:   postureCheckID,
 			Name: postureCheckName,
 			Checks: posture.ChecksDefinition{
 				NBVersionCheck: &posture.NBVersionCheck{
@@ -56,7 +57,8 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 		assert.Len(t, checks, 1)
 
 		// should not be possible to create posture check with non uniq name
-		_, err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+		err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+			ID:   "new-id",
 			Name: postureCheckName,
 			Checks: posture.ChecksDefinition{
 				GeoLocationCheck: &posture.GeoLocationCheck{
@@ -71,20 +73,23 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 		assert.Error(t, err)
 
 		// admins can update posture checks
-		postureCheck.Checks = posture.ChecksDefinition{
-			NBVersionCheck: &posture.NBVersionCheck{
-				MinVersion: "0.27.0",
+		err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+			ID:   postureCheckID,
+			Name: postureCheckName,
+			Checks: posture.ChecksDefinition{
+				NBVersionCheck: &posture.NBVersionCheck{
+					MinVersion: "0.27.0",
+				},
 			},
-		}
-		_, err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheck)
+		})
 		assert.NoError(t, err)
 
 		// users should not be able to delete posture checks
-		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheck.ID, regularUserID)
+		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheckID, regularUserID)
 		assert.Error(t, err)
 
 		// admin should be able to delete posture checks
-		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheck.ID, adminUserID)
+		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheckID, adminUserID)
 		assert.NoError(t, err)
 		checks, err = am.ListPostureChecks(context.Background(), account.Id, adminUserID)
 		assert.NoError(t, err)
@@ -144,22 +149,9 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
 	})
 
-	postureCheckA := &posture.Checks{
-		Name:      "postureCheckA",
-		AccountID: account.Id,
-		Checks: posture.ChecksDefinition{
-			ProcessCheck: &posture.ProcessCheck{
-				Processes: []posture.Process{
-					{LinuxPath: "/usr/bin/netbird", MacPath: "/usr/local/bin/netbird"},
-				},
-			},
-		},
-	}
-	postureCheckA, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckA)
-	require.NoError(t, err)
-
-	postureCheckB := &posture.Checks{
-		Name:      "postureCheckB",
+	postureCheck := posture.Checks{
+		ID:        "postureCheck",
+		Name:      "postureCheck",
 		AccountID: account.Id,
 		Checks: posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{
@@ -176,7 +168,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		postureCheckB, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -194,12 +186,12 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		postureCheckB.Checks = posture.ChecksDefinition{
+		postureCheck.Checks = posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{
 				MinVersion: "0.29.0",
 			},
 		}
-		_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -209,10 +201,12 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	policy := &Policy{
+	policy := Policy{
+		ID:      "policyA",
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
+				ID:            xid.New().String(),
 				Enabled:       true,
 				Sources:       []string{"groupA"},
 				Destinations:  []string{"groupA"},
@@ -220,7 +214,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 				Action:        PolicyTrafficActionAccept,
 			},
 		},
-		SourcePostureChecks: []string{postureCheckB.ID},
+		SourcePostureChecks: []string{postureCheck.ID},
 	}
 
 	// Linking posture check to policy should trigger update account peers and send peer update
@@ -231,7 +225,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		policy, err = manager.SavePolicy(context.Background(), account.Id, userID, policy)
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		select {
@@ -243,7 +237,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 
 	// Updating linked posture checks should update account peers and send peer update
 	t.Run("updating linked to posture check with peers", func(t *testing.T) {
-		postureCheckB.Checks = posture.ChecksDefinition{
+		postureCheck.Checks = posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{
 				MinVersion: "0.29.0",
 			},
@@ -260,7 +254,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -270,6 +264,25 @@ 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{})
@@ -279,7 +292,8 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}()
 
 		policy.SourcePostureChecks = []string{}
-		_, err := manager.SavePolicy(context.Background(), account.Id, userID, policy)
+
+		err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 		assert.NoError(t, err)
 
 		select {
@@ -297,7 +311,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		err := manager.DeletePostureChecks(context.Background(), account.Id, postureCheckA.ID, userID)
+		err := manager.DeletePostureChecks(context.Background(), account.Id, "postureCheck", userID)
 		assert.NoError(t, err)
 
 		select {
@@ -307,15 +321,17 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+	err = manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 	assert.NoError(t, err)
 
 	// Updating linked posture check to policy with no peers should not trigger account peers update and not send peer update
 	t.Run("updating linked posture check to policy with no peers", func(t *testing.T) {
-		_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+		policy = Policy{
+			ID:      "policyB",
 			Enabled: true,
 			Rules: []*PolicyRule{
 				{
+					ID:            xid.New().String(),
 					Enabled:       true,
 					Sources:       []string{"groupB"},
 					Destinations:  []string{"groupC"},
@@ -323,8 +339,9 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-			SourcePostureChecks: []string{postureCheckB.ID},
-		})
+			SourcePostureChecks: []string{postureCheck.ID},
+		}
+		err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 		assert.NoError(t, err)
 
 		done := make(chan struct{})
@@ -333,12 +350,12 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		postureCheckB.Checks = posture.ChecksDefinition{
+		postureCheck.Checks = posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{
 				MinVersion: "0.29.0",
 			},
 		}
-		_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -355,11 +372,12 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		t.Cleanup(func() {
 			manager.peersUpdateManager.CloseChannel(context.Background(), peer2.ID)
 		})
-
-		_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+		policy = Policy{
+			ID:      "policyB",
 			Enabled: true,
 			Rules: []*PolicyRule{
 				{
+					ID:            xid.New().String(),
 					Enabled:       true,
 					Sources:       []string{"groupB"},
 					Destinations:  []string{"groupA"},
@@ -367,8 +385,10 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-			SourcePostureChecks: []string{postureCheckB.ID},
-		})
+			SourcePostureChecks: []string{postureCheck.ID},
+		}
+
+		err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, true)
 		assert.NoError(t, err)
 
 		done := make(chan struct{})
@@ -377,12 +397,12 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		postureCheckB.Checks = posture.ChecksDefinition{
+		postureCheck.Checks = posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{
 				MinVersion: "0.29.0",
 			},
 		}
-		_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -392,10 +412,11 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 		}
 	})
 
-	// Updating linked client posture check to policy where source has peers but destination does not,
-	// should trigger account peers update and send peer update
+	// 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) {
-		_, err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
+		policy = Policy{
+			ID:      "policyB",
 			Enabled: true,
 			Rules: []*PolicyRule{
 				{
@@ -406,8 +427,50 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 					Action:        PolicyTrafficActionAccept,
 				},
 			},
-			SourcePostureChecks: []string{postureCheckB.ID},
-		})
+			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) {
+		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{})
@@ -416,7 +479,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		postureCheckB.Checks = posture.ChecksDefinition{
+		postureCheck.Checks = posture.ChecksDefinition{
 			ProcessCheck: &posture.ProcessCheck{
 				Processes: []posture.Process{
 					{
@@ -425,7 +488,7 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 				},
 			},
 		}
-		_, err = manager.SavePostureChecks(context.Background(), account.Id, userID, postureCheckB)
+		err := manager.SavePostureChecks(context.Background(), account.Id, userID, &postureCheck)
 		assert.NoError(t, err)
 
 		select {
@@ -436,120 +499,80 @@ func TestPostureCheckAccountPeersUpdate(t *testing.T) {
 	})
 }
 
-func TestArePostureCheckChangesAffectPeers(t *testing.T) {
-	manager, err := createManager(t)
-	require.NoError(t, err, "failed to create account manager")
-
-	account, err := initTestPostureChecksAccount(manager)
-	require.NoError(t, err, "failed to init testing account")
-
-	groupA := &group.Group{
-		ID:        "groupA",
-		AccountID: account.Id,
-		Peers:     []string{"peer1"},
-	}
-
-	groupB := &group.Group{
-		ID:        "groupB",
-		AccountID: account.Id,
-		Peers:     []string{},
-	}
-	err = manager.Store.SaveGroups(context.Background(), LockingStrengthUpdate, []*group.Group{groupA, groupB})
-	require.NoError(t, err, "failed to save groups")
-
-	postureCheckA := &posture.Checks{
-		Name:      "checkA",
-		AccountID: account.Id,
-		Checks: posture.ChecksDefinition{
-			NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.33.1"},
-		},
-	}
-	postureCheckA, err = manager.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheckA)
-	require.NoError(t, err, "failed to save postureCheckA")
-
-	postureCheckB := &posture.Checks{
-		Name:      "checkB",
-		AccountID: account.Id,
-		Checks: posture.ChecksDefinition{
-			NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.33.1"},
-		},
-	}
-	postureCheckB, err = manager.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheckB)
-	require.NoError(t, err, "failed to save postureCheckB")
-
-	policy := &Policy{
-		AccountID: account.Id,
-		Rules: []*PolicyRule{
+func TestArePostureCheckChangesAffectingPeers(t *testing.T) {
+	account := &Account{
+		Policies: []*Policy{
 			{
-				Enabled:      true,
-				Sources:      []string{"groupA"},
-				Destinations: []string{"groupA"},
+				ID: "policyA",
+				Rules: []*PolicyRule{
+					{
+						Enabled:      true,
+						Sources:      []string{"groupA"},
+						Destinations: []string{"groupA"},
+					},
+				},
+				SourcePostureChecks: []string{"checkA"},
+			},
+		},
+		Groups: map[string]*group.Group{
+			"groupA": {
+				ID:    "groupA",
+				Peers: []string{"peer1"},
+			},
+			"groupB": {
+				ID:    "groupB",
+				Peers: []string{},
+			},
+		},
+		PostureChecks: []*posture.Checks{
+			{
+				ID: "checkA",
+			},
+			{
+				ID: "checkB",
 			},
 		},
-		SourcePostureChecks: []string{postureCheckA.ID},
 	}
 
-	policy, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
-	require.NoError(t, err, "failed to save policy")
-
 	t.Run("posture check exists and is linked to policy with peers", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
-		require.NoError(t, err)
+		result := arePostureCheckChangesAffectingPeers(account, "checkA", true)
 		assert.True(t, result)
 	})
 
 	t.Run("posture check exists but is not linked to any policy", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckB.ID)
-		require.NoError(t, err)
+		result := arePostureCheckChangesAffectingPeers(account, "checkB", true)
 		assert.False(t, result)
 	})
 
 	t.Run("posture check does not exist", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, "unknown")
-		require.NoError(t, err)
+		result := arePostureCheckChangesAffectingPeers(account, "unknown", false)
 		assert.False(t, result)
 	})
 
 	t.Run("posture check is linked to policy with no peers in source groups", func(t *testing.T) {
-		policy.Rules[0].Sources = []string{"groupB"}
-		policy.Rules[0].Destinations = []string{"groupA"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
-		require.NoError(t, err, "failed to update policy")
-
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
-		require.NoError(t, err)
+		account.Policies[0].Rules[0].Sources = []string{"groupB"}
+		account.Policies[0].Rules[0].Destinations = []string{"groupA"}
+		result := arePostureCheckChangesAffectingPeers(account, "checkA", true)
 		assert.True(t, result)
 	})
 
 	t.Run("posture check is linked to policy with no peers in destination groups", func(t *testing.T) {
-		policy.Rules[0].Sources = []string{"groupA"}
-		policy.Rules[0].Destinations = []string{"groupB"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
-		require.NoError(t, err, "failed to update policy")
-
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
-		require.NoError(t, err)
+		account.Policies[0].Rules[0].Sources = []string{"groupA"}
+		account.Policies[0].Rules[0].Destinations = []string{"groupB"}
+		result := arePostureCheckChangesAffectingPeers(account, "checkA", true)
 		assert.True(t, result)
 	})
 
-	t.Run("posture check is linked to policy but no peers in groups", func(t *testing.T) {
-		groupA.Peers = []string{}
-		err = manager.Store.SaveGroup(context.Background(), LockingStrengthUpdate, groupA)
-		require.NoError(t, err, "failed to save groups")
-
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
-		require.NoError(t, err)
+	t.Run("posture check is linked to policy with non-existent group", func(t *testing.T) {
+		account.Policies[0].Rules[0].Sources = []string{"nonExistentGroup"}
+		account.Policies[0].Rules[0].Destinations = []string{"nonExistentGroup"}
+		result := arePostureCheckChangesAffectingPeers(account, "checkA", true)
 		assert.False(t, result)
 	})
 
-	t.Run("posture check is linked to policy with non-existent group", func(t *testing.T) {
-		policy.Rules[0].Sources = []string{"nonExistentGroup"}
-		policy.Rules[0].Destinations = []string{"nonExistentGroup"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
-		require.NoError(t, err, "failed to update policy")
-
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
-		require.NoError(t, err)
+	t.Run("posture check is linked to policy but no peers in groups", func(t *testing.T) {
+		account.Groups["groupA"].Peers = []string{}
+		result := arePostureCheckChangesAffectingPeers(account, "checkA", true)
 		assert.False(t, result)
 	})
 }

management/server/route.go

@@ -237,8 +237,8 @@ func (am *DefaultAccountManager) CreateRoute(ctx context.Context, accountID stri
 		return nil, err
 	}
 
-	if am.isRouteChangeAffectPeers(account, &newRoute) {
-		am.updateAccountPeers(ctx, accountID)
+	if isRouteChangeAffectPeers(account, &newRoute) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	am.StoreEvent(ctx, userID, string(newRoute.ID), accountID, activity.RouteCreated, newRoute.EventMeta())
@@ -323,8 +323,8 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
 		return err
 	}
 
-	if am.isRouteChangeAffectPeers(account, oldRoute) || am.isRouteChangeAffectPeers(account, routeToSave) {
-		am.updateAccountPeers(ctx, accountID)
+	if isRouteChangeAffectPeers(account, oldRoute) || isRouteChangeAffectPeers(account, routeToSave) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	am.StoreEvent(ctx, userID, string(routeToSave.ID), accountID, activity.RouteUpdated, routeToSave.EventMeta())
@@ -355,8 +355,8 @@ func (am *DefaultAccountManager) DeleteRoute(ctx context.Context, accountID stri
 
 	am.StoreEvent(ctx, userID, string(routy.ID), accountID, activity.RouteRemoved, routy.EventMeta())
 
-	if am.isRouteChangeAffectPeers(account, routy) {
-		am.updateAccountPeers(ctx, accountID)
+	if isRouteChangeAffectPeers(account, routy) {
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
@@ -651,6 +651,6 @@ func getProtoPortInfo(rule *RouteFirewallRule) *proto.PortInfo {
 
 // isRouteChangeAffectPeers checks if a given route affects peers by determining
 // if it has a routing peer, distribution, or peer groups that include peers
-func (am *DefaultAccountManager) isRouteChangeAffectPeers(account *Account, route *route.Route) bool {
-	return am.anyGroupHasPeers(account, route.Groups) || am.anyGroupHasPeers(account, route.PeerGroups) || route.Peer != ""
+func isRouteChangeAffectPeers(account *Account, route *route.Route) bool {
+	return anyGroupHasPeers(account, route.Groups) || anyGroupHasPeers(account, route.PeerGroups) || route.Peer != ""
 }

management/server/route_test.go

@@ -1091,7 +1091,7 @@ func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
 	require.NoError(t, err)
 	assert.Len(t, peer4Routes.Routes, 1, "HA route should have 1 server route")
 
-	groups, err := am.Store.GetAccountGroups(context.Background(), LockingStrengthShare, account.Id)
+	groups, err := am.ListGroups(context.Background(), account.Id)
 	require.NoError(t, err)
 	var groupHA1, groupHA2 *nbgroup.Group
 	for _, group := range groups {
@@ -1214,11 +1214,12 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 
 	defaultRule := rules[0]
 	newPolicy := defaultRule.Copy()
+	newPolicy.ID = xid.New().String()
 	newPolicy.Name = "peer1 only"
 	newPolicy.Rules[0].Sources = []string{newGroup.ID}
 	newPolicy.Rules[0].Destinations = []string{newGroup.ID}
 
-	_, err = am.SavePolicy(context.Background(), account.Id, userID, newPolicy)
+	err = am.SavePolicy(context.Background(), account.Id, userID, newPolicy, false)
 	require.NoError(t, err)
 
 	err = am.DeletePolicy(context.Background(), account.Id, defaultRule.ID, userID)
@@ -1937,6 +1938,26 @@ 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/setupkey.go

@@ -4,17 +4,18 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
+	"fmt"
 	"hash/fnv"
-	"slices"
 	"strconv"
 	"strings"
 	"time"
 	"unicode/utf8"
 
 	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
-	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -228,43 +229,32 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	var setupKey *SetupKey
-	var plainKey string
-	var eventsToStore []func()
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, autoGroups); err != nil {
-			return err
-		}
-
-		setupKey, plainKey = GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
-		setupKey.AccountID = accountID
-
-		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, autoGroups, nil, setupKey)
-		eventsToStore = append(eventsToStore, events...)
-
-		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, setupKey)
-	})
-	if err != nil {
+	if err := validateSetupKeyAutoGroups(account, autoGroups); err != nil {
 		return nil, err
 	}
 
+	setupKey, plainKey := GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
+	account.SetupKeys[setupKey.Key] = setupKey
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
+		return nil, status.Errorf(status.Internal, "failed adding account key")
+	}
+
 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-	for _, storeEvent := range eventsToStore {
-		storeEvent()
+
+	for _, g := range setupKey.AutoGroups {
+		group := account.GetGroup(g)
+		if group != nil {
+			am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.GroupAddedToSetupKey,
+				map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": setupKey.Name})
+		} else {
+			log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+		}
 	}
 
 	// for the creation return the plain key to the caller
@@ -278,56 +268,43 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 // (e.g. the key itself, creation date, ID, etc).
 // These properties are overwritten: Name, AutoGroups, Revoked. The rest is copied from the existing key.
 func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID string, keyToSave *SetupKey, userID string) (*SetupKey, error) {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
 	if keyToSave == nil {
 		return nil, status.Errorf(status.InvalidArgument, "provided setup key to update is nil")
 	}
 
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
 	var oldKey *SetupKey
-	var newKey *SetupKey
-	var eventsToStore []func()
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups); err != nil {
-			return err
+	for _, key := range account.SetupKeys {
+		if key.Id == keyToSave.Id {
+			oldKey = key.Copy()
+			break
 		}
+	}
+	if oldKey == nil {
+		return nil, status.Errorf(status.NotFound, "setup key not found")
+	}
 
-		oldKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyToSave.Id)
-		if err != nil {
-			return err
-		}
+	if err := validateSetupKeyAutoGroups(account, keyToSave.AutoGroups); err != nil {
+		return nil, err
+	}
 
-		// only auto groups, revoked status, and name can be updated for now
-		newKey = oldKey.Copy()
-		newKey.Name = keyToSave.Name
-		newKey.AutoGroups = keyToSave.AutoGroups
-		newKey.Revoked = keyToSave.Revoked
-		newKey.UpdatedAt = time.Now().UTC()
+	// only auto groups, revoked status, and name can be updated for now
+	newKey := oldKey.Copy()
+	newKey.Name = keyToSave.Name
+	newKey.AutoGroups = keyToSave.AutoGroups
+	newKey.Revoked = keyToSave.Revoked
+	newKey.UpdatedAt = time.Now().UTC()
 
-		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
-		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+	account.SetupKeys[newKey.Key] = newKey
 
-		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, addedGroups, removedGroups, oldKey)
-		eventsToStore = append(eventsToStore, events...)
-
-		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, newKey)
-	})
-	if err != nil {
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return nil, err
 	}
 
@@ -335,9 +312,30 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		am.StoreEvent(ctx, userID, newKey.Id, accountID, activity.SetupKeyRevoked, newKey.EventMeta())
 	}
 
-	for _, storeEvent := range eventsToStore {
-		storeEvent()
-	}
+	defer func() {
+		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
+		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+		for _, g := range removedGroups {
+			group := account.GetGroup(g)
+			if group != nil {
+				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupRemovedFromSetupKey,
+					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
+			} else {
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+			}
+
+		}
+
+		for _, g := range addedGroups {
+			group := account.GetGroup(g)
+			if group != nil {
+				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupAddedToSetupKey,
+					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
+			} else {
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+			}
+		}
+	}()
 
 	return newKey, nil
 }
@@ -349,15 +347,16 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.NewUnauthorizedToViewSetupKeysError()
 	}
 
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	setupKeys, err := am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
 	}
 
-	return am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
+	return setupKeys, nil
 }
 
 // GetSetupKey looks up a SetupKey by KeyID, returns NotFound error if not found.
@@ -367,12 +366,8 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.NewUnauthorizedToViewSetupKeysError()
 	}
 
 	setupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
@@ -392,29 +387,21 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID, userID, keyID string) error {
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get user: %w", err)
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return status.NewUnauthorizedToViewSetupKeysError()
 	}
 
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
-	var deletedSetupKey *SetupKey
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyID)
-		if err != nil {
-			return err
-		}
-
-		return transaction.DeleteSetupKey(ctx, LockingStrengthUpdate, accountID, keyID)
-	})
+	deletedSetupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get setup key: %w", err)
+	}
+
+	err = am.Store.DeleteSetupKey(ctx, accountID, keyID)
+	if err != nil {
+		return fmt.Errorf("failed to delete setup key: %w", err)
 	}
 
 	am.StoreEvent(ctx, userID, keyID, accountID, activity.SetupKeyDeleted, deletedSetupKey.EventMeta())
@@ -422,62 +409,15 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }
 
-func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) error {
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, autoGroupIDs)
-	if err != nil {
-		return err
-	}
-
-	for _, groupID := range autoGroupIDs {
-		group, ok := groups[groupID]
+func validateSetupKeyAutoGroups(account *Account, autoGroups []string) error {
+	for _, group := range autoGroups {
+		g, ok := account.Groups[group]
 		if !ok {
-			return status.Errorf(status.NotFound, "group not found: %s", groupID)
+			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
 		}
-
-		if group.IsGroupAll() {
-			return status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
+		if g.Name == "All" {
+			return status.Errorf(status.InvalidArgument, "can't add All group to the setup key")
 		}
 	}
-
 	return nil
 }
-
-// prepareSetupKeyEvents prepares a list of event functions to be stored.
-func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string, key *SetupKey) []func() {
-	var eventsToStore []func()
-
-	modifiedGroups := slices.Concat(addedGroups, removedGroups)
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, modifiedGroups)
-	if err != nil {
-		log.WithContext(ctx).Debugf("failed to get groups for setup key events: %v", err)
-		return nil
-	}
-
-	for _, g := range removedGroups {
-		group, ok := groups[g]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromSetupKey activity: group not found", g)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupRemovedFromSetupKey, meta)
-		})
-	}
-
-	for _, g := range addedGroups {
-		group, ok := groups[g]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToSetupKey activity: group not found", g)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupAddedToSetupKey, meta)
-		})
-	}
-
-	return eventsToStore
-}

management/server/setupkey_test.go

@@ -390,7 +390,8 @@ func TestSetupKeyAccountPeersUpdate(t *testing.T) {
 	})
 	assert.NoError(t, err)
 
-	policy := &Policy{
+	policy := Policy{
+		ID:      "policy",
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -402,7 +403,7 @@ func TestSetupKeyAccountPeersUpdate(t *testing.T) {
 			},
 		},
 	}
-	_, err = manager.SavePolicy(context.Background(), account.Id, userID, policy)
+	err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 	require.NoError(t, err)
 
 	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)

management/server/sql_store.go

@@ -33,13 +33,12 @@ import (
 )
 
 const (
-	storeSqliteFileName         = "store.db"
-	idQueryCondition            = "id = ?"
-	keyQueryCondition           = "key = ?"
-	accountAndIDQueryCondition  = "account_id = ? and id = ?"
-	accountAndIDsQueryCondition = "account_id = ? AND id IN ?"
-	accountIDCondition          = "account_id = ?"
-	peerNotFoundFMT             = "peer %s not found"
+	storeSqliteFileName        = "store.db"
+	idQueryCondition           = "id = ?"
+	keyQueryCondition          = "key = ?"
+	accountAndIDQueryCondition = "account_id = ? and id = ?"
+	accountIDCondition         = "account_id = ?"
+	peerNotFoundFMT            = "peer %s not found"
 )
 
 // SqlStore represents an account storage backed by a Sql DB persisted to disk
@@ -70,17 +69,9 @@ 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.WithContext(ctx).Infof("Set max open db connections to %d", conns)
+	log.Infof("Set max open db connections to %d", conns)
 
 	if err := migrate(ctx, db); err != nil {
 		return nil, fmt.Errorf("migrate: %w", err)
@@ -305,7 +296,7 @@ func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.
 	peerCopy := peer.Copy()
 	peerCopy.AccountID = accountID
 
-	err := s.db.Transaction(func(tx *gorm.DB) error {
+	err := s.db.WithContext(ctx).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)
@@ -340,7 +331,7 @@ func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID
 	}
 
 	fieldsToUpdate := []string{"domain", "domain_category", "is_domain_primary_account"}
-	result := s.db.Model(&Account{}).
+	result := s.db.WithContext(ctx).Model(&Account{}).
 		Select(fieldsToUpdate).
 		Where(idQueryCondition, accountID).
 		Updates(&accountCopy)
@@ -412,19 +403,14 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 		}
 		usersToSave = append(usersToSave, *user)
 	}
-	err := s.db.Session(&gorm.Session{FullSaveAssociations: true}).
+	return 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
+	result := s.db.WithContext(ctx).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)
 	}
@@ -437,7 +423,7 @@ func (s *SqlStore) SaveGroups(ctx context.Context, lockStrength LockingStrength,
 		return nil
 	}
 
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&groups)
+	result := s.db.WithContext(ctx).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)
 	}
@@ -466,7 +452,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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("id").
+	result := s.db.WithContext(ctx).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)
@@ -483,13 +469,12 @@ 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.Select("account_id").First(&key, keyQueryCondition, setupKey)
+	result := s.db.WithContext(ctx).Select("account_id").First(&key, keyQueryCondition, setupKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(setupKey)
+			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account by setup key from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get account by setup key from store")
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if key.AccountID == "" {
@@ -544,7 +529,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.Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
 		Preload(clause.Associations).First(&user, idQueryCondition, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -556,9 +541,9 @@ func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStre
 	return &user, nil
 }
 
-func (s *SqlStore) GetAccountUsers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*User, error) {
+func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*User, error) {
 	var users []*User
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&users, accountIDCondition, accountID)
+	result := s.db.Find(&users, accountIDCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
@@ -570,15 +555,15 @@ func (s *SqlStore) GetAccountUsers(ctx context.Context, lockStrength LockingStre
 	return users, nil
 }
 
-func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error) {
+func (s *SqlStore) GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error) {
 	var groups []*nbgroup.Group
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountIDCondition, accountID)
+	result := s.db.Find(&groups, accountIDCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account groups from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get account groups from the store")
+		log.WithContext(ctx).Errorf("error when getting groups from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "issue getting groups from store")
 	}
 
 	return groups, nil
@@ -677,7 +662,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.Select("account_id").First(&user, idQueryCondition, userID)
+	result := s.db.WithContext(ctx).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")
@@ -694,7 +679,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.Select("account_id").First(&peer, idQueryCondition, peerID)
+	result := s.db.WithContext(ctx).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")
@@ -711,7 +696,8 @@ 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.Select("account_id").First(&peer, keyQueryCondition, peerKey)
+
+	result := s.db.WithContext(ctx).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")
@@ -729,7 +715,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.Model(&peer).Select("account_id").Where(keyQueryCondition, peerKey).First(&accountID)
+	result := s.db.WithContext(ctx).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")
@@ -755,13 +741,12 @@ 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.Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
+	result := s.db.WithContext(ctx).Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return "", status.NewSetupKeyNotFoundError(setupKey)
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account ID by setup key from store: %v", result.Error)
-		return "", status.Errorf(status.Internal, "failed to get account ID by setup key from store")
+		return "", status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if accountID == "" {
@@ -775,7 +760,7 @@ func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength
 	var ipJSONStrings []string
 
 	// Fetch the IP addresses as JSON strings
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Where("account_id = ?", accountID).
 		Pluck("ip", &ipJSONStrings)
 	if result.Error != nil {
@@ -800,7 +785,8 @@ 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
+
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Where("account_id = ?", accountID).
 		Pluck("dns_label", &labels)
 
@@ -817,7 +803,8 @@ 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.Model(&Account{}).Where(idQueryCondition, accountID).First(&accountNetwork).Error; err != nil {
+
+	if err := s.db.WithContext(ctx).Model(&Account{}).Where(idQueryCondition, accountID).First(&accountNetwork).Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, status.NewAccountNotFoundError(accountID)
 		}
@@ -828,7 +815,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.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&peer, keyQueryCondition, peerKey)
+	result := s.db.WithContext(ctx).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")
@@ -841,7 +828,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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Where(idQueryCondition, accountID).First(&accountSettings).Error; err != nil {
+	if err := s.db.WithContext(ctx).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")
 		}
@@ -853,7 +840,8 @@ 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.First(&user, accountAndIDQueryCondition, accountID, userID)
+
+	result := s.db.WithContext(ctx).First(&user, accountAndIDQueryCondition, accountID, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.NewUserNotFoundError(userID)
@@ -984,20 +972,19 @@ 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.Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
 		First(&setupKey, keyQueryCondition, key)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(key)
+			return nil, status.Errorf(status.NotFound, "setup key not found")
 		}
-		log.WithContext(ctx).Errorf("failed to get setup key by secret from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get setup key by secret from store")
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 	return &setupKey, nil
 }
 
 func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error {
-	result := s.db.Model(&SetupKey{}).
+	result := s.db.WithContext(ctx).Model(&SetupKey{}).
 		Where(idQueryCondition, setupKeyID).
 		Updates(map[string]interface{}{
 			"used_times": gorm.Expr("used_times + 1"),
@@ -1009,7 +996,7 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string
 	}
 
 	if result.RowsAffected == 0 {
-		return status.NewSetupKeyNotFoundError(setupKeyID)
+		return status.Errorf(status.NotFound, "setup key not found")
 	}
 
 	return nil
@@ -1017,7 +1004,8 @@ 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.Where("account_id = ? AND name = ?", accountID, "All").First(&group)
+
+	result := s.db.WithContext(ctx).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")
@@ -1042,12 +1030,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.Where(accountAndIDQueryCondition, accountId, groupID).First(&group)
+
+	result := s.db.WithContext(ctx).Where(accountAndIDQueryCondition, accountId, groupID).First(&group)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return status.NewGroupNotFoundError(groupID)
+			return status.Errorf(status.NotFound, "group not found for account")
 		}
-
 		return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
 	}
 
@@ -1068,62 +1056,27 @@ 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.Where("user_id = ?", userID), lockStrength, accountID)
+	return getRecords[*nbpeer.Peer](s.db.WithContext(ctx).Where("user_id = ?", userID), lockStrength, accountID)
 }
 
 func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error {
-	if err := s.db.Create(peer).Error; err != nil {
+	if err := s.db.WithContext(ctx).Create(peer).Error; err != nil {
 		return status.Errorf(status.Internal, "issue adding peer to account: %s", err)
 	}
 
 	return nil
 }
 
-// GetPeerByID retrieves a peer by its ID and account ID.
-func (s *SqlStore) GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID, peerID string) (*nbpeer.Peer, error) {
-	var peer *nbpeer.Peer
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&peer, accountAndIDQueryCondition, accountID, peerID)
+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"))
 	if result.Error != nil {
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "peer not found")
-		}
-		log.WithContext(ctx).Errorf("failed to get peer from store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get peer from store")
-	}
-
-	return peer, nil
-}
-
-// GetPeersByIDs retrieves peers by their IDs and account ID.
-func (s *SqlStore) GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error) {
-	var peers []*nbpeer.Peer
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&peers, accountAndIDsQueryCondition, accountID, peerIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get peers by ID's from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get peers by ID's from the store")
-	}
-
-	peersMap := make(map[string]*nbpeer.Peer)
-	for _, peer := range peers {
-		peersMap[peer.ID] = peer
-	}
-
-	return peersMap, nil
-}
-
-func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, lockStrength LockingStrength, accountId string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Model(&Account{}).Where(idQueryCondition, accountId).Update("network_serial", gorm.Expr("network_serial + 1"))
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to increment network serial count in store: %v", result.Error)
-		return status.Errorf(status.Internal, "failed to increment network serial count in store")
+		return status.Errorf(status.Internal, "issue incrementing network serial count: %s", result.Error)
 	}
 	return nil
 }
 
 func (s *SqlStore) ExecuteInTransaction(ctx context.Context, operation func(store Store) error) error {
-	tx := s.db.Begin()
+	tx := s.db.WithContext(ctx).Begin()
 	if tx.Error != nil {
 		return tx.Error
 	}
@@ -1138,8 +1091,7 @@ func (s *SqlStore) ExecuteInTransaction(ctx context.Context, operation func(stor
 
 func (s *SqlStore) withTx(tx *gorm.DB) Store {
 	return &SqlStore{
-		db:          tx,
-		storeEngine: s.storeEngine,
+		db: tx,
 	}
 }
 
@@ -1149,14 +1101,14 @@ 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
 		First(&accountDNSSettings, idQueryCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewAccountNotFoundError(accountID)
+			return nil, status.Errorf(status.NotFound, "dns settings not found")
 		}
-		log.WithContext(ctx).Errorf("failed to get dns settings from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get dns settings from store")
+		return nil, status.Errorf(status.Internal, "failed to get dns settings from store: %v", result.Error)
 	}
 	return &accountDNSSettings.DNSSettings, nil
 }
@@ -1164,7 +1116,8 @@ 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+
+	result := s.db.WithContext(ctx).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) {
@@ -1179,7 +1132,8 @@ 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select("domain", "domain_category").
+
+	result := s.db.WithContext(ctx).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) {
@@ -1192,352 +1146,94 @@ 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, accountID, groupID string) (*nbgroup.Group, error) {
-	var group *nbgroup.Group
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&group, accountAndIDQueryCondition, accountID, groupID)
-	if err := result.Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.NewGroupNotFoundError(groupID)
-		}
-		log.WithContext(ctx).Errorf("failed to get group from store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get group from store")
-	}
-
-	return group, nil
+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)
 }
 
 // GetGroupByName retrieves a group by name and account ID.
-func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStrength, accountID, groupName string) (*nbgroup.Group, error) {
+func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStrength, groupName, accountID string) (*nbgroup.Group, error) {
 	var group nbgroup.Group
 
 	// 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.Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations)
+	query := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations)
 	if s.storeEngine == PostgresStoreEngine {
 		query = query.Order("json_array_length(peers::json) DESC")
 	} else {
 		query = query.Order("json_array_length(peers) DESC")
 	}
 
-	result := query.First(&group, "account_id = ? AND name = ?", accountID, groupName)
+	result := query.First(&group, "name = ? and account_id = ?", groupName, accountID)
 	if err := result.Error; err != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewGroupNotFoundError(groupName)
+			return nil, status.Errorf(status.NotFound, "group not found")
 		}
-		log.WithContext(ctx).Errorf("failed to get group by name from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get group by name from store")
+		return nil, status.Errorf(status.Internal, "failed to get group from store: %s", result.Error)
 	}
 	return &group, nil
 }
 
-// GetGroupsByIDs retrieves groups by their IDs and account ID.
-func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error) {
-	var groups []*nbgroup.Group
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountAndIDsQueryCondition, accountID, groupIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get groups by ID's from store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get groups by ID's from store")
-	}
-
-	groupsMap := make(map[string]*nbgroup.Group)
-	for _, group := range groups {
-		groupsMap[group.ID] = group
-	}
-
-	return groupsMap, nil
-}
-
 // SaveGroup saves a group to the store.
 func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
 	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to save group to store: %v", result.Error)
-		return status.Errorf(status.Internal, "failed to save group to store")
+		return status.Errorf(status.Internal, "failed to save group to store: %v", result.Error)
 	}
 	return nil
 }
 
-// DeleteGroup deletes a group from the database.
-func (s *SqlStore) DeleteGroup(ctx context.Context, lockStrength LockingStrength, accountID, groupID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&nbgroup.Group{}, accountAndIDQueryCondition, accountID, groupID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to delete group from store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to delete group from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewGroupNotFoundError(groupID)
-	}
-
-	return nil
-}
-
-// DeleteGroups deletes groups from the database.
-func (s *SqlStore) DeleteGroups(ctx context.Context, strength LockingStrength, accountID string, groupIDs []string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(strength)}).
-		Delete(&nbgroup.Group{}, accountAndIDsQueryCondition, accountID, groupIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to delete groups from store: %v", result.Error)
-		return status.Errorf(status.Internal, "failed to delete groups from store: %v", result.Error)
-	}
-
-	return nil
-}
-
 // GetAccountPolicies retrieves policies for an account.
 func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) {
-	var policies []*Policy
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Preload(clause.Associations).Find(&policies, accountIDCondition, accountID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to get policies from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get policies from store")
-	}
-
-	return policies, nil
+	return getRecords[*Policy](s.db.WithContext(ctx).Preload(clause.Associations), lockStrength, accountID)
 }
 
 // GetPolicyByID retrieves a policy by its ID and account ID.
-func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*Policy, error) {
-	var policy *Policy
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations).
-		First(&policy, accountAndIDQueryCondition, accountID, policyID)
-	if err := result.Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.NewPolicyNotFoundError(policyID)
-		}
-		log.WithContext(ctx).Errorf("failed to get policy from store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get policy from store")
-	}
-
-	return policy, nil
-}
-
-func (s *SqlStore) CreatePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Create(policy)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to create policy in store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to create policy in store")
-	}
-
-	return nil
-}
-
-// SavePolicy saves a policy to the database.
-func (s *SqlStore) SavePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error {
-	result := s.db.Session(&gorm.Session{FullSaveAssociations: true}).
-		Clauses(clause.Locking{Strength: string(lockStrength)}).Save(policy)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to save policy to the store: %s", err)
-		return status.Errorf(status.Internal, "failed to save policy to store")
-	}
-	return nil
-}
-
-func (s *SqlStore) DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&Policy{}, accountAndIDQueryCondition, accountID, policyID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
-		return status.Errorf(status.Internal, "failed to delete policy from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewPolicyNotFoundError(policyID)
-	}
-
-	return nil
+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)
 }
 
 // GetAccountPostureChecks retrieves posture checks for an account.
 func (s *SqlStore) GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error) {
-	var postureChecks []*posture.Checks
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&postureChecks, accountIDCondition, accountID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get posture checks from store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get posture checks from store")
-	}
-
-	return postureChecks, nil
+	return getRecords[*posture.Checks](s.db.WithContext(ctx), lockStrength, accountID)
 }
 
 // GetPostureChecksByID retrieves posture checks by their ID and account ID.
-func (s *SqlStore) GetPostureChecksByID(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) (*posture.Checks, error) {
-	var postureCheck *posture.Checks
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&postureCheck, accountAndIDQueryCondition, accountID, postureChecksID)
-	if result.Error != nil {
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewPostureChecksNotFoundError(postureChecksID)
-		}
-		log.WithContext(ctx).Errorf("failed to get posture check from store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get posture check from store")
-	}
-
-	return postureCheck, nil
-}
-
-// GetPostureChecksByIDs retrieves posture checks by their IDs and account ID.
-func (s *SqlStore) GetPostureChecksByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, postureChecksIDs []string) (map[string]*posture.Checks, error) {
-	var postureChecks []*posture.Checks
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&postureChecks, accountAndIDsQueryCondition, accountID, postureChecksIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get posture checks by ID's from store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get posture checks by ID's from store")
-	}
-
-	postureChecksMap := make(map[string]*posture.Checks)
-	for _, postureCheck := range postureChecks {
-		postureChecksMap[postureCheck.ID] = postureCheck
-	}
-
-	return postureChecksMap, nil
-}
-
-// SavePostureChecks saves a posture checks to the database.
-func (s *SqlStore) SavePostureChecks(ctx context.Context, lockStrength LockingStrength, postureCheck *posture.Checks) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(postureCheck)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to save posture checks to store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to save posture checks to store")
-	}
-
-	return nil
-}
-
-// DeletePostureChecks deletes a posture checks from the database.
-func (s *SqlStore) DeletePostureChecks(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&posture.Checks{}, accountAndIDQueryCondition, accountID, postureChecksID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to delete posture checks from store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to delete posture checks from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewPostureChecksNotFoundError(postureChecksID)
-	}
-
-	return nil
+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)
 }
 
 // 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, lockStrength, accountID)
+	return getRecords[*route.Route](s.db.WithContext(ctx), 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, lockStrength, routeID, accountID)
+	return getRecordByID[route.Route](s.db.WithContext(ctx), lockStrength, routeID, accountID)
 }
 
 // GetAccountSetupKeys retrieves setup keys for an account.
 func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error) {
-	var setupKeys []*SetupKey
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Find(&setupKeys, accountIDCondition, accountID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to get setup keys from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get setup keys from store")
-	}
-
-	return setupKeys, nil
+	return getRecords[*SetupKey](s.db.WithContext(ctx), lockStrength, accountID)
 }
 
 // GetSetupKeyByID retrieves a setup key by its ID and account ID.
-func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error) {
-	var setupKey *SetupKey
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
-	if err := result.Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(setupKeyID)
-		}
-		log.WithContext(ctx).Errorf("failed to get setup key from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get setup key from store")
-	}
-
-	return setupKey, nil
-}
-
-// SaveSetupKey saves a setup key to the database.
-func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(setupKey)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to save setup key to store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to save setup key to store")
-	}
-
-	return nil
-}
-
-// DeleteSetupKey deletes a setup key from the database.
-func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(&SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to delete setup key from store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to delete setup key from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewSetupKeyNotFoundError(keyID)
-	}
-
-	return nil
+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)
 }
 
 // GetAccountNameServerGroups retrieves name server groups for an account.
 func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
-	var nsGroups []*nbdns.NameServerGroup
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&nsGroups, accountIDCondition, accountID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to get name server groups from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get name server groups from store")
-	}
-
-	return nsGroups, nil
+	return getRecords[*nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, accountID)
 }
 
 // GetNameServerGroupByID retrieves a name server group by its ID and account ID.
-func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, accountID, nsGroupID string) (*nbdns.NameServerGroup, error) {
-	var nsGroup *nbdns.NameServerGroup
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&nsGroup, accountAndIDQueryCondition, accountID, nsGroupID)
-	if err := result.Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.NewNameServerGroupNotFoundError(nsGroupID)
-		}
-		log.WithContext(ctx).Errorf("failed to get name server group from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get name server group from store")
-	}
-
-	return nsGroup, nil
+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)
 }
 
-// SaveNameServerGroup saves a name server group to the database.
-func (s *SqlStore) SaveNameServerGroup(ctx context.Context, lockStrength LockingStrength, nameServerGroup *nbdns.NameServerGroup) error {
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(nameServerGroup)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to save name server group to the store: %s", err)
-		return status.Errorf(status.Internal, "failed to save name server group to store")
-	}
-	return nil
-}
-
-// DeleteNameServerGroup deletes a name server group from the database.
-func (s *SqlStore) DeleteNameServerGroup(ctx context.Context, lockStrength LockingStrength, accountID, nsGroupID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(&nbdns.NameServerGroup{}, accountAndIDQueryCondition, accountID, nsGroupID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to delete name server group from the store: %s", err)
-		return status.Errorf(status.Internal, "failed to delete name server group from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewNameServerGroupNotFoundError(nsGroupID)
-	}
-
-	return nil
+func (s *SqlStore) DeleteSetupKey(ctx context.Context, accountID, keyID string) error {
+	return deleteRecordByID[SetupKey](s.db.WithContext(ctx), LockingStrengthUpdate, keyID, accountID)
 }
 
 // getRecords retrieves records from the database based on the account ID.
@@ -1573,17 +1269,19 @@ func getRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, a
 	return &record, nil
 }
 
-// SaveDNSSettings saves the DNS settings to the store.
-func (s *SqlStore) SaveDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *DNSSettings) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
-		Where(idQueryCondition, accountID).Updates(&AccountDNSSettings{DNSSettings: *settings})
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to save dns settings to store: %v", result.Error)
-		return status.Errorf(status.Internal, "failed to save dns settings to store")
+// deleteRecordByID deletes a record by its ID and account ID from the database.
+func deleteRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, accountID string) error {
+	var record T
+	result := db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(record, accountAndIDQueryCondition, accountID, recordID)
+	if err := result.Error; err != nil {
+		parts := strings.Split(fmt.Sprintf("%T", record), ".")
+		recordType := parts[len(parts)-1]
+
+		return status.Errorf(status.Internal, "failed to delete %s from store: %v", recordType, err)
 	}
 
 	if result.RowsAffected == 0 {
-		return status.NewAccountNotFoundError(accountID)
+		return status.Errorf(status.NotFound, "record not found")
 	}
 
 	return nil

management/server/sql_store_test.go

@@ -14,12 +14,12 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	nbdns "github.com/netbirdio/netbird/dns"
-	nbgroup "github.com/netbirdio/netbird/management/server/group"
-	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	nbdns "github.com/netbirdio/netbird/dns"
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
+
 	route2 "github.com/netbirdio/netbird/route"
 
 	"github.com/netbirdio/netbird/management/server/status"
@@ -1181,7 +1181,7 @@ func TestSqlite_CreateAndGetObjectInTransaction(t *testing.T) {
 			t.Fatal("failed to save group")
 			return err
 		}
-		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthUpdate, group.AccountID, group.ID)
+		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthUpdate, group.ID, group.AccountID)
 		if err != nil {
 			t.Fatal("failed to get group")
 			return err
@@ -1201,7 +1201,7 @@ func TestSqlite_GetAccoundUsers(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	account, err := store.GetAccount(context.Background(), accountID)
 	require.NoError(t, err)
-	users, err := store.GetAccountUsers(context.Background(), LockingStrengthShare, accountID)
+	users, err := store.GetAccountUsers(context.Background(), accountID)
 	require.NoError(t, err)
 	require.Len(t, users, len(account.Users))
 }
@@ -1260,9 +1260,9 @@ func TestSqlite_GetGroupByName(t *testing.T) {
 	}
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 
-	group, err := store.GetGroupByName(context.Background(), LockingStrengthShare, accountID, "All")
+	group, err := store.GetGroupByName(context.Background(), LockingStrengthShare, "All", accountID)
 	require.NoError(t, err)
-	require.True(t, group.IsGroupAll())
+	require.Equal(t, "All", group.Name)
 }
 
 func Test_DeleteSetupKeySuccessfully(t *testing.T) {
@@ -1274,7 +1274,7 @@ func Test_DeleteSetupKeySuccessfully(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	setupKeyID := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
 
-	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, setupKeyID)
+	err = store.DeleteSetupKey(context.Background(), accountID, setupKeyID)
 	require.NoError(t, err)
 
 	_, err = store.GetSetupKeyByID(context.Background(), LockingStrengthShare, setupKeyID, accountID)
@@ -1290,762 +1290,6 @@ func Test_DeleteSetupKeyFailsForNonExistingKey(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	nonExistingKeyID := "non-existing-key-id"
 
-	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, nonExistingKeyID)
+	err = store.DeleteSetupKey(context.Background(), accountID, nonExistingKeyID)
 	require.Error(t, err)
 }
-
-func TestSqlStore_GetGroupsByIDs(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	tests := []struct {
-		name          string
-		groupIDs      []string
-		expectedCount int
-	}{
-		{
-			name:          "retrieve existing groups by existing IDs",
-			groupIDs:      []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
-			expectedCount: 2,
-		},
-		{
-			name:          "empty group IDs list",
-			groupIDs:      []string{},
-			expectedCount: 0,
-		},
-		{
-			name:          "non-existing group IDs",
-			groupIDs:      []string{"nonexistent1", "nonexistent2"},
-			expectedCount: 0,
-		},
-		{
-			name:          "mixed existing and non-existing group IDs",
-			groupIDs:      []string{"cfefqs706sqkneg59g4g", "nonexistent"},
-			expectedCount: 1,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			groups, err := store.GetGroupsByIDs(context.Background(), LockingStrengthShare, accountID, tt.groupIDs)
-			require.NoError(t, err)
-			require.Len(t, groups, tt.expectedCount)
-		})
-	}
-}
-
-func TestSqlStore_SaveGroup(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	group := &nbgroup.Group{
-		ID:        "group-id",
-		AccountID: accountID,
-		Issued:    "api",
-		Peers:     []string{"peer1", "peer2"},
-	}
-	err = store.SaveGroup(context.Background(), LockingStrengthUpdate, group)
-	require.NoError(t, err)
-
-	savedGroup, err := store.GetGroupByID(context.Background(), LockingStrengthShare, accountID, "group-id")
-	require.NoError(t, err)
-	require.Equal(t, savedGroup, group)
-}
-
-func TestSqlStore_SaveGroups(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	groups := []*nbgroup.Group{
-		{
-			ID:        "group-1",
-			AccountID: accountID,
-			Issued:    "api",
-			Peers:     []string{"peer1", "peer2"},
-		},
-		{
-			ID:        "group-2",
-			AccountID: accountID,
-			Issued:    "integration",
-			Peers:     []string{"peer3", "peer4"},
-		},
-	}
-	err = store.SaveGroups(context.Background(), LockingStrengthUpdate, groups)
-	require.NoError(t, err)
-}
-
-func TestSqlStore_DeleteGroup(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	tests := []struct {
-		name        string
-		groupID     string
-		expectError bool
-	}{
-		{
-			name:        "delete existing group",
-			groupID:     "cfefqs706sqkneg59g4g",
-			expectError: false,
-		},
-		{
-			name:        "delete non-existing group",
-			groupID:     "non-existing-group-id",
-			expectError: true,
-		},
-		{
-			name:        "delete with empty group ID",
-			groupID:     "",
-			expectError: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err := store.DeleteGroup(context.Background(), LockingStrengthUpdate, accountID, tt.groupID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-			} else {
-				require.NoError(t, err)
-
-				group, err := store.GetGroupByID(context.Background(), LockingStrengthShare, accountID, tt.groupID)
-				require.Error(t, err)
-				require.Nil(t, group)
-			}
-		})
-	}
-}
-
-func TestSqlStore_DeleteGroups(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	tests := []struct {
-		name        string
-		groupIDs    []string
-		expectError bool
-	}{
-		{
-			name:        "delete multiple existing groups",
-			groupIDs:    []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
-			expectError: false,
-		},
-		{
-			name:        "delete non-existing groups",
-			groupIDs:    []string{"non-existing-id-1", "non-existing-id-2"},
-			expectError: false,
-		},
-		{
-			name:        "delete with empty group IDs list",
-			groupIDs:    []string{},
-			expectError: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err := store.DeleteGroups(context.Background(), LockingStrengthUpdate, accountID, tt.groupIDs)
-			if tt.expectError {
-				require.Error(t, err)
-			} else {
-				require.NoError(t, err)
-
-				for _, groupID := range tt.groupIDs {
-					group, err := store.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
-					require.Error(t, err)
-					require.Nil(t, group)
-				}
-			}
-		})
-	}
-}
-
-func TestSqlStore_GetPeerByID(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store_policy_migrate.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	tests := []struct {
-		name        string
-		peerID      string
-		expectError bool
-	}{
-		{
-			name:        "retrieve existing peer",
-			peerID:      "cfefqs706sqkneg59g4g",
-			expectError: false,
-		},
-		{
-			name:        "retrieve non-existing peer",
-			peerID:      "non-existing",
-			expectError: true,
-		},
-		{
-			name:        "retrieve with empty peer ID",
-			peerID:      "",
-			expectError: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			peer, err := store.GetPeerByID(context.Background(), LockingStrengthShare, accountID, tt.peerID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-				require.Nil(t, peer)
-			} else {
-				require.NoError(t, err)
-				require.NotNil(t, peer)
-				require.Equal(t, tt.peerID, peer.ID)
-			}
-		})
-	}
-}
-
-func TestSqlStore_GetPeersByIDs(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store_policy_migrate.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	tests := []struct {
-		name          string
-		peerIDs       []string
-		expectedCount int
-	}{
-		{
-			name:          "retrieve existing peers by existing IDs",
-			peerIDs:       []string{"cfefqs706sqkneg59g4g", "cfeg6sf06sqkneg59g50"},
-			expectedCount: 2,
-		},
-		{
-			name:          "empty peer IDs list",
-			peerIDs:       []string{},
-			expectedCount: 0,
-		},
-		{
-			name:          "non-existing peer IDs",
-			peerIDs:       []string{"nonexistent1", "nonexistent2"},
-			expectedCount: 0,
-		},
-		{
-			name:          "mixed existing and non-existing peer IDs",
-			peerIDs:       []string{"cfeg6sf06sqkneg59g50", "nonexistent"},
-			expectedCount: 1,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			peers, err := store.GetPeersByIDs(context.Background(), LockingStrengthShare, accountID, tt.peerIDs)
-			require.NoError(t, err)
-			require.Len(t, peers, tt.expectedCount)
-		})
-	}
-}
-
-func TestSqlStore_GetPostureChecksByID(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	tests := []struct {
-		name            string
-		postureChecksID string
-		expectError     bool
-	}{
-		{
-			name:            "retrieve existing posture checks",
-			postureChecksID: "csplshq7qv948l48f7t0",
-			expectError:     false,
-		},
-		{
-			name:            "retrieve non-existing posture checks",
-			postureChecksID: "non-existing",
-			expectError:     true,
-		},
-		{
-			name:            "retrieve with empty posture checks ID",
-			postureChecksID: "",
-			expectError:     true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			postureChecks, err := store.GetPostureChecksByID(context.Background(), LockingStrengthShare, accountID, tt.postureChecksID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-				require.Nil(t, postureChecks)
-			} else {
-				require.NoError(t, err)
-				require.NotNil(t, postureChecks)
-				require.Equal(t, tt.postureChecksID, postureChecks.ID)
-			}
-		})
-	}
-}
-
-func TestSqlStore_GetPostureChecksByIDs(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	tests := []struct {
-		name            string
-		postureCheckIDs []string
-		expectedCount   int
-	}{
-		{
-			name:            "retrieve existing posture checks by existing IDs",
-			postureCheckIDs: []string{"csplshq7qv948l48f7t0", "cspnllq7qv95uq1r4k90"},
-			expectedCount:   2,
-		},
-		{
-			name:            "empty posture check IDs list",
-			postureCheckIDs: []string{},
-			expectedCount:   0,
-		},
-		{
-			name:            "non-existing posture check IDs",
-			postureCheckIDs: []string{"nonexistent1", "nonexistent2"},
-			expectedCount:   0,
-		},
-		{
-			name:            "mixed existing and non-existing posture check IDs",
-			postureCheckIDs: []string{"cspnllq7qv95uq1r4k90", "nonexistent"},
-			expectedCount:   1,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			groups, err := store.GetPostureChecksByIDs(context.Background(), LockingStrengthShare, accountID, tt.postureCheckIDs)
-			require.NoError(t, err)
-			require.Len(t, groups, tt.expectedCount)
-		})
-	}
-}
-
-func TestSqlStore_SavePostureChecks(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	postureChecks := &posture.Checks{
-		ID:        "posture-checks-id",
-		AccountID: accountID,
-		Checks: posture.ChecksDefinition{
-			NBVersionCheck: &posture.NBVersionCheck{
-				MinVersion: "0.31.0",
-			},
-			OSVersionCheck: &posture.OSVersionCheck{
-				Ios: &posture.MinVersionCheck{
-					MinVersion: "13.0.1",
-				},
-				Linux: &posture.MinKernelVersionCheck{
-					MinKernelVersion: "5.3.3-dev",
-				},
-			},
-			GeoLocationCheck: &posture.GeoLocationCheck{
-				Locations: []posture.Location{
-					{
-						CountryCode: "DE",
-						CityName:    "Berlin",
-					},
-				},
-				Action: posture.CheckActionAllow,
-			},
-		},
-	}
-	err = store.SavePostureChecks(context.Background(), LockingStrengthUpdate, postureChecks)
-	require.NoError(t, err)
-
-	savePostureChecks, err := store.GetPostureChecksByID(context.Background(), LockingStrengthShare, accountID, "posture-checks-id")
-	require.NoError(t, err)
-	require.Equal(t, savePostureChecks, postureChecks)
-}
-
-func TestSqlStore_DeletePostureChecks(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	tests := []struct {
-		name            string
-		postureChecksID string
-		expectError     bool
-	}{
-		{
-			name:            "delete existing posture checks",
-			postureChecksID: "csplshq7qv948l48f7t0",
-			expectError:     false,
-		},
-		{
-			name:            "delete non-existing posture checks",
-			postureChecksID: "non-existing-posture-checks-id",
-			expectError:     true,
-		},
-		{
-			name:            "delete with empty posture checks ID",
-			postureChecksID: "",
-			expectError:     true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err = store.DeletePostureChecks(context.Background(), LockingStrengthUpdate, accountID, tt.postureChecksID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-			} else {
-				require.NoError(t, err)
-				group, err := store.GetPostureChecksByID(context.Background(), LockingStrengthShare, accountID, tt.postureChecksID)
-				require.Error(t, err)
-				require.Nil(t, group)
-			}
-		})
-	}
-}
-
-func TestSqlStore_GetPolicyByID(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	tests := []struct {
-		name        string
-		policyID    string
-		expectError bool
-	}{
-		{
-			name:        "retrieve existing policy",
-			policyID:    "cs1tnh0hhcjnqoiuebf0",
-			expectError: false,
-		},
-		{
-			name:        "retrieve non-existing policy checks",
-			policyID:    "non-existing",
-			expectError: true,
-		},
-		{
-			name:        "retrieve with empty policy ID",
-			policyID:    "",
-			expectError: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			policy, err := store.GetPolicyByID(context.Background(), LockingStrengthShare, accountID, tt.policyID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-				require.Nil(t, policy)
-			} else {
-				require.NoError(t, err)
-				require.NotNil(t, policy)
-				require.Equal(t, tt.policyID, policy.ID)
-			}
-		})
-	}
-}
-
-func TestSqlStore_CreatePolicy(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	policy := &Policy{
-		ID:        "policy-id",
-		AccountID: accountID,
-		Enabled:   true,
-		Rules: []*PolicyRule{
-			{
-				Enabled:       true,
-				Sources:       []string{"groupA"},
-				Destinations:  []string{"groupC"},
-				Bidirectional: true,
-				Action:        PolicyTrafficActionAccept,
-			},
-		},
-	}
-	err = store.CreatePolicy(context.Background(), LockingStrengthUpdate, policy)
-	require.NoError(t, err)
-
-	savePolicy, err := store.GetPolicyByID(context.Background(), LockingStrengthShare, accountID, policy.ID)
-	require.NoError(t, err)
-	require.Equal(t, savePolicy, policy)
-
-}
-
-func TestSqlStore_SavePolicy(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	policyID := "cs1tnh0hhcjnqoiuebf0"
-
-	policy, err := store.GetPolicyByID(context.Background(), LockingStrengthShare, accountID, policyID)
-	require.NoError(t, err)
-
-	policy.Enabled = false
-	policy.Description = "policy"
-	err = store.SavePolicy(context.Background(), LockingStrengthUpdate, policy)
-	require.NoError(t, err)
-
-	savePolicy, err := store.GetPolicyByID(context.Background(), LockingStrengthShare, accountID, policy.ID)
-	require.NoError(t, err)
-	require.Equal(t, savePolicy, policy)
-}
-
-func TestSqlStore_DeletePolicy(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	policyID := "cs1tnh0hhcjnqoiuebf0"
-
-	err = store.DeletePolicy(context.Background(), LockingStrengthShare, accountID, policyID)
-	require.NoError(t, err)
-
-	policy, err := store.GetPolicyByID(context.Background(), LockingStrengthShare, accountID, policyID)
-	require.Error(t, err)
-	require.Nil(t, policy)
-}
-
-func TestSqlStore_GetDNSSettings(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	tests := []struct {
-		name        string
-		accountID   string
-		expectError bool
-	}{
-		{
-			name:        "retrieve existing account dns settings",
-			accountID:   "bf1c8084-ba50-4ce7-9439-34653001fc3b",
-			expectError: false,
-		},
-		{
-			name:        "retrieve non-existing account dns settings",
-			accountID:   "non-existing",
-			expectError: true,
-		},
-		{
-			name:        "retrieve dns settings with empty account ID",
-			accountID:   "",
-			expectError: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthShare, tt.accountID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-				require.Nil(t, dnsSettings)
-			} else {
-				require.NoError(t, err)
-				require.NotNil(t, dnsSettings)
-			}
-		})
-	}
-}
-
-func TestSqlStore_SaveDNSSettings(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthShare, accountID)
-	require.NoError(t, err)
-
-	dnsSettings.DisabledManagementGroups = []string{"groupA", "groupB"}
-	err = store.SaveDNSSettings(context.Background(), LockingStrengthUpdate, accountID, dnsSettings)
-	require.NoError(t, err)
-
-	saveDNSSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthShare, accountID)
-	require.NoError(t, err)
-	require.Equal(t, saveDNSSettings, dnsSettings)
-}
-
-func TestSqlStore_GetAccountNameServerGroups(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	tests := []struct {
-		name          string
-		accountID     string
-		expectedCount int
-	}{
-		{
-			name:          "retrieve name server groups by existing account ID",
-			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
-			expectedCount: 1,
-		},
-		{
-			name:          "non-existing account ID",
-			accountID:     "nonexistent",
-			expectedCount: 0,
-		},
-		{
-			name:          "empty account ID",
-			accountID:     "",
-			expectedCount: 0,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			peers, err := store.GetAccountNameServerGroups(context.Background(), LockingStrengthShare, tt.accountID)
-			require.NoError(t, err)
-			require.Len(t, peers, tt.expectedCount)
-		})
-	}
-
-}
-
-func TestSqlStore_GetNameServerByID(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	tests := []struct {
-		name        string
-		nsGroupID   string
-		expectError bool
-	}{
-		{
-			name:        "retrieve existing nameserver group",
-			nsGroupID:   "csqdelq7qv97ncu7d9t0",
-			expectError: false,
-		},
-		{
-			name:        "retrieve non-existing nameserver group",
-			nsGroupID:   "non-existing",
-			expectError: true,
-		},
-		{
-			name:        "retrieve with empty nameserver group ID",
-			nsGroupID:   "",
-			expectError: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthShare, accountID, tt.nsGroupID)
-			if tt.expectError {
-				require.Error(t, err)
-				sErr, ok := status.FromError(err)
-				require.True(t, ok)
-				require.Equal(t, sErr.Type(), status.NotFound)
-				require.Nil(t, nsGroup)
-			} else {
-				require.NoError(t, err)
-				require.NotNil(t, nsGroup)
-				require.Equal(t, tt.nsGroupID, nsGroup.ID)
-			}
-		})
-	}
-}
-
-func TestSqlStore_SaveNameServerGroup(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-
-	nsGroup := &nbdns.NameServerGroup{
-		ID:        "ns-group-id",
-		AccountID: accountID,
-		Name:      "NS Group",
-		NameServers: []nbdns.NameServer{
-			{
-				IP:     netip.MustParseAddr("8.8.8.8"),
-				NSType: 1,
-				Port:   53,
-			},
-		},
-		Groups:               []string{"groupA"},
-		Primary:              true,
-		Enabled:              true,
-		SearchDomainsEnabled: false,
-	}
-
-	err = store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, nsGroup)
-	require.NoError(t, err)
-
-	saveNSGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthShare, accountID, nsGroup.ID)
-	require.NoError(t, err)
-	require.Equal(t, saveNSGroup, nsGroup)
-}
-
-func TestSqlStore_DeleteNameServerGroup(t *testing.T) {
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "testdata/extended-store.sql", t.TempDir())
-	t.Cleanup(cleanup)
-	require.NoError(t, err)
-
-	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
-	nsGroupID := "csqdelq7qv97ncu7d9t0"
-
-	err = store.DeleteNameServerGroup(context.Background(), LockingStrengthShare, accountID, nsGroupID)
-	require.NoError(t, err)
-
-	nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthShare, accountID, nsGroupID)
-	require.Error(t, err)
-	require.Nil(t, nsGroup)
-}

management/server/status/error.go

@@ -102,55 +102,25 @@ func NewPeerLoginExpiredError() error {
 }
 
 // NewSetupKeyNotFoundError creates a new Error with NotFound type for a missing setup key
-func NewSetupKeyNotFoundError(setupKeyID string) error {
-	return Errorf(NotFound, "setup key: %s not found", setupKeyID)
+func NewSetupKeyNotFoundError(err error) error {
+	return Errorf(NotFound, "setup key not found: %s", err)
 }
 
 func NewGetAccountFromStoreError(err error) error {
 	return Errorf(Internal, "issue getting account from store: %s", err)
 }
 
-// NewUserNotPartOfAccountError creates a new Error with PermissionDenied type for a user not being part of an account
-func NewUserNotPartOfAccountError() error {
-	return Errorf(PermissionDenied, "user is not part of this account")
-}
-
 // NewGetUserFromStoreError creates a new Error with Internal type for an issue getting user from store
 func NewGetUserFromStoreError() error {
 	return Errorf(Internal, "issue getting user from store")
 }
 
-// NewAdminPermissionError creates a new Error with PermissionDenied type for actions requiring admin role.
-func NewAdminPermissionError() error {
-	return Errorf(PermissionDenied, "admin role required to perform this action")
-}
-
 // NewInvalidKeyIDError creates a new Error with InvalidArgument type for an issue getting a setup key
 func NewInvalidKeyIDError() error {
 	return Errorf(InvalidArgument, "invalid key ID")
 }
 
-// NewGetAccountError creates a new Error with Internal type for an issue getting account
-func NewGetAccountError(err error) error {
-	return Errorf(Internal, "error getting account: %s", err)
-}
-
-// NewGroupNotFoundError creates a new Error with NotFound type for a missing group
-func NewGroupNotFoundError(groupID string) error {
-	return Errorf(NotFound, "group: %s not found", groupID)
-}
-
-// NewPostureChecksNotFoundError creates a new Error with NotFound type for a missing posture checks
-func NewPostureChecksNotFoundError(postureChecksID string) error {
-	return Errorf(NotFound, "posture checks: %s not found", postureChecksID)
-}
-
-// NewPolicyNotFoundError creates a new Error with NotFound type for a missing policy
-func NewPolicyNotFoundError(policyID string) error {
-	return Errorf(NotFound, "policy: %s not found", policyID)
-}
-
-// NewNameServerGroupNotFoundError creates a new Error with NotFound type for a missing name server group
-func NewNameServerGroupNotFoundError(nsGroupID string) error {
-	return Errorf(NotFound, "nameserver group: %s not found", nsGroupID)
+// NewUnauthorizedToViewSetupKeysError creates a new Error with Unauthorized type for an issue getting a setup key
+func NewUnauthorizedToViewSetupKeysError() error {
+	return Errorf(Unauthorized, "only users with admin power can view setup keys")
 }

management/server/store.go

@@ -59,11 +59,10 @@ type Store interface {
 	SaveAccount(ctx context.Context, account *Account) error
 	DeleteAccount(ctx context.Context, account *Account) error
 	UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error
-	SaveDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *DNSSettings) error
 
 	GetUserByTokenID(ctx context.Context, tokenID string) (*User, error)
 	GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*User, error)
-	GetAccountUsers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*User, error)
+	GetAccountUsers(ctx context.Context, accountID string) ([]*User, error)
 	SaveUsers(accountID string, users map[string]*User) error
 	SaveUser(ctx context.Context, lockStrength LockingStrength, user *User) error
 	SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error
@@ -71,27 +70,18 @@ type Store interface {
 	DeleteHashedPAT2TokenIDIndex(hashedToken string) error
 	DeleteTokenID2UserIDIndex(tokenID string) error
 
-	GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error)
+	GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error)
 	GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error)
 	GetGroupByName(ctx context.Context, lockStrength LockingStrength, groupName, accountID string) (*nbgroup.Group, error)
-	GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error)
 	SaveGroups(ctx context.Context, lockStrength LockingStrength, groups []*nbgroup.Group) error
 	SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error
-	DeleteGroup(ctx context.Context, lockStrength LockingStrength, accountID, groupID string) error
-	DeleteGroups(ctx context.Context, strength LockingStrength, accountID string, groupIDs []string) error
 
 	GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error)
-	GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*Policy, error)
-	CreatePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error
-	SavePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error
-	DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error
+	GetPolicyByID(ctx context.Context, lockStrength LockingStrength, policyID string, accountID string) (*Policy, error)
 
 	GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 	GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error)
-	GetPostureChecksByID(ctx context.Context, lockStrength LockingStrength, accountID, postureCheckID string) (*posture.Checks, error)
-	GetPostureChecksByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, postureChecksIDs []string) (map[string]*posture.Checks, error)
-	SavePostureChecks(ctx context.Context, lockStrength LockingStrength, postureCheck *posture.Checks) error
-	DeletePostureChecks(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) error
+	GetPostureChecksByID(ctx context.Context, lockStrength LockingStrength, postureCheckID string, accountID string) (*posture.Checks, error)
 
 	GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountId string) ([]string, error)
 	AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error
@@ -99,8 +89,6 @@ type Store interface {
 	AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error
 	GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error)
 	GetUserPeers(ctx context.Context, lockStrength LockingStrength, accountID, userID string) ([]*nbpeer.Peer, error)
-	GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) (*nbpeer.Peer, error)
-	GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error)
 	SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error
 	SavePeerStatus(accountID, peerID string, status nbpeer.PeerStatus) error
 	SavePeerLocation(accountID string, peer *nbpeer.Peer) error
@@ -108,20 +96,16 @@ type Store interface {
 	GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error)
 	IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error
 	GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error)
-	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error)
-	SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error
-	DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error
+	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error)
 
 	GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error)
 	GetRouteByID(ctx context.Context, lockStrength LockingStrength, routeID string, accountID string) (*route.Route, error)
 
 	GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*dns.NameServerGroup, error)
 	GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, nameServerGroupID string, accountID string) (*dns.NameServerGroup, error)
-	SaveNameServerGroup(ctx context.Context, lockStrength LockingStrength, nameServerGroup *dns.NameServerGroup) error
-	DeleteNameServerGroup(ctx context.Context, lockStrength LockingStrength, accountID, nameServerGroupID string) error
 
 	GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountId string) ([]net.IP, error)
-	IncrementNetworkSerial(ctx context.Context, lockStrength LockingStrength, accountId string) error
+	IncrementNetworkSerial(ctx context.Context, accountId string) error
 	GetAccountNetwork(ctx context.Context, lockStrength LockingStrength, accountId string) (*Network, error)
 
 	GetInstallationID() string
@@ -140,6 +124,7 @@ type Store interface {
 	// This is also a method of metrics.DataSource interface.
 	GetStoreEngine() StoreEngine
 	ExecuteInTransaction(ctx context.Context, f func(store Store) error) error
+	DeleteSetupKey(ctx context.Context, accountID, keyID string) error
 }
 
 type StoreEngine string

management/server/telemetry/updatechannel_metrics.go

@@ -18,6 +18,7 @@ type UpdateChannelMetrics struct {
 	getAllConnectedPeersDurationMicro metric.Int64Histogram
 	getAllConnectedPeers              metric.Int64Histogram
 	hasChannelDurationMicro           metric.Int64Histogram
+	networkMapDiffDurationMicro       metric.Int64Histogram
 	ctx                               context.Context
 }
 
@@ -63,6 +64,11 @@ func NewUpdateChannelMetrics(ctx context.Context, meter metric.Meter) (*UpdateCh
 		return nil, err
 	}
 
+	networkMapDiffDurationMicro, err := meter.Int64Histogram("management.updatechannel.networkmap.diff.duration.micro")
+	if err != nil {
+		return nil, err
+	}
+
 	return &UpdateChannelMetrics{
 		createChannelDurationMicro:        createChannelDurationMicro,
 		closeChannelDurationMicro:         closeChannelDurationMicro,
@@ -72,6 +78,7 @@ func NewUpdateChannelMetrics(ctx context.Context, meter metric.Meter) (*UpdateCh
 		getAllConnectedPeersDurationMicro: getAllConnectedPeersDurationMicro,
 		getAllConnectedPeers:              getAllConnectedPeers,
 		hasChannelDurationMicro:           hasChannelDurationMicro,
+		networkMapDiffDurationMicro:       networkMapDiffDurationMicro,
 		ctx:                               ctx,
 	}, nil
 }
@@ -111,3 +118,8 @@ func (metrics *UpdateChannelMetrics) CountGetAllConnectedPeersDuration(duration
 func (metrics *UpdateChannelMetrics) CountHasChannelDuration(duration time.Duration) {
 	metrics.hasChannelDurationMicro.Record(metrics.ctx, duration.Microseconds())
 }
+
+// CountNetworkMapDiffDurationMicro counts the duration of the NetworkMapDiff method
+func (metrics *UpdateChannelMetrics) CountNetworkMapDiffDurationMicro(duration time.Duration) {
+	metrics.networkMapDiffDurationMicro.Record(metrics.ctx, duration.Microseconds())
+}

management/server/testdata/extended-store.sql

@@ -34,7 +34,4 @@ INSERT INTO personal_access_tokens VALUES('9dj38s35-63fb-11ec-90d6-0242ac120003'
 INSERT INTO "groups" VALUES('cfefqs706sqkneg59g4g','bf1c8084-ba50-4ce7-9439-34653001fc3b','All','api','[]',0,'');
 INSERT INTO "groups" VALUES('cfefqs706sqkneg59g3g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup1','api','[]',0,'');
 INSERT INTO "groups" VALUES('cfefqs706sqkneg59g2g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup2','api','[]',0,'');
-INSERT INTO posture_checks VALUES('csplshq7qv948l48f7t0','NetBird Version > 0.32.0','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"NBVersionCheck":{"MinVersion":"0.31.0"}}');
-INSERT INTO posture_checks VALUES('cspnllq7qv95uq1r4k90','Allow Berlin and Deny local network 172.16.1.0/24','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"GeoLocationCheck":{"Locations":[{"CountryCode":"DE","CityName":"Berlin"}],"Action":"allow"},"PeerNetworkRangeCheck":{"Action":"deny","Ranges":["172.16.1.0/24"]}}');
-INSERT INTO name_server_groups VALUES('csqdelq7qv97ncu7d9t0','bf1c8084-ba50-4ce7-9439-34653001fc3b','Google DNS','Google DNS Servers','[{"IP":"8.8.8.8","NSType":1,"Port":53},{"IP":"8.8.4.4","NSType":1,"Port":53}]','["cfefqs706sqkneg59g2g"]',1,'[]',1,0);
 INSERT INTO installations VALUES(1,'');

management/server/updatechannel.go

@@ -2,12 +2,16 @@ package server
 
 import (
 	"context"
+	"fmt"
+	"runtime/debug"
 	"sync"
 	"time"
 
+	"github.com/r3labs/diff/v3"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/proto"
+	"github.com/netbirdio/netbird/management/server/differs"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 )
 
@@ -21,6 +25,8 @@ 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
@@ -30,9 +36,10 @@ type PeersUpdateManager struct {
 // NewPeersUpdateManager returns a new instance of PeersUpdateManager
 func NewPeersUpdateManager(metrics telemetry.AppMetrics) *PeersUpdateManager {
 	return &PeersUpdateManager{
-		peerChannels: make(map[string]chan *UpdateMessage),
-		channelsMux:  &sync.RWMutex{},
-		metrics:      metrics,
+		peerChannels:      make(map[string]chan *UpdateMessage),
+		peerUpdateMessage: make(map[string]*UpdateMessage),
+		channelsMux:       &sync.RWMutex{},
+		metrics:           metrics,
 	}
 }
 
@@ -41,6 +48,15 @@ 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() {
@@ -50,6 +66,16 @@ 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 {
@@ -82,6 +108,7 @@ 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)
@@ -96,12 +123,10 @@ func (p *PeersUpdateManager) closeChannel(ctx context.Context, peerID string) {
 	if channel, ok := p.peerChannels[peerID]; ok {
 		delete(p.peerChannels, peerID)
 		close(channel)
-
-		log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
-		return
+		delete(p.peerUpdateMessage, peerID)
 	}
 
-	log.WithContext(ctx).Debugf("closing updates channel: peer %s has no channel", peerID)
+	log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
 }
 
 // CloseChannels closes updates channel for each given peer
@@ -175,3 +200,79 @@ 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, p.metrics)
+		if err != nil {
+			log.WithContext(ctx).Errorf("error checking for SyncResponse updates: %v", err)
+			return true
+		}
+		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, metric telemetry.AppMetrics) (isNew bool, err error) {
+	startTime := time.Now()
+
+	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)
+	}
+
+	if metric != nil {
+		metric.UpdateChannelMetrics().CountNetworkMapDiffDurationMicro(time.Since(startTime))
+	}
+
+	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,10 +2,21 @@ package server
 
 import (
 	"context"
+	"net"
+	"net/netip"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
+	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"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+	nbroute "github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/util"
 )
 
 // var peersUpdater *PeersUpdateManager
@@ -77,3 +88,474 @@ 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 {
+		metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Run(tt.name, func(t *testing.T) {
+			p := NewPeersUpdateManager(metrics)
+			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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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, nil)
+		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,
+	}
+}

management/server/user.go

@@ -9,16 +9,14 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/management/server/activity"
-	nbContext "github.com/netbirdio/netbird/management/server/context"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/integration_reference"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/status"
+	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -105,11 +103,6 @@ func (u *User) IsAdminOrServiceUser() bool {
 	return u.HasAdminPower() || u.IsServiceUser
 }
 
-// IsRegularUser checks if the user is a regular user.
-func (u *User) IsRegularUser() bool {
-	return !u.HasAdminPower() && !u.IsServiceUser
-}
-
 // ToUserInfo converts a User object to a UserInfo object.
 func (u *User) ToUserInfo(userData *idp.UserData, settings *Settings) (*UserInfo, error) {
 	autoGroups := u.AutoGroups
@@ -494,7 +487,7 @@ func (am *DefaultAccountManager) deleteRegularUser(ctx context.Context, account
 
 	am.StoreEvent(ctx, initiatorUserID, targetUserID, account.Id, activity.UserDeleted, meta)
 	if updateAccountPeers {
-		am.updateAccountPeers(ctx, account.Id)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	return nil
@@ -835,7 +828,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 	}
 
 	if account.Settings.GroupsPropagationEnabled && areUsersLinkedToPeers(account, userIDs) {
-		am.updateAccountPeers(ctx, account.Id)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	for _, storeEvent := range eventsToStore {
@@ -1107,9 +1100,6 @@ func (am *DefaultAccountManager) GetUsersFromAccount(ctx context.Context, accoun
 func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, account *Account, peers []*nbpeer.Peer) error {
 	var peerIDs []string
 	for _, peer := range peers {
-		// nolint:staticcheck
-		ctx = context.WithValue(ctx, nbContext.PeerIDKey, peer.Key)
-
 		if peer.Status.LoginExpired {
 			continue
 		}
@@ -1117,11 +1107,8 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
 		peer.MarkLoginExpired(true)
 		account.UpdatePeer(peer)
 		if err := am.Store.SavePeerStatus(account.Id, peer.ID, *peer.Status); err != nil {
-			return fmt.Errorf("failed saving peer status for peer %s: %s", peer.ID, err)
+			return err
 		}
-
-		log.WithContext(ctx).Tracef("mark peer %s login expired", peer.ID)
-
 		am.StoreEvent(
 			ctx,
 			peer.UserID, peer.ID, account.Id,
@@ -1132,7 +1119,7 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
 	if len(peerIDs) != 0 {
 		// this will trigger peer disconnect from the management service
 		am.peersUpdateManager.CloseChannels(ctx, peerIDs)
-		am.updateAccountPeers(ctx, account.Id)
+		am.updateAccountPeers(ctx, account)
 	}
 	return nil
 }
@@ -1240,7 +1227,7 @@ func (am *DefaultAccountManager) DeleteRegularUsers(ctx context.Context, account
 	}
 
 	if updateAccountPeers {
-		am.updateAccountPeers(ctx, accountID)
+		am.updateAccountPeers(ctx, account)
 	}
 
 	for targetUserID, meta := range deletedUsersMeta {

management/server/user_test.go

@@ -1279,7 +1279,8 @@ func TestUserAccountPeersUpdate(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	policy := &Policy{
+	policy := Policy{
+		ID:      "policy",
 		Enabled: true,
 		Rules: []*PolicyRule{
 			{
@@ -1291,7 +1292,7 @@ func TestUserAccountPeersUpdate(t *testing.T) {
 			},
 		},
 	}
-	_, err = manager.SavePolicy(context.Background(), account.Id, userID, policy)
+	err = manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
 	require.NoError(t, err)
 
 	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)

relay/client/client.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -448,11 +449,11 @@ func (c *Client) writeTo(connReference *Conn, id string, dstID []byte, payload [
 	conn, ok := c.conns[id]
 	c.mu.Unlock()
 	if !ok {
-		return 0, net.ErrClosed
+		return 0, io.EOF
 	}
 
 	if conn.conn != connReference {
-		return 0, net.ErrClosed
+		return 0, io.EOF
 	}
 
 	// todo: use buffer pool instead of create new transport msg.
@@ -507,7 +508,7 @@ func (c *Client) closeConn(connReference *Conn, id string) error {
 
 	container, ok := c.conns[id]
 	if !ok {
-		return net.ErrClosed
+		return fmt.Errorf("connection already closed")
 	}
 
 	if container.conn != connReference {

relay/client/conn.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"io"
 	"net"
 	"time"
 )
@@ -39,7 +40,7 @@ func (c *Conn) Write(p []byte) (n int, err error) {
 func (c *Conn) Read(b []byte) (n int, err error) {
 	msg, ok := <-c.messageChan
 	if !ok {
-		return 0, net.ErrClosed
+		return 0, io.EOF
 	}
 
 	n = copy(b, msg.Payload)

relay/client/manager.go

@@ -16,7 +16,6 @@ import (
 
 var (
 	relayCleanupInterval = 60 * time.Second
-	keepUnusedServerTime = 5 * time.Second
 
 	ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
 )
@@ -28,13 +27,10 @@ type RelayTrack struct {
 	sync.RWMutex
 	relayClient *Client
 	err         error
-	created     time.Time
 }
 
 func NewRelayTrack() *RelayTrack {
-	return &RelayTrack{
-		created: time.Now(),
-	}
+	return &RelayTrack{}
 }
 
 type OnServerCloseListener func()
@@ -306,18 +302,6 @@ 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,9 +288,8 @@ func TestForeginAutoClose(t *testing.T) {
 		t.Fatalf("failed to close connection: %s", err)
 	}
 
-	timeout := relayCleanupInterval + keepUnusedServerTime + 1*time.Second
-	t.Logf("waiting for relay cleanup: %s", timeout)
-	time.Sleep(timeout)
+	t.Logf("waiting for relay cleanup: %s", relayCleanupInterval+1*time.Second)
+	time.Sleep(relayCleanupInterval + 1*time.Second)
 	if len(mgr.relayClients) != 0 {
 		t.Errorf("expected 0, got %d", len(mgr.relayClients))
 	}

relay/client/picker_test.go

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

relay/server/listener/ws/conn.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -99,7 +100,7 @@ func (c *Conn) isClosed() bool {
 
 func (c *Conn) ioErrHandling(err error) error {
 	if c.isClosed() {
-		return net.ErrClosed
+		return io.EOF
 	}
 
 	var wErr *websocket.CloseError
@@ -107,7 +108,7 @@ func (c *Conn) ioErrHandling(err error) error {
 		return err
 	}
 	if wErr.Code == websocket.StatusNormalClosure {
-		return net.ErrClosed
+		return io.EOF
 	}
 	return err
 }

relay/server/listener/ws/listener.go

@@ -63,14 +63,13 @@ 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", connRemoteAddr, err)
+		log.Errorf("failed to accept ws connection from %s: %s", r.RemoteAddr, err)
 		return
 	}
 
-	rAddr, err := net.ResolveTCPAddr("tcp", connRemoteAddr)
+	rAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
 	if err != nil {
 		err = wsConn.Close(websocket.StatusInternalError, "internal error")
 		if err != nil {
@@ -91,10 +90,3 @@ 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"))
-}

relay/server/peer.go

@@ -2,7 +2,7 @@ package server
 
 import (
 	"context"
-	"errors"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -57,7 +57,7 @@ func (p *Peer) Work() {
 	for {
 		n, err := p.conn.Read(buf)
 		if err != nil {
-			if !errors.Is(err, net.ErrClosed) {
+			if err != io.EOF {
 				p.log.Errorf("failed to read message: %s", err)
 			}
 			return

util/net/net.go

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