code cleaning

.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/manager_linux.go

@@ -83,11 +83,9 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early to ensure cleanup of chains
-	go func() {
-		if err := stateManager.PersistState(context.Background()); err != nil {
-			log.Errorf("failed to persist state: %v", err)
-		}
-	}()
+	if err := stateManager.PersistState(context.Background()); err != nil {
+		log.Errorf("failed to persist state: %v", err)
+	}
 
 	return nil
 }

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/manager_linux.go

@@ -99,11 +99,9 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early
-	go func() {
-		if err := stateManager.PersistState(context.Background()); err != nil {
-			log.Errorf("failed to persist state: %v", err)
-		}
-	}()
+	if err := stateManager.PersistState(context.Background()); err != nil {
+		log.Errorf("failed to persist state: %v", err)
+	}
 
 	return nil
 }

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) {

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/config.go

@@ -164,7 +164,7 @@ func UpdateOrCreateConfig(input ConfigInput) (*Config, error) {
 		if err != nil {
 			return nil, err
 		}
-		err = util.WriteJsonWithRestrictedPermission(context.Background(), input.ConfigPath, cfg)
+		err = util.WriteJsonWithRestrictedPermission(input.ConfigPath, cfg)
 		return cfg, err
 	}
 
@@ -185,7 +185,7 @@ func CreateInMemoryConfig(input ConfigInput) (*Config, error) {
 
 // WriteOutConfig write put the prepared config to the given path
 func WriteOutConfig(path string, config *Config) error {
-	return util.WriteJson(context.Background(), path, config)
+	return util.WriteJson(path, config)
 }
 
 // createNewConfig creates a new config generating a new Wireguard key and saving to file
@@ -215,7 +215,7 @@ func update(input ConfigInput) (*Config, error) {
 	}
 
 	if updated {
-		if err := util.WriteJson(context.Background(), input.ConfigPath, config); err != nil {
+		if err := util.WriteJson(input.ConfigPath, config); err != nil {
 			return nil, 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
@@ -230,8 +229,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 		c.statusRecorder.MarkSignalConnected()
 
 		relayURLs, token := parseRelayInfo(loginResp)
-		relayManager := relayClient.NewManager(engineCtx, []string{"rels://temp-relay-quic.relay.netbird.io:443"}, myPrivateKey.PublicKey().String())
-		c.statusRecorder.SetRelayMgr(relayManager)
+		relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String())
 		if len(relayURLs) > 0 {
 			if token != nil {
 				if err := relayManager.UpdateToken(token); err != nil {
@@ -242,7 +240,9 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 			log.Infof("connecting to the Relay service(s): %s", strings.Join(relayURLs, ", "))
 			if err = relayManager.Serve(); err != nil {
 				log.Error(err)
+				return wrapErr(err)
 			}
+			c.statusRecorder.SetRelayMgr(relayManager)
 		}
 
 		peerConfig := loginResp.GetPeerConfig()

client/internal/dns/server.go

@@ -326,13 +326,9 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	// persist dns state right away
 	ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
 	defer cancel()
-
-	// don't block
-	go func() {
-		if err := s.stateManager.PersistState(ctx); err != nil {
-			log.Errorf("Failed to persist dns state: %v", err)
-		}
-	}()
+	if err := s.stateManager.PersistState(ctx); err != nil {
+		log.Errorf("Failed to persist dns state: %v", err)
+	}
 
 	if s.searchDomainNotifier != nil {
 		s.searchDomainNotifier.onNewSearchDomains(s.SearchDomains())

client/internal/engine.go

@@ -38,6 +38,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 
+
 	nbssh "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
 	nbdns "github.com/netbirdio/netbird/dns"
@@ -170,7 +171,7 @@ type Engine struct {
 
 	relayManager *relayClient.Manager
 	stateManager *statemanager.Manager
-	srWatcher    *guard.SRWatcher
+	srWatcher *guard.SRWatcher
 }
 
 // Peer is an instance of the Connection Peer
@@ -537,7 +538,6 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 
 		relayMsg := wCfg.GetRelay()
 		if relayMsg != nil {
-			// when we receive token we expect valid address list too
 			c := &auth.Token{
 				Payload:   relayMsg.GetTokenPayload(),
 				Signature: relayMsg.GetTokenSignature(),
@@ -546,16 +546,9 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 				log.Errorf("failed to update relay token: %v", err)
 				return fmt.Errorf("update relay token: %w", err)
 			}
-
-			e.relayManager.UpdateServerURLs([]string{"rels://temp-relay-quic.relay.netbird.io:443"})
-
-			// Just in case the agent started with an MGM server where the relay was disabled but was later enabled.
-			// We can ignore all errors because the guard will manage the reconnection retries.
-			_ = e.relayManager.Serve()
-		} else {
-			e.relayManager.UpdateServerURLs(nil)
 		}
 
+		// todo update relay address in the relay manager
 		// todo update signal
 	}
 
@@ -648,10 +641,6 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 }
 
 func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
-	if e.wgInterface == nil {
-		return errors.New("wireguard interface is not initialized")
-	}
-
 	if e.wgInterface.Address().String() != conf.Address {
 		oldAddr := e.wgInterface.Address().String()
 		log.Debugf("updating peer address from %s to %s", oldAddr, conf.Address)

client/internal/peer/conn.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"math/rand"
 	"net"
-	"os"
 	"runtime"
 	"strings"
 	"sync"
@@ -65,14 +64,6 @@ type ConnConfig struct {
 	ICEConfig icemaker.Config
 }
 
-type WorkerCallbacks struct {
-	OnRelayReadyCallback func(info RelayConnInfo)
-	OnRelayStatusChanged func(ConnStatus)
-
-	OnICEConnReadyCallback func(ConnPriority, ICEConnInfo)
-	OnICEStatusChanged     func(ConnStatus)
-}
-
 type Conn struct {
 	log            *log.Entry
 	mu             sync.Mutex
@@ -134,32 +125,16 @@ 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)
 
-	relayIsSupportedLocally := conn.workerRelay.RelayIsSupportedLocally()
-	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally, wFns)
+	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, conn, relayManager)
+	conn.workerICE, err = NewWorkerICE(ctx, connLog, config, conn, signaler, iFaceDiscover, statusRecorder, conn.workerRelay.RelayIsSupportedLocally())
 	if err != nil {
 		return nil, err
 	}
 
 	conn.handshaker = NewHandshaker(ctx, connLog, config, signaler, conn.workerICE, conn.workerRelay)
 
-	conn.handshaker.AddOnNewOfferListener(conn.workerRelay.OnNewOffer)
-	if os.Getenv("NB_FORCE_RELAY") != "true" {
-		conn.handshaker.AddOnNewOfferListener(conn.workerICE.OnNewOffer)
-	}
-
 	conn.guard = guard.NewGuard(connLog, ctrl, conn.isConnectedOnAllWay, config.Timeout, srWatcher)
 
 	go conn.handshaker.Listen()
@@ -301,7 +276,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 +284,18 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		return
 	}
 
-	if remoteConnNil(conn.log, iceConnInfo.RemoteConn) {
-		conn.log.Errorf("remote ICE connection is nil")
-		return
-	}
+	/*
+		// temporarily disabled the check
+		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.Priority {
 		conn.statusICE.Set(StatusConnected)
 		conn.updateIceState(iceConnInfo)
 		return
@@ -366,14 +345,14 @@ func (conn *Conn) iCEConnectionIsReady(priority ConnPriority, iceConnInfo ICECon
 		return
 	}
 	wgConfigWorkaround()
-	conn.currentConnPriority = priority
+	conn.currentConnPriority = iceConnInfo.Priority
 	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 +360,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 +380,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() != StatusDisconnected
+	conn.statusICE.Set(StatusDisconnected)
 
 	conn.guard.SetICEConnDisconnected(changed)
 
@@ -442,7 +421,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 +444,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,15 +715,6 @@ func (conn *Conn) logTraceConnState() {
 	}
 }
 
-func (conn *Conn) setRelayedProxy(proxy wgproxy.Proxy) {
-	if conn.wgProxyRelay != nil {
-		if err := conn.wgProxyRelay.CloseConn(); err != nil {
-			conn.log.Warnf("failed to close deprecated wg proxy conn: %v", err)
-		}
-	}
-	conn.wgProxyRelay = proxy
-}
-
 func isController(config ConnConfig) bool {
 	return config.LocalKey > config.Key
 }

client/internal/peer/handshaker.go

@@ -3,6 +3,7 @@ package peer
 import (
 	"context"
 	"errors"
+	"os"
 	"sync"
 
 	log "github.com/sirupsen/logrus"
@@ -58,7 +59,7 @@ type Handshaker struct {
 }
 
 func NewHandshaker(ctx context.Context, log *log.Entry, config ConnConfig, signaler *Signaler, ice *WorkerICE, relay *WorkerRelay) *Handshaker {
-	return &Handshaker{
+	hs := &Handshaker{
 		ctx:            ctx,
 		log:            log,
 		config:         config,
@@ -68,10 +69,12 @@ func NewHandshaker(ctx context.Context, log *log.Entry, config ConnConfig, signa
 		remoteOffersCh: make(chan OfferAnswer),
 		remoteAnswerCh: make(chan OfferAnswer),
 	}
-}
 
-func (h *Handshaker) AddOnNewOfferListener(offer func(remoteOfferAnswer *OfferAnswer)) {
-	h.onNewOfferListeners = append(h.onNewOfferListeners, offer)
+	hs.onNewOfferListeners = append(hs.onNewOfferListeners, hs.relay.OnNewOffer)
+	if os.Getenv("NB_FORCE_RELAY") != "true" {
+		hs.onNewOfferListeners = append(hs.onNewOfferListeners, hs.ice.OnNewOffer)
+	}
+	return hs
 }
 
 func (h *Handshaker) Listen() {

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
 }
 
@@ -676,23 +669,25 @@ func (d *Status) GetRelayStates() []relay.ProbeResult {
 	// extend the list of stun, turn servers with relay address
 	relayStates := slices.Clone(d.relayStates)
 
+	var relayState relay.ProbeResult
+
 	// if the server connection is not established then we will use the general address
 	// in case of connection we will use the instance specific address
 	instanceAddr, err := d.relayMgr.RelayInstanceAddress()
 	if err != nil {
 		// TODO add their status
-		for _, r := range d.relayMgr.ServerURLs() {
-			relayStates = append(relayStates, relay.ProbeResult{
-				URI: r,
-				Err: err,
-			})
+		if errors.Is(err, relayClient.ErrRelayClientNotConnected) {
+			for _, r := range d.relayMgr.ServerURLs() {
+				relayStates = append(relayStates, relay.ProbeResult{
+					URI: r,
+				})
+			}
+			return relayStates
 		}
-		return relayStates
+		relayState.Err = err
 	}
 
-	relayState := relay.ProbeResult{
-		URI: instanceAddr,
-	}
+	relayState.URI = instanceAddr
 	return append(relayStates, relayState)
 }
 
@@ -760,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

@@ -29,22 +29,18 @@ type ICEConnInfo struct {
 	LocalIceCandidateEndpoint  string
 	Relayed                    bool
 	RelayedOnLocal             bool
-}
-
-type WorkerICECallbacks struct {
-	OnConnReady     func(ConnPriority, ICEConnInfo)
-	OnStatusChanged func(ConnStatus)
+	Priority                   ConnPriority
 }
 
 type WorkerICE struct {
 	ctx               context.Context
 	log               *log.Entry
 	config            ConnConfig
+	conn              *Conn
 	signaler          *Signaler
 	iFaceDiscover     stdnet.ExternalIFaceDiscover
 	statusRecorder    *Status
 	hasRelayOnLocally bool
-	conn              WorkerICECallbacks
 
 	agent    *ice.Agent
 	muxAgent sync.Mutex
@@ -55,21 +51,18 @@ 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, conn *Conn, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool) (*WorkerICE, error) {
 	w := &WorkerICE{
 		ctx:               ctx,
 		log:               log,
 		config:            config,
+		conn:              conn,
 		signaler:          signaler,
 		iFaceDiscover:     ifaceDiscover,
 		statusRecorder:    statusRecorder,
 		hasRelayOnLocally: hasRelayOnLocally,
-		conn:              callBacks,
 	}
 
 	localUfrag, localPwd, err := icemaker.GenerateICECredentials()
@@ -91,10 +84,16 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		return
 	}
 
-	var preferredCandidateTypes []ice.CandidateType
+	var (
+		preferredCandidateTypes []ice.CandidateType
+		selectedPriority        ConnPriority
+	)
+
 	if w.hasRelayOnLocally && remoteOfferAnswer.RelaySrvAddress != "" {
+		selectedPriority = connPriorityICEP2P
 		preferredCandidateTypes = icemaker.CandidateTypesP2P()
 	} else {
+		selectedPriority = connPriorityICETurn
 		preferredCandidateTypes = icemaker.CandidateTypes()
 	}
 
@@ -153,9 +152,10 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		RemoteIceCandidateEndpoint: fmt.Sprintf("%s:%d", pair.Remote.Address(), pair.Remote.Port()),
 		Relayed:                    isRelayed(pair),
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
+		Priority:                   selectedPriority,
 	}
 	w.log.Debugf("on ICE conn read to use ready")
-	go w.conn.OnConnReady(selectedPriority(pair), ci)
+	w.conn.iCEConnectionIsReady(ci)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -193,7 +193,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)
 	}
 }
@@ -213,18 +214,15 @@ 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)
-			}
-			w.closeAgent(agentCancel)
-		default:
-			return
+		if state == ice.ConnectionStateFailed || state == ice.ConnectionStateDisconnected {
+			w.conn.onWorkerICEStateDisconnected()
+
+			w.muxAgent.Lock()
+			agentCancel()
+			_ = agent.Close()
+			w.agent = nil
+
+			w.muxAgent.Unlock()
 		}
 	})
 	if err != nil {
@@ -250,17 +248,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)
@@ -390,11 +377,3 @@ func isRelayed(pair *ice.CandidatePair) bool {
 	}
 	return false
 }
-
-func selectedPriority(pair *ice.CandidatePair) ConnPriority {
-	if isRelayed(pair) {
-		return connPriorityICETurn
-	} else {
-		return connPriorityICEP2P
-	}
-}

client/internal/peer/worker_relay.go

@@ -34,7 +34,7 @@ type WorkerRelay struct {
 	isController bool
 	config       ConnConfig
 	relayManager relayClient.ManagerService
-	callBacks    WorkerRelayCallbacks
+	conn         *Conn
 
 	relayedConn      net.Conn
 	relayLock        sync.Mutex
@@ -45,13 +45,13 @@ 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, conn *Conn, relayManager *relayClient.Manager) *WorkerRelay {
 	r := &WorkerRelay{
 		log:          log,
 		isController: ctrl,
 		config:       config,
 		relayManager: relayManager,
-		callBacks:    callbacks,
+		conn:         conn,
 	}
 	return r
 }
@@ -95,7 +95,7 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	}
 
 	w.log.Debugf("peer conn opened via Relay: %s", srv)
-	go w.callBacks.OnConnReady(RelayConnInfo{
+	go w.conn.relayConnectionIsReady(RelayConnInfo{
 		relayedConn:     relayedConn,
 		rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
 		rosenpassAddr:   remoteOfferAnswer.RosenpassAddr,
@@ -187,7 +187,7 @@ func (w *WorkerRelay) wgStateCheck(ctx context.Context, ctxCancel context.Cancel
 					w.relayLock.Lock()
 					_ = w.relayedConn.Close()
 					w.relayLock.Unlock()
-					w.callBacks.OnDisconnected()
+					go w.conn.onWorkerRelayStateDisconnected()
 					return
 				}
 
@@ -232,5 +232,5 @@ func (w *WorkerRelay) onRelayMGDisconnected() {
 	if w.ctxCancelWgWatch != nil {
 		w.ctxCancelWgWatch()
 	}
-	go w.callBacks.OnDisconnected()
+	go w.conn.onWorkerRelayStateDisconnected()
 }

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/client_test.go

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

client/internal/routemanager/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/internal/statemanager/manager.go

@@ -16,7 +16,6 @@ import (
 	"golang.org/x/exp/maps"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
-	"github.com/netbirdio/netbird/util"
 )
 
 // State interface defines the methods that all state types must implement
@@ -179,14 +178,25 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		return nil
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
 
 	done := make(chan error, 1)
 
-	start := time.Now()
 	go func() {
-		done <- util.WriteJsonWithRestrictedPermission(ctx, m.filePath, m.states)
+		data, err := json.MarshalIndent(m.states, "", "  ")
+		if err != nil {
+			done <- fmt.Errorf("marshal states: %w", err)
+			return
+		}
+
+		// nolint:gosec
+		if err := os.WriteFile(m.filePath, data, 0640); err != nil {
+			done <- fmt.Errorf("write state file: %w", err)
+			return
+		}
+
+		done <- nil
 	}()
 
 	select {
@@ -198,7 +208,7 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		}
 	}
 
-	log.Debugf("persisted shutdown states: %v, took %v", maps.Keys(m.dirty), time.Since(start))
+	log.Debugf("persisted shutdown states: %v", maps.Keys(m.dirty))
 
 	clear(m.dirty)
 

client/internal/statemanager/path.go

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

client/server/server.go

@@ -626,8 +626,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,7 +71,6 @@ require (
 	github.com/pion/transport/v3 v3.0.1
 	github.com/pion/turn/v3 v3.0.1
 	github.com/prometheus/client_golang v1.19.1
-	github.com/quic-go/quic-go v0.48.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
@@ -152,13 +151,11 @@ require (
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/go-redis/redis/v8 v8.11.5 // indirect
-	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/go-text/render v0.1.0 // indirect
 	github.com/go-text/typesetting v0.1.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/google/btree v1.1.2 // indirect
-	github.com/google/pprof v0.0.0-20211214055906-6f57359322fd // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.3 // indirect
@@ -219,7 +216,6 @@ require (
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect
 	go.opentelemetry.io/otel/sdk v1.26.0 // indirect
 	go.opentelemetry.io/otel/trace v1.26.0 // indirect
-	go.uber.org/mock v0.4.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/image v0.18.0 // indirect
 	golang.org/x/mod v0.17.0 // indirect
@@ -240,7 +236,7 @@ replace github.com/kardianos/service => github.com/netbirdio/service v0.0.0-2024
 
 replace github.com/getlantern/systray => github.com/netbirdio/systray v0.0.0-20231030152038-ef1ed2a27949
 
-replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-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

@@ -400,7 +400,6 @@ github.com/hashicorp/serf v0.8.2/go.mod h1:6hOLApaqBFA1NXqRQAsxw9QxuDEvNxSQRwA/J
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
-github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
@@ -522,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=
@@ -606,8 +605,6 @@ github.com/prometheus/common v0.53.0 h1:U2pL9w9nmJwJDa4qqLQ3ZaePJ6ZTwt7cMD3AG3+a
 github.com/prometheus/common v0.53.0/go.mod h1:BrxBKv3FWBIGXw89Mg1AeBq7FSyRzXWI3l3e7W3RN5U=
 github.com/prometheus/procfs v0.15.0 h1:A82kmvXJq2jTu5YUhSGNlYoxh85zLnKgPz4bMZgI5Ek=
 github.com/prometheus/procfs v0.15.0/go.mod h1:Y0RJ/Y5g5wJpkTisOtqwDSo4HwhGmLB4VQSw2sQJLHk=
-github.com/quic-go/quic-go v0.48.1 h1:y/8xmfWI9qmGTc+lBr4jKRUWLGSlSigv847ULJ4hYXA=
-github.com/quic-go/quic-go v0.48.1/go.mod h1:yBgs3rWBOADpga7F+jJsb6Ybg1LSYiQvwWlLX+/6HMs=
 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=
@@ -756,8 +753,6 @@ go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v8
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
-go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
-go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
@@ -968,7 +963,6 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

management/server/account.go

@@ -1249,7 +1249,7 @@ func (am *DefaultAccountManager) peerInactivityExpirationJob(ctx context.Context
 
 		account, err := am.Store.GetAccount(ctx, accountID)
 		if err != nil {
-			log.Errorf("failed getting account %s expiring peers", accountID)
+			log.Errorf("failed getting account %s expiring peers", account.Id)
 			return account.GetNextInactivePeerExpiration()
 		}
 
@@ -2029,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)
 		}
@@ -2059,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)
 			}
@@ -2129,7 +2129,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 	if settings.GroupsPropagationEnabled {
 		account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 		if err != nil {
-			return status.NewGetAccountError(err)
+			return fmt.Errorf("error getting account: %w", err)
 		}
 
 		if areGroupChangesAffectPeers(account, addNewGroups) || areGroupChangesAffectPeers(account, removeOldGroups) {
@@ -2290,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)
@@ -2314,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)

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 {
@@ -2773,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/file_store.go

@@ -223,7 +223,7 @@ func restore(ctx context.Context, file string) (*FileStore, error) {
 // It is recommended to call it with locking FileStore.mux
 func (s *FileStore) persist(ctx context.Context, file string) error {
 	start := time.Now()
-	err := util.WriteJson(context.Background(), file, s)
+	err := util.WriteJson(file, s)
 	if err != nil {
 		return err
 	}

management/server/group.go

@@ -6,11 +6,10 @@ import (
 	"fmt"
 	"slices"
 
-	"github.com/rs/xid"
-	log "github.com/sirupsen/logrus"
-
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
+	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
@@ -28,12 +27,17 @@ func (e *GroupLinkError) Error() string {
 
 // CheckGroupPermissions validates if a user has the necessary permissions to view groups
 func (am *DefaultAccountManager) CheckGroupPermissions(ctx context.Context, accountID, userID string) error {
+	settings, err := am.Store.GetAccountSettings(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return err
+	}
+
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return err
 	}
 
-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+	if (!user.IsAdminOrServiceUser() && settings.RegularUsersViewBlocked) || user.AccountID != accountID {
 		return status.Errorf(status.PermissionDenied, "groups are blocked for users")
 	}
 
@@ -55,7 +59,7 @@ func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID, us
 		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

management/server/group/group.go

@@ -49,8 +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"
-}

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/peers_handler.go

@@ -184,26 +184,14 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 
 	dnsDomain := h.accountManager.GetDNSDomain()
 
-	peers, err := h.accountManager.GetPeers(r.Context(), accountID, userID)
-	if err != nil {
-		util.WriteError(r.Context(), err, w)
-		return
-	}
-
-	groupsMap := map[string]*nbgroup.Group{}
-	groups, _ := h.accountManager.GetAllGroups(r.Context(), accountID, userID)
-	for _, group := range groups {
-		groupsMap[group.ID] = group
-	}
-
-	respBody := make([]*api.PeerBatch, 0, len(peers))
-	for _, peer := range peers {
+	respBody := make([]*api.PeerBatch, 0, len(account.Peers))
+	for _, peer := range account.Peers {
 		peerToReturn, err := h.checkPeerStatus(peer)
 		if err != nil {
 			util.WriteError(r.Context(), err, w)
 			return
 		}
-		groupMinimumInfo := toGroupsInfo(groupsMap, peer.ID)
+		groupMinimumInfo := toGroupsInfo(account.Groups, peer.ID)
 
 		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, 0))
 	}
@@ -316,7 +304,7 @@ func peerToAccessiblePeer(peer *nbpeer.Peer, dnsDomain string) api.AccessiblePee
 }
 
 func toGroupsInfo(groups map[string]*nbgroup.Group, peerID string) []api.GroupMinimum {
-	groupsInfo := []api.GroupMinimum{}
+	var groupsInfo []api.GroupMinimum
 	groupsChecked := make(map[string]struct{})
 	for _, group := range groups {
 		_, ok := groupsChecked[group.ID]

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/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/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)
@@ -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,7 +265,7 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 		return nil, err
 	}
 
-	if peerLabelUpdated || requiresPeerUpdates {
+	if peerLabelUpdated {
 		am.updateAccountPeers(ctx, account)
 	}
 
@@ -589,7 +586,7 @@ 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()
@@ -642,7 +639,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)
@@ -659,7 +656,7 @@ 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 {
@@ -669,7 +666,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 
 	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
@@ -687,7 +684,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 
 	validPeersMap, err := am.GetValidatedPeers(account)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to get validated peers: %w", err)
+		return nil, nil, nil, err
 	}
 	postureChecks = am.getPeerPostureChecks(account, peer)
 
@@ -767,7 +764,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
 	}

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

management/server/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).Errorf("issue getting 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: %v", g, err)
-			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: %v", g, err)
-			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/sql_store.go

@@ -69,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)
@@ -300,11 +292,13 @@ func (s *SqlStore) GetInstallationID() string {
 }
 
 func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
+	startTime := time.Now()
+
 	// To maintain data integrity, we create a copy of the peer's to prevent unintended updates to other fields.
 	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)
@@ -325,6 +319,9 @@ func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.
 	})
 
 	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return err
 	}
 
@@ -332,6 +329,8 @@ func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.
 }
 
 func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error {
+	startTime := time.Now()
+
 	accountCopy := Account{
 		Domain:                 domain,
 		DomainCategory:         category,
@@ -339,11 +338,14 @@ 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)
 	if result.Error != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return result.Error
 	}
 
@@ -355,6 +357,8 @@ func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID
 }
 
 func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
+	startTime := time.Now()
+
 	var peerCopy nbpeer.Peer
 	peerCopy.Status = &peerStatus
 
@@ -367,6 +371,9 @@ func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.Pe
 		Where(accountAndIDQueryCondition, accountID, peerID).
 		Updates(&peerCopy)
 	if result.Error != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return result.Error
 	}
 
@@ -378,6 +385,8 @@ func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.Pe
 }
 
 func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.Peer) error {
+	startTime := time.Now()
+
 	// To maintain data integrity, we create a copy of the peer's location to prevent unintended updates to other fields.
 	var peerCopy nbpeer.Peer
 	// Since the location field has been migrated to JSON serialization,
@@ -389,6 +398,9 @@ func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.P
 		Updates(peerCopy)
 
 	if result.Error != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return result.Error
 	}
 
@@ -402,6 +414,8 @@ func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.P
 // SaveUsers saves the given list of users to the database.
 // It updates existing users if a conflict occurs.
 func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
+	startTime := time.Now()
+
 	usersToSave := make([]User, 0, len(users))
 	for _, user := range users {
 		user.AccountID = accountID
@@ -415,6 +429,9 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 		Clauses(clause.OnConflict{UpdateAll: true}).
 		Create(&usersToSave).Error
 	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "failed to save users to store: %v", err)
 	}
 
@@ -423,8 +440,13 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 
 // 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)
+	startTime := time.Now()
+
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
 	if result.Error != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "failed to save user to store: %v", result.Error)
 	}
 	return nil
@@ -432,12 +454,17 @@ func (s *SqlStore) SaveUser(ctx context.Context, lockStrength LockingStrength, u
 
 // SaveGroups saves the given list of groups to the database.
 func (s *SqlStore) SaveGroups(ctx context.Context, lockStrength LockingStrength, groups []*nbgroup.Group) error {
+	startTime := time.Now()
+
 	if len(groups) == 0 {
 		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 {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "failed to save groups to store: %v", result.Error)
 	}
 	return nil
@@ -464,8 +491,10 @@ func (s *SqlStore) GetAccountByPrivateDomain(ctx context.Context, domain string)
 }
 
 func (s *SqlStore) GetAccountIDByPrivateDomain(ctx context.Context, lockStrength LockingStrength, domain string) (string, error) {
+	startTime := time.Now()
+
 	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)
@@ -473,6 +502,9 @@ func (s *SqlStore) GetAccountIDByPrivateDomain(ctx context.Context, lockStrength
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return "", status.Errorf(status.NotFound, "account not found: provided domain is not registered or is not private")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		log.WithContext(ctx).Errorf("error when getting account from the store: %s", result.Error)
 		return "", status.NewGetAccountFromStoreError(result.Error)
 	}
@@ -481,14 +513,18 @@ func (s *SqlStore) GetAccountIDByPrivateDomain(ctx context.Context, lockStrength
 }
 
 func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*Account, error) {
+	startTime := time.Now()
+
 	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")
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if key.AccountID == "" {
@@ -499,12 +535,17 @@ func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*
 }
 
 func (s *SqlStore) GetTokenIDByHashedToken(ctx context.Context, hashedToken string) (string, error) {
+	startTime := time.Now()
+
 	var token PersonalAccessToken
 	result := s.db.First(&token, "hashed_token = ?", hashedToken)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		log.WithContext(ctx).Errorf("error when getting token from the store: %s", result.Error)
 		return "", status.NewGetAccountFromStoreError(result.Error)
 	}
@@ -513,12 +554,17 @@ func (s *SqlStore) GetTokenIDByHashedToken(ctx context.Context, hashedToken stri
 }
 
 func (s *SqlStore) GetUserByTokenID(ctx context.Context, tokenID string) (*User, error) {
+	startTime := time.Now()
+
 	var token PersonalAccessToken
 	result := s.db.First(&token, idQueryCondition, tokenID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		log.WithContext(ctx).Errorf("error when getting token from the store: %s", result.Error)
 		return nil, status.NewGetAccountFromStoreError(result.Error)
 	}
@@ -542,13 +588,18 @@ func (s *SqlStore) GetUserByTokenID(ctx context.Context, tokenID string) (*User,
 }
 
 func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*User, error) {
+	startTime := time.Now()
+
 	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) {
 			return nil, status.NewUserNotFoundError(userID)
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.NewGetUserFromStoreError()
 	}
 
@@ -556,12 +607,17 @@ func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStre
 }
 
 func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*User, error) {
+	startTime := time.Now()
+
 	var users []*User
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		log.WithContext(ctx).Errorf("error when getting users from the store: %s", result.Error)
 		return nil, status.Errorf(status.Internal, "issue getting users from store")
 	}
@@ -569,15 +625,20 @@ func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*Us
 	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) {
+	startTime := time.Now()
+
 	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")
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
+		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
@@ -675,12 +736,17 @@ func (s *SqlStore) GetAccount(ctx context.Context, accountID string) (*Account,
 }
 
 func (s *SqlStore) GetAccountByUser(ctx context.Context, userID string) (*Account, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.NewGetAccountFromStoreError(result.Error)
 	}
 
@@ -692,12 +758,17 @@ func (s *SqlStore) GetAccountByUser(ctx context.Context, userID string) (*Accoun
 }
 
 func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*Account, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.NewGetAccountFromStoreError(result.Error)
 	}
 
@@ -709,12 +780,17 @@ func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*Acco
 }
 
 func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (*Account, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.NewGetAccountFromStoreError(result.Error)
 	}
 
@@ -726,13 +802,18 @@ func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (
 }
 
 func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string) (string, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return "", status.NewGetAccountFromStoreError(result.Error)
 	}
 
@@ -740,12 +821,17 @@ func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string)
 }
 
 func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
+	startTime := time.Now()
+
 	var accountID string
 	result := s.db.Model(&User{}).Select("account_id").Where(idQueryCondition, userID).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return "", status.NewGetAccountFromStoreError(result.Error)
 	}
 
@@ -753,14 +839,18 @@ func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
 }
 
 func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string) (string, error) {
+	startTime := time.Now()
+
 	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")
+		if errors.Is(result.Error, context.Canceled) {
+			return "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
+		return "", status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if accountID == "" {
@@ -771,16 +861,21 @@ func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string)
 }
 
 func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountID string) ([]net.IP, error) {
+	startTime := time.Now()
+
 	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 {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "no peers found for the account")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.Errorf(status.Internal, "issue getting IPs from store: %s", result.Error)
 	}
 
@@ -798,8 +893,10 @@ func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength
 }
 
 func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string) ([]string, error) {
+	startTime := time.Now()
+
 	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)
 
@@ -807,6 +904,9 @@ func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength Lock
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "no peers found for the account")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		log.WithContext(ctx).Errorf("error when getting dns labels from the store: %s", result.Error)
 		return nil, status.Errorf(status.Internal, "issue getting dns labels from store: %s", result.Error)
 	}
@@ -815,23 +915,33 @@ func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength Lock
 }
 
 func (s *SqlStore) GetAccountNetwork(ctx context.Context, lockStrength LockingStrength, accountID string) (*Network, error) {
+	startTime := time.Now()
+
 	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)
 		}
+		if errors.Is(err, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.Errorf(status.Internal, "issue getting network from store: %s", err)
 	}
 	return accountNetwork.Network, nil
 }
 
 func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.Errorf(status.Internal, "issue getting peer from store: %s", result.Error)
 	}
 
@@ -839,11 +949,16 @@ func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength Locking
 }
 
 func (s *SqlStore) GetAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*Settings, error) {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(err, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.Errorf(status.Internal, "issue getting settings from store: %s", err)
 	}
 	return accountSettings.Settings, nil
@@ -851,13 +966,17 @@ 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 {
+	startTime := time.Now()
+
 	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)
 		}
-
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.NewGetUserFromStoreError()
 	}
 	user.LastLogin = lastLogin
@@ -866,6 +985,8 @@ func (s *SqlStore) SaveUserLastLogin(ctx context.Context, accountID, userID stri
 }
 
 func (s *SqlStore) GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
+	startTime := time.Now()
+
 	definitionJSON, err := json.Marshal(checks)
 	if err != nil {
 		return nil, err
@@ -874,6 +995,9 @@ func (s *SqlStore) GetPostureCheckByChecksDefinition(accountID string, checks *p
 	var postureCheck posture.Checks
 	err = s.db.Where("account_id = ? AND checks = ?", accountID, string(definitionJSON)).First(&postureCheck).Error
 	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, err
 	}
 
@@ -983,21 +1107,27 @@ func NewPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore,
 }
 
 func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error) {
+	startTime := time.Now()
+
 	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")
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 	return &setupKey, nil
 }
 
 func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error {
-	result := s.db.Model(&SetupKey{}).
+	startTime := time.Now()
+
+	result := s.db.WithContext(ctx).Model(&SetupKey{}).
 		Where(idQueryCondition, setupKeyID).
 		Updates(map[string]interface{}{
 			"used_times": gorm.Expr("used_times + 1"),
@@ -1005,23 +1135,31 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string
 		})
 
 	if result.Error != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue incrementing setup key usage count: %s", result.Error)
 	}
 
 	if result.RowsAffected == 0 {
-		return status.NewSetupKeyNotFoundError(setupKeyID)
+		return status.Errorf(status.NotFound, "setup key not found")
 	}
 
 	return nil
 }
 
 func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error {
+	startTime := time.Now()
+
 	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")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue finding group 'All': %s", result.Error)
 	}
 
@@ -1034,6 +1172,9 @@ func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peer
 	group.Peers = append(group.Peers, peerID)
 
 	if err := s.db.Save(&group).Error; err != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue updating group 'All': %s", err)
 	}
 
@@ -1041,13 +1182,17 @@ func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peer
 }
 
 func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountId string, peerId string, groupID string) error {
+	startTime := time.Now()
+
 	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.Errorf(status.NotFound, "group not found for account")
 		}
-
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
 	}
 
@@ -1060,6 +1205,9 @@ func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountId string, peerId
 	group.Peers = append(group.Peers, peerId)
 
 	if err := s.db.Save(&group).Error; err != nil {
+		if errors.Is(result.Error, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue updating group: %s", err)
 	}
 
@@ -1068,11 +1216,16 @@ 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 {
+	startTime := time.Now()
+
+	if err := s.db.WithContext(ctx).Create(peer).Error; err != nil {
+		if errors.Is(err, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return status.Errorf(status.Internal, "issue adding peer to account: %s", err)
 	}
 
@@ -1080,15 +1233,20 @@ func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) erro
 }
 
 func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, accountId string) error {
-	result := s.db.Model(&Account{}).Where(idQueryCondition, accountId).Update("network_serial", gorm.Expr("network_serial + 1"))
+	startTime := time.Now()
+
+	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, context.Canceled) {
+			return status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		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
 	}
@@ -1112,13 +1270,18 @@ func (s *SqlStore) GetDB() *gorm.DB {
 }
 
 func (s *SqlStore) GetAccountDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*DNSSettings, error) {
+	startTime := time.Now()
+
 	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.Errorf(status.NotFound, "dns settings not found")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return nil, status.Errorf(status.Internal, "failed to get dns settings from store: %v", result.Error)
 	}
 	return &accountDNSSettings.DNSSettings, nil
@@ -1126,13 +1289,18 @@ 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) {
+	startTime := time.Now()
+
 	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) {
 			return false, nil
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return false, status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return false, result.Error
 	}
 
@@ -1141,13 +1309,18 @@ 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) {
+	startTime := time.Now()
+
 	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) {
 			return "", "", status.Errorf(status.NotFound, "account not found")
 		}
+		if errors.Is(result.Error, context.Canceled) {
+			return "", "", status.NewStoreContextCanceledError(time.Since(startTime))
+		}
 		return "", "", status.Errorf(status.Internal, "failed to get domain category from store: %v", result.Error)
 	}
 
@@ -1156,7 +1329,7 @@ func (s *SqlStore) GetAccountDomainAndCategory(ctx context.Context, lockStrength
 
 // GetGroupByID retrieves a group by ID and account ID.
 func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error) {
-	return getRecordByID[nbgroup.Group](s.db.Preload(clause.Associations), lockStrength, groupID, accountID)
+	return getRecordByID[nbgroup.Group](s.db.WithContext(ctx).Preload(clause.Associations), lockStrength, groupID, accountID)
 }
 
 // GetGroupByName retrieves a group by name and account ID.
@@ -1165,7 +1338,7 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 
 	// TODO: This fix is accepted for now, but if we need to handle this more frequently
 	// we may need to reconsider changing the types.
-	query := s.db.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 {
@@ -1182,26 +1355,9 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 	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, "account_id = ? AND id in ?", accountID, groupIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get groups by ID's from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get groups by ID's from the 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 {
 		return status.Errorf(status.Internal, "failed to save group to store: %v", result.Error)
 	}
@@ -1210,97 +1366,56 @@ func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength,
 
 // GetAccountPolicies retrieves policies for an account.
 func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) {
-	return getRecords[*Policy](s.db.Preload(clause.Associations), lockStrength, accountID)
+	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, policyID string, accountID string) (*Policy, error) {
-	return getRecordByID[Policy](s.db.Preload(clause.Associations), lockStrength, policyID, accountID)
+	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) {
-	return getRecords[*posture.Checks](s.db, lockStrength, accountID)
+	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, postureCheckID string, accountID string) (*posture.Checks, error) {
-	return getRecordByID[posture.Checks](s.db, lockStrength, postureCheckID, accountID)
+	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) {
-	return getRecords[*nbdns.NameServerGroup](s.db, lockStrength, accountID)
+	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, nsGroupID string, accountID string) (*nbdns.NameServerGroup, error) {
-	return getRecordByID[nbdns.NameServerGroup](s.db, lockStrength, nsGroupID, accountID)
+	return getRecordByID[nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, nsGroupID, accountID)
+}
+
+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.
@@ -1335,3 +1450,21 @@ func getRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, a
 	}
 	return &record, nil
 }
+
+// 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.Errorf(status.NotFound, "record not found")
+	}
+
+	return nil
+}

management/server/sql_store_test.go

@@ -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,6 +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)
 }

management/server/status/error.go

@@ -103,29 +103,19 @@ 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")
-}
-
 // NewStoreContextCanceledError creates a new Error with Internal type for a canceled store context
 func NewStoreContextCanceledError(duration time.Duration) error {
 	return Errorf(Internal, "store access: context canceled after %v", duration)
@@ -136,7 +126,7 @@ 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)
+// 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

@@ -70,10 +70,9 @@ 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
 
@@ -97,9 +96,7 @@ 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)
@@ -127,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/updatechannel.go

@@ -96,12 +96,9 @@ 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
 	}
 
-	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

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
@@ -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,

relay/client/client.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -10,7 +11,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	auth "github.com/netbirdio/netbird/relay/auth/hmac"
-	"github.com/netbirdio/netbird/relay/client/dialer/quic"
 	"github.com/netbirdio/netbird/relay/client/dialer/ws"
 	"github.com/netbirdio/netbird/relay/healthcheck"
 	"github.com/netbirdio/netbird/relay/messages"
@@ -94,6 +94,10 @@ func (cc *connContainer) writeMsg(msg Msg) {
 	case cc.messages <- msg:
 	case <-cc.ctx.Done():
 		msg.Free()
+	default:
+		msg.Free()
+		cc.log.Infof("message queue is full")
+		// todo consider to close the connection
 	}
 }
 
@@ -137,7 +141,7 @@ type Client struct {
 	instanceURL      *RelayAddr
 	muInstanceURL    sync.Mutex
 
-	onDisconnectListener func(string)
+	onDisconnectListener func()
 	onConnectedListener  func()
 	listenerMutex        sync.Mutex
 }
@@ -176,7 +180,8 @@ func (c *Client) Connect() error {
 		return nil
 	}
 
-	if err := c.connect(); err != nil {
+	err := c.connect()
+	if err != nil {
 		return err
 	}
 
@@ -229,7 +234,7 @@ func (c *Client) ServerInstanceURL() (string, error) {
 }
 
 // SetOnDisconnectListener sets a function that will be called when the connection to the relay server is closed.
-func (c *Client) SetOnDisconnectListener(fn func(string)) {
+func (c *Client) SetOnDisconnectListener(fn func()) {
 	c.listenerMutex.Lock()
 	defer c.listenerMutex.Unlock()
 	c.onDisconnectListener = fn
@@ -260,20 +265,14 @@ func (c *Client) Close() error {
 }
 
 func (c *Client) connect() error {
-	var conn net.Conn
-	var err error
-	if c.connectionURL == "rels://temp-relay-quic.relay.netbird.io:443" {
-		log.Infof("connecting to relay server %s using quic protocol", c.connectionURL)
-		conn, err = quic.Dial(c.connectionURL)
-	} else {
-		conn, err = ws.Dial(c.connectionURL)
-	}
+	conn, err := ws.Dial(c.connectionURL)
 	if err != nil {
 		return err
 	}
 	c.relayConn = conn
 
-	if err = c.handShake(); err != nil {
+	err = c.handShake()
+	if err != nil {
 		cErr := conn.Close()
 		if cErr != nil {
 			c.log.Errorf("failed to close connection: %s", cErr)
@@ -347,7 +346,7 @@ func (c *Client) readLoop(relayConn net.Conn) {
 			c.log.Infof("start to Relay read loop exit")
 			c.mu.Lock()
 			if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
-				c.log.Errorf("failed to read message from relay server: %s", errExit)
+				c.log.Debugf("failed to read message from relay server: %s", errExit)
 			}
 			c.mu.Unlock()
 			break
@@ -450,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.
@@ -509,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 {
@@ -556,7 +555,7 @@ func (c *Client) notifyDisconnected() {
 	if c.onDisconnectListener == nil {
 		return
 	}
-	go c.onDisconnectListener(c.connectionURL)
+	go c.onDisconnectListener()
 }
 
 func (c *Client) notifyConnected() {

relay/client/client_test.go

@@ -551,7 +551,7 @@ func TestCloseByServer(t *testing.T) {
 	}
 
 	disconnected := make(chan struct{})
-	relayClient.SetOnDisconnectListener(func(_ string) {
+	relayClient.SetOnDisconnectListener(func() {
 		log.Infof("client disconnected")
 		close(disconnected)
 	})

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/dialer/quic/conn.go

@@ -1,87 +0,0 @@
-package quic
-
-import (
-	"context"
-	"net"
-	"time"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-type QuicAddr struct {
-	addr string
-}
-
-func (a QuicAddr) Network() string {
-	return "quic"
-}
-
-func (a QuicAddr) String() string {
-	return a.addr
-}
-
-type Conn struct {
-	session    quic.Connection
-	remoteAddr QuicAddr
-	ctx        context.Context
-}
-
-func NewConn(session quic.Connection, serverAddress string) net.Conn {
-	return &Conn{
-		session:    session,
-		remoteAddr: QuicAddr{addr: serverAddress},
-		ctx:        context.Background(),
-	}
-}
-
-func (c *Conn) Read(b []byte) (n int, err error) {
-	dgram, err := c.session.ReceiveDatagram(c.ctx)
-	if err != nil {
-		log.Errorf("failed to read from QUIC session: %v", err)
-		return 0, err
-	}
-
-	n = copy(b, dgram)
-	return n, nil
-}
-
-func (c *Conn) Write(b []byte) (int, error) {
-	err := c.session.SendDatagram(b)
-	if err != nil {
-		log.Errorf("failed to write to QUIC stream: %v", err)
-		return 0, err
-	}
-	return len(b), nil
-}
-
-func (c *Conn) RemoteAddr() net.Addr {
-	if c.session != nil {
-		return c.session.RemoteAddr()
-	}
-	return c.remoteAddr
-}
-
-func (c *Conn) LocalAddr() net.Addr {
-	if c.session != nil {
-		return c.session.LocalAddr()
-	}
-	return QuicAddr{addr: "unknown"}
-}
-
-func (c *Conn) SetReadDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetWriteDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetDeadline(t time.Time) error {
-
-	return nil
-}
-
-func (c *Conn) Close() error {
-	return c.session.CloseWithError(0, "normal closure")
-}

relay/client/dialer/quic/quic.go

@@ -1,60 +0,0 @@
-package quic
-
-import (
-	"context"
-	"crypto/tls"
-	"fmt"
-	"net"
-	"strings"
-	"time"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	dialTimeout = 30 * time.Second
-)
-
-func Dial(address string) (net.Conn, error) {
-	quicURL, err := prepareURL(address)
-	if err != nil {
-		return nil, err
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
-	defer cancel()
-
-	tlsConf := &tls.Config{
-		InsecureSkipVerify: true,           // Set to true only for testing
-		NextProtos:         []string{"h2"}, // Ensure this matches the server's ALPN
-	}
-
-	quicConfig := &quic.Config{
-		KeepAlivePeriod: 15 * time.Second,
-		MaxIdleTimeout:  60 * time.Second,
-		EnableDatagrams: true,
-	}
-
-	// todo add support for custom dialer
-
-	session, err := quic.DialAddr(ctx, quicURL, tlsConf, quicConfig)
-	if err != nil {
-		log.Errorf("failed to dial to Relay server via QUIC '%s': %s", quicURL, err)
-		return nil, err
-	}
-
-	conn := NewConn(session, address)
-	return conn, nil
-}
-
-func prepareURL(address string) (string, error) {
-	if !strings.HasPrefix(address, "rel://") && !strings.HasPrefix(address, "rels://") {
-		return "", fmt.Errorf("unsupported scheme: %s", address)
-	}
-
-	if strings.HasPrefix(address, "rels://") {
-		return address[7:], nil
-	}
-	return address[6:], nil
-}

relay/client/dialer/ws/ws.go

@@ -2,7 +2,6 @@ package ws
 
 import (
 	"context"
-	"crypto/tls"
 	"fmt"
 	"net"
 	"net/http"
@@ -60,10 +59,6 @@ func httpClientNbDialer() *http.Client {
 		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 			return customDialer.DialContext(ctx, network, addr)
 		},
-		// Set up a TLS configuration that skips certificate verification
-		TLSClientConfig: &tls.Config{
-			InsecureSkipVerify: true, // This accepts invalid TLS certificates
-		},
 	}
 
 	return &http.Client{

relay/client/guard.go

@@ -4,120 +4,65 @@ import (
 	"context"
 	"time"
 
-	"github.com/cenkalti/backoff/v4"
 	log "github.com/sirupsen/logrus"
 )
 
 var (
-	reconnectingTimeout = 60 * time.Second
+	reconnectingTimeout = 5 * time.Second
 )
 
 // Guard manage the reconnection tries to the Relay server in case of disconnection event.
 type Guard struct {
-	// OnNewRelayClient is a channel that is used to notify the relay client about a new relay client instance.
-	OnNewRelayClient chan *Client
-	serverPicker     *ServerPicker
+	ctx         context.Context
+	relayClient *Client
 }
 
 // NewGuard creates a new guard for the relay client.
-func NewGuard(sp *ServerPicker) *Guard {
+func NewGuard(context context.Context, relayClient *Client) *Guard {
 	g := &Guard{
-		OnNewRelayClient: make(chan *Client, 1),
-		serverPicker:     sp,
+		ctx:         context,
+		relayClient: relayClient,
 	}
 	return g
 }
 
-// StartReconnectTrys is called when the relay client is disconnected from the relay server.
-// It attempts to reconnect to the relay server. The function first tries a quick reconnect
-// to the same server that was used before, if the server URL is still valid. If the quick
-// reconnect fails, it starts a ticker to periodically attempt server picking until it
-// succeeds or the context is done.
-//
-// Parameters:
-// - ctx: The context to control the lifecycle of the reconnection attempts.
-// - relayClient: The relay client instance that was disconnected.
+// OnDisconnected is called when the relay client is disconnected from the relay server. It will trigger the reconnection
 // todo prevent multiple reconnection instances. In the current usage it should not happen, but it is better to prevent
-func (g *Guard) StartReconnectTrys(ctx context.Context, relayClient *Client) {
-	if relayClient == nil {
-		goto RETRY
-	}
-	if g.isServerURLStillValid(relayClient) && g.quickReconnect(ctx, relayClient) {
+func (g *Guard) OnDisconnected() {
+	if g.quickReconnect() {
 		return
 	}
 
-RETRY:
-	ticker := exponentTicker(ctx)
+	ticker := time.NewTicker(reconnectingTimeout)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ticker.C:
-			if err := g.retry(ctx); err != nil {
-				log.Errorf("failed to pick new Relay server: %s", err)
+			err := g.relayClient.Connect()
+			if err != nil {
+				log.Errorf("failed to reconnect to relay server: %s", err)
 				continue
 			}
 			return
-		case <-ctx.Done():
+		case <-g.ctx.Done():
 			return
 		}
 	}
 }
 
-func (g *Guard) retry(ctx context.Context) error {
-	log.Infof("try to pick up a new Relay server")
-	relayClient, err := g.serverPicker.PickServer(ctx)
-	if err != nil {
-		return err
-	}
-
-	// prevent to work with a deprecated Relay client instance
-	g.drainRelayClientChan()
-
-	g.OnNewRelayClient <- relayClient
-	return nil
-}
-
-func (g *Guard) quickReconnect(parentCtx context.Context, rc *Client) bool {
-	ctx, cancel := context.WithTimeout(parentCtx, 1500*time.Millisecond)
+func (g *Guard) quickReconnect() bool {
+	ctx, cancel := context.WithTimeout(g.ctx, 1500*time.Millisecond)
 	defer cancel()
 	<-ctx.Done()
 
-	if parentCtx.Err() != nil {
+	if g.ctx.Err() != nil {
 		return false
 	}
-	log.Infof("try to reconnect to Relay server: %s", rc.connectionURL)
 
-	if err := rc.Connect(); err != nil {
+	if err := g.relayClient.Connect(); err != nil {
 		log.Errorf("failed to reconnect to relay server: %s", err)
 		return false
 	}
 	return true
 }
-
-func (g *Guard) drainRelayClientChan() {
-	select {
-	case <-g.OnNewRelayClient:
-	default:
-	}
-}
-
-func (g *Guard) isServerURLStillValid(rc *Client) bool {
-	for _, url := range g.serverPicker.ServerURLs.Load().([]string) {
-		if url == rc.connectionURL {
-			return true
-		}
-	}
-	return false
-}
-
-func exponentTicker(ctx context.Context) *backoff.Ticker {
-	bo := backoff.WithContext(&backoff.ExponentialBackOff{
-		InitialInterval: 2 * time.Second,
-		Multiplier:      2,
-		MaxInterval:     reconnectingTimeout,
-		Clock:           backoff.SystemClock,
-	}, ctx)
-
-	return backoff.NewTicker(bo)
-}

relay/client/manager.go

@@ -57,15 +57,12 @@ type ManagerService interface {
 // relay servers will be closed if there is no active connection. Periodically the manager will check if there is any
 // unused relay connection and close it.
 type Manager struct {
-	ctx          context.Context
-	peerID       string
-	running      bool
-	tokenStore   *relayAuth.TokenStore
-	serverPicker *ServerPicker
+	ctx        context.Context
+	serverURLs []string
+	peerID     string
+	tokenStore *relayAuth.TokenStore
 
-	relayClient *Client
-	// the guard logic can overwrite the relayClient variable, this mutex protect the usage of the variable
-	relayClientMu  sync.Mutex
+	relayClient    *Client
 	reconnectGuard *Guard
 
 	relayClients      map[string]*RelayTrack
@@ -79,54 +76,48 @@ type Manager struct {
 // NewManager creates a new manager instance.
 // The serverURL address can be empty. In this case, the manager will not serve.
 func NewManager(ctx context.Context, serverURLs []string, peerID string) *Manager {
-	tokenStore := &relayAuth.TokenStore{}
-
-	m := &Manager{
-		ctx:        ctx,
-		peerID:     peerID,
-		tokenStore: tokenStore,
-		serverPicker: &ServerPicker{
-			TokenStore: tokenStore,
-			PeerID:     peerID,
-		},
+	return &Manager{
+		ctx:                     ctx,
+		serverURLs:              serverURLs,
+		peerID:                  peerID,
+		tokenStore:              &relayAuth.TokenStore{},
 		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),
 	}
-	m.serverPicker.ServerURLs.Store(serverURLs)
-	m.reconnectGuard = NewGuard(m.serverPicker)
-	return m
 }
 
-// Serve starts the manager, attempting to establish a connection with the relay server.
-// If the connection fails, it will keep trying to reconnect in the background.
-// Additionally, it starts a cleanup loop to remove unused relay connections.
-// The manager will automatically reconnect to the relay server in case of disconnection.
+// Serve starts the manager. It will establish a connection to the relay server and start the relay cleanup loop for
+// the unused relay connections. The manager will automatically reconnect to the relay server in case of disconnection.
 func (m *Manager) Serve() error {
-	if m.running {
+	if m.relayClient != nil {
 		return fmt.Errorf("manager already serving")
 	}
-	m.running = true
-	log.Debugf("starting relay client manager with %v relay servers", m.serverPicker.ServerURLs.Load())
+	log.Debugf("starting relay client manager with %v relay servers", m.serverURLs)
 
-	client, err := m.serverPicker.PickServer(m.ctx)
-	if err != nil {
-		go m.reconnectGuard.StartReconnectTrys(m.ctx, nil)
-	} else {
-		m.storeClient(client)
+	sp := ServerPicker{
+		TokenStore: m.tokenStore,
+		PeerID:     m.peerID,
 	}
 
-	go m.listenGuardEvent(m.ctx)
-	go m.startCleanupLoop()
-	return err
+	client, err := sp.PickServer(m.ctx, m.serverURLs)
+	if err != nil {
+		return err
+	}
+	m.relayClient = client
+
+	m.reconnectGuard = NewGuard(m.ctx, m.relayClient)
+	m.relayClient.SetOnConnectedListener(m.onServerConnected)
+	m.relayClient.SetOnDisconnectListener(func() {
+		m.onServerDisconnected(client.connectionURL)
+	})
+	m.startCleanupLoop()
+	return nil
 }
 
 // OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
 // established via the relay server. If the peer is on a different relay server, the manager will establish a new
 // connection to the relay server. It returns back with a net.Conn what represent the remote peer connection.
 func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return nil, ErrRelayClientNotConnected
 	}
@@ -155,9 +146,6 @@ func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 
 // Ready returns true if the home Relay client is connected to the relay server.
 func (m *Manager) Ready() bool {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return false
 	}
@@ -171,13 +159,6 @@ func (m *Manager) SetOnReconnectedListener(f func()) {
 // AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection
 // closed.
 func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
-	if m.relayClient == nil {
-		return ErrRelayClientNotConnected
-	}
-
 	foreign, err := m.isForeignServer(serverAddress)
 	if err != nil {
 		return err
@@ -196,9 +177,6 @@ func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServ
 // RelayInstanceAddress returns the address of the permanent relay server. It could change if the network connection is
 // lost. This address will be sent to the target peer to choose the common relay server for the communication.
 func (m *Manager) RelayInstanceAddress() (string, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return "", ErrRelayClientNotConnected
 	}
@@ -207,18 +185,13 @@ func (m *Manager) RelayInstanceAddress() (string, error) {
 
 // ServerURLs returns the addresses of the relay servers.
 func (m *Manager) ServerURLs() []string {
-	return m.serverPicker.ServerURLs.Load().([]string)
+	return m.serverURLs
 }
 
 // HasRelayAddress returns true if the manager is serving. With this method can check if the peer can communicate with
 // Relay service.
 func (m *Manager) HasRelayAddress() bool {
-	return len(m.serverPicker.ServerURLs.Load().([]string)) > 0
-}
-
-func (m *Manager) UpdateServerURLs(serverURLs []string) {
-	log.Infof("update relay server URLs: %v", serverURLs)
-	m.serverPicker.ServerURLs.Store(serverURLs)
+	return len(m.serverURLs) > 0
 }
 
 // UpdateToken updates the token in the token store.
@@ -272,7 +245,9 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 		return nil, err
 	}
 	// if connection closed then delete the relay client from the list
-	relayClient.SetOnDisconnectListener(m.onServerDisconnected)
+	relayClient.SetOnDisconnectListener(func() {
+		m.onServerDisconnected(serverAddress)
+	})
 	rt.relayClient = relayClient
 	rt.Unlock()
 
@@ -290,37 +265,14 @@ func (m *Manager) onServerConnected() {
 	go m.onReconnectedListenerFn()
 }
 
-// onServerDisconnected start to reconnection for home server only
 func (m *Manager) onServerDisconnected(serverAddress string) {
-	m.relayClientMu.Lock()
 	if serverAddress == m.relayClient.connectionURL {
-		go m.reconnectGuard.StartReconnectTrys(m.ctx, m.relayClient)
+		go m.reconnectGuard.OnDisconnected()
 	}
-	m.relayClientMu.Unlock()
 
 	m.notifyOnDisconnectListeners(serverAddress)
 }
 
-func (m *Manager) listenGuardEvent(ctx context.Context) {
-	for {
-		select {
-		case rc := <-m.reconnectGuard.OnNewRelayClient:
-			m.storeClient(rc)
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
-func (m *Manager) storeClient(client *Client) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
-	m.relayClient = client
-	m.relayClient.SetOnConnectedListener(m.onServerConnected)
-	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
-}
-
 func (m *Manager) isForeignServer(address string) (bool, error) {
 	rAddr, err := m.relayClient.ServerInstanceURL()
 	if err != nil {
@@ -330,16 +282,22 @@ func (m *Manager) isForeignServer(address string) (bool, error) {
 }
 
 func (m *Manager) startCleanupLoop() {
-	ticker := time.NewTicker(relayCleanupInterval)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-m.ctx.Done():
-			return
-		case <-ticker.C:
-			m.cleanUpUnusedRelays()
-		}
+	if m.ctx.Err() != nil {
+		return
 	}
+
+	ticker := time.NewTicker(relayCleanupInterval)
+	go func() {
+		defer ticker.Stop()
+		for {
+			select {
+			case <-m.ctx.Done():
+				return
+			case <-ticker.C:
+				m.cleanUpUnusedRelays()
+			}
+		}
+	}()
 }
 
 func (m *Manager) cleanUpUnusedRelays() {

relay/client/picker.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync/atomic"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -13,13 +12,10 @@ import (
 )
 
 const (
+	connectionTimeout    = 30 * time.Second
 	maxConcurrentServers = 7
 )
 
-var (
-	connectionTimeout = 30 * time.Second
-)
-
 type connResult struct {
 	RelayClient *Client
 	Url         string
@@ -28,22 +24,20 @@ type connResult struct {
 
 type ServerPicker struct {
 	TokenStore *auth.TokenStore
-	ServerURLs atomic.Value
 	PeerID     string
 }
 
-func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
+func (sp *ServerPicker) PickServer(parentCtx context.Context, urls []string) (*Client, error) {
 	ctx, cancel := context.WithTimeout(parentCtx, connectionTimeout)
 	defer cancel()
 
-	totalServers := len(sp.ServerURLs.Load().([]string))
+	totalServers := len(urls)
 
 	connResultChan := make(chan connResult, totalServers)
 	successChan := make(chan connResult, 1)
 	concurrentLimiter := make(chan struct{}, maxConcurrentServers)
 
-	log.Debugf("pick server from list: %v", sp.ServerURLs.Load().([]string))
-	for _, url := range sp.ServerURLs.Load().([]string) {
+	for _, url := range urls {
 		// todo check if we have a successful connection so we do not need to connect to other servers
 		concurrentLimiter <- struct{}{}
 		go func(url string) {
@@ -84,7 +78,7 @@ func (sp *ServerPicker) processConnResults(resultChan chan connResult, successCh
 	for numOfResults := 0; numOfResults < cap(resultChan); numOfResults++ {
 		cr := <-resultChan
 		if cr.Err != nil {
-			log.Tracef("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
+			log.Debugf("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
 			continue
 		}
 		log.Infof("connected to Relay server: %s", cr.Url)

relay/client/picker_test.go

@@ -7,19 +7,16 @@ import (
 )
 
 func TestServerPicker_UnavailableServers(t *testing.T) {
-	connectionTimeout = 5 * time.Second
-
 	sp := ServerPicker{
 		TokenStore: nil,
 		PeerID:     "test",
 	}
-	sp.ServerURLs.Store([]string{"rel://dummy1", "rel://dummy2"})
 
 	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
 	defer cancel()
 
 	go func() {
-		_, err := sp.PickServer(ctx)
+		_, err := sp.PickServer(ctx, []string{"rel://dummy1", "rel://dummy2"})
 		if err == nil {
 			t.Error(err)
 		}

relay/server/listener/quic/conn.go

@@ -1,112 +0,0 @@
-package quic
-
-import (
-	"context"
-	"errors"
-	"io"
-	"net"
-	"sync"
-	"time"
-
-	"github.com/quic-go/quic-go"
-)
-
-const (
-	writeTimeout = 10 * time.Second
-)
-
-type Conn struct {
-	session   quic.Connection
-	closed    bool
-	closedMu  sync.Mutex
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-}
-
-func NewConn(session quic.Connection) *Conn {
-	ctx, cancel := context.WithCancel(context.Background())
-	return &Conn{
-		session:   session,
-		ctx:       ctx,
-		ctxCancel: cancel,
-	}
-}
-
-func (c *Conn) Read(b []byte) (n int, err error) {
-	if c.isClosed() {
-		return 0, io.EOF
-	}
-
-	dgram, err := c.session.ReceiveDatagram(c.ctx)
-	if err != nil {
-		return 0, c.ioErrHandling(err)
-	}
-	// Copy data to b, ensuring we don’t exceed the size of b
-	n = copy(b, dgram)
-	return n, nil
-}
-
-func (c *Conn) Write(b []byte) (int, error) {
-	err := c.session.SendDatagram(b)
-	return len(b), err
-}
-
-func (c *Conn) LocalAddr() net.Addr {
-	return c.session.LocalAddr()
-}
-
-func (c *Conn) RemoteAddr() net.Addr {
-	return c.session.RemoteAddr()
-}
-
-func (c *Conn) SetReadDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetWriteDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) Close() error {
-	c.closedMu.Lock()
-	if c.closed {
-		c.closedMu.Unlock()
-		return nil
-	}
-	c.closed = true
-	c.closedMu.Unlock()
-
-	c.ctxCancel() // Cancel the context
-
-	sessionErr := c.session.CloseWithError(0, "normal closure")
-	return sessionErr
-}
-
-func (c *Conn) isClosed() bool {
-	c.closedMu.Lock()
-	defer c.closedMu.Unlock()
-	return c.closed
-}
-
-func (c *Conn) ioErrHandling(err error) error {
-	if c.isClosed() {
-		return io.EOF
-	}
-
-	// Handle QUIC-specific errors
-	if err == nil {
-		return nil
-	}
-
-	// Check if the connection was closed remotely
-	var appErr *quic.ApplicationError
-	if errors.As(err, &appErr) && appErr.ErrorCode == 0 { // 0 is normal closure
-		return io.EOF
-	}
-
-	return err
-}

relay/server/listener/quic/listener.go

@@ -1,66 +0,0 @@
-package quic
-
-import (
-	"context"
-	"crypto/tls"
-	"errors"
-	"fmt"
-	"net"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-type Listener struct {
-	// Address is the address to listen on
-	Address string
-	// TLSConfig is the TLS configuration for the server
-	TLSConfig *tls.Config
-
-	listener *quic.Listener
-	acceptFn func(conn net.Conn)
-}
-
-func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
-	l.acceptFn = acceptFn
-
-	quicCfg := &quic.Config{
-		EnableDatagrams: true,
-	}
-	listener, err := quic.ListenAddr(l.Address, l.TLSConfig, quicCfg)
-	if err != nil {
-		return fmt.Errorf("failed to create QUIC listener: %v", err)
-	}
-
-	l.listener = listener
-	log.Infof("QUIC server listening on address: %s", l.Address)
-
-	for {
-		session, err := listener.Accept(context.Background())
-		if err != nil {
-			if errors.Is(err, quic.ErrServerClosed) {
-				return nil
-			}
-
-			log.Errorf("Failed to accept QUIC session: %v", err)
-			continue
-		}
-
-		log.Infof("QUIC client connected from: %s", session.RemoteAddr())
-		conn := NewConn(session)
-		l.acceptFn(conn)
-	}
-}
-
-func (l *Listener) Shutdown(ctx context.Context) error {
-	if l.listener == nil {
-		return nil
-	}
-
-	log.Infof("stopping QUIC listener")
-	if err := l.listener.Close(); err != nil {
-		return fmt.Errorf("listener shutdown failed: %v", err)
-	}
-	log.Infof("QUIC listener stopped")
-	return nil
-}

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

@@ -23,7 +23,6 @@ type Listener struct {
 
 	server   *http.Server
 	acceptFn func(conn net.Conn)
-	log      *log.Entry
 }
 
 func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
@@ -64,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 {
@@ -89,15 +87,6 @@ func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	log.Infof("WS client connected from: %s", rAddr)
-
 	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"
@@ -16,8 +16,6 @@ import (
 
 const (
 	bufferSize = 8820
-
-	errCloseConn = "failed to close connection to peer: %s"
 )
 
 // Peer represents a peer connection
@@ -48,12 +46,6 @@ func NewPeer(metrics *metrics.Metrics, id []byte, conn net.Conn, store *Store) *
 // It manages the protocol (healthcheck, transport, close). Read the message and determine the message type and handle
 // the message accordingly.
 func (p *Peer) Work() {
-	defer func() {
-		if err := p.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
-			p.log.Errorf(errCloseConn, err)
-		}
-	}()
-
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -65,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
@@ -105,7 +97,7 @@ func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *
 	case messages.MsgTypeClose:
 		p.log.Infof("peer exited gracefully")
 		if err := p.conn.Close(); err != nil {
-			log.Errorf(errCloseConn, err)
+			log.Errorf("failed to close connection to peer: %s", err)
 		}
 	default:
 		p.log.Warnf("received unexpected message type: %s", msgType)
@@ -129,8 +121,9 @@ func (p *Peer) CloseGracefully(ctx context.Context) {
 		p.log.Errorf("failed to send close message to peer: %s", p.String())
 	}
 
-	if err := p.conn.Close(); err != nil {
-		p.log.Errorf(errCloseConn, err)
+	err = p.conn.Close()
+	if err != nil {
+		p.log.Errorf("failed to close connection to peer: %s", err)
 	}
 }
 
@@ -139,7 +132,7 @@ func (p *Peer) Close() {
 	defer p.connMu.Unlock()
 
 	if err := p.conn.Close(); err != nil {
-		p.log.Errorf(errCloseConn, err)
+		p.log.Errorf("failed to close connection to peer: %s", err)
 	}
 }
 

relay/server/server.go

@@ -2,25 +2,13 @@ package server
 
 import (
 	"context"
-	"crypto/rand"
-	"crypto/rsa"
 	"crypto/tls"
-	"crypto/x509"
-	"crypto/x509/pkix"
-	"encoding/pem"
-	"math/big"
-	"net"
-	"sync"
-	"time"
 
-	"github.com/hashicorp/go-multierror"
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel/metric"
 
-	nberrors "github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/relay/auth"
 	"github.com/netbirdio/netbird/relay/server/listener"
-	"github.com/netbirdio/netbird/relay/server/listener/quic"
 	"github.com/netbirdio/netbird/relay/server/listener/ws"
 )
 
@@ -36,8 +24,8 @@ type ListenerConfig struct {
 // It is the gate between the WebSocket listener and the Relay server logic.
 // In a new HTTP connection, the server will accept the connection and pass it to the Relay server via the Accept method.
 type Server struct {
-	relay     *Relay
-	listeners []listener.Listener
+	relay      *Relay
+	wSListener listener.Listener
 }
 
 // NewServer creates a new relay server instance.
@@ -51,120 +39,38 @@ func NewServer(meter metric.Meter, exposedAddress string, tlsSupport bool, authV
 		return nil, err
 	}
 	return &Server{
-		relay:     relay,
-		listeners: make([]listener.Listener, 0, 2),
+		relay: relay,
 	}, nil
 }
 
 // Listen starts the relay server.
 func (r *Server) Listen(cfg ListenerConfig) error {
-	wSListener := &ws.Listener{
+	r.wSListener = &ws.Listener{
 		Address:   cfg.Address,
 		TLSConfig: cfg.TLSConfig,
 	}
-	r.listeners = append(r.listeners, wSListener)
 
-	quicListener := &quic.Listener{
-		Address: cfg.Address,
+	wslErr := r.wSListener.Listen(r.relay.Accept)
+	if wslErr != nil {
+		log.Errorf("failed to bind ws server: %s", wslErr)
 	}
 
-	if cfg.TLSConfig != nil {
-		quicListener.TLSConfig = cfg.TLSConfig
-	} else {
-		tlsConfig, err := generateTestTLSConfig()
-		if err != nil {
-			return err
-		}
-		quicListener.TLSConfig = tlsConfig
-	}
-	r.listeners = append(r.listeners, quicListener)
-
-	errChan := make(chan error, len(r.listeners))
-	wg := sync.WaitGroup{}
-	for _, l := range r.listeners {
-		wg.Add(1)
-		go func(listener listener.Listener) {
-			defer wg.Done()
-			errChan <- listener.Listen(r.relay.Accept)
-		}(l)
-	}
-
-	wg.Wait()
-	close(errChan)
-	var multiErr *multierror.Error
-	for err := range errChan {
-		multiErr = multierror.Append(multiErr, err)
-	}
-
-	return nberrors.FormatErrorOrNil(multiErr)
+	return wslErr
 }
 
 // Shutdown stops the relay server. If there are active connections, they will be closed gracefully. In case of a context,
 // the connections will be forcefully closed.
-func (r *Server) Shutdown(ctx context.Context) error {
-	var multiErr *multierror.Error
-	for _, l := range r.listeners {
-		if err := l.Shutdown(ctx); err != nil {
-			multiErr = multierror.Append(multiErr, err)
-		}
+func (r *Server) Shutdown(ctx context.Context) (err error) {
+	// stop service new connections
+	if r.wSListener != nil {
+		err = r.wSListener.Shutdown(ctx)
 	}
 
 	r.relay.Shutdown(ctx)
-	return nberrors.FormatErrorOrNil(multiErr)
+	return
 }
 
 // InstanceURL returns the instance URL of the relay server.
 func (r *Server) InstanceURL() string {
 	return r.relay.instanceURL
 }
-
-// GenerateTestTLSConfig creates a self-signed certificate for testing
-func generateTestTLSConfig() (*tls.Config, error) {
-	log.Infof("generating test TLS config")
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-	if err != nil {
-		return nil, err
-	}
-
-	template := x509.Certificate{
-		SerialNumber: big.NewInt(1),
-		Subject: pkix.Name{
-			Organization: []string{"Test Organization"},
-		},
-		NotBefore: time.Now(),
-		NotAfter:  time.Now().Add(time.Hour * 24 * 180), // Valid for 180 days
-		KeyUsage:  x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
-		ExtKeyUsage: []x509.ExtKeyUsage{
-			x509.ExtKeyUsageServerAuth,
-		},
-		BasicConstraintsValid: true,
-		DNSNames:              []string{"localhost"},
-		IPAddresses:           []net.IP{net.ParseIP("192.168.0.10")},
-	}
-
-	// Create certificate
-	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
-	if err != nil {
-		return nil, err
-	}
-
-	certPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "CERTIFICATE",
-		Bytes: certDER,
-	})
-
-	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "RSA PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
-	})
-
-	tlsCert, err := tls.X509KeyPair(certPEM, privateKeyPEM)
-	if err != nil {
-		return nil, err
-	}
-
-	return &tls.Config{
-		Certificates: []tls.Certificate{tlsCert},
-		NextProtos:   []string{"netbird-relay"},
-	}, nil
-}

util/file.go

@@ -15,7 +15,7 @@ import (
 )
 
 // WriteJsonWithRestrictedPermission writes JSON config object to a file. Enforces permission on the parent directory
-func WriteJsonWithRestrictedPermission(ctx context.Context, file string, obj interface{}) error {
+func WriteJsonWithRestrictedPermission(file string, obj interface{}) error {
 	configDir, configFileName, err := prepareConfigFileDir(file)
 	if err != nil {
 		return err
@@ -26,18 +26,18 @@ func WriteJsonWithRestrictedPermission(ctx context.Context, file string, obj int
 		return err
 	}
 
-	return writeJson(ctx, file, obj, configDir, configFileName)
+	return writeJson(file, obj, configDir, configFileName)
 }
 
 // WriteJson writes JSON config object to a file creating parent directories if required
 // The output JSON is pretty-formatted
-func WriteJson(ctx context.Context, file string, obj interface{}) error {
+func WriteJson(file string, obj interface{}) error {
 	configDir, configFileName, err := prepareConfigFileDir(file)
 	if err != nil {
 		return err
 	}
 
-	return writeJson(ctx, file, obj, configDir, configFileName)
+	return writeJson(file, obj, configDir, configFileName)
 }
 
 // DirectWriteJson writes JSON config object to a file creating parent directories if required without creating a temporary file
@@ -79,11 +79,7 @@ func DirectWriteJson(ctx context.Context, file string, obj interface{}) error {
 	return nil
 }
 
-func writeJson(ctx context.Context, file string, obj interface{}, configDir string, configFileName string) error {
-	// Check context before expensive operations
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
+func writeJson(file string, obj interface{}, configDir string, configFileName string) error {
 
 	// make it pretty
 	bs, err := json.MarshalIndent(obj, "", "    ")
@@ -91,10 +87,6 @@ func writeJson(ctx context.Context, file string, obj interface{}, configDir stri
 		return err
 	}
 
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
 	tempFile, err := os.CreateTemp(configDir, ".*"+configFileName)
 	if err != nil {
 		return err
@@ -119,11 +111,6 @@ func writeJson(ctx context.Context, file string, obj interface{}, configDir stri
 		return err
 	}
 
-	// Check context again
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
 	err = os.Rename(tempFileName, file)
 	if err != nil {
 		return err

util/file_test.go

@@ -1,7 +1,6 @@
 package util
 
 import (
-	"context"
 	"crypto/md5"
 	"encoding/hex"
 	"io"
@@ -40,7 +39,7 @@ func TestConfigJSON(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			tmpDir := t.TempDir()
 
-			err := WriteJson(context.Background(), tmpDir+"/testconfig.json", tt.config)
+			err := WriteJson(tmpDir+"/testconfig.json", tt.config)
 			require.NoError(t, err)
 
 			read, err := ReadJson(tmpDir+"/testconfig.json", &TestConfig{})
@@ -74,7 +73,7 @@ func TestCopyFileContents(t *testing.T) {
 			src := tmpDir + "/copytest_src"
 			dst := tmpDir + "/copytest_dst"
 
-			err := WriteJson(context.Background(), src, tt.srcContent)
+			err := WriteJson(src, tt.srcContent)
 			require.NoError(t, err)
 
 			err = CopyFileContents(src, dst)
@@ -128,7 +127,7 @@ func TestHandleConfigFileWithoutFullPath(t *testing.T) {
 				_ = os.Remove(cfgFile)
 			}()
 
-			err := WriteJson(context.Background(), cfgFile, tt.config)
+			err := WriteJson(cfgFile, tt.config)
 			require.NoError(t, err)
 
 			read, err := ReadJson(cfgFile, &TestConfig{})

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"
 )