Add missing store tests

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>

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

@@ -52,47 +52,6 @@ jobs:
       - name: Test
         run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 10m -p 1 ./...
 
-  benchmark:
-    strategy:
-      fail-fast: false
-      matrix:
-        arch: [ '386','amd64' ]
-        store: [ 'sqlite', 'postgres' ]
-    runs-on: ubuntu-22.04
-    steps:
-      - name: Install Go
-        uses: actions/setup-go@v5
-        with:
-          go-version: "1.23.x"
-
-
-      - name: Cache Go modules
-        uses: actions/cache@v4
-        with:
-          path: ~/go/pkg/mod
-          key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
-          restore-keys: |
-            ${{ runner.os }}-go-
-
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Install dependencies
-        run: sudo apt update && sudo apt install -y -q libgtk-3-dev libayatana-appindicator3-dev libgl1-mesa-dev xorg-dev gcc-multilib libpcap-dev
-
-      - name: Install 32-bit libpcap
-        if: matrix.arch == '386'
-        run: sudo dpkg --add-architecture i386 && sudo apt update && sudo apt-get install -y libpcap0.8-dev:i386
-
-      - name: Install modules
-        run: go mod tidy
-
-      - name: check git status
-        run: git --no-pager diff --exit-code
-
-      - name: Test
-        run: CGO_ENABLED=1 GOARCH=${{ matrix.arch }} NETBIRD_STORE_ENGINE=${{ matrix.store }} CI=true go test -run=^$ -bench=. -exec 'sudo --preserve-env=CI,NETBIRD_STORE_ENGINE' -timeout 10m -p 1 ./...
-
   test_client_on_docker:
     runs-on: ubuntu-20.04
     steps:

.github/workflows/release.yml

@@ -9,7 +9,7 @@ on:
   pull_request:
 
 env:
-  SIGN_PIPE_VER: "v0.0.17"
+  SIGN_PIPE_VER: "v0.0.16"
   GORELEASER_VER: "v2.3.2"
   PRODUCT_NAME: "NetBird"
   COPYRIGHT: "Wiretrustee UG (haftungsbeschreankt)"

README.md

@@ -17,12 +17,8 @@
        <img src="https://img.shields.io/badge/license-BSD--3-blue" />
      </a> 
     <br>
-    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">
+    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2p5zwhm4g-8fHollzrQa5y4PZF5AEpvQ">
         <img src="https://img.shields.io/badge/slack-@netbird-red.svg?logo=slack"/>
-     </a>  
-     <br>
-    <a href="https://gurubase.io/g/netbird">
-        <img src="https://img.shields.io/badge/Gurubase-Ask%20NetBird%20Guru-006BFF"/>
      </a>    
   </p>
 </div>
@@ -34,7 +30,7 @@
   <br/>
   See <a href="https://netbird.io/docs/">Documentation</a>
   <br/>
-   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">Slack channel</a>
+   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2p5zwhm4g-8fHollzrQa5y4PZF5AEpvQ">Slack channel</a>
   <br/>
  
 </strong>

client/anonymize/anonymize.go

@@ -12,8 +12,6 @@ import (
 	"strings"
 )
 
-const anonTLD = ".domain"
-
 type Anonymizer struct {
 	ipAnonymizer     map[netip.Addr]netip.Addr
 	domainAnonymizer map[string]string
@@ -85,39 +83,29 @@ func (a *Anonymizer) AnonymizeIPString(ip string) string {
 }
 
 func (a *Anonymizer) AnonymizeDomain(domain string) string {
-	baseDomain := domain
-	hasDot := strings.HasSuffix(domain, ".")
-	if hasDot {
-		baseDomain = domain[:len(domain)-1]
-	}
-
-	if strings.HasSuffix(baseDomain, "netbird.io") ||
-		strings.HasSuffix(baseDomain, "netbird.selfhosted") ||
-		strings.HasSuffix(baseDomain, "netbird.cloud") ||
-		strings.HasSuffix(baseDomain, "netbird.stage") ||
-		strings.HasSuffix(baseDomain, anonTLD) {
+	if strings.HasSuffix(domain, "netbird.io") ||
+		strings.HasSuffix(domain, "netbird.selfhosted") ||
+		strings.HasSuffix(domain, "netbird.cloud") ||
+		strings.HasSuffix(domain, "netbird.stage") ||
+		strings.HasSuffix(domain, ".domain") {
 		return domain
 	}
 
-	parts := strings.Split(baseDomain, ".")
+	parts := strings.Split(domain, ".")
 	if len(parts) < 2 {
 		return domain
 	}
 
-	baseForLookup := parts[len(parts)-2] + "." + parts[len(parts)-1]
+	baseDomain := parts[len(parts)-2] + "." + parts[len(parts)-1]
 
-	anonymized, ok := a.domainAnonymizer[baseForLookup]
+	anonymized, ok := a.domainAnonymizer[baseDomain]
 	if !ok {
-		anonymizedBase := "anon-" + generateRandomString(5) + anonTLD
-		a.domainAnonymizer[baseForLookup] = anonymizedBase
+		anonymizedBase := "anon-" + generateRandomString(5) + ".domain"
+		a.domainAnonymizer[baseDomain] = anonymizedBase
 		anonymized = anonymizedBase
 	}
 
-	result := strings.Replace(baseDomain, baseForLookup, anonymized, 1)
-	if hasDot {
-		result += "."
-	}
-	return result
+	return strings.Replace(domain, baseDomain, anonymized, 1)
 }
 
 func (a *Anonymizer) AnonymizeURI(uri string) string {
@@ -164,9 +152,9 @@ func (a *Anonymizer) AnonymizeString(str string) string {
 	return str
 }
 
-// AnonymizeSchemeURI finds and anonymizes URIs with ws, wss, rel, rels, stun, stuns, turn, and turns schemes.
+// AnonymizeSchemeURI finds and anonymizes URIs with stun, stuns, turn, and turns schemes.
 func (a *Anonymizer) AnonymizeSchemeURI(text string) string {
-	re := regexp.MustCompile(`(?i)\b(wss?://|rels?://|stuns?:|turns?:|https?://)\S+\b`)
+	re := regexp.MustCompile(`(?i)\b(stuns?:|turns?:|https?://)\S+\b`)
 
 	return re.ReplaceAllStringFunc(text, a.AnonymizeURI)
 }
@@ -180,10 +168,10 @@ func (a *Anonymizer) AnonymizeDNSLogLine(logEntry string) string {
 		parts := strings.Split(match, `"`)
 		if len(parts) >= 2 {
 			domain := parts[1]
-			if strings.HasSuffix(domain, anonTLD) {
+			if strings.HasSuffix(domain, ".domain") {
 				return match
 			}
-			randomDomain := generateRandomString(10) + anonTLD
+			randomDomain := generateRandomString(10) + ".domain"
 			return strings.Replace(match, domain, randomDomain, 1)
 		}
 		return match

client/anonymize/anonymize_test.go

@@ -67,36 +67,18 @@ func TestAnonymizeDomain(t *testing.T) {
 			`^anon-[a-zA-Z0-9]+\.domain$`,
 			true,
 		},
-		{
-			"Domain with Trailing Dot",
-			"example.com.",
-			`^anon-[a-zA-Z0-9]+\.domain.$`,
-			true,
-		},
 		{
 			"Subdomain",
 			"sub.example.com",
 			`^sub\.anon-[a-zA-Z0-9]+\.domain$`,
 			true,
 		},
-		{
-			"Subdomain with Trailing Dot",
-			"sub.example.com.",
-			`^sub\.anon-[a-zA-Z0-9]+\.domain.$`,
-			true,
-		},
 		{
 			"Protected Domain",
 			"netbird.io",
 			`^netbird\.io$`,
 			false,
 		},
-		{
-			"Protected Domain with Trailing Dot",
-			"netbird.io.",
-			`^netbird\.io.$`,
-			false,
-		},
 	}
 
 	for _, tc := range tests {
@@ -158,16 +140,8 @@ func TestAnonymizeSchemeURI(t *testing.T) {
 		expect string
 	}{
 		{"STUN URI in text", "Connection made via stun:example.com", `Connection made via stun:anon-[a-zA-Z0-9]+\.domain`},
-		{"STUNS URI in message", "Secure connection to stuns:example.com:443", `Secure connection to stuns:anon-[a-zA-Z0-9]+\.domain:443`},
 		{"TURN URI in log", "Failed attempt turn:some.example.com:3478?transport=tcp: retrying", `Failed attempt turn:some.anon-[a-zA-Z0-9]+\.domain:3478\?transport=tcp: retrying`},
-		{"TURNS URI in message", "Secure connection to turns:example.com:5349", `Secure connection to turns:anon-[a-zA-Z0-9]+\.domain:5349`},
-		{"HTTP URI in text", "Visit http://example.com for more", `Visit http://anon-[a-zA-Z0-9]+\.domain for more`},
-		{"HTTPS URI in CAPS", "Visit HTTPS://example.com for more", `Visit https://anon-[a-zA-Z0-9]+\.domain for more`},
 		{"HTTPS URI in message", "Visit https://example.com for more", `Visit https://anon-[a-zA-Z0-9]+\.domain for more`},
-		{"WS URI in log", "Connection established to ws://example.com:8080", `Connection established to ws://anon-[a-zA-Z0-9]+\.domain:8080`},
-		{"WSS URI in message", "Secure connection to wss://example.com", `Secure connection to wss://anon-[a-zA-Z0-9]+\.domain`},
-		{"Rel URI in text", "Relaying to rel://example.com", `Relaying to rel://anon-[a-zA-Z0-9]+\.domain`},
-		{"Rels URI in message", "Relaying to rels://example.com", `Relaying to rels://anon-[a-zA-Z0-9]+\.domain`},
 	}
 
 	for _, tc := range tests {

client/cmd/debug.go

@@ -3,7 +3,6 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -62,15 +61,6 @@ var forCmd = &cobra.Command{
 	RunE:    runForDuration,
 }
 
-var persistenceCmd = &cobra.Command{
-	Use:     "persistence [on|off]",
-	Short:   "Set network map memory persistence",
-	Long:    `Configure whether the latest network map should persist in memory. When enabled, the last known network map will be kept in memory.`,
-	Example: "  netbird debug persistence on",
-	Args:    cobra.ExactArgs(1),
-	RunE:    setNetworkMapPersistence,
-}
-
 func debugBundle(cmd *cobra.Command, _ []string) error {
 	conn, err := getClient(cmd)
 	if err != nil {
@@ -181,13 +171,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 
 	time.Sleep(1 * time.Second)
 
-	// Enable network map persistence before bringing the service up
-	if _, err := client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
-		Enabled: true,
-	}); err != nil {
-		return fmt.Errorf("failed to enable network map persistence: %v", status.Convert(err).Message())
-	}
-
 	if _, err := client.Up(cmd.Context(), &proto.UpRequest{}); err != nil {
 		return fmt.Errorf("failed to up: %v", status.Convert(err).Message())
 	}
@@ -217,13 +200,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to bundle debug: %v", status.Convert(err).Message())
 	}
 
-	// Disable network map persistence after creating the debug bundle
-	if _, err := client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
-		Enabled: false,
-	}); err != nil {
-		return fmt.Errorf("failed to disable network map persistence: %v", status.Convert(err).Message())
-	}
-
 	if stateWasDown {
 		if _, err := client.Down(cmd.Context(), &proto.DownRequest{}); err != nil {
 			return fmt.Errorf("failed to down: %v", status.Convert(err).Message())
@@ -243,34 +219,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-func setNetworkMapPersistence(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		if err := conn.Close(); err != nil {
-			log.Errorf(errCloseConnection, err)
-		}
-	}()
-
-	persistence := strings.ToLower(args[0])
-	if persistence != "on" && persistence != "off" {
-		return fmt.Errorf("invalid persistence value: %s. Use 'on' or 'off'", args[0])
-	}
-
-	client := proto.NewDaemonServiceClient(conn)
-	_, err = client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
-		Enabled: persistence == "on",
-	})
-	if err != nil {
-		return fmt.Errorf("failed to set network map persistence: %v", status.Convert(err).Message())
-	}
-
-	cmd.Printf("Network map persistence set to: %s\n", persistence)
-	return nil
-}
-
 func getStatusOutput(cmd *cobra.Command) string {
 	var statusOutputString string
 	statusResp, err := getStatus(cmd.Context())

client/cmd/pprof.go

@@ -1,33 +0,0 @@
-//go:build pprof
-// +build pprof
-
-package cmd
-
-import (
-	"net/http"
-	_ "net/http/pprof"
-	"os"
-
-	log "github.com/sirupsen/logrus"
-)
-
-func init() {
-	addr := pprofAddr()
-	go pprof(addr)
-}
-
-func pprofAddr() string {
-	listenAddr := os.Getenv("NB_PPROF_ADDR")
-	if listenAddr == "" {
-		return "localhost:6969"
-	}
-
-	return listenAddr
-}
-
-func pprof(listenAddr string) {
-	log.Infof("listening pprof on: %s\n", listenAddr)
-	if err := http.ListenAndServe(listenAddr, nil); err != nil {
-		log.Fatalf("Failed to start pprof: %v", err)
-	}
-}

client/cmd/root.go

@@ -155,7 +155,6 @@ func init() {
 	debugCmd.AddCommand(logCmd)
 	logCmd.AddCommand(logLevelCmd)
 	debugCmd.AddCommand(forCmd)
-	debugCmd.AddCommand(persistenceCmd)
 
 	upCmd.PersistentFlags().StringSliceVar(&natExternalIPs, externalIPMapFlag, nil,
 		`Sets external IPs maps between local addresses and interfaces.`+

client/cmd/state.go

@@ -1,181 +0,0 @@
-package cmd
-
-import (
-	"fmt"
-
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
-	"google.golang.org/grpc/status"
-
-	"github.com/netbirdio/netbird/client/proto"
-)
-
-var (
-	allFlag bool
-)
-
-var stateCmd = &cobra.Command{
-	Use:   "state",
-	Short: "Manage daemon state",
-	Long:  "Provides commands for managing and inspecting the Netbird daemon state.",
-}
-
-var stateListCmd = &cobra.Command{
-	Use:     "list",
-	Aliases: []string{"ls"},
-	Short:   "List all stored states",
-	Long:    "Lists all registered states with their status and basic information.",
-	Example: "  netbird state list",
-	RunE:    stateList,
-}
-
-var stateCleanCmd = &cobra.Command{
-	Use:   "clean [state-name]",
-	Short: "Clean stored states",
-	Long: `Clean specific state or all states. The daemon must not be running.
-This will perform cleanup operations and remove the state.`,
-	Example: `  netbird state clean dns_state
-  netbird state clean --all`,
-	RunE: stateClean,
-	PreRunE: func(cmd *cobra.Command, args []string) error {
-		// Check mutual exclusivity between --all flag and state-name argument
-		if allFlag && len(args) > 0 {
-			return fmt.Errorf("cannot specify both --all flag and state name")
-		}
-		if !allFlag && len(args) != 1 {
-			return fmt.Errorf("requires a state name argument or --all flag")
-		}
-		return nil
-	},
-}
-
-var stateDeleteCmd = &cobra.Command{
-	Use:   "delete [state-name]",
-	Short: "Delete stored states",
-	Long: `Delete specific state or all states from storage. The daemon must not be running.
-This will remove the state without performing any cleanup operations.`,
-	Example: `  netbird state delete dns_state
-  netbird state delete --all`,
-	RunE: stateDelete,
-	PreRunE: func(cmd *cobra.Command, args []string) error {
-		// Check mutual exclusivity between --all flag and state-name argument
-		if allFlag && len(args) > 0 {
-			return fmt.Errorf("cannot specify both --all flag and state name")
-		}
-		if !allFlag && len(args) != 1 {
-			return fmt.Errorf("requires a state name argument or --all flag")
-		}
-		return nil
-	},
-}
-
-func init() {
-	rootCmd.AddCommand(stateCmd)
-	stateCmd.AddCommand(stateListCmd, stateCleanCmd, stateDeleteCmd)
-
-	stateCleanCmd.Flags().BoolVarP(&allFlag, "all", "a", false, "Clean all states")
-	stateDeleteCmd.Flags().BoolVarP(&allFlag, "all", "a", false, "Delete all states")
-}
-
-func stateList(cmd *cobra.Command, _ []string) error {
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		if err := conn.Close(); err != nil {
-			log.Errorf(errCloseConnection, err)
-		}
-	}()
-
-	client := proto.NewDaemonServiceClient(conn)
-	resp, err := client.ListStates(cmd.Context(), &proto.ListStatesRequest{})
-	if err != nil {
-		return fmt.Errorf("failed to list states: %v", status.Convert(err).Message())
-	}
-
-	cmd.Printf("\nStored states:\n\n")
-	for _, state := range resp.States {
-		cmd.Printf("- %s\n", state.Name)
-	}
-
-	return nil
-}
-
-func stateClean(cmd *cobra.Command, args []string) error {
-	var stateName string
-	if !allFlag {
-		stateName = args[0]
-	}
-
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		if err := conn.Close(); err != nil {
-			log.Errorf(errCloseConnection, err)
-		}
-	}()
-
-	client := proto.NewDaemonServiceClient(conn)
-	resp, err := client.CleanState(cmd.Context(), &proto.CleanStateRequest{
-		StateName: stateName,
-		All:       allFlag,
-	})
-	if err != nil {
-		return fmt.Errorf("failed to clean state: %v", status.Convert(err).Message())
-	}
-
-	if resp.CleanedStates == 0 {
-		cmd.Println("No states were cleaned")
-		return nil
-	}
-
-	if allFlag {
-		cmd.Printf("Successfully cleaned %d states\n", resp.CleanedStates)
-	} else {
-		cmd.Printf("Successfully cleaned state %q\n", stateName)
-	}
-
-	return nil
-}
-
-func stateDelete(cmd *cobra.Command, args []string) error {
-	var stateName string
-	if !allFlag {
-		stateName = args[0]
-	}
-
-	conn, err := getClient(cmd)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		if err := conn.Close(); err != nil {
-			log.Errorf(errCloseConnection, err)
-		}
-	}()
-
-	client := proto.NewDaemonServiceClient(conn)
-	resp, err := client.DeleteState(cmd.Context(), &proto.DeleteStateRequest{
-		StateName: stateName,
-		All:       allFlag,
-	})
-	if err != nil {
-		return fmt.Errorf("failed to delete state: %v", status.Convert(err).Message())
-	}
-
-	if resp.DeletedStates == 0 {
-		cmd.Println("No states were deleted")
-		return nil
-	}
-
-	if allFlag {
-		cmd.Printf("Successfully deleted %d states\n", resp.DeletedStates)
-	} else {
-		cmd.Printf("Successfully deleted state %q\n", stateName)
-	}
-
-	return nil
-}

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

@@ -37,11 +37,6 @@ func (s *ipList) UnmarshalJSON(data []byte) error {
 		return err
 	}
 	s.ips = temp.IPs
-
-	if temp.IPs == nil {
-		temp.IPs = make(map[string]struct{})
-	}
-
 	return nil
 }
 
@@ -94,10 +89,5 @@ func (s *ipsetStore) UnmarshalJSON(data []byte) error {
 		return err
 	}
 	s.ipsets = temp.IPSets
-
-	if temp.IPSets == nil {
-		temp.IPSets = make(map[string]*ipList)
-	}
-
 	return nil
 }

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
 }
@@ -199,7 +197,7 @@ func (m *Manager) AllowNetbird() error {
 
 	var chain *nftables.Chain
 	for _, c := range chains {
-		if c.Table.Name == tableNameFilter && c.Name == chainNameInput {
+		if c.Table.Name == tableNameFilter && c.Name == chainNameForward {
 			chain = c
 			break
 		}
@@ -276,7 +274,7 @@ func (m *Manager) resetNetbirdInputRules() error {
 
 func (m *Manager) deleteNetbirdInputRules(chains []*nftables.Chain) {
 	for _, c := range chains {
-		if c.Table.Name == tableNameFilter && c.Name == chainNameInput {
+		if c.Table.Name == "filter" && c.Name == "INPUT" {
 			rules, err := m.rConn.GetRules(c.Table, c)
 			if err != nil {
 				log.Errorf("get rules for chain %q: %v", c.Name, err)
@@ -351,9 +349,7 @@ func (m *Manager) applyAllowNetbirdRules(chain *nftables.Chain) {
 				Register: 1,
 				Data:     ifname(m.wgIface.Name()),
 			},
-			&expr.Verdict{
-				Kind: expr.VerdictAccept,
-			},
+			&expr.Verdict{},
 		},
 		UserData: []byte(allowNetbirdInputRuleID),
 	}

client/firewall/nftables/manager_linux_test.go

@@ -1,11 +1,9 @@
 package nftables
 
 import (
-	"bytes"
 	"fmt"
 	"net"
 	"net/netip"
-	"os/exec"
 	"testing"
 	"time"
 
@@ -227,105 +225,3 @@ func TestNFtablesCreatePerformance(t *testing.T) {
 		})
 	}
 }
-
-func runIptablesSave(t *testing.T) (string, string) {
-	t.Helper()
-	var stdout, stderr bytes.Buffer
-	cmd := exec.Command("iptables-save")
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	err := cmd.Run()
-	require.NoError(t, err, "iptables-save failed to run")
-
-	return stdout.String(), stderr.String()
-}
-
-func verifyIptablesOutput(t *testing.T, stdout, stderr string) {
-	t.Helper()
-	// Check for any incompatibility warnings
-	require.NotContains(t,
-		stderr,
-		"incompatible",
-		"iptables-save produced compatibility warning. Full stderr: %s",
-		stderr,
-	)
-
-	// Verify standard tables are present
-	expectedTables := []string{
-		"*filter",
-		"*nat",
-		"*mangle",
-	}
-
-	for _, table := range expectedTables {
-		require.Contains(t,
-			stdout,
-			table,
-			"iptables-save output missing expected table: %s\nFull stdout: %s",
-			table,
-			stdout,
-		)
-	}
-}
-
-func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
-	if check() != NFTABLES {
-		t.Skip("nftables not supported on this system")
-	}
-
-	if _, err := exec.LookPath("iptables-save"); err != nil {
-		t.Skipf("iptables-save not available on this system: %v", err)
-	}
-
-	// First ensure iptables-nft tables exist by running iptables-save
-	stdout, stderr := runIptablesSave(t)
-	verifyIptablesOutput(t, stdout, stderr)
-
-	manager, err := Create(ifaceMock)
-	require.NoError(t, err, "failed to create manager")
-	require.NoError(t, manager.Init(nil))
-
-	t.Cleanup(func() {
-		err := manager.Reset(nil)
-		require.NoError(t, err, "failed to reset manager state")
-
-		// Verify iptables output after reset
-		stdout, stderr := runIptablesSave(t)
-		verifyIptablesOutput(t, stdout, stderr)
-	})
-
-	ip := net.ParseIP("100.96.0.1")
-	_, err = manager.AddPeerFiltering(
-		ip,
-		fw.ProtocolTCP,
-		nil,
-		&fw.Port{Values: []int{80}},
-		fw.RuleDirectionIN,
-		fw.ActionAccept,
-		"",
-		"test rule",
-	)
-	require.NoError(t, err, "failed to add peer filtering rule")
-
-	_, err = manager.AddRouteFiltering(
-		[]netip.Prefix{netip.MustParsePrefix("192.168.2.0/24")},
-		netip.MustParsePrefix("10.1.0.0/24"),
-		fw.ProtocolTCP,
-		nil,
-		&fw.Port{Values: []int{443}},
-		fw.ActionAccept,
-	)
-	require.NoError(t, err, "failed to add route filtering rule")
-
-	pair := fw.RouterPair{
-		Source:      netip.MustParsePrefix("192.168.1.0/24"),
-		Destination: netip.MustParsePrefix("10.0.0.0/24"),
-		Masquerade:  true,
-	}
-	err = manager.AddNatRule(pair)
-	require.NoError(t, err, "failed to add NAT rule")
-
-	stdout, stderr = runIptablesSave(t)
-	verifyIptablesOutput(t, stdout, stderr)
-}

client/firewall/nftables/state.go

@@ -0,0 +1 @@
+package nftables

client/firewall/uspfilter/uspfilter.go

@@ -239,7 +239,7 @@ func (m *Manager) DeletePeerRule(rule firewall.Rule) error {
 // SetLegacyManagement doesn't need to be implemented for this manager
 func (m *Manager) SetLegacyManagement(isLegacy bool) error {
 	if m.nativeFirewall == nil {
-		return nil
+		return errRouteNotSupported
 	}
 	return m.nativeFirewall.SetLegacyManagement(isLegacy)
 }

client/iface/bind/control_android.go

@@ -1,12 +0,0 @@
-package bind
-
-import (
-	wireguard "golang.zx2c4.com/wireguard/conn"
-
-	nbnet "github.com/netbirdio/netbird/util/net"
-)
-
-func init() {
-	// ControlFns is not thread safe and should only be modified during init.
-	*wireguard.ControlFns = append(*wireguard.ControlFns, nbnet.ControlProtectSocket)
-}

client/iface/bind/udp_mux.go

@@ -162,13 +162,12 @@ func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
 		params.Logger.Warn("UDPMuxDefault should not listening on unspecified address, use NewMultiUDPMuxFromPort instead")
 		var networks []ice.NetworkType
 		switch {
+		case addr.IP.To4() != nil:
+			networks = []ice.NetworkType{ice.NetworkTypeUDP4}
 
 		case addr.IP.To16() != nil:
 			networks = []ice.NetworkType{ice.NetworkTypeUDP4, ice.NetworkTypeUDP6}
 
-		case addr.IP.To4() != nil:
-			networks = []ice.NetworkType{ice.NetworkTypeUDP4}
-
 		default:
 			params.Logger.Errorf("LocalAddr expected IPV4 or IPV6, got %T", params.UDPConn.LocalAddr())
 		}

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

@@ -40,8 +40,6 @@ type ConnectClient struct {
 	statusRecorder *peer.Status
 	engine         *Engine
 	engineMutex    sync.Mutex
-
-	persistNetworkMap bool
 }
 
 func NewConnectClient(
@@ -159,8 +157,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		engineCtx, cancel := context.WithCancel(c.ctx)
 		defer func() {
-			_, err := state.Status()
-			c.statusRecorder.MarkManagementDisconnected(err)
+			c.statusRecorder.MarkManagementDisconnected(state.err)
 			c.statusRecorder.CleanLocalPeerState()
 			cancel()
 		}()
@@ -234,7 +231,6 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		relayURLs, token := parseRelayInfo(loginResp)
 		relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String())
-		c.statusRecorder.SetRelayMgr(relayManager)
 		if len(relayURLs) > 0 {
 			if token != nil {
 				if err := relayManager.UpdateToken(token); err != nil {
@@ -245,7 +241,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()
@@ -260,7 +258,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 
 		c.engineMutex.Lock()
 		c.engine = NewEngineWithProbes(engineCtx, cancel, signalClient, mgmClient, relayManager, engineConfig, mobileDependency, c.statusRecorder, probes, checks)
-		c.engine.SetNetworkMapPersistence(c.persistNetworkMap)
+
 		c.engineMutex.Unlock()
 
 		if err := c.engine.Start(); err != nil {
@@ -338,19 +336,6 @@ func (c *ConnectClient) Engine() *Engine {
 	return e
 }
 
-// Status returns the current client status
-func (c *ConnectClient) Status() StatusType {
-	if c == nil {
-		return StatusIdle
-	}
-	status, err := CtxGetState(c.ctx).Status()
-	if err != nil {
-		return StatusIdle
-	}
-
-	return status
-}
-
 func (c *ConnectClient) Stop() error {
 	if c == nil {
 		return nil
@@ -377,22 +362,6 @@ func (c *ConnectClient) isContextCancelled() bool {
 	}
 }
 
-// SetNetworkMapPersistence enables or disables network map persistence.
-// When enabled, the last received network map will be stored and can be retrieved
-// through the Engine's getLatestNetworkMap method. When disabled, any stored
-// network map will be cleared. This functionality is primarily used for debugging
-// and should not be enabled during normal operation.
-func (c *ConnectClient) SetNetworkMapPersistence(enabled bool) {
-	c.engineMutex.Lock()
-	c.persistNetworkMap = enabled
-	c.engineMutex.Unlock()
-
-	engine := c.Engine()
-	if engine != nil {
-		engine.SetNetworkMapPersistence(enabled)
-	}
-}
-
 // createEngineConfig converts configuration received from Management Service to EngineConfig
 func createEngineConfig(key wgtypes.Key, config *Config, peerConfig *mgmProto.PeerConfig) (*EngineConfig, error) {
 	nm := false

client/internal/dns/server.go

@@ -7,6 +7,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/miekg/dns"
 	"github.com/mitchellh/hashstructure/v2"
@@ -322,12 +323,12 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 		log.Error(err)
 	}
 
-	go func() {
-		// persist dns state right away
-		if err := s.stateManager.PersistState(s.ctx); err != nil {
-			log.Errorf("Failed to persist dns state: %v", err)
-		}
-	}()
+	// persist dns state right away
+	ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
+	defer cancel()
+	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())
@@ -532,11 +533,12 @@ func (s *DefaultServer) upstreamCallbacks(
 			l.Errorf("Failed to apply nameserver deactivation on the host: %v", err)
 		}
 
-		go func() {
-			if err := s.stateManager.PersistState(s.ctx); err != nil {
-				l.Errorf("Failed to persist dns state: %v", err)
-			}
-		}()
+		// persist dns state right away
+		ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
+		defer cancel()
+		if err := s.stateManager.PersistState(ctx); err != nil {
+			l.Errorf("Failed to persist dns state: %v", err)
+		}
 
 		if runtime.GOOS == "android" && nsGroup.Primary && len(s.hostsDNSHolder.get()) > 0 {
 			s.addHostRootZone()

client/internal/dns/server_test.go

@@ -782,7 +782,7 @@ func TestDNSPermanent_matchOnly(t *testing.T) {
 						Port:   53,
 					},
 				},
-				Domains: []string{"google.com"},
+				Domains: []string{"customdomain.com"},
 				Primary: false,
 			},
 		},
@@ -804,7 +804,7 @@ func TestDNSPermanent_matchOnly(t *testing.T) {
 	if ips[0] != zoneRecords[0].RData {
 		t.Fatalf("invalid zone record: %v", err)
 	}
-	_, err = resolver.LookupHost(context.Background(), "google.com")
+	_, err = resolver.LookupHost(context.Background(), "customdomain.com")
 	if err != nil {
 		t.Errorf("failed to resolve: %s", err)
 	}

client/internal/engine.go

@@ -11,7 +11,6 @@ import (
 	"reflect"
 	"runtime"
 	"slices"
-	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -21,7 +20,6 @@ import (
 	"github.com/pion/stun/v2"
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"google.golang.org/protobuf/proto"
 
 	"github.com/netbirdio/netbird/client/firewall"
 	"github.com/netbirdio/netbird/client/firewall/manager"
@@ -40,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"
@@ -172,11 +171,7 @@ type Engine struct {
 
 	relayManager *relayClient.Manager
 	stateManager *statemanager.Manager
-	srWatcher    *guard.SRWatcher
-
-	// Network map persistence
-	persistNetworkMap bool
-	latestNetworkMap  *mgmProto.NetworkMap
+	srWatcher *guard.SRWatcher
 }
 
 // Peer is an instance of the Connection Peer
@@ -302,7 +297,7 @@ func (e *Engine) Stop() error {
 	if err := e.stateManager.Stop(ctx); err != nil {
 		return fmt.Errorf("failed to stop state manager: %w", err)
 	}
-	if err := e.stateManager.PersistState(context.Background()); err != nil {
+	if err := e.stateManager.PersistState(ctx); err != nil {
 		log.Errorf("failed to persist state: %v", err)
 	}
 
@@ -354,17 +349,8 @@ func (e *Engine) Start() error {
 	}
 	e.dnsServer = dnsServer
 
-	e.routeManager = routemanager.NewManager(
-		e.ctx,
-		e.config.WgPrivateKey.PublicKey().String(),
-		e.config.DNSRouteInterval,
-		e.wgInterface,
-		e.statusRecorder,
-		e.relayManager,
-		initialRoutes,
-		e.stateManager,
-	)
-	beforePeerHook, afterPeerHook, err := e.routeManager.Init()
+	e.routeManager = routemanager.NewManager(e.ctx, e.config.WgPrivateKey.PublicKey().String(), e.config.DNSRouteInterval, e.wgInterface, e.statusRecorder, e.relayManager, initialRoutes)
+	beforePeerHook, afterPeerHook, err := e.routeManager.Init(e.stateManager)
 	if err != nil {
 		log.Errorf("Failed to initialize route manager: %s", err)
 	} else {
@@ -552,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(),
@@ -561,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(relayMsg.Urls)
-
-			// 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
 	}
 
@@ -578,22 +556,13 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 		return err
 	}
 
-	nm := update.GetNetworkMap()
-	if nm == nil {
-		return nil
+	if update.GetNetworkMap() != nil {
+		// only apply new changes and ignore old ones
+		err := e.updateNetworkMap(update.GetNetworkMap())
+		if err != nil {
+			return err
+		}
 	}
-
-	// Store network map if persistence is enabled
-	if e.persistNetworkMap {
-		e.latestNetworkMap = nm
-		log.Debugf("network map persisted with serial %d", nm.GetSerial())
-	}
-
-	// only apply new changes and ignore old ones
-	if err := e.updateNetworkMap(nm); err != nil {
-		return err
-	}
-
 	return nil
 }
 
@@ -672,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)
@@ -1514,59 +1479,8 @@ func (e *Engine) stopDNSServer() {
 	e.statusRecorder.UpdateDNSStates(nsGroupStates)
 }
 
-// SetNetworkMapPersistence enables or disables network map persistence
-func (e *Engine) SetNetworkMapPersistence(enabled bool) {
-	e.syncMsgMux.Lock()
-	defer e.syncMsgMux.Unlock()
-
-	if enabled == e.persistNetworkMap {
-		return
-	}
-	e.persistNetworkMap = enabled
-	log.Debugf("Network map persistence is set to %t", enabled)
-
-	if !enabled {
-		e.latestNetworkMap = nil
-	}
-}
-
-// GetLatestNetworkMap returns the stored network map if persistence is enabled
-func (e *Engine) GetLatestNetworkMap() (*mgmProto.NetworkMap, error) {
-	e.syncMsgMux.Lock()
-	defer e.syncMsgMux.Unlock()
-
-	if !e.persistNetworkMap {
-		return nil, errors.New("network map persistence is disabled")
-	}
-
-	if e.latestNetworkMap == nil {
-		//nolint:nilnil
-		return nil, nil
-	}
-
-	// Create a deep copy to avoid external modifications
-	nm, ok := proto.Clone(e.latestNetworkMap).(*mgmProto.NetworkMap)
-	if !ok {
-
-		return nil, fmt.Errorf("failed to clone network map")
-	}
-
-	return nm, nil
-}
-
 // isChecksEqual checks if two slices of checks are equal.
 func isChecksEqual(checks []*mgmProto.Checks, oChecks []*mgmProto.Checks) bool {
-	for _, check := range checks {
-		sort.Slice(check.Files, func(i, j int) bool {
-			return check.Files[i] < check.Files[j]
-		})
-	}
-	for _, oCheck := range oChecks {
-		sort.Slice(oCheck.Files, func(i, j int) bool {
-			return oCheck.Files[i] < oCheck.Files[j]
-		})
-	}
-
 	return slices.EqualFunc(checks, oChecks, func(checks, oChecks *mgmProto.Checks) bool {
 		return slices.Equal(checks.Files, oChecks.Files)
 	})

client/internal/engine_test.go

@@ -245,15 +245,12 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 		nil)
 
 	wgIface := &iface.MockWGIface{
-		NameFunc: func() string { return "utun102" },
 		RemovePeerFunc: func(peerKey string) error {
 			return nil
 		},
 	}
 	engine.wgInterface = wgIface
-	engine.routeManager = routemanager.NewManager(ctx, key.PublicKey().String(), time.Minute, engine.wgInterface, engine.statusRecorder, relayMgr, nil, nil)
-	_, _, err = engine.routeManager.Init()
-	require.NoError(t, err)
+	engine.routeManager = routemanager.NewManager(ctx, key.PublicKey().String(), time.Minute, engine.wgInterface, engine.statusRecorder, relayMgr, nil)
 	engine.dnsServer = &dns.MockServer{
 		UpdateDNSServerFunc: func(serial uint64, update nbdns.Config) error { return nil },
 	}
@@ -1009,99 +1006,6 @@ func Test_ParseNATExternalIPMappings(t *testing.T) {
 	}
 }
 
-func Test_CheckFilesEqual(t *testing.T) {
-	testCases := []struct {
-		name         string
-		inputChecks1 []*mgmtProto.Checks
-		inputChecks2 []*mgmtProto.Checks
-		expectedBool bool
-	}{
-		{
-			name: "Equal Files In Equal Order Should Return True",
-			inputChecks1: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile2",
-					},
-				},
-			},
-			inputChecks2: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile2",
-					},
-				},
-			},
-			expectedBool: true,
-		},
-		{
-			name: "Equal Files In Reverse Order Should Return True",
-			inputChecks1: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile2",
-					},
-				},
-			},
-			inputChecks2: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile2",
-						"testfile1",
-					},
-				},
-			},
-			expectedBool: true,
-		},
-		{
-			name: "Unequal Files Should Return False",
-			inputChecks1: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile2",
-					},
-				},
-			},
-			inputChecks2: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile3",
-					},
-				},
-			},
-			expectedBool: false,
-		},
-		{
-			name: "Compared With Empty Should Return False",
-			inputChecks1: []*mgmtProto.Checks{
-				{
-					Files: []string{
-						"testfile1",
-						"testfile2",
-					},
-				},
-			},
-			inputChecks2: []*mgmtProto.Checks{
-				{
-					Files: []string{},
-				},
-			},
-			expectedBool: false,
-		},
-	}
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			result := isChecksEqual(testCase.inputChecks1, testCase.inputChecks2)
-			assert.Equal(t, testCase.expectedBool, result, "result should match expected bool")
-		})
-	}
-}
-
 func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey string, i int, mgmtAddr string, signalAddr string) (*Engine, error) {
 	key, err := wgtypes.GeneratePrivateKey()
 	if err != nil {

client/internal/peer/conn.go

@@ -83,6 +83,7 @@ type Conn struct {
 	signaler       *Signaler
 	relayManager   *relayClient.Manager
 	allowedIP      net.IP
+	allowedNet     string
 	handshaker     *Handshaker
 
 	onConnected    func(remoteWireGuardKey string, remoteRosenpassPubKey []byte, wireGuardIP string, remoteRosenpassAddr string)
@@ -110,7 +111,7 @@ type Conn struct {
 // NewConn creates a new not opened Conn to the remote peer.
 // To establish a connection run Conn.Open
 func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Status, signaler *Signaler, iFaceDiscover stdnet.ExternalIFaceDiscover, relayManager *relayClient.Manager, srWatcher *guard.SRWatcher) (*Conn, error) {
-	allowedIP, _, err := net.ParseCIDR(config.WgConfig.AllowedIps)
+	allowedIP, allowedNet, err := net.ParseCIDR(config.WgConfig.AllowedIps)
 	if err != nil {
 		log.Errorf("failed to parse allowedIPS: %v", err)
 		return nil, err
@@ -128,6 +129,7 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 		signaler:       signaler,
 		relayManager:   relayManager,
 		allowedIP:      allowedIP,
+		allowedNet:     allowedNet.String(),
 		statusRelay:    NewAtomicConnStatus(),
 		statusICE:      NewAtomicConnStatus(),
 	}
@@ -592,7 +594,7 @@ func (conn *Conn) doOnConnected(remoteRosenpassPubKey []byte, remoteRosenpassAdd
 	}
 
 	if conn.onConnected != nil {
-		conn.onConnected(conn.config.Key, remoteRosenpassPubKey, conn.allowedIP.String(), remoteRosenpassAddr)
+		conn.onConnected(conn.config.Key, remoteRosenpassPubKey, conn.allowedNet, remoteRosenpassAddr)
 	}
 }
 

client/internal/peer/status.go

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

client/internal/peer/worker_ice.go

@@ -46,6 +46,8 @@ type WorkerICE struct {
 	hasRelayOnLocally bool
 	conn              WorkerICECallbacks
 
+	selectedPriority ConnPriority
+
 	agent    *ice.Agent
 	muxAgent sync.Mutex
 
@@ -93,8 +95,10 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 
 	var preferredCandidateTypes []ice.CandidateType
 	if w.hasRelayOnLocally && remoteOfferAnswer.RelaySrvAddress != "" {
+		w.selectedPriority = connPriorityICEP2P
 		preferredCandidateTypes = icemaker.CandidateTypesP2P()
 	} else {
+		w.selectedPriority = connPriorityICETurn
 		preferredCandidateTypes = icemaker.CandidateTypes()
 	}
 
@@ -155,7 +159,7 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
 	}
 	w.log.Debugf("on ICE conn read to use ready")
-	go w.conn.OnConnReady(selectedPriority(pair), ci)
+	go w.conn.OnConnReady(w.selectedPriority, ci)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -390,11 +394,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/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/manager.go

@@ -32,7 +32,7 @@ import (
 
 // Manager is a route manager interface
 type Manager interface {
-	Init() (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error)
+	Init(*statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error)
 	UpdateRoutes(updateSerial uint64, newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap, error)
 	TriggerSelection(route.HAMap)
 	GetRouteSelector() *routeselector.RouteSelector
@@ -59,7 +59,6 @@ type DefaultManager struct {
 	routeRefCounter      *refcounter.RouteRefCounter
 	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter
 	dnsRouteInterval     time.Duration
-	stateManager         *statemanager.Manager
 }
 
 func NewManager(
@@ -70,7 +69,6 @@ func NewManager(
 	statusRecorder *peer.Status,
 	relayMgr *relayClient.Manager,
 	initialRoutes []*route.Route,
-	stateManager *statemanager.Manager,
 ) *DefaultManager {
 	mCTX, cancel := context.WithCancel(ctx)
 	notifier := notifier.NewNotifier()
@@ -82,12 +80,12 @@ func NewManager(
 		dnsRouteInterval: dnsRouteInterval,
 		clientNetworks:   make(map[route.HAUniqueID]*clientNetwork),
 		relayMgr:         relayMgr,
+		routeSelector:    routeselector.NewRouteSelector(),
 		sysOps:           sysOps,
 		statusRecorder:   statusRecorder,
 		wgInterface:      wgInterface,
 		pubKey:           pubKey,
 		notifier:         notifier,
-		stateManager:     stateManager,
 	}
 
 	dm.routeRefCounter = refcounter.New(
@@ -123,7 +121,7 @@ func NewManager(
 }
 
 // Init sets up the routing
-func (m *DefaultManager) Init() (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (m *DefaultManager) Init(stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	if nbnet.CustomRoutingDisabled() {
 		return nil, nil, nil
 	}
@@ -139,38 +137,14 @@ func (m *DefaultManager) Init() (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error)
 
 	ips := resolveURLsToIPs(initialAddresses)
 
-	beforePeerHook, afterPeerHook, err := m.sysOps.SetupRouting(ips, m.stateManager)
+	beforePeerHook, afterPeerHook, err := m.sysOps.SetupRouting(ips, stateManager)
 	if err != nil {
 		return nil, nil, fmt.Errorf("setup routing: %w", err)
 	}
-
-	m.routeSelector = m.initSelector()
-
 	log.Info("Routing setup complete")
 	return beforePeerHook, afterPeerHook, nil
 }
 
-func (m *DefaultManager) initSelector() *routeselector.RouteSelector {
-	var state *SelectorState
-	m.stateManager.RegisterState(state)
-
-	// restore selector state if it exists
-	if err := m.stateManager.LoadState(state); err != nil {
-		log.Warnf("failed to load state: %v", err)
-		return routeselector.NewRouteSelector()
-	}
-
-	if state := m.stateManager.GetState(state); state != nil {
-		if selector, ok := state.(*SelectorState); ok {
-			return (*routeselector.RouteSelector)(selector)
-		}
-
-		log.Warnf("failed to convert state with type %T to SelectorState", state)
-	}
-
-	return routeselector.NewRouteSelector()
-}
-
 func (m *DefaultManager) EnableServerRouter(firewall firewall.Manager) error {
 	var err error
 	m.serverRouter, err = newServerRouter(m.ctx, m.wgInterface, firewall, m.statusRecorder)
@@ -278,10 +252,6 @@ func (m *DefaultManager) TriggerSelection(networks route.HAMap) {
 		go clientNetworkWatcher.peersStateAndUpdateWatcher()
 		clientNetworkWatcher.sendUpdateToClientNetworkWatcher(routesUpdate{routes: routes})
 	}
-
-	if err := m.stateManager.UpdateState((*SelectorState)(m.routeSelector)); err != nil {
-		log.Errorf("failed to update state: %v", err)
-	}
 }
 
 // stopObsoleteClients stops the client network watcher for the networks that are not in the new list

client/internal/routemanager/manager_test.go

@@ -424,9 +424,9 @@ func TestManagerUpdateRoutes(t *testing.T) {
 
 			statusRecorder := peer.NewRecorder("https://mgm")
 			ctx := context.TODO()
-			routeManager := NewManager(ctx, localPeerKey, 0, wgInterface, statusRecorder, nil, nil, nil)
+			routeManager := NewManager(ctx, localPeerKey, 0, wgInterface, statusRecorder, nil, nil)
 
-			_, _, err = routeManager.Init()
+			_, _, err = routeManager.Init(nil)
 
 			require.NoError(t, err, "should init route manager")
 			defer routeManager.Stop(nil)

client/internal/routemanager/mock.go

@@ -21,7 +21,7 @@ type MockManager struct {
 	StopFunc             func(manager *statemanager.Manager)
 }
 
-func (m *MockManager) Init() (net.AddHookFunc, net.RemoveHookFunc, error) {
+func (m *MockManager) Init(*statemanager.Manager) (net.AddHookFunc, net.RemoveHookFunc, error) {
 	return nil, nil, nil
 }
 

client/internal/routemanager/refcounter/refcounter.go

@@ -47,9 +47,10 @@ type RemoveFunc[Key, O any] func(key Key, out O) error
 type Counter[Key comparable, I, O any] struct {
 	// refCountMap keeps track of the reference Ref for keys
 	refCountMap map[Key]Ref[O]
-	mu          sync.Mutex
+	refCountMu  sync.Mutex
 	// idMap keeps track of the keys associated with an ID for removal
 	idMap  map[string][]Key
+	idMu   sync.Mutex
 	add    AddFunc[Key, I, O]
 	remove RemoveFunc[Key, O]
 }
@@ -71,14 +72,13 @@ func New[Key comparable, I, O any](add AddFunc[Key, I, O], remove RemoveFunc[Key
 }
 
 // LoadData loads the data from the existing counter
-// The passed counter should not be used any longer after calling this function.
 func (rm *Counter[Key, I, O]) LoadData(
 	existingCounter *Counter[Key, I, O],
 ) {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
-	existingCounter.mu.Lock()
-	defer existingCounter.mu.Unlock()
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
+	rm.idMu.Lock()
+	defer rm.idMu.Unlock()
 
 	rm.refCountMap = existingCounter.refCountMap
 	rm.idMap = existingCounter.idMap
@@ -87,8 +87,8 @@ func (rm *Counter[Key, I, O]) LoadData(
 // Get retrieves the current reference count and associated data for a key.
 // If the key doesn't exist, it returns a zero value Ref and false.
 func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
 
 	ref, ok := rm.refCountMap[key]
 	return ref, ok
@@ -97,13 +97,9 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
 // Increment increments the reference count for the given key.
 // If this is the first reference to the key, the AddFunc is called.
 func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
 
-	return rm.increment(key, in)
-}
-
-func (rm *Counter[Key, I, O]) increment(key Key, in I) (Ref[O], error) {
 	ref := rm.refCountMap[key]
 	logCallerF("Increasing ref count [%d -> %d] for key %v with In [%v] Out [%v]", ref.Count, ref.Count+1, key, in, ref.Out)
 
@@ -130,10 +126,10 @@ func (rm *Counter[Key, I, O]) increment(key Key, in I) (Ref[O], error) {
 // IncrementWithID increments the reference count for the given key and groups it under the given ID.
 // If this is the first reference to the key, the AddFunc is called.
 func (rm *Counter[Key, I, O]) IncrementWithID(id string, key Key, in I) (Ref[O], error) {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.idMu.Lock()
+	defer rm.idMu.Unlock()
 
-	ref, err := rm.increment(key, in)
+	ref, err := rm.Increment(key, in)
 	if err != nil {
 		return ref, fmt.Errorf("with ID: %w", err)
 	}
@@ -145,12 +141,9 @@ func (rm *Counter[Key, I, O]) IncrementWithID(id string, key Key, in I) (Ref[O],
 // Decrement decrements the reference count for the given key.
 // If the reference count reaches 0, the RemoveFunc is called.
 func (rm *Counter[Key, I, O]) Decrement(key Key) (Ref[O], error) {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
-	return rm.decrement(key)
-}
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
 
-func (rm *Counter[Key, I, O]) decrement(key Key) (Ref[O], error) {
 	ref, ok := rm.refCountMap[key]
 	if !ok {
 		logCallerF("No reference found for key %v", key)
@@ -175,12 +168,12 @@ func (rm *Counter[Key, I, O]) decrement(key Key) (Ref[O], error) {
 // DecrementWithID decrements the reference count for all keys associated with the given ID.
 // If the reference count reaches 0, the RemoveFunc is called.
 func (rm *Counter[Key, I, O]) DecrementWithID(id string) error {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.idMu.Lock()
+	defer rm.idMu.Unlock()
 
 	var merr *multierror.Error
 	for _, key := range rm.idMap[id] {
-		if _, err := rm.decrement(key); err != nil {
+		if _, err := rm.Decrement(key); err != nil {
 			merr = multierror.Append(merr, err)
 		}
 	}
@@ -191,8 +184,10 @@ func (rm *Counter[Key, I, O]) DecrementWithID(id string) error {
 
 // Flush removes all references and calls RemoveFunc for each key.
 func (rm *Counter[Key, I, O]) Flush() error {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
+	rm.idMu.Lock()
+	defer rm.idMu.Unlock()
 
 	var merr *multierror.Error
 	for key := range rm.refCountMap {
@@ -211,8 +206,10 @@ func (rm *Counter[Key, I, O]) Flush() error {
 
 // Clear removes all references without calling RemoveFunc.
 func (rm *Counter[Key, I, O]) Clear() {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
+	rm.refCountMu.Lock()
+	defer rm.refCountMu.Unlock()
+	rm.idMu.Lock()
+	defer rm.idMu.Unlock()
 
 	clear(rm.refCountMap)
 	clear(rm.idMap)
@@ -220,8 +217,10 @@ 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.mu.Lock()
-	defer rm.mu.Unlock()
+	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"`
@@ -234,9 +233,6 @@ func (rm *Counter[Key, I, O]) MarshalJSON() ([]byte, error) {
 
 // UnmarshalJSON implements the json.Unmarshaler interface for Counter.
 func (rm *Counter[Key, I, O]) UnmarshalJSON(data []byte) error {
-	rm.mu.Lock()
-	defer rm.mu.Unlock()
-
 	var temp struct {
 		RefCountMap map[Key]Ref[O]   `json:"refCountMap"`
 		IDMap       map[string][]Key `json:"idMap"`
@@ -247,13 +243,6 @@ func (rm *Counter[Key, I, O]) UnmarshalJSON(data []byte) error {
 	rm.refCountMap = temp.RefCountMap
 	rm.idMap = temp.IDMap
 
-	if temp.RefCountMap == nil {
-		temp.RefCountMap = map[Key]Ref[O]{}
-	}
-	if temp.IDMap == nil {
-		temp.IDMap = map[string][]Key{}
-	}
-
 	return nil
 }
 

client/internal/routemanager/state.go

@@ -1,19 +0,0 @@
-package routemanager
-
-import (
-	"github.com/netbirdio/netbird/client/internal/routeselector"
-)
-
-type SelectorState routeselector.RouteSelector
-
-func (s *SelectorState) Name() string {
-	return "routeselector_state"
-}
-
-func (s *SelectorState) MarshalJSON() ([]byte, error) {
-	return (*routeselector.RouteSelector)(s).MarshalJSON()
-}
-
-func (s *SelectorState) UnmarshalJSON(data []byte) error {
-	return (*routeselector.RouteSelector)(s).UnmarshalJSON(data)
-}

client/internal/routemanager/systemops/state.go

@@ -2,28 +2,31 @@ package systemops
 
 import (
 	"net/netip"
+	"sync"
 
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 )
 
-type ShutdownState ExclusionCounter
+type ShutdownState struct {
+	Counter *ExclusionCounter `json:"counter,omitempty"`
+	mu      sync.RWMutex
+}
 
 func (s *ShutdownState) Name() string {
 	return "route_state"
 }
 
 func (s *ShutdownState) Cleanup() error {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.Counter == nil {
+		return nil
+	}
+
 	sysops := NewSysOps(nil, nil)
 	sysops.refCounter = refcounter.New[netip.Prefix, struct{}, Nexthop](nil, sysops.removeFromRouteTable)
-	sysops.refCounter.LoadData((*ExclusionCounter)(s))
+	sysops.refCounter.LoadData(s.Counter)
 
 	return sysops.refCounter.Flush()
 }
-
-func (s *ShutdownState) MarshalJSON() ([]byte, error) {
-	return (*ExclusionCounter)(s).MarshalJSON()
-}
-
-func (s *ShutdownState) UnmarshalJSON(data []byte) error {
-	return (*ExclusionCounter)(s).UnmarshalJSON(data)
-}

client/internal/routemanager/systemops/systemops_generic.go

@@ -57,19 +57,30 @@ func (r *SysOps) setupRefCounter(initAddresses []net.IP, stateManager *statemana
 				return nexthop, refcounter.ErrIgnore
 			}
 
+			r.updateState(stateManager)
+
 			return nexthop, err
 		},
-		r.removeFromRouteTable,
+		func(prefix netip.Prefix, nexthop Nexthop) error {
+			// remove from state even if we have trouble removing it from the route table
+			// it could be already gone
+			r.updateState(stateManager)
+
+			return r.removeFromRouteTable(prefix, nexthop)
+		},
 	)
 
 	r.refCounter = refCounter
 
-	return r.setupHooks(initAddresses, stateManager)
+	return r.setupHooks(initAddresses)
 }
 
-// updateState updates state on every change so it will be persisted regularly
 func (r *SysOps) updateState(stateManager *statemanager.Manager) {
-	if err := stateManager.UpdateState((*ShutdownState)(r.refCounter)); err != nil {
+	state := getState(stateManager)
+
+	state.Counter = r.refCounter
+
+	if err := stateManager.UpdateState(state); err != nil {
 		log.Errorf("failed to update state: %v", err)
 	}
 }
@@ -325,7 +336,7 @@ func (r *SysOps) genericRemoveVPNRoute(prefix netip.Prefix, intf *net.Interface)
 	return r.removeFromRouteTable(prefix, nextHop)
 }
 
-func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) setupHooks(initAddresses []net.IP) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	beforeHook := func(connID nbnet.ConnectionID, ip net.IP) error {
 		prefix, err := util.GetPrefixFromIP(ip)
 		if err != nil {
@@ -336,8 +347,6 @@ func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager *statemanager.M
 			return fmt.Errorf("adding route reference: %v", err)
 		}
 
-		r.updateState(stateManager)
-
 		return nil
 	}
 	afterHook := func(connID nbnet.ConnectionID) error {
@@ -345,8 +354,6 @@ func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager *statemanager.M
 			return fmt.Errorf("remove route reference: %w", err)
 		}
 
-		r.updateState(stateManager)
-
 		return nil
 	}
 
@@ -525,3 +532,14 @@ func isVpnRoute(addr netip.Addr, vpnRoutes []netip.Prefix, localRoutes []netip.P
 	// Return true if the longest matching prefix is from vpnRoutes
 	return isVpn, longestPrefix
 }
+
+func getState(stateManager *statemanager.Manager) *ShutdownState {
+	var shutdownState *ShutdownState
+	if state := stateManager.GetState(shutdownState); state != nil {
+		shutdownState = state.(*ShutdownState)
+	} else {
+		shutdownState = &ShutdownState{}
+	}
+
+	return shutdownState
+}

client/internal/routemanager/systemops/systemops_linux.go

@@ -55,7 +55,7 @@ type ruleParams struct {
 
 // isLegacy determines whether to use the legacy routing setup
 func isLegacy() bool {
-	return os.Getenv("NB_USE_LEGACY_ROUTING") == "true" || nbnet.CustomRoutingDisabled() || nbnet.SkipSocketMark()
+	return os.Getenv("NB_USE_LEGACY_ROUTING") == "true" || nbnet.CustomRoutingDisabled()
 }
 
 // setIsLegacy sets the legacy routing setup
@@ -92,6 +92,17 @@ func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager
 		return r.setupRefCounter(initAddresses, stateManager)
 	}
 
+	if err = addRoutingTableName(); err != nil {
+		log.Errorf("Error adding routing table name: %v", err)
+	}
+
+	originalValues, err := sysctl.Setup(r.wgInterface)
+	if err != nil {
+		log.Errorf("Error setting up sysctl: %v", err)
+		sysctlFailed = true
+	}
+	originalSysctl = originalValues
+
 	defer func() {
 		if err != nil {
 			if cleanErr := r.CleanupRouting(stateManager); cleanErr != nil {
@@ -112,17 +123,6 @@ func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager
 		}
 	}
 
-	if err = addRoutingTableName(); err != nil {
-		log.Errorf("Error adding routing table name: %v", err)
-	}
-
-	originalValues, err := sysctl.Setup(r.wgInterface)
-	if err != nil {
-		log.Errorf("Error setting up sysctl: %v", err)
-		sysctlFailed = true
-	}
-	originalSysctl = originalValues
-
 	return nil, nil, nil
 }
 
@@ -450,7 +450,7 @@ func addRule(params ruleParams) error {
 	rule.Invert = params.invert
 	rule.SuppressPrefixlen = params.suppressPrefix
 
-	if err := netlink.RuleAdd(rule); err != nil && !errors.Is(err, syscall.EEXIST) {
+	if err := netlink.RuleAdd(rule); err != nil && !errors.Is(err, syscall.EEXIST) && !errors.Is(err, syscall.EAFNOSUPPORT) {
 		return fmt.Errorf("add routing rule: %w", err)
 	}
 
@@ -467,7 +467,7 @@ func removeRule(params ruleParams) error {
 	rule.Priority = params.priority
 	rule.SuppressPrefixlen = params.suppressPrefix
 
-	if err := netlink.RuleDel(rule); err != nil && !errors.Is(err, syscall.ENOENT) {
+	if err := netlink.RuleDel(rule); err != nil && !errors.Is(err, syscall.ENOENT) && !errors.Is(err, syscall.EAFNOSUPPORT) {
 		return fmt.Errorf("remove routing rule: %w", err)
 	}
 

client/internal/routemanager/systemops/systemops_windows.go

@@ -230,13 +230,10 @@ func (rm *RouteMonitor) parseUpdate(row *MIB_IPFORWARD_ROW2, notificationType MI
 	if idx != 0 {
 		intf, err := net.InterfaceByIndex(idx)
 		if err != nil {
-			log.Warnf("failed to get interface name for index %d: %v", idx, err)
-			update.Interface = &net.Interface{
-				Index: idx,
-			}
-		} else {
-			update.Interface = intf
+			return update, fmt.Errorf("get interface name: %w", err)
 		}
+
+		update.Interface = intf
 	}
 
 	log.Tracef("Received route update with destination %v, next hop %v, interface %v", row.DestinationPrefix, row.NextHop, update.Interface)

client/internal/routeselector/routeselector.go

@@ -1,10 +1,8 @@
 package routeselector
 
 import (
-	"encoding/json"
 	"fmt"
 	"slices"
-	"sync"
 
 	"github.com/hashicorp/go-multierror"
 	"golang.org/x/exp/maps"
@@ -14,7 +12,6 @@ import (
 )
 
 type RouteSelector struct {
-	mu             sync.RWMutex
 	selectedRoutes map[route.NetID]struct{}
 	selectAll      bool
 }
@@ -29,9 +26,6 @@ func NewRouteSelector() *RouteSelector {
 
 // SelectRoutes updates the selected routes based on the provided route IDs.
 func (rs *RouteSelector) SelectRoutes(routes []route.NetID, appendRoute bool, allRoutes []route.NetID) error {
-	rs.mu.Lock()
-	defer rs.mu.Unlock()
-
 	if !appendRoute {
 		rs.selectedRoutes = map[route.NetID]struct{}{}
 	}
@@ -52,9 +46,6 @@ func (rs *RouteSelector) SelectRoutes(routes []route.NetID, appendRoute bool, al
 
 // SelectAllRoutes sets the selector to select all routes.
 func (rs *RouteSelector) SelectAllRoutes() {
-	rs.mu.Lock()
-	defer rs.mu.Unlock()
-
 	rs.selectAll = true
 	rs.selectedRoutes = map[route.NetID]struct{}{}
 }
@@ -62,9 +53,6 @@ func (rs *RouteSelector) SelectAllRoutes() {
 // DeselectRoutes removes specific routes from the selection.
 // If the selector is in "select all" mode, it will transition to "select specific" mode.
 func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.NetID) error {
-	rs.mu.Lock()
-	defer rs.mu.Unlock()
-
 	if rs.selectAll {
 		rs.selectAll = false
 		rs.selectedRoutes = map[route.NetID]struct{}{}
@@ -88,18 +76,12 @@ func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.
 
 // DeselectAllRoutes deselects all routes, effectively disabling route selection.
 func (rs *RouteSelector) DeselectAllRoutes() {
-	rs.mu.Lock()
-	defer rs.mu.Unlock()
-
 	rs.selectAll = false
 	rs.selectedRoutes = map[route.NetID]struct{}{}
 }
 
 // IsSelected checks if a specific route is selected.
 func (rs *RouteSelector) IsSelected(routeID route.NetID) bool {
-	rs.mu.RLock()
-	defer rs.mu.RUnlock()
-
 	if rs.selectAll {
 		return true
 	}
@@ -109,9 +91,6 @@ func (rs *RouteSelector) IsSelected(routeID route.NetID) bool {
 
 // FilterSelected removes unselected routes from the provided map.
 func (rs *RouteSelector) FilterSelected(routes route.HAMap) route.HAMap {
-	rs.mu.RLock()
-	defer rs.mu.RUnlock()
-
 	if rs.selectAll {
 		return maps.Clone(routes)
 	}
@@ -124,49 +103,3 @@ func (rs *RouteSelector) FilterSelected(routes route.HAMap) route.HAMap {
 	}
 	return filtered
 }
-
-// MarshalJSON implements the json.Marshaler interface
-func (rs *RouteSelector) MarshalJSON() ([]byte, error) {
-	rs.mu.RLock()
-	defer rs.mu.RUnlock()
-
-	return json.Marshal(struct {
-		SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
-		SelectAll      bool                     `json:"select_all"`
-	}{
-		SelectAll:      rs.selectAll,
-		SelectedRoutes: rs.selectedRoutes,
-	})
-}
-
-// UnmarshalJSON implements the json.Unmarshaler interface
-// If the JSON is empty or null, it will initialize like a NewRouteSelector.
-func (rs *RouteSelector) UnmarshalJSON(data []byte) error {
-	rs.mu.Lock()
-	defer rs.mu.Unlock()
-
-	// Check for null or empty JSON
-	if len(data) == 0 || string(data) == "null" {
-		rs.selectedRoutes = map[route.NetID]struct{}{}
-		rs.selectAll = true
-		return nil
-	}
-
-	var temp struct {
-		SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
-		SelectAll      bool                     `json:"select_all"`
-	}
-
-	if err := json.Unmarshal(data, &temp); err != nil {
-		return err
-	}
-
-	rs.selectedRoutes = temp.SelectedRoutes
-	rs.selectAll = temp.SelectAll
-
-	if rs.selectedRoutes == nil {
-		rs.selectedRoutes = map[route.NetID]struct{}{}
-	}
-
-	return nil
-}

client/internal/statemanager/manager.go

@@ -16,39 +16,14 @@ import (
 	"golang.org/x/exp/maps"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
-	"github.com/netbirdio/netbird/util"
-)
-
-const (
-	errStateNotRegistered = "state %s not registered"
-	errLoadStateFile      = "load state file: %w"
 )
 
 // State interface defines the methods that all state types must implement
 type State interface {
 	Name() string
-}
-
-// CleanableState interface extends State with cleanup capability
-type CleanableState interface {
-	State
 	Cleanup() error
 }
 
-// RawState wraps raw JSON data for unregistered states
-type RawState struct {
-	data json.RawMessage
-}
-
-func (r *RawState) Name() string {
-	return "" // This is a placeholder implementation
-}
-
-// MarshalJSON implements json.Marshaler to preserve the original JSON
-func (r *RawState) MarshalJSON() ([]byte, error) {
-	return r.data, nil
-}
-
 // Manager handles the persistence and management of various states
 type Manager struct {
 	mu     sync.Mutex
@@ -98,15 +73,15 @@ func (m *Manager) Stop(ctx context.Context) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if m.cancel == nil {
-		return nil
-	}
-	m.cancel()
+	if m.cancel != nil {
+		m.cancel()
 
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-m.done:
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-m.done:
+			return nil
+		}
 	}
 
 	return nil
@@ -164,7 +139,7 @@ func (m *Manager) setState(name string, state State) error {
 	defer m.mu.Unlock()
 
 	if _, exists := m.states[name]; !exists {
-		return fmt.Errorf(errStateNotRegistered, name)
+		return fmt.Errorf("state %s not registered", name)
 	}
 
 	m.states[name] = state
@@ -173,63 +148,6 @@ func (m *Manager) setState(name string, state State) error {
 	return nil
 }
 
-// DeleteStateByName handles deletion of states without cleanup.
-// It doesn't require the state to be registered.
-func (m *Manager) DeleteStateByName(stateName string) error {
-	if m == nil {
-		return nil
-	}
-
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	rawStates, err := m.loadStateFile(false)
-	if err != nil {
-		return fmt.Errorf(errLoadStateFile, err)
-	}
-	if rawStates == nil {
-		return nil
-	}
-
-	if _, exists := rawStates[stateName]; !exists {
-		return fmt.Errorf("state %s not found", stateName)
-	}
-
-	// Mark state as deleted by setting it to nil and marking it dirty
-	m.states[stateName] = nil
-	m.dirty[stateName] = struct{}{}
-
-	return nil
-}
-
-// DeleteAllStates removes all states.
-func (m *Manager) DeleteAllStates() (int, error) {
-	if m == nil {
-		return 0, nil
-	}
-
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	rawStates, err := m.loadStateFile(false)
-	if err != nil {
-		return 0, fmt.Errorf(errLoadStateFile, err)
-	}
-	if rawStates == nil {
-		return 0, nil
-	}
-
-	count := len(rawStates)
-
-	// Mark all states as deleted and dirty
-	for name := range rawStates {
-		m.states[name] = nil
-		m.dirty[name] = struct{}{}
-	}
-
-	return count, nil
-}
-
 func (m *Manager) periodicStateSave(ctx context.Context) {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
@@ -260,18 +178,25 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		return nil
 	}
 
-	bs, err := marshalWithPanicRecovery(m.states)
-	if err != nil {
-		return fmt.Errorf("marshal states: %w", err)
-	}
-
-	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.WriteBytesWithRestrictedPermission(ctx, m.filePath, bs)
+		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 {
@@ -283,175 +208,63 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		}
 	}
 
-	log.Debugf("persisted states: %v, took %v", maps.Keys(m.dirty), time.Since(start))
+	log.Debugf("persisted shutdown states: %v", maps.Keys(m.dirty))
 
 	clear(m.dirty)
 
 	return nil
 }
 
-// loadStateFile reads and unmarshals the state file into a map of raw JSON messages
-func (m *Manager) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage, error) {
+// loadState loads the existing state from the state file
+func (m *Manager) loadState() error {
 	data, err := os.ReadFile(m.filePath)
 	if err != nil {
 		if errors.Is(err, fs.ErrNotExist) {
 			log.Debug("state file does not exist")
-			return nil, nil // nolint:nilnil
+			return nil
 		}
-		return nil, fmt.Errorf("read state file: %w", err)
+		return fmt.Errorf("read state file: %w", err)
 	}
 
 	var rawStates map[string]json.RawMessage
 	if err := json.Unmarshal(data, &rawStates); err != nil {
-		if deleteCorrupt {
-			log.Warn("State file appears to be corrupted, attempting to delete it", err)
-			if err := os.Remove(m.filePath); err != nil {
-				log.Errorf("Failed to delete corrupted state file: %v", err)
-			} else {
-				log.Info("State file deleted")
-			}
+		log.Warn("State file appears to be corrupted, attempting to delete it")
+		if err := os.Remove(m.filePath); err != nil {
+			log.Errorf("Failed to delete corrupted state file: %v", err)
+		} else {
+			log.Info("State file deleted")
 		}
-		return nil, fmt.Errorf("unmarshal states: %w", err)
+		return fmt.Errorf("unmarshal states: %w", err)
 	}
 
-	return rawStates, nil
-}
+	var merr *multierror.Error
 
-// loadSingleRawState unmarshals a raw state into a concrete state object
-func (m *Manager) loadSingleRawState(name string, rawState json.RawMessage) (State, error) {
-	stateType, ok := m.stateTypes[name]
-	if !ok {
-		return nil, fmt.Errorf(errStateNotRegistered, name)
-	}
+	for name, rawState := range rawStates {
+		stateType, ok := m.stateTypes[name]
+		if !ok {
+			merr = multierror.Append(merr, fmt.Errorf("unknown state type: %s", name))
+			continue
+		}
 
-	if string(rawState) == "null" {
-		return nil, nil //nolint:nilnil
-	}
+		if string(rawState) == "null" {
+			continue
+		}
 
-	statePtr := reflect.New(stateType).Interface().(State)
-	if err := json.Unmarshal(rawState, statePtr); err != nil {
-		return nil, fmt.Errorf("unmarshal state %s: %w", name, err)
-	}
+		statePtr := reflect.New(stateType).Interface().(State)
+		if err := json.Unmarshal(rawState, statePtr); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("unmarshal state %s: %w", name, err))
+			continue
+		}
 
-	return statePtr, nil
-}
-
-// LoadState loads a specific state from the state file
-func (m *Manager) LoadState(state State) error {
-	if m == nil {
-		return nil
-	}
-
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	rawStates, err := m.loadStateFile(false)
-	if err != nil {
-		return err
-	}
-	if rawStates == nil {
-		return nil
-	}
-
-	name := state.Name()
-	rawState, exists := rawStates[name]
-	if !exists {
-		return nil
-	}
-
-	loadedState, err := m.loadSingleRawState(name, rawState)
-	if err != nil {
-		return err
-	}
-
-	m.states[name] = loadedState
-	if loadedState != nil {
+		m.states[name] = statePtr
 		log.Debugf("loaded state: %s", name)
 	}
 
-	return nil
+	return nberrors.FormatErrorOrNil(merr)
 }
 
-// cleanupSingleState handles the cleanup of a specific state and returns any error.
-// The caller must hold the mutex.
-func (m *Manager) cleanupSingleState(name string, rawState json.RawMessage) error {
-	// For unregistered states, preserve the raw JSON
-	if _, registered := m.stateTypes[name]; !registered {
-		m.states[name] = &RawState{data: rawState}
-		return nil
-	}
-
-	// Load the state
-	loadedState, err := m.loadSingleRawState(name, rawState)
-	if err != nil {
-		return err
-	}
-
-	if loadedState == nil {
-		return nil
-	}
-
-	// Check if state supports cleanup
-	cleanableState, isCleanable := loadedState.(CleanableState)
-	if !isCleanable {
-		// If it doesn't support cleanup, keep it as-is
-		m.states[name] = loadedState
-		return nil
-	}
-
-	// Perform cleanup
-	log.Infof("cleaning up state %s", name)
-	if err := cleanableState.Cleanup(); err != nil {
-		// On cleanup error, preserve the state
-		m.states[name] = loadedState
-		return fmt.Errorf("cleanup state: %w", err)
-	}
-
-	// Successfully cleaned up - mark for deletion
-	m.states[name] = nil
-	m.dirty[name] = struct{}{}
-	return nil
-}
-
-// CleanupStateByName loads and cleans up a specific state by name if it implements CleanableState.
-// Returns an error if the state doesn't exist, isn't registered, or cleanup fails.
-func (m *Manager) CleanupStateByName(name string) error {
-	if m == nil {
-		return nil
-	}
-
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	// Check if state is registered
-	if _, registered := m.stateTypes[name]; !registered {
-		return fmt.Errorf(errStateNotRegistered, name)
-	}
-
-	// Load raw states from file
-	rawStates, err := m.loadStateFile(false)
-	if err != nil {
-		return err
-	}
-	if rawStates == nil {
-		return nil
-	}
-
-	// Check if state exists in file
-	rawState, exists := rawStates[name]
-	if !exists {
-		return nil
-	}
-
-	if err := m.cleanupSingleState(name, rawState); err != nil {
-		return fmt.Errorf("%s: %w", name, err)
-	}
-
-	return nil
-}
-
-// PerformCleanup retrieves all states from the state file and calls Cleanup on registered states that support it.
-// Unregistered states are preserved in their original state.
+// PerformCleanup retrieves all states from the state file for the registered states and calls Cleanup on them.
+// If the cleanup is successful, the state is marked for deletion.
 func (m *Manager) PerformCleanup() error {
 	if m == nil {
 		return nil
@@ -460,63 +273,26 @@ func (m *Manager) PerformCleanup() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	// Load raw states from file
-	rawStates, err := m.loadStateFile(true)
-	if err != nil {
-		return fmt.Errorf(errLoadStateFile, err)
-	}
-	if rawStates == nil {
-		return nil
+	if err := m.loadState(); err != nil {
+		log.Warnf("Failed to load state during cleanup: %v", err)
 	}
 
 	var merr *multierror.Error
+	for name, state := range m.states {
+		if state == nil {
+			// If no state was found in the state file, we don't mark the state dirty nor return an error
+			continue
+		}
 
-	// Process each state in the file
-	for name, rawState := range rawStates {
-		if err := m.cleanupSingleState(name, rawState); err != nil {
-			merr = multierror.Append(merr, fmt.Errorf("%s: %w", name, err))
+		log.Infof("client was not shut down properly, cleaning up %s", name)
+		if err := state.Cleanup(); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("cleanup state for %s: %w", name, err))
+		} else {
+			// mark for deletion on cleanup success
+			m.states[name] = nil
+			m.dirty[name] = struct{}{}
 		}
 	}
 
 	return nberrors.FormatErrorOrNil(merr)
 }
-
-// GetSavedStateNames returns all state names that are currently saved in the state file.
-func (m *Manager) GetSavedStateNames() ([]string, error) {
-	if m == nil {
-		return nil, nil
-	}
-
-	rawStates, err := m.loadStateFile(false)
-	if err != nil {
-		return nil, fmt.Errorf(errLoadStateFile, err)
-	}
-	if rawStates == nil {
-		return nil, nil
-	}
-
-	var states []string
-	for name, state := range rawStates {
-		if len(state) != 0 && string(state) != "null" {
-			states = append(states, name)
-		}
-	}
-
-	return states, nil
-}
-
-func marshalWithPanicRecovery(v any) ([]byte, error) {
-	var bs []byte
-	var err error
-
-	func() {
-		defer func() {
-			if r := recover(); r != nil {
-				err = fmt.Errorf("panic during marshal: %v", r)
-			}
-		}()
-		bs, err = json.Marshal(v)
-	}()
-
-	return bs, err
-}

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/proto/daemon.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.26.0
-// 	protoc        v4.23.4
+// 	protoc        v3.21.12
 // source: daemon.proto
 
 package proto
@@ -2103,434 +2103,6 @@ func (*SetLogLevelResponse) Descriptor() ([]byte, []int) {
 	return file_daemon_proto_rawDescGZIP(), []int{30}
 }
 
-// State represents a daemon state entry
-type State struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
-}
-
-func (x *State) Reset() {
-	*x = State{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[31]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *State) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*State) ProtoMessage() {}
-
-func (x *State) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[31]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use State.ProtoReflect.Descriptor instead.
-func (*State) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{31}
-}
-
-func (x *State) GetName() string {
-	if x != nil {
-		return x.Name
-	}
-	return ""
-}
-
-// ListStatesRequest is empty as it requires no parameters
-type ListStatesRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-}
-
-func (x *ListStatesRequest) Reset() {
-	*x = ListStatesRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[32]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *ListStatesRequest) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*ListStatesRequest) ProtoMessage() {}
-
-func (x *ListStatesRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[32]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use ListStatesRequest.ProtoReflect.Descriptor instead.
-func (*ListStatesRequest) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{32}
-}
-
-// ListStatesResponse contains a list of states
-type ListStatesResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	States []*State `protobuf:"bytes,1,rep,name=states,proto3" json:"states,omitempty"`
-}
-
-func (x *ListStatesResponse) Reset() {
-	*x = ListStatesResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[33]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *ListStatesResponse) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*ListStatesResponse) ProtoMessage() {}
-
-func (x *ListStatesResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[33]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use ListStatesResponse.ProtoReflect.Descriptor instead.
-func (*ListStatesResponse) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{33}
-}
-
-func (x *ListStatesResponse) GetStates() []*State {
-	if x != nil {
-		return x.States
-	}
-	return nil
-}
-
-// CleanStateRequest for cleaning states
-type CleanStateRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	StateName string `protobuf:"bytes,1,opt,name=state_name,json=stateName,proto3" json:"state_name,omitempty"`
-	All       bool   `protobuf:"varint,2,opt,name=all,proto3" json:"all,omitempty"`
-}
-
-func (x *CleanStateRequest) Reset() {
-	*x = CleanStateRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[34]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *CleanStateRequest) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*CleanStateRequest) ProtoMessage() {}
-
-func (x *CleanStateRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[34]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use CleanStateRequest.ProtoReflect.Descriptor instead.
-func (*CleanStateRequest) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{34}
-}
-
-func (x *CleanStateRequest) GetStateName() string {
-	if x != nil {
-		return x.StateName
-	}
-	return ""
-}
-
-func (x *CleanStateRequest) GetAll() bool {
-	if x != nil {
-		return x.All
-	}
-	return false
-}
-
-// CleanStateResponse contains the result of the clean operation
-type CleanStateResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	CleanedStates int32 `protobuf:"varint,1,opt,name=cleaned_states,json=cleanedStates,proto3" json:"cleaned_states,omitempty"`
-}
-
-func (x *CleanStateResponse) Reset() {
-	*x = CleanStateResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[35]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *CleanStateResponse) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*CleanStateResponse) ProtoMessage() {}
-
-func (x *CleanStateResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[35]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use CleanStateResponse.ProtoReflect.Descriptor instead.
-func (*CleanStateResponse) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{35}
-}
-
-func (x *CleanStateResponse) GetCleanedStates() int32 {
-	if x != nil {
-		return x.CleanedStates
-	}
-	return 0
-}
-
-// DeleteStateRequest for deleting states
-type DeleteStateRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	StateName string `protobuf:"bytes,1,opt,name=state_name,json=stateName,proto3" json:"state_name,omitempty"`
-	All       bool   `protobuf:"varint,2,opt,name=all,proto3" json:"all,omitempty"`
-}
-
-func (x *DeleteStateRequest) Reset() {
-	*x = DeleteStateRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[36]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *DeleteStateRequest) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*DeleteStateRequest) ProtoMessage() {}
-
-func (x *DeleteStateRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[36]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use DeleteStateRequest.ProtoReflect.Descriptor instead.
-func (*DeleteStateRequest) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{36}
-}
-
-func (x *DeleteStateRequest) GetStateName() string {
-	if x != nil {
-		return x.StateName
-	}
-	return ""
-}
-
-func (x *DeleteStateRequest) GetAll() bool {
-	if x != nil {
-		return x.All
-	}
-	return false
-}
-
-// DeleteStateResponse contains the result of the delete operation
-type DeleteStateResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	DeletedStates int32 `protobuf:"varint,1,opt,name=deleted_states,json=deletedStates,proto3" json:"deleted_states,omitempty"`
-}
-
-func (x *DeleteStateResponse) Reset() {
-	*x = DeleteStateResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[37]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *DeleteStateResponse) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*DeleteStateResponse) ProtoMessage() {}
-
-func (x *DeleteStateResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[37]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use DeleteStateResponse.ProtoReflect.Descriptor instead.
-func (*DeleteStateResponse) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{37}
-}
-
-func (x *DeleteStateResponse) GetDeletedStates() int32 {
-	if x != nil {
-		return x.DeletedStates
-	}
-	return 0
-}
-
-type SetNetworkMapPersistenceRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	Enabled bool `protobuf:"varint,1,opt,name=enabled,proto3" json:"enabled,omitempty"`
-}
-
-func (x *SetNetworkMapPersistenceRequest) Reset() {
-	*x = SetNetworkMapPersistenceRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[38]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *SetNetworkMapPersistenceRequest) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*SetNetworkMapPersistenceRequest) ProtoMessage() {}
-
-func (x *SetNetworkMapPersistenceRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[38]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use SetNetworkMapPersistenceRequest.ProtoReflect.Descriptor instead.
-func (*SetNetworkMapPersistenceRequest) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{38}
-}
-
-func (x *SetNetworkMapPersistenceRequest) GetEnabled() bool {
-	if x != nil {
-		return x.Enabled
-	}
-	return false
-}
-
-type SetNetworkMapPersistenceResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-}
-
-func (x *SetNetworkMapPersistenceResponse) Reset() {
-	*x = SetNetworkMapPersistenceResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_daemon_proto_msgTypes[39]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
-}
-
-func (x *SetNetworkMapPersistenceResponse) String() string {
-	return protoimpl.X.MessageStringOf(x)
-}
-
-func (*SetNetworkMapPersistenceResponse) ProtoMessage() {}
-
-func (x *SetNetworkMapPersistenceResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_daemon_proto_msgTypes[39]
-	if protoimpl.UnsafeEnabled && x != nil {
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		if ms.LoadMessageInfo() == nil {
-			ms.StoreMessageInfo(mi)
-		}
-		return ms
-	}
-	return mi.MessageOf(x)
-}
-
-// Deprecated: Use SetNetworkMapPersistenceResponse.ProtoReflect.Descriptor instead.
-func (*SetNetworkMapPersistenceResponse) Descriptor() ([]byte, []int) {
-	return file_daemon_proto_rawDescGZIP(), []int{39}
-}
-
 var File_daemon_proto protoreflect.FileDescriptor
 
 var file_daemon_proto_rawDesc = []byte{
@@ -2827,116 +2399,66 @@ var file_daemon_proto_rawDesc = []byte{
 	0x32, 0x10, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76,
 	0x65, 0x6c, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x22, 0x15, 0x0a, 0x13, 0x53, 0x65, 0x74,
 	0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x1b, 0x0a, 0x05, 0x53, 0x74, 0x61, 0x74, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d,
-	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x13, 0x0a,
-	0x11, 0x4c, 0x69, 0x73, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65,
-	0x73, 0x74, 0x22, 0x3b, 0x0a, 0x12, 0x4c, 0x69, 0x73, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x73,
-	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74,
-	0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0d, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f,
-	0x6e, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x65, 0x73, 0x22,
-	0x44, 0x0a, 0x11, 0x43, 0x6c, 0x65, 0x61, 0x6e, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x74, 0x61, 0x74, 0x65, 0x5f, 0x6e, 0x61,
-	0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x74, 0x61, 0x74, 0x65, 0x4e,
-	0x61, 0x6d, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x61, 0x6c, 0x6c, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08,
-	0x52, 0x03, 0x61, 0x6c, 0x6c, 0x22, 0x3b, 0x0a, 0x12, 0x43, 0x6c, 0x65, 0x61, 0x6e, 0x53, 0x74,
-	0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x0e, 0x63,
-	0x6c, 0x65, 0x61, 0x6e, 0x65, 0x64, 0x5f, 0x73, 0x74, 0x61, 0x74, 0x65, 0x73, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x05, 0x52, 0x0d, 0x63, 0x6c, 0x65, 0x61, 0x6e, 0x65, 0x64, 0x53, 0x74, 0x61, 0x74,
-	0x65, 0x73, 0x22, 0x45, 0x0a, 0x12, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53, 0x74, 0x61, 0x74,
-	0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x74, 0x61, 0x74,
-	0x65, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x74,
-	0x61, 0x74, 0x65, 0x4e, 0x61, 0x6d, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x61, 0x6c, 0x6c, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x08, 0x52, 0x03, 0x61, 0x6c, 0x6c, 0x22, 0x3c, 0x0a, 0x13, 0x44, 0x65, 0x6c,
-	0x65, 0x74, 0x65, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x12, 0x25, 0x0a, 0x0e, 0x64, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x64, 0x5f, 0x73, 0x74, 0x61, 0x74,
-	0x65, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0d, 0x64, 0x65, 0x6c, 0x65, 0x74, 0x65,
-	0x64, 0x53, 0x74, 0x61, 0x74, 0x65, 0x73, 0x22, 0x3b, 0x0a, 0x1f, 0x53, 0x65, 0x74, 0x4e, 0x65,
-	0x74, 0x77, 0x6f, 0x72, 0x6b, 0x4d, 0x61, 0x70, 0x50, 0x65, 0x72, 0x73, 0x69, 0x73, 0x74, 0x65,
-	0x6e, 0x63, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x65, 0x6e,
-	0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x65, 0x6e, 0x61,
-	0x62, 0x6c, 0x65, 0x64, 0x22, 0x22, 0x0a, 0x20, 0x53, 0x65, 0x74, 0x4e, 0x65, 0x74, 0x77, 0x6f,
-	0x72, 0x6b, 0x4d, 0x61, 0x70, 0x50, 0x65, 0x72, 0x73, 0x69, 0x73, 0x74, 0x65, 0x6e, 0x63, 0x65,
-	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x2a, 0x62, 0x0a, 0x08, 0x4c, 0x6f, 0x67, 0x4c,
-	0x65, 0x76, 0x65, 0x6c, 0x12, 0x0b, 0x0a, 0x07, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10,
-	0x00, 0x12, 0x09, 0x0a, 0x05, 0x50, 0x41, 0x4e, 0x49, 0x43, 0x10, 0x01, 0x12, 0x09, 0x0a, 0x05,
-	0x46, 0x41, 0x54, 0x41, 0x4c, 0x10, 0x02, 0x12, 0x09, 0x0a, 0x05, 0x45, 0x52, 0x52, 0x4f, 0x52,
-	0x10, 0x03, 0x12, 0x08, 0x0a, 0x04, 0x57, 0x41, 0x52, 0x4e, 0x10, 0x04, 0x12, 0x08, 0x0a, 0x04,
-	0x49, 0x4e, 0x46, 0x4f, 0x10, 0x05, 0x12, 0x09, 0x0a, 0x05, 0x44, 0x45, 0x42, 0x55, 0x47, 0x10,
-	0x06, 0x12, 0x09, 0x0a, 0x05, 0x54, 0x52, 0x41, 0x43, 0x45, 0x10, 0x07, 0x32, 0x81, 0x09, 0x0a,
-	0x0d, 0x44, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x36,
-	0x0a, 0x05, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x12, 0x14, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e,
-	0x2e, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e,
-	0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4b, 0x0a, 0x0c, 0x57, 0x61, 0x69, 0x74, 0x53, 0x53,
-	0x4f, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x12, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e,
-	0x57, 0x61, 0x69, 0x74, 0x53, 0x53, 0x4f, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x71, 0x75,
-	0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x57, 0x61, 0x69,
-	0x74, 0x53, 0x53, 0x4f, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73,
-	0x65, 0x22, 0x00, 0x12, 0x2d, 0x0a, 0x02, 0x55, 0x70, 0x12, 0x11, 0x2e, 0x64, 0x61, 0x65, 0x6d,
-	0x6f, 0x6e, 0x2e, 0x55, 0x70, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x12, 0x2e, 0x64,
-	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x55, 0x70, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x00, 0x12, 0x39, 0x0a, 0x06, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x15, 0x2e, 0x64,
-	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x75,
-	0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x74, 0x61,
-	0x74, 0x75, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x33, 0x0a,
-	0x04, 0x44, 0x6f, 0x77, 0x6e, 0x12, 0x13, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44,
-	0x6f, 0x77, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x14, 0x2e, 0x64, 0x61, 0x65,
-	0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x6f, 0x77, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x00, 0x12, 0x42, 0x0a, 0x09, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12,
-	0x18, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6e, 0x66,
-	0x69, 0x67, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x64, 0x61, 0x65, 0x6d,
-	0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x52, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x45, 0x0a, 0x0a, 0x4c, 0x69, 0x73, 0x74, 0x52, 0x6f,
-	0x75, 0x74, 0x65, 0x73, 0x12, 0x19, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x69,
-	0x73, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a,
-	0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x69, 0x73, 0x74, 0x52, 0x6f, 0x75,
-	0x74, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4b, 0x0a,
-	0x0c, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x12, 0x1b, 0x2e,
-	0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75,
-	0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61, 0x65,
-	0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73,
-	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4d, 0x0a, 0x0e, 0x44, 0x65,
-	0x73, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x12, 0x1b, 0x2e, 0x64,
-	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74,
-	0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61, 0x65, 0x6d,
-	0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52,
-	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x44, 0x65, 0x62,
-	0x75, 0x67, 0x42, 0x75, 0x6e, 0x64, 0x6c, 0x65, 0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f,
-	0x6e, 0x2e, 0x44, 0x65, 0x62, 0x75, 0x67, 0x42, 0x75, 0x6e, 0x64, 0x6c, 0x65, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x65,
-	0x62, 0x75, 0x67, 0x42, 0x75, 0x6e, 0x64, 0x6c, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73,
-	0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x47, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76,
-	0x65, 0x6c, 0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x4c,
-	0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b,
-	0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65,
-	0x76, 0x65, 0x6c, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a,
-	0x0b, 0x53, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x1a, 0x2e, 0x64,
-	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65,
-	0x6c, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f,
-	0x6e, 0x2e, 0x53, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x45, 0x0a, 0x0a, 0x4c, 0x69, 0x73, 0x74, 0x53,
-	0x74, 0x61, 0x74, 0x65, 0x73, 0x12, 0x19, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c,
-	0x69, 0x73, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
-	0x1a, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x69, 0x73, 0x74, 0x53, 0x74,
-	0x61, 0x74, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x45,
-	0x0a, 0x0a, 0x43, 0x6c, 0x65, 0x61, 0x6e, 0x53, 0x74, 0x61, 0x74, 0x65, 0x12, 0x19, 0x2e, 0x64,
-	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x43, 0x6c, 0x65, 0x61, 0x6e, 0x53, 0x74, 0x61, 0x74, 0x65,
+	0x2a, 0x62, 0x0a, 0x08, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x0b, 0x0a, 0x07,
+	0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x09, 0x0a, 0x05, 0x50, 0x41, 0x4e,
+	0x49, 0x43, 0x10, 0x01, 0x12, 0x09, 0x0a, 0x05, 0x46, 0x41, 0x54, 0x41, 0x4c, 0x10, 0x02, 0x12,
+	0x09, 0x0a, 0x05, 0x45, 0x52, 0x52, 0x4f, 0x52, 0x10, 0x03, 0x12, 0x08, 0x0a, 0x04, 0x57, 0x41,
+	0x52, 0x4e, 0x10, 0x04, 0x12, 0x08, 0x0a, 0x04, 0x49, 0x4e, 0x46, 0x4f, 0x10, 0x05, 0x12, 0x09,
+	0x0a, 0x05, 0x44, 0x45, 0x42, 0x55, 0x47, 0x10, 0x06, 0x12, 0x09, 0x0a, 0x05, 0x54, 0x52, 0x41,
+	0x43, 0x45, 0x10, 0x07, 0x32, 0xb8, 0x06, 0x0a, 0x0d, 0x44, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x53,
+	0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x36, 0x0a, 0x05, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x12,
+	0x14, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65,
+	0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c,
+	0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4b,
+	0x0a, 0x0c, 0x57, 0x61, 0x69, 0x74, 0x53, 0x53, 0x4f, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x12, 0x1b,
+	0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x57, 0x61, 0x69, 0x74, 0x53, 0x53, 0x4f, 0x4c,
+	0x6f, 0x67, 0x69, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61,
+	0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x57, 0x61, 0x69, 0x74, 0x53, 0x53, 0x4f, 0x4c, 0x6f, 0x67, 0x69,
+	0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x2d, 0x0a, 0x02, 0x55,
+	0x70, 0x12, 0x11, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x55, 0x70, 0x52, 0x65, 0x71,
+	0x75, 0x65, 0x73, 0x74, 0x1a, 0x12, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x55, 0x70,
+	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x39, 0x0a, 0x06, 0x53, 0x74,
+	0x61, 0x74, 0x75, 0x73, 0x12, 0x15, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x74,
+	0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x64, 0x61,
+	0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f,
+	0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x33, 0x0a, 0x04, 0x44, 0x6f, 0x77, 0x6e, 0x12, 0x13, 0x2e,
+	0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x6f, 0x77, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65,
+	0x73, 0x74, 0x1a, 0x14, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x6f, 0x77, 0x6e,
+	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x42, 0x0a, 0x09, 0x47, 0x65,
+	0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x18, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e,
+	0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
+	0x74, 0x1a, 0x19, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x43, 0x6f,
+	0x6e, 0x66, 0x69, 0x67, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x45,
+	0x0a, 0x0a, 0x4c, 0x69, 0x73, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x12, 0x19, 0x2e, 0x64,
+	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x4c, 0x69, 0x73, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73,
 	0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e,
-	0x2e, 0x43, 0x6c, 0x65, 0x61, 0x6e, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f,
-	0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x53,
-	0x74, 0x61, 0x74, 0x65, 0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x65,
-	0x6c, 0x65, 0x74, 0x65, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
-	0x1a, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65,
-	0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12,
-	0x6f, 0x0a, 0x18, 0x53, 0x65, 0x74, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x4d, 0x61, 0x70,
-	0x50, 0x65, 0x72, 0x73, 0x69, 0x73, 0x74, 0x65, 0x6e, 0x63, 0x65, 0x12, 0x27, 0x2e, 0x64, 0x61,
-	0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x74, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x4d,
-	0x61, 0x70, 0x50, 0x65, 0x72, 0x73, 0x69, 0x73, 0x74, 0x65, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x1a, 0x28, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65,
-	0x74, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x4d, 0x61, 0x70, 0x50, 0x65, 0x72, 0x73, 0x69,
-	0x73, 0x74, 0x65, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00,
-	0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
-	0x6f, 0x33,
+	0x2e, 0x4c, 0x69, 0x73, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f,
+	0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x4b, 0x0a, 0x0c, 0x53, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52,
+	0x6f, 0x75, 0x74, 0x65, 0x73, 0x12, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53,
+	0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65,
+	0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65,
+	0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
+	0x22, 0x00, 0x12, 0x4d, 0x0a, 0x0e, 0x44, 0x65, 0x73, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f,
+	0x75, 0x74, 0x65, 0x73, 0x12, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65,
+	0x6c, 0x65, 0x63, 0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
+	0x74, 0x1a, 0x1c, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x6c, 0x65, 0x63,
+	0x74, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22,
+	0x00, 0x12, 0x48, 0x0a, 0x0b, 0x44, 0x65, 0x62, 0x75, 0x67, 0x42, 0x75, 0x6e, 0x64, 0x6c, 0x65,
+	0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x65, 0x62, 0x75, 0x67, 0x42,
+	0x75, 0x6e, 0x64, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x64,
+	0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x44, 0x65, 0x62, 0x75, 0x67, 0x42, 0x75, 0x6e, 0x64, 0x6c,
+	0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x47,
+	0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65,
+	0x6d, 0x6f, 0x6e, 0x2e, 0x47, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52,
+	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e,
+	0x47, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x73, 0x70, 0x6f,
+	0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x0b, 0x53, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x4c,
+	0x65, 0x76, 0x65, 0x6c, 0x12, 0x1a, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65,
+	0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
+	0x1a, 0x1b, 0x2e, 0x64, 0x61, 0x65, 0x6d, 0x6f, 0x6e, 0x2e, 0x53, 0x65, 0x74, 0x4c, 0x6f, 0x67,
+	0x4c, 0x65, 0x76, 0x65, 0x6c, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42,
+	0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f,
+	0x33,
 }
 
 var (
@@ -2952,59 +2474,50 @@ func file_daemon_proto_rawDescGZIP() []byte {
 }
 
 var file_daemon_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
-var file_daemon_proto_msgTypes = make([]protoimpl.MessageInfo, 41)
+var file_daemon_proto_msgTypes = make([]protoimpl.MessageInfo, 32)
 var file_daemon_proto_goTypes = []interface{}{
-	(LogLevel)(0),                            // 0: daemon.LogLevel
-	(*LoginRequest)(nil),                     // 1: daemon.LoginRequest
-	(*LoginResponse)(nil),                    // 2: daemon.LoginResponse
-	(*WaitSSOLoginRequest)(nil),              // 3: daemon.WaitSSOLoginRequest
-	(*WaitSSOLoginResponse)(nil),             // 4: daemon.WaitSSOLoginResponse
-	(*UpRequest)(nil),                        // 5: daemon.UpRequest
-	(*UpResponse)(nil),                       // 6: daemon.UpResponse
-	(*StatusRequest)(nil),                    // 7: daemon.StatusRequest
-	(*StatusResponse)(nil),                   // 8: daemon.StatusResponse
-	(*DownRequest)(nil),                      // 9: daemon.DownRequest
-	(*DownResponse)(nil),                     // 10: daemon.DownResponse
-	(*GetConfigRequest)(nil),                 // 11: daemon.GetConfigRequest
-	(*GetConfigResponse)(nil),                // 12: daemon.GetConfigResponse
-	(*PeerState)(nil),                        // 13: daemon.PeerState
-	(*LocalPeerState)(nil),                   // 14: daemon.LocalPeerState
-	(*SignalState)(nil),                      // 15: daemon.SignalState
-	(*ManagementState)(nil),                  // 16: daemon.ManagementState
-	(*RelayState)(nil),                       // 17: daemon.RelayState
-	(*NSGroupState)(nil),                     // 18: daemon.NSGroupState
-	(*FullStatus)(nil),                       // 19: daemon.FullStatus
-	(*ListRoutesRequest)(nil),                // 20: daemon.ListRoutesRequest
-	(*ListRoutesResponse)(nil),               // 21: daemon.ListRoutesResponse
-	(*SelectRoutesRequest)(nil),              // 22: daemon.SelectRoutesRequest
-	(*SelectRoutesResponse)(nil),             // 23: daemon.SelectRoutesResponse
-	(*IPList)(nil),                           // 24: daemon.IPList
-	(*Route)(nil),                            // 25: daemon.Route
-	(*DebugBundleRequest)(nil),               // 26: daemon.DebugBundleRequest
-	(*DebugBundleResponse)(nil),              // 27: daemon.DebugBundleResponse
-	(*GetLogLevelRequest)(nil),               // 28: daemon.GetLogLevelRequest
-	(*GetLogLevelResponse)(nil),              // 29: daemon.GetLogLevelResponse
-	(*SetLogLevelRequest)(nil),               // 30: daemon.SetLogLevelRequest
-	(*SetLogLevelResponse)(nil),              // 31: daemon.SetLogLevelResponse
-	(*State)(nil),                            // 32: daemon.State
-	(*ListStatesRequest)(nil),                // 33: daemon.ListStatesRequest
-	(*ListStatesResponse)(nil),               // 34: daemon.ListStatesResponse
-	(*CleanStateRequest)(nil),                // 35: daemon.CleanStateRequest
-	(*CleanStateResponse)(nil),               // 36: daemon.CleanStateResponse
-	(*DeleteStateRequest)(nil),               // 37: daemon.DeleteStateRequest
-	(*DeleteStateResponse)(nil),              // 38: daemon.DeleteStateResponse
-	(*SetNetworkMapPersistenceRequest)(nil),  // 39: daemon.SetNetworkMapPersistenceRequest
-	(*SetNetworkMapPersistenceResponse)(nil), // 40: daemon.SetNetworkMapPersistenceResponse
-	nil,                                      // 41: daemon.Route.ResolvedIPsEntry
-	(*durationpb.Duration)(nil),              // 42: google.protobuf.Duration
-	(*timestamppb.Timestamp)(nil),            // 43: google.protobuf.Timestamp
+	(LogLevel)(0),                 // 0: daemon.LogLevel
+	(*LoginRequest)(nil),          // 1: daemon.LoginRequest
+	(*LoginResponse)(nil),         // 2: daemon.LoginResponse
+	(*WaitSSOLoginRequest)(nil),   // 3: daemon.WaitSSOLoginRequest
+	(*WaitSSOLoginResponse)(nil),  // 4: daemon.WaitSSOLoginResponse
+	(*UpRequest)(nil),             // 5: daemon.UpRequest
+	(*UpResponse)(nil),            // 6: daemon.UpResponse
+	(*StatusRequest)(nil),         // 7: daemon.StatusRequest
+	(*StatusResponse)(nil),        // 8: daemon.StatusResponse
+	(*DownRequest)(nil),           // 9: daemon.DownRequest
+	(*DownResponse)(nil),          // 10: daemon.DownResponse
+	(*GetConfigRequest)(nil),      // 11: daemon.GetConfigRequest
+	(*GetConfigResponse)(nil),     // 12: daemon.GetConfigResponse
+	(*PeerState)(nil),             // 13: daemon.PeerState
+	(*LocalPeerState)(nil),        // 14: daemon.LocalPeerState
+	(*SignalState)(nil),           // 15: daemon.SignalState
+	(*ManagementState)(nil),       // 16: daemon.ManagementState
+	(*RelayState)(nil),            // 17: daemon.RelayState
+	(*NSGroupState)(nil),          // 18: daemon.NSGroupState
+	(*FullStatus)(nil),            // 19: daemon.FullStatus
+	(*ListRoutesRequest)(nil),     // 20: daemon.ListRoutesRequest
+	(*ListRoutesResponse)(nil),    // 21: daemon.ListRoutesResponse
+	(*SelectRoutesRequest)(nil),   // 22: daemon.SelectRoutesRequest
+	(*SelectRoutesResponse)(nil),  // 23: daemon.SelectRoutesResponse
+	(*IPList)(nil),                // 24: daemon.IPList
+	(*Route)(nil),                 // 25: daemon.Route
+	(*DebugBundleRequest)(nil),    // 26: daemon.DebugBundleRequest
+	(*DebugBundleResponse)(nil),   // 27: daemon.DebugBundleResponse
+	(*GetLogLevelRequest)(nil),    // 28: daemon.GetLogLevelRequest
+	(*GetLogLevelResponse)(nil),   // 29: daemon.GetLogLevelResponse
+	(*SetLogLevelRequest)(nil),    // 30: daemon.SetLogLevelRequest
+	(*SetLogLevelResponse)(nil),   // 31: daemon.SetLogLevelResponse
+	nil,                           // 32: daemon.Route.ResolvedIPsEntry
+	(*durationpb.Duration)(nil),   // 33: google.protobuf.Duration
+	(*timestamppb.Timestamp)(nil), // 34: google.protobuf.Timestamp
 }
 var file_daemon_proto_depIdxs = []int32{
-	42, // 0: daemon.LoginRequest.dnsRouteInterval:type_name -> google.protobuf.Duration
+	33, // 0: daemon.LoginRequest.dnsRouteInterval:type_name -> google.protobuf.Duration
 	19, // 1: daemon.StatusResponse.fullStatus:type_name -> daemon.FullStatus
-	43, // 2: daemon.PeerState.connStatusUpdate:type_name -> google.protobuf.Timestamp
-	43, // 3: daemon.PeerState.lastWireguardHandshake:type_name -> google.protobuf.Timestamp
-	42, // 4: daemon.PeerState.latency:type_name -> google.protobuf.Duration
+	34, // 2: daemon.PeerState.connStatusUpdate:type_name -> google.protobuf.Timestamp
+	34, // 3: daemon.PeerState.lastWireguardHandshake:type_name -> google.protobuf.Timestamp
+	33, // 4: daemon.PeerState.latency:type_name -> google.protobuf.Duration
 	16, // 5: daemon.FullStatus.managementState:type_name -> daemon.ManagementState
 	15, // 6: daemon.FullStatus.signalState:type_name -> daemon.SignalState
 	14, // 7: daemon.FullStatus.localPeerState:type_name -> daemon.LocalPeerState
@@ -3012,48 +2525,39 @@ var file_daemon_proto_depIdxs = []int32{
 	17, // 9: daemon.FullStatus.relays:type_name -> daemon.RelayState
 	18, // 10: daemon.FullStatus.dns_servers:type_name -> daemon.NSGroupState
 	25, // 11: daemon.ListRoutesResponse.routes:type_name -> daemon.Route
-	41, // 12: daemon.Route.resolvedIPs:type_name -> daemon.Route.ResolvedIPsEntry
+	32, // 12: daemon.Route.resolvedIPs:type_name -> daemon.Route.ResolvedIPsEntry
 	0,  // 13: daemon.GetLogLevelResponse.level:type_name -> daemon.LogLevel
 	0,  // 14: daemon.SetLogLevelRequest.level:type_name -> daemon.LogLevel
-	32, // 15: daemon.ListStatesResponse.states:type_name -> daemon.State
-	24, // 16: daemon.Route.ResolvedIPsEntry.value:type_name -> daemon.IPList
-	1,  // 17: daemon.DaemonService.Login:input_type -> daemon.LoginRequest
-	3,  // 18: daemon.DaemonService.WaitSSOLogin:input_type -> daemon.WaitSSOLoginRequest
-	5,  // 19: daemon.DaemonService.Up:input_type -> daemon.UpRequest
-	7,  // 20: daemon.DaemonService.Status:input_type -> daemon.StatusRequest
-	9,  // 21: daemon.DaemonService.Down:input_type -> daemon.DownRequest
-	11, // 22: daemon.DaemonService.GetConfig:input_type -> daemon.GetConfigRequest
-	20, // 23: daemon.DaemonService.ListRoutes:input_type -> daemon.ListRoutesRequest
-	22, // 24: daemon.DaemonService.SelectRoutes:input_type -> daemon.SelectRoutesRequest
-	22, // 25: daemon.DaemonService.DeselectRoutes:input_type -> daemon.SelectRoutesRequest
-	26, // 26: daemon.DaemonService.DebugBundle:input_type -> daemon.DebugBundleRequest
-	28, // 27: daemon.DaemonService.GetLogLevel:input_type -> daemon.GetLogLevelRequest
-	30, // 28: daemon.DaemonService.SetLogLevel:input_type -> daemon.SetLogLevelRequest
-	33, // 29: daemon.DaemonService.ListStates:input_type -> daemon.ListStatesRequest
-	35, // 30: daemon.DaemonService.CleanState:input_type -> daemon.CleanStateRequest
-	37, // 31: daemon.DaemonService.DeleteState:input_type -> daemon.DeleteStateRequest
-	39, // 32: daemon.DaemonService.SetNetworkMapPersistence:input_type -> daemon.SetNetworkMapPersistenceRequest
-	2,  // 33: daemon.DaemonService.Login:output_type -> daemon.LoginResponse
-	4,  // 34: daemon.DaemonService.WaitSSOLogin:output_type -> daemon.WaitSSOLoginResponse
-	6,  // 35: daemon.DaemonService.Up:output_type -> daemon.UpResponse
-	8,  // 36: daemon.DaemonService.Status:output_type -> daemon.StatusResponse
-	10, // 37: daemon.DaemonService.Down:output_type -> daemon.DownResponse
-	12, // 38: daemon.DaemonService.GetConfig:output_type -> daemon.GetConfigResponse
-	21, // 39: daemon.DaemonService.ListRoutes:output_type -> daemon.ListRoutesResponse
-	23, // 40: daemon.DaemonService.SelectRoutes:output_type -> daemon.SelectRoutesResponse
-	23, // 41: daemon.DaemonService.DeselectRoutes:output_type -> daemon.SelectRoutesResponse
-	27, // 42: daemon.DaemonService.DebugBundle:output_type -> daemon.DebugBundleResponse
-	29, // 43: daemon.DaemonService.GetLogLevel:output_type -> daemon.GetLogLevelResponse
-	31, // 44: daemon.DaemonService.SetLogLevel:output_type -> daemon.SetLogLevelResponse
-	34, // 45: daemon.DaemonService.ListStates:output_type -> daemon.ListStatesResponse
-	36, // 46: daemon.DaemonService.CleanState:output_type -> daemon.CleanStateResponse
-	38, // 47: daemon.DaemonService.DeleteState:output_type -> daemon.DeleteStateResponse
-	40, // 48: daemon.DaemonService.SetNetworkMapPersistence:output_type -> daemon.SetNetworkMapPersistenceResponse
-	33, // [33:49] is the sub-list for method output_type
-	17, // [17:33] is the sub-list for method input_type
-	17, // [17:17] is the sub-list for extension type_name
-	17, // [17:17] is the sub-list for extension extendee
-	0,  // [0:17] is the sub-list for field type_name
+	24, // 15: daemon.Route.ResolvedIPsEntry.value:type_name -> daemon.IPList
+	1,  // 16: daemon.DaemonService.Login:input_type -> daemon.LoginRequest
+	3,  // 17: daemon.DaemonService.WaitSSOLogin:input_type -> daemon.WaitSSOLoginRequest
+	5,  // 18: daemon.DaemonService.Up:input_type -> daemon.UpRequest
+	7,  // 19: daemon.DaemonService.Status:input_type -> daemon.StatusRequest
+	9,  // 20: daemon.DaemonService.Down:input_type -> daemon.DownRequest
+	11, // 21: daemon.DaemonService.GetConfig:input_type -> daemon.GetConfigRequest
+	20, // 22: daemon.DaemonService.ListRoutes:input_type -> daemon.ListRoutesRequest
+	22, // 23: daemon.DaemonService.SelectRoutes:input_type -> daemon.SelectRoutesRequest
+	22, // 24: daemon.DaemonService.DeselectRoutes:input_type -> daemon.SelectRoutesRequest
+	26, // 25: daemon.DaemonService.DebugBundle:input_type -> daemon.DebugBundleRequest
+	28, // 26: daemon.DaemonService.GetLogLevel:input_type -> daemon.GetLogLevelRequest
+	30, // 27: daemon.DaemonService.SetLogLevel:input_type -> daemon.SetLogLevelRequest
+	2,  // 28: daemon.DaemonService.Login:output_type -> daemon.LoginResponse
+	4,  // 29: daemon.DaemonService.WaitSSOLogin:output_type -> daemon.WaitSSOLoginResponse
+	6,  // 30: daemon.DaemonService.Up:output_type -> daemon.UpResponse
+	8,  // 31: daemon.DaemonService.Status:output_type -> daemon.StatusResponse
+	10, // 32: daemon.DaemonService.Down:output_type -> daemon.DownResponse
+	12, // 33: daemon.DaemonService.GetConfig:output_type -> daemon.GetConfigResponse
+	21, // 34: daemon.DaemonService.ListRoutes:output_type -> daemon.ListRoutesResponse
+	23, // 35: daemon.DaemonService.SelectRoutes:output_type -> daemon.SelectRoutesResponse
+	23, // 36: daemon.DaemonService.DeselectRoutes:output_type -> daemon.SelectRoutesResponse
+	27, // 37: daemon.DaemonService.DebugBundle:output_type -> daemon.DebugBundleResponse
+	29, // 38: daemon.DaemonService.GetLogLevel:output_type -> daemon.GetLogLevelResponse
+	31, // 39: daemon.DaemonService.SetLogLevel:output_type -> daemon.SetLogLevelResponse
+	28, // [28:40] is the sub-list for method output_type
+	16, // [16:28] is the sub-list for method input_type
+	16, // [16:16] is the sub-list for extension type_name
+	16, // [16:16] is the sub-list for extension extendee
+	0,  // [0:16] is the sub-list for field type_name
 }
 
 func init() { file_daemon_proto_init() }
@@ -3434,114 +2938,6 @@ func file_daemon_proto_init() {
 				return nil
 			}
 		}
-		file_daemon_proto_msgTypes[31].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*State); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[32].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*ListStatesRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[33].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*ListStatesResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[34].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*CleanStateRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[35].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*CleanStateResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[36].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*DeleteStateRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[37].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*DeleteStateResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[38].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*SetNetworkMapPersistenceRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_daemon_proto_msgTypes[39].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*SetNetworkMapPersistenceResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
 	}
 	file_daemon_proto_msgTypes[0].OneofWrappers = []interface{}{}
 	type x struct{}
@@ -3550,7 +2946,7 @@ func file_daemon_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: file_daemon_proto_rawDesc,
 			NumEnums:      1,
-			NumMessages:   41,
+			NumMessages:   32,
 			NumExtensions: 0,
 			NumServices:   1,
 		},

client/proto/daemon.proto

@@ -45,20 +45,7 @@ service DaemonService {
 
   // SetLogLevel sets the log level of the daemon
   rpc SetLogLevel(SetLogLevelRequest) returns (SetLogLevelResponse) {}
-
-  // List all states
-  rpc ListStates(ListStatesRequest) returns (ListStatesResponse) {}
-
-  // Clean specific state or all states
-  rpc CleanState(CleanStateRequest) returns (CleanStateResponse) {}
-
-  // Delete specific state or all states
-  rpc DeleteState(DeleteStateRequest) returns (DeleteStateResponse) {}
-
-  // SetNetworkMapPersistence enables or disables network map persistence
-  rpc SetNetworkMapPersistence(SetNetworkMapPersistenceRequest) returns (SetNetworkMapPersistenceResponse) {}
-}
-
+};
 
 message LoginRequest {
   // setupKey wiretrustee setup key.
@@ -306,46 +293,4 @@ message SetLogLevelRequest {
 }
 
 message SetLogLevelResponse {
-}
-
-// State represents a daemon state entry
-message State {
-  string name = 1;
-}
-
-// ListStatesRequest is empty as it requires no parameters
-message ListStatesRequest {}
-
-// ListStatesResponse contains a list of states
-message ListStatesResponse {
-  repeated State states = 1;
-}
-
-// CleanStateRequest for cleaning states
-message CleanStateRequest {
-  string state_name = 1;
-  bool all = 2;
-}
-
-// CleanStateResponse contains the result of the clean operation
-message CleanStateResponse {
-  int32 cleaned_states = 1;
-}
-
-// DeleteStateRequest for deleting states
-message DeleteStateRequest {
-  string state_name = 1;
-  bool all = 2;
-}
-
-// DeleteStateResponse contains the result of the delete operation
-message DeleteStateResponse {
-  int32 deleted_states = 1;
-}
-
-
-message SetNetworkMapPersistenceRequest {
-  bool enabled = 1;
-}
-
-message SetNetworkMapPersistenceResponse {}
+}

client/proto/daemon_grpc.pb.go

@@ -43,14 +43,6 @@ type DaemonServiceClient interface {
 	GetLogLevel(ctx context.Context, in *GetLogLevelRequest, opts ...grpc.CallOption) (*GetLogLevelResponse, error)
 	// SetLogLevel sets the log level of the daemon
 	SetLogLevel(ctx context.Context, in *SetLogLevelRequest, opts ...grpc.CallOption) (*SetLogLevelResponse, error)
-	// List all states
-	ListStates(ctx context.Context, in *ListStatesRequest, opts ...grpc.CallOption) (*ListStatesResponse, error)
-	// Clean specific state or all states
-	CleanState(ctx context.Context, in *CleanStateRequest, opts ...grpc.CallOption) (*CleanStateResponse, error)
-	// Delete specific state or all states
-	DeleteState(ctx context.Context, in *DeleteStateRequest, opts ...grpc.CallOption) (*DeleteStateResponse, error)
-	// SetNetworkMapPersistence enables or disables network map persistence
-	SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error)
 }
 
 type daemonServiceClient struct {
@@ -169,42 +161,6 @@ func (c *daemonServiceClient) SetLogLevel(ctx context.Context, in *SetLogLevelRe
 	return out, nil
 }
 
-func (c *daemonServiceClient) ListStates(ctx context.Context, in *ListStatesRequest, opts ...grpc.CallOption) (*ListStatesResponse, error) {
-	out := new(ListStatesResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/ListStates", in, out, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return out, nil
-}
-
-func (c *daemonServiceClient) CleanState(ctx context.Context, in *CleanStateRequest, opts ...grpc.CallOption) (*CleanStateResponse, error) {
-	out := new(CleanStateResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/CleanState", in, out, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return out, nil
-}
-
-func (c *daemonServiceClient) DeleteState(ctx context.Context, in *DeleteStateRequest, opts ...grpc.CallOption) (*DeleteStateResponse, error) {
-	out := new(DeleteStateResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/DeleteState", in, out, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return out, nil
-}
-
-func (c *daemonServiceClient) SetNetworkMapPersistence(ctx context.Context, in *SetNetworkMapPersistenceRequest, opts ...grpc.CallOption) (*SetNetworkMapPersistenceResponse, error) {
-	out := new(SetNetworkMapPersistenceResponse)
-	err := c.cc.Invoke(ctx, "/daemon.DaemonService/SetNetworkMapPersistence", in, out, opts...)
-	if err != nil {
-		return nil, err
-	}
-	return out, nil
-}
-
 // DaemonServiceServer is the server API for DaemonService service.
 // All implementations must embed UnimplementedDaemonServiceServer
 // for forward compatibility
@@ -234,14 +190,6 @@ type DaemonServiceServer interface {
 	GetLogLevel(context.Context, *GetLogLevelRequest) (*GetLogLevelResponse, error)
 	// SetLogLevel sets the log level of the daemon
 	SetLogLevel(context.Context, *SetLogLevelRequest) (*SetLogLevelResponse, error)
-	// List all states
-	ListStates(context.Context, *ListStatesRequest) (*ListStatesResponse, error)
-	// Clean specific state or all states
-	CleanState(context.Context, *CleanStateRequest) (*CleanStateResponse, error)
-	// Delete specific state or all states
-	DeleteState(context.Context, *DeleteStateRequest) (*DeleteStateResponse, error)
-	// SetNetworkMapPersistence enables or disables network map persistence
-	SetNetworkMapPersistence(context.Context, *SetNetworkMapPersistenceRequest) (*SetNetworkMapPersistenceResponse, error)
 	mustEmbedUnimplementedDaemonServiceServer()
 }
 
@@ -285,18 +233,6 @@ func (UnimplementedDaemonServiceServer) GetLogLevel(context.Context, *GetLogLeve
 func (UnimplementedDaemonServiceServer) SetLogLevel(context.Context, *SetLogLevelRequest) (*SetLogLevelResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method SetLogLevel not implemented")
 }
-func (UnimplementedDaemonServiceServer) ListStates(context.Context, *ListStatesRequest) (*ListStatesResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method ListStates not implemented")
-}
-func (UnimplementedDaemonServiceServer) CleanState(context.Context, *CleanStateRequest) (*CleanStateResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method CleanState not implemented")
-}
-func (UnimplementedDaemonServiceServer) DeleteState(context.Context, *DeleteStateRequest) (*DeleteStateResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method DeleteState not implemented")
-}
-func (UnimplementedDaemonServiceServer) SetNetworkMapPersistence(context.Context, *SetNetworkMapPersistenceRequest) (*SetNetworkMapPersistenceResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method SetNetworkMapPersistence not implemented")
-}
 func (UnimplementedDaemonServiceServer) mustEmbedUnimplementedDaemonServiceServer() {}
 
 // UnsafeDaemonServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -526,78 +462,6 @@ func _DaemonService_SetLogLevel_Handler(srv interface{}, ctx context.Context, de
 	return interceptor(ctx, in, info, handler)
 }
 
-func _DaemonService_ListStates_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(ListStatesRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(DaemonServiceServer).ListStates(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: "/daemon.DaemonService/ListStates",
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).ListStates(ctx, req.(*ListStatesRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
-func _DaemonService_CleanState_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(CleanStateRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(DaemonServiceServer).CleanState(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: "/daemon.DaemonService/CleanState",
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).CleanState(ctx, req.(*CleanStateRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
-func _DaemonService_DeleteState_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(DeleteStateRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(DaemonServiceServer).DeleteState(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: "/daemon.DaemonService/DeleteState",
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).DeleteState(ctx, req.(*DeleteStateRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
-func _DaemonService_SetNetworkMapPersistence_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(SetNetworkMapPersistenceRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(DaemonServiceServer).SetNetworkMapPersistence(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: "/daemon.DaemonService/SetNetworkMapPersistence",
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(DaemonServiceServer).SetNetworkMapPersistence(ctx, req.(*SetNetworkMapPersistenceRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
 // DaemonService_ServiceDesc is the grpc.ServiceDesc for DaemonService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -653,22 +517,6 @@ var DaemonService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "SetLogLevel",
 			Handler:    _DaemonService_SetLogLevel_Handler,
 		},
-		{
-			MethodName: "ListStates",
-			Handler:    _DaemonService_ListStates_Handler,
-		},
-		{
-			MethodName: "CleanState",
-			Handler:    _DaemonService_CleanState_Handler,
-		},
-		{
-			MethodName: "DeleteState",
-			Handler:    _DaemonService_DeleteState_Handler,
-		},
-		{
-			MethodName: "SetNetworkMapPersistence",
-			Handler:    _DaemonService_SetNetworkMapPersistence_Handler,
-		},
 	},
 	Streams:  []grpc.StreamDesc{},
 	Metadata: "daemon.proto",

client/server/debug.go

@@ -5,44 +5,32 @@ package server
 import (
 	"archive/zip"
 	"bufio"
-	"bytes"
 	"context"
-	"encoding/json"
-	"errors"
 	"fmt"
 	"io"
-	"io/fs"
 	"net"
 	"net/netip"
 	"os"
-	"path/filepath"
 	"sort"
 	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"google.golang.org/protobuf/encoding/protojson"
 
 	"github.com/netbirdio/netbird/client/anonymize"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
-	"github.com/netbirdio/netbird/client/internal/statemanager"
 	"github.com/netbirdio/netbird/client/proto"
-	mgmProto "github.com/netbirdio/netbird/management/proto"
 )
 
 const readmeContent = `Netbird debug bundle
 This debug bundle contains the following files:
 
 status.txt: Anonymized status information of the NetBird client.
-client.log: Most recent, anonymized client log file of the NetBird client.
-netbird.err: Most recent, anonymized stderr log file of the NetBird client.
-netbird.out: Most recent, anonymized stdout log file of the NetBird client.
+client.log: Most recent, anonymized log file of the NetBird client.
 routes.txt: Anonymized system routes, if --system-info flag was provided.
 interfaces.txt: Anonymized network interface information, if --system-info flag was provided.
 config.txt: Anonymized configuration information of the NetBird client.
-network_map.json: Anonymized network map containing peer configurations, routes, DNS settings, and firewall rules.
-state.json: Anonymized client state dump containing netbird states.
 
 
 Anonymization Process
@@ -62,32 +50,8 @@ Domains
 All domain names (except for the netbird domains) are replaced with randomly generated strings ending in ".domain". Anonymized domains are consistent across all files in the bundle.
 Reoccuring domain names are replaced with the same anonymized domain.
 
-Network Map
-The network_map.json file contains the following anonymized information:
-- Peer configurations (addresses, FQDNs, DNS settings)
-- Remote and offline peer information (allowed IPs, FQDNs)
-- Routes (network ranges, associated domains)
-- DNS configuration (nameservers, domains, custom zones)
-- Firewall rules (peer IPs, source/destination ranges)
-
-SSH keys in the network map are replaced with a placeholder value. All IP addresses and domains in the network map follow the same anonymization rules as described above.
-
-State File
-The state.json file contains anonymized internal state information of the NetBird client, including:
-- DNS settings and configuration
-- Firewall rules
-- Exclusion routes
-- Route selection
-- Other internal states that may be present
-
-The state file follows the same anonymization rules as other files:
-- IP addresses (both individual and CIDR ranges) are anonymized while preserving their structure
-- Domain names are consistently anonymized
-- Technical identifiers and non-sensitive data remain unchanged
-
 Routes
 For anonymized routes, the IP addresses are replaced as described above. The prefix length remains unchanged. Note that for prefixes, the anonymized IP might not be a network address, but the prefix length is still correct.
-
 Network Interfaces
 The interfaces.txt file contains information about network interfaces, including:
 - Interface name
@@ -108,12 +72,6 @@ The config.txt file contains anonymized configuration information of the NetBird
 Other non-sensitive configuration options are included without anonymization.
 `
 
-const (
-	clientLogFile = "client.log"
-	errorLogFile  = "netbird.err"
-	stdoutLogFile = "netbird.out"
-)
-
 // DebugBundle creates a debug bundle and returns the location.
 func (s *Server) DebugBundle(_ context.Context, req *proto.DebugBundleRequest) (resp *proto.DebugBundleResponse, err error) {
 	s.mutex.Lock()
@@ -161,27 +119,19 @@ func (s *Server) createArchive(bundlePath *os.File, req *proto.DebugBundleReques
 	seedFromStatus(anonymizer, &status)
 
 	if err := s.addConfig(req, anonymizer, archive); err != nil {
-		log.Errorf("Failed to add config to debug bundle: %v", err)
+		return fmt.Errorf("add config: %w", err)
 	}
 
 	if req.GetSystemInfo() {
 		if err := s.addRoutes(req, anonymizer, archive); err != nil {
-			log.Errorf("Failed to add routes to debug bundle: %v", err)
+			return fmt.Errorf("add routes: %w", err)
 		}
 
 		if err := s.addInterfaces(req, anonymizer, archive); err != nil {
-			log.Errorf("Failed to add interfaces to debug bundle: %v", err)
+			return fmt.Errorf("add interfaces: %w", err)
 		}
 	}
 
-	if err := s.addNetworkMap(req, anonymizer, archive); err != nil {
-		return fmt.Errorf("add network map: %w", err)
-	}
-
-	if err := s.addStateFile(req, anonymizer, archive); err != nil {
-		log.Errorf("Failed to add state file to debug bundle: %v", err)
-	}
-
 	if err := s.addLogfile(req, anonymizer, archive); err != nil {
 		return fmt.Errorf("add log file: %w", err)
 	}
@@ -270,16 +220,15 @@ func (s *Server) addCommonConfigFields(configContent *strings.Builder) {
 }
 
 func (s *Server) addRoutes(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
-	routes, err := systemops.GetRoutesFromTable()
-	if err != nil {
-		return fmt.Errorf("get routes: %w", err)
-	}
-
-	// TODO: get routes including nexthop
-	routesContent := formatRoutes(routes, req.GetAnonymize(), anonymizer)
-	routesReader := strings.NewReader(routesContent)
-	if err := addFileToZip(archive, routesReader, "routes.txt"); err != nil {
-		return fmt.Errorf("add routes file to zip: %w", err)
+	if routes, err := systemops.GetRoutesFromTable(); err != nil {
+		log.Errorf("Failed to get routes: %v", err)
+	} else {
+		// TODO: get routes including nexthop
+		routesContent := formatRoutes(routes, req.GetAnonymize(), anonymizer)
+		routesReader := strings.NewReader(routesContent)
+		if err := addFileToZip(archive, routesReader, "routes.txt"); err != nil {
+			return fmt.Errorf("add routes file to zip: %w", err)
+		}
 	}
 	return nil
 }
@@ -299,106 +248,14 @@ func (s *Server) addInterfaces(req *proto.DebugBundleRequest, anonymizer *anonym
 	return nil
 }
 
-func (s *Server) addNetworkMap(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
-	networkMap, err := s.getLatestNetworkMap()
+func (s *Server) addLogfile(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) (err error) {
+	logFile, err := os.Open(s.logFile)
 	if err != nil {
-		// Skip if network map is not available, but log it
-		log.Debugf("skipping empty network map in debug bundle: %v", err)
-		return nil
-	}
-
-	if req.GetAnonymize() {
-		if err := anonymizeNetworkMap(networkMap, anonymizer); err != nil {
-			return fmt.Errorf("anonymize network map: %w", err)
-		}
-	}
-
-	options := protojson.MarshalOptions{
-		EmitUnpopulated: true,
-		UseProtoNames:   true,
-		Indent:          "  ",
-		AllowPartial:    true,
-	}
-
-	jsonBytes, err := options.Marshal(networkMap)
-	if err != nil {
-		return fmt.Errorf("generate json: %w", err)
-	}
-
-	if err := addFileToZip(archive, bytes.NewReader(jsonBytes), "network_map.json"); err != nil {
-		return fmt.Errorf("add network map to zip: %w", err)
-	}
-
-	return nil
-}
-
-func (s *Server) addStateFile(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
-	path := statemanager.GetDefaultStatePath()
-	if path == "" {
-		return nil
-	}
-
-	data, err := os.ReadFile(path)
-	if err != nil {
-		if errors.Is(err, fs.ErrNotExist) {
-			return nil
-		}
-		return fmt.Errorf("read state file: %w", err)
-	}
-
-	if req.GetAnonymize() {
-		var rawStates map[string]json.RawMessage
-		if err := json.Unmarshal(data, &rawStates); err != nil {
-			return fmt.Errorf("unmarshal states: %w", err)
-		}
-
-		if err := anonymizeStateFile(&rawStates, anonymizer); err != nil {
-			return fmt.Errorf("anonymize state file: %w", err)
-		}
-
-		bs, err := json.MarshalIndent(rawStates, "", "  ")
-		if err != nil {
-			return fmt.Errorf("marshal states: %w", err)
-		}
-		data = bs
-	}
-
-	if err := addFileToZip(archive, bytes.NewReader(data), "state.json"); err != nil {
-		return fmt.Errorf("add state file to zip: %w", err)
-	}
-
-	return nil
-}
-
-func (s *Server) addLogfile(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
-	logDir := filepath.Dir(s.logFile)
-
-	if err := s.addSingleLogfile(s.logFile, clientLogFile, req, anonymizer, archive); err != nil {
-		return fmt.Errorf("add client log file to zip: %w", err)
-	}
-
-	errLogPath := filepath.Join(logDir, errorLogFile)
-	if err := s.addSingleLogfile(errLogPath, errorLogFile, req, anonymizer, archive); err != nil {
-		log.Warnf("Failed to add %s to zip: %v", errorLogFile, err)
-	}
-
-	stdoutLogPath := filepath.Join(logDir, stdoutLogFile)
-	if err := s.addSingleLogfile(stdoutLogPath, stdoutLogFile, req, anonymizer, archive); err != nil {
-		log.Warnf("Failed to add %s to zip: %v", stdoutLogFile, err)
-	}
-
-	return nil
-}
-
-// addSingleLogfile adds a single log file to the archive
-func (s *Server) addSingleLogfile(logPath, targetName string, req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
-	logFile, err := os.Open(logPath)
-	if err != nil {
-		return fmt.Errorf("open log file %s: %w", targetName, err)
+		return fmt.Errorf("open log file: %w", err)
 	}
 	defer func() {
 		if err := logFile.Close(); err != nil {
-			log.Errorf("Failed to close log file %s: %v", targetName, err)
+			log.Errorf("Failed to close original log file: %v", err)
 		}
 	}()
 
@@ -407,55 +264,45 @@ func (s *Server) addSingleLogfile(logPath, targetName string, req *proto.DebugBu
 		var writer *io.PipeWriter
 		logReader, writer = io.Pipe()
 
-		go anonymizeLog(logFile, writer, anonymizer)
+		go s.anonymize(logFile, writer, anonymizer)
 	} else {
 		logReader = logFile
 	}
-
-	if err := addFileToZip(archive, logReader, targetName); err != nil {
-		return fmt.Errorf("add %s to zip: %w", targetName, err)
+	if err := addFileToZip(archive, logReader, "client.log"); err != nil {
+		return fmt.Errorf("add log file to zip: %w", err)
 	}
 
 	return nil
 }
 
-// getLatestNetworkMap returns the latest network map from the engine if network map persistence is enabled
-func (s *Server) getLatestNetworkMap() (*mgmProto.NetworkMap, error) {
-	if s.connectClient == nil {
-		return nil, errors.New("connect client is not initialized")
-	}
+func (s *Server) anonymize(reader io.Reader, writer *io.PipeWriter, anonymizer *anonymize.Anonymizer) {
+	defer func() {
+		// always nil
+		_ = writer.Close()
+	}()
 
-	engine := s.connectClient.Engine()
-	if engine == nil {
-		return nil, errors.New("engine is not initialized")
+	scanner := bufio.NewScanner(reader)
+	for scanner.Scan() {
+		line := anonymizer.AnonymizeString(scanner.Text())
+		if _, err := writer.Write([]byte(line + "\n")); err != nil {
+			writer.CloseWithError(fmt.Errorf("anonymize write: %w", err))
+			return
+		}
 	}
-
-	networkMap, err := engine.GetLatestNetworkMap()
-	if err != nil {
-		return nil, fmt.Errorf("get latest network map: %w", err)
+	if err := scanner.Err(); err != nil {
+		writer.CloseWithError(fmt.Errorf("anonymize scan: %w", err))
+		return
 	}
-
-	if networkMap == nil {
-		return nil, errors.New("network map is not available")
-	}
-
-	return networkMap, nil
 }
 
 // GetLogLevel gets the current logging level for the server.
 func (s *Server) GetLogLevel(_ context.Context, _ *proto.GetLogLevelRequest) (*proto.GetLogLevelResponse, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	level := ParseLogLevel(log.GetLevel().String())
 	return &proto.GetLogLevelResponse{Level: level}, nil
 }
 
 // SetLogLevel sets the logging level for the server.
 func (s *Server) SetLogLevel(_ context.Context, req *proto.SetLogLevelRequest) (*proto.SetLogLevelResponse, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	level, err := log.ParseLevel(req.Level.String())
 	if err != nil {
 		return nil, fmt.Errorf("invalid log level: %w", err)
@@ -466,20 +313,6 @@ func (s *Server) SetLogLevel(_ context.Context, req *proto.SetLogLevelRequest) (
 	return &proto.SetLogLevelResponse{}, nil
 }
 
-// SetNetworkMapPersistence sets the network map persistence for the server.
-func (s *Server) SetNetworkMapPersistence(_ context.Context, req *proto.SetNetworkMapPersistenceRequest) (*proto.SetNetworkMapPersistenceResponse, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
-	enabled := req.GetEnabled()
-	s.persistNetworkMap = enabled
-	if s.connectClient != nil {
-		s.connectClient.SetNetworkMapPersistence(enabled)
-	}
-
-	return &proto.SetNetworkMapPersistenceResponse{}, nil
-}
-
 func addFileToZip(archive *zip.Writer, reader io.Reader, filename string) error {
 	header := &zip.FileHeader{
 		Name:     filename,
@@ -625,26 +458,6 @@ func formatInterfaces(interfaces []net.Interface, anonymize bool, anonymizer *an
 	return builder.String()
 }
 
-func anonymizeLog(reader io.Reader, writer *io.PipeWriter, anonymizer *anonymize.Anonymizer) {
-	defer func() {
-		// always nil
-		_ = writer.Close()
-	}()
-
-	scanner := bufio.NewScanner(reader)
-	for scanner.Scan() {
-		line := anonymizer.AnonymizeString(scanner.Text())
-		if _, err := writer.Write([]byte(line + "\n")); err != nil {
-			writer.CloseWithError(fmt.Errorf("anonymize write: %w", err))
-			return
-		}
-	}
-	if err := scanner.Err(); err != nil {
-		writer.CloseWithError(fmt.Errorf("anonymize scan: %w", err))
-		return
-	}
-}
-
 func anonymizeNATExternalIPs(ips []string, anonymizer *anonymize.Anonymizer) []string {
 	anonymizedIPs := make([]string, len(ips))
 	for i, ip := range ips {
@@ -671,248 +484,3 @@ func anonymizeNATExternalIPs(ips []string, anonymizer *anonymize.Anonymizer) []s
 	}
 	return anonymizedIPs
 }
-
-func anonymizeNetworkMap(networkMap *mgmProto.NetworkMap, anonymizer *anonymize.Anonymizer) error {
-	if networkMap.PeerConfig != nil {
-		anonymizePeerConfig(networkMap.PeerConfig, anonymizer)
-	}
-
-	for _, peer := range networkMap.RemotePeers {
-		anonymizeRemotePeer(peer, anonymizer)
-	}
-
-	for _, peer := range networkMap.OfflinePeers {
-		anonymizeRemotePeer(peer, anonymizer)
-	}
-
-	for _, r := range networkMap.Routes {
-		anonymizeRoute(r, anonymizer)
-	}
-
-	if networkMap.DNSConfig != nil {
-		anonymizeDNSConfig(networkMap.DNSConfig, anonymizer)
-	}
-
-	for _, rule := range networkMap.FirewallRules {
-		anonymizeFirewallRule(rule, anonymizer)
-	}
-
-	for _, rule := range networkMap.RoutesFirewallRules {
-		anonymizeRouteFirewallRule(rule, anonymizer)
-	}
-
-	return nil
-}
-
-func anonymizePeerConfig(config *mgmProto.PeerConfig, anonymizer *anonymize.Anonymizer) {
-	if config == nil {
-		return
-	}
-
-	if addr, err := netip.ParseAddr(config.Address); err == nil {
-		config.Address = anonymizer.AnonymizeIP(addr).String()
-	}
-
-	if config.SshConfig != nil && len(config.SshConfig.SshPubKey) > 0 {
-		config.SshConfig.SshPubKey = []byte("ssh-placeholder-key")
-	}
-
-	config.Dns = anonymizer.AnonymizeString(config.Dns)
-	config.Fqdn = anonymizer.AnonymizeDomain(config.Fqdn)
-}
-
-func anonymizeRemotePeer(peer *mgmProto.RemotePeerConfig, anonymizer *anonymize.Anonymizer) {
-	if peer == nil {
-		return
-	}
-
-	for i, ip := range peer.AllowedIps {
-		// Try to parse as prefix first (CIDR)
-		if prefix, err := netip.ParsePrefix(ip); err == nil {
-			anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-			peer.AllowedIps[i] = fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-		} else if addr, err := netip.ParseAddr(ip); err == nil {
-			peer.AllowedIps[i] = anonymizer.AnonymizeIP(addr).String()
-		}
-	}
-
-	peer.Fqdn = anonymizer.AnonymizeDomain(peer.Fqdn)
-
-	if peer.SshConfig != nil && len(peer.SshConfig.SshPubKey) > 0 {
-		peer.SshConfig.SshPubKey = []byte("ssh-placeholder-key")
-	}
-}
-
-func anonymizeRoute(route *mgmProto.Route, anonymizer *anonymize.Anonymizer) {
-	if route == nil {
-		return
-	}
-
-	if prefix, err := netip.ParsePrefix(route.Network); err == nil {
-		anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-		route.Network = fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-	}
-
-	for i, domain := range route.Domains {
-		route.Domains[i] = anonymizer.AnonymizeDomain(domain)
-	}
-
-	route.NetID = anonymizer.AnonymizeString(route.NetID)
-}
-
-func anonymizeDNSConfig(config *mgmProto.DNSConfig, anonymizer *anonymize.Anonymizer) {
-	if config == nil {
-		return
-	}
-
-	anonymizeNameServerGroups(config.NameServerGroups, anonymizer)
-	anonymizeCustomZones(config.CustomZones, anonymizer)
-}
-
-func anonymizeNameServerGroups(groups []*mgmProto.NameServerGroup, anonymizer *anonymize.Anonymizer) {
-	for _, group := range groups {
-		anonymizeServers(group.NameServers, anonymizer)
-		anonymizeDomains(group.Domains, anonymizer)
-	}
-}
-
-func anonymizeServers(servers []*mgmProto.NameServer, anonymizer *anonymize.Anonymizer) {
-	for _, server := range servers {
-		if addr, err := netip.ParseAddr(server.IP); err == nil {
-			server.IP = anonymizer.AnonymizeIP(addr).String()
-		}
-	}
-}
-
-func anonymizeDomains(domains []string, anonymizer *anonymize.Anonymizer) {
-	for i, domain := range domains {
-		domains[i] = anonymizer.AnonymizeDomain(domain)
-	}
-}
-
-func anonymizeCustomZones(zones []*mgmProto.CustomZone, anonymizer *anonymize.Anonymizer) {
-	for _, zone := range zones {
-		zone.Domain = anonymizer.AnonymizeDomain(zone.Domain)
-		anonymizeRecords(zone.Records, anonymizer)
-	}
-}
-
-func anonymizeRecords(records []*mgmProto.SimpleRecord, anonymizer *anonymize.Anonymizer) {
-	for _, record := range records {
-		record.Name = anonymizer.AnonymizeDomain(record.Name)
-		anonymizeRData(record, anonymizer)
-	}
-}
-
-func anonymizeRData(record *mgmProto.SimpleRecord, anonymizer *anonymize.Anonymizer) {
-	switch record.Type {
-	case 1, 28: // A or AAAA record
-		if addr, err := netip.ParseAddr(record.RData); err == nil {
-			record.RData = anonymizer.AnonymizeIP(addr).String()
-		}
-	default:
-		record.RData = anonymizer.AnonymizeString(record.RData)
-	}
-}
-
-func anonymizeFirewallRule(rule *mgmProto.FirewallRule, anonymizer *anonymize.Anonymizer) {
-	if rule == nil {
-		return
-	}
-
-	if addr, err := netip.ParseAddr(rule.PeerIP); err == nil {
-		rule.PeerIP = anonymizer.AnonymizeIP(addr).String()
-	}
-}
-
-func anonymizeRouteFirewallRule(rule *mgmProto.RouteFirewallRule, anonymizer *anonymize.Anonymizer) {
-	if rule == nil {
-		return
-	}
-
-	for i, sourceRange := range rule.SourceRanges {
-		if prefix, err := netip.ParsePrefix(sourceRange); err == nil {
-			anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-			rule.SourceRanges[i] = fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-		}
-	}
-
-	if prefix, err := netip.ParsePrefix(rule.Destination); err == nil {
-		anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-		rule.Destination = fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-	}
-}
-
-func anonymizeStateFile(rawStates *map[string]json.RawMessage, anonymizer *anonymize.Anonymizer) error {
-	for name, rawState := range *rawStates {
-		if string(rawState) == "null" {
-			continue
-		}
-
-		var state map[string]any
-		if err := json.Unmarshal(rawState, &state); err != nil {
-			return fmt.Errorf("unmarshal state %s: %w", name, err)
-		}
-
-		state = anonymizeValue(state, anonymizer).(map[string]any)
-
-		bs, err := json.Marshal(state)
-		if err != nil {
-			return fmt.Errorf("marshal state %s: %w", name, err)
-		}
-
-		(*rawStates)[name] = bs
-	}
-
-	return nil
-}
-
-func anonymizeValue(value any, anonymizer *anonymize.Anonymizer) any {
-	switch v := value.(type) {
-	case string:
-		return anonymizeString(v, anonymizer)
-	case map[string]any:
-		return anonymizeMap(v, anonymizer)
-	case []any:
-		return anonymizeSlice(v, anonymizer)
-	}
-	return value
-}
-
-func anonymizeString(v string, anonymizer *anonymize.Anonymizer) string {
-	if prefix, err := netip.ParsePrefix(v); err == nil {
-		anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-		return fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-	}
-	if ip, err := netip.ParseAddr(v); err == nil {
-		return anonymizer.AnonymizeIP(ip).String()
-	}
-	return anonymizer.AnonymizeString(v)
-}
-
-func anonymizeMap(v map[string]any, anonymizer *anonymize.Anonymizer) map[string]any {
-	result := make(map[string]any, len(v))
-	for key, val := range v {
-		newKey := anonymizeMapKey(key, anonymizer)
-		result[newKey] = anonymizeValue(val, anonymizer)
-	}
-	return result
-}
-
-func anonymizeMapKey(key string, anonymizer *anonymize.Anonymizer) string {
-	if prefix, err := netip.ParsePrefix(key); err == nil {
-		anonIP := anonymizer.AnonymizeIP(prefix.Addr())
-		return fmt.Sprintf("%s/%d", anonIP, prefix.Bits())
-	}
-	if ip, err := netip.ParseAddr(key); err == nil {
-		return anonymizer.AnonymizeIP(ip).String()
-	}
-	return key
-}
-
-func anonymizeSlice(v []any, anonymizer *anonymize.Anonymizer) []any {
-	for i, val := range v {
-		v[i] = anonymizeValue(val, anonymizer)
-	}
-	return v
-}

client/server/debug_test.go

@@ -1,430 +0,0 @@
-package server
-
-import (
-	"encoding/json"
-	"net"
-	"strings"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/anonymize"
-	mgmProto "github.com/netbirdio/netbird/management/proto"
-)
-
-func TestAnonymizeStateFile(t *testing.T) {
-	testState := map[string]json.RawMessage{
-		"null_state": json.RawMessage("null"),
-		"test_state": mustMarshal(map[string]any{
-			// Test simple fields
-			"public_ip":      "203.0.113.1",
-			"private_ip":     "192.168.1.1",
-			"protected_ip":   "100.64.0.1",
-			"well_known_ip":  "8.8.8.8",
-			"ipv6_addr":      "2001:db8::1",
-			"private_ipv6":   "fd00::1",
-			"domain":         "test.example.com",
-			"uri":            "stun:stun.example.com:3478",
-			"uri_with_ip":    "turn:203.0.113.1:3478",
-			"netbird_domain": "device.netbird.cloud",
-
-			// Test CIDR ranges
-			"public_cidr":       "203.0.113.0/24",
-			"private_cidr":      "192.168.0.0/16",
-			"protected_cidr":    "100.64.0.0/10",
-			"ipv6_cidr":         "2001:db8::/32",
-			"private_ipv6_cidr": "fd00::/8",
-
-			// Test nested structures
-			"nested": map[string]any{
-				"ip":     "203.0.113.2",
-				"domain": "nested.example.com",
-				"more_nest": map[string]any{
-					"ip":     "203.0.113.3",
-					"domain": "deep.example.com",
-				},
-			},
-
-			// Test arrays
-			"string_array": []any{
-				"203.0.113.4",
-				"test1.example.com",
-				"test2.example.com",
-			},
-			"object_array": []any{
-				map[string]any{
-					"ip":     "203.0.113.5",
-					"domain": "array1.example.com",
-				},
-				map[string]any{
-					"ip":     "203.0.113.6",
-					"domain": "array2.example.com",
-				},
-			},
-
-			// Test multiple occurrences of same value
-			"duplicate_ip":     "203.0.113.1",      // Same as public_ip
-			"duplicate_domain": "test.example.com", // Same as domain
-
-			// Test URIs with various schemes
-			"stun_uri":  "stun:stun.example.com:3478",
-			"turns_uri": "turns:turns.example.com:5349",
-			"http_uri":  "http://web.example.com:80",
-			"https_uri": "https://secure.example.com:443",
-
-			// Test strings that might look like IPs but aren't
-			"not_ip":         "300.300.300.300",
-			"partial_ip":     "192.168",
-			"ip_like_string": "1234.5678",
-
-			// Test mixed content strings
-			"mixed_content": "Server at 203.0.113.1 (test.example.com) on port 80",
-
-			// Test empty and special values
-			"empty_string":  "",
-			"null_value":    nil,
-			"numeric_value": 42,
-			"boolean_value": true,
-		}),
-		"route_state": mustMarshal(map[string]any{
-			"routes": []any{
-				map[string]any{
-					"network": "203.0.113.0/24",
-					"gateway": "203.0.113.1",
-					"domains": []any{
-						"route1.example.com",
-						"route2.example.com",
-					},
-				},
-				map[string]any{
-					"network": "2001:db8::/32",
-					"gateway": "2001:db8::1",
-					"domains": []any{
-						"route3.example.com",
-						"route4.example.com",
-					},
-				},
-			},
-			// Test map with IP/CIDR keys
-			"refCountMap": map[string]any{
-				"203.0.113.1/32": map[string]any{
-					"Count": 1,
-					"Out": map[string]any{
-						"IP": "192.168.0.1",
-						"Intf": map[string]any{
-							"Name":  "eth0",
-							"Index": 1,
-						},
-					},
-				},
-				"2001:db8::1/128": map[string]any{
-					"Count": 1,
-					"Out": map[string]any{
-						"IP": "fe80::1",
-						"Intf": map[string]any{
-							"Name":  "eth0",
-							"Index": 1,
-						},
-					},
-				},
-				"10.0.0.1/32": map[string]any{ // private IP should remain unchanged
-					"Count": 1,
-					"Out": map[string]any{
-						"IP": "192.168.0.1",
-					},
-				},
-			},
-		}),
-	}
-
-	anonymizer := anonymize.NewAnonymizer(anonymize.DefaultAddresses())
-
-	// Pre-seed the domains we need to verify in the test assertions
-	anonymizer.AnonymizeDomain("test.example.com")
-	anonymizer.AnonymizeDomain("nested.example.com")
-	anonymizer.AnonymizeDomain("deep.example.com")
-	anonymizer.AnonymizeDomain("array1.example.com")
-
-	err := anonymizeStateFile(&testState, anonymizer)
-	require.NoError(t, err)
-
-	// Helper function to unmarshal and get nested values
-	var state map[string]any
-	err = json.Unmarshal(testState["test_state"], &state)
-	require.NoError(t, err)
-
-	// Test null state remains unchanged
-	require.Equal(t, "null", string(testState["null_state"]))
-
-	// Basic assertions
-	assert.NotEqual(t, "203.0.113.1", state["public_ip"])
-	assert.Equal(t, "192.168.1.1", state["private_ip"])  // Private IP unchanged
-	assert.Equal(t, "100.64.0.1", state["protected_ip"]) // Protected IP unchanged
-	assert.Equal(t, "8.8.8.8", state["well_known_ip"])   // Well-known IP unchanged
-	assert.NotEqual(t, "2001:db8::1", state["ipv6_addr"])
-	assert.Equal(t, "fd00::1", state["private_ipv6"]) // Private IPv6 unchanged
-	assert.NotEqual(t, "test.example.com", state["domain"])
-	assert.True(t, strings.HasSuffix(state["domain"].(string), ".domain"))
-	assert.Equal(t, "device.netbird.cloud", state["netbird_domain"]) // Netbird domain unchanged
-
-	// CIDR ranges
-	assert.NotEqual(t, "203.0.113.0/24", state["public_cidr"])
-	assert.Contains(t, state["public_cidr"], "/24")           // Prefix preserved
-	assert.Equal(t, "192.168.0.0/16", state["private_cidr"])  // Private CIDR unchanged
-	assert.Equal(t, "100.64.0.0/10", state["protected_cidr"]) // Protected CIDR unchanged
-	assert.NotEqual(t, "2001:db8::/32", state["ipv6_cidr"])
-	assert.Contains(t, state["ipv6_cidr"], "/32") // IPv6 prefix preserved
-
-	// Nested structures
-	nested := state["nested"].(map[string]any)
-	assert.NotEqual(t, "203.0.113.2", nested["ip"])
-	assert.NotEqual(t, "nested.example.com", nested["domain"])
-	moreNest := nested["more_nest"].(map[string]any)
-	assert.NotEqual(t, "203.0.113.3", moreNest["ip"])
-	assert.NotEqual(t, "deep.example.com", moreNest["domain"])
-
-	// Arrays
-	strArray := state["string_array"].([]any)
-	assert.NotEqual(t, "203.0.113.4", strArray[0])
-	assert.NotEqual(t, "test1.example.com", strArray[1])
-	assert.True(t, strings.HasSuffix(strArray[1].(string), ".domain"))
-
-	objArray := state["object_array"].([]any)
-	firstObj := objArray[0].(map[string]any)
-	assert.NotEqual(t, "203.0.113.5", firstObj["ip"])
-	assert.NotEqual(t, "array1.example.com", firstObj["domain"])
-
-	// Duplicate values should be anonymized consistently
-	assert.Equal(t, state["public_ip"], state["duplicate_ip"])
-	assert.Equal(t, state["domain"], state["duplicate_domain"])
-
-	// URIs
-	assert.NotContains(t, state["stun_uri"], "stun.example.com")
-	assert.NotContains(t, state["turns_uri"], "turns.example.com")
-	assert.NotContains(t, state["http_uri"], "web.example.com")
-	assert.NotContains(t, state["https_uri"], "secure.example.com")
-
-	// Non-IP strings should remain unchanged
-	assert.Equal(t, "300.300.300.300", state["not_ip"])
-	assert.Equal(t, "192.168", state["partial_ip"])
-	assert.Equal(t, "1234.5678", state["ip_like_string"])
-
-	// Mixed content should have IPs and domains replaced
-	mixedContent := state["mixed_content"].(string)
-	assert.NotContains(t, mixedContent, "203.0.113.1")
-	assert.NotContains(t, mixedContent, "test.example.com")
-	assert.Contains(t, mixedContent, "Server at ")
-	assert.Contains(t, mixedContent, " on port 80")
-
-	// Special values should remain unchanged
-	assert.Equal(t, "", state["empty_string"])
-	assert.Nil(t, state["null_value"])
-	assert.Equal(t, float64(42), state["numeric_value"])
-	assert.Equal(t, true, state["boolean_value"])
-
-	// Check route state
-	var routeState map[string]any
-	err = json.Unmarshal(testState["route_state"], &routeState)
-	require.NoError(t, err)
-
-	routes := routeState["routes"].([]any)
-	route1 := routes[0].(map[string]any)
-	assert.NotEqual(t, "203.0.113.0/24", route1["network"])
-	assert.Contains(t, route1["network"], "/24")
-	assert.NotEqual(t, "203.0.113.1", route1["gateway"])
-	domains := route1["domains"].([]any)
-	assert.True(t, strings.HasSuffix(domains[0].(string), ".domain"))
-	assert.True(t, strings.HasSuffix(domains[1].(string), ".domain"))
-
-	// Check map keys are anonymized
-	refCountMap := routeState["refCountMap"].(map[string]any)
-	hasPublicIPKey := false
-	hasIPv6Key := false
-	hasPrivateIPKey := false
-	for key := range refCountMap {
-		if strings.Contains(key, "203.0.113.1") {
-			hasPublicIPKey = true
-		}
-		if strings.Contains(key, "2001:db8::1") {
-			hasIPv6Key = true
-		}
-		if key == "10.0.0.1/32" {
-			hasPrivateIPKey = true
-		}
-	}
-	assert.False(t, hasPublicIPKey, "public IP in key should be anonymized")
-	assert.False(t, hasIPv6Key, "IPv6 in key should be anonymized")
-	assert.True(t, hasPrivateIPKey, "private IP in key should remain unchanged")
-}
-
-func mustMarshal(v any) json.RawMessage {
-	data, err := json.Marshal(v)
-	if err != nil {
-		panic(err)
-	}
-	return data
-}
-
-func TestAnonymizeNetworkMap(t *testing.T) {
-	networkMap := &mgmProto.NetworkMap{
-		PeerConfig: &mgmProto.PeerConfig{
-			Address: "203.0.113.5",
-			Dns:     "1.2.3.4",
-			Fqdn:    "peer1.corp.example.com",
-			SshConfig: &mgmProto.SSHConfig{
-				SshPubKey: []byte("ssh-rsa AAAAB3NzaC1..."),
-			},
-		},
-		RemotePeers: []*mgmProto.RemotePeerConfig{
-			{
-				AllowedIps: []string{
-					"203.0.113.1/32",
-					"2001:db8:1234::1/128",
-					"192.168.1.1/32",
-					"100.64.0.1/32",
-					"10.0.0.1/32",
-				},
-				Fqdn: "peer2.corp.example.com",
-				SshConfig: &mgmProto.SSHConfig{
-					SshPubKey: []byte("ssh-rsa AAAAB3NzaC2..."),
-				},
-			},
-		},
-		Routes: []*mgmProto.Route{
-			{
-				Network: "197.51.100.0/24",
-				Domains: []string{"prod.example.com", "staging.example.com"},
-				NetID:   "net-123abc",
-			},
-		},
-		DNSConfig: &mgmProto.DNSConfig{
-			NameServerGroups: []*mgmProto.NameServerGroup{
-				{
-					NameServers: []*mgmProto.NameServer{
-						{IP: "8.8.8.8"},
-						{IP: "1.1.1.1"},
-						{IP: "203.0.113.53"},
-					},
-					Domains: []string{"example.com", "internal.example.com"},
-				},
-			},
-			CustomZones: []*mgmProto.CustomZone{
-				{
-					Domain: "custom.example.com",
-					Records: []*mgmProto.SimpleRecord{
-						{
-							Name:  "www.custom.example.com",
-							Type:  1,
-							RData: "203.0.113.10",
-						},
-						{
-							Name:  "internal.custom.example.com",
-							Type:  1,
-							RData: "192.168.1.10",
-						},
-					},
-				},
-			},
-		},
-	}
-
-	// Create anonymizer with test addresses
-	anonymizer := anonymize.NewAnonymizer(anonymize.DefaultAddresses())
-
-	// Anonymize the network map
-	err := anonymizeNetworkMap(networkMap, anonymizer)
-	require.NoError(t, err)
-
-	// Test PeerConfig anonymization
-	peerCfg := networkMap.PeerConfig
-	require.NotEqual(t, "203.0.113.5", peerCfg.Address)
-
-	// Verify DNS and FQDN are properly anonymized
-	require.NotEqual(t, "1.2.3.4", peerCfg.Dns)
-	require.NotEqual(t, "peer1.corp.example.com", peerCfg.Fqdn)
-	require.True(t, strings.HasSuffix(peerCfg.Fqdn, ".domain"))
-
-	// Verify SSH key is replaced
-	require.Equal(t, []byte("ssh-placeholder-key"), peerCfg.SshConfig.SshPubKey)
-
-	// Test RemotePeers anonymization
-	remotePeer := networkMap.RemotePeers[0]
-
-	// Verify FQDN is anonymized
-	require.NotEqual(t, "peer2.corp.example.com", remotePeer.Fqdn)
-	require.True(t, strings.HasSuffix(remotePeer.Fqdn, ".domain"))
-
-	// Check that public IPs are anonymized but private IPs are preserved
-	for _, allowedIP := range remotePeer.AllowedIps {
-		ip, _, err := net.ParseCIDR(allowedIP)
-		require.NoError(t, err)
-
-		if ip.IsPrivate() || isInCGNATRange(ip) {
-			require.Contains(t, []string{
-				"192.168.1.1/32",
-				"100.64.0.1/32",
-				"10.0.0.1/32",
-			}, allowedIP)
-		} else {
-			require.NotContains(t, []string{
-				"203.0.113.1/32",
-				"2001:db8:1234::1/128",
-			}, allowedIP)
-		}
-	}
-
-	// Test Routes anonymization
-	route := networkMap.Routes[0]
-	require.NotEqual(t, "197.51.100.0/24", route.Network)
-	for _, domain := range route.Domains {
-		require.True(t, strings.HasSuffix(domain, ".domain"))
-		require.NotContains(t, domain, "example.com")
-	}
-
-	// Test DNS config anonymization
-	dnsConfig := networkMap.DNSConfig
-	nameServerGroup := dnsConfig.NameServerGroups[0]
-
-	// Verify well-known DNS servers are preserved
-	require.Equal(t, "8.8.8.8", nameServerGroup.NameServers[0].IP)
-	require.Equal(t, "1.1.1.1", nameServerGroup.NameServers[1].IP)
-
-	// Verify public DNS server is anonymized
-	require.NotEqual(t, "203.0.113.53", nameServerGroup.NameServers[2].IP)
-
-	// Verify domains are anonymized
-	for _, domain := range nameServerGroup.Domains {
-		require.True(t, strings.HasSuffix(domain, ".domain"))
-		require.NotContains(t, domain, "example.com")
-	}
-
-	// Test CustomZones anonymization
-	customZone := dnsConfig.CustomZones[0]
-	require.True(t, strings.HasSuffix(customZone.Domain, ".domain"))
-	require.NotContains(t, customZone.Domain, "example.com")
-
-	// Verify records are properly anonymized
-	for _, record := range customZone.Records {
-		require.True(t, strings.HasSuffix(record.Name, ".domain"))
-		require.NotContains(t, record.Name, "example.com")
-
-		ip := net.ParseIP(record.RData)
-		if ip != nil {
-			if !ip.IsPrivate() {
-				require.NotEqual(t, "203.0.113.10", record.RData)
-			} else {
-				require.Equal(t, "192.168.1.10", record.RData)
-			}
-		}
-	}
-}
-
-// Helper function to check if IP is in CGNAT range
-func isInCGNATRange(ip net.IP) bool {
-	cgnat := net.IPNet{
-		IP:   net.ParseIP("100.64.0.0"),
-		Mask: net.CIDRMask(10, 32),
-	}
-	return cgnat.Contains(ip)
-}

client/server/server.go

@@ -68,8 +68,6 @@ type Server struct {
 	relayProbe  *internal.Probe
 	wgProbe     *internal.Probe
 	lastProbe   time.Time
-
-	persistNetworkMap bool
 }
 
 type oauthAuthFlow struct {
@@ -198,7 +196,6 @@ func (s *Server) connectWithRetryRuns(ctx context.Context, config *internal.Conf
 	runOperation := func() error {
 		log.Tracef("running client connection")
 		s.connectClient = internal.NewConnectClient(ctx, config, statusRecorder)
-		s.connectClient.SetNetworkMapPersistence(s.persistNetworkMap)
 
 		probes := internal.ProbeHolder{
 			MgmProbe:    s.mgmProbe,

client/server/state.go

@@ -5,112 +5,12 @@ import (
 	"fmt"
 
 	"github.com/hashicorp/go-multierror"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
-	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
-	"github.com/netbirdio/netbird/client/proto"
 )
 
-// ListStates returns a list of all saved states
-func (s *Server) ListStates(_ context.Context, _ *proto.ListStatesRequest) (*proto.ListStatesResponse, error) {
-	mgr := statemanager.New(statemanager.GetDefaultStatePath())
-
-	stateNames, err := mgr.GetSavedStateNames()
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to get saved state names: %v", err)
-	}
-
-	states := make([]*proto.State, 0, len(stateNames))
-	for _, name := range stateNames {
-		states = append(states, &proto.State{
-			Name: name,
-		})
-	}
-
-	return &proto.ListStatesResponse{
-		States: states,
-	}, nil
-}
-
-// CleanState handles cleaning of states (performing cleanup operations)
-func (s *Server) CleanState(ctx context.Context, req *proto.CleanStateRequest) (*proto.CleanStateResponse, error) {
-	if s.connectClient.Status() == internal.StatusConnected || s.connectClient.Status() == internal.StatusConnecting {
-		return nil, status.Errorf(codes.FailedPrecondition, "cannot clean state while connecting or connected, run 'netbird down' first.")
-	}
-
-	if req.All {
-		// Reuse existing cleanup logic for all states
-		if err := restoreResidualState(ctx); err != nil {
-			return nil, status.Errorf(codes.Internal, "failed to clean all states: %v", err)
-		}
-
-		// Get count of cleaned states
-		mgr := statemanager.New(statemanager.GetDefaultStatePath())
-		stateNames, err := mgr.GetSavedStateNames()
-		if err != nil {
-			return nil, status.Errorf(codes.Internal, "failed to get state count: %v", err)
-		}
-
-		return &proto.CleanStateResponse{
-			CleanedStates: int32(len(stateNames)),
-		}, nil
-	}
-
-	// Handle single state cleanup
-	mgr := statemanager.New(statemanager.GetDefaultStatePath())
-	registerStates(mgr)
-
-	if err := mgr.CleanupStateByName(req.StateName); err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to clean state %s: %v", req.StateName, err)
-	}
-
-	if err := mgr.PersistState(ctx); err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to persist state changes: %v", err)
-	}
-
-	return &proto.CleanStateResponse{
-		CleanedStates: 1,
-	}, nil
-}
-
-// DeleteState handles deletion of states without cleanup
-func (s *Server) DeleteState(ctx context.Context, req *proto.DeleteStateRequest) (*proto.DeleteStateResponse, error) {
-	if s.connectClient.Status() == internal.StatusConnected || s.connectClient.Status() == internal.StatusConnecting {
-		return nil, status.Errorf(codes.FailedPrecondition, "cannot clean state while connecting or connected, run 'netbird down' first.")
-	}
-
-	mgr := statemanager.New(statemanager.GetDefaultStatePath())
-
-	var count int
-	var err error
-
-	if req.All {
-		count, err = mgr.DeleteAllStates()
-	} else {
-		err = mgr.DeleteStateByName(req.StateName)
-		if err == nil {
-			count = 1
-		}
-	}
-
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to delete state: %v", err)
-	}
-
-	// Persist the changes
-	if err := mgr.PersistState(ctx); err != nil {
-		return nil, status.Errorf(codes.Internal, "failed to persist state changes: %v", err)
-	}
-
-	return &proto.DeleteStateResponse{
-		DeletedStates: int32(count),
-	}, nil
-}
-
-// restoreResidualState checks if the client was not shut down in a clean way and restores residual if required.
+// restoreResidualConfig checks if the client was not shut down in a clean way and restores residual state if required.
 // Otherwise, we might not be able to connect to the management server to retrieve new config.
 func restoreResidualState(ctx context.Context) error {
 	path := statemanager.GetDefaultStatePath()
@@ -124,7 +24,6 @@ func restoreResidualState(ctx context.Context) error {
 	registerStates(mgr)
 
 	var merr *multierror.Error
-
 	if err := mgr.PerformCleanup(); err != nil {
 		merr = multierror.Append(merr, fmt.Errorf("perform cleanup: %w", err))
 	}

client/system/info.go

@@ -61,14 +61,6 @@ type Info struct {
 	Files              []File // for posture checks
 }
 
-// StaticInfo is an object that contains machine information that does not change
-type StaticInfo struct {
-	SystemSerialNumber string
-	SystemProductName  string
-	SystemManufacturer string
-	Environment        Environment
-}
-
 // extractUserAgent extracts Netbird's agent (client) name and version from the outgoing context
 func extractUserAgent(ctx context.Context) string {
 	md, hasMeta := metadata.FromOutgoingContext(ctx)

client/system/info_darwin.go

@@ -10,12 +10,13 @@ import (
 	"os/exec"
 	"runtime"
 	"strings"
-	"time"
 
 	"golang.org/x/sys/unix"
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/system/detect_cloud"
+	"github.com/netbirdio/netbird/client/system/detect_platform"
 	"github.com/netbirdio/netbird/version"
 )
 
@@ -40,10 +41,11 @@ func GetInfo(ctx context.Context) *Info {
 		log.Warnf("failed to discover network addresses: %s", err)
 	}
 
-	start := time.Now()
-	si := updateStaticInfo()
-	if time.Since(start) > 1*time.Second {
-		log.Warnf("updateStaticInfo took %s", time.Since(start))
+	serialNum, prodName, manufacturer := sysInfo()
+
+	env := Environment{
+		Cloud:    detect_cloud.Detect(ctx),
+		Platform: detect_platform.Detect(ctx),
 	}
 
 	gio := &Info{
@@ -55,10 +57,10 @@ func GetInfo(ctx context.Context) *Info {
 		CPUs:               runtime.NumCPU(),
 		KernelVersion:      release,
 		NetworkAddresses:   addrs,
-		SystemSerialNumber: si.SystemSerialNumber,
-		SystemProductName:  si.SystemProductName,
-		SystemManufacturer: si.SystemManufacturer,
-		Environment:        si.Environment,
+		SystemSerialNumber: serialNum,
+		SystemProductName:  prodName,
+		SystemManufacturer: manufacturer,
+		Environment:        env,
 	}
 
 	systemHostname, _ := os.Hostname()

client/system/info_linux.go

@@ -1,4 +1,5 @@
 //go:build !android
+// +build !android
 
 package system
 
@@ -15,13 +16,30 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/zcalusic/sysinfo"
 
+	"github.com/netbirdio/netbird/client/system/detect_cloud"
+	"github.com/netbirdio/netbird/client/system/detect_platform"
 	"github.com/netbirdio/netbird/version"
 )
 
-var (
-	// it is override in tests
-	getSystemInfo = defaultSysInfoImplementation
-)
+type SysInfoGetter interface {
+	GetSysInfo() SysInfo
+}
+
+type SysInfoWrapper struct {
+	si sysinfo.SysInfo
+}
+
+func (s SysInfoWrapper) GetSysInfo() SysInfo {
+	s.si.GetSysInfo()
+	return SysInfo{
+		ChassisSerial: s.si.Chassis.Serial,
+		ProductSerial: s.si.Product.Serial,
+		BoardSerial:   s.si.Board.Serial,
+		ProductName:   s.si.Product.Name,
+		BoardName:     s.si.Board.Name,
+		ProductVendor: s.si.Product.Vendor,
+	}
+}
 
 // GetInfo retrieves and parses the system information
 func GetInfo(ctx context.Context) *Info {
@@ -47,10 +65,12 @@ func GetInfo(ctx context.Context) *Info {
 		log.Warnf("failed to discover network addresses: %s", err)
 	}
 
-	start := time.Now()
-	si := updateStaticInfo()
-	if time.Since(start) > 1*time.Second {
-		log.Warnf("updateStaticInfo took %s", time.Since(start))
+	si := SysInfoWrapper{}
+	serialNum, prodName, manufacturer := sysInfo(si.GetSysInfo())
+
+	env := Environment{
+		Cloud:    detect_cloud.Detect(ctx),
+		Platform: detect_platform.Detect(ctx),
 	}
 
 	gio := &Info{
@@ -65,10 +85,10 @@ func GetInfo(ctx context.Context) *Info {
 		UIVersion:          extractUserAgent(ctx),
 		KernelVersion:      osInfo[1],
 		NetworkAddresses:   addrs,
-		SystemSerialNumber: si.SystemSerialNumber,
-		SystemProductName:  si.SystemProductName,
-		SystemManufacturer: si.SystemManufacturer,
-		Environment:        si.Environment,
+		SystemSerialNumber: serialNum,
+		SystemProductName:  prodName,
+		SystemManufacturer: manufacturer,
+		Environment:        env,
 	}
 
 	return gio
@@ -88,9 +108,9 @@ func _getInfo() string {
 	return out.String()
 }
 
-func sysInfo() (string, string, string) {
+func sysInfo(si SysInfo) (string, string, string) {
 	isascii := regexp.MustCompile("^[[:ascii:]]+$")
-	si := getSystemInfo()
+
 	serials := []string{si.ChassisSerial, si.ProductSerial}
 	serial := ""
 
@@ -121,16 +141,3 @@ func sysInfo() (string, string, string) {
 	}
 	return serial, name, manufacturer
 }
-
-func defaultSysInfoImplementation() SysInfo {
-	si := sysinfo.SysInfo{}
-	si.GetSysInfo()
-	return SysInfo{
-		ChassisSerial: si.Chassis.Serial,
-		ProductSerial: si.Product.Serial,
-		BoardSerial:   si.Board.Serial,
-		ProductName:   si.Product.Name,
-		BoardName:     si.Board.Name,
-		ProductVendor: si.Product.Vendor,
-	}
-}

client/system/info_windows.go

@@ -6,12 +6,13 @@ import (
 	"os"
 	"runtime"
 	"strings"
-	"time"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/yusufpapurcu/wmi"
 	"golang.org/x/sys/windows/registry"
 
+	"github.com/netbirdio/netbird/client/system/detect_cloud"
+	"github.com/netbirdio/netbird/client/system/detect_platform"
 	"github.com/netbirdio/netbird/version"
 )
 
@@ -41,10 +42,24 @@ func GetInfo(ctx context.Context) *Info {
 		log.Warnf("failed to discover network addresses: %s", err)
 	}
 
-	start := time.Now()
-	si := updateStaticInfo()
-	if time.Since(start) > 1*time.Second {
-		log.Warnf("updateStaticInfo took %s", time.Since(start))
+	serialNum, err := sysNumber()
+	if err != nil {
+		log.Warnf("failed to get system serial number: %s", err)
+	}
+
+	prodName, err := sysProductName()
+	if err != nil {
+		log.Warnf("failed to get system product name: %s", err)
+	}
+
+	manufacturer, err := sysManufacturer()
+	if err != nil {
+		log.Warnf("failed to get system manufacturer: %s", err)
+	}
+
+	env := Environment{
+		Cloud:    detect_cloud.Detect(ctx),
+		Platform: detect_platform.Detect(ctx),
 	}
 
 	gio := &Info{
@@ -56,10 +71,10 @@ func GetInfo(ctx context.Context) *Info {
 		CPUs:               runtime.NumCPU(),
 		KernelVersion:      buildVersion,
 		NetworkAddresses:   addrs,
-		SystemSerialNumber: si.SystemSerialNumber,
-		SystemProductName:  si.SystemProductName,
-		SystemManufacturer: si.SystemManufacturer,
-		Environment:        si.Environment,
+		SystemSerialNumber: serialNum,
+		SystemProductName:  prodName,
+		SystemManufacturer: manufacturer,
+		Environment:        env,
 	}
 
 	systemHostname, _ := os.Hostname()
@@ -70,26 +85,6 @@ func GetInfo(ctx context.Context) *Info {
 	return gio
 }
 
-func sysInfo() (serialNumber string, productName string, manufacturer string) {
-	var err error
-	serialNumber, err = sysNumber()
-	if err != nil {
-		log.Warnf("failed to get system serial number: %s", err)
-	}
-
-	productName, err = sysProductName()
-	if err != nil {
-		log.Warnf("failed to get system product name: %s", err)
-	}
-
-	manufacturer, err = sysManufacturer()
-	if err != nil {
-		log.Warnf("failed to get system manufacturer: %s", err)
-	}
-
-	return serialNumber, productName, manufacturer
-}
-
 func getOSNameAndVersion() (string, string) {
 	var dst []Win32_OperatingSystem
 	query := wmi.CreateQuery(&dst, "")

client/system/static_info.go

@@ -1,46 +0,0 @@
-//go:build (linux && !android) || windows || (darwin && !ios)
-
-package system
-
-import (
-	"context"
-	"sync"
-	"time"
-
-	"github.com/netbirdio/netbird/client/system/detect_cloud"
-	"github.com/netbirdio/netbird/client/system/detect_platform"
-)
-
-var (
-	staticInfo StaticInfo
-	once       sync.Once
-)
-
-func init() {
-	go func() {
-		_ = updateStaticInfo()
-	}()
-}
-
-func updateStaticInfo() StaticInfo {
-	once.Do(func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		wg := sync.WaitGroup{}
-		wg.Add(3)
-		go func() {
-			staticInfo.SystemSerialNumber, staticInfo.SystemProductName, staticInfo.SystemManufacturer = sysInfo()
-			wg.Done()
-		}()
-		go func() {
-			staticInfo.Environment.Cloud = detect_cloud.Detect(ctx)
-			wg.Done()
-		}()
-		go func() {
-			staticInfo.Environment.Platform = detect_platform.Detect(ctx)
-			wg.Done()
-		}()
-		wg.Wait()
-	})
-	return staticInfo
-}

client/system/sysinfo_linux_test.go

@@ -183,10 +183,7 @@ func Test_sysInfo(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			getSystemInfo = func() SysInfo {
-				return tt.sysInfo
-			}
-			gotSerialNum, gotProdName, gotManufacturer := sysInfo()
+			gotSerialNum, gotProdName, gotManufacturer := sysInfo(tt.sysInfo)
 			if gotSerialNum != tt.wantSerialNum {
 				t.Errorf("sysInfo() gotSerialNum = %v, want %v", gotSerialNum, tt.wantSerialNum)
 			}

client/testdata/store.sql

@@ -31,6 +31,9 @@ INSERT INTO accounts VALUES('bf1c8084-ba50-4ce7-9439-34653001fc3b','','2024-10-0
 INSERT INTO setup_keys VALUES('','bf1c8084-ba50-4ce7-9439-34653001fc3b','A2C8E62B-38F5-4553-B31E-DD66C696CEBB','Default key','reusable','2021-08-19 20:46:20.005936822+02:00','2321-09-18 20:46:20.005936822+02:00','2021-08-19 20:46:20.005936822+02:00',0,0,'0001-01-01 00:00:00+00:00','[]',0,0);
 INSERT INTO users VALUES('edafee4e-63fb-11ec-90d6-0242ac120003','bf1c8084-ba50-4ce7-9439-34653001fc3b','admin',0,0,'','[]',0,'0001-01-01 00:00:00+00:00','2024-10-02 21:28:24.830506+02:00','api',0,'');
 INSERT INTO users VALUES('f4f6d672-63fb-11ec-90d6-0242ac120003','bf1c8084-ba50-4ce7-9439-34653001fc3b','user',0,0,'','[]',0,'0001-01-01 00:00:00+00:00','2024-10-02 21:28:24.830506+02:00','api',0,'');
+INSERT INTO "groups" VALUES('cs1tnh0hhcjnqoiuebeg','bf1c8084-ba50-4ce7-9439-34653001fc3b','All','api','["cfvprsrlo1hqoo49ohog", "cg3161rlo1hs9cq94gdg", "cg05lnblo1hkg2j514p0"]',0,'');
+INSERT INTO policies VALUES('cs1tnh0hhcjnqoiuebf0','bf1c8084-ba50-4ce7-9439-34653001fc3b','Default','This is a default rule that allows connections between all the resources',1,'[]');
+INSERT INTO policy_rules VALUES('cs387mkv2d4bgq41b6n0','cs1tnh0hhcjnqoiuebf0','Default','This is a default rule that allows connections between all the resources',1,'accept','["cs1tnh0hhcjnqoiuebeg"]','["cs1tnh0hhcjnqoiuebeg"]',1,'all',NULL,NULL);
 INSERT INTO installations VALUES(1,'');
 
 COMMIT;

go.mod

@@ -25,7 +25,7 @@ require (
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	golang.zx2c4.com/wireguard/windows v0.5.3
 	google.golang.org/grpc v1.64.1
-	google.golang.org/protobuf v1.34.2
+	google.golang.org/protobuf v1.34.1
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 )
 
@@ -224,7 +224,7 @@ require (
 	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240509183442-62759503f434 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240515191416-fc5f0ca64291 // indirect
 	gopkg.in/square/go-jose.v2 v2.6.0 // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/tomb.v2 v2.0.0-20161208151619-d5d1b5820637 // indirect
@@ -236,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-20241125150134-f9cdce5e32e9
+replace golang.zx2c4.com/wireguard => github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73
 
 replace github.com/cloudflare/circl => github.com/cunicu/circl v0.0.0-20230801113412-fec58fc7b5f6
 

go.sum

@@ -527,8 +527,8 @@ github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502 h1:3tHlFmhTdX9ax
 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-20241125150134-f9cdce5e32e9 h1:Pu/7EukijT09ynHUOzQYW7cC3M/BKU8O4qyN/TvTGoY=
-github.com/netbirdio/wireguard-go v0.0.0-20241125150134-f9cdce5e32e9/go.mod h1:tkCQ4FQXmpAgYVh++1cq16/dH4QJtmvpRv19DWGAHSA=
+github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73 h1:jayg97LH/jJlvpIHVxueTfa+tfQ+FY8fy2sIhCwkz0g=
+github.com/netbirdio/wireguard-go v0.0.0-20241107152827-57d8513b5f73/go.mod h1:tkCQ4FQXmpAgYVh++1cq16/dH4QJtmvpRv19DWGAHSA=
 github.com/nicksnyder/go-i18n/v2 v2.4.0 h1:3IcvPOAvnCKwNm0TB0dLDTuawWEj+ax/RERNC+diLMM=
 github.com/nicksnyder/go-i18n/v2 v2.4.0/go.mod h1:nxYSZE9M0bf3Y70gPQjN9ha7XNHX7gMc814+6wVyEI4=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
@@ -1151,8 +1151,8 @@ google.golang.org/genproto v0.0.0-20210402141018-6c239bbf2bb1/go.mod h1:9lPAdzaE
 google.golang.org/genproto v0.0.0-20210602131652-f16073e35f0c/go.mod h1:UODoCrxHCcBojKKwX1terBiRUaqAsFqJiF615XL43r0=
 google.golang.org/genproto/googleapis/api v0.0.0-20240509183442-62759503f434 h1:OpXbo8JnN8+jZGPrL4SSfaDjSCjupr8lXyBAbexEm/U=
 google.golang.org/genproto/googleapis/api v0.0.0-20240509183442-62759503f434/go.mod h1:FfiGhwUm6CJviekPrc0oJ+7h29e+DmWU6UtjX0ZvI7Y=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 h1:pPJltXNxVzT4pK9yD8vR9X75DaWYYmLGMsEvBfFQZzQ=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240515191416-fc5f0ca64291 h1:AgADTJarZTBqgjiUzRgfaBchgYB3/WFTC80GPwsMcRI=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240515191416-fc5f0ca64291/go.mod h1:EfXuqaE1J41VCDicxHzUDm+8rk+7ZdXzHV0IhO/I6s0=
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
 google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
@@ -1189,8 +1189,8 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
-google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
+google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
+google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

infrastructure_files/getting-started-with-zitadel.sh

@@ -530,7 +530,7 @@ renderCaddyfile() {
 {
   debug
 	servers :80,:443 {
-    protocols h1 h2c h3
+    protocols h1 h2c
   }
 }
 
@@ -788,7 +788,6 @@ services:
     networks: [ netbird ]
     ports:
       - '443:443'
-      - '443:443/udp'
       - '80:80'
       - '8080:8080'
     volumes:

management/server/account_test.go

@@ -6,17 +6,13 @@ import (
 	b64 "encoding/base64"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net"
-	"os"
 	"reflect"
-	"strconv"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/golang-jwt/jwt"
-	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
@@ -405,7 +401,14 @@ func TestAccount_GetPeerNetworkMap(t *testing.T) {
 	}
 
 	for _, testCase := range tt {
-		account := newAccountWithId(context.Background(), "account-1", userID, "netbird.io")
+		store := newStore(t)
+
+		err := newAccountWithId(context.Background(), store, "account-1", userID, "netbird.io")
+		require.NoError(t, err, "failed to create account")
+
+		account, err := store.GetAccount(context.Background(), "account-1")
+		require.NoError(t, err, "failed to get account")
+
 		account.UpdateSettings(&testCase.accountSettings)
 		account.Network = network
 		account.Peers = testCase.peers
@@ -423,6 +426,8 @@ func TestAccount_GetPeerNetworkMap(t *testing.T) {
 		networkMap := account.GetPeerNetworkMap(context.Background(), testCase.peerID, customZone, validatedPeers, nil)
 		assert.Len(t, networkMap.Peers, len(testCase.expectedPeers))
 		assert.Len(t, networkMap.OfflinePeers, len(testCase.expectedOfflinePeers))
+
+		store.Close(context.Background())
 	}
 }
 
@@ -430,27 +435,35 @@ func TestNewAccount(t *testing.T) {
 	domain := "netbird.io"
 	userId := "account_creator"
 	accountID := "account_id"
-	account := newAccountWithId(context.Background(), accountID, userId, domain)
+
+	store := newStore(t)
+	defer store.Close(context.Background())
+
+	err := newAccountWithId(context.Background(), store, accountID, userId, domain)
+	require.NoError(t, err, "failed to create account")
+
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err, "failed to get account")
 	verifyNewAccountHasDefaultFields(t, account, userId, domain, []string{userId})
 }
 
-func TestAccountManager_GetOrCreateAccountByUser(t *testing.T) {
+func TestAccountManager_GetOrCreateAccountIDByUser(t *testing.T) {
 	manager, err := createManager(t)
 	if err != nil {
 		t.Fatal(err)
 		return
 	}
 
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "")
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userID, "")
 	if err != nil {
 		t.Fatal(err)
 	}
-	if account == nil {
+	if accountID == "" {
 		t.Fatalf("expected to create an account for a user %s", userID)
 		return
 	}
 
-	account, err = manager.Store.GetAccountByUser(context.Background(), userID)
+	account, err := manager.Store.GetAccountByUser(context.Background(), userID)
 	if err != nil {
 		t.Errorf("expected to get existing account after creation, no account was found for a user %s", userID)
 		return
@@ -673,15 +686,12 @@ func TestDefaultAccountManager_GetGroupsFromTheToken(t *testing.T) {
 	userId := "user-id"
 	domain := "test.domain"
 
-	_ = newAccountWithId(context.Background(), "", userId, domain)
 	manager, err := createManager(t)
 	require.NoError(t, err, "unable to create account manager")
 
 	accountID, err := manager.GetAccountIDByUserID(context.Background(), userId, domain)
 	require.NoError(t, err, "create init user failed")
-	// as initAccount was created without account id we have to take the id after account initialization
-	// that happens inside the GetAccountIDByUserID where the id is getting generated
-	// it is important to set the id as it help to avoid creating additional account with empty Id and re-pointing indices to it
+
 	initAccount, err := manager.Store.GetAccount(context.Background(), accountID)
 	require.NoError(t, err, "get init account failed")
 
@@ -697,44 +707,53 @@ func TestDefaultAccountManager_GetGroupsFromTheToken(t *testing.T) {
 		accountID, _, err := manager.GetAccountIDFromToken(context.Background(), claims)
 		require.NoError(t, err, "get account by token failed")
 
-		account, err := manager.Store.GetAccount(context.Background(), accountID)
-		require.NoError(t, err, "get account failed")
+		accountGroups, err := manager.Store.GetAccountGroups(context.Background(), LockingStrengthShare, accountID)
+		require.NoError(t, err, "failed to get account groups")
 
-		require.Len(t, account.Groups, 1, "only ALL group should exists")
+		require.Len(t, accountGroups, 1, "only ALL group should exists")
 	})
 
 	t.Run("JWT groups enabled without claim name", func(t *testing.T) {
 		initAccount.Settings.JWTGroupsEnabled = true
-		err := manager.Store.SaveAccount(context.Background(), initAccount)
-		require.NoError(t, err, "save account failed")
-		require.Len(t, manager.Store.GetAllAccounts(context.Background()), 1, "only one account should exist")
+		_, err = manager.UpdateAccountSettings(context.Background(), accountID, userId, initAccount.Settings)
+		require.NoError(t, err, "failed to update account settings")
+
+		totalAccounts, err := manager.Store.GetTotalAccounts(context.Background())
+		require.NoError(t, err, "failed to get total accounts")
+		require.Equal(t, int64(1), totalAccounts, "only one account should exist")
 
 		accountID, _, err := manager.GetAccountIDFromToken(context.Background(), claims)
 		require.NoError(t, err, "get account by token failed")
 
-		account, err := manager.Store.GetAccount(context.Background(), accountID)
-		require.NoError(t, err, "get account failed")
+		accountGroups, err := manager.Store.GetAccountGroups(context.Background(), LockingStrengthShare, accountID)
+		require.NoError(t, err, "failed to get account groups")
 
-		require.Len(t, account.Groups, 1, "if group claim is not set no group added from JWT")
+		require.Len(t, accountGroups, 1, "if group claim is not set no group added from JWT")
 	})
 
 	t.Run("JWT groups enabled", func(t *testing.T) {
 		initAccount.Settings.JWTGroupsEnabled = true
 		initAccount.Settings.JWTGroupsClaimName = "idp-groups"
-		err := manager.Store.SaveAccount(context.Background(), initAccount)
-		require.NoError(t, err, "save account failed")
-		require.Len(t, manager.Store.GetAllAccounts(context.Background()), 1, "only one account should exist")
+		_, err = manager.UpdateAccountSettings(context.Background(), accountID, userId, initAccount.Settings)
+		require.NoError(t, err, "failed to update account settings")
+
+		totalAccounts, err := manager.Store.GetTotalAccounts(context.Background())
+		require.NoError(t, err, "failed to get total accounts")
+		require.Equal(t, int64(1), totalAccounts, "only one account should exist")
 
 		accountID, _, err := manager.GetAccountIDFromToken(context.Background(), claims)
 		require.NoError(t, err, "get account by token failed")
 
-		account, err := manager.Store.GetAccount(context.Background(), accountID)
-		require.NoError(t, err, "get account failed")
+		exists, err := manager.Store.AccountExists(context.Background(), LockingStrengthShare, accountID)
+		require.NoError(t, err, "failed to check account existence")
+		require.True(t, exists, "account should exist")
 
-		require.Len(t, account.Groups, 3, "groups should be added to the account")
+		accountGroups, err := manager.GetAllGroups(context.Background(), accountID, userId)
+		require.NoError(t, err, "failed to get account groups")
+		require.Len(t, accountGroups, 3, "groups should be added to the account")
 
 		groupsByNames := map[string]*group.Group{}
-		for _, g := range account.Groups {
+		for _, g := range accountGroups {
 			groupsByNames[g.Name] = g
 		}
 
@@ -750,62 +769,55 @@ func TestDefaultAccountManager_GetGroupsFromTheToken(t *testing.T) {
 	})
 }
 
-func TestAccountManager_GetAccountFromPAT(t *testing.T) {
+func TestAccountManager_GetAccountInfoFromPAT(t *testing.T) {
 	store := newStore(t)
-	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+	err := newAccountWithId(context.Background(), store, "account_id", "testuser", "")
+	require.NoError(t, err, "failed to create account")
 
 	token := "nbp_9999EUDNdkeusjentDLSJEn1902u84390W6W"
 	hashedToken := sha256.Sum256([]byte(token))
 	encodedHashedToken := b64.StdEncoding.EncodeToString(hashedToken[:])
-	account.Users["someUser"] = &User{
-		Id: "someUser",
-		PATs: map[string]*PersonalAccessToken{
-			"tokenId": {
-				ID:          "tokenId",
-				HashedToken: encodedHashedToken,
-			},
-		},
-	}
-	err := store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Fatalf("Error when saving account: %s", err)
+
+	userPAT := &PersonalAccessToken{
+		ID:          "tokenId",
+		UserID:      "testuser",
+		HashedToken: encodedHashedToken,
+		CreatedAt:   time.Now().UTC(),
 	}
+	err = store.SavePAT(context.Background(), LockingStrengthUpdate, userPAT)
+	require.NoError(t, err, "failed to save PAT")
 
 	am := DefaultAccountManager{
 		Store: store,
 	}
 
-	account, user, pat, err := am.GetAccountFromPAT(context.Background(), token)
+	user, pat, _, _, err := am.GetAccountInfoFromPAT(context.Background(), token)
 	if err != nil {
 		t.Fatalf("Error when getting Account from PAT: %s", err)
 	}
 
-	assert.Equal(t, "account_id", account.Id)
-	assert.Equal(t, "someUser", user.Id)
-	assert.Equal(t, account.Users["someUser"].PATs["tokenId"], pat)
+	assert.Equal(t, "account_id", user.AccountID)
+	assert.Equal(t, "testuser", user.Id)
+	assert.Equal(t, userPAT, pat)
 }
 
 func TestDefaultAccountManager_MarkPATUsed(t *testing.T) {
 	store := newStore(t)
-	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+	err := newAccountWithId(context.Background(), store, "account_id", "testuser", "")
+	require.NoError(t, err, "failed to create account")
 
 	token := "nbp_9999EUDNdkeusjentDLSJEn1902u84390W6W"
 	hashedToken := sha256.Sum256([]byte(token))
 	encodedHashedToken := b64.StdEncoding.EncodeToString(hashedToken[:])
-	account.Users["someUser"] = &User{
-		Id: "someUser",
-		PATs: map[string]*PersonalAccessToken{
-			"tokenId": {
-				ID:          "tokenId",
-				HashedToken: encodedHashedToken,
-				LastUsed:    time.Time{},
-			},
-		},
-	}
-	err := store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Fatalf("Error when saving account: %s", err)
+
+	userPAT := &PersonalAccessToken{
+		ID:          "tokenId",
+		UserID:      "someUser",
+		HashedToken: encodedHashedToken,
+		LastUsed:    time.Time{},
 	}
+	err = store.SavePAT(context.Background(), LockingStrengthUpdate, userPAT)
+	require.NoError(t, err, "failed to save PAT")
 
 	am := DefaultAccountManager{
 		Store: store,
@@ -816,11 +828,10 @@ func TestDefaultAccountManager_MarkPATUsed(t *testing.T) {
 		t.Fatalf("Error when marking PAT used: %s", err)
 	}
 
-	account, err = am.Store.GetAccount(context.Background(), "account_id")
-	if err != nil {
-		t.Fatalf("Error when getting account: %s", err)
-	}
-	assert.True(t, !account.Users["someUser"].PATs["tokenId"].LastUsed.IsZero())
+	userPAT, err = store.GetPATByID(context.Background(), LockingStrengthShare, userPAT.UserID, userPAT.ID)
+	require.NoError(t, err, "failed to get PAT")
+
+	assert.True(t, !userPAT.LastUsed.IsZero())
 }
 
 func TestAccountManager_PrivateAccount(t *testing.T) {
@@ -831,15 +842,15 @@ func TestAccountManager_PrivateAccount(t *testing.T) {
 	}
 
 	userId := "test_user"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userId, "")
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userId, "")
 	if err != nil {
 		t.Fatal(err)
 	}
-	if account == nil {
+	if accountID == "" {
 		t.Fatalf("expected to create an account for a user %s", userId)
 	}
 
-	account, err = manager.Store.GetAccountByUser(context.Background(), userId)
+	account, err := manager.Store.GetAccountByUser(context.Background(), userId)
 	if err != nil {
 		t.Errorf("expected to get existing account after creation, no account was found for a user %s", userId)
 	}
@@ -858,32 +869,22 @@ func TestAccountManager_SetOrUpdateDomain(t *testing.T) {
 
 	userId := "test_user"
 	domain := "hotmail.com"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userId, domain)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if account == nil {
-		t.Fatalf("expected to create an account for a user %s", userId)
-	}
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userId, domain)
+	require.NoError(t, err, "failed to get or create account by user")
+	require.NotEmptyf(t, accountID, "expected to create an account for a user %s", userId)
 
-	if account != nil && account.Domain != domain {
-		t.Errorf("setting account domain failed, expected %s, got %s", domain, account.Domain)
-	}
+	accDomain, _, err := manager.Store.GetAccountDomainAndCategory(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account domain and category")
+	require.Equal(t, domain, accDomain, "expected account domain to match")
 
 	domain = "gmail.com"
 
-	account, err = manager.GetOrCreateAccountByUser(context.Background(), userId, domain)
-	if err != nil {
-		t.Fatalf("got the following error while retrieving existing acc: %v", err)
-	}
+	accountID, err = manager.GetOrCreateAccountIDByUser(context.Background(), userId, domain)
+	require.NoError(t, err, "failed to get or create account by user")
 
-	if account == nil {
-		t.Fatalf("expected to get an account for a user %s", userId)
-	}
-
-	if account != nil && account.Domain != domain {
-		t.Errorf("updating domain. expected %s got %s", domain, account.Domain)
-	}
+	accDomain, _, err = manager.Store.GetAccountDomainAndCategory(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account domain and category")
+	require.Equal(t, domain, accDomain, "expected account domain to match")
 }
 
 func TestAccountManager_GetAccountByUserID(t *testing.T) {
@@ -915,12 +916,11 @@ func TestAccountManager_GetAccountByUserID(t *testing.T) {
 }
 
 func createAccount(am *DefaultAccountManager, accountID, userID, domain string) (*Account, error) {
-	account := newAccountWithId(context.Background(), accountID, userID, domain)
-	err := am.Store.SaveAccount(context.Background(), account)
+	err := newAccountWithId(context.Background(), am.Store, accountID, userID, domain)
 	if err != nil {
 		return nil, err
 	}
-	return account, nil
+	return am.Store.GetAccount(context.Background(), accountID)
 }
 
 func TestAccountManager_GetAccount(t *testing.T) {
@@ -1042,7 +1042,7 @@ func BenchmarkTest_GetAccountWithclaims(b *testing.B) {
 	}
 
 	b.Run("public without account ID", func(b *testing.B) {
-		// b.ResetTimer()
+		//b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), publicClaims)
 			if err != nil {
@@ -1052,7 +1052,7 @@ func BenchmarkTest_GetAccountWithclaims(b *testing.B) {
 	})
 
 	b.Run("private without account ID", func(b *testing.B) {
-		// b.ResetTimer()
+		//b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), claims)
 			if err != nil {
@@ -1063,7 +1063,7 @@ func BenchmarkTest_GetAccountWithclaims(b *testing.B) {
 
 	b.Run("private with account ID", func(b *testing.B) {
 		claims.AccountId = id
-		// b.ResetTimer()
+		//b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			_, err := am.getAccountIDWithAuthorizationClaims(context.Background(), claims)
 			if err != nil {
@@ -1167,23 +1167,18 @@ func TestAccountManager_AddPeerWithUserID(t *testing.T) {
 		return
 	}
 
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "netbird.cloud")
-	if err != nil {
-		t.Fatal(err)
-	}
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userID, "netbird.cloud")
+	require.NoError(t, err, "failed to get or create account by user")
 
-	serial := account.Network.CurrentSerial() // should be 0
+	network, err := manager.Store.GetAccountNetwork(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account network")
 
-	if account.Network.Serial != 0 {
-		t.Errorf("expecting account network to have an initial Serial=0")
-		return
-	}
+	serial := network.CurrentSerial() // should be 0
+	require.Equal(t, 0, int(serial), "expected account network to have an initial Serial=0")
 
 	key, err := wgtypes.GeneratePrivateKey()
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
+	require.NoError(t, err, "failed to generate private key")
+
 	expectedPeerKey := key.PublicKey().String()
 	expectedUserID := userID
 
@@ -1191,16 +1186,10 @@ func TestAccountManager_AddPeerWithUserID(t *testing.T) {
 		Key:  expectedPeerKey,
 		Meta: nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
 	})
-	if err != nil {
-		t.Errorf("expecting peer to be added, got failure %v, account users: %v", err, account.CreatedBy)
-		return
-	}
+	require.NoError(t, err, "failed to add peer")
 
-	account, err = manager.Store.GetAccount(context.Background(), account.Id)
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
+	account, err := manager.Store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err, "failed to get account")
 
 	if peer.Key != expectedPeerKey {
 		t.Errorf("expecting just added peer to have key = %s, got %s", expectedPeerKey, peer.Key)
@@ -1476,7 +1465,6 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 		return
 	}
 
-	userID := "account_creator"
 	account, err := createAccount(manager, "test_account", userID, "netbird.cloud")
 	if err != nil {
 		t.Fatal(err)
@@ -1505,7 +1493,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 		return
 	}
 
-	err = manager.DeletePeer(context.Background(), account.Id, peerKey, userID)
+	err = manager.DeletePeer(context.Background(), account.Id, peer.ID, userID)
 	if err != nil {
 		return
 	}
@@ -1527,7 +1515,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 	assert.Equal(t, peer.Name, ev.Meta["name"])
 	assert.Equal(t, peer.FQDN(account.Domain), ev.Meta["fqdn"])
 	assert.Equal(t, userID, ev.InitiatorID)
-	assert.Equal(t, peer.IP.String(), ev.TargetID)
+	assert.Equal(t, peer.ID, ev.TargetID)
 	assert.Equal(t, peer.IP.String(), fmt.Sprint(ev.Meta["ip"]))
 }
 
@@ -1857,16 +1845,15 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
 	accountID, err := manager.GetAccountIDByUserID(context.Background(), userID, "")
 	require.NoError(t, err, "unable to get the account")
 
-	account, err := manager.Store.GetAccount(context.Background(), accountID)
-	require.NoError(t, err, "unable to get the account")
-
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
-	account, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
-		PeerLoginExpiration:        time.Hour,
-		PeerLoginExpirationEnabled: true,
-	})
+	settings, err := manager.GetAccountSettings(context.Background(), accountID, userID)
+	require.NoError(t, err, "unable to get account settings")
+
+	settings.PeerLoginExpirationEnabled = true
+	settings.PeerLoginExpiration = time.Hour
+	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, settings)
 	require.NoError(t, err, "expecting to update account settings successfully but got error")
 
 	wg := &sync.WaitGroup{}
@@ -1883,11 +1870,11 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
 	// disable expiration first
 	update := peer.Copy()
 	update.LoginExpirationEnabled = false
-	_, err = manager.UpdatePeer(context.Background(), account.Id, userID, update)
+	_, err = manager.UpdatePeer(context.Background(), accountID, userID, update)
 	require.NoError(t, err, "unable to update peer")
 	// enabling expiration should trigger the routine
 	update.LoginExpirationEnabled = true
-	_, err = manager.UpdatePeer(context.Background(), account.Id, userID, update)
+	_, err = manager.UpdatePeer(context.Background(), accountID, userID, update)
 	require.NoError(t, err, "unable to update peer")
 
 	failed := waitTimeout(wg, time.Second)
@@ -1911,10 +1898,14 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
 		LoginExpirationEnabled: true,
 	})
 	require.NoError(t, err, "unable to add peer")
-	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
-		PeerLoginExpiration:        time.Hour,
-		PeerLoginExpirationEnabled: true,
-	})
+
+	settings, err := manager.GetAccountSettings(context.Background(), accountID, userID)
+	require.NoError(t, err, "unable to get account settings")
+
+	settings.PeerLoginExpirationEnabled = true
+	settings.PeerLoginExpiration = time.Hour
+
+	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, settings)
 	require.NoError(t, err, "expecting to update account settings successfully but got error")
 
 	wg := &sync.WaitGroup{}
@@ -1931,11 +1922,8 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
 	accountID, err = manager.GetAccountIDByUserID(context.Background(), userID, "")
 	require.NoError(t, err, "unable to get the account")
 
-	account, err := manager.Store.GetAccount(context.Background(), accountID)
-	require.NoError(t, err, "unable to get the account")
-
 	// when we mark peer as connected, the peer login expiration routine should trigger
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
 	failed := waitTimeout(wg, time.Second)
@@ -1966,7 +1954,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 	account, err := manager.Store.GetAccount(context.Background(), accountID)
 	require.NoError(t, err, "unable to get the account")
 
-	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, account)
+	err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), true, nil, accountID)
 	require.NoError(t, err, "unable to mark peer connected")
 
 	wg := &sync.WaitGroup{}
@@ -1979,11 +1967,15 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 			wg.Done()
 		},
 	}
+
 	// enabling PeerLoginExpirationEnabled should trigger the expiration job
-	account, err = manager.UpdateAccountSettings(context.Background(), account.Id, userID, &Settings{
-		PeerLoginExpiration:        time.Hour,
-		PeerLoginExpirationEnabled: true,
-	})
+	settings, err := manager.GetAccountSettings(context.Background(), accountID, userID)
+	require.NoError(t, err, "unable to get account settings")
+
+	settings.PeerLoginExpirationEnabled = true
+	settings.PeerLoginExpiration = time.Hour
+
+	_, err = manager.UpdateAccountSettings(context.Background(), account.Id, userID, settings)
 	require.NoError(t, err, "expecting to update account settings successfully but got error")
 
 	failed := waitTimeout(wg, time.Second)
@@ -1993,10 +1985,8 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
 	wg.Add(1)
 
 	// disabling PeerLoginExpirationEnabled should trigger cancel
-	_, err = manager.UpdateAccountSettings(context.Background(), account.Id, userID, &Settings{
-		PeerLoginExpiration:        time.Hour,
-		PeerLoginExpirationEnabled: false,
-	})
+	settings.PeerLoginExpirationEnabled = false
+	_, err = manager.UpdateAccountSettings(context.Background(), account.Id, userID, settings)
 	require.NoError(t, err, "expecting to update account settings successfully but got error")
 	failed = waitTimeout(wg, time.Second)
 	if failed {
@@ -2011,30 +2001,29 @@ func TestDefaultAccountManager_UpdateAccountSettings(t *testing.T) {
 	accountID, err := manager.GetAccountIDByUserID(context.Background(), userID, "")
 	require.NoError(t, err, "unable to create an account")
 
-	updated, err := manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
-		PeerLoginExpiration:        time.Hour,
-		PeerLoginExpirationEnabled: false,
-	})
-	require.NoError(t, err, "expecting to update account settings successfully but got error")
-	assert.False(t, updated.Settings.PeerLoginExpirationEnabled)
-	assert.Equal(t, updated.Settings.PeerLoginExpiration, time.Hour)
+	settings, err := manager.GetAccountSettings(context.Background(), accountID, userID)
+	require.NoError(t, err, "unable to get account settings")
 
-	settings, err := manager.Store.GetAccountSettings(context.Background(), LockingStrengthShare, accountID)
+	settings.PeerLoginExpirationEnabled = false
+	settings.PeerLoginExpiration = time.Hour
+
+	updatedSettings, err := manager.UpdateAccountSettings(context.Background(), accountID, userID, settings)
+	require.NoError(t, err, "expecting to update account settings successfully but got error")
+	assert.False(t, updatedSettings.PeerLoginExpirationEnabled)
+	assert.Equal(t, updatedSettings.PeerLoginExpiration, time.Hour)
+
+	settings, err = manager.Store.GetAccountSettings(context.Background(), LockingStrengthShare, accountID)
 	require.NoError(t, err, "unable to get account settings")
 
 	assert.False(t, settings.PeerLoginExpirationEnabled)
 	assert.Equal(t, settings.PeerLoginExpiration, time.Hour)
 
-	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
-		PeerLoginExpiration:        time.Second,
-		PeerLoginExpirationEnabled: false,
-	})
+	settings.PeerLoginExpiration = time.Second
+	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, settings)
 	require.Error(t, err, "expecting to fail when providing PeerLoginExpiration less than one hour")
 
-	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &Settings{
-		PeerLoginExpiration:        time.Hour * 24 * 181,
-		PeerLoginExpirationEnabled: false,
-	})
+	settings.PeerLoginExpiration = time.Hour * 24 * 181
+	_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, settings)
 	require.Error(t, err, "expecting to fail when providing PeerLoginExpiration more than 180 days")
 }
 
@@ -2986,218 +2975,3 @@ func peerShouldReceiveUpdate(t *testing.T, updateMessage <-chan *UpdateMessage)
 		t.Error("Timed out waiting for update message")
 	}
 }
-
-func BenchmarkSyncAndMarkPeer(b *testing.B) {
-	benchCases := []struct {
-		name   string
-		peers  int
-		groups int
-		// We need different expectations for CI/CD and local runs because of the different performance characteristics
-		minMsPerOpLocal float64
-		maxMsPerOpLocal float64
-		minMsPerOpCICD  float64
-		maxMsPerOpCICD  float64
-	}{
-		{"Small", 50, 5, 1, 3, 4, 10},
-		{"Medium", 500, 100, 7, 13, 10, 60},
-		{"Large", 5000, 200, 65, 80, 60, 170},
-		{"Small single", 50, 10, 1, 3, 4, 60},
-		{"Medium single", 500, 10, 7, 13, 10, 26},
-		{"Large 5", 5000, 15, 65, 80, 60, 170},
-	}
-
-	log.SetOutput(io.Discard)
-	defer log.SetOutput(os.Stderr)
-	for _, bc := range benchCases {
-		b.Run(bc.name, func(b *testing.B) {
-			manager, accountID, _, err := setupTestAccountManager(b, bc.peers, bc.groups)
-			if err != nil {
-				b.Fatalf("Failed to setup test account manager: %v", err)
-			}
-			ctx := context.Background()
-			account, err := manager.Store.GetAccount(ctx, accountID)
-			if err != nil {
-				b.Fatalf("Failed to get account: %v", err)
-			}
-			peerChannels := make(map[string]chan *UpdateMessage)
-			for peerID := range account.Peers {
-				peerChannels[peerID] = make(chan *UpdateMessage, channelBufferSize)
-			}
-			manager.peersUpdateManager.peerChannels = peerChannels
-
-			b.ResetTimer()
-			start := time.Now()
-			for i := 0; i < b.N; i++ {
-				_, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1})
-				assert.NoError(b, err)
-			}
-
-			duration := time.Since(start)
-			msPerOp := float64(duration.Nanoseconds()) / float64(b.N) / 1e6
-			b.ReportMetric(msPerOp, "ms/op")
-
-			minExpected := bc.minMsPerOpLocal
-			maxExpected := bc.maxMsPerOpLocal
-			if os.Getenv("CI") == "true" {
-				minExpected = bc.minMsPerOpCICD
-				maxExpected = bc.maxMsPerOpCICD
-			}
-
-			if msPerOp < minExpected {
-				b.Fatalf("Benchmark %s failed: too fast (%.2f ms/op, minimum %.2f ms/op)", bc.name, msPerOp, minExpected)
-			}
-
-			if msPerOp > maxExpected {
-				b.Fatalf("Benchmark %s failed: too slow (%.2f ms/op, maximum %.2f ms/op)", bc.name, msPerOp, maxExpected)
-			}
-		})
-	}
-}
-
-func BenchmarkLoginPeer_ExistingPeer(b *testing.B) {
-	benchCases := []struct {
-		name   string
-		peers  int
-		groups int
-		// We need different expectations for CI/CD and local runs because of the different performance characteristics
-		minMsPerOpLocal float64
-		maxMsPerOpLocal float64
-		minMsPerOpCICD  float64
-		maxMsPerOpCICD  float64
-	}{
-		{"Small", 50, 5, 102, 110, 102, 120},
-		{"Medium", 500, 100, 105, 140, 105, 170},
-		{"Large", 5000, 200, 160, 200, 160, 270},
-		{"Small single", 50, 10, 102, 110, 102, 120},
-		{"Medium single", 500, 10, 105, 140, 105, 170},
-		{"Large 5", 5000, 15, 160, 200, 160, 270},
-	}
-
-	log.SetOutput(io.Discard)
-	defer log.SetOutput(os.Stderr)
-	for _, bc := range benchCases {
-		b.Run(bc.name, func(b *testing.B) {
-			manager, accountID, _, err := setupTestAccountManager(b, bc.peers, bc.groups)
-			if err != nil {
-				b.Fatalf("Failed to setup test account manager: %v", err)
-			}
-			ctx := context.Background()
-			account, err := manager.Store.GetAccount(ctx, accountID)
-			if err != nil {
-				b.Fatalf("Failed to get account: %v", err)
-			}
-			peerChannels := make(map[string]chan *UpdateMessage)
-			for peerID := range account.Peers {
-				peerChannels[peerID] = make(chan *UpdateMessage, channelBufferSize)
-			}
-			manager.peersUpdateManager.peerChannels = peerChannels
-
-			b.ResetTimer()
-			start := time.Now()
-			for i := 0; i < b.N; i++ {
-				_, _, _, err := manager.LoginPeer(context.Background(), PeerLogin{
-					WireGuardPubKey: account.Peers["peer-1"].Key,
-					SSHKey:          "someKey",
-					Meta:            nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)},
-					UserID:          "regular_user",
-					SetupKey:        "",
-					ConnectionIP:    net.IP{1, 1, 1, 1},
-				})
-				assert.NoError(b, err)
-			}
-
-			duration := time.Since(start)
-			msPerOp := float64(duration.Nanoseconds()) / float64(b.N) / 1e6
-			b.ReportMetric(msPerOp, "ms/op")
-
-			minExpected := bc.minMsPerOpLocal
-			maxExpected := bc.maxMsPerOpLocal
-			if os.Getenv("CI") == "true" {
-				minExpected = bc.minMsPerOpCICD
-				maxExpected = bc.maxMsPerOpCICD
-			}
-
-			if msPerOp < minExpected {
-				b.Fatalf("Benchmark %s failed: too fast (%.2f ms/op, minimum %.2f ms/op)", bc.name, msPerOp, minExpected)
-			}
-
-			if msPerOp > maxExpected {
-				b.Fatalf("Benchmark %s failed: too slow (%.2f ms/op, maximum %.2f ms/op)", bc.name, msPerOp, maxExpected)
-			}
-		})
-	}
-}
-
-func BenchmarkLoginPeer_NewPeer(b *testing.B) {
-	benchCases := []struct {
-		name   string
-		peers  int
-		groups int
-		// We need different expectations for CI/CD and local runs because of the different performance characteristics
-		minMsPerOpLocal float64
-		maxMsPerOpLocal float64
-		minMsPerOpCICD  float64
-		maxMsPerOpCICD  float64
-	}{
-		{"Small", 50, 5, 107, 120, 107, 140},
-		{"Medium", 500, 100, 105, 140, 105, 170},
-		{"Large", 5000, 200, 180, 220, 180, 320},
-		{"Small single", 50, 10, 107, 120, 105, 140},
-		{"Medium single", 500, 10, 105, 140, 105, 170},
-		{"Large 5", 5000, 15, 180, 220, 180, 320},
-	}
-
-	log.SetOutput(io.Discard)
-	defer log.SetOutput(os.Stderr)
-	for _, bc := range benchCases {
-		b.Run(bc.name, func(b *testing.B) {
-			manager, accountID, _, err := setupTestAccountManager(b, bc.peers, bc.groups)
-			if err != nil {
-				b.Fatalf("Failed to setup test account manager: %v", err)
-			}
-			ctx := context.Background()
-			account, err := manager.Store.GetAccount(ctx, accountID)
-			if err != nil {
-				b.Fatalf("Failed to get account: %v", err)
-			}
-			peerChannels := make(map[string]chan *UpdateMessage)
-			for peerID := range account.Peers {
-				peerChannels[peerID] = make(chan *UpdateMessage, channelBufferSize)
-			}
-			manager.peersUpdateManager.peerChannels = peerChannels
-
-			b.ResetTimer()
-			start := time.Now()
-			for i := 0; i < b.N; i++ {
-				_, _, _, err := manager.LoginPeer(context.Background(), PeerLogin{
-					WireGuardPubKey: "some-new-key" + strconv.Itoa(i),
-					SSHKey:          "someKey",
-					Meta:            nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)},
-					UserID:          "regular_user",
-					SetupKey:        "",
-					ConnectionIP:    net.IP{1, 1, 1, 1},
-				})
-				assert.NoError(b, err)
-			}
-
-			duration := time.Since(start)
-			msPerOp := float64(duration.Nanoseconds()) / float64(b.N) / 1e6
-			b.ReportMetric(msPerOp, "ms/op")
-
-			minExpected := bc.minMsPerOpLocal
-			maxExpected := bc.maxMsPerOpLocal
-			if os.Getenv("CI") == "true" {
-				minExpected = bc.minMsPerOpCICD
-				maxExpected = bc.maxMsPerOpCICD
-			}
-
-			if msPerOp < minExpected {
-				b.Fatalf("Benchmark %s failed: too fast (%.2f ms/op, minimum %.2f ms/op)", bc.name, msPerOp, minExpected)
-			}
-
-			if msPerOp > maxExpected {
-				b.Fatalf("Benchmark %s failed: too slow (%.2f ms/op, maximum %.2f ms/op)", bc.name, msPerOp, maxExpected)
-			}
-		})
-	}
-}

management/server/activity/codes.go

@@ -148,9 +148,6 @@ const (
 	AccountPeerInactivityExpirationDurationUpdated Activity = 67
 
 	SetupKeyDeleted Activity = 68
-
-	UserGroupPropagationEnabled  Activity = 69
-	UserGroupPropagationDisabled Activity = 70
 )
 
 var activityMap = map[Activity]Code{
@@ -225,9 +222,6 @@ var activityMap = map[Activity]Code{
 	AccountPeerInactivityExpirationDisabled:        {"Account peer inactivity expiration disabled", "account.peer.inactivity.expiration.disable"},
 	AccountPeerInactivityExpirationDurationUpdated: {"Account peer inactivity expiration duration updated", "account.peer.inactivity.expiration.update"},
 	SetupKeyDeleted: {"Setup key deleted", "setupkey.delete"},
-
-	UserGroupPropagationEnabled:  {"User group propagation enabled", "account.setting.group.propagation.enable"},
-	UserGroupPropagationDisabled: {"User group propagation disabled", "account.setting.group.propagation.disable"},
 }
 
 // StringCode returns a string code of the activity

management/server/dns.go

@@ -161,7 +161,7 @@ func (am *DefaultAccountManager) SaveDNSSettings(ctx context.Context, accountID
 	return nil
 }
 
-// prepareDNSSettingsEvents prepares a list of event functions to be stored.
+// prepareGroupEvents prepares a list of event functions to be stored.
 func (am *DefaultAccountManager) prepareDNSSettingsEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string) []func() {
 	var eventsToStore []func()
 

management/server/dns_test.go

@@ -39,12 +39,12 @@ func TestGetDNSSettings(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestDNSAccount(t, am)
+	accountID, err := initTestDNSAccount(t, am)
 	if err != nil {
 		t.Fatal("failed to init testing account")
 	}
 
-	dnsSettings, err := am.GetDNSSettings(context.Background(), account.Id, dnsAdminUserID)
+	dnsSettings, err := am.GetDNSSettings(context.Background(), accountID, dnsAdminUserID)
 	if err != nil {
 		t.Fatalf("Got an error when trying to retrieve the DNS settings with an admin user, err: %s", err)
 	}
@@ -53,16 +53,12 @@ func TestGetDNSSettings(t *testing.T) {
 		t.Fatal("DNS settings for new accounts shouldn't return nil")
 	}
 
-	account.DNSSettings = DNSSettings{
+	err = am.Store.SaveDNSSettings(context.Background(), LockingStrengthUpdate, accountID, &DNSSettings{
 		DisabledManagementGroups: []string{group1ID},
-	}
+	})
+	require.NoError(t, err, "failed to update DNS settings")
 
-	err = am.Store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Error("failed to save testing account with new DNS settings")
-	}
-
-	dnsSettings, err = am.GetDNSSettings(context.Background(), account.Id, dnsAdminUserID)
+	dnsSettings, err = am.GetDNSSettings(context.Background(), accountID, dnsAdminUserID)
 	if err != nil {
 		t.Errorf("Got an error when trying to retrieve the DNS settings with an admin user, err: %s", err)
 	}
@@ -71,7 +67,7 @@ func TestGetDNSSettings(t *testing.T) {
 		t.Errorf("DNS settings should have one disabled mgmt group, groups: %s", dnsSettings.DisabledManagementGroups)
 	}
 
-	_, err = am.GetDNSSettings(context.Background(), account.Id, dnsRegularUserID)
+	_, err = am.GetDNSSettings(context.Background(), accountID, dnsRegularUserID)
 	if err == nil {
 		t.Errorf("An error should be returned when getting the DNS settings with a regular user")
 	}
@@ -126,12 +122,12 @@ func TestSaveDNSSettings(t *testing.T) {
 				t.Error("failed to create account manager")
 			}
 
-			account, err := initTestDNSAccount(t, am)
+			accountID, err := initTestDNSAccount(t, am)
 			if err != nil {
 				t.Error("failed to init testing account")
 			}
 
-			err = am.SaveDNSSettings(context.Background(), account.Id, testCase.userID, testCase.inputSettings)
+			err = am.SaveDNSSettings(context.Background(), accountID, testCase.userID, testCase.inputSettings)
 			if err != nil {
 				if testCase.shouldFail {
 					return
@@ -139,7 +135,7 @@ func TestSaveDNSSettings(t *testing.T) {
 				t.Error(err)
 			}
 
-			updatedAccount, err := am.Store.GetAccount(context.Background(), account.Id)
+			updatedAccount, err := am.Store.GetAccount(context.Background(), accountID)
 			if err != nil {
 				t.Errorf("should be able to retrieve updated account, got err: %s", err)
 			}
@@ -158,17 +154,17 @@ func TestGetNetworkMap_DNSConfigSync(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestDNSAccount(t, am)
+	accountID, err := initTestDNSAccount(t, am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
-	peer1, err := account.FindPeerByPubKey(dnsPeer1Key)
+	peer1, err := am.Store.GetPeerByPeerPubKey(context.Background(), LockingStrengthShare, dnsPeer1Key)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
-	peer2, err := account.FindPeerByPubKey(dnsPeer2Key)
+	peer2, err := am.Store.GetPeerByPeerPubKey(context.Background(), LockingStrengthShare, dnsPeer2Key)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
@@ -179,11 +175,13 @@ func TestGetNetworkMap_DNSConfigSync(t *testing.T) {
 	require.True(t, newAccountDNSConfig.DNSConfig.ServiceEnable, "default DNS config should have local DNS service enabled")
 	require.Len(t, newAccountDNSConfig.DNSConfig.NameServerGroups, 0, "updated DNS config should have no nameserver groups since peer 1 is NS for the only existing NS group")
 
-	dnsSettings := account.DNSSettings.Copy()
+	accountDNSSettings, err := am.Store.GetAccountDNSSettings(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account DNS settings")
+
+	dnsSettings := accountDNSSettings.Copy()
 	dnsSettings.DisabledManagementGroups = append(dnsSettings.DisabledManagementGroups, dnsGroup1ID)
-	account.DNSSettings = dnsSettings
-	err = am.Store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+	err = am.Store.SaveDNSSettings(context.Background(), LockingStrengthUpdate, accountID, &dnsSettings)
+	require.NoError(t, err, "failed to update DNS settings")
 
 	updatedAccountDNSConfig, err := am.GetNetworkMap(context.Background(), peer1.ID)
 	require.NoError(t, err)
@@ -222,7 +220,7 @@ func createDNSStore(t *testing.T) (Store, error) {
 	return store, nil
 }
 
-func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error) {
+func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (string, error) {
 	t.Helper()
 	peer1 := &nbpeer.Peer{
 		Key:  dnsPeer1Key,
@@ -257,64 +255,65 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 
 	domain := "example.com"
 
-	account := newAccountWithId(context.Background(), dnsAccountID, dnsAdminUserID, domain)
-
-	account.Users[dnsRegularUserID] = &User{
-		Id:   dnsRegularUserID,
-		Role: UserRoleUser,
+	err := newAccountWithId(context.Background(), am.Store, dnsAccountID, dnsAdminUserID, domain)
+	if err != nil {
+		return "", err
 	}
 
-	err := am.Store.SaveAccount(context.Background(), account)
+	err = am.Store.SaveUser(context.Background(), LockingStrengthUpdate, &User{
+		Id:        dnsRegularUserID,
+		AccountID: dnsAccountID,
+		Role:      UserRoleUser,
+	})
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
 	savedPeer1, _, _, err := am.AddPeer(context.Background(), "", dnsAdminUserID, peer1)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	_, _, _, err = am.AddPeer(context.Background(), "", dnsAdminUserID, peer2)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	account, err = am.Store.GetAccount(context.Background(), account.Id)
+	peer1, err = am.Store.GetPeerByPeerPubKey(context.Background(), LockingStrengthShare, peer1.Key)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	peer1, err = account.FindPeerByPubKey(peer1.Key)
+	_, err = am.Store.GetPeerByPeerPubKey(context.Background(), LockingStrengthShare, peer2.Key)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	_, err = account.FindPeerByPubKey(peer2.Key)
+	err = am.Store.SaveGroups(context.Background(), LockingStrengthUpdate, []*group.Group{
+		{
+			ID:        dnsGroup1ID,
+			AccountID: dnsAccountID,
+			Peers:     []string{peer1.ID},
+			Name:      dnsGroup1ID,
+		},
+		{
+			ID:        dnsGroup2ID,
+			AccountID: dnsAccountID,
+			Name:      dnsGroup2ID,
+		},
+	})
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	newGroup1 := &group.Group{
-		ID:    dnsGroup1ID,
-		Peers: []string{peer1.ID},
-		Name:  dnsGroup1ID,
-	}
-
-	newGroup2 := &group.Group{
-		ID:   dnsGroup2ID,
-		Name: dnsGroup2ID,
-	}
-
-	account.Groups[newGroup1.ID] = newGroup1
-	account.Groups[newGroup2.ID] = newGroup2
-
-	allGroup, err := account.GetGroupAll()
+	allGroup, err := am.Store.GetGroupByName(context.Background(), LockingStrengthShare, dnsAccountID, "All")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	account.NameServerGroups[dnsNSGroup1] = &dns.NameServerGroup{
-		ID:   dnsNSGroup1,
-		Name: "ns-group-1",
+	err = am.Store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, &dns.NameServerGroup{
+		ID:        dnsNSGroup1,
+		AccountID: dnsAccountID,
+		Name:      "ns-group-1",
 		NameServers: []dns.NameServer{{
 			IP:     netip.MustParseAddr(savedPeer1.IP.String()),
 			NSType: dns.UDPNameServerType,
@@ -323,14 +322,12 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 		Primary: true,
 		Enabled: true,
 		Groups:  []string{allGroup.ID},
-	}
-
-	err = am.Store.SaveAccount(context.Background(), account)
+	})
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	return am.Store.GetAccount(context.Background(), account.Id)
+	return dnsAccountID, nil
 }
 
 func generateTestData(size int) nbdns.Config {

management/server/ephemeral.go

@@ -20,10 +20,10 @@ var (
 )
 
 type ephemeralPeer struct {
-	id       string
-	account  *Account
-	deadline time.Time
-	next     *ephemeralPeer
+	id        string
+	accountID string
+	deadline  time.Time
+	next      *ephemeralPeer
 }
 
 // todo: consider to remove peer from ephemeral list when the peer has been deleted via API. If we do not do it
@@ -104,12 +104,6 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
 
 	log.WithContext(ctx).Tracef("add peer to ephemeral list: %s", peer.ID)
 
-	a, err := e.store.GetAccountByPeerID(context.Background(), peer.ID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to add peer to ephemeral list: %s", err)
-		return
-	}
-
 	e.peersLock.Lock()
 	defer e.peersLock.Unlock()
 
@@ -117,7 +111,7 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
 		return
 	}
 
-	e.addPeer(peer.ID, a, newDeadLine())
+	e.addPeer(peer.AccountID, peer.ID, newDeadLine())
 	if e.timer == nil {
 		e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() {
 			e.cleanup(ctx)
@@ -126,17 +120,21 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
 }
 
 func (e *EphemeralManager) loadEphemeralPeers(ctx context.Context) {
-	accounts := e.store.GetAllAccounts(context.Background())
+	peers, err := e.store.GetAllEphemeralPeers(ctx, LockingStrengthShare)
+	if err != nil {
+		log.WithContext(ctx).Debugf("failed to load ephemeral peers: %s", err)
+		return
+	}
+
 	t := newDeadLine()
 	count := 0
-	for _, a := range accounts {
-		for id, p := range a.Peers {
-			if p.Ephemeral {
-				count++
-				e.addPeer(id, a, t)
-			}
+	for _, p := range peers {
+		if p.Ephemeral {
+			count++
+			e.addPeer(p.AccountID, p.ID, t)
 		}
 	}
+
 	log.WithContext(ctx).Debugf("loaded ephemeral peer(s): %d", count)
 }
 
@@ -170,18 +168,18 @@ func (e *EphemeralManager) cleanup(ctx context.Context) {
 
 	for id, p := range deletePeers {
 		log.WithContext(ctx).Debugf("delete ephemeral peer: %s", id)
-		err := e.accountManager.DeletePeer(ctx, p.account.Id, id, activity.SystemInitiator)
+		err := e.accountManager.DeletePeer(ctx, p.accountID, id, activity.SystemInitiator)
 		if err != nil {
 			log.WithContext(ctx).Errorf("failed to delete ephemeral peer: %s", err)
 		}
 	}
 }
 
-func (e *EphemeralManager) addPeer(id string, account *Account, deadline time.Time) {
+func (e *EphemeralManager) addPeer(accountID string, peerID string, deadline time.Time) {
 	ep := &ephemeralPeer{
-		id:       id,
-		account:  account,
-		deadline: deadline,
+		id:        peerID,
+		accountID: accountID,
+		deadline:  deadline,
 	}
 
 	if e.headPeer == nil {

management/server/ephemeral_test.go

@@ -7,25 +7,12 @@ import (
 	"time"
 
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
-	"github.com/netbirdio/netbird/management/server/status"
+	"github.com/stretchr/testify/require"
 )
 
 type MockStore struct {
 	Store
-	account *Account
-}
-
-func (s *MockStore) GetAllAccounts(_ context.Context) []*Account {
-	return []*Account{s.account}
-}
-
-func (s *MockStore) GetAccountByPeerID(_ context.Context, peerId string) (*Account, error) {
-	_, ok := s.account.Peers[peerId]
-	if ok {
-		return s.account, nil
-	}
-
-	return nil, status.NewPeerNotFoundError(peerId)
+	accountID string
 }
 
 type MocAccountManager struct {
@@ -33,9 +20,8 @@ type MocAccountManager struct {
 	store *MockStore
 }
 
-func (a MocAccountManager) DeletePeer(_ context.Context, accountID, peerID, userID string) error {
-	delete(a.store.account.Peers, peerID)
-	return nil //nolint:nil
+func (a MocAccountManager) DeletePeer(_ context.Context, accountID, peerID, _ string) error {
+	return a.store.DeletePeer(context.Background(), LockingStrengthUpdate, accountID, peerID)
 }
 
 func TestNewManager(t *testing.T) {
@@ -44,23 +30,26 @@ func TestNewManager(t *testing.T) {
 		return startTime
 	}
 
-	store := &MockStore{}
+	store := &MockStore{
+		Store: newStore(t),
+	}
 	am := MocAccountManager{
 		store: store,
 	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
-	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	err := seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	require.NoError(t, err, "failed to seed peers")
 
 	mgr := NewEphemeralManager(store, am)
 	mgr.loadEphemeralPeers(context.Background())
 	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
-	if len(store.account.Peers) != numberOfPeers {
-		t.Errorf("failed to cleanup ephemeral peers, expected: %d, result: %d", numberOfPeers, len(store.account.Peers))
-	}
+	peers, err := store.GetAccountPeers(context.Background(), LockingStrengthShare, store.accountID)
+	require.NoError(t, err, "failed to get account peers")
+	require.Equal(t, numberOfPeers, len(peers), "failed to cleanup ephemeral peers")
 }
 
 func TestNewManagerPeerConnected(t *testing.T) {
@@ -69,26 +58,32 @@ func TestNewManagerPeerConnected(t *testing.T) {
 		return startTime
 	}
 
-	store := &MockStore{}
+	store := &MockStore{
+		Store: newStore(t),
+	}
 	am := MocAccountManager{
 		store: store,
 	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
-	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	err := seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	require.NoError(t, err, "failed to seed peers")
 
 	mgr := NewEphemeralManager(store, am)
 	mgr.loadEphemeralPeers(context.Background())
-	mgr.OnPeerConnected(context.Background(), store.account.Peers["ephemeral_peer_0"])
+
+	peer, err := am.store.GetPeerByID(context.Background(), LockingStrengthShare, store.accountID, "ephemeral_peer_0")
+	require.NoError(t, err, "failed to get peer")
+
+	mgr.OnPeerConnected(context.Background(), peer)
 
 	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
-	expected := numberOfPeers + 1
-	if len(store.account.Peers) != expected {
-		t.Errorf("failed to cleanup ephemeral peers, expected: %d, result: %d", expected, len(store.account.Peers))
-	}
+	peers, err := store.GetAccountPeers(context.Background(), LockingStrengthShare, store.accountID)
+	require.NoError(t, err, "failed to get account peers")
+	require.Equal(t, numberOfPeers+1, len(peers), "failed to cleanup ephemeral peers")
 }
 
 func TestNewManagerPeerDisconnected(t *testing.T) {
@@ -97,50 +92,73 @@ func TestNewManagerPeerDisconnected(t *testing.T) {
 		return startTime
 	}
 
-	store := &MockStore{}
+	store := &MockStore{
+		Store: newStore(t),
+	}
 	am := MocAccountManager{
 		store: store,
 	}
 
 	numberOfPeers := 5
 	numberOfEphemeralPeers := 3
-	seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	err := seedPeers(store, numberOfPeers, numberOfEphemeralPeers)
+	require.NoError(t, err, "failed to seed peers")
 
 	mgr := NewEphemeralManager(store, am)
 	mgr.loadEphemeralPeers(context.Background())
-	for _, v := range store.account.Peers {
-		mgr.OnPeerConnected(context.Background(), v)
 
+	peers, err := store.GetAccountPeers(context.Background(), LockingStrengthShare, store.accountID)
+	require.NoError(t, err, "failed to get account peers")
+	for _, v := range peers {
+		mgr.OnPeerConnected(context.Background(), v)
 	}
-	mgr.OnPeerDisconnected(context.Background(), store.account.Peers["ephemeral_peer_0"])
+
+	peer, err := am.store.GetPeerByID(context.Background(), LockingStrengthShare, store.accountID, "ephemeral_peer_0")
+	require.NoError(t, err, "failed to get peer")
+	mgr.OnPeerDisconnected(context.Background(), peer)
 
 	startTime = startTime.Add(ephemeralLifeTime + 1)
 	mgr.cleanup(context.Background())
 
+	peers, err = store.GetAccountPeers(context.Background(), LockingStrengthShare, store.accountID)
+	require.NoError(t, err, "failed to get account peers")
 	expected := numberOfPeers + numberOfEphemeralPeers - 1
-	if len(store.account.Peers) != expected {
-		t.Errorf("failed to cleanup ephemeral peers, expected: %d, result: %d", expected, len(store.account.Peers))
-	}
+	require.Equal(t, expected, len(peers), "failed to cleanup ephemeral peers")
 }
 
-func seedPeers(store *MockStore, numberOfPeers int, numberOfEphemeralPeers int) {
-	store.account = newAccountWithId(context.Background(), "my account", "", "")
+func seedPeers(store *MockStore, numberOfPeers int, numberOfEphemeralPeers int) error {
+	accountID := "my account"
+	err := newAccountWithId(context.Background(), store, accountID, "", "")
+	if err != nil {
+		return err
+	}
+	store.accountID = accountID
 
 	for i := 0; i < numberOfPeers; i++ {
 		peerId := fmt.Sprintf("peer_%d", i)
 		p := &nbpeer.Peer{
 			ID:        peerId,
+			AccountID: accountID,
 			Ephemeral: false,
 		}
-		store.account.Peers[p.ID] = p
+		err = store.AddPeerToAccount(context.Background(), p)
+		if err != nil {
+			return err
+		}
 	}
 
 	for i := 0; i < numberOfEphemeralPeers; i++ {
 		peerId := fmt.Sprintf("ephemeral_peer_%d", i)
 		p := &nbpeer.Peer{
 			ID:        peerId,
+			AccountID: accountID,
 			Ephemeral: true,
 		}
-		store.account.Peers[p.ID] = p
+		err = store.AddPeerToAccount(context.Background(), p)
+		if err != nil {
+			return err
+		}
 	}
+
+	return nil
 }

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,6 +27,11 @@ 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
@@ -37,7 +41,7 @@ func (am *DefaultAccountManager) CheckGroupPermissions(ctx context.Context, acco
 		return status.NewUserNotPartOfAccountError()
 	}
 
-	if user.IsRegularUser() {
+	if user.IsRegularUser() && settings.RegularUsersViewBlocked {
 		return status.NewAdminPermissionError()
 	}
 
@@ -211,9 +215,48 @@ func difference(a, b []string) []string {
 
 // DeleteGroup object of the peers.
 func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, userID, groupID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-	return am.DeleteGroups(ctx, accountID, userID, []string{groupID})
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	if err != nil {
+		return err
+	}
+
+	if user.AccountID != accountID {
+		return status.NewUserNotPartOfAccountError()
+	}
+
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
+	}
+
+	var group *nbgroup.Group
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		group, err = transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, groupID)
+		if err != nil {
+			return err
+		}
+
+		if group.IsGroupAll() {
+			return status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed")
+		}
+
+		if err = validateDeleteGroup(ctx, transaction, group, userID); err != nil {
+			return err
+		}
+
+		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
+			return err
+		}
+
+		return transaction.DeleteGroup(ctx, LockingStrengthUpdate, accountID, groupID)
+	})
+	if err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, userID, groupID, accountID, activity.GroupDeleted, group.EventMeta())
+
+	return nil
 }
 
 // DeleteGroups deletes groups from an account.
@@ -242,14 +285,13 @@ func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountID, us
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
 		for _, groupID := range groupIDs {
-			group, err := transaction.GetGroupByID(ctx, LockingStrengthUpdate, accountID, groupID)
+			group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, groupID)
 			if err != nil {
-				allErrors = errors.Join(allErrors, err)
 				continue
 			}
 
 			if err := validateDeleteGroup(ctx, transaction, group, userID); err != nil {
-				allErrors = errors.Join(allErrors, err)
+				allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete group %s: %w", groupID, err))
 				continue
 			}
 
@@ -276,15 +318,12 @@ func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountID, us
 
 // GroupAddPeer appends peer to the group
 func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, groupID, peerID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	var group *nbgroup.Group
 	var updateAccountPeers bool
 	var err error
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthUpdate, accountID, groupID)
+		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
 		if err != nil {
 			return err
 		}
@@ -317,15 +356,12 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr
 
 // GroupDeletePeer removes peer from the group
 func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, groupID, peerID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	var group *nbgroup.Group
 	var updateAccountPeers bool
 	var err error
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthUpdate, accountID, groupID)
+		group, err = transaction.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID)
 		if err != nil {
 			return err
 		}
@@ -394,17 +430,13 @@ func validateDeleteGroup(ctx context.Context, transaction Store, group *nbgroup.
 	if group.Issued == nbgroup.GroupIssuedIntegration {
 		executingUser, err := transaction.GetUserByUserID(ctx, LockingStrengthShare, userID)
 		if err != nil {
-			return err
+			return status.Errorf(status.NotFound, "user not found")
 		}
 		if executingUser.Role != UserRoleAdmin || !executingUser.IsServiceUser {
 			return status.Errorf(status.PermissionDenied, "only service users with admin power can delete integration group")
 		}
 	}
 
-	if group.IsGroupAll() {
-		return status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed")
-	}
-
 	if isLinked, linkedRoute := isGroupLinkedToRoute(ctx, transaction, group.AccountID, group.ID); isLinked {
 		return &GroupLinkError{"route", string(linkedRoute.NetID)}
 	}
@@ -566,15 +598,6 @@ func areGroupChangesAffectPeers(ctx context.Context, transaction Store, accountI
 	return false, nil
 }
 
-func (am *DefaultAccountManager) anyGroupHasPeers(account *Account, groupIDs []string) bool {
-	for _, groupID := range groupIDs {
-		if group, exists := account.Groups[groupID]; exists && group.HasPeers() {
-			return true
-		}
-	}
-	return false
-}
-
 // anyGroupHasPeers checks if any of the given groups in the account have peers.
 func anyGroupHasPeers(ctx context.Context, transaction Store, accountID string, groupIDs []string) (bool, error) {
 	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, groupIDs)

management/server/group_test.go

@@ -208,7 +208,7 @@ func TestDefaultAccountManager_DeleteGroups(t *testing.T) {
 		{
 			name:            "delete non-existent group",
 			groupIDs:        []string{"non-existent-group"},
-			expectedReasons: []string{"group: non-existent-group not found"},
+			expectedDeleted: []string{"non-existent-group"},
 		},
 		{
 			name:               "delete multiple groups with mixed results",
@@ -328,25 +328,30 @@ func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *A
 	}
 
 	routeResource := &route.Route{
-		ID:     "example route",
-		Groups: []string{groupForRoute.ID},
+		ID:        "example route",
+		AccountID: accountID,
+		Groups:    []string{groupForRoute.ID},
 	}
 
 	routePeerGroupResource := &route.Route{
 		ID:         "example route with peer groups",
+		AccountID:  accountID,
 		PeerGroups: []string{groupForRoute2.ID},
 	}
 
 	nameServerGroup := &nbdns.NameServerGroup{
-		ID:     "example name server group",
-		Groups: []string{groupForNameServerGroups.ID},
+		ID:        "example name server group",
+		AccountID: accountID,
+		Groups:    []string{groupForNameServerGroups.ID},
 	}
 
 	policy := &Policy{
-		ID: "example policy",
+		ID:        "example policy",
+		AccountID: accountID,
 		Rules: []*PolicyRule{
 			{
 				ID:           "example policy rule",
+				PolicyID:     "example policy",
 				Destinations: []string{groupForPolicies.ID},
 			},
 		},
@@ -354,35 +359,60 @@ func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *A
 
 	setupKey := &SetupKey{
 		Id:         "example setup key",
+		AccountID:  accountID,
 		AutoGroups: []string{groupForSetupKeys.ID},
 	}
 
 	user := &User{
 		Id:         "example user",
+		AccountID:  accountID,
 		AutoGroups: []string{groupForUsers.ID},
 	}
-	account := newAccountWithId(context.Background(), accountID, groupAdminUserID, domain)
-	account.Routes[routeResource.ID] = routeResource
-	account.Routes[routePeerGroupResource.ID] = routePeerGroupResource
-	account.NameServerGroups[nameServerGroup.ID] = nameServerGroup
-	account.Policies = append(account.Policies, policy)
-	account.SetupKeys[setupKey.Id] = setupKey
-	account.Users[user.Id] = user
 
-	err := am.Store.SaveAccount(context.Background(), account)
+	err := newAccountWithId(context.Background(), am.Store, accountID, groupAdminUserID, domain)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForRoute)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForRoute2)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForNameServerGroups)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForPolicies)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForSetupKeys)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForUsers)
-	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForIntegration)
+	err = am.Store.SaveRoute(context.Background(), LockingStrengthUpdate, routeResource)
+	if err != nil {
+		return nil, nil, err
+	}
 
-	acc, err := am.Store.GetAccount(context.Background(), account.Id)
+	err = am.Store.SaveRoute(context.Background(), LockingStrengthUpdate, routePeerGroupResource)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.Store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, nameServerGroup)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.Store.CreatePolicy(context.Background(), LockingStrengthUpdate, policy)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.Store.SaveSetupKey(context.Background(), LockingStrengthUpdate, setupKey)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.Store.SaveUser(context.Background(), LockingStrengthUpdate, user)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = am.SaveGroups(context.Background(), accountID, groupAdminUserID, []*nbgroup.Group{
+		groupForRoute, groupForRoute2, groupForNameServerGroups, groupForPolicies,
+		groupForSetupKeys, groupForUsers, groupForIntegration,
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
+	acc, err := am.Store.GetAccount(context.Background(), accountID)
 	if err != nil {
 		return nil, nil, err
 	}

management/server/http/accounts_handler.go

@@ -100,13 +100,13 @@ func (h *AccountsHandler) UpdateAccount(w http.ResponseWriter, r *http.Request)
 		settings.JWTAllowGroups = *req.Settings.JwtAllowGroups
 	}
 
-	updatedAccount, err := h.accountManager.UpdateAccountSettings(r.Context(), accountID, userID, settings)
+	updatedSettings, err := h.accountManager.UpdateAccountSettings(r.Context(), accountID, userID, settings)
 	if err != nil {
 		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	resp := toAccountResponse(updatedAccount.Id, updatedAccount.Settings)
+	resp := toAccountResponse(accountID, updatedSettings)
 
 	util.WriteJSONObject(r.Context(), w, &resp)
 }

management/server/http/accounts_handler_test.go

@@ -29,7 +29,7 @@ func initAccountsTestData(account *server.Account, admin *server.User) *Accounts
 			GetAccountSettingsFunc: func(ctx context.Context, accountID string, userID string) (*server.Settings, error) {
 				return account.Settings, nil
 			},
-			UpdateAccountSettingsFunc: func(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
+			UpdateAccountSettingsFunc: func(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Settings, error) {
 				halfYearLimit := 180 * 24 * time.Hour
 				if newSettings.PeerLoginExpiration > halfYearLimit {
 					return nil, status.Errorf(status.InvalidArgument, "peer login expiration can't be larger than 180 days")
@@ -39,9 +39,7 @@ func initAccountsTestData(account *server.Account, admin *server.User) *Accounts
 					return nil, status.Errorf(status.InvalidArgument, "peer login expiration can't be smaller than one hour")
 				}
 
-				accCopy := account.Copy()
-				accCopy.UpdateSettings(newSettings)
-				return accCopy, nil
+				return newSettings.Copy(), nil
 			},
 		},
 		claimsExtractor: jwtclaims.NewClaimsExtractor(

management/server/http/api/openapi.yml

@@ -439,13 +439,17 @@ components:
               example: 5
           required:
             - accessible_peers_count
-    SetupKeyBase:
+    SetupKey:
       type: object
       properties:
         id:
           description: Setup Key ID
           type: string
           example: 2531583362
+        key:
+          description: Setup Key value
+          type: string
+          example: A616097E-FCF0-48FA-9354-CA4A61142761
         name:
           description: Setup key name identifier
           type: string
@@ -514,31 +518,22 @@ components:
         - updated_at
         - usage_limit
         - ephemeral
-    SetupKeyClear:
-      allOf:
-        - $ref: '#/components/schemas/SetupKeyBase'
-        - type: object
-          properties:
-            key:
-              description: Setup Key as plain text
-              type: string
-              example: A616097E-FCF0-48FA-9354-CA4A61142761
-          required:
-            - key
-    SetupKey:
-      allOf:
-        - $ref: '#/components/schemas/SetupKeyBase'
-        - type: object
-          properties:
-            key:
-              description: Setup Key as secret
-              type: string
-              example: A6160****
-          required:
-            - key
     SetupKeyRequest:
       type: object
       properties:
+        name:
+          description: Setup Key name
+          type: string
+          example: Default key
+        type:
+          description: Setup key type, one-off for single time usage and reusable
+          type: string
+          example: reusable
+        expires_in:
+          description: Expiration time in seconds, 0 will mean the key never expires
+          type: integer
+          minimum: 0
+          example: 86400
         revoked:
           description: Setup key revocation status
           type: boolean
@@ -549,9 +544,21 @@ components:
           items:
             type: string
             example: "ch8i4ug6lnn4g9hqv7m0"
+        usage_limit:
+          description: A number of times this key can be used. The value of 0 indicates the unlimited usage.
+          type: integer
+          example: 0
+        ephemeral:
+          description: Indicate that the peer will be ephemeral or not
+          type: boolean
+          example: true
       required:
+        - name
+        - type
+        - expires_in
         - revoked
         - auto_groups
+        - usage_limit
     CreateSetupKeyRequest:
       type: object
       properties:
@@ -1936,7 +1943,7 @@ paths:
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/SetupKeyClear'
+                $ref: '#/components/schemas/SetupKey'
         '400':
           "$ref": "#/components/responses/bad_request"
         '401':

management/server/http/api/types.gen.go

@@ -1062,94 +1062,7 @@ type SetupKey struct {
 	// Id Setup Key ID
 	Id string `json:"id"`
 
-	// Key Setup Key as secret
-	Key string `json:"key"`
-
-	// LastUsed Setup key last usage date
-	LastUsed time.Time `json:"last_used"`
-
-	// Name Setup key name identifier
-	Name string `json:"name"`
-
-	// Revoked Setup key revocation status
-	Revoked bool `json:"revoked"`
-
-	// State Setup key status, "valid", "overused","expired" or "revoked"
-	State string `json:"state"`
-
-	// Type Setup key type, one-off for single time usage and reusable
-	Type string `json:"type"`
-
-	// UpdatedAt Setup key last update date
-	UpdatedAt time.Time `json:"updated_at"`
-
-	// UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage.
-	UsageLimit int `json:"usage_limit"`
-
-	// UsedTimes Usage count of setup key
-	UsedTimes int `json:"used_times"`
-
-	// Valid Setup key validity status
-	Valid bool `json:"valid"`
-}
-
-// SetupKeyBase defines model for SetupKeyBase.
-type SetupKeyBase struct {
-	// AutoGroups List of group IDs to auto-assign to peers registered with this key
-	AutoGroups []string `json:"auto_groups"`
-
-	// Ephemeral Indicate that the peer will be ephemeral or not
-	Ephemeral bool `json:"ephemeral"`
-
-	// Expires Setup Key expiration date
-	Expires time.Time `json:"expires"`
-
-	// Id Setup Key ID
-	Id string `json:"id"`
-
-	// LastUsed Setup key last usage date
-	LastUsed time.Time `json:"last_used"`
-
-	// Name Setup key name identifier
-	Name string `json:"name"`
-
-	// Revoked Setup key revocation status
-	Revoked bool `json:"revoked"`
-
-	// State Setup key status, "valid", "overused","expired" or "revoked"
-	State string `json:"state"`
-
-	// Type Setup key type, one-off for single time usage and reusable
-	Type string `json:"type"`
-
-	// UpdatedAt Setup key last update date
-	UpdatedAt time.Time `json:"updated_at"`
-
-	// UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage.
-	UsageLimit int `json:"usage_limit"`
-
-	// UsedTimes Usage count of setup key
-	UsedTimes int `json:"used_times"`
-
-	// Valid Setup key validity status
-	Valid bool `json:"valid"`
-}
-
-// SetupKeyClear defines model for SetupKeyClear.
-type SetupKeyClear struct {
-	// AutoGroups List of group IDs to auto-assign to peers registered with this key
-	AutoGroups []string `json:"auto_groups"`
-
-	// Ephemeral Indicate that the peer will be ephemeral or not
-	Ephemeral bool `json:"ephemeral"`
-
-	// Expires Setup Key expiration date
-	Expires time.Time `json:"expires"`
-
-	// Id Setup Key ID
-	Id string `json:"id"`
-
-	// Key Setup Key as plain text
+	// Key Setup Key value
 	Key string `json:"key"`
 
 	// LastUsed Setup key last usage date
@@ -1185,8 +1098,23 @@ type SetupKeyRequest struct {
 	// AutoGroups List of group IDs to auto-assign to peers registered with this key
 	AutoGroups []string `json:"auto_groups"`
 
+	// Ephemeral Indicate that the peer will be ephemeral or not
+	Ephemeral *bool `json:"ephemeral,omitempty"`
+
+	// ExpiresIn Expiration time in seconds, 0 will mean the key never expires
+	ExpiresIn int `json:"expires_in"`
+
+	// Name Setup Key name
+	Name string `json:"name"`
+
 	// Revoked Setup key revocation status
 	Revoked bool `json:"revoked"`
+
+	// Type Setup key type, one-off for single time usage and reusable
+	Type string `json:"type"`
+
+	// UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage.
+	UsageLimit int `json:"usage_limit"`
 }
 
 // User defines model for User.

management/server/http/handler.go

@@ -47,7 +47,7 @@ func APIHandler(ctx context.Context, accountManager s.AccountManager, LocationMa
 	)
 
 	authMiddleware := middleware.NewAuthMiddleware(
-		accountManager.GetAccountFromPAT,
+		accountManager.GetAccountInfoFromPAT,
 		jwtValidator.ValidateAndParse,
 		accountManager.MarkPATUsed,
 		accountManager.CheckUserAccessByJWTGroups,

management/server/http/middleware/auth_middleware.go

@@ -19,8 +19,8 @@ import (
 	"github.com/netbirdio/netbird/management/server/status"
 )
 
-// GetAccountFromPATFunc function
-type GetAccountFromPATFunc func(ctx context.Context, token string) (*server.Account, *server.User, *server.PersonalAccessToken, error)
+// GetAccountInfoFromPATFunc function
+type GetAccountInfoFromPATFunc func(ctx context.Context, token string) (user *server.User, pat *server.PersonalAccessToken, domain string, category string, err error)
 
 // ValidateAndParseTokenFunc function
 type ValidateAndParseTokenFunc func(ctx context.Context, token string) (*jwt.Token, error)
@@ -33,7 +33,7 @@ type CheckUserAccessByJWTGroupsFunc func(ctx context.Context, claims jwtclaims.A
 
 // AuthMiddleware middleware to verify personal access tokens (PAT) and JWT tokens
 type AuthMiddleware struct {
-	getAccountFromPAT          GetAccountFromPATFunc
+	getAccountInfoFromPAT      GetAccountInfoFromPATFunc
 	validateAndParseToken      ValidateAndParseTokenFunc
 	markPATUsed                MarkPATUsedFunc
 	checkUserAccessByJWTGroups CheckUserAccessByJWTGroupsFunc
@@ -47,7 +47,7 @@ const (
 )
 
 // NewAuthMiddleware instance constructor
-func NewAuthMiddleware(getAccountFromPAT GetAccountFromPATFunc, validateAndParseToken ValidateAndParseTokenFunc,
+func NewAuthMiddleware(getAccountInfoFromPAT GetAccountInfoFromPATFunc, validateAndParseToken ValidateAndParseTokenFunc,
 	markPATUsed MarkPATUsedFunc, checkUserAccessByJWTGroups CheckUserAccessByJWTGroupsFunc, claimsExtractor *jwtclaims.ClaimsExtractor,
 	audience string, userIdClaim string) *AuthMiddleware {
 	if userIdClaim == "" {
@@ -55,7 +55,7 @@ func NewAuthMiddleware(getAccountFromPAT GetAccountFromPATFunc, validateAndParse
 	}
 
 	return &AuthMiddleware{
-		getAccountFromPAT:          getAccountFromPAT,
+		getAccountInfoFromPAT:      getAccountInfoFromPAT,
 		validateAndParseToken:      validateAndParseToken,
 		markPATUsed:                markPATUsed,
 		checkUserAccessByJWTGroups: checkUserAccessByJWTGroups,
@@ -151,13 +151,11 @@ func (m *AuthMiddleware) verifyUserAccess(ctx context.Context, validatedToken *j
 // CheckPATFromRequest checks if the PAT is valid
 func (m *AuthMiddleware) checkPATFromRequest(w http.ResponseWriter, r *http.Request, auth []string) error {
 	token, err := getTokenFromPATRequest(auth)
-
-	// If an error occurs, call the error handler and return an error
 	if err != nil {
-		return fmt.Errorf("Error extracting token: %w", err)
+		return fmt.Errorf("error extracting token: %w", err)
 	}
 
-	account, user, pat, err := m.getAccountFromPAT(r.Context(), token)
+	user, pat, accDomain, accCategory, err := m.getAccountInfoFromPAT(r.Context(), token)
 	if err != nil {
 		return fmt.Errorf("invalid Token: %w", err)
 	}
@@ -172,9 +170,9 @@ func (m *AuthMiddleware) checkPATFromRequest(w http.ResponseWriter, r *http.Requ
 
 	claimMaps := jwt.MapClaims{}
 	claimMaps[m.userIDClaim] = user.Id
-	claimMaps[m.audience+jwtclaims.AccountIDSuffix] = account.Id
-	claimMaps[m.audience+jwtclaims.DomainIDSuffix] = account.Domain
-	claimMaps[m.audience+jwtclaims.DomainCategorySuffix] = account.DomainCategory
+	claimMaps[m.audience+jwtclaims.AccountIDSuffix] = user.AccountID
+	claimMaps[m.audience+jwtclaims.DomainIDSuffix] = accDomain
+	claimMaps[m.audience+jwtclaims.DomainCategorySuffix] = accCategory
 	jwtToken := jwt.NewWithClaims(jwt.SigningMethodHS256, claimMaps)
 	newRequest := r.WithContext(context.WithValue(r.Context(), jwtclaims.TokenUserProperty, jwtToken)) //nolint
 	// Update the current request with the new context information.

management/server/http/middleware/auth_middleware_test.go

@@ -33,7 +33,8 @@ var testAccount = &server.Account{
 	Domain: domain,
 	Users: map[string]*server.User{
 		userID: {
-			Id: userID,
+			Id:        userID,
+			AccountID: accountID,
 			PATs: map[string]*server.PersonalAccessToken{
 				tokenID: {
 					ID:             tokenID,
@@ -49,11 +50,11 @@ var testAccount = &server.Account{
 	},
 }
 
-func mockGetAccountFromPAT(_ context.Context, token string) (*server.Account, *server.User, *server.PersonalAccessToken, error) {
+func mockGetAccountInfoFromPAT(_ context.Context, token string) (user *server.User, pat *server.PersonalAccessToken, domain string, category string, err error) {
 	if token == PAT {
-		return testAccount, testAccount.Users[userID], testAccount.Users[userID].PATs[tokenID], nil
+		return testAccount.Users[userID], testAccount.Users[userID].PATs[tokenID], testAccount.Domain, testAccount.DomainCategory, nil
 	}
-	return nil, nil, nil, fmt.Errorf("PAT invalid")
+	return nil, nil, "", "", fmt.Errorf("PAT invalid")
 }
 
 func mockValidateAndParseToken(_ context.Context, token string) (*jwt.Token, error) {
@@ -165,7 +166,7 @@ func TestAuthMiddleware_Handler(t *testing.T) {
 	)
 
 	authMiddleware := NewAuthMiddleware(
-		mockGetAccountFromPAT,
+		mockGetAccountInfoFromPAT,
 		mockValidateAndParseToken,
 		mockMarkPATUsed,
 		mockCheckUserAccessByJWTGroups,

management/server/http/peers_handler.go

@@ -48,8 +48,8 @@ func (h *PeersHandler) checkPeerStatus(peer *nbpeer.Peer) (*nbpeer.Peer, error)
 	return peerToReturn, nil
 }
 
-func (h *PeersHandler) getPeer(ctx context.Context, account *server.Account, peerID, userID string, w http.ResponseWriter) {
-	peer, err := h.accountManager.GetPeer(ctx, account.Id, peerID, userID)
+func (h *PeersHandler) getPeer(ctx context.Context, accountID, peerID, userID string, w http.ResponseWriter) {
+	peer, err := h.accountManager.GetPeer(ctx, accountID, peerID, userID)
 	if err != nil {
 		util.WriteError(ctx, err, w)
 		return
@@ -62,11 +62,16 @@ func (h *PeersHandler) getPeer(ctx context.Context, account *server.Account, pee
 	}
 	dnsDomain := h.accountManager.GetDNSDomain()
 
-	groupsInfo := toGroupsInfo(account.Groups, peer.ID)
-
-	validPeers, err := h.accountManager.GetValidatedPeers(account)
+	peerGroups, err := h.accountManager.GetPeerGroups(ctx, accountID, peer.ID)
 	if err != nil {
-		log.WithContext(ctx).Errorf("failed to list appreoved peers: %v", err)
+		util.WriteError(ctx, err, w)
+		return
+	}
+	groupsInfo := toGroupsInfo(peerGroups)
+
+	validPeers, err := h.accountManager.GetValidatedPeers(ctx, accountID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to list approved peers: %v", err)
 		util.WriteError(ctx, fmt.Errorf("internal error"), w)
 		return
 	}
@@ -75,7 +80,7 @@ func (h *PeersHandler) getPeer(ctx context.Context, account *server.Account, pee
 	util.WriteJSONObject(ctx, w, toSinglePeerResponse(peerToReturn, groupsInfo, dnsDomain, valid))
 }
 
-func (h *PeersHandler) updatePeer(ctx context.Context, account *server.Account, userID, peerID string, w http.ResponseWriter, r *http.Request) {
+func (h *PeersHandler) updatePeer(ctx context.Context, accountID, userID, peerID string, w http.ResponseWriter, r *http.Request) {
 	req := &api.PeerRequest{}
 	err := json.NewDecoder(r.Body).Decode(&req)
 	if err != nil {
@@ -99,16 +104,21 @@ func (h *PeersHandler) updatePeer(ctx context.Context, account *server.Account,
 		}
 	}
 
-	peer, err := h.accountManager.UpdatePeer(ctx, account.Id, userID, update)
+	peer, err := h.accountManager.UpdatePeer(ctx, accountID, userID, update)
 	if err != nil {
 		util.WriteError(ctx, err, w)
 		return
 	}
 	dnsDomain := h.accountManager.GetDNSDomain()
 
-	groupMinimumInfo := toGroupsInfo(account.Groups, peer.ID)
+	peerGroups, err := h.accountManager.GetPeerGroups(ctx, accountID, peer.ID)
+	if err != nil {
+		util.WriteError(ctx, err, w)
+		return
+	}
+	groupMinimumInfo := toGroupsInfo(peerGroups)
 
-	validPeers, err := h.accountManager.GetValidatedPeers(account)
+	validPeers, err := h.accountManager.GetValidatedPeers(ctx, accountID)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to list appreoved peers: %v", err)
 		util.WriteError(ctx, fmt.Errorf("internal error"), w)
@@ -149,18 +159,11 @@ func (h *PeersHandler) HandlePeer(w http.ResponseWriter, r *http.Request) {
 	case http.MethodDelete:
 		h.deletePeer(r.Context(), accountID, userID, peerID, w)
 		return
-	case http.MethodGet, http.MethodPut:
-		account, err := h.accountManager.GetAccountByID(r.Context(), accountID, userID)
-		if err != nil {
-			util.WriteError(r.Context(), err, w)
-			return
-		}
-
-		if r.Method == http.MethodGet {
-			h.getPeer(r.Context(), account, peerID, userID, w)
-		} else {
-			h.updatePeer(r.Context(), account, userID, peerID, w, r)
-		}
+	case http.MethodGet:
+		h.getPeer(r.Context(), accountID, peerID, userID, w)
+		return
+	case http.MethodPut:
+		h.updatePeer(r.Context(), accountID, userID, peerID, w, r)
 		return
 	default:
 		util.WriteError(r.Context(), status.Errorf(status.NotFound, "unknown METHOD"), w)
@@ -176,25 +179,13 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	account, err := h.accountManager.GetAccountByID(r.Context(), accountID, userID)
-	if err != nil {
-		util.WriteError(r.Context(), err, w)
-		return
-	}
-
-	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
-	}
+	dnsDomain := h.accountManager.GetDNSDomain()
 
 	respBody := make([]*api.PeerBatch, 0, len(peers))
 	for _, peer := range peers {
@@ -203,12 +194,18 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 			util.WriteError(r.Context(), err, w)
 			return
 		}
-		groupMinimumInfo := toGroupsInfo(groupsMap, peer.ID)
+
+		peerGroups, err := h.accountManager.GetPeerGroups(r.Context(), accountID, peer.ID)
+		if err != nil {
+			util.WriteError(r.Context(), err, w)
+			return
+		}
+		groupMinimumInfo := toGroupsInfo(peerGroups)
 
 		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, 0))
 	}
 
-	validPeersMap, err := h.accountManager.GetValidatedPeers(account)
+	validPeersMap, err := h.accountManager.GetValidatedPeers(r.Context(), accountID)
 	if err != nil {
 		log.WithContext(r.Context()).Errorf("failed to list appreoved peers: %v", err)
 		util.WriteError(r.Context(), fmt.Errorf("internal error"), w)
@@ -271,16 +268,16 @@ func (h *PeersHandler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request
 		}
 	}
 
-	dnsDomain := h.accountManager.GetDNSDomain()
-
-	validPeers, err := h.accountManager.GetValidatedPeers(account)
+	validPeers, err := h.accountManager.GetValidatedPeers(r.Context(), accountID)
 	if err != nil {
 		log.WithContext(r.Context()).Errorf("failed to list approved peers: %v", err)
 		util.WriteError(r.Context(), fmt.Errorf("internal error"), w)
 		return
 	}
 
-	customZone := account.GetPeersCustomZone(r.Context(), h.accountManager.GetDNSDomain())
+	dnsDomain := h.accountManager.GetDNSDomain()
+
+	customZone := account.GetPeersCustomZone(r.Context(), dnsDomain)
 	netMap := account.GetPeerNetworkMap(r.Context(), peerID, customZone, validPeers, nil)
 
 	util.WriteJSONObject(r.Context(), w, toAccessiblePeers(netMap, dnsDomain))
@@ -315,26 +312,14 @@ func peerToAccessiblePeer(peer *nbpeer.Peer, dnsDomain string) api.AccessiblePee
 	}
 }
 
-func toGroupsInfo(groups map[string]*nbgroup.Group, peerID string) []api.GroupMinimum {
-	groupsInfo := []api.GroupMinimum{}
-	groupsChecked := make(map[string]struct{})
+func toGroupsInfo(groups []*nbgroup.Group) []api.GroupMinimum {
+	groupsInfo := make([]api.GroupMinimum, 0, len(groups))
 	for _, group := range groups {
-		_, ok := groupsChecked[group.ID]
-		if ok {
-			continue
-		}
-		groupsChecked[group.ID] = struct{}{}
-		for _, pk := range group.Peers {
-			if pk == peerID {
-				info := api.GroupMinimum{
-					Id:         group.ID,
-					Name:       group.Name,
-					PeersCount: len(group.Peers),
-				}
-				groupsInfo = append(groupsInfo, info)
-				break
-			}
-		}
+		groupsInfo = append(groupsInfo, api.GroupMinimum{
+			Id:         group.ID,
+			Name:       group.Name,
+			PeersCount: len(group.Peers),
+		})
 	}
 	return groupsInfo
 }

management/server/http/peers_handler_test.go

@@ -39,6 +39,68 @@ const (
 )
 
 func initTestMetaData(peers ...*nbpeer.Peer) *PeersHandler {
+
+	peersMap := make(map[string]*nbpeer.Peer)
+	for _, peer := range peers {
+		peersMap[peer.ID] = peer.Copy()
+	}
+
+	policy := &server.Policy{
+		ID:        "policy",
+		AccountID: "test_id",
+		Name:      "policy",
+		Enabled:   true,
+		Rules: []*server.PolicyRule{
+			{
+				ID:            "rule",
+				Name:          "rule",
+				Enabled:       true,
+				Action:        "accept",
+				Destinations:  []string{"group1"},
+				Sources:       []string{"group1"},
+				Bidirectional: true,
+				Protocol:      "all",
+				Ports:         []string{"80"},
+			},
+		},
+	}
+
+	srvUser := server.NewRegularUser(serviceUser)
+	srvUser.IsServiceUser = true
+
+	account := &server.Account{
+		Id:     "test_id",
+		Domain: "hotmail.com",
+		Peers:  peersMap,
+		Users: map[string]*server.User{
+			adminUser:   server.NewAdminUser(adminUser),
+			regularUser: server.NewRegularUser(regularUser),
+			serviceUser: srvUser,
+		},
+		Groups: map[string]*nbgroup.Group{
+			"group1": {
+				ID:        "group1",
+				AccountID: "test_id",
+				Name:      "group1",
+				Issued:    "api",
+				Peers:     maps.Keys(peersMap),
+			},
+		},
+		Settings: &server.Settings{
+			PeerLoginExpirationEnabled: true,
+			PeerLoginExpiration:        time.Hour,
+		},
+		Policies: []*server.Policy{policy},
+		Network: &server.Network{
+			Identifier: "ciclqisab2ss43jdn8q0",
+			Net: net.IPNet{
+				IP:   net.ParseIP("100.67.0.0"),
+				Mask: net.IPv4Mask(255, 255, 0, 0),
+			},
+			Serial: 51,
+		},
+	}
+
 	return &PeersHandler{
 		accountManager: &mock_server.MockAccountManager{
 			UpdatePeerFunc: func(_ context.Context, accountID, userID string, update *nbpeer.Peer) (*nbpeer.Peer, error) {
@@ -67,74 +129,31 @@ func initTestMetaData(peers ...*nbpeer.Peer) *PeersHandler {
 			GetPeersFunc: func(_ context.Context, accountID, userID string) ([]*nbpeer.Peer, error) {
 				return peers, nil
 			},
+			GetPeerGroupsFunc: func(ctx context.Context, accountID, peerID string) ([]*nbgroup.Group, error) {
+				peersID := make([]string, len(peers))
+				for _, peer := range peers {
+					peersID = append(peersID, peer.ID)
+				}
+				return []*nbgroup.Group{
+					{
+						ID:        "group1",
+						AccountID: accountID,
+						Name:      "group1",
+						Issued:    "api",
+						Peers:     peersID,
+					},
+				}, nil
+			},
 			GetDNSDomainFunc: func() string {
 				return "netbird.selfhosted"
 			},
 			GetAccountIDFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (string, string, error) {
 				return claims.AccountId, claims.UserId, nil
 			},
+			GetAccountFunc: func(ctx context.Context, accountID string) (*server.Account, error) {
+				return account, nil
+			},
 			GetAccountByIDFunc: func(ctx context.Context, accountID string, userID string) (*server.Account, error) {
-				peersMap := make(map[string]*nbpeer.Peer)
-				for _, peer := range peers {
-					peersMap[peer.ID] = peer.Copy()
-				}
-
-				policy := &server.Policy{
-					ID:        "policy",
-					AccountID: accountID,
-					Name:      "policy",
-					Enabled:   true,
-					Rules: []*server.PolicyRule{
-						{
-							ID:            "rule",
-							Name:          "rule",
-							Enabled:       true,
-							Action:        "accept",
-							Destinations:  []string{"group1"},
-							Sources:       []string{"group1"},
-							Bidirectional: true,
-							Protocol:      "all",
-							Ports:         []string{"80"},
-						},
-					},
-				}
-
-				srvUser := server.NewRegularUser(serviceUser)
-				srvUser.IsServiceUser = true
-
-				account := &server.Account{
-					Id:     accountID,
-					Domain: "hotmail.com",
-					Peers:  peersMap,
-					Users: map[string]*server.User{
-						adminUser:   server.NewAdminUser(adminUser),
-						regularUser: server.NewRegularUser(regularUser),
-						serviceUser: srvUser,
-					},
-					Groups: map[string]*nbgroup.Group{
-						"group1": {
-							ID:        "group1",
-							AccountID: accountID,
-							Name:      "group1",
-							Issued:    "api",
-							Peers:     maps.Keys(peersMap),
-						},
-					},
-					Settings: &server.Settings{
-						PeerLoginExpirationEnabled: true,
-						PeerLoginExpiration:        time.Hour,
-					},
-					Policies: []*server.Policy{policy},
-					Network: &server.Network{
-						Identifier: "ciclqisab2ss43jdn8q0",
-						Net: net.IPNet{
-							IP:   net.ParseIP("100.67.0.0"),
-							Mask: net.IPv4Mask(255, 255, 0, 0),
-						},
-						Serial: 51,
-					},
-				}
-
 				return account, nil
 			},
 			HasConnectedChannelFunc: func(peerID string) bool {

management/server/http/policies_handler.go

@@ -128,13 +128,8 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
 		Description: req.Description,
 	}
 	for _, rule := range req.Rules {
-		var ruleID string
-		if rule.Id != nil {
-			ruleID = *rule.Id
-		}
-
 		pr := server.PolicyRule{
-			ID:            ruleID,
+			ID:            policyID, // TODO: when policy can contain multiple rules, need refactor
 			PolicyID:      policyID,
 			Name:          rule.Name,
 			Destinations:  rule.Destinations,

management/server/http/setupkeys_handler.go

@@ -137,6 +137,11 @@ func (h *SetupKeysHandler) UpdateSetupKey(w http.ResponseWriter, r *http.Request
 		return
 	}
 
+	if req.Name == "" {
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "setup key name field is invalid: %s", req.Name), w)
+		return
+	}
+
 	if req.AutoGroups == nil {
 		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "setup key AutoGroups field is invalid"), w)
 		return
@@ -145,6 +150,7 @@ func (h *SetupKeysHandler) UpdateSetupKey(w http.ResponseWriter, r *http.Request
 	newKey := &server.SetupKey{}
 	newKey.AutoGroups = req.AutoGroups
 	newKey.Revoked = req.Revoked
+	newKey.Name = req.Name
 	newKey.Id = keyID
 
 	newKey, err = h.accountManager.SaveSetupKey(r.Context(), accountID, newKey, userID)

management/server/integrated_validator.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/account"
@@ -73,6 +75,39 @@ func (am *DefaultAccountManager) GroupValidation(ctx context.Context, accountID
 	return true, nil
 }
 
-func (am *DefaultAccountManager) GetValidatedPeers(account *Account) (map[string]struct{}, error) {
-	return am.integratedPeerValidator.GetValidatedPeers(account.Id, account.Groups, account.Peers, account.Settings.Extra)
+func (am *DefaultAccountManager) GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, error) {
+	var err error
+	var groups []*nbgroup.Group
+	var peers []*nbpeer.Peer
+	var settings *Settings
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		groups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return err
+		}
+
+		peers, err = transaction.GetAccountPeers(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return err
+		}
+
+		settings, err = transaction.GetAccountSettings(ctx, LockingStrengthShare, accountID)
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	groupsMap := make(map[string]*nbgroup.Group, len(groups))
+	for _, group := range groups {
+		groupsMap[group.ID] = group
+	}
+
+	peersMap := make(map[string]*nbpeer.Peer, len(peers))
+	for _, peer := range peers {
+		peersMap[peer.ID] = peer
+	}
+
+	return am.integratedPeerValidator.GetValidatedPeers(accountID, groupsMap, peersMap, settings.Extra)
 }

management/server/jwtclaims/jwtValidator.go

@@ -77,8 +77,6 @@ type JWTValidator struct {
 	options Options
 }
 
-var keyNotFound = errors.New("unable to find appropriate key")
-
 // NewJWTValidator constructor
 func NewJWTValidator(ctx context.Context, issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) (*JWTValidator, error) {
 	keys, err := getPemKeys(ctx, keysLocation)
@@ -126,18 +124,12 @@ func NewJWTValidator(ctx context.Context, issuer string, audienceList []string,
 			}
 
 			publicKey, err := getPublicKey(ctx, token, keys)
-			if err == nil {
-				return publicKey, nil
+			if err != nil {
+				log.WithContext(ctx).Errorf("getPublicKey error: %s", err)
+				return nil, err
 			}
 
-			msg := fmt.Sprintf("getPublicKey error: %s", err)
-			if errors.Is(err, keyNotFound) && !idpSignkeyRefreshEnabled {
-				msg = fmt.Sprintf("getPublicKey error: %s. You can enable key refresh by setting HttpServerConfig.IdpSignKeyRefreshEnabled to true in your management.json file and restart the service", err)
-			}
-
-			log.WithContext(ctx).Error(msg)
-
-			return nil, err
+			return publicKey, nil
 		},
 		EnableAuthOnOptions: false,
 	}
@@ -237,7 +229,7 @@ func getPublicKey(ctx context.Context, token *jwt.Token, jwks *Jwks) (interface{
 		log.WithContext(ctx).Debugf("Key Type: %s not yet supported, please raise ticket!", jwks.Keys[k].Kty)
 	}
 
-	return nil, keyNotFound
+	return nil, errors.New("unable to find appropriate key")
 }
 
 func getPublicKeyFromECDSA(jwk JSONWebKey) (publicKey *ecdsa.PublicKey, err error) {
@@ -318,3 +310,4 @@ func getMaxAgeFromCacheHeader(ctx context.Context, cacheControl string) int {
 
 	return 0
 }
+

management/server/mock_server/account_mock.go

@@ -22,9 +22,9 @@ import (
 )
 
 type MockAccountManager struct {
-	GetOrCreateAccountByUserFunc func(ctx context.Context, userId, domain string) (*server.Account, error)
-	GetAccountFunc               func(ctx context.Context, accountID string) (*server.Account, error)
-	CreateSetupKeyFunc           func(ctx context.Context, accountId string, keyName string, keyType server.SetupKeyType,
+	GetOrCreateAccountIDByUserFunc func(ctx context.Context, userId, domain string) (string, error)
+	GetAccountFunc                 func(ctx context.Context, accountID string) (*server.Account, error)
+	CreateSetupKeyFunc             func(ctx context.Context, accountId string, keyName string, keyType server.SetupKeyType,
 		expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool) (*server.SetupKey, error)
 	GetSetupKeyFunc                     func(ctx context.Context, accountID, userID, keyID string) (*server.SetupKey, error)
 	AccountExistsFunc                   func(ctx context.Context, accountID string) (bool, error)
@@ -47,13 +47,14 @@ type MockAccountManager struct {
 	DeleteGroupsFunc                    func(ctx context.Context, accountId, userId string, groupIDs []string) error
 	GroupAddPeerFunc                    func(ctx context.Context, accountID, groupID, peerID string) error
 	GroupDeletePeerFunc                 func(ctx context.Context, accountID, groupID, peerID string) error
+	GetPeerGroupsFunc                   func(ctx context.Context, accountID, peerID string) ([]*group.Group, error)
 	DeleteRuleFunc                      func(ctx context.Context, accountID, ruleID, userID string) error
 	GetPolicyFunc                       func(ctx context.Context, accountID, policyID, userID string) (*server.Policy, error)
 	SavePolicyFunc                      func(ctx context.Context, accountID, userID string, policy *server.Policy) (*server.Policy, error)
 	DeletePolicyFunc                    func(ctx context.Context, accountID, policyID, userID string) error
 	ListPoliciesFunc                    func(ctx context.Context, accountID, userID string) ([]*server.Policy, error)
 	GetUsersFromAccountFunc             func(ctx context.Context, accountID, userID string) ([]*server.UserInfo, error)
-	GetAccountFromPATFunc               func(ctx context.Context, pat string) (*server.Account, *server.User, *server.PersonalAccessToken, error)
+	GetAccountInfoFromPATFunc           func(ctx context.Context, token string) (*server.User, *server.PersonalAccessToken, string, string, error)
 	MarkPATUsedFunc                     func(ctx context.Context, pat string) error
 	UpdatePeerMetaFunc                  func(ctx context.Context, peerID string, meta nbpeer.PeerSystemMeta) error
 	UpdatePeerFunc                      func(ctx context.Context, accountID, userID string, peer *nbpeer.Peer) (*nbpeer.Peer, error)
@@ -88,9 +89,9 @@ type MockAccountManager struct {
 	GetDNSSettingsFunc                  func(ctx context.Context, accountID, userID string) (*server.DNSSettings, error)
 	SaveDNSSettingsFunc                 func(ctx context.Context, accountID, userID string, dnsSettingsToSave *server.DNSSettings) error
 	GetPeerFunc                         func(ctx context.Context, accountID, peerID, userID string) (*nbpeer.Peer, error)
-	UpdateAccountSettingsFunc           func(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Account, error)
+	UpdateAccountSettingsFunc           func(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Settings, error)
 	LoginPeerFunc                       func(ctx context.Context, login server.PeerLogin) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
-	SyncPeerFunc                        func(ctx context.Context, sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
+	SyncPeerFunc                        func(ctx context.Context, sync server.PeerSync, accountID string) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
 	InviteUserFunc                      func(ctx context.Context, accountID string, initiatorUserID string, targetUserEmail string) error
 	GetAllConnectedPeersFunc            func() (map[string]struct{}, error)
 	HasConnectedChannelFunc             func(peerID string) bool
@@ -130,7 +131,12 @@ func (am *MockAccountManager) OnPeerDisconnected(_ context.Context, accountID st
 	panic("implement me")
 }
 
-func (am *MockAccountManager) GetValidatedPeers(account *server.Account) (map[string]struct{}, error) {
+func (am *MockAccountManager) GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, error) {
+	account, err := am.GetAccountFunc(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+
 	approvedPeers := make(map[string]struct{})
 	for id := range account.Peers {
 		approvedPeers[id] = struct{}{}
@@ -170,16 +176,16 @@ func (am *MockAccountManager) DeletePeer(ctx context.Context, accountID, peerID,
 	return status.Errorf(codes.Unimplemented, "method DeletePeer is not implemented")
 }
 
-// GetOrCreateAccountByUser mock implementation of GetOrCreateAccountByUser from server.AccountManager interface
-func (am *MockAccountManager) GetOrCreateAccountByUser(
+// GetOrCreateAccountIDByUser mock implementation of GetOrCreateAccountByUser from server.AccountManager interface
+func (am *MockAccountManager) GetOrCreateAccountIDByUser(
 	ctx context.Context, userId, domain string,
-) (*server.Account, error) {
-	if am.GetOrCreateAccountByUserFunc != nil {
-		return am.GetOrCreateAccountByUserFunc(ctx, userId, domain)
+) (string, error) {
+	if am.GetOrCreateAccountIDByUserFunc != nil {
+		return am.GetOrCreateAccountIDByUserFunc(ctx, userId, domain)
 	}
-	return nil, status.Errorf(
+	return "", status.Errorf(
 		codes.Unimplemented,
-		"method GetOrCreateAccountByUser is not implemented",
+		"method GetOrCreateAccountIDByUser is not implemented",
 	)
 }
 
@@ -221,19 +227,19 @@ func (am *MockAccountManager) GetAccountIDByUserID(ctx context.Context, userId,
 }
 
 // MarkPeerConnected mock implementation of MarkPeerConnected from server.AccountManager interface
-func (am *MockAccountManager) MarkPeerConnected(ctx context.Context, peerKey string, connected bool, realIP net.IP, account *server.Account) error {
+func (am *MockAccountManager) MarkPeerConnected(ctx context.Context, peerKey string, connected bool, realIP net.IP, accountID string) error {
 	if am.MarkPeerConnectedFunc != nil {
 		return am.MarkPeerConnectedFunc(ctx, peerKey, connected, realIP)
 	}
 	return status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
 }
 
-// GetAccountFromPAT mock implementation of GetAccountFromPAT from server.AccountManager interface
-func (am *MockAccountManager) GetAccountFromPAT(ctx context.Context, pat string) (*server.Account, *server.User, *server.PersonalAccessToken, error) {
-	if am.GetAccountFromPATFunc != nil {
-		return am.GetAccountFromPATFunc(ctx, pat)
+// GetAccountInfoFromPAT mock implementation of GetAccountInfoFromPAT from server.AccountManager interface
+func (am *MockAccountManager) GetAccountInfoFromPAT(ctx context.Context, pat string) (*server.User, *server.PersonalAccessToken, string, string, error) {
+	if am.GetAccountInfoFromPATFunc != nil {
+		return am.GetAccountInfoFromPATFunc(ctx, pat)
 	}
-	return nil, nil, nil, status.Errorf(codes.Unimplemented, "method GetAccountFromPAT is not implemented")
+	return nil, nil, "", "", status.Errorf(codes.Unimplemented, "method GetAccountInfoFromPAT is not implemented")
 }
 
 // DeleteAccount mock implementation of DeleteAccount from server.AccountManager interface
@@ -666,7 +672,7 @@ func (am *MockAccountManager) GetPeer(ctx context.Context, accountID, peerID, us
 }
 
 // UpdateAccountSettings mocks UpdateAccountSettings of the AccountManager interface
-func (am *MockAccountManager) UpdateAccountSettings(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
+func (am *MockAccountManager) UpdateAccountSettings(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Settings, error) {
 	if am.UpdateAccountSettingsFunc != nil {
 		return am.UpdateAccountSettingsFunc(ctx, accountID, userID, newSettings)
 	}
@@ -682,9 +688,9 @@ func (am *MockAccountManager) LoginPeer(ctx context.Context, login server.PeerLo
 }
 
 // SyncPeer mocks SyncPeer of the AccountManager interface
-func (am *MockAccountManager) SyncPeer(ctx context.Context, sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
+func (am *MockAccountManager) SyncPeer(ctx context.Context, sync server.PeerSync, accountID string) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
 	if am.SyncPeerFunc != nil {
-		return am.SyncPeerFunc(ctx, sync, account)
+		return am.SyncPeerFunc(ctx, sync, accountID)
 	}
 	return nil, nil, nil, status.Errorf(codes.Unimplemented, "method SyncPeer is not implemented")
 }
@@ -831,3 +837,11 @@ func (am *MockAccountManager) GetAccount(ctx context.Context, accountID string)
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method GetAccount is not implemented")
 }
+
+// GetPeerGroups mocks GetPeerGroups of the AccountManager interface
+func (am *MockAccountManager) GetPeerGroups(ctx context.Context, accountID, peerID string) ([]*group.Group, error) {
+	if am.GetPeerGroupsFunc != nil {
+		return am.GetPeerGroupsFunc(ctx, accountID, peerID)
+	}
+	return nil, status.Errorf(codes.Unimplemented, "method GetPeerGroups is not implemented")
+}

management/server/nameserver_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -381,14 +382,14 @@ func TestCreateNameServerGroup(t *testing.T) {
 				t.Error("failed to create account manager")
 			}
 
-			account, err := initTestNSAccount(t, am)
+			accountID, err := initTestNSAccount(t, am)
 			if err != nil {
 				t.Error("failed to init testing account")
 			}
 
 			outNSGroup, err := am.CreateNameServerGroup(
 				context.Background(),
-				account.Id,
+				accountID,
 				testCase.inputArgs.name,
 				testCase.inputArgs.description,
 				testCase.inputArgs.nameServers,
@@ -609,20 +610,16 @@ func TestSaveNameServerGroup(t *testing.T) {
 				t.Error("failed to create account manager")
 			}
 
-			account, err := initTestNSAccount(t, am)
+			accountID, err := initTestNSAccount(t, am)
 			if err != nil {
 				t.Error("failed to init testing account")
 			}
 
-			account.NameServerGroups[testCase.existingNSGroup.ID] = testCase.existingNSGroup
-
-			err = am.Store.SaveAccount(context.Background(), account)
-			if err != nil {
-				t.Error("account should be saved")
-			}
+			testCase.existingNSGroup.AccountID = accountID
+			err = am.Store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, testCase.existingNSGroup)
+			require.NoError(t, err, "failed to save existing nameserver group")
 
 			var nsGroupToSave *nbdns.NameServerGroup
-
 			if !testCase.skipCopying {
 				nsGroupToSave = testCase.existingNSGroup.Copy()
 
@@ -651,22 +648,17 @@ func TestSaveNameServerGroup(t *testing.T) {
 				}
 			}
 
-			err = am.SaveNameServerGroup(context.Background(), account.Id, userID, nsGroupToSave)
-
+			err = am.SaveNameServerGroup(context.Background(), accountID, userID, nsGroupToSave)
 			testCase.errFunc(t, err)
 
 			if !testCase.shouldCreate {
 				return
 			}
 
-			account, err = am.Store.GetAccount(context.Background(), account.Id)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			savedNSGroup, saved := account.NameServerGroups[testCase.expectedNSGroup.ID]
-			require.True(t, saved)
+			savedNSGroup, err := am.Store.GetNameServerGroupByID(context.Background(), LockingStrengthShare, accountID, testCase.expectedNSGroup.ID)
+			require.NoError(t, err, "failed to get saved nameserver group")
 
+			testCase.expectedNSGroup.AccountID = accountID
 			if !testCase.expectedNSGroup.IsEqual(savedNSGroup) {
 				t.Errorf("new nameserver group didn't match expected group:\nGot %#v\nExpected:%#v\n", savedNSGroup, testCase.expectedNSGroup)
 			}
@@ -703,32 +695,25 @@ func TestDeleteNameServerGroup(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestNSAccount(t, am)
+	accountID, err := initTestNSAccount(t, am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
-	account.NameServerGroups[testingNSGroup.ID] = testingNSGroup
+	testingNSGroup.AccountID = accountID
+	err = am.Store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, testingNSGroup)
+	require.NoError(t, err, "failed to save nameserver group")
 
-	err = am.Store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Error("failed to save account")
-	}
-
-	err = am.DeleteNameServerGroup(context.Background(), account.Id, testingNSGroup.ID, userID)
+	err = am.DeleteNameServerGroup(context.Background(), accountID, testingNSGroup.ID, userID)
 	if err != nil {
 		t.Error("deleting nameserver group failed with error: ", err)
 	}
 
-	savedAccount, err := am.Store.GetAccount(context.Background(), account.Id)
-	if err != nil {
-		t.Error("failed to retrieve saved account with error: ", err)
-	}
-
-	_, found := savedAccount.NameServerGroups[testingNSGroup.ID]
-	if found {
-		t.Error("nameserver group shouldn't be found after delete")
-	}
+	_, err = am.Store.GetNameServerGroupByID(context.Background(), LockingStrengthShare, accountID, testingNSGroup.ID)
+	require.NotNil(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok, "error should be a status error")
+	assert.Equal(t, status.NotFound, sErr.Type(), "nameserver group shouldn't be found after delete")
 }
 
 func TestGetNameServerGroup(t *testing.T) {
@@ -738,12 +723,12 @@ func TestGetNameServerGroup(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestNSAccount(t, am)
+	accountID, err := initTestNSAccount(t, am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
-	foundGroup, err := am.GetNameServerGroup(context.Background(), account.Id, testUserID, existingNSGroupID)
+	foundGroup, err := am.GetNameServerGroup(context.Background(), accountID, testUserID, existingNSGroupID)
 	if err != nil {
 		t.Error("getting existing nameserver group failed with error: ", err)
 	}
@@ -752,7 +737,7 @@ func TestGetNameServerGroup(t *testing.T) {
 		t.Error("got a nil group while getting nameserver group with ID")
 	}
 
-	_, err = am.GetNameServerGroup(context.Background(), account.Id, testUserID, "not existing")
+	_, err = am.GetNameServerGroup(context.Background(), accountID, testUserID, "not existing")
 	if err == nil {
 		t.Error("getting not existing nameserver group should return error, got nil")
 	}
@@ -784,8 +769,12 @@ func createNSStore(t *testing.T) (Store, error) {
 	return store, nil
 }
 
-func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error) {
+func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (string, error) {
 	t.Helper()
+	accountID := "testingAcc"
+	userID := testUserID
+	domain := "example.com"
+
 	peer1 := &nbpeer.Peer{
 		Key:  nsGroupPeer1Key,
 		Name: "test-host1@netbird.io",
@@ -816,6 +805,7 @@ func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error
 	}
 	existingNSGroup := nbdns.NameServerGroup{
 		ID:          existingNSGroupID,
+		AccountID:   accountID,
 		Name:        existingNSGroupName,
 		Description: "",
 		NameServers: []nbdns.NameServer{
@@ -834,42 +824,42 @@ func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error
 		Enabled: true,
 	}
 
-	accountID := "testingAcc"
-	userID := testUserID
-	domain := "example.com"
-
-	account := newAccountWithId(context.Background(), accountID, userID, domain)
-
-	account.NameServerGroups[existingNSGroup.ID] = &existingNSGroup
-
-	newGroup1 := &nbgroup.Group{
-		ID:   group1ID,
-		Name: group1ID,
-	}
-
-	newGroup2 := &nbgroup.Group{
-		ID:   group2ID,
-		Name: group2ID,
-	}
-
-	account.Groups[newGroup1.ID] = newGroup1
-	account.Groups[newGroup2.ID] = newGroup2
-
-	err := am.Store.SaveAccount(context.Background(), account)
+	err := newAccountWithId(context.Background(), am.Store, accountID, userID, domain)
 	if err != nil {
-		return nil, err
+		return "", err
+	}
+
+	err = am.Store.SaveNameServerGroup(context.Background(), LockingStrengthUpdate, &existingNSGroup)
+	if err != nil {
+		return "", err
+	}
+
+	err = am.Store.SaveGroups(context.Background(), LockingStrengthUpdate, []*nbgroup.Group{
+		{
+			ID:        group1ID,
+			AccountID: accountID,
+			Name:      group1ID,
+		},
+		{
+			ID:        group2ID,
+			AccountID: accountID,
+			Name:      group2ID,
+		},
+	})
+	if err != nil {
+		return "", err
 	}
 
 	_, _, _, err = am.AddPeer(context.Background(), "", userID, peer1)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	_, _, _, err = am.AddPeer(context.Background(), "", userID, peer2)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	return account, nil
+	return accountID, nil
 }
 
 func TestValidateDomain(t *testing.T) {

management/server/peer/peer.go

@@ -44,7 +44,7 @@ type Peer struct {
 	// CreatedAt records the time the peer was created
 	CreatedAt time.Time
 	// Indicate ephemeral peer attribute
-	Ephemeral bool
+	Ephemeral bool `gorm:"index"`
 	// Geo location based on connection IP
 	Location Location `gorm:"embedded;embeddedPrefix:location_"`
 }

management/server/peer_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	nbAccount "github.com/netbirdio/netbird/management/server/account"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
@@ -22,7 +23,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"
@@ -468,21 +468,25 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 	accountID := "test_account"
 	adminUser := "account_creator"
 	someUser := "some_user"
-	account := newAccountWithId(context.Background(), accountID, adminUser, "")
-	account.Users[someUser] = &User{
-		Id:   someUser,
-		Role: UserRoleUser,
-	}
-	account.Settings.RegularUsersViewBlocked = false
+	err = newAccountWithId(context.Background(), manager.Store, accountID, adminUser, "")
+	require.NoError(t, err, "failed to create account")
 
-	err = manager.Store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Fatal(err)
-		return
-	}
+	err = manager.Store.SaveUser(context.Background(), LockingStrengthUpdate, &User{
+		Id:        someUser,
+		AccountID: accountID,
+		Role:      UserRoleUser,
+	})
+	require.NoError(t, err, "failed to create user")
+
+	settings, err := manager.Store.GetAccountSettings(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account settings")
+
+	settings.RegularUsersViewBlocked = false
+	err = manager.Store.SaveAccountSettings(context.Background(), LockingStrengthUpdate, accountID, settings)
+	require.NoError(t, err, "failed to save account settings")
 
 	// two peers one added by a regular user and one with a setup key
-	setupKey, err := manager.CreateSetupKey(context.Background(), account.Id, "test-key", SetupKeyReusable, time.Hour, nil, 999, adminUser, false)
+	setupKey, err := manager.CreateSetupKey(context.Background(), accountID, "test-key", SetupKeyReusable, time.Hour, nil, 999, adminUser, false)
 	if err != nil {
 		t.Fatal("error creating setup key")
 		return
@@ -536,7 +540,10 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
 	assert.NotNil(t, peer)
 
 	// delete the all-to-all policy so that user's peer1 has no access to peer2
-	for _, policy := range account.Policies {
+	accountPolicies, err := manager.Store.GetAccountPolicies(context.Background(), LockingStrengthShare, accountID)
+	require.NoError(t, err, "failed to get account policies")
+
+	for _, policy := range accountPolicies {
 		err = manager.DeletePolicy(context.Background(), accountID, policy.ID, adminUser)
 		if err != nil {
 			t.Fatal(err)
@@ -655,21 +662,33 @@ func TestDefaultAccountManager_GetPeers(t *testing.T) {
 			accountID := "test_account"
 			adminUser := "account_creator"
 			someUser := "some_user"
-			account := newAccountWithId(context.Background(), accountID, adminUser, "")
-			account.Users[someUser] = &User{
+
+			err = newAccountWithId(context.Background(), manager.Store, accountID, adminUser, "")
+			require.NoError(t, err, "failed to create account")
+
+			err = manager.Store.SaveUser(context.Background(), LockingStrengthUpdate, &User{
 				Id:            someUser,
+				AccountID:     accountID,
 				Role:          testCase.role,
 				IsServiceUser: testCase.isServiceUser,
-			}
-			account.Policies = []*Policy{}
-			account.Settings.RegularUsersViewBlocked = testCase.limitedViewSettings
+			})
+			require.NoError(t, err, "failed to create user")
 
-			err = manager.Store.SaveAccount(context.Background(), account)
-			if err != nil {
-				t.Fatal(err)
-				return
+			accountPolicies, err := manager.Store.GetAccountPolicies(context.Background(), LockingStrengthShare, accountID)
+			require.NoError(t, err, "failed to get account policies")
+
+			for _, policy := range accountPolicies {
+				err = manager.DeletePolicy(context.Background(), accountID, policy.ID, adminUser)
+				require.NoError(t, err, "failed to delete policy")
 			}
 
+			settings, err := manager.Store.GetAccountSettings(context.Background(), LockingStrengthShare, accountID)
+			require.NoError(t, err, "failed to get account settings")
+
+			settings.RegularUsersViewBlocked = testCase.limitedViewSettings
+			err = manager.Store.SaveAccountSettings(context.Background(), LockingStrengthUpdate, accountID, settings)
+			require.NoError(t, err, "failed to save account settings")
+
 			peerKey1, err := wgtypes.GeneratePrivateKey()
 			if err != nil {
 				t.Fatal(err)
@@ -725,10 +744,18 @@ func setupTestAccountManager(b *testing.B, peers int, groups int) (*DefaultAccou
 	adminUser := "account_creator"
 	regularUser := "regular_user"
 
-	account := newAccountWithId(context.Background(), accountID, adminUser, "")
-	account.Users[regularUser] = &User{
-		Id:   regularUser,
-		Role: UserRoleUser,
+	err = newAccountWithId(context.Background(), manager.Store, accountID, adminUser, "")
+	if err != nil {
+		return nil, "", "", err
+	}
+
+	err = manager.Store.SaveUser(context.Background(), LockingStrengthUpdate, &User{
+		Id:        regularUser,
+		AccountID: accountID,
+		Role:      UserRoleUser,
+	})
+	if err != nil {
+		return nil, "", "", err
 	}
 
 	// Create peers
@@ -742,31 +769,40 @@ func setupTestAccountManager(b *testing.B, peers int, groups int) (*DefaultAccou
 			Status:   &nbpeer.PeerStatus{},
 			UserID:   regularUser,
 		}
-		account.Peers[peer.ID] = peer
+		err = manager.Store.SavePeer(context.Background(), LockingStrengthUpdate, accountID, peer)
+		if err != nil {
+			return nil, "", "", err
+		}
 	}
 
 	// Create groups and policies
-	account.Policies = make([]*Policy, 0, groups)
 	for i := 0; i < groups; i++ {
 		groupID := fmt.Sprintf("group-%d", i)
 		group := &nbgroup.Group{
-			ID:   groupID,
-			Name: fmt.Sprintf("Group %d", i),
+			ID:        groupID,
+			AccountID: accountID,
+			Name:      fmt.Sprintf("Group %d", i),
 		}
 		for j := 0; j < peers/groups; j++ {
 			peerIndex := i*(peers/groups) + j
 			group.Peers = append(group.Peers, fmt.Sprintf("peer-%d", peerIndex))
 		}
-		account.Groups[groupID] = group
+
+		err = manager.Store.SaveGroup(context.Background(), LockingStrengthUpdate, group)
+		if err != nil {
+			return nil, "", "", err
+		}
 
 		// Create a policy for this group
 		policy := &Policy{
-			ID:      fmt.Sprintf("policy-%d", i),
-			Name:    fmt.Sprintf("Policy for Group %d", i),
-			Enabled: true,
+			ID:        fmt.Sprintf("policy-%d", i),
+			AccountID: accountID,
+			Name:      fmt.Sprintf("Policy for Group %d", i),
+			Enabled:   true,
 			Rules: []*PolicyRule{
 				{
 					ID:            fmt.Sprintf("rule-%d", i),
+					PolicyID:      fmt.Sprintf("policy-%d", i),
 					Name:          fmt.Sprintf("Rule for Group %d", i),
 					Enabled:       true,
 					Sources:       []string{groupID},
@@ -777,22 +813,23 @@ func setupTestAccountManager(b *testing.B, peers int, groups int) (*DefaultAccou
 				},
 			},
 		}
-		account.Policies = append(account.Policies, policy)
+
+		err = manager.Store.SavePolicy(context.Background(), LockingStrengthUpdate, policy)
+		if err != nil {
+			return nil, "", "", err
+		}
 	}
 
-	account.PostureChecks = []*posture.Checks{
-		{
-			ID:   "PostureChecksAll",
-			Name: "All",
-			Checks: posture.ChecksDefinition{
-				NBVersionCheck: &posture.NBVersionCheck{
-					MinVersion: "0.0.1",
-				},
+	err = manager.Store.SavePostureChecks(context.Background(), LockingStrengthUpdate, &posture.Checks{
+		ID:        "PostureChecksAll",
+		AccountID: accountID,
+		Name:      "All",
+		Checks: posture.ChecksDefinition{
+			NBVersionCheck: &posture.NBVersionCheck{
+				MinVersion: "0.0.1",
 			},
 		},
-	}
-
-	err = manager.Store.SaveAccount(context.Background(), account)
+	})
 	if err != nil {
 		return nil, "", "", err
 	}
@@ -833,23 +870,19 @@ func BenchmarkGetPeers(b *testing.B) {
 		})
 	}
 }
+
 func BenchmarkUpdateAccountPeers(b *testing.B) {
 	benchCases := []struct {
 		name   string
 		peers  int
 		groups int
-		// We need different expectations for CI/CD and local runs because of the different performance characteristics
-		minMsPerOpLocal float64
-		maxMsPerOpLocal float64
-		minMsPerOpCICD  float64
-		maxMsPerOpCICD  float64
 	}{
-		{"Small", 50, 5, 90, 120, 90, 120},
-		{"Medium", 500, 100, 110, 140, 120, 200},
-		{"Large", 5000, 200, 800, 1300, 2500, 3600},
-		{"Small single", 50, 10, 90, 120, 90, 120},
-		{"Medium single", 500, 10, 110, 170, 120, 200},
-		{"Large 5", 5000, 15, 1300, 1800, 5000, 6000},
+		{"Small", 50, 5},
+		{"Medium", 500, 10},
+		{"Large", 5000, 20},
+		{"Small single", 50, 1},
+		{"Medium single", 500, 1},
+		{"Large 5", 5000, 5},
 	}
 
 	log.SetOutput(io.Discard)
@@ -885,23 +918,8 @@ func BenchmarkUpdateAccountPeers(b *testing.B) {
 			}
 
 			duration := time.Since(start)
-			msPerOp := float64(duration.Nanoseconds()) / float64(b.N) / 1e6
-			b.ReportMetric(msPerOp, "ms/op")
-
-			minExpected := bc.minMsPerOpLocal
-			maxExpected := bc.maxMsPerOpLocal
-			if os.Getenv("CI") == "true" {
-				minExpected = bc.minMsPerOpCICD
-				maxExpected = bc.maxMsPerOpCICD
-			}
-
-			if msPerOp < minExpected {
-				b.Fatalf("Benchmark %s failed: too fast (%.2f ms/op, minimum %.2f ms/op)", bc.name, msPerOp, minExpected)
-			}
-
-			if msPerOp > maxExpected {
-				b.Fatalf("Benchmark %s failed: too slow (%.2f ms/op, maximum %.2f ms/op)", bc.name, msPerOp, maxExpected)
-			}
+			b.ReportMetric(float64(duration.Nanoseconds())/float64(b.N)/1e6, "ms/op")
+			b.ReportMetric(0, "ns/op")
 		})
 	}
 }
@@ -1465,7 +1483,8 @@ func TestPeerAccountPeersUpdate(t *testing.T) {
 	// 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{
-			Enabled: true,
+			AccountID: account.Id,
+			Enabled:   true,
 			Rules: []*PolicyRule{
 				{
 					Enabled:       true,

management/server/personal_access_token.go

@@ -41,6 +41,7 @@ type PersonalAccessToken struct {
 func (t *PersonalAccessToken) Copy() *PersonalAccessToken {
 	return &PersonalAccessToken{
 		ID:             t.ID,
+		UserID:         t.UserID,
 		Name:           t.Name,
 		HashedToken:    t.HashedToken,
 		ExpirationDate: t.ExpirationDate,
@@ -58,7 +59,7 @@ type PersonalAccessTokenGenerated struct {
 
 // CreateNewPAT will generate a new PersonalAccessToken that can be assigned to a User.
 // Additionally, it will return the token in plain text once, to give to the user and only save a hashed version
-func CreateNewPAT(name string, expirationInDays int, createdBy string) (*PersonalAccessTokenGenerated, error) {
+func CreateNewPAT(name string, expirationInDays int, targetID, createdBy string) (*PersonalAccessTokenGenerated, error) {
 	hashedToken, plainToken, err := generateNewToken()
 	if err != nil {
 		return nil, err
@@ -67,6 +68,7 @@ func CreateNewPAT(name string, expirationInDays int, createdBy string) (*Persona
 	return &PersonalAccessTokenGenerated{
 		PersonalAccessToken: PersonalAccessToken{
 			ID:             xid.New().String(),
+			UserID:         targetID,
 			Name:           name,
 			HashedToken:    hashedToken,
 			ExpirationDate: currentTime.AddDate(0, 0, expirationInDays),

management/server/policy.go

@@ -435,7 +435,7 @@ func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, po
 	var updateAccountPeers bool
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		policy, err = transaction.GetPolicyByID(ctx, LockingStrengthUpdate, accountID, policyID)
+		policy, err = transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policyID)
 		if err != nil {
 			return err
 		}
@@ -502,6 +502,8 @@ func arePolicyChangesAffectPeers(ctx context.Context, transaction Store, account
 		if hasPeers {
 			return true, nil
 		}
+
+		return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
 	}
 
 	return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
@@ -532,7 +534,7 @@ func validatePolicy(ctx context.Context, transaction Store, accountID string, po
 	for i, rule := range policy.Rules {
 		ruleCopy := rule.Copy()
 		if ruleCopy.ID == "" {
-			ruleCopy.ID = policy.ID // TODO: when policy can contain multiple rules, need refactor
+			ruleCopy.ID = xid.New().String()
 			ruleCopy.PolicyID = policy.ID
 		}
 

management/server/posture_checks.go

@@ -2,16 +2,15 @@ package server
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"slices"
 
-	"github.com/rs/xid"
-	"golang.org/x/exp/maps"
-
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/status"
+	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
 )
 
 func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error) {
@@ -33,9 +32,6 @@ func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID
 
 // SavePostureChecks saves a posture check.
 func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) (*posture.Checks, error) {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
@@ -89,9 +85,6 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
 
 // DeletePostureChecks deletes a posture check by ID.
 func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accountID, postureChecksID, userID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return err
@@ -151,21 +144,38 @@ func (am *DefaultAccountManager) ListPostureChecks(ctx context.Context, accountI
 }
 
 // getPeerPostureChecks returns the posture checks applied for a given peer.
-func (am *DefaultAccountManager) getPeerPostureChecks(account *Account, peerID string) ([]*posture.Checks, error) {
+func (am *DefaultAccountManager) getPeerPostureChecks(ctx context.Context, accountID string, peerID string) ([]*posture.Checks, error) {
 	peerPostureChecks := make(map[string]*posture.Checks)
 
-	if len(account.PostureChecks) == 0 {
-		return nil, nil
-	}
-
-	for _, policy := range account.Policies {
-		if !policy.Enabled || len(policy.SourcePostureChecks) == 0 {
-			continue
+	err := am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		postureChecks, err := transaction.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return err
 		}
 
-		if err := addPolicyPostureChecks(account, peerID, policy, peerPostureChecks); err != nil {
-			return nil, err
+		if len(postureChecks) == 0 {
+			return nil
 		}
+
+		policies, err := transaction.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return err
+		}
+
+		for _, policy := range policies {
+			if !policy.Enabled {
+				continue
+			}
+
+			if err = addPolicyPostureChecks(ctx, transaction, accountID, peerID, policy, peerPostureChecks); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	return maps.Values(peerPostureChecks), nil
@@ -226,8 +236,8 @@ func validatePostureChecks(ctx context.Context, transaction Store, accountID str
 }
 
 // addPolicyPostureChecks adds posture checks from a policy to the peer posture checks map if the peer is in the policy's source groups.
-func addPolicyPostureChecks(account *Account, peerID string, policy *Policy, peerPostureChecks map[string]*posture.Checks) error {
-	isInGroup, err := isPeerInPolicySourceGroups(account, peerID, policy)
+func addPolicyPostureChecks(ctx context.Context, transaction Store, accountID, peerID string, policy *Policy, peerPostureChecks map[string]*posture.Checks) error {
+	isInGroup, err := isPeerInPolicySourceGroups(ctx, transaction, accountID, peerID, policy)
 	if err != nil {
 		return err
 	}
@@ -237,9 +247,9 @@ func addPolicyPostureChecks(account *Account, peerID string, policy *Policy, pee
 	}
 
 	for _, sourcePostureCheckID := range policy.SourcePostureChecks {
-		postureCheck := account.getPostureChecks(sourcePostureCheckID)
-		if postureCheck == nil {
-			return errors.New("failed to add policy posture checks: posture checks not found")
+		postureCheck, err := transaction.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, sourcePostureCheckID)
+		if err != nil {
+			return err
 		}
 		peerPostureChecks[sourcePostureCheckID] = postureCheck
 	}
@@ -248,16 +258,17 @@ func addPolicyPostureChecks(account *Account, peerID string, policy *Policy, pee
 }
 
 // isPeerInPolicySourceGroups checks if a peer is present in any of the policy rule source groups.
-func isPeerInPolicySourceGroups(account *Account, peerID string, policy *Policy) (bool, error) {
+func isPeerInPolicySourceGroups(ctx context.Context, transaction Store, accountID, peerID string, policy *Policy) (bool, error) {
 	for _, rule := range policy.Rules {
 		if !rule.Enabled {
 			continue
 		}
 
 		for _, sourceGroup := range rule.Sources {
-			group := account.GetGroup(sourceGroup)
-			if group == nil {
-				return false, fmt.Errorf("failed to check peer in policy source group: group not found")
+			group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, accountID, sourceGroup)
+			if err != nil {
+				log.WithContext(ctx).Debugf("failed to check peer in policy source group: %v", err)
+				return false, fmt.Errorf("failed to check peer in policy source group: %w", err)
 			}
 
 			if slices.Contains(group.Peers, peerID) {

management/server/posture_checks_test.go

@@ -25,22 +25,22 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestPostureChecksAccount(am)
+	accountID, err := initTestPostureChecksAccount(am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
 	t.Run("Generic posture check flow", func(t *testing.T) {
 		// regular users can not create checks
-		_, err = am.SavePostureChecks(context.Background(), account.Id, regularUserID, &posture.Checks{})
+		_, err = am.SavePostureChecks(context.Background(), accountID, regularUserID, &posture.Checks{})
 		assert.Error(t, err)
 
 		// regular users cannot list check
-		_, err = am.ListPostureChecks(context.Background(), account.Id, regularUserID)
+		_, err = am.ListPostureChecks(context.Background(), accountID, regularUserID)
 		assert.Error(t, err)
 
 		// should be possible to create posture check with uniq name
-		postureCheck, err := am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+		postureCheck, err := am.SavePostureChecks(context.Background(), accountID, adminUserID, &posture.Checks{
 			Name: postureCheckName,
 			Checks: posture.ChecksDefinition{
 				NBVersionCheck: &posture.NBVersionCheck{
@@ -51,12 +51,12 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 		assert.NoError(t, err)
 
 		// admin users can list check
-		checks, err := am.ListPostureChecks(context.Background(), account.Id, adminUserID)
+		checks, err := am.ListPostureChecks(context.Background(), accountID, adminUserID)
 		assert.NoError(t, err)
 		assert.Len(t, checks, 1)
 
 		// should not be possible to create posture check with non uniq name
-		_, err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
+		_, err = am.SavePostureChecks(context.Background(), accountID, adminUserID, &posture.Checks{
 			Name: postureCheckName,
 			Checks: posture.ChecksDefinition{
 				GeoLocationCheck: &posture.GeoLocationCheck{
@@ -76,45 +76,48 @@ func TestDefaultAccountManager_PostureCheck(t *testing.T) {
 				MinVersion: "0.27.0",
 			},
 		}
-		_, err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheck)
+		_, err = am.SavePostureChecks(context.Background(), accountID, adminUserID, postureCheck)
 		assert.NoError(t, err)
 
 		// users should not be able to delete posture checks
-		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheck.ID, regularUserID)
+		err = am.DeletePostureChecks(context.Background(), accountID, postureCheck.ID, regularUserID)
 		assert.Error(t, err)
 
 		// admin should be able to delete posture checks
-		err = am.DeletePostureChecks(context.Background(), account.Id, postureCheck.ID, adminUserID)
+		err = am.DeletePostureChecks(context.Background(), accountID, postureCheck.ID, adminUserID)
 		assert.NoError(t, err)
-		checks, err = am.ListPostureChecks(context.Background(), account.Id, adminUserID)
+		checks, err = am.ListPostureChecks(context.Background(), accountID, adminUserID)
 		assert.NoError(t, err)
 		assert.Len(t, checks, 0)
 	})
 }
 
-func initTestPostureChecksAccount(am *DefaultAccountManager) (*Account, error) {
+func initTestPostureChecksAccount(am *DefaultAccountManager) (string, error) {
 	accountID := "testingAccount"
 	domain := "example.com"
 
-	admin := &User{
-		Id:   adminUserID,
-		Role: UserRoleAdmin,
-	}
-	user := &User{
-		Id:   regularUserID,
-		Role: UserRoleUser,
-	}
-
-	account := newAccountWithId(context.Background(), accountID, groupAdminUserID, domain)
-	account.Users[admin.Id] = admin
-	account.Users[user.Id] = user
-
-	err := am.Store.SaveAccount(context.Background(), account)
+	err := newAccountWithId(context.Background(), am.Store, accountID, groupAdminUserID, domain)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	return am.Store.GetAccount(context.Background(), account.Id)
+	err = am.Store.SaveUsers(context.Background(), LockingStrengthUpdate, []*User{
+		{
+			Id:        adminUserID,
+			AccountID: accountID,
+			Role:      UserRoleAdmin,
+		},
+		{
+			Id:        regularUserID,
+			AccountID: accountID,
+			Role:      UserRoleUser,
+		},
+	})
+	if err != nil {
+		return "", err
+	}
+
+	return accountID, nil
 }
 
 func TestPostureCheckAccountPeersUpdate(t *testing.T) {
@@ -440,18 +443,18 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 	manager, err := createManager(t)
 	require.NoError(t, err, "failed to create account manager")
 
-	account, err := initTestPostureChecksAccount(manager)
+	accountID, err := initTestPostureChecksAccount(manager)
 	require.NoError(t, err, "failed to init testing account")
 
 	groupA := &group.Group{
 		ID:        "groupA",
-		AccountID: account.Id,
+		AccountID: accountID,
 		Peers:     []string{"peer1"},
 	}
 
 	groupB := &group.Group{
 		ID:        "groupB",
-		AccountID: account.Id,
+		AccountID: accountID,
 		Peers:     []string{},
 	}
 	err = manager.Store.SaveGroups(context.Background(), LockingStrengthUpdate, []*group.Group{groupA, groupB})
@@ -459,26 +462,26 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 
 	postureCheckA := &posture.Checks{
 		Name:      "checkA",
-		AccountID: account.Id,
+		AccountID: accountID,
 		Checks: posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.33.1"},
 		},
 	}
-	postureCheckA, err = manager.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheckA)
+	postureCheckA, err = manager.SavePostureChecks(context.Background(), accountID, adminUserID, postureCheckA)
 	require.NoError(t, err, "failed to save postureCheckA")
 
 	postureCheckB := &posture.Checks{
 		Name:      "checkB",
-		AccountID: account.Id,
+		AccountID: accountID,
 		Checks: posture.ChecksDefinition{
 			NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.33.1"},
 		},
 	}
-	postureCheckB, err = manager.SavePostureChecks(context.Background(), account.Id, adminUserID, postureCheckB)
+	postureCheckB, err = manager.SavePostureChecks(context.Background(), accountID, adminUserID, postureCheckB)
 	require.NoError(t, err, "failed to save postureCheckB")
 
 	policy := &Policy{
-		AccountID: account.Id,
+		AccountID: accountID,
 		Rules: []*PolicyRule{
 			{
 				Enabled:      true,
@@ -489,23 +492,23 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 		SourcePostureChecks: []string{postureCheckA.ID},
 	}
 
-	policy, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
+	policy, err = manager.SavePolicy(context.Background(), accountID, adminUserID, policy)
 	require.NoError(t, err, "failed to save policy")
 
 	t.Run("posture check exists and is linked to policy with peers", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckA.ID)
 		require.NoError(t, err)
 		assert.True(t, result)
 	})
 
 	t.Run("posture check exists but is not linked to any policy", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckB.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckB.ID)
 		require.NoError(t, err)
 		assert.False(t, result)
 	})
 
 	t.Run("posture check does not exist", func(t *testing.T) {
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, "unknown")
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, "unknown")
 		require.NoError(t, err)
 		assert.False(t, result)
 	})
@@ -513,10 +516,10 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 	t.Run("posture check is linked to policy with no peers in source groups", func(t *testing.T) {
 		policy.Rules[0].Sources = []string{"groupB"}
 		policy.Rules[0].Destinations = []string{"groupA"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
+		_, err = manager.SavePolicy(context.Background(), accountID, adminUserID, policy)
 		require.NoError(t, err, "failed to update policy")
 
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckA.ID)
 		require.NoError(t, err)
 		assert.True(t, result)
 	})
@@ -524,10 +527,10 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 	t.Run("posture check is linked to policy with no peers in destination groups", func(t *testing.T) {
 		policy.Rules[0].Sources = []string{"groupA"}
 		policy.Rules[0].Destinations = []string{"groupB"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
+		_, err = manager.SavePolicy(context.Background(), accountID, adminUserID, policy)
 		require.NoError(t, err, "failed to update policy")
 
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckA.ID)
 		require.NoError(t, err)
 		assert.True(t, result)
 	})
@@ -537,7 +540,7 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 		err = manager.Store.SaveGroup(context.Background(), LockingStrengthUpdate, groupA)
 		require.NoError(t, err, "failed to save groups")
 
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckA.ID)
 		require.NoError(t, err)
 		assert.False(t, result)
 	})
@@ -545,10 +548,10 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
 	t.Run("posture check is linked to policy with non-existent group", func(t *testing.T) {
 		policy.Rules[0].Sources = []string{"nonExistentGroup"}
 		policy.Rules[0].Destinations = []string{"nonExistentGroup"}
-		_, err = manager.SavePolicy(context.Background(), account.Id, adminUserID, policy)
+		_, err = manager.SavePolicy(context.Background(), accountID, adminUserID, policy)
 		require.NoError(t, err, "failed to update policy")
 
-		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, account.Id, postureCheckA.ID)
+		result, err := arePostureCheckChangesAffectPeers(context.Background(), manager.Store, accountID, postureCheckA.ID)
 		require.NoError(t, err)
 		assert.False(t, result)
 	})

management/server/route.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"unicode/utf8"
 
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 
@@ -52,17 +53,46 @@ func (am *DefaultAccountManager) GetRoute(ctx context.Context, accountID string,
 		return nil, err
 	}
 
-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
-		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can view Network Routes")
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	return am.Store.GetRouteByID(ctx, LockingStrengthShare, string(routeID), accountID)
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
+	return am.Store.GetRouteByID(ctx, LockingStrengthShare, accountID, string(routeID))
+}
+
+// GetRoutesByPrefixOrDomains return list of routes by account and route prefix
+func getRoutesByPrefixOrDomains(ctx context.Context, transaction Store, accountID string, prefix netip.Prefix, domains domain.List) ([]*route.Route, error) {
+	accountRoutes, err := transaction.GetAccountRoutes(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	routes := make([]*route.Route, 0)
+	for _, r := range accountRoutes {
+		dynamic := r.IsDynamic()
+		if dynamic && r.Domains.PunycodeString() == domains.PunycodeString() ||
+			!dynamic && r.Network.String() == prefix.String() {
+			routes = append(routes, r)
+		}
+	}
+
+	return routes, nil
 }
 
 // checkRoutePrefixOrDomainsExistForPeers checks if a route with a given prefix exists for a single peer or multiple peer groups.
-func (am *DefaultAccountManager) checkRoutePrefixOrDomainsExistForPeers(account *Account, peerID string, routeID route.ID, peerGroupIDs []string, prefix netip.Prefix, domains domain.List) error {
+func checkRoutePrefixOrDomainsExistForPeers(ctx context.Context, transaction Store, accountID string, checkRoute *route.Route, groupsMap map[string]*nbgroup.Group) error {
 	// routes can have both peer and peer_groups
-	routesWithPrefix := account.GetRoutesByPrefixOrDomains(prefix, domains)
+	prefix := checkRoute.Network
+	domains := checkRoute.Domains
+
+	routesWithPrefix, err := getRoutesByPrefixOrDomains(ctx, transaction, accountID, prefix, domains)
+	if err != nil {
+		return err
+	}
 
 	// lets remember all the peers and the peer groups from routesWithPrefix
 	seenPeers := make(map[string]bool)
@@ -71,18 +101,24 @@ func (am *DefaultAccountManager) checkRoutePrefixOrDomainsExistForPeers(account
 	for _, prefixRoute := range routesWithPrefix {
 		// we skip route(s) with the same network ID as we want to allow updating of the existing route
 		// when creating a new route routeID is newly generated so nothing will be skipped
-		if routeID == prefixRoute.ID {
+		if checkRoute.ID == prefixRoute.ID {
 			continue
 		}
 
 		if prefixRoute.Peer != "" {
 			seenPeers[string(prefixRoute.ID)] = true
 		}
+
+		peerGroupsMap, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, prefixRoute.PeerGroups)
+		if err != nil {
+			return err
+		}
+
 		for _, groupID := range prefixRoute.PeerGroups {
 			seenPeerGroups[groupID] = true
 
-			group := account.GetGroup(groupID)
-			if group == nil {
+			group, ok := peerGroupsMap[groupID]
+			if !ok || group == nil {
 				return status.Errorf(
 					status.InvalidArgument, "failed to add route with %s - peer group %s doesn't exist",
 					getRouteDescriptor(prefix, domains), groupID,
@@ -95,12 +131,13 @@ func (am *DefaultAccountManager) checkRoutePrefixOrDomainsExistForPeers(account
 		}
 	}
 
-	if peerID != "" {
+	if peerID := checkRoute.Peer; peerID != "" {
 		// check that peerID exists and is not in any route as single peer or part of the group
-		peer := account.GetPeer(peerID)
-		if peer == nil {
+		_, err = transaction.GetPeerByID(context.Background(), LockingStrengthShare, accountID, peerID)
+		if err != nil {
 			return status.Errorf(status.InvalidArgument, "peer with ID %s not found", peerID)
 		}
+
 		if _, ok := seenPeers[peerID]; ok {
 			return status.Errorf(status.AlreadyExists,
 				"failed to add route with %s - peer %s already has this route", getRouteDescriptor(prefix, domains), peerID)
@@ -108,9 +145,8 @@ func (am *DefaultAccountManager) checkRoutePrefixOrDomainsExistForPeers(account
 	}
 
 	// check that peerGroupIDs are not in any route peerGroups list
-	for _, groupID := range peerGroupIDs {
-		group := account.GetGroup(groupID) // we validated the group existence before entering this function, no need to check again.
-
+	for _, groupID := range checkRoute.PeerGroups {
+		group := groupsMap[groupID] // we validated the group existence before entering this function, no need to check again.
 		if _, ok := seenPeerGroups[groupID]; ok {
 			return status.Errorf(
 				status.AlreadyExists, "failed to add route with %s - peer group %s already has this route",
@@ -118,12 +154,18 @@ func (am *DefaultAccountManager) checkRoutePrefixOrDomainsExistForPeers(account
 		}
 
 		// check that the peers from peerGroupIDs groups are not the same peers we saw in routesWithPrefix
+		peersMap, err := transaction.GetPeersByIDs(ctx, LockingStrengthShare, accountID, group.Peers)
+		if err != nil {
+			return err
+		}
+
 		for _, id := range group.Peers {
 			if _, ok := seenPeers[id]; ok {
-				peer := account.GetPeer(id)
-				if peer == nil {
-					return status.Errorf(status.InvalidArgument, "peer with ID %s not found", peerID)
+				peer, ok := peersMap[id]
+				if !ok || peer == nil {
+					return status.Errorf(status.InvalidArgument, "peer with ID %s not found", id)
 				}
+
 				return status.Errorf(status.AlreadyExists,
 					"failed to add route with %s - peer %s from the group %s already has this route",
 					getRouteDescriptor(prefix, domains), peer.Name, group.Name)
@@ -146,104 +188,63 @@ func (am *DefaultAccountManager) CreateRoute(ctx context.Context, accountID stri
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}
 
-	// Do not allow non-Linux peers
-	if peer := account.GetPeer(peerID); peer != nil {
-		if peer.Meta.GoOS != "linux" {
-			return nil, status.Errorf(status.InvalidArgument, "non-linux peers are not supported as network routes")
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
+	}
+
+	var newRoute *route.Route
+	var updateAccountPeers bool
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		newRoute = &route.Route{
+			ID:                  route.ID(xid.New().String()),
+			AccountID:           accountID,
+			Network:             prefix,
+			Domains:             domains,
+			KeepRoute:           keepRoute,
+			NetID:               netID,
+			Description:         description,
+			Peer:                peerID,
+			PeerGroups:          peerGroupIDs,
+			NetworkType:         networkType,
+			Masquerade:          masquerade,
+			Metric:              metric,
+			Enabled:             enabled,
+			Groups:              groups,
+			AccessControlGroups: accessControlGroupIDs,
 		}
-	}
 
-	if len(domains) > 0 && prefix.IsValid() {
-		return nil, status.Errorf(status.InvalidArgument, "domains and network should not be provided at the same time")
-	}
+		if err = validateRoute(ctx, transaction, accountID, newRoute); err != nil {
+			return err
+		}
 
-	if len(domains) == 0 && !prefix.IsValid() {
-		return nil, status.Errorf(status.InvalidArgument, "invalid Prefix")
-	}
-
-	if len(domains) > 0 {
-		prefix = getPlaceholderIP()
-	}
-
-	if peerID != "" && len(peerGroupIDs) != 0 {
-		return nil, status.Errorf(
-			status.InvalidArgument,
-			"peer with ID %s and peers group %s should not be provided at the same time",
-			peerID, peerGroupIDs)
-	}
-
-	var newRoute route.Route
-	newRoute.ID = route.ID(xid.New().String())
-
-	if len(peerGroupIDs) > 0 {
-		err = validateGroups(peerGroupIDs, account.Groups)
+		updateAccountPeers, err = areRouteChangesAffectPeers(ctx, transaction, newRoute)
 		if err != nil {
-			return nil, err
+			return err
 		}
-	}
 
-	if len(accessControlGroupIDs) > 0 {
-		err = validateGroups(accessControlGroupIDs, account.Groups)
-		if err != nil {
-			return nil, err
+		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
+			return err
 		}
-	}
 
-	err = am.checkRoutePrefixOrDomainsExistForPeers(account, peerID, newRoute.ID, peerGroupIDs, prefix, domains)
+		return transaction.SaveRoute(ctx, LockingStrengthUpdate, newRoute)
+	})
 	if err != nil {
 		return nil, err
 	}
 
-	if metric < route.MinMetric || metric > route.MaxMetric {
-		return nil, status.Errorf(status.InvalidArgument, "metric should be between %d and %d", route.MinMetric, route.MaxMetric)
-	}
-
-	if utf8.RuneCountInString(string(netID)) > route.MaxNetIDChar || netID == "" {
-		return nil, status.Errorf(status.InvalidArgument, "identifier should be between 1 and %d", route.MaxNetIDChar)
-	}
-
-	err = validateGroups(groups, account.Groups)
-	if err != nil {
-		return nil, err
-	}
-
-	newRoute.Peer = peerID
-	newRoute.PeerGroups = peerGroupIDs
-	newRoute.Network = prefix
-	newRoute.Domains = domains
-	newRoute.NetworkType = networkType
-	newRoute.Description = description
-	newRoute.NetID = netID
-	newRoute.Masquerade = masquerade
-	newRoute.Metric = metric
-	newRoute.Enabled = enabled
-	newRoute.Groups = groups
-	newRoute.KeepRoute = keepRoute
-	newRoute.AccessControlGroups = accessControlGroupIDs
-
-	if account.Routes == nil {
-		account.Routes = make(map[route.ID]*route.Route)
-	}
-
-	account.Routes[newRoute.ID] = &newRoute
-
-	account.Network.IncSerial()
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
-		return nil, err
-	}
-
-	if am.isRouteChangeAffectPeers(account, &newRoute) {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
 	am.StoreEvent(ctx, userID, string(newRoute.ID), accountID, activity.RouteCreated, newRoute.EventMeta())
 
-	return &newRoute, nil
+	if updateAccountPeers {
+		am.updateAccountPeers(ctx, accountID)
+	}
+
+	return newRoute, nil
 }
 
 // SaveRoute saves route
@@ -251,10 +252,151 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	if err != nil {
+		return err
+	}
+
+	if user.AccountID != accountID {
+		return status.NewUserNotPartOfAccountError()
+	}
+
+	var oldRoute *route.Route
+	var oldRouteAffectsPeers bool
+	var newRouteAffectsPeers bool
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		if err = validateRoute(ctx, transaction, accountID, routeToSave); err != nil {
+			return err
+		}
+
+		oldRoute, err = transaction.GetRouteByID(ctx, LockingStrengthUpdate, accountID, string(routeToSave.ID))
+		if err != nil {
+			return err
+		}
+
+		oldRouteAffectsPeers, err = areRouteChangesAffectPeers(ctx, transaction, oldRoute)
+		if err != nil {
+			return err
+		}
+
+		newRouteAffectsPeers, err = areRouteChangesAffectPeers(ctx, transaction, routeToSave)
+		if err != nil {
+			return err
+		}
+		routeToSave.AccountID = accountID
+
+		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
+			return err
+		}
+
+		return transaction.SaveRoute(ctx, LockingStrengthUpdate, routeToSave)
+	})
+	if err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, userID, string(routeToSave.ID), accountID, activity.RouteUpdated, routeToSave.EventMeta())
+
+	if oldRouteAffectsPeers || newRouteAffectsPeers {
+		am.updateAccountPeers(ctx, accountID)
+	}
+
+	return nil
+}
+
+// DeleteRoute deletes route with routeID
+func (am *DefaultAccountManager) DeleteRoute(ctx context.Context, accountID string, routeID route.ID, userID string) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	if err != nil {
+		return err
+	}
+
+	if user.AccountID != accountID {
+		return status.NewUserNotPartOfAccountError()
+	}
+
+	var route *route.Route
+	var updateAccountPeers bool
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		route, err = transaction.GetRouteByID(ctx, LockingStrengthUpdate, accountID, string(routeID))
+		if err != nil {
+			return err
+		}
+
+		updateAccountPeers, err = areRouteChangesAffectPeers(ctx, transaction, route)
+		if err != nil {
+			return err
+		}
+
+		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
+			return err
+		}
+
+		return transaction.DeleteRoute(ctx, LockingStrengthUpdate, accountID, string(routeID))
+	})
+
+	am.StoreEvent(ctx, userID, string(route.ID), accountID, activity.RouteRemoved, route.EventMeta())
+
+	if updateAccountPeers {
+		am.updateAccountPeers(ctx, accountID)
+	}
+
+	return nil
+}
+
+// ListRoutes returns a list of routes from account
+func (am *DefaultAccountManager) ListRoutes(ctx context.Context, accountID, userID string) ([]*route.Route, error) {
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	if err != nil {
+		return nil, err
+	}
+
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
+	return am.Store.GetAccountRoutes(ctx, LockingStrengthShare, accountID)
+}
+
+func validateRoute(ctx context.Context, transaction Store, accountID string, routeToSave *route.Route) error {
 	if routeToSave == nil {
 		return status.Errorf(status.InvalidArgument, "route provided is nil")
 	}
 
+	if err := validateRouteProperties(routeToSave); err != nil {
+		return err
+	}
+
+	if routeToSave.Peer != "" {
+		peer, err := transaction.GetPeerByID(ctx, LockingStrengthShare, accountID, routeToSave.Peer)
+		if err != nil {
+			return err
+		}
+
+		if peer.Meta.GoOS != "linux" {
+			return status.Errorf(status.InvalidArgument, "non-linux peers are not supported as network routes")
+		}
+	}
+
+	groupsMap, err := validateRouteGroups(ctx, transaction, accountID, routeToSave)
+	if err != nil {
+		return err
+	}
+
+	return checkRoutePrefixOrDomainsExistForPeers(ctx, transaction, accountID, routeToSave, groupsMap)
+}
+
+// Helper to validate route properties.
+func validateRouteProperties(routeToSave *route.Route) error {
 	if routeToSave.Metric < route.MinMetric || routeToSave.Metric > route.MaxMetric {
 		return status.Errorf(status.InvalidArgument, "metric should be between %d and %d", route.MinMetric, route.MaxMetric)
 	}
@@ -263,18 +405,6 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
 		return status.Errorf(status.InvalidArgument, "identifier should be between 1 and %d", route.MaxNetIDChar)
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
-	if err != nil {
-		return err
-	}
-
-	// Do not allow non-Linux peers
-	if peer := account.GetPeer(routeToSave.Peer); peer != nil {
-		if peer.Meta.GoOS != "linux" {
-			return status.Errorf(status.InvalidArgument, "non-linux peers are not supported as network routes")
-		}
-	}
-
 	if len(routeToSave.Domains) > 0 && routeToSave.Network.IsValid() {
 		return status.Errorf(status.InvalidArgument, "domains and network should not be provided at the same time")
 	}
@@ -291,89 +421,34 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
 		return status.Errorf(status.InvalidArgument, "peer with ID and peer groups should not be provided at the same time")
 	}
 
-	if len(routeToSave.PeerGroups) > 0 {
-		err = validateGroups(routeToSave.PeerGroups, account.Groups)
-		if err != nil {
-			return err
-		}
-	}
-
-	if len(routeToSave.AccessControlGroups) > 0 {
-		err = validateGroups(routeToSave.AccessControlGroups, account.Groups)
-		if err != nil {
-			return err
-		}
-	}
-
-	err = am.checkRoutePrefixOrDomainsExistForPeers(account, routeToSave.Peer, routeToSave.ID, routeToSave.Copy().PeerGroups, routeToSave.Network, routeToSave.Domains)
-	if err != nil {
-		return err
-	}
-
-	err = validateGroups(routeToSave.Groups, account.Groups)
-	if err != nil {
-		return err
-	}
-
-	oldRoute := account.Routes[routeToSave.ID]
-	account.Routes[routeToSave.ID] = routeToSave
-
-	account.Network.IncSerial()
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
-		return err
-	}
-
-	if am.isRouteChangeAffectPeers(account, oldRoute) || am.isRouteChangeAffectPeers(account, routeToSave) {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
-	am.StoreEvent(ctx, userID, string(routeToSave.ID), accountID, activity.RouteUpdated, routeToSave.EventMeta())
-
 	return nil
 }
 
-// DeleteRoute deletes route with routeID
-func (am *DefaultAccountManager) DeleteRoute(ctx context.Context, accountID string, routeID route.ID, userID string) error {
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
-	account, err := am.Store.GetAccount(ctx, accountID)
-	if err != nil {
-		return err
-	}
-
-	routy := account.Routes[routeID]
-	if routy == nil {
-		return status.Errorf(status.NotFound, "route with ID %s doesn't exist", routeID)
-	}
-	delete(account.Routes, routeID)
-
-	account.Network.IncSerial()
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
-		return err
-	}
-
-	am.StoreEvent(ctx, userID, string(routy.ID), accountID, activity.RouteRemoved, routy.EventMeta())
-
-	if am.isRouteChangeAffectPeers(account, routy) {
-		am.updateAccountPeers(ctx, accountID)
-	}
-
-	return nil
-}
-
-// ListRoutes returns a list of routes from account
-func (am *DefaultAccountManager) ListRoutes(ctx context.Context, accountID, userID string) ([]*route.Route, error) {
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+// validateRouteGroups validates the route groups and returns the validated groups map.
+func validateRouteGroups(ctx context.Context, transaction Store, accountID string, routeToSave *route.Route) (map[string]*nbgroup.Group, error) {
+	groupsToValidate := slices.Concat(routeToSave.Groups, routeToSave.PeerGroups, routeToSave.AccessControlGroups)
+	groupsMap, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, groupsToValidate)
 	if err != nil {
 		return nil, err
 	}
 
-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
-		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can view Network Routes")
+	if len(routeToSave.PeerGroups) > 0 {
+		if err = validateGroups(routeToSave.PeerGroups, groupsMap); err != nil {
+			return nil, err
+		}
 	}
 
-	return am.Store.GetAccountRoutes(ctx, LockingStrengthShare, accountID)
+	if len(routeToSave.AccessControlGroups) > 0 {
+		if err = validateGroups(routeToSave.AccessControlGroups, groupsMap); err != nil {
+			return nil, err
+		}
+	}
+
+	if err = validateGroups(routeToSave.Groups, groupsMap); err != nil {
+		return nil, err
+	}
+
+	return groupsMap, nil
 }
 
 func toProtocolRoute(route *route.Route) *proto.Route {
@@ -417,84 +492,27 @@ func (a *Account) getPeerRoutesFirewallRules(ctx context.Context, peerID string,
 			continue
 		}
 
-		distributionPeers := a.getDistributionGroupsPeers(route)
+		policies := getAllRoutePoliciesFromGroups(a, route.AccessControlGroups)
+		for _, policy := range policies {
+			if !policy.Enabled {
+				continue
+			}
 
-		for _, accessGroup := range route.AccessControlGroups {
-			policies := getAllRoutePoliciesFromGroups(a, []string{accessGroup})
-			rules := a.getRouteFirewallRules(ctx, peerID, policies, route, validatedPeersMap, distributionPeers)
-			routesFirewallRules = append(routesFirewallRules, rules...)
+			for _, rule := range policy.Rules {
+				if !rule.Enabled {
+					continue
+				}
+
+				distributionGroupPeers, _ := a.getAllPeersFromGroups(ctx, route.Groups, peerID, nil, validatedPeersMap)
+				rules := generateRouteFirewallRules(ctx, route, rule, distributionGroupPeers, firewallRuleDirectionIN)
+				routesFirewallRules = append(routesFirewallRules, rules...)
+			}
 		}
 	}
 
 	return routesFirewallRules
 }
 
-func (a *Account) getRouteFirewallRules(ctx context.Context, peerID string, policies []*Policy, route *route.Route, validatedPeersMap map[string]struct{}, distributionPeers map[string]struct{}) []*RouteFirewallRule {
-	var fwRules []*RouteFirewallRule
-	for _, policy := range policies {
-		if !policy.Enabled {
-			continue
-		}
-
-		for _, rule := range policy.Rules {
-			if !rule.Enabled {
-				continue
-			}
-
-			rulePeers := a.getRulePeers(rule, peerID, distributionPeers, validatedPeersMap)
-			rules := generateRouteFirewallRules(ctx, route, rule, rulePeers, firewallRuleDirectionIN)
-			fwRules = append(fwRules, rules...)
-		}
-	}
-	return fwRules
-}
-
-func (a *Account) getRulePeers(rule *PolicyRule, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*nbpeer.Peer {
-	distPeersWithPolicy := make(map[string]struct{})
-	for _, id := range rule.Sources {
-		group := a.Groups[id]
-		if group == nil {
-			continue
-		}
-
-		for _, pID := range group.Peers {
-			if pID == peerID {
-				continue
-			}
-			_, distPeer := distributionPeers[pID]
-			_, valid := validatedPeersMap[pID]
-			if distPeer && valid {
-				distPeersWithPolicy[pID] = struct{}{}
-			}
-		}
-	}
-
-	distributionGroupPeers := make([]*nbpeer.Peer, 0, len(distPeersWithPolicy))
-	for pID := range distPeersWithPolicy {
-		peer := a.Peers[pID]
-		if peer == nil {
-			continue
-		}
-		distributionGroupPeers = append(distributionGroupPeers, peer)
-	}
-	return distributionGroupPeers
-}
-
-func (a *Account) getDistributionGroupsPeers(route *route.Route) map[string]struct{} {
-	distPeers := make(map[string]struct{})
-	for _, id := range route.Groups {
-		group := a.Groups[id]
-		if group == nil {
-			continue
-		}
-
-		for _, pID := range group.Peers {
-			distPeers[pID] = struct{}{}
-		}
-	}
-	return distPeers
-}
-
 func getDefaultPermit(route *route.Route) []*RouteFirewallRule {
 	var rules []*RouteFirewallRule
 
@@ -706,8 +724,21 @@ func getProtoPortInfo(rule *RouteFirewallRule) *proto.PortInfo {
 	return &portInfo
 }
 
-// isRouteChangeAffectPeers checks if a given route affects peers by determining
-// if it has a routing peer, distribution, or peer groups that include peers
-func (am *DefaultAccountManager) isRouteChangeAffectPeers(account *Account, route *route.Route) bool {
-	return am.anyGroupHasPeers(account, route.Groups) || am.anyGroupHasPeers(account, route.PeerGroups) || route.Peer != ""
+// areRouteChangesAffectPeers checks if a given route affects peers by determining
+// if it has a routing peer, distribution, or peer groups that include peers.
+func areRouteChangesAffectPeers(ctx context.Context, transaction Store, route *route.Route) (bool, error) {
+	if route.Peer != "" {
+		return true, nil
+	}
+
+	hasPeers, err := anyGroupHasPeers(ctx, transaction, route.AccountID, route.Groups)
+	if err != nil {
+		return false, err
+	}
+
+	if hasPeers {
+		return true, nil
+	}
+
+	return anyGroupHasPeers(ctx, transaction, route.AccountID, route.PeerGroups)
 }

management/server/route_test.go

@@ -5,10 +5,11 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
-	"sort"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/rs/xid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -428,22 +429,22 @@ func TestCreateRoute(t *testing.T) {
 				t.Error("failed to create account manager")
 			}
 
-			account, err := initTestRouteAccount(t, am)
+			accountID, err := initTestRouteAccount(t, am)
 			if err != nil {
 				t.Errorf("failed to init testing account: %s", err)
 			}
 
 			if testCase.createInitRoute {
-				groupAll, errInit := account.GetGroupAll()
+				groupAll, errInit := am.Store.GetGroupByName(context.Background(), LockingStrengthShare, accountID, "All")
 				require.NoError(t, errInit)
-				_, errInit = am.CreateRoute(context.Background(), account.Id, existingNetwork, 1, nil, "", []string{routeGroup3, routeGroup4}, "", existingRouteID, false, 1000, []string{groupAll.ID}, []string{}, true, userID, false)
+
+				_, errInit = am.CreateRoute(context.Background(), accountID, existingNetwork, 1, nil, "", []string{routeGroup3, routeGroup4}, "", existingRouteID, false, 1000, []string{groupAll.ID}, []string{}, true, userID, false)
 				require.NoError(t, errInit)
-				_, errInit = am.CreateRoute(context.Background(), account.Id, netip.Prefix{}, 3, existingDomains, "", []string{routeGroup3, routeGroup4}, "", existingRouteID, false, 1000, []string{groupAll.ID}, []string{groupAll.ID}, true, userID, false)
+				_, errInit = am.CreateRoute(context.Background(), accountID, netip.Prefix{}, 3, existingDomains, "", []string{routeGroup3, routeGroup4}, "", existingRouteID, false, 1000, []string{groupAll.ID}, []string{groupAll.ID}, true, userID, false)
 				require.NoError(t, errInit)
 			}
 
-			outRoute, err := am.CreateRoute(context.Background(), account.Id, testCase.inputArgs.network, testCase.inputArgs.networkType, testCase.inputArgs.domains, testCase.inputArgs.peerKey, testCase.inputArgs.peerGroupIDs, testCase.inputArgs.description, testCase.inputArgs.netID, testCase.inputArgs.masquerade, testCase.inputArgs.metric, testCase.inputArgs.groups, testCase.inputArgs.accessControlGroups, testCase.inputArgs.enabled, userID, testCase.inputArgs.keepRoute)
-
+			outRoute, err := am.CreateRoute(context.Background(), accountID, testCase.inputArgs.network, testCase.inputArgs.networkType, testCase.inputArgs.domains, testCase.inputArgs.peerKey, testCase.inputArgs.peerGroupIDs, testCase.inputArgs.description, testCase.inputArgs.netID, testCase.inputArgs.masquerade, testCase.inputArgs.metric, testCase.inputArgs.groups, testCase.inputArgs.accessControlGroups, testCase.inputArgs.enabled, userID, testCase.inputArgs.keepRoute)
 			testCase.errFunc(t, err)
 
 			if !testCase.shouldCreate {
@@ -452,6 +453,7 @@ func TestCreateRoute(t *testing.T) {
 
 			// assign generated ID
 			testCase.expectedRoute.ID = outRoute.ID
+			testCase.expectedRoute.AccountID = accountID
 
 			if !testCase.expectedRoute.IsEqual(outRoute) {
 				t.Errorf("new route didn't match expected route:\nGot %#v\nExpected:%#v\n", outRoute, testCase.expectedRoute)
@@ -918,14 +920,15 @@ func TestSaveRoute(t *testing.T) {
 				t.Error("failed to create account manager")
 			}
 
-			account, err := initTestRouteAccount(t, am)
+			accountID, err := initTestRouteAccount(t, am)
 			if err != nil {
 				t.Error("failed to init testing account")
 			}
 
 			if testCase.createInitRoute {
-				account.Routes["initRoute"] = &route.Route{
+				initRoute := &route.Route{
 					ID:          "initRoute",
+					AccountID:   accountID,
 					Network:     existingNetwork,
 					NetID:       existingRouteID,
 					NetworkType: route.IPv4Network,
@@ -936,14 +939,13 @@ func TestSaveRoute(t *testing.T) {
 					Enabled:     true,
 					Groups:      []string{routeGroup1},
 				}
+				err = am.Store.SaveRoute(context.Background(), LockingStrengthUpdate, initRoute)
+				require.NoError(t, err, "failed to save init route")
 			}
 
-			account.Routes[testCase.existingRoute.ID] = testCase.existingRoute
-
-			err = am.Store.SaveAccount(context.Background(), account)
-			if err != nil {
-				t.Error("account should be saved")
-			}
+			testCase.existingRoute.AccountID = accountID
+			err = am.Store.SaveRoute(context.Background(), LockingStrengthUpdate, testCase.existingRoute)
+			require.NoError(t, err, "failed to save existing route")
 
 			var routeToSave *route.Route
 
@@ -978,7 +980,7 @@ func TestSaveRoute(t *testing.T) {
 				}
 			}
 
-			err = am.SaveRoute(context.Background(), account.Id, userID, routeToSave)
+			err = am.SaveRoute(context.Background(), accountID, userID, routeToSave)
 
 			testCase.errFunc(t, err)
 
@@ -986,14 +988,10 @@ func TestSaveRoute(t *testing.T) {
 				return
 			}
 
-			account, err = am.Store.GetAccount(context.Background(), account.Id)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			savedRoute, saved := account.Routes[testCase.expectedRoute.ID]
-			require.True(t, saved)
+			savedRoute, err := am.GetRoute(context.Background(), accountID, testCase.existingRoute.ID, userID)
+			require.NoError(t, err, "failed to get saved route")
 
+			testCase.expectedRoute.AccountID = accountID
 			if !testCase.expectedRoute.IsEqual(savedRoute) {
 				t.Errorf("new route didn't match expected route:\nGot %#v\nExpected:%#v\n", savedRoute, testCase.expectedRoute)
 			}
@@ -1002,50 +1000,48 @@ func TestSaveRoute(t *testing.T) {
 }
 
 func TestDeleteRoute(t *testing.T) {
-	testingRoute := &route.Route{
-		ID:          "testingRoute",
-		Network:     netip.MustParsePrefix("192.168.0.0/16"),
-		Domains:     domain.List{"domain1", "domain2"},
-		KeepRoute:   true,
-		NetworkType: route.IPv4Network,
-		Peer:        peer1Key,
-		Description: "super",
-		Masquerade:  false,
-		Metric:      9999,
-		Enabled:     true,
-	}
-
 	am, err := createRouterManager(t)
 	if err != nil {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestRouteAccount(t, am)
+	accountID, err := initTestRouteAccount(t, am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
 
-	account.Routes[testingRoute.ID] = testingRoute
+	err = am.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
+		ID:        "GroupA",
+		AccountID: accountID,
+		Name:      "GroupA",
+	})
+	require.NoError(t, err, "failed to save group")
 
-	err = am.Store.SaveAccount(context.Background(), account)
-	if err != nil {
-		t.Error("failed to save account")
+	testingRoute := &route.Route{
+		Network:     netip.MustParsePrefix("192.168.0.0/16"),
+		NetID:       route.NetID("12345678901234567890qw"),
+		Groups:      []string{"GroupA"},
+		KeepRoute:   true,
+		NetworkType: route.IPv4Network,
+		Peer:        peer1ID,
+		Description: "super",
+		Masquerade:  false,
+		Metric:      9999,
+		Enabled:     true,
 	}
+	createdRoute, err := am.CreateRoute(context.Background(), accountID, testingRoute.Network, testingRoute.NetworkType, testingRoute.Domains, peer1ID, []string{}, testingRoute.Description, testingRoute.NetID, testingRoute.Masquerade, testingRoute.Metric, testingRoute.Groups, testingRoute.AccessControlGroups, true, userID, testingRoute.KeepRoute)
+	require.NoError(t, err, "failed to create route")
 
-	err = am.DeleteRoute(context.Background(), account.Id, testingRoute.ID, userID)
+	err = am.DeleteRoute(context.Background(), accountID, createdRoute.ID, userID)
 	if err != nil {
 		t.Error("deleting route failed with error: ", err)
 	}
 
-	savedAccount, err := am.Store.GetAccount(context.Background(), account.Id)
-	if err != nil {
-		t.Error("failed to retrieve saved account with error: ", err)
-	}
-
-	_, found := savedAccount.Routes[testingRoute.ID]
-	if found {
-		t.Error("route shouldn't be found after delete")
-	}
+	_, err = am.GetRoute(context.Background(), accountID, testingRoute.ID, userID)
+	require.NotNil(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, sErr.Type())
 }
 
 func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
@@ -1067,16 +1063,14 @@ func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestRouteAccount(t, am)
-	if err != nil {
-		t.Error("failed to init testing account")
-	}
+	accountID, err := initTestRouteAccount(t, am)
+	require.NoError(t, err, "failed to init testing account")
 
 	newAccountRoutes, err := am.GetNetworkMap(context.Background(), peer1ID)
 	require.NoError(t, err)
 	require.Len(t, newAccountRoutes.Routes, 0, "new accounts should have no routes")
 
-	newRoute, err := am.CreateRoute(context.Background(), account.Id, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, baseRoute.Peer, baseRoute.PeerGroups, baseRoute.Description, baseRoute.NetID, baseRoute.Masquerade, baseRoute.Metric, baseRoute.Groups, baseRoute.AccessControlGroups, baseRoute.Enabled, userID, baseRoute.KeepRoute)
+	newRoute, err := am.CreateRoute(context.Background(), accountID, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, baseRoute.Peer, baseRoute.PeerGroups, baseRoute.Description, baseRoute.NetID, baseRoute.Masquerade, baseRoute.Metric, baseRoute.Groups, baseRoute.AccessControlGroups, baseRoute.Enabled, userID, baseRoute.KeepRoute)
 	require.NoError(t, err)
 	require.Equal(t, newRoute.Enabled, true)
 
@@ -1092,7 +1086,7 @@ func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
 	require.NoError(t, err)
 	assert.Len(t, peer4Routes.Routes, 1, "HA route should have 1 server route")
 
-	groups, err := am.Store.GetAccountGroups(context.Background(), LockingStrengthShare, account.Id)
+	groups, err := am.Store.GetAccountGroups(context.Background(), LockingStrengthShare, accountID)
 	require.NoError(t, err)
 	var groupHA1, groupHA2 *nbgroup.Group
 	for _, group := range groups {
@@ -1104,21 +1098,21 @@ func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
 		}
 	}
 
-	err = am.GroupDeletePeer(context.Background(), account.Id, groupHA1.ID, peer2ID)
+	err = am.GroupDeletePeer(context.Background(), accountID, groupHA1.ID, peer2ID)
 	require.NoError(t, err)
 
 	peer2RoutesAfterDelete, err := am.GetNetworkMap(context.Background(), peer2ID)
 	require.NoError(t, err)
 	assert.Len(t, peer2RoutesAfterDelete.Routes, 2, "after peer deletion group should have 2 client routes")
 
-	err = am.GroupDeletePeer(context.Background(), account.Id, groupHA2.ID, peer4ID)
+	err = am.GroupDeletePeer(context.Background(), accountID, groupHA2.ID, peer4ID)
 	require.NoError(t, err)
 
 	peer2RoutesAfterDelete, err = am.GetNetworkMap(context.Background(), peer2ID)
 	require.NoError(t, err)
 	assert.Len(t, peer2RoutesAfterDelete.Routes, 1, "after peer deletion group should have only 1 route")
 
-	err = am.GroupAddPeer(context.Background(), account.Id, groupHA2.ID, peer4ID)
+	err = am.GroupAddPeer(context.Background(), accountID, groupHA2.ID, peer4ID)
 	require.NoError(t, err)
 
 	peer1RoutesAfterAdd, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1129,7 +1123,7 @@ func TestGetNetworkMap_RouteSyncPeerGroups(t *testing.T) {
 	require.NoError(t, err)
 	assert.Len(t, peer2RoutesAfterAdd.Routes, 2, "HA route should have 2 client routes")
 
-	err = am.DeleteRoute(context.Background(), account.Id, newRoute.ID, userID)
+	err = am.DeleteRoute(context.Background(), accountID, newRoute.ID, userID)
 	require.NoError(t, err)
 
 	peer1DeletedRoute, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1159,7 +1153,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestRouteAccount(t, am)
+	accountID, err := initTestRouteAccount(t, am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
@@ -1168,7 +1162,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, newAccountRoutes.Routes, 0, "new accounts should have no routes")
 
-	createdRoute, err := am.CreateRoute(context.Background(), account.Id, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, peer1ID, []string{}, baseRoute.Description, baseRoute.NetID, baseRoute.Masquerade, baseRoute.Metric, baseRoute.Groups, baseRoute.AccessControlGroups, false, userID, baseRoute.KeepRoute)
+	createdRoute, err := am.CreateRoute(context.Background(), accountID, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, peer1ID, []string{}, baseRoute.Description, baseRoute.NetID, baseRoute.Masquerade, baseRoute.Metric, baseRoute.Groups, baseRoute.AccessControlGroups, false, userID, baseRoute.KeepRoute)
 	require.NoError(t, err)
 
 	noDisabledRoutes, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1182,7 +1176,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 	expectedRoute := enabledRoute.Copy()
 	expectedRoute.Peer = peer1Key
 
-	err = am.SaveRoute(context.Background(), account.Id, userID, enabledRoute)
+	err = am.SaveRoute(context.Background(), accountID, userID, enabledRoute)
 	require.NoError(t, err)
 
 	peer1Routes, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1194,7 +1188,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, peer2Routes.Routes, 0, "no routes for peers not in the distribution group")
 
-	err = am.GroupAddPeer(context.Background(), account.Id, routeGroup1, peer2ID)
+	err = am.GroupAddPeer(context.Background(), accountID, routeGroup1, peer2ID)
 	require.NoError(t, err)
 
 	peer2Routes, err = am.GetNetworkMap(context.Background(), peer2ID)
@@ -1207,10 +1201,10 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 		Name:  "peer1 group",
 		Peers: []string{peer1ID},
 	}
-	err = am.SaveGroup(context.Background(), account.Id, userID, newGroup)
+	err = am.SaveGroup(context.Background(), accountID, userID, newGroup)
 	require.NoError(t, err)
 
-	rules, err := am.ListPolicies(context.Background(), account.Id, "testingUser")
+	rules, err := am.ListPolicies(context.Background(), accountID, "testingUser")
 	require.NoError(t, err)
 
 	defaultRule := rules[0]
@@ -1219,10 +1213,10 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 	newPolicy.Rules[0].Sources = []string{newGroup.ID}
 	newPolicy.Rules[0].Destinations = []string{newGroup.ID}
 
-	_, err = am.SavePolicy(context.Background(), account.Id, userID, newPolicy)
+	_, err = am.SavePolicy(context.Background(), accountID, userID, newPolicy)
 	require.NoError(t, err)
 
-	err = am.DeletePolicy(context.Background(), account.Id, defaultRule.ID, userID)
+	err = am.DeletePolicy(context.Background(), accountID, defaultRule.ID, userID)
 	require.NoError(t, err)
 
 	peer1GroupRoutes, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1233,7 +1227,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, peer2GroupRoutes.Routes, 0, "we should not receive routes for peer2")
 
-	err = am.DeleteRoute(context.Background(), account.Id, enabledRoute.ID, userID)
+	err = am.DeleteRoute(context.Background(), accountID, enabledRoute.ID, userID)
 	require.NoError(t, err)
 
 	peer1DeletedRoute, err := am.GetNetworkMap(context.Background(), peer1ID)
@@ -1267,179 +1261,104 @@ func createRouterStore(t *testing.T) (Store, error) {
 	return store, nil
 }
 
-func initTestRouteAccount(t *testing.T, am *DefaultAccountManager) (*Account, error) {
+func initTestRouteAccount(t *testing.T, am *DefaultAccountManager) (string, error) {
 	t.Helper()
 
 	accountID := "testingAcc"
 	domain := "example.com"
 
-	account := newAccountWithId(context.Background(), accountID, userID, domain)
-	err := am.Store.SaveAccount(context.Background(), account)
+	err := newAccountWithId(context.Background(), am.Store, accountID, userID, domain)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	ips := account.getTakenIPs()
-	peer1IP, err := AllocatePeerIP(account.Network.Net, ips)
+	createPeer := func(peerID, peerKey, peerName, dnsLabel, kernel, core, platform, os string) (*nbpeer.Peer, error) {
+		ips, err := am.Store.GetTakenIPs(context.Background(), LockingStrengthShare, accountID)
+		if err != nil {
+			return nil, err
+		}
+
+		network, err := am.Store.GetAccountNetwork(context.Background(), LockingStrengthShare, accountID)
+		if err != nil {
+			return nil, err
+		}
+
+		peerIP, err := AllocatePeerIP(network.Net, ips)
+		if err != nil {
+			return nil, err
+		}
+
+		peer := &nbpeer.Peer{
+			IP:        peerIP,
+			AccountID: accountID,
+			ID:        peerID,
+			Key:       peerKey,
+			Name:      peerName,
+			DNSLabel:  dnsLabel,
+			UserID:    userID,
+			Meta: nbpeer.PeerSystemMeta{
+				Hostname:  peerName,
+				GoOS:      strings.ToLower(kernel),
+				Kernel:    kernel,
+				Core:      core,
+				Platform:  platform,
+				OS:        os,
+				WtVersion: "development",
+				UIVersion: "development",
+			},
+			Status: &nbpeer.PeerStatus{},
+		}
+		if err := am.Store.AddPeerToAccount(context.Background(), peer); err != nil {
+			return nil, err
+		}
+		return peer, nil
+	}
+
+	// Create peers
+	peer1, err := createPeer(peer1ID, peer1Key, "test-host1@netbird.io", "test-host1", "Linux", "21.04", "x86_64", "Ubuntu")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
-
-	peer1 := &nbpeer.Peer{
-		IP:       peer1IP,
-		ID:       peer1ID,
-		Key:      peer1Key,
-		Name:     "test-host1@netbird.io",
-		DNSLabel: "test-host1",
-		UserID:   userID,
-		Meta: nbpeer.PeerSystemMeta{
-			Hostname:  "test-host1@netbird.io",
-			GoOS:      "linux",
-			Kernel:    "Linux",
-			Core:      "21.04",
-			Platform:  "x86_64",
-			OS:        "Ubuntu",
-			WtVersion: "development",
-			UIVersion: "development",
-		},
-		Status: &nbpeer.PeerStatus{},
-	}
-	account.Peers[peer1.ID] = peer1
-
-	ips = account.getTakenIPs()
-	peer2IP, err := AllocatePeerIP(account.Network.Net, ips)
+	peer2, err := createPeer(peer2ID, peer2Key, "test-host2@netbird.io", "test-host2", "Linux", "21.04", "x86_64", "Ubuntu")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
-
-	peer2 := &nbpeer.Peer{
-		IP:       peer2IP,
-		ID:       peer2ID,
-		Key:      peer2Key,
-		Name:     "test-host2@netbird.io",
-		DNSLabel: "test-host2",
-		UserID:   userID,
-		Meta: nbpeer.PeerSystemMeta{
-			Hostname:  "test-host2@netbird.io",
-			GoOS:      "linux",
-			Kernel:    "Linux",
-			Core:      "21.04",
-			Platform:  "x86_64",
-			OS:        "Ubuntu",
-			WtVersion: "development",
-			UIVersion: "development",
-		},
-		Status: &nbpeer.PeerStatus{},
-	}
-	account.Peers[peer2.ID] = peer2
-
-	ips = account.getTakenIPs()
-	peer3IP, err := AllocatePeerIP(account.Network.Net, ips)
+	peer3, err := createPeer(peer3ID, peer3Key, "test-host3@netbird.io", "test-host3", "Darwin", "13.4.1", "arm64", "darwin")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
-
-	peer3 := &nbpeer.Peer{
-		IP:       peer3IP,
-		ID:       peer3ID,
-		Key:      peer3Key,
-		Name:     "test-host3@netbird.io",
-		DNSLabel: "test-host3",
-		UserID:   userID,
-		Meta: nbpeer.PeerSystemMeta{
-			Hostname:  "test-host3@netbird.io",
-			GoOS:      "darwin",
-			Kernel:    "Darwin",
-			Core:      "13.4.1",
-			Platform:  "arm64",
-			OS:        "darwin",
-			WtVersion: "development",
-			UIVersion: "development",
-		},
-		Status: &nbpeer.PeerStatus{},
-	}
-	account.Peers[peer3.ID] = peer3
-
-	ips = account.getTakenIPs()
-	peer4IP, err := AllocatePeerIP(account.Network.Net, ips)
+	peer4, err := createPeer(peer4ID, peer4Key, "test-host4@netbird.io", "test-host4", "Linux", "21.04", "x86_64", "Ubuntu")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
-
-	peer4 := &nbpeer.Peer{
-		IP:       peer4IP,
-		ID:       peer4ID,
-		Key:      peer4Key,
-		Name:     "test-host4@netbird.io",
-		DNSLabel: "test-host4",
-		UserID:   userID,
-		Meta: nbpeer.PeerSystemMeta{
-			Hostname:  "test-host4@netbird.io",
-			GoOS:      "linux",
-			Kernel:    "Linux",
-			Core:      "21.04",
-			Platform:  "x86_64",
-			OS:        "Ubuntu",
-			WtVersion: "development",
-			UIVersion: "development",
-		},
-		Status: &nbpeer.PeerStatus{},
-	}
-	account.Peers[peer4.ID] = peer4
-
-	ips = account.getTakenIPs()
-	peer5IP, err := AllocatePeerIP(account.Network.Net, ips)
+	peer5, err := createPeer(peer5ID, peer5Key, "test-host5@netbird.io", "test-host5", "Linux", "21.04", "x86_64", "Ubuntu")
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	peer5 := &nbpeer.Peer{
-		IP:       peer5IP,
-		ID:       peer5ID,
-		Key:      peer5Key,
-		Name:     "test-host5@netbird.io",
-		DNSLabel: "test-host5",
-		UserID:   userID,
-		Meta: nbpeer.PeerSystemMeta{
-			Hostname:  "test-host5@netbird.io",
-			GoOS:      "linux",
-			Kernel:    "Linux",
-			Core:      "21.04",
-			Platform:  "x86_64",
-			OS:        "Ubuntu",
-			WtVersion: "development",
-			UIVersion: "development",
-		},
-		Status: &nbpeer.PeerStatus{},
+	groupAll, err := am.GetGroupByName(context.Background(), "All", accountID)
+	if err != nil {
+		return "", err
 	}
-	account.Peers[peer5.ID] = peer5
 
-	err = am.Store.SaveAccount(context.Background(), account)
-	if err != nil {
-		return nil, err
-	}
-	groupAll, err := account.GetGroupAll()
-	if err != nil {
-		return nil, err
-	}
 	err = am.GroupAddPeer(context.Background(), accountID, groupAll.ID, peer1ID)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	err = am.GroupAddPeer(context.Background(), accountID, groupAll.ID, peer2ID)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	err = am.GroupAddPeer(context.Background(), accountID, groupAll.ID, peer3ID)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	err = am.GroupAddPeer(context.Background(), accountID, groupAll.ID, peer4ID)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	newGroup := []*nbgroup.Group{
+	newGroups := []*nbgroup.Group{
 		{
 			ID:    routeGroup1,
 			Name:  routeGroup1,
@@ -1471,15 +1390,12 @@ func initTestRouteAccount(t *testing.T, am *DefaultAccountManager) (*Account, er
 			Peers: []string{peer1.ID, peer4.ID},
 		},
 	}
-
-	for _, group := range newGroup {
-		err = am.SaveGroup(context.Background(), accountID, userID, group)
-		if err != nil {
-			return nil, err
-		}
+	err = am.SaveGroups(context.Background(), accountID, userID, newGroups)
+	if err != nil {
+		return "", err
 	}
 
-	return am.Store.GetAccount(context.Background(), account.Id)
+	return accountID, nil
 }
 
 func TestAccount_getPeersRoutesFirewall(t *testing.T) {
@@ -1487,8 +1403,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 		peerBIp = "100.65.80.39"
 		peerCIp = "100.65.254.139"
 		peerHIp = "100.65.29.55"
-		peerJIp = "100.65.29.65"
-		peerKIp = "100.65.29.66"
 	)
 
 	account := &Account{
@@ -1544,16 +1458,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 				IP:     net.ParseIP(peerHIp),
 				Status: &nbpeer.PeerStatus{},
 			},
-			"peerJ": {
-				ID:     "peerJ",
-				IP:     net.ParseIP(peerJIp),
-				Status: &nbpeer.PeerStatus{},
-			},
-			"peerK": {
-				ID:     "peerK",
-				IP:     net.ParseIP(peerKIp),
-				Status: &nbpeer.PeerStatus{},
-			},
 		},
 		Groups: map[string]*nbgroup.Group{
 			"routingPeer1": {
@@ -1580,11 +1484,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 				Name:  "Route2",
 				Peers: []string{},
 			},
-			"route4": {
-				ID:    "route4",
-				Name:  "route4",
-				Peers: []string{},
-			},
 			"finance": {
 				ID:   "finance",
 				Name: "Finance",
@@ -1602,28 +1501,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 					"peerB",
 				},
 			},
-			"qa": {
-				ID:   "qa",
-				Name: "QA",
-				Peers: []string{
-					"peerJ",
-					"peerK",
-				},
-			},
-			"restrictQA": {
-				ID:   "restrictQA",
-				Name: "restrictQA",
-				Peers: []string{
-					"peerJ",
-				},
-			},
-			"unrestrictedQA": {
-				ID:   "unrestrictedQA",
-				Name: "unrestrictedQA",
-				Peers: []string{
-					"peerK",
-				},
-			},
 			"contractors": {
 				ID:    "contractors",
 				Name:  "Contractors",
@@ -1671,19 +1548,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 				Groups:              []string{"contractors"},
 				AccessControlGroups: []string{},
 			},
-			"route4": {
-				ID:                  "route4",
-				Network:             netip.MustParsePrefix("192.168.10.0/16"),
-				NetID:               "route4",
-				NetworkType:         route.IPv4Network,
-				PeerGroups:          []string{"routingPeer1"},
-				Description:         "Route4",
-				Masquerade:          false,
-				Metric:              9999,
-				Enabled:             true,
-				Groups:              []string{"qa"},
-				AccessControlGroups: []string{"route4"},
-			},
 		},
 		Policies: []*Policy{
 			{
@@ -1738,49 +1602,6 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 					},
 				},
 			},
-			{
-				ID:      "RuleRoute4",
-				Name:    "RuleRoute4",
-				Enabled: true,
-				Rules: []*PolicyRule{
-					{
-						ID:            "RuleRoute4",
-						Name:          "RuleRoute4",
-						Bidirectional: true,
-						Enabled:       true,
-						Protocol:      PolicyRuleProtocolTCP,
-						Action:        PolicyTrafficActionAccept,
-						Ports:         []string{"80"},
-						Sources: []string{
-							"restrictQA",
-						},
-						Destinations: []string{
-							"route4",
-						},
-					},
-				},
-			},
-			{
-				ID:      "RuleRoute5",
-				Name:    "RuleRoute5",
-				Enabled: true,
-				Rules: []*PolicyRule{
-					{
-						ID:            "RuleRoute5",
-						Name:          "RuleRoute5",
-						Bidirectional: true,
-						Enabled:       true,
-						Protocol:      PolicyRuleProtocolALL,
-						Action:        PolicyTrafficActionAccept,
-						Sources: []string{
-							"unrestrictedQA",
-						},
-						Destinations: []string{
-							"route4",
-						},
-					},
-				},
-			},
 		},
 	}
 
@@ -1805,7 +1626,7 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 
 	t.Run("check peer routes firewall rules", func(t *testing.T) {
 		routesFirewallRules := account.getPeerRoutesFirewallRules(context.Background(), "peerA", validatedPeers)
-		assert.Len(t, routesFirewallRules, 4)
+		assert.Len(t, routesFirewallRules, 2)
 
 		expectedRoutesFirewallRules := []*RouteFirewallRule{
 			{
@@ -1831,32 +1652,12 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 				Port:        320,
 			},
 		}
-		additionalFirewallRule := []*RouteFirewallRule{
-			{
-				SourceRanges: []string{
-					fmt.Sprintf(AllowedIPsFormat, peerJIp),
-				},
-				Action:      "accept",
-				Destination: "192.168.10.0/16",
-				Protocol:    "tcp",
-				Port:        80,
-			},
-			{
-				SourceRanges: []string{
-					fmt.Sprintf(AllowedIPsFormat, peerKIp),
-				},
-				Action:      "accept",
-				Destination: "192.168.10.0/16",
-				Protocol:    "all",
-			},
-		}
-
-		assert.ElementsMatch(t, orderRuleSourceRanges(routesFirewallRules), orderRuleSourceRanges(append(expectedRoutesFirewallRules, additionalFirewallRule...)))
+		assert.ElementsMatch(t, routesFirewallRules, expectedRoutesFirewallRules)
 
 		// peerD is also the routing peer for route1, should contain same routes firewall rules as peerA
 		routesFirewallRules = account.getPeerRoutesFirewallRules(context.Background(), "peerD", validatedPeers)
 		assert.Len(t, routesFirewallRules, 2)
-		assert.ElementsMatch(t, orderRuleSourceRanges(routesFirewallRules), orderRuleSourceRanges(expectedRoutesFirewallRules))
+		assert.ElementsMatch(t, routesFirewallRules, expectedRoutesFirewallRules)
 
 		// peerE is a single routing peer for route 2 and route 3
 		routesFirewallRules = account.getPeerRoutesFirewallRules(context.Background(), "peerE", validatedPeers)
@@ -1885,7 +1686,7 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 				IsDynamic:    true,
 			},
 		}
-		assert.ElementsMatch(t, orderRuleSourceRanges(routesFirewallRules), orderRuleSourceRanges(expectedRoutesFirewallRules))
+		assert.ElementsMatch(t, routesFirewallRules, expectedRoutesFirewallRules)
 
 		// peerC is part of route1 distribution groups but should not receive the routes firewall rules
 		routesFirewallRules = account.getPeerRoutesFirewallRules(context.Background(), "peerC", validatedPeers)
@@ -1894,22 +1695,14 @@ func TestAccount_getPeersRoutesFirewall(t *testing.T) {
 
 }
 
-// orderList is a helper function to sort a list of strings
-func orderRuleSourceRanges(ruleList []*RouteFirewallRule) []*RouteFirewallRule {
-	for _, rule := range ruleList {
-		sort.Strings(rule.SourceRanges)
-	}
-	return ruleList
-}
-
 func TestRouteAccountPeersUpdate(t *testing.T) {
 	manager, err := createRouterManager(t)
 	require.NoError(t, err, "failed to create account manager")
 
-	account, err := initTestRouteAccount(t, manager)
+	accountID, err := initTestRouteAccount(t, manager)
 	require.NoError(t, err, "failed to init testing account")
 
-	err = manager.SaveGroups(context.Background(), account.Id, userID, []*nbgroup.Group{
+	err = manager.SaveGroups(context.Background(), accountID, userID, []*nbgroup.Group{
 		{
 			ID:    "groupA",
 			Name:  "GroupA",
@@ -1955,7 +1748,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 		}()
 
 		_, err := manager.CreateRoute(
-			context.Background(), account.Id, route.Network, route.NetworkType, route.Domains, route.Peer,
+			context.Background(), accountID, route.Network, route.NetworkType, route.Domains, route.Peer,
 			route.PeerGroups, route.Description, route.NetID, route.Masquerade, route.Metric,
 			route.Groups, []string{}, true, userID, route.KeepRoute,
 		)
@@ -1991,7 +1784,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 		}()
 
 		_, err := manager.CreateRoute(
-			context.Background(), account.Id, route.Network, route.NetworkType, route.Domains, route.Peer,
+			context.Background(), accountID, route.Network, route.NetworkType, route.Domains, route.Peer,
 			route.PeerGroups, route.Description, route.NetID, route.Masquerade, route.Metric,
 			route.Groups, []string{}, true, userID, route.KeepRoute,
 		)
@@ -2027,7 +1820,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 		}()
 
 		newRoute, err := manager.CreateRoute(
-			context.Background(), account.Id, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, baseRoute.Peer,
+			context.Background(), accountID, baseRoute.Network, baseRoute.NetworkType, baseRoute.Domains, baseRoute.Peer,
 			baseRoute.PeerGroups, baseRoute.Description, baseRoute.NetID, baseRoute.Masquerade, baseRoute.Metric,
 			baseRoute.Groups, []string{}, true, userID, baseRoute.KeepRoute,
 		)
@@ -2051,7 +1844,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		err := manager.SaveRoute(context.Background(), account.Id, userID, &baseRoute)
+		err := manager.SaveRoute(context.Background(), accountID, userID, &baseRoute)
 		require.NoError(t, err)
 
 		select {
@@ -2069,7 +1862,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		err := manager.DeleteRoute(context.Background(), account.Id, baseRoute.ID, userID)
+		err := manager.DeleteRoute(context.Background(), accountID, baseRoute.ID, userID)
 		require.NoError(t, err)
 
 		select {
@@ -2093,7 +1886,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			Groups:      []string{routeGroup1},
 		}
 		_, err := manager.CreateRoute(
-			context.Background(), account.Id, newRoute.Network, newRoute.NetworkType, newRoute.Domains, newRoute.Peer,
+			context.Background(), accountID, newRoute.Network, newRoute.NetworkType, newRoute.Domains, newRoute.Peer,
 			newRoute.PeerGroups, newRoute.Description, newRoute.NetID, newRoute.Masquerade, newRoute.Metric,
 			newRoute.Groups, []string{}, true, userID, newRoute.KeepRoute,
 		)
@@ -2105,7 +1898,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
+		err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
 			ID:    "groupB",
 			Name:  "GroupB",
 			Peers: []string{peer1ID},
@@ -2133,7 +1926,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			Groups:      []string{"groupC"},
 		}
 		_, err := manager.CreateRoute(
-			context.Background(), account.Id, newRoute.Network, newRoute.NetworkType, newRoute.Domains, newRoute.Peer,
+			context.Background(), accountID, newRoute.Network, newRoute.NetworkType, newRoute.Domains, newRoute.Peer,
 			newRoute.PeerGroups, newRoute.Description, newRoute.NetID, newRoute.Masquerade, newRoute.Metric,
 			newRoute.Groups, []string{}, true, userID, newRoute.KeepRoute,
 		)
@@ -2145,7 +1938,7 @@ func TestRouteAccountPeersUpdate(t *testing.T) {
 			close(done)
 		}()
 
-		err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
+		err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
 			ID:    "groupC",
 			Name:  "GroupC",
 			Peers: []string{peer1ID},

management/server/setupkey.go

@@ -12,10 +12,9 @@ import (
 	"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 (
@@ -277,7 +276,7 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 // SaveSetupKey saves the provided SetupKey to the database overriding the existing one.
 // Due to the unique nature of a SetupKey certain properties must not be overwritten
 // (e.g. the key itself, creation date, ID, etc).
-// These properties are overwritten: AutoGroups, Revoked (only from false to true), and the UpdatedAt. The rest is copied from the existing key.
+// 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) {
 	if keyToSave == nil {
 		return nil, status.Errorf(status.InvalidArgument, "provided setup key to update is nil")
@@ -313,12 +312,9 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 			return err
 		}
 
-		if oldKey.Revoked && !keyToSave.Revoked {
-			return status.Errorf(status.InvalidArgument, "can't un-revoke a revoked setup key")
-		}
-
-		// only auto groups, revoked status (from false to true) can be updated
+		// 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()
@@ -379,7 +375,7 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 		return nil, status.NewAdminPermissionError()
 	}
 
-	setupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyID)
+	setupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
 	if err != nil {
 		return nil, err
 	}

management/server/setupkey_test.go

@@ -25,12 +25,12 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 	}
 
 	userID := "testingUser"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "")
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userID, "")
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = manager.SaveGroups(context.Background(), account.Id, userID, []*nbgroup.Group{
+	err = manager.SaveGroups(context.Background(), accountID, userID, []*nbgroup.Group{
 		{
 			ID:    "group_1",
 			Name:  "group_name_1",
@@ -49,16 +49,18 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 	expiresIn := time.Hour
 	keyName := "my-test-key"
 
-	key, err := manager.CreateSetupKey(context.Background(), account.Id, keyName, SetupKeyReusable, expiresIn, []string{},
+	key, err := manager.CreateSetupKey(context.Background(), accountID, keyName, SetupKeyReusable, expiresIn, []string{},
 		SetupKeyUnlimitedUsage, userID, false)
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	autoGroups := []string{"group_1", "group_2"}
+	newKeyName := "my-new-test-key"
 	revoked := true
-	newKey, err := manager.SaveSetupKey(context.Background(), account.Id, &SetupKey{
+	newKey, err := manager.SaveSetupKey(context.Background(), accountID, &SetupKey{
 		Id:         key.Id,
+		Name:       newKeyName,
 		Revoked:    revoked,
 		AutoGroups: autoGroups,
 	}, userID)
@@ -66,27 +68,28 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	assertKey(t, newKey, keyName, revoked, "reusable", 0, key.CreatedAt, key.ExpiresAt,
+	assertKey(t, newKey, newKeyName, revoked, "reusable", 0, key.CreatedAt, key.ExpiresAt,
 		key.Id, time.Now().UTC(), autoGroups, true)
 
 	// check the corresponding events that should have been generated
-	ev := getEvent(t, account.Id, manager, activity.SetupKeyRevoked)
+	ev := getEvent(t, accountID, manager, activity.SetupKeyRevoked)
 
 	assert.NotNil(t, ev)
-	assert.Equal(t, account.Id, ev.AccountID)
-	assert.Equal(t, keyName, ev.Meta["name"])
+	assert.Equal(t, accountID, ev.AccountID)
+	assert.Equal(t, newKeyName, ev.Meta["name"])
 	assert.Equal(t, fmt.Sprint(key.Type), fmt.Sprint(ev.Meta["type"]))
 	assert.NotEmpty(t, ev.Meta["key"])
 	assert.Equal(t, userID, ev.InitiatorID)
 	assert.Equal(t, key.Id, ev.TargetID)
 
-	groupAll, err := account.GetGroupAll()
+	groupAll, err := manager.GetGroupByName(context.Background(), "All", accountID)
 	assert.NoError(t, err)
 
 	// saving setup key with All group assigned to auto groups should return error
 	autoGroups = append(autoGroups, groupAll.ID)
-	_, err = manager.SaveSetupKey(context.Background(), account.Id, &SetupKey{
+	_, err = manager.SaveSetupKey(context.Background(), accountID, &SetupKey{
 		Id:         key.Id,
+		Name:       newKeyName,
 		Revoked:    revoked,
 		AutoGroups: autoGroups,
 	}, userID)
@@ -100,12 +103,12 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 	}
 
 	userID := "testingUser"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "")
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userID, "")
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
+	err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
 		ID:    "group_1",
 		Name:  "group_name_1",
 		Peers: []string{},
@@ -114,7 +117,7 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
+	err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
 		ID:    "group_2",
 		Name:  "group_name_2",
 		Peers: []string{},
@@ -123,7 +126,7 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	groupAll, err := account.GetGroupAll()
+	groupAll, err := manager.GetGroupByName(context.Background(), "All", accountID)
 	assert.NoError(t, err)
 
 	type testCase struct {
@@ -167,7 +170,7 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 
 	for _, tCase := range []testCase{testCase1, testCase2, testCase3} {
 		t.Run(tCase.name, func(t *testing.T) {
-			key, err := manager.CreateSetupKey(context.Background(), account.Id, tCase.expectedKeyName, SetupKeyReusable, expiresIn,
+			key, err := manager.CreateSetupKey(context.Background(), accountID, tCase.expectedKeyName, SetupKeyReusable, expiresIn,
 				tCase.expectedGroups, SetupKeyUnlimitedUsage, userID, false)
 
 			if tCase.expectedFailure {
@@ -186,10 +189,10 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 				tCase.expectedUpdatedAt, tCase.expectedGroups, false)
 
 			// check the corresponding events that should have been generated
-			ev := getEvent(t, account.Id, manager, activity.SetupKeyCreated)
+			ev := getEvent(t, accountID, manager, activity.SetupKeyCreated)
 
 			assert.NotNil(t, ev)
-			assert.Equal(t, account.Id, ev.AccountID)
+			assert.Equal(t, accountID, ev.AccountID)
 			assert.Equal(t, tCase.expectedKeyName, ev.Meta["name"])
 			assert.Equal(t, tCase.expectedType, fmt.Sprint(ev.Meta["type"]))
 			assert.NotEmpty(t, ev.Meta["key"])
@@ -205,46 +208,27 @@ func TestGetSetupKeys(t *testing.T) {
 	}
 
 	userID := "testingUser"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "")
+	accountID, err := manager.GetOrCreateAccountIDByUser(context.Background(), userID, "")
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	plainKey, err := manager.CreateSetupKey(context.Background(), account.Id, "key1", SetupKeyReusable, time.Hour, nil, SetupKeyUnlimitedUsage, userID, false)
+	err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
+		ID:    "group_1",
+		Name:  "group_name_1",
+		Peers: []string{},
+	})
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	type testCase struct {
-		name            string
-		keyId           string
-		expectedFailure bool
-	}
-
-	testCase1 := testCase{
-		name:            "Should get existing Setup Key",
-		keyId:           plainKey.Id,
-		expectedFailure: false,
-	}
-	testCase2 := testCase{
-		name:            "Should fail to get non-existent Setup Key",
-		keyId:           "some key",
-		expectedFailure: true,
-	}
-
-	for _, tCase := range []testCase{testCase1, testCase2} {
-		t.Run(tCase.name, func(t *testing.T) {
-			key, err := manager.GetSetupKey(context.Background(), account.Id, userID, tCase.keyId)
-
-			if tCase.expectedFailure {
-				if err == nil {
-					t.Fatal("expected to fail")
-				}
-				return
-			}
-
-			assert.NotEqual(t, plainKey.Key, key.Key)
-		})
+	err = manager.SaveGroup(context.Background(), accountID, userID, &nbgroup.Group{
+		ID:    "group_2",
+		Name:  "group_name_2",
+		Peers: []string{},
+	})
+	if err != nil {
+		t.Fatal(err)
 	}
 }
 
@@ -464,31 +448,3 @@ func TestSetupKeyAccountPeersUpdate(t *testing.T) {
 		}
 	})
 }
-
-func TestDefaultAccountManager_CreateSetupKey_ShouldNotAllowToUpdateRevokedKey(t *testing.T) {
-	manager, err := createManager(t)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	userID := "testingUser"
-	account, err := manager.GetOrCreateAccountByUser(context.Background(), userID, "")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	key, err := manager.CreateSetupKey(context.Background(), account.Id, "testName", SetupKeyReusable, time.Hour, nil, SetupKeyUnlimitedUsage, userID, false)
-	assert.NoError(t, err)
-
-	// revoke the key
-	updateKey := key.Copy()
-	updateKey.Revoked = true
-	_, err = manager.SaveSetupKey(context.Background(), account.Id, updateKey, userID)
-	assert.NoError(t, err)
-
-	// re-activate revoked key
-	updateKey.Revoked = false
-	_, err = manager.SaveSetupKey(context.Background(), account.Id, updateKey, userID)
-	assert.Error(t, err, "should not allow to update revoked key")
-
-}

management/server/sql_store.go

@@ -261,7 +261,7 @@ func (s *SqlStore) DeleteAccount(ctx context.Context, account *Account) error {
 			return result.Error
 		}
 
-		result = tx.Select(clause.Associations).Delete(account.UsersG, "account_id = ?", account.Id)
+		result = tx.Debug().Select(clause.Associations).Delete(account.UsersG, "account_id = ?", account.Id)
 		if result.Error != nil {
 			return result.Error
 		}
@@ -300,12 +300,12 @@ func (s *SqlStore) GetInstallationID() string {
 	return installation.InstallationIDValue
 }
 
-func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
+func (s *SqlStore) SavePeer(ctx context.Context, lockStrength LockingStrength, accountID string, peer *nbpeer.Peer) error {
 	// 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.Clauses(clause.Locking{Strength: string(lockStrength)}).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)
@@ -332,30 +332,33 @@ func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.
 	return nil
 }
 
-func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error {
+func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, lockStrength LockingStrength, accountID string, domain string, category string, isPrimaryDomain *bool) error {
 	accountCopy := Account{
-		Domain:                 domain,
-		DomainCategory:         category,
-		IsDomainPrimaryAccount: isPrimaryDomain,
+		Domain:         domain,
+		DomainCategory: category,
 	}
 
-	fieldsToUpdate := []string{"domain", "domain_category", "is_domain_primary_account"}
-	result := s.db.Model(&Account{}).
-		Select(fieldsToUpdate).
-		Where(idQueryCondition, accountID).
-		Updates(&accountCopy)
+	fieldsToUpdate := []string{"domain", "domain_category"}
+	if isPrimaryDomain != nil {
+		accountCopy.IsDomainPrimaryAccount = *isPrimaryDomain
+		fieldsToUpdate = append(fieldsToUpdate, "is_domain_primary_account")
+	}
+
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).Select(fieldsToUpdate).
+		Where(idQueryCondition, accountID).Updates(&accountCopy)
 	if result.Error != nil {
-		return result.Error
+		log.WithContext(ctx).Errorf("failed to update account domain attributes in store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to update account domain attributes in store")
 	}
 
 	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "account %s", accountID)
+		return status.NewAccountNotFoundError(accountID)
 	}
 
 	return nil
 }
 
-func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
+func (s *SqlStore) SavePeerStatus(ctx context.Context, lockStrength LockingStrength, accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
 	var peerCopy nbpeer.Peer
 	peerCopy.Status = &peerStatus
 
@@ -363,7 +366,7 @@ func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.Pe
 		"peer_status_last_seen", "peer_status_connected",
 		"peer_status_login_expired", "peer_status_required_approval",
 	}
-	result := s.db.Model(&nbpeer.Peer{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Select(fieldsToUpdate).
 		Where(accountAndIDQueryCondition, accountID, peerID).
 		Updates(&peerCopy)
@@ -378,14 +381,14 @@ func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.Pe
 	return nil
 }
 
-func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.Peer) error {
+func (s *SqlStore) SavePeerLocation(ctx context.Context, lockStrength LockingStrength, accountID string, peerWithLocation *nbpeer.Peer) error {
 	// 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,
 	// updating the struct ensures the correct data format is inserted into the database.
 	peerCopy.Location = peerWithLocation.Location
 
-	result := s.db.Model(&nbpeer.Peer{}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
 		Where(accountAndIDQueryCondition, accountID, peerWithLocation.ID).
 		Updates(peerCopy)
 
@@ -401,24 +404,16 @@ 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 {
-	usersToSave := make([]User, 0, len(users))
-	for _, user := range users {
-		user.AccountID = accountID
-		for id, pat := range user.PATs {
-			pat.ID = id
-			user.PATsG = append(user.PATsG, *pat)
-		}
-		usersToSave = append(usersToSave, *user)
-	}
-	err := s.db.Session(&gorm.Session{FullSaveAssociations: true}).
-		Clauses(clause.OnConflict{UpdateAll: true}).
-		Create(&usersToSave).Error
-	if err != nil {
-		return status.Errorf(status.Internal, "failed to save users to store: %v", err)
+func (s *SqlStore) SaveUsers(ctx context.Context, lockStrength LockingStrength, users []*User) error {
+	if len(users) == 0 {
+		return nil
 	}
 
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&users)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save users to store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to save users to store")
+	}
 	return nil
 }
 
@@ -426,7 +421,8 @@ func (s *SqlStore) SaveUsers(accountID string, users map[string]*User) error {
 func (s *SqlStore) SaveUser(ctx context.Context, lockStrength LockingStrength, user *User) error {
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(user)
 	if result.Error != nil {
-		return status.Errorf(status.Internal, "failed to save user to store: %v", result.Error)
+		log.WithContext(ctx).Errorf("failed to save user to store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to save user to store")
 	}
 	return nil
 }
@@ -444,16 +440,6 @@ func (s *SqlStore) SaveGroups(ctx context.Context, lockStrength LockingStrength,
 	return nil
 }
 
-// DeleteHashedPAT2TokenIDIndex is noop in SqlStore
-func (s *SqlStore) DeleteHashedPAT2TokenIDIndex(hashedToken string) error {
-	return nil
-}
-
-// DeleteTokenID2UserIDIndex is noop in SqlStore
-func (s *SqlStore) DeleteTokenID2UserIDIndex(tokenID string) error {
-	return nil
-}
-
 func (s *SqlStore) GetAccountByPrivateDomain(ctx context.Context, domain string) (*Account, error) {
 	accountID, err := s.GetAccountIDByPrivateDomain(ctx, LockingStrengthShare, domain)
 	if err != nil {
@@ -513,30 +499,17 @@ func (s *SqlStore) GetTokenIDByHashedToken(ctx context.Context, hashedToken stri
 	return token.ID, nil
 }
 
-func (s *SqlStore) GetUserByTokenID(ctx context.Context, tokenID string) (*User, error) {
-	var token PersonalAccessToken
-	result := s.db.First(&token, idQueryCondition, tokenID)
+func (s *SqlStore) GetUserByPATID(ctx context.Context, lockStrength LockingStrength, patID string) (*User, error) {
+	var user User
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Joins("JOIN personal_access_tokens ON personal_access_tokens.user_id = users.id").
+		Where("personal_access_tokens.id = ?", patID).First(&user)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
+			return nil, status.NewPATNotFoundError(patID)
 		}
-		log.WithContext(ctx).Errorf("error when getting token from the store: %s", result.Error)
-		return nil, status.NewGetAccountFromStoreError(result.Error)
-	}
-
-	if token.UserID == "" {
-		return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
-	}
-
-	var user User
-	result = s.db.Preload("PATsG").First(&user, idQueryCondition, token.UserID)
-	if result.Error != nil {
-		return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
-	}
-
-	user.PATs = make(map[string]*PersonalAccessToken, len(user.PATsG))
-	for _, pat := range user.PATsG {
-		user.PATs[pat.ID] = pat.Copy()
+		log.WithContext(ctx).Errorf("failed to get token user from the store: %s", result.Error)
+		return nil, status.NewGetUserFromStoreError()
 	}
 
 	return &user, nil
@@ -544,8 +517,7 @@ func (s *SqlStore) GetUserByTokenID(ctx context.Context, tokenID string) (*User,
 
 func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*User, error) {
 	var user User
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Preload(clause.Associations).First(&user, idQueryCondition, userID)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&user, idQueryCondition, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.NewUserNotFoundError(userID)
@@ -556,6 +528,25 @@ func (s *SqlStore) GetUserByUserID(ctx context.Context, lockStrength LockingStre
 	return &user, nil
 }
 
+func (s *SqlStore) DeleteUser(ctx context.Context, lockStrength LockingStrength, accountID, userID string) error {
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		result := tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Delete(&PersonalAccessToken{}, "user_id = ?", userID)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		return tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Delete(&User{}, accountAndIDQueryCondition, accountID, userID).Error
+	})
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to delete user from the store: %s", err)
+		return status.Errorf(status.Internal, "failed to delete user from store")
+	}
+
+	return nil
+}
+
 func (s *SqlStore) GetAccountUsers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*User, error) {
 	var users []*User
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&users, accountIDCondition, accountID)
@@ -740,9 +731,10 @@ func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string)
 	return accountID, nil
 }
 
-func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
+func (s *SqlStore) GetAccountIDByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (string, error) {
 	var accountID string
-	result := s.db.Model(&User{}).Select("account_id").Where(idQueryCondition, userID).First(&accountID)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).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")
@@ -753,6 +745,20 @@ func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
 	return accountID, nil
 }
 
+func (s *SqlStore) GetAccountIDByPeerID(ctx context.Context, lockStrength LockingStrength, peerID string) (string, error) {
+	var accountID string
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&nbpeer.Peer{}).
+		Select("account_id").Where(idQueryCondition, peerID).First(&accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "peer %s account not found", peerID)
+		}
+		return "", status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return accountID, nil
+}
+
 func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string) (string, error) {
 	var accountID string
 	result := s.db.Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
@@ -831,7 +837,7 @@ func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength Locking
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&peer, keyQueryCondition, peerKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "peer not found")
+			return nil, status.NewPeerNotFoundError(peerKey)
 		}
 		return nil, status.Errorf(status.Internal, "issue getting peer from store: %s", result.Error)
 	}
@@ -850,6 +856,31 @@ func (s *SqlStore) GetAccountSettings(ctx context.Context, lockStrength LockingS
 	return accountSettings.Settings, nil
 }
 
+func (s *SqlStore) GetAccountCreatedBy(ctx context.Context, lockStrength LockingStrength, accountID string) (string, error) {
+	var createdBy string
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+		Select("created_by").First(&createdBy, idQueryCondition, accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.NewAccountNotFoundError(accountID)
+		}
+		return "", status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return createdBy, nil
+}
+
+func (s *SqlStore) GetTotalAccounts(ctx context.Context) (int64, error) {
+	var count int64
+	result := s.db.Model(&Account{}).Count(&count)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get total accounts from store: %s", result.Error)
+		return 0, status.Errorf(status.Internal, "failed to get total accounts from store")
+	}
+
+	return count, nil
+}
+
 // SaveUserLastLogin stores the last login time for a user in DB.
 func (s *SqlStore) SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error {
 	var user User
@@ -1066,9 +1097,29 @@ func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountId string, peerId
 	return nil
 }
 
+// GetAccountPeers retrieves peers for an account.
+func (s *SqlStore) GetAccountPeers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error) {
+	var peers []*nbpeer.Peer
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&peers, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get peers from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get peers from store")
+	}
+
+	return peers, nil
+}
+
 // 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)
+	var peers []*nbpeer.Peer
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Find(&peers, "account_id = ? AND user_id = ?", accountID, userID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get peers from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get peers from store")
+	}
+
+	return peers, nil
 }
 
 func (s *SqlStore) AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error {
@@ -1086,7 +1137,7 @@ func (s *SqlStore) GetPeerByID(ctx context.Context, lockStrength LockingStrength
 		First(&peer, accountAndIDQueryCondition, accountID, peerID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "peer not found")
+			return nil, status.NewPeerNotFoundError(peerID)
 		}
 		log.WithContext(ctx).Errorf("failed to get peer from store: %s", result.Error)
 		return nil, status.Errorf(status.Internal, "failed to get peer from store")
@@ -1112,6 +1163,68 @@ func (s *SqlStore) GetPeersByIDs(ctx context.Context, lockStrength LockingStreng
 	return peersMap, nil
 }
 
+// GetAccountPeersWithExpiration retrieves a list of peers that have login expiration enabled and added by a user.
+func (s *SqlStore) GetAccountPeersWithExpiration(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error) {
+	var peers []*nbpeer.Peer
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Where("login_expiration_enabled = ? AND user_id IS NOT NULL AND user_id != ''", true).
+		Find(&peers, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get peers with expiration from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get peers with expiration from store")
+	}
+
+	return peers, nil
+}
+
+// GetAccountPeersWithInactivity retrieves a list of peers that have login expiration enabled and added by a user.
+func (s *SqlStore) GetAccountPeersWithInactivity(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error) {
+	var peers []*nbpeer.Peer
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Where("inactivity_expiration_enabled = ? AND user_id IS NOT NULL AND user_id != ''", true).
+		Find(&peers, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get peers with inactivity from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get peers with inactivity from store")
+	}
+
+	return peers, nil
+}
+
+// GetAllEphemeralPeers retrieves all peers with Ephemeral set to true across all accounts, optimized for batch processing.
+func (s *SqlStore) GetAllEphemeralPeers(ctx context.Context, lockStrength LockingStrength) ([]*nbpeer.Peer, error) {
+	var allEphemeralPeers, batchPeers []*nbpeer.Peer
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Where("ephemeral = ?", true).
+		FindInBatches(&batchPeers, 1000, func(tx *gorm.DB, batch int) error {
+			allEphemeralPeers = append(allEphemeralPeers, batchPeers...)
+			return nil
+		})
+
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to retrieve ephemeral peers: %s", result.Error)
+		return nil, fmt.Errorf("failed to retrieve ephemeral peers")
+	}
+
+	return allEphemeralPeers, nil
+}
+
+// DeletePeer removes a peer from the store.
+func (s *SqlStore) DeletePeer(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Delete(&nbpeer.Peer{}, accountAndIDQueryCondition, accountID, peerID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to delete peer from the store: %s", err)
+		return status.Errorf(status.Internal, "failed to delete peer from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewPeerNotFoundError(peerID)
+	}
+
+	return nil
+}
+
 func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, lockStrength LockingStrength, accountId string) error {
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		Model(&Account{}).Where(idQueryCondition, accountId).Update("network_serial", gorm.Expr("network_serial + 1"))
@@ -1123,7 +1236,6 @@ func (s *SqlStore) IncrementNetworkSerial(ctx context.Context, lockStrength Lock
 }
 
 func (s *SqlStore) ExecuteInTransaction(ctx context.Context, operation func(store Store) error) error {
-	startTime := time.Now()
 	tx := s.db.Begin()
 	if tx.Error != nil {
 		return tx.Error
@@ -1134,15 +1246,7 @@ func (s *SqlStore) ExecuteInTransaction(ctx context.Context, operation func(stor
 		tx.Rollback()
 		return err
 	}
-
-	err = tx.Commit().Error
-
-	log.WithContext(ctx).Tracef("transaction took %v", time.Since(startTime))
-	if s.metrics != nil {
-		s.metrics.StoreMetrics().CountTransactionDuration(time.Since(startTime))
-	}
-
-	return err
+	return tx.Commit().Error
 }
 
 func (s *SqlStore) withTx(tx *gorm.DB) Store {
@@ -1288,7 +1392,7 @@ func (s *SqlStore) DeleteGroups(ctx context.Context, strength LockingStrength, a
 		Delete(&nbgroup.Group{}, accountAndIDsQueryCondition, accountID, groupIDs)
 	if result.Error != nil {
 		log.WithContext(ctx).Errorf("failed to delete groups from store: %v", result.Error)
-		return status.Errorf(status.Internal, "failed to delete groups from store")
+		return status.Errorf(status.Internal, "failed to delete groups from store: %v", result.Error)
 	}
 
 	return nil
@@ -1345,17 +1449,21 @@ func (s *SqlStore) SavePolicy(ctx context.Context, lockStrength LockingStrength,
 }
 
 func (s *SqlStore) DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&Policy{}, accountAndIDQueryCondition, accountID, policyID)
-	if err := result.Error; err != nil {
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		result := tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Delete(&PolicyRule{}, "policy_id = ?", policyID)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		return tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Delete(&Policy{}, accountAndIDQueryCondition, accountID, policyID).Error
+	})
+	if err != nil {
 		log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
 		return status.Errorf(status.Internal, "failed to delete policy from store")
 	}
 
-	if result.RowsAffected == 0 {
-		return status.NewPolicyNotFoundError(policyID)
-	}
-
 	return nil
 }
 
@@ -1433,12 +1541,58 @@ func (s *SqlStore) DeletePostureChecks(ctx context.Context, lockStrength Locking
 
 // 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)
+	var routes []*route.Route
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Find(&routes, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get routes from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get routes from store")
+	}
+
+	return routes, nil
 }
 
 // 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)
+func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrength, accountID string, routeID string) (*route.Route, error) {
+	var route *route.Route
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		First(&route, accountAndIDQueryCondition, accountID, routeID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.NewRouteNotFoundError(routeID)
+		}
+		log.WithContext(ctx).Errorf("failed to get route from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get route from store")
+	}
+
+	return route, nil
+}
+
+// SaveRoute saves a route to the database.
+func (s *SqlStore) SaveRoute(ctx context.Context, lockStrength LockingStrength, route *route.Route) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(route)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to save route to the store: %s", err)
+		return status.Errorf(status.Internal, "failed to save route to store")
+	}
+
+	return nil
+}
+
+// DeleteRoute deletes a route from the database.
+func (s *SqlStore) DeleteRoute(ctx context.Context, lockStrength LockingStrength, accountID, routeID string) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Delete(&route.Route{}, accountAndIDQueryCondition, accountID, routeID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to delete route from the store: %s", err)
+		return status.Errorf(status.Internal, "failed to delete route from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewRouteNotFoundError(routeID)
+	}
+
+	return nil
 }
 
 // GetAccountSetupKeys retrieves setup keys for an account.
@@ -1499,7 +1653,7 @@ func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStren
 // GetAccountNameServerGroups retrieves name server groups for an account.
 func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
 	var nsGroups []*nbdns.NameServerGroup
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&nsGroups, accountIDCondition, accountID)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&nsGroups, accountIDCondition, accountID)
 	if err := result.Error; err != nil {
 		log.WithContext(ctx).Errorf("failed to get name server groups from the store: %s", err)
 		return nil, status.Errorf(status.Internal, "failed to get name server groups from store")
@@ -1511,7 +1665,7 @@ func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength
 // GetNameServerGroupByID retrieves a name server group by its ID and account ID.
 func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, accountID, nsGroupID string) (*nbdns.NameServerGroup, error) {
 	var nsGroup *nbdns.NameServerGroup
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		First(&nsGroup, accountAndIDQueryCondition, accountID, nsGroupID)
 	if err := result.Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -1526,7 +1680,7 @@ func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength Lock
 
 // SaveNameServerGroup saves a name server group to the database.
 func (s *SqlStore) SaveNameServerGroup(ctx context.Context, lockStrength LockingStrength, nameServerGroup *nbdns.NameServerGroup) error {
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(nameServerGroup)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(nameServerGroup)
 	if err := result.Error; err != nil {
 		log.WithContext(ctx).Errorf("failed to save name server group to the store: %s", err)
 		return status.Errorf(status.Internal, "failed to save name server group to store")
@@ -1549,39 +1703,6 @@ func (s *SqlStore) DeleteNameServerGroup(ctx context.Context, lockStrength Locki
 	return nil
 }
 
-// getRecords retrieves records from the database based on the account ID.
-func getRecords[T any](db *gorm.DB, lockStrength LockingStrength, accountID string) ([]T, error) {
-	var record []T
-
-	result := db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&record, accountIDCondition, accountID)
-	if err := result.Error; err != nil {
-		parts := strings.Split(fmt.Sprintf("%T", record), ".")
-		recordType := parts[len(parts)-1]
-
-		return nil, status.Errorf(status.Internal, "failed to get account %ss from store: %v", recordType, err)
-	}
-
-	return record, nil
-}
-
-// getRecordByID retrieves a record by its ID and account ID from the database.
-func getRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, accountID string) (*T, error) {
-	var record T
-
-	result := db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&record, accountAndIDQueryCondition, accountID, recordID)
-	if err := result.Error; err != nil {
-		parts := strings.Split(fmt.Sprintf("%T", record), ".")
-		recordType := parts[len(parts)-1]
-
-		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "%s not found", recordType)
-		}
-		return nil, status.Errorf(status.Internal, "failed to get %s from store: %v", recordType, err)
-	}
-	return &record, nil
-}
-
 // SaveDNSSettings saves the DNS settings to the store.
 func (s *SqlStore) SaveDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *DNSSettings) error {
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
@@ -1597,3 +1718,119 @@ func (s *SqlStore) SaveDNSSettings(ctx context.Context, lockStrength LockingStre
 
 	return nil
 }
+
+// SaveAccountSettings stores the account settings in DB.
+func (s *SqlStore) SaveAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *Settings) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Model(&Account{}).
+		Select("*").Where(idQueryCondition, accountID).Updates(&AccountSettings{Settings: settings})
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save account settings to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save account settings to store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewAccountNotFoundError(accountID)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) CreateAccount(ctx context.Context, lockStrength LockingStrength, account *Account) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Create(&account)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save new account in store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save new account in store")
+	}
+	return nil
+}
+
+// GetPATByHashedToken returns a PersonalAccessToken by its hashed token.
+func (s *SqlStore) GetPATByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken string) (*PersonalAccessToken, error) {
+	var pat PersonalAccessToken
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).First(&pat, "hashed_token = ?", hashedToken)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewPATNotFoundError(hashedToken)
+		}
+		log.WithContext(ctx).Errorf("failed to get pat by hash from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get pat by hash from store")
+	}
+
+	return &pat, nil
+}
+
+// GetPATByID retrieves a personal access token by its ID and user ID.
+func (s *SqlStore) GetPATByID(ctx context.Context, lockStrength LockingStrength, userID string, patID string) (*PersonalAccessToken, error) {
+	var pat PersonalAccessToken
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		First(&pat, "id = ? AND user_id = ?", patID, userID)
+	if err := result.Error; err != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewPATNotFoundError(patID)
+		}
+		log.WithContext(ctx).Errorf("failed to get pat from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get pat from store")
+	}
+
+	return &pat, nil
+}
+
+// GetUserPATs retrieves personal access tokens for a user.
+func (s *SqlStore) GetUserPATs(ctx context.Context, lockStrength LockingStrength, userID string) ([]*PersonalAccessToken, error) {
+	var pats []*PersonalAccessToken
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&pats, "user_id = ?", userID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get user pat's from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get user pat's from store")
+	}
+
+	return pats, nil
+}
+
+// MarkPATUsed marks a personal access token as used.
+func (s *SqlStore) MarkPATUsed(ctx context.Context, lockStrength LockingStrength, patID string) error {
+	patCopy := PersonalAccessToken{
+		LastUsed: time.Now().UTC(),
+	}
+
+	fieldsToUpdate := []string{"last_used"}
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Select(fieldsToUpdate).
+		Where(idQueryCondition, patID).Updates(&patCopy)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to mark pat as used: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to mark pat as used")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewPATNotFoundError(patID)
+	}
+
+	return nil
+}
+
+// SavePAT saves a personal access token to the database.
+func (s *SqlStore) SavePAT(ctx context.Context, lockStrength LockingStrength, pat *PersonalAccessToken) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(pat)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to save pat to the store: %s", err)
+		return status.Errorf(status.Internal, "failed to save pat to store")
+	}
+
+	return nil
+}
+
+// DeletePAT deletes a personal access token from the database.
+func (s *SqlStore) DeletePAT(ctx context.Context, lockStrength LockingStrength, userID, patID string) error {
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Delete(&PersonalAccessToken{}, "user_id = ? AND id = ?", userID, patID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to delete pat from the store: %s", err)
+		return status.Errorf(status.Internal, "failed to delete pat from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewPATNotFoundError(patID)
+	}
+
+	return nil
+}

management/server/status/error.go

@@ -86,9 +86,14 @@ func NewAccountNotFoundError(accountKey string) error {
 	return Errorf(NotFound, "account not found: %s", accountKey)
 }
 
+// NewPeerNotPartOfAccountError creates a new Error with PermissionDenied type for a peer not being part of an account
+func NewPeerNotPartOfAccountError() error {
+	return Errorf(PermissionDenied, "peer is not part of this account")
+}
+
 // NewUserNotFoundError creates a new Error with NotFound type for a missing user
 func NewUserNotFoundError(userKey string) error {
-	return Errorf(NotFound, "user not found: %s", userKey)
+	return Errorf(NotFound, "user: %s not found", userKey)
 }
 
 // NewPeerNotRegisteredError creates a new Error with NotFound type for a missing peer
@@ -154,3 +159,22 @@ func NewPolicyNotFoundError(policyID string) error {
 func NewNameServerGroupNotFoundError(nsGroupID string) error {
 	return Errorf(NotFound, "nameserver group: %s not found", nsGroupID)
 }
+
+// NewServiceUserRoleInvalidError creates a new Error with InvalidArgument type for creating a service user with owner role
+func NewServiceUserRoleInvalidError() error {
+	return Errorf(InvalidArgument, "can't create a service user with owner role")
+}
+
+// NewOwnerDeletePermissionError creates a new Error with PermissionDenied type for attempting
+// to delete a user with the owner role.
+func NewOwnerDeletePermissionError() error {
+	return Errorf(PermissionDenied, "can't delete a user with the owner role")
+}
+
+func NewPATNotFoundError(patID string) error {
+	return Errorf(NotFound, "PAT: %s not found", patID)
+}
+
+func NewRouteNotFoundError(routeID string) error {
+	return Errorf(NotFound, "route: %s not found", routeID)
+}

management/server/store.go

@@ -48,32 +48,43 @@ type Store interface {
 	GetAccountByUser(ctx context.Context, userID string) (*Account, error)
 	GetAccountByPeerPubKey(ctx context.Context, peerKey string) (*Account, error)
 	GetAccountIDByPeerPubKey(ctx context.Context, peerKey string) (string, error)
-	GetAccountIDByUserID(userID string) (string, error)
+	GetAccountIDByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (string, error)
 	GetAccountIDBySetupKey(ctx context.Context, peerKey string) (string, error)
+	GetAccountIDByPeerID(ctx context.Context, lockStrength LockingStrength, peerID string) (string, error)
 	GetAccountByPeerID(ctx context.Context, peerID string) (*Account, error)
 	GetAccountBySetupKey(ctx context.Context, setupKey string) (*Account, error) // todo use key hash later
 	GetAccountByPrivateDomain(ctx context.Context, domain string) (*Account, error)
 	GetAccountIDByPrivateDomain(ctx context.Context, lockStrength LockingStrength, domain string) (string, error)
 	GetAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*Settings, error)
 	GetAccountDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*DNSSettings, error)
+	GetAccountCreatedBy(ctx context.Context, lockStrength LockingStrength, accountID string) (string, error)
+	GetTotalAccounts(ctx context.Context) (int64, error)
 	SaveAccount(ctx context.Context, account *Account) error
 	DeleteAccount(ctx context.Context, account *Account) error
-	UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error
+	UpdateAccountDomainAttributes(ctx context.Context, lockStrength LockingStrength, accountID string, domain string, category string, isPrimaryDomain *bool) error
 	SaveDNSSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *DNSSettings) error
+	SaveAccountSettings(ctx context.Context, lockStrength LockingStrength, accountID string, settings *Settings) error
+	CreateAccount(ctx context.Context, lockStrength LockingStrength, account *Account) error
 
-	GetUserByTokenID(ctx context.Context, tokenID string) (*User, error)
+	GetUserByPATID(ctx context.Context, lockStrength LockingStrength, patID string) (*User, error)
 	GetUserByUserID(ctx context.Context, lockStrength LockingStrength, userID string) (*User, error)
 	GetAccountUsers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*User, error)
-	SaveUsers(accountID string, users map[string]*User) error
+	SaveUsers(ctx context.Context, lockStrength LockingStrength, users []*User) error
 	SaveUser(ctx context.Context, lockStrength LockingStrength, user *User) error
 	SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error
+	DeleteUser(ctx context.Context, lockStrength LockingStrength, accountID, userID string) error
 	GetTokenIDByHashedToken(ctx context.Context, secret string) (string, error)
-	DeleteHashedPAT2TokenIDIndex(hashedToken string) error
-	DeleteTokenID2UserIDIndex(tokenID string) error
+
+	GetPATByID(ctx context.Context, lockStrength LockingStrength, userID, patID string) (*PersonalAccessToken, error)
+	GetUserPATs(ctx context.Context, lockStrength LockingStrength, userID string) ([]*PersonalAccessToken, error)
+	GetPATByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken string) (*PersonalAccessToken, error)
+	MarkPATUsed(ctx context.Context, lockStrength LockingStrength, patID string) error
+	SavePAT(ctx context.Context, strength LockingStrength, pat *PersonalAccessToken) error
+	DeletePAT(ctx context.Context, strength LockingStrength, userID, patID string) error
 
 	GetAccountGroups(ctx context.Context, lockStrength LockingStrength, 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)
+	GetGroupByName(ctx context.Context, lockStrength LockingStrength, accountID, groupName 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
@@ -98,12 +109,17 @@ type Store interface {
 	AddPeerToGroup(ctx context.Context, accountId string, peerId string, groupID string) error
 	AddPeerToAccount(ctx context.Context, peer *nbpeer.Peer) error
 	GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error)
+	GetAccountPeers(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error)
 	GetUserPeers(ctx context.Context, lockStrength LockingStrength, accountID, userID string) ([]*nbpeer.Peer, error)
 	GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) (*nbpeer.Peer, error)
 	GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error)
-	SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error
-	SavePeerStatus(accountID, peerID string, status nbpeer.PeerStatus) error
-	SavePeerLocation(accountID string, peer *nbpeer.Peer) error
+	GetAccountPeersWithExpiration(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error)
+	GetAccountPeersWithInactivity(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error)
+	GetAllEphemeralPeers(ctx context.Context, lockStrength LockingStrength) ([]*nbpeer.Peer, error)
+	SavePeer(ctx context.Context, lockStrength LockingStrength, accountID string, peer *nbpeer.Peer) error
+	SavePeerStatus(ctx context.Context, lockStrength LockingStrength, accountID, peerID string, status nbpeer.PeerStatus) error
+	SavePeerLocation(ctx context.Context, lockStrength LockingStrength, accountID string, peer *nbpeer.Peer) error
+	DeletePeer(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) error
 
 	GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error)
 	IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error
@@ -113,7 +129,9 @@ type Store interface {
 	DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) 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)
+	GetRouteByID(ctx context.Context, lockStrength LockingStrength, accountID, routeID string) (*route.Route, error)
+	SaveRoute(ctx context.Context, lockStrength LockingStrength, route *route.Route) error
+	DeleteRoute(ctx context.Context, lockStrength LockingStrength, accountID, routeID string) error
 
 	GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*dns.NameServerGroup, error)
 	GetNameServerGroupByID(ctx context.Context, lockStrength LockingStrength, nameServerGroupID string, accountID string) (*dns.NameServerGroup, error)

management/server/telemetry/accountmanager_metrics.go

@@ -13,7 +13,6 @@ type AccountManagerMetrics struct {
 	updateAccountPeersDurationMs metric.Float64Histogram
 	getPeerNetworkMapDurationMs  metric.Float64Histogram
 	networkMapObjectCount        metric.Int64Histogram
-	peerMetaUpdateCount          metric.Int64Counter
 }
 
 // NewAccountManagerMetrics creates an instance of AccountManagerMetrics
@@ -45,17 +44,11 @@ func NewAccountManagerMetrics(ctx context.Context, meter metric.Meter) (*Account
 		return nil, err
 	}
 
-	peerMetaUpdateCount, err := meter.Int64Counter("management.account.peer.meta.update.counter", metric.WithUnit("1"))
-	if err != nil {
-		return nil, err
-	}
-
 	return &AccountManagerMetrics{
 		ctx:                          ctx,
 		getPeerNetworkMapDurationMs:  getPeerNetworkMapDurationMs,
 		updateAccountPeersDurationMs: updateAccountPeersDurationMs,
 		networkMapObjectCount:        networkMapObjectCount,
-		peerMetaUpdateCount:          peerMetaUpdateCount,
 	}, nil
 
 }
@@ -74,8 +67,3 @@ func (metrics *AccountManagerMetrics) CountGetPeerNetworkMapDuration(duration ti
 func (metrics *AccountManagerMetrics) CountNetworkMapObjects(count int64) {
 	metrics.networkMapObjectCount.Record(metrics.ctx, count)
 }
-
-// CountPeerMetUpdate counts the number of peer meta updates
-func (metrics *AccountManagerMetrics) CountPeerMetUpdate() {
-	metrics.peerMetaUpdateCount.Add(metrics.ctx, 1)
-}

management/server/telemetry/store_metrics.go

@@ -13,7 +13,6 @@ type StoreMetrics struct {
 	globalLockAcquisitionDurationMs    metric.Int64Histogram
 	persistenceDurationMicro           metric.Int64Histogram
 	persistenceDurationMs              metric.Int64Histogram
-	transactionDurationMs              metric.Int64Histogram
 	ctx                                context.Context
 }
 
@@ -41,17 +40,11 @@ func NewStoreMetrics(ctx context.Context, meter metric.Meter) (*StoreMetrics, er
 		return nil, err
 	}
 
-	transactionDurationMs, err := meter.Int64Histogram("management.store.transaction.duration.ms")
-	if err != nil {
-		return nil, err
-	}
-
 	return &StoreMetrics{
 		globalLockAcquisitionDurationMicro: globalLockAcquisitionDurationMicro,
 		globalLockAcquisitionDurationMs:    globalLockAcquisitionDurationMs,
 		persistenceDurationMicro:           persistenceDurationMicro,
 		persistenceDurationMs:              persistenceDurationMs,
-		transactionDurationMs:              transactionDurationMs,
 		ctx:                                ctx,
 	}, nil
 }
@@ -67,8 +60,3 @@ func (metrics *StoreMetrics) CountPersistenceDuration(duration time.Duration) {
 	metrics.persistenceDurationMicro.Record(metrics.ctx, duration.Microseconds())
 	metrics.persistenceDurationMs.Record(metrics.ctx, duration.Milliseconds())
 }
-
-// CountTransactionDuration counts the duration of a store persistence operation
-func (metrics *StoreMetrics) CountTransactionDuration(duration time.Duration) {
-	metrics.transactionDurationMs.Record(metrics.ctx, duration.Milliseconds())
-}

management/server/testdata/extended-store.sql

@@ -37,4 +37,5 @@ INSERT INTO "groups" VALUES('cfefqs706sqkneg59g2g','bf1c8084-ba50-4ce7-9439-3465
 INSERT INTO posture_checks VALUES('csplshq7qv948l48f7t0','NetBird Version > 0.32.0','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"NBVersionCheck":{"MinVersion":"0.31.0"}}');
 INSERT INTO posture_checks VALUES('cspnllq7qv95uq1r4k90','Allow Berlin and Deny local network 172.16.1.0/24','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"GeoLocationCheck":{"Locations":[{"CountryCode":"DE","CityName":"Berlin"}],"Action":"allow"},"PeerNetworkRangeCheck":{"Action":"deny","Ranges":["172.16.1.0/24"]}}');
 INSERT INTO name_server_groups VALUES('csqdelq7qv97ncu7d9t0','bf1c8084-ba50-4ce7-9439-34653001fc3b','Google DNS','Google DNS Servers','[{"IP":"8.8.8.8","NSType":1,"Port":53},{"IP":"8.8.4.4","NSType":1,"Port":53}]','["cfefqs706sqkneg59g2g"]',1,'[]',1,0);
+INSERT INTO routes VALUES('ct03t427qv97vmtmglog','bf1c8084-ba50-4ce7-9439-34653001fc3b','"10.10.0.0/16"',NULL,0,'aws-eu-central-1-vpc','Production VPC in Frankfurt','ct03r5q7qv97vmtmglng',NULL,1,1,9999,1,'["cfefqs706sqkneg59g2g"]',NULL);
 INSERT INTO installations VALUES(1,'');

management/server/testdata/store.sql

@@ -25,7 +25,7 @@ CREATE INDEX `idx_routes_account_id` ON `routes`(`account_id`);
 CREATE INDEX `idx_name_server_groups_account_id` ON `name_server_groups`(`account_id`);
 CREATE INDEX `idx_posture_checks_account_id` ON `posture_checks`(`account_id`);
 
-INSERT INTO accounts VALUES('bf1c8084-ba50-4ce7-9439-34653001fc3b','','2024-10-02 16:03:06.778746+02:00','test.com','private',1,'af1c8024-ha40-4ce2-9418-34653101fc3c','{"IP":"100.64.0.0","Mask":"//8AAA=="}','',0,'[]',0,86400000000000,0,0,0,'',NULL,NULL,NULL);
+INSERT INTO accounts VALUES('bf1c8084-ba50-4ce7-9439-34653001fc3b','edafee4e-63fb-11ec-90d6-0242ac120003','2024-10-02 16:03:06.778746+02:00','test.com','private',1,'af1c8024-ha40-4ce2-9418-34653101fc3c','{"IP":"100.64.0.0","Mask":"//8AAA=="}','',0,'[]',0,86400000000000,0,0,0,'',NULL,NULL,NULL);
 INSERT INTO "groups" VALUES('cs1tnh0hhcjnqoiuebeg','bf1c8084-ba50-4ce7-9439-34653001fc3b','All','api','[]',0,'');
 INSERT INTO setup_keys VALUES('','bf1c8084-ba50-4ce7-9439-34653001fc3b','A2C8E62B-38F5-4553-B31E-DD66C696CEBB','Default key','reusable','2021-08-19 20:46:20.005936822+02:00','2321-09-18 20:46:20.005936822+02:00','2021-08-19 20:46:20.005936822+02:00',0,0,'0001-01-01 00:00:00+00:00','["cs1tnh0hhcjnqoiuebeg"]',0,0);
 INSERT INTO users VALUES('edafee4e-63fb-11ec-90d6-0242ac120003','bf1c8084-ba50-4ce7-9439-34653001fc3b','admin',0,0,'','[]',0,'0001-01-01 00:00:00+00:00','2024-10-02 16:03:06.779156+02:00','api',0,'');