Merge remote-tracking branch 'origin/main' into fix_debug_upload_url_from_mgmt

# Conflicts:
#	management/server/activity/codes.go
#	management/server/store/sql_store.go
#	management/server/store/sql_store_test.go
#	upload-server/server/server.go

management/cmd/proxy/proxy.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"strings"
 	"text/tabwriter"
+	"unicode"
 
 	"github.com/spf13/cobra"
 
@@ -68,8 +69,8 @@ func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.R
 
 	toDisconnect := 0
 	w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
-	_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tACCOUNT\tSTATUS\tLAST SEEN")
-	_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t------\t---------")
+	_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tVERSION\tACCOUNT\tSTATUS\tLAST SEEN")
+	_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t-------\t------\t---------")
 
 	for _, p := range proxies {
 		if p.Status != rpproxy.StatusDisconnected {
@@ -80,11 +81,16 @@ func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.R
 		if p.AccountID != nil {
 			account = *p.AccountID
 		}
+		version := "-"
+		if p.Version != "" {
+			version = sanitizeReportedValue(p.Version)
+		}
 
-		_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
-			p.ID,
+		_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
+			sanitizeReportedValue(p.ID),
 			p.ClusterAddress,
 			p.IPAddress,
+			version,
 			account,
 			p.Status,
 			p.LastSeen.Format("2006-01-02 15:04:05"),
@@ -139,3 +145,16 @@ func confirmDisconnectAll(out io.Writer, in io.Reader) (bool, error) {
 
 	return strings.EqualFold(strings.TrimSpace(scanner.Text()), disconnectAllConfirmation), nil
 }
+
+// sanitizeReportedValue replaces non-printable characters in a value the proxy
+// reports about itself. Both the id and the version arrive unvalidated over
+// gRPC, so a tab would forge a column, a carriage return or ANSI escape would
+// redraw the operator's terminal, and U+202E would reverse the rest of the line.
+func sanitizeReportedValue(s string) string {
+	return strings.Map(func(r rune) rune {
+		if unicode.IsPrint(r) {
+			return r
+		}
+		return '\uFFFD'
+	}, s)
+}

management/cmd/proxy/proxy_test.go

@@ -35,6 +35,7 @@ func seedProxies(t *testing.T, ctx context.Context, s store.Store) {
 			SessionID:      "session-1",
 			ClusterAddress: "cluster-a.example.com",
 			IPAddress:      "10.0.0.1",
+			Version:        "0.60.0",
 			LastSeen:       time.Now(),
 			Status:         rpproxy.StatusConnected,
 		},
@@ -89,6 +90,7 @@ func TestRunDisconnectAllWithConfirmation(t *testing.T) {
 	require.Contains(t, output, "proxy-2")
 	require.Contains(t, output, "proxy-3")
 	require.Contains(t, output, "cluster-a.example.com")
+	require.Contains(t, output, "0.60.0")
 	require.Contains(t, output, "account-1")
 	require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
 	require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
@@ -178,3 +180,40 @@ func TestRunDisconnectAllEmpty(t *testing.T) {
 	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, false))
 	require.Contains(t, out.String(), "No reverse proxy instances found.")
 }
+
+func TestRunDisconnectAllEscapesProxyReportedFields(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+
+	// A proxy reports its own id and version on connect, so both reach this
+	// listing unvalidated. Carriage returns, tabs and ANSI escapes would let
+	// a malicious proxy redraw the table or forge a row on the operator's
+	// terminal; U+202E would reverse the rendering of the rest of the line.
+	require.NoError(t, s.SaveProxy(ctx, &rpproxy.Proxy{
+		ID:             "proxy-\r\x1b[2Kevil",
+		SessionID:      "session-1",
+		ClusterAddress: "cluster-a.example.com",
+		IPAddress:      "10.0.0.1",
+		Version:        "0.60.0\tfake\rcolumn\u202e",
+		LastSeen:       time.Now(),
+		Status:         rpproxy.StatusConnected,
+	}))
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(disconnectAllConfirmation+"\n"), true, false))
+
+	output := out.String()
+	for _, forbidden := range []string{"\r", "\x1b", "\u202e"} {
+		require.NotContains(t, output, forbidden, "listing must not carry proxy-reported control characters")
+	}
+	// The table has one data row; a smuggled tab would add a phantom column.
+	var dataRow string
+	for _, line := range strings.Split(output, "\n") {
+		if strings.Contains(line, "evil") {
+			dataRow = line
+		}
+	}
+	require.NotEmpty(t, dataRow, "listing should still show the proxy row")
+	require.NotContains(t, dataRow, "\t", "tabwriter output should not carry a smuggled column separator")
+	require.Contains(t, dataRow, "0.60.0", "the printable part of the version should survive")
+}

management/internals/controllers/network_map/nmaptest/runner.go

@@ -245,7 +245,7 @@ func computeMode(t *testing.T, ctx context.Context, mode Mode, nmData *networkma
 		peerGroups := maps.Keys(nmData.GetPeerGroups(peerID))
 		resp := mgmtgrpc.ToComponentSyncResponse(ctx, nil, nil, nil, peer, nil, nil, components, nil,
 			dnsDomain, nil, nmData.AccountSettings, nil, peerGroups, dnsFwdPort)
-		res, err := networkmap.EnvelopeToNetworkMap(ctx, resp.NetworkMapEnvelope, peer.Key, dnsDomain)
+		res, err := networkmap.EnvelopeToNetworkMap(ctx, resp.NetworkMapEnvelope, peer.Key, dnsDomain, false)
 		require.NoError(t, err, "expand envelope")
 		return res.NetworkMap
 	default:

management/internals/modules/agentnetwork/handlers/handlers_test.go

@@ -9,6 +9,7 @@ import (
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	"go.uber.org/mock/gomock"
 	"github.com/gorilla/mux"
@@ -17,6 +18,7 @@ import (
 
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
 	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
 	"github.com/netbirdio/netbird/management/server/account"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/permissions"
@@ -29,6 +31,9 @@ import (
 const (
 	testAccountID = "acc-1"
 	testUserID    = "user-bob"
+	// testClusterAddress is the shared proxy cluster the settings tests pin
+	// their gateway to; the fixture seeds a connected private-capable proxy for it.
+	testClusterAddress = "eu.proxy.netbird.io"
 )
 
 // agentNetworkHandlerFixture builds a real agentnetwork.Manager with
@@ -75,6 +80,12 @@ func newAgentNetworkHandlerFixture(t *testing.T) *agentNetworkHandlerFixture {
 	manager := agentnetwork.NewManager(st, perms, accounts, nil)
 	h := &handler{manager: manager}
 
+	// The labeled bootstrap validates its proxy_address against the live
+	// clusters, so seed the shared cluster these tests pin to as a real,
+	// private-capable one — the wire-shape assertions then run through the
+	// validated path rather than the "nothing connected yet" carve-out.
+	seedSharedPrivateCluster(t, st, testClusterAddress)
+
 	router := mux.NewRouter()
 	router.HandleFunc("/agent-network/providers", h.createProvider).Methods("POST")
 	router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET")
@@ -268,3 +279,21 @@ func TestConsumptionHandler_PopulatedAccountListsRows(t *testing.T) {
 	assert.Equal(t, groupRow.WindowStartUtc, userRow.WindowStartUtc,
 		"rows recorded in the same window must share the aligned window_start_utc")
 }
+
+// seedSharedPrivateCluster registers a connected, NetBird-operated proxy
+// with private capabilities (the `private` capability) so
+// clusterAddr is a cluster any account may pin its agent-network gateway to.
+func seedSharedPrivateCluster(t *testing.T, st store.Store, clusterAddr string) {
+	t.Helper()
+	private := true
+	now := time.Now().UTC()
+	require.NoError(t, st.SaveProxy(context.Background(), &rpproxy.Proxy{
+		ID:             "shared-proxy-" + clusterAddr,
+		SessionID:      "shared-session",
+		ClusterAddress: clusterAddr,
+		LastSeen:       now,
+		ConnectedAt:    &now,
+		Status:         rpproxy.StatusConnected,
+		Capabilities:   rpproxy.Capabilities{Private: &private},
+	}), "seeding the shared proxy cluster must succeed")
+}

management/internals/modules/agentnetwork/manager.go

@@ -1036,6 +1036,18 @@ func (m *managerImpl) bootstrapSelfAddressed(ctx context.Context, settings *type
 	if err != nil {
 		return status.Errorf(status.InvalidArgument, "invalid endpoint: %s", err)
 	}
+	if err := m.requireHostNotForeign(ctx, settings.AccountID, hostname); err != nil {
+		return err
+	}
+	// Another account's labeled pin beneath this hostname makes it their
+	// cluster: a proxy serving them there would never serve this endpoint.
+	// The domain unique index already arbitrates two endpoints on one name.
+	if err := m.requireNotClaimedByOtherAccount(ctx, settings.AccountID, hostname, m.store.HasGatewayClusterPinnedByOtherAccount); err != nil {
+		return err
+	}
+	if err := m.validateGatewayCluster(ctx, settings.AccountID, hostname); err != nil {
+		return err
+	}
 
 	settings.Domain = hostname
 	settings.ProxyAddress = hostname
@@ -1054,6 +1066,99 @@ func (m *managerImpl) bootstrapSelfAddressed(ctx context.Context, settings *type
 	return nil
 }
 
+// validateGatewayCluster rejects a bootstrap pinned to a cluster that cannot
+// serve the account's gateway — a labeled endpoint beneath the cluster and a
+// self-addressed one on the very address a proxy declares alike, since the
+// service behind either is the same private one.
+//
+// The synthesised gateway service is unconditionally private
+// (buildAccountService): agents reach it over the WireGuard tunnel and are
+// authorised by ValidateTunnelPeer against the policies' source groups, and
+// its single target is the cluster itself with DirectUpstream. Only a cluster
+// with private capabilities can serve that. Management reports it per cluster
+// as the `private` capability, the same flag the dashboard renders as
+// supports_private when it gates NetBird-only services.
+//
+// Without this check the bootstrap happily pins to any cluster the caller
+// names, including one without private capabilities — and the endpoint it
+// allocates is immutable, so the account is left with a dead gateway that only
+// a DeleteSettings/re-bootstrap can undo.
+//
+// Whether management knows the cluster is decided on the proxy rows
+// themselves, never on how fresh their heartbeats are: a cluster's rows
+// outlive its proxies' liveness (only the stale-proxy reaper removes them), so
+// a cluster that exists stays judged as one. Judging on liveness instead would
+// make the same centralised cluster pass or fail depending on whether its
+// proxies happened to have heartbeated in the last couple of minutes.
+//
+// The single opening left is a cluster management holds no proxy row for at
+// all: pinning ahead of a proxy's first connection is a legitimate order — the
+// dedicated path claims an address the same way, before any proxy declares it.
+func (m *managerImpl) validateGatewayCluster(ctx context.Context, accountID, clusterAddr string) error {
+	declared, err := m.accountClusterSpellings(ctx, accountID, clusterAddr)
+	if err != nil {
+		return err
+	}
+	if len(declared) == 0 {
+		// No proxy has ever declared this address: an address-first pin.
+		return nil
+	}
+
+	// A cluster management knows has to prove it can serve the gateway, and
+	// only a live proxy reporting the capability proves that. Both an explicit false and an
+	// unreported capability (nothing live in the cluster, or proxies predating
+	// capability reporting) fail here: unusable and unproven are the same
+	// answer for a decision that cannot be revisited later.
+	//
+	// The capability is read per declared spelling and taken as any-true, the
+	// same way it aggregates over a cluster's proxies: the store matches
+	// cluster_address exactly, so a host two proxies spelled differently must
+	// not come back unproven just because it was asked about under one of them.
+	for _, address := range declared {
+		if private := m.store.GetClusterSupportsPrivate(ctx, address); private != nil && *private {
+			return nil
+		}
+	}
+
+	return status.Errorf(status.InvalidArgument,
+		"proxy cluster %s has no private capabilities: the agent network gateway requires a reverse proxy cluster "+
+			"with private capabilities", clusterAddr)
+}
+
+// accountClusterSpellings returns every proxy cluster address in the account's
+// view — its own (BYOP) clusters plus the shared ones — that names the same
+// host as clusterAddr. Empty means management holds no proxy row for that host
+// in this account's view.
+//
+// A proxy declares its cluster address as the operator spelled it, so identity
+// is compared on the normalised form rather than byte-equal — an in-memory pass
+// over the account's clusters, not a query. What comes back is the stored
+// spelling, because the capability lookup matches cluster_address exactly and
+// would silently find nothing under a spelling the store never held. The
+// cluster listing is not gated on heartbeats, so this answer does not change
+// while a cluster's proxies are merely offline.
+func (m *managerImpl) accountClusterSpellings(ctx context.Context, accountID, clusterAddr string) ([]string, error) {
+	clusters, err := m.store.GetProxyClusters(ctx, accountID)
+	if err != nil {
+		return nil, fmt.Errorf("list proxy clusters: %w", err)
+	}
+
+	var spellings []string
+	for _, cluster := range clusters {
+		normalized, err := types.NormalizeHostname(cluster.Address)
+		if err != nil {
+			// An address declared in a shape we cannot normalise is not one an
+			// endpoint can be allocated beneath.
+			log.WithContext(ctx).Debugf("skipping unusable proxy cluster address %q: %s", cluster.Address, err)
+			continue
+		}
+		if normalized == clusterAddr {
+			spellings = append(spellings, cluster.Address)
+		}
+	}
+	return spellings, nil
+}
+
 // bootstrapLabeled allocates a labeled endpoint one label beneath the given
 // cluster address: Domain = <label>.<proxyAddress>, served by whichever proxy
 // declares the parent. Labels are adjective-noun tuples; a candidate is
@@ -1065,6 +1170,20 @@ func (m *managerImpl) bootstrapLabeled(ctx context.Context, settings *types.Sett
 	if err != nil {
 		return status.Errorf(status.InvalidArgument, "invalid proxy_address: %s", err)
 	}
+	if err := m.requireHostNotForeign(ctx, settings.AccountID, parent); err != nil {
+		return err
+	}
+	// Another account's endpoint at this exact hostname means the proxy that
+	// declares it is theirs, so nothing would serve a label beneath it. Other
+	// accounts' labeled pins under the same cluster are not asked about: a
+	// shared cluster carries many of them by design.
+	if err := m.requireNotClaimedByOtherAccount(ctx, settings.AccountID, parent, m.store.HasGatewayEndpointByOtherAccount); err != nil {
+		return err
+	}
+
+	if err := m.validateGatewayCluster(ctx, settings.AccountID, parent); err != nil {
+		return err
+	}
 
 	for attempt := 1; attempt <= maxDomainAllocationAttempts; attempt++ {
 		label := labelgen.PickTuple()
@@ -1111,6 +1230,41 @@ func (m *managerImpl) bootstrapLabeled(ctx context.Context, settings *types.Sett
 	return fmt.Errorf("allocate agent network endpoint for account %s: %d attempts exhausted", settings.AccountID, maxDomainAllocationAttempts)
 }
 
+// requireHostNotForeign refuses to pin the account's gateway onto a host that
+// another account's proxy declares. The pin's proxy_address is what selects
+// the proxy that serves the endpoint, and an account-scoped proxy only ever
+// receives its own account's mappings, so such a pin could never be served —
+// and the endpoint it assigns is immutable. Shared proxies are not foreign, and
+// a host no proxy has declared stays pinnable: claiming the address before the
+// proxy's first connection is the documented order.
+func (m *managerImpl) requireHostNotForeign(ctx context.Context, accountID, host string) error {
+	foreign, err := m.store.HasForeignAccountProxyAtHost(ctx, host, accountID)
+	if err != nil {
+		return fmt.Errorf("check proxy host ownership: %w", err)
+	}
+	if foreign {
+		return errHostNotAvailable(host)
+	}
+	return nil
+}
+
+// requireNotClaimedByOtherAccount refuses the pin when another account's
+// gateway settings already claim the host in the shape claimed answers for.
+func (m *managerImpl) requireNotClaimedByOtherAccount(ctx context.Context, accountID, host string, claimed func(context.Context, string, string) (bool, error)) error {
+	taken, err := claimed(ctx, host, accountID)
+	if err != nil {
+		return fmt.Errorf("check agent network gateway claims at host: %w", err)
+	}
+	if taken {
+		return errHostNotAvailable(host)
+	}
+	return nil
+}
+
+func errHostNotAvailable(host string) error {
+	return status.Errorf(status.InvalidArgument, "proxy cluster %s is not available to this account", host)
+}
+
 // isUniqueConstraintError reports whether err is a database unique-constraint
 // violation, matched on the driver message because CreateAgentNetworkSettings
 // deliberately returns the driver error unwrapped.

management/internals/modules/agentnetwork/settings_bootstrap_test.go

@@ -5,12 +5,14 @@ import (
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
-	"go.uber.org/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
@@ -70,6 +72,57 @@ func (f *bootstrapFixture) createSettings(ctx context.Context, accountID, userID
 	return f.manager.CreateSettings(ctx, userID, types.DefaultSettings(accountID), proxyAddress, endpoint)
 }
 
+func ptrTo[T any](v T) *T { return &v }
+
+// seedProxy registers a proxy in clusterAddr, heartbeating now, so the labeled
+// bootstrap path has a real cluster to validate against. accountID empty makes
+// it a shared (NetBird-operated) cluster; private mirrors the capability an
+// proxy with private capabilities reports, nil an unreported one.
+func (f *bootstrapFixture) seedProxy(t *testing.T, proxyID, accountID, clusterAddr string, private *bool) {
+	t.Helper()
+	f.seedProxyAt(t, proxyID, accountID, clusterAddr, private, time.Now().UTC())
+}
+
+// seedProxyAt is seedProxy with an explicit last-seen, for cases that need a
+// proxy whose heartbeat has aged past the active window while its row (and so
+// its cluster) is still on record.
+func (f *bootstrapFixture) seedProxyAt(t *testing.T, proxyID, accountID, clusterAddr string, private *bool, lastSeen time.Time) {
+	t.Helper()
+	p := &proxy.Proxy{
+		ID:             proxyID,
+		ClusterAddress: clusterAddr,
+		Status:         proxy.StatusConnected,
+		LastSeen:       lastSeen,
+		Capabilities:   proxy.Capabilities{Private: private},
+	}
+	if accountID != "" {
+		p.AccountID = &accountID
+	}
+	require.NoError(t, f.store.SaveProxy(context.Background(), p), "seeding a proxy must succeed")
+}
+
+// seedPrivateCluster is the common case: a shared cluster with a connected
+// proxy that has private capabilities, which is what a bootstrap requires.
+func (f *bootstrapFixture) seedPrivateCluster(t *testing.T, clusterAddr string) {
+	t.Helper()
+	f.seedProxy(t, "proxy-"+clusterAddr, "", clusterAddr, ptrTo(true))
+}
+
+// requireForeignClusterRefusal asserts the refusal a pin onto another
+// account's host gets, and that it left no row behind.
+func (f *bootstrapFixture) requireForeignClusterRefusal(t *testing.T, err error, accountID string) {
+	t.Helper()
+	require.Error(t, err, "another account's host must be refused")
+	var sErr *status.Error
+	require.ErrorAs(t, err, &sErr)
+	assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
+	assert.Contains(t, err.Error(), "not available to this account",
+		"the error must say the host is not the account's to use")
+
+	_, err = f.store.GetAgentNetworkSettings(context.Background(), store.LockingStrengthNone, accountID)
+	assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
+}
+
 // TestCreateSettingsRequiresPermission pins the gate: bootstrap assigns the
 // account's immutable endpoint, a settings write requiring the settings
 // Create permission — and a denial leaves no row behind.
@@ -94,6 +147,7 @@ func TestCreateSettingsRequiresPermission(t *testing.T) {
 func TestCreateSettingsLabeled(t *testing.T) {
 	ctx := context.Background()
 	f := newBootstrapFixture(t)
+	f.seedPrivateCluster(t, "cluster1.example.com")
 	f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
 
 	created, err := f.createSettings(ctx, "account1", "user1", "Cluster1.Example.com", "")
@@ -167,6 +221,7 @@ func TestCreateSettingsIdentityFieldValidation(t *testing.T) {
 func TestCreateSettingsConflictsOnSecondBootstrap(t *testing.T) {
 	ctx := context.Background()
 	f := newBootstrapFixture(t)
+	f.seedPrivateCluster(t, "cluster1.example.com")
 	f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
 
 	first, err := f.createSettings(ctx, "account1", "user1", "cluster1.example.com", "")
@@ -230,3 +285,279 @@ func TestCreateProviderHasNoSettingsSideEffects(t *testing.T) {
 	_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
 	assert.Error(t, err, "provider create must not conjure a settings row")
 }
+
+// TestCreateSettingsRejectsOfflineCluster is the guard against deciding on
+// heartbeat freshness. A centralised cluster is refused while its proxies are
+// live; the same cluster must stay refused once they stop heartbeating, which
+// takes only a couple of minutes (proxyActiveThreshold). Judging on liveness
+// would turn "wait for the proxy to go quiet" into a way to pin the account's
+// immutable endpoint to a cluster that can never serve it.
+func TestCreateSettingsRejectsOfflineCluster(t *testing.T) {
+	ctx := context.Background()
+	notPrivate := false
+
+	cases := map[string]*bool{
+		"centralised proxy gone quiet": &notPrivate,
+		// A cluster that could serve the gateway still has to have something
+		// live in it to prove so at bootstrap: refusing is the safe direction
+		// (reconnect the proxy and retry) where accepting is permanent.
+		"private proxy gone quiet": ptrTo(true),
+	}
+	for name, private := range cases {
+		t.Run(name, func(t *testing.T) {
+			f := newBootstrapFixture(t)
+			f.seedProxyAt(t, "proxy1", "", "offline.example.com", private,
+				time.Now().UTC().Add(-time.Hour))
+			f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+			_, err := f.createSettings(ctx, "account1", "user1", "offline.example.com", "")
+			require.Error(t, err, "a known cluster with nothing live in it must be rejected")
+			var sErr *status.Error
+			require.ErrorAs(t, err, &sErr)
+			assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
+			assert.Contains(t, err.Error(), "private capabilities",
+				"the error must say private capabilities are what is missing")
+
+			_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
+			assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
+		})
+	}
+}
+
+// TestCreateSettingsRequiresPrivateCluster pins the capability gate: the
+// synthesised gateway service is always private, so a live cluster whose
+// proxies lack private capabilities cannot serve it and must not
+// become the account's immutable endpoint.
+func TestCreateSettingsRequiresPrivateCluster(t *testing.T) {
+	ctx := context.Background()
+	f := newBootstrapFixture(t)
+	notPrivate := false
+	f.seedProxy(t, "proxy1", "", "central.example.com", &notPrivate)
+	f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+	_, err := f.createSettings(ctx, "account1", "user1", "central.example.com", "")
+	require.Error(t, err, "a cluster without private capabilities must be rejected")
+	var sErr *status.Error
+	require.ErrorAs(t, err, &sErr)
+	assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
+	assert.Contains(t, err.Error(), "private capabilities", "the error must name what the cluster is missing")
+
+	_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
+	assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
+}
+
+// TestCreateSettingsAcceptsOwnPrivateCluster pins the BYOP happy path: the
+// account's own cluster with a connected private-capable proxy is a valid pin.
+func TestCreateSettingsAcceptsOwnPrivateCluster(t *testing.T) {
+	ctx := context.Background()
+	f := newBootstrapFixture(t)
+	f.seedProxy(t, "proxy1", "account1", "byop.account1.example.com", ptrTo(true))
+	f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+	created, err := f.createSettings(ctx, "account1", "user1", "byop.account1.example.com", "")
+	require.NoError(t, err, "the account's own private cluster must be accepted")
+	assert.Equal(t, "byop.account1.example.com", created.ProxyAddress)
+}
+
+// TestCreateSettingsMatchesClusterCasing pins that a cluster spelled with
+// capitals in the store is still recognised as the same cluster the normalised
+// proxy_address names, in both directions: a private cluster is accepted and a
+// centralised one is refused, whatever the casing. The comparison is in memory
+// over the account's cluster list; the capability lookup is still asked under
+// the spelling the store actually holds, which is what an exact match needs.
+func TestCreateSettingsMatchesClusterCasing(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("own private cluster is found", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "proxy1", "", "EU.Proxy.Example.com", ptrTo(true))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		created, err := f.createSettings(ctx, "account1", "user1", "eu.proxy.example.com", "")
+		require.NoError(t, err, "a private cluster declared with capitals must still be accepted")
+		assert.Equal(t, "eu.proxy.example.com", created.ProxyAddress)
+	})
+
+	t.Run("non-private cluster is still refused", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "proxy1", "", "Central.Example.com", ptrTo(false))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		_, err := f.createSettings(ctx, "account1", "user1", "central.example.com", "")
+		require.Error(t, err, "casing must not become a way past the capability check")
+		assert.Contains(t, err.Error(), "private capabilities")
+	})
+}
+
+// TestCreateSettingsRejectsForeignCluster pins tenant consistency on the pin:
+// an account may not pin its gateway onto a host another account's proxy
+// declares. That proxy only ever receives its own account's mappings, so the
+// pin could never be served, and the endpoint it assigns is immutable.
+// Ownership is decided on the proxy rows, not on heartbeat freshness — a
+// cluster whose proxies are merely offline is still somebody's — and on the
+// normalised host, since proxies declare their address as the operator
+// spelled it.
+func TestCreateSettingsRejectsForeignCluster(t *testing.T) {
+	ctx := context.Background()
+
+	cases := map[string]struct {
+		spelling string
+		lastSeen time.Time
+	}{
+		"live":            {"byop.account2.example.com", time.Now().UTC()},
+		"offline":         {"byop.account2.example.com", time.Now().UTC().Add(-time.Hour)},
+		"spelled in caps": {"BYOP.Account2.Example.com", time.Now().UTC()},
+	}
+	for name, tc := range cases {
+		t.Run("labeled "+name, func(t *testing.T) {
+			f := newBootstrapFixture(t)
+			f.seedProxyAt(t, "proxy1", "account2", tc.spelling, ptrTo(true), tc.lastSeen)
+			f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+			_, err := f.createSettings(ctx, "account1", "user1", "byop.account2.example.com", "")
+			f.requireForeignClusterRefusal(t, err, "account1")
+		})
+		t.Run("self-addressed "+name, func(t *testing.T) {
+			f := newBootstrapFixture(t)
+			f.seedProxyAt(t, "proxy1", "account2", tc.spelling, ptrTo(true), tc.lastSeen)
+			f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+			_, err := f.createSettings(ctx, "account1", "user1", "", "byop.account2.example.com")
+			f.requireForeignClusterRefusal(t, err, "account1")
+		})
+	}
+}
+
+// TestCreateSettingsSharedClusterStaysPinnable pins the constraint the
+// ownership check must respect: a shared (NetBird-operated) cluster is not
+// anybody's, so any number of accounts pin their gateways to it — including
+// an account that also runs a proxy of its own elsewhere.
+func TestCreateSettingsSharedClusterStaysPinnable(t *testing.T) {
+	ctx := context.Background()
+	f := newBootstrapFixture(t)
+	f.seedProxy(t, "shared", "", "eu.proxy.netbird.io", ptrTo(true))
+	f.seedProxy(t, "own", "account1", "byop.account1.example.com", ptrTo(true))
+
+	for _, account := range []string{"account1", "account2"} {
+		f.expectPermission(account, "user", modules.AgentNetworkSettings, operations.Create, true)
+		created, err := f.createSettings(ctx, account, "user", "eu.proxy.netbird.io", "")
+		require.NoError(t, err, "a shared cluster must stay pinnable by %s", account)
+		assert.Equal(t, "eu.proxy.netbird.io", created.ProxyAddress)
+	}
+}
+
+// TestCreateSettingsOwnClusterIsPinnable is the BYOP order in both directions:
+// the account's own proxy is not a competing claim, whether the pin is labeled
+// beneath its cluster or self-addressed onto the very host it declares.
+func TestCreateSettingsOwnClusterIsPinnable(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("labeled", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "own", "account1", "byop.account1.example.com", ptrTo(true))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		created, err := f.createSettings(ctx, "account1", "user1", "byop.account1.example.com", "")
+		require.NoError(t, err, "the account's own cluster must be pinnable")
+		assert.True(t, strings.HasSuffix(created.Domain, ".byop.account1.example.com"))
+	})
+	t.Run("self-addressed", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "own", "account1", "gw.account1.example.com", ptrTo(true))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		created, err := f.createSettings(ctx, "account1", "user1", "", "gw.account1.example.com")
+		require.NoError(t, err, "the host the account's own proxy declares must be pinnable")
+		assert.Equal(t, "gw.account1.example.com", created.ProxyAddress)
+	})
+}
+
+// TestCreateSettingsUnknownHostIsPinnable pins the address-first order: a host
+// no proxy has ever declared is nobody's, so the pin goes through and the
+// proxy is deployed after.
+func TestCreateSettingsUnknownHostIsPinnable(t *testing.T) {
+	ctx := context.Background()
+	f := newBootstrapFixture(t)
+	f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+	created, err := f.createSettings(ctx, "account1", "user1", "future.example.com", "")
+	require.NoError(t, err, "a host no proxy has declared must stay pinnable")
+	assert.Equal(t, "future.example.com", created.ProxyAddress)
+}
+
+// TestCreateSettingsRejectsHostAnotherAccountPinned covers claims made by pins
+// rather than proxies, which the proxy-row check cannot see. A labeled pin
+// beneath a host makes that host the other account's cluster, so a
+// self-addressed endpoint on it would never be served; a self-addressed
+// endpoint on a host makes the proxy declaring it theirs, so a label beneath
+// it would never be served either. Neither is a shared-cluster shape: many
+// labeled pins under one cluster are asked about in neither direction.
+func TestCreateSettingsRejectsHostAnotherAccountPinned(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("self-addressed onto another account's cluster", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.expectPermission("account2", "user2", modules.AgentNetworkSettings, operations.Create, true)
+		_, err := f.createSettings(ctx, "account2", "user2", "gw.example.com", "")
+		require.NoError(t, err, "account2's labeled pin beneath the host must go through first")
+
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+		_, err = f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
+		f.requireForeignClusterRefusal(t, err, "account1")
+	})
+
+	t.Run("labeled beneath another account's endpoint", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.expectPermission("account2", "user2", modules.AgentNetworkSettings, operations.Create, true)
+		_, err := f.createSettings(ctx, "account2", "user2", "", "gw.example.com")
+		require.NoError(t, err, "account2's self-addressed endpoint must go through first")
+
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+		_, err = f.createSettings(ctx, "account1", "user1", "gw.example.com", "")
+		f.requireForeignClusterRefusal(t, err, "account1")
+	})
+
+	t.Run("labeled beside another account's labeled pin stays allowed", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		for _, account := range []string{"account1", "account2"} {
+			f.expectPermission(account, "user", modules.AgentNetworkSettings, operations.Create, true)
+			_, err := f.createSettings(ctx, account, "user", "eu.proxy.netbird.io", "")
+			require.NoError(t, err, "labeled pins under one cluster are the shared-cluster shape and must not refuse each other")
+		}
+	})
+}
+
+// TestCreateSettingsSelfAddressedRequiresPrivateCluster pins that the
+// capability gate applies to a self-addressed endpoint too: the service behind
+// it is the same private one, so a proxy that already declares the hostname
+// must have private capabilities, whether the account's own or a shared cluster's. A
+// hostname no proxy declares yet stays claimable (TestCreateSettingsSelfAddressed).
+func TestCreateSettingsSelfAddressedRequiresPrivateCluster(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("centralised proxy at the hostname is refused", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "central", "", "gw.example.com", ptrTo(false))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		_, err := f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
+		require.Error(t, err, "a self-addressed endpoint on a centralised proxy can never be served")
+		var sErr *status.Error
+		require.ErrorAs(t, err, &sErr)
+		assert.Equal(t, status.InvalidArgument, sErr.Type())
+		assert.Contains(t, err.Error(), "private capabilities")
+
+		_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
+		assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
+	})
+
+	t.Run("private proxy at the hostname is accepted", func(t *testing.T) {
+		f := newBootstrapFixture(t)
+		f.seedProxy(t, "private", "", "gw.example.com", ptrTo(true))
+		f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
+
+		created, err := f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
+		require.NoError(t, err)
+		assert.Equal(t, "gw.example.com", created.ProxyAddress)
+	})
+}

management/internals/modules/reverseproxy/domain/domain.go

@@ -1,5 +1,13 @@
 package domain
 
+import "time"
+
+// ValidationTTL is the time available to validate a custom domain registration.
+const ValidationTTL = 48 * time.Hour
+
+// ID identifies a custom domain registration.
+type ID string
+
 type Type string
 
 const (
@@ -8,12 +16,13 @@ const (
 )
 
 type Domain struct {
-	ID            string `gorm:"unique;primaryKey;autoIncrement"`
-	Domain        string `gorm:"unique"` // Domain records must be unique, this avoids domain reuse across accounts.
-	AccountID     string `gorm:"index"`
-	TargetCluster string // The proxy cluster this domain should be validated against
-	Type          Type   `gorm:"-"`
-	Validated     bool
+	ID                  string `gorm:"unique;primaryKey;autoIncrement"`
+	Domain              string `gorm:"unique"` // Domain records must be unique, this avoids domain reuse across accounts.
+	AccountID           string `gorm:"index"`
+	TargetCluster       string // The proxy cluster this domain should be validated against
+	Type                Type   `gorm:"-"`
+	Validated           bool
+	ValidationExpiresAt *time.Time `gorm:"index"`
 	// SupportsCustomPorts is populated at query time for free domains from the
 	// proxy cluster capabilities. Not persisted.
 	SupportsCustomPorts *bool `gorm:"-"`
@@ -36,7 +45,12 @@ func (d *Domain) EventMeta() map[string]any {
 	}
 }
 
+// Copy returns a copy with an independent validation deadline.
 func (d *Domain) Copy() *Domain {
 	dCopy := *d
+	if d.ValidationExpiresAt != nil {
+		expiresAt := *d.ValidationExpiresAt
+		dCopy.ValidationExpiresAt = &expiresAt
+	}
 	return &dCopy
 }

management/internals/modules/reverseproxy/domain/manager/expiration.go

@@ -0,0 +1,73 @@
+package manager
+
+import (
+	"context"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	"github.com/netbirdio/netbird/management/server/activity"
+)
+
+const (
+	validationCleanupInterval = 60 * time.Minute
+	validationCleanupBatch    = 100
+)
+
+// RunValidationCleanup removes expired registrations on startup and hourly until cancellation.
+func (m Manager) RunValidationCleanup(ctx context.Context) {
+	ticker := time.NewTicker(validationCleanupInterval)
+	defer ticker.Stop()
+	for {
+		m.cleanupExpiredDomains(ctx, time.Now().UTC())
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+	}
+}
+
+func (m Manager) cleanupExpiredDomains(ctx context.Context, now time.Time) {
+	var afterID domain.ID
+	for ctx.Err() == nil {
+		domains, err := m.store.GetExpiredCustomDomains(ctx, now, afterID, validationCleanupBatch)
+		if err != nil {
+			if ctx.Err() == nil {
+				log.WithContext(ctx).WithError(err).Error("list expired custom domain registrations")
+			}
+			return
+		}
+		for _, d := range domains {
+			if ctx.Err() != nil {
+				return
+			}
+			m.deleteExpiredDomain(ctx, d, now)
+			afterID = domain.ID(d.ID)
+		}
+		if len(domains) < validationCleanupBatch {
+			return
+		}
+	}
+}
+
+func (m Manager) deleteExpiredDomain(ctx context.Context, d *domain.Domain, now time.Time) {
+	deleted, err := m.store.DeleteExpiredCustomDomain(ctx, d, now)
+	if err != nil {
+		if ctx.Err() == nil {
+			log.WithContext(ctx).WithFields(log.Fields{"accountID": d.AccountID, "domainID": d.ID}).
+				WithError(err).Warn("could not expire custom domain registration")
+		}
+		return
+	}
+	if !deleted {
+		return
+	}
+	meta := d.EventMeta()
+	if d.ValidationExpiresAt != nil {
+		meta["validation_expires_at"] = d.ValidationExpiresAt.UTC().Format(time.RFC3339)
+	}
+	m.accountManager.StoreEvent(ctx, activity.SystemInitiator, d.ID, d.AccountID,
+		activity.CustomDomainValidationExpired, meta)
+}

management/internals/modules/reverseproxy/domain/manager/expiration_test.go

@@ -0,0 +1,274 @@
+package manager
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	"github.com/netbirdio/netbird/management/server/activity"
+	"github.com/netbirdio/netbird/management/server/mock_server"
+	nbstore "github.com/netbirdio/netbird/management/server/store"
+)
+
+func TestValidateDomain_ExpiredRegistration(t *testing.T) {
+	env := setupDomainTest(t)
+	ctx := context.Background()
+	d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "expired.example.com", testCluster)
+	require.NoError(t, err)
+	expiresAt := time.Now().Add(-time.Second)
+	db := env.store.(*nbstore.SqlStore).GetDB()
+	require.NoError(t, db.Model(&domain.Domain{}).Where("id = ?", d.ID).
+		Update("validation_expires_at", expiresAt).Error)
+	env.resolver.set("validation.expired.example.com", testCluster)
+
+	env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
+
+	stored := storedDomain(t, env.store, accountA, d.Domain)
+	require.NotNil(t, stored)
+	assert.False(t, stored.Validated, "an expired registration must not become usable before cleanup runs")
+}
+
+func TestCreateDomain_ValidationDeadline(t *testing.T) {
+	env := setupClockDomainTest(t)
+	synctest.Test(t, func(t *testing.T) {
+		ctx := context.Background()
+		createdAt := time.Now().UTC()
+		d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "pending.example.com", testCluster)
+		require.NoError(t, err)
+		require.NotNil(t, d.ValidationExpiresAt)
+		assert.Equal(t, createdAt.Add(48*time.Hour), *d.ValidationExpiresAt, "new registrations get 48 hours")
+
+		time.Sleep(time.Hour)
+		env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
+		stored := storedDomain(t, env.store, accountA, d.Domain)
+		require.NotNil(t, stored)
+		require.NotNil(t, stored.ValidationExpiresAt)
+		assert.WithinDuration(t, *d.ValidationExpiresAt, *stored.ValidationExpiresAt, 0, "failed validation must not extend the deadline")
+	})
+}
+
+func TestCleanupExpiredDomains_Boundaries(t *testing.T) {
+	env := setupDomainTest(t)
+	events := captureDomainEvents(env)
+	ctx := context.Background()
+	now := time.Now().UTC().Truncate(time.Second)
+	tests := []struct {
+		name      string
+		expiresAt time.Time
+		validated bool
+		deleted   bool
+	}{
+		{"expired", now.Add(-time.Second), false, true},
+		{"deadline", now, false, true},
+		{"pending", now.Add(time.Second), false, false},
+		{"validated", now.Add(-time.Hour), true, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			d := createExpiringDomain(t, env, tt.name+".example.com", tt.expiresAt)
+			if tt.validated {
+				require.NoError(t, env.store.(*nbstore.SqlStore).GetDB().Model(d).Update("validated", true).Error)
+			}
+			env.manager.cleanupExpiredDomains(ctx, now)
+			stored := storedDomain(t, env.store, accountA, d.Domain)
+			if !tt.deleted {
+				assert.NotNil(t, stored, "pending and validated registrations must survive cleanup")
+				return
+			}
+			assert.Nil(t, stored, "expired unused registrations must be removed")
+			replacement, err := env.manager.CreateDomain(ctx, accountB, accountBUser, d.Domain, testCluster)
+			require.NoError(t, err)
+			assert.NotEqual(t, d.ID, replacement.ID, "the released name must receive a fresh registration")
+			assert.False(t, replacement.Validated, "the new account must validate its own registration")
+		})
+	}
+	got := events.get()
+	require.Len(t, got, 2, "only successful expiration deletions emit events")
+	for _, event := range got {
+		assert.Equal(t, activity.CustomDomainValidationExpired, event.Activity, "use the requested expiration event")
+		assert.Equal(t, activity.SystemInitiator, event.InitiatorID, "cleanup is attributed to the system")
+		assert.Equal(t, accountA, event.AccountID, "expiration belongs to the original account")
+		assert.NotEmpty(t, event.TargetID, "retain the deleted domain ID")
+		assert.NotEmpty(t, event.Meta["domain"], "retain the deleted domain name")
+		assert.NotEmpty(t, event.Meta["validation_expires_at"], "include the validation deadline")
+	}
+}
+
+func TestCleanupExpiredDomains_ContinuesPastProtectedBatch(t *testing.T) {
+	env := setupDomainTest(t)
+	ctx := context.Background()
+	now := time.Now().UTC()
+	for i := range validationCleanupBatch {
+		d := createExpiringDomain(t, env, fmt.Sprintf("protected-%d.example.com", i), now.Add(-time.Hour))
+		require.NoError(t, env.store.CreateService(ctx, &rpservice.Service{
+			ID: fmt.Sprintf("service-%d", i), AccountID: accountA, Domain: "app." + d.Domain,
+		}))
+	}
+	unprotected := createExpiringDomain(t, env, "unused.example.com", now.Add(-time.Hour))
+	env.manager.cleanupExpiredDomains(ctx, now)
+	assert.Nil(t, storedDomain(t, env.store, accountA, unprotected.Domain), "protected registrations must not starve later batches")
+	remaining, err := env.store.ListCustomDomains(ctx, accountA)
+	require.NoError(t, err)
+	assert.Len(t, remaining, validationCleanupBatch, "all registrations with dependent services must survive")
+}
+
+func TestCleanupExpiredDomains_ConcurrentWorkers(t *testing.T) {
+	env := setupDomainTest(t)
+	events := captureDomainEvents(env)
+	now := time.Now().UTC()
+	d := createExpiringDomain(t, env, "concurrent.example.com", now.Add(-time.Hour))
+	var workers sync.WaitGroup
+	for range 2 {
+		workers.Go(func() { env.manager.cleanupExpiredDomains(context.Background(), now) })
+	}
+	workers.Wait()
+	assert.Nil(t, storedDomain(t, env.store, accountA, d.Domain), "one worker must remove the expired registration")
+	assert.Len(t, events.get(), 1, "only the worker that deletes the row may emit the event")
+}
+
+func TestRunValidationCleanup_HourlyAndRestart(t *testing.T) {
+	env := setupClockDomainTest(t)
+	synctest.Test(t, func(t *testing.T) {
+		events := captureDomainEvents(env)
+		now := time.Now().UTC()
+		startup := createExpiringDomain(t, env, "startup.example.com", now.Add(-time.Hour))
+		hourly := createExpiringDomain(t, env, "hourly.example.com", now.Add(time.Minute))
+		ctx, cancel := context.WithCancel(context.Background())
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			env.manager.RunValidationCleanup(ctx)
+		}()
+		synctest.Wait()
+		assert.Nil(t, storedDomain(t, env.store, accountA, startup.Domain), "startup must collect overdue registrations")
+		time.Sleep(59 * time.Minute)
+		synctest.Wait()
+		assert.NotNil(t, storedDomain(t, env.store, accountA, hourly.Domain), "cleanup must wait for the 60-minute interval")
+		time.Sleep(time.Minute)
+		synctest.Wait()
+		assert.Nil(t, storedDomain(t, env.store, accountA, hourly.Domain), "the hourly scan must collect expired registrations")
+		cancel()
+		<-done
+
+		offline := createExpiringDomain(t, env, "offline.example.com", time.Now().UTC().Add(time.Minute))
+		time.Sleep(2 * time.Hour)
+		assert.NotNil(t, storedDomain(t, env.store, accountA, offline.Domain), "a stopped worker must not continue deleting")
+		ctx, cancel = context.WithCancel(context.Background())
+		done = make(chan struct{})
+		go func() {
+			defer close(done)
+			env.manager.RunValidationCleanup(ctx)
+		}()
+		synctest.Wait()
+		assert.Nil(t, storedDomain(t, env.store, accountA, offline.Domain), "restart must use the persisted deadline")
+		cancel()
+		<-done
+		assert.Len(t, events.get(), 3, "each deletion should emit an expiration event")
+	})
+}
+
+type blockingDomainResolver struct {
+	started chan struct{}
+	release chan struct{}
+}
+
+func (r blockingDomainResolver) LookupCNAME(context.Context, string) (string, error) {
+	close(r.started)
+	<-r.release
+	return testCluster + ".", nil
+}
+
+func TestValidateDomain_DeadlinePassesDuringLookup(t *testing.T) {
+	for _, cleanup := range []bool{false, true} {
+		t.Run(fmt.Sprintf("cleanup=%t", cleanup), func(t *testing.T) {
+			env := setupClockDomainTest(t)
+			synctest.Test(t, func(t *testing.T) {
+				events := captureDomainEvents(env)
+				ctx := context.Background()
+				d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "late.example.com", testCluster)
+				require.NoError(t, err)
+				resolver := blockingDomainResolver{started: make(chan struct{}), release: make(chan struct{})}
+				env.manager.validator.Resolver = resolver
+				done := make(chan struct{})
+				go func() {
+					defer close(done)
+					env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
+				}()
+				<-resolver.started
+				time.Sleep(48 * time.Hour)
+				if cleanup {
+					env.manager.cleanupExpiredDomains(ctx, time.Now().UTC())
+					_, err = env.store.CreateCustomDomain(ctx, accountB, d.Domain, testCluster, false)
+					require.NoError(t, err)
+				}
+				close(resolver.release)
+				<-done
+				owner := accountA
+				if cleanup {
+					assert.Nil(t, storedDomain(t, env.store, accountA, d.Domain), "late validation must not restore the old claim")
+					owner = accountB
+				}
+				stored := storedDomain(t, env.store, owner, d.Domain)
+				require.NotNil(t, stored)
+				assert.False(t, stored.Validated, "late validation must not validate either claim")
+				for _, event := range events.get() {
+					assert.NotEqual(t, activity.DomainValidated, event.Activity, "a rejected write must not emit a validation event")
+				}
+			})
+		})
+	}
+}
+
+func setupClockDomainTest(t *testing.T) *domainTestEnv {
+	t.Helper()
+	// Network driver watchers cannot share cancellation channels across synctest bubbles.
+	// Store boundary and concurrency tests still exercise the selected database engine.
+	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
+	return setupDomainTest(t)
+}
+
+func createExpiringDomain(t *testing.T, env *domainTestEnv, name string, expiresAt time.Time) *domain.Domain {
+	t.Helper()
+	d, err := env.store.CreateCustomDomain(context.Background(), accountA, name, testCluster, false)
+	require.NoError(t, err)
+	require.NoError(t, env.store.(*nbstore.SqlStore).GetDB().Model(d).Update("validation_expires_at", expiresAt).Error)
+	d.ValidationExpiresAt = &expiresAt
+	return d
+}
+
+type domainEvents struct {
+	mu     sync.Mutex
+	events []*activity.Event
+}
+
+func captureDomainEvents(env *domainTestEnv) *domainEvents {
+	events := &domainEvents{}
+	env.manager.accountManager = &mock_server.MockAccountManager{
+		StoreEventFunc: func(_ context.Context, initiator, target, account string, code activity.ActivityDescriber, meta map[string]any) {
+			if code == activity.DomainAdded {
+				return
+			}
+			events.mu.Lock()
+			defer events.mu.Unlock()
+			events.events = append(events.events, &activity.Event{
+				InitiatorID: initiator, TargetID: target, AccountID: account,
+				Activity: code.(activity.Activity), Meta: meta,
+			})
+		},
+	}
+	return events
+}
+
+func (e *domainEvents) get() []*activity.Event {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	return append([]*activity.Event(nil), e.events...)
+}

management/internals/modules/reverseproxy/domain/manager/manager.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 
@@ -18,6 +19,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	nbstore "github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
+	nbdomain "github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/status"
 )
 
@@ -32,6 +34,8 @@ type store interface {
 	CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error)
 	UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error)
 	DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error
+	GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error)
+	DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error)
 }
 
 type proxyManager interface {
@@ -106,12 +110,13 @@ func (m Manager) GetDomains(ctx context.Context, accountID, userID string) ([]*d
 	// Add custom domains.
 	for _, d := range domains {
 		cd := &domain.Domain{
-			ID:            d.ID,
-			Domain:        d.Domain,
-			AccountID:     accountID,
-			TargetCluster: d.TargetCluster,
-			Type:          domain.TypeCustom,
-			Validated:     d.Validated,
+			ID:                  d.ID,
+			Domain:              d.Domain,
+			AccountID:           accountID,
+			TargetCluster:       d.TargetCluster,
+			Type:                domain.TypeCustom,
+			Validated:           d.Validated,
+			ValidationExpiresAt: d.ValidationExpiresAt,
 		}
 		if d.TargetCluster != "" {
 			cd.SupportsCustomPorts = m.proxyManager.ClusterSupportsCustomPorts(ctx, d.TargetCluster)
@@ -126,6 +131,7 @@ func (m Manager) GetDomains(ctx context.Context, accountID, userID string) ([]*d
 	return ret, nil
 }
 
+// CreateDomain registers a normalized custom domain and attempts DNS validation.
 func (m Manager) CreateDomain(ctx context.Context, accountID, userID, domainName, targetCluster string) (*domain.Domain, error) {
 	ok, ctx, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Services, operations.Create)
 	if err != nil {
@@ -135,6 +141,15 @@ func (m Manager) CreateDomain(ctx context.Context, accountID, userID, domainName
 		return nil, status.NewPermissionDeniedError()
 	}
 
+	parsed, err := nbdomain.FromString(strings.TrimSuffix(domainName, "."))
+	if err != nil {
+		return nil, status.Errorf(status.InvalidArgument, "invalid domain: %v", err)
+	}
+	domainName = parsed.PunycodeString()
+	if !nbdomain.IsValidDomainNoWildcard(domainName) {
+		return nil, status.Errorf(status.InvalidArgument, "invalid domain format")
+	}
+
 	// Verify the target cluster is in the available clusters for this account
 	allowList, err := m.getClusterAllowList(ctx, accountID)
 	if err != nil {
@@ -243,6 +258,14 @@ func (m Manager) ValidateDomain(ctx context.Context, accountID, userID, domainID
 		}).WithError(err).Error("get custom domain from store")
 		return
 	}
+	if d.Validated {
+		return
+	}
+	if d.ValidationExpiresAt == nil || !time.Now().Before(*d.ValidationExpiresAt) {
+		log.WithFields(log.Fields{"accountID": accountID, "domainID": domainID}).
+			Debug("custom domain validation window has expired")
+		return
+	}
 
 	// Validate only against the domain's target cluster
 	targetCluster := d.TargetCluster
@@ -263,20 +286,21 @@ func (m Manager) ValidateDomain(ctx context.Context, accountID, userID, domainID
 	}).Info("validating domain against target cluster")
 
 	if m.validator.IsValid(context.Background(), d.Domain, []string{targetCluster}) {
-		log.WithFields(log.Fields{
-			"accountID": accountID,
-			"domainID":  domainID,
-			"domain":    d.Domain,
-		}).Info("domain validated successfully")
 		d.Validated = true
 		if _, err := m.store.UpdateCustomDomain(context.Background(), accountID, d); err != nil {
-			log.WithFields(log.Fields{
+			entry := log.WithFields(log.Fields{
 				"accountID": accountID,
 				"domainID":  domainID,
-				"domain":    d.Domain,
-			}).WithError(err).Error("update custom domain in store")
+			}).WithError(err)
+			if sErr, ok := status.FromError(err); ok && sErr.Type() == status.PreconditionFailed {
+				entry.Debug("custom domain registration is no longer pending validation")
+				return
+			}
+			entry.Error("update custom domain in store")
 			return
 		}
+		log.WithFields(log.Fields{"accountID": accountID, "domainID": domainID}).
+			Info("custom domain validated successfully")
 
 		m.accountManager.StoreEvent(context.Background(), userID, domainID, accountID, activity.DomainValidated, d.EventMeta())
 	} else {

management/internals/modules/reverseproxy/domain/manager/manager_realstore_test.go

@@ -99,7 +99,7 @@ func setupDomainTest(t *testing.T) *domainTestEnv {
 	proxyMgr, err := proxymanager.NewManager(testStore, noop.NewMeterProvider().Meter(""))
 	require.NoError(t, err)
 
-	_, err = proxyMgr.Connect(ctx, "proxy-1", "session-1", testCluster, "127.0.0.1", nil, nil)
+	_, err = proxyMgr.Connect(ctx, "proxy-1", "session-1", testCluster, "127.0.0.1", "", nil, nil)
 	require.NoError(t, err)
 
 	resolver := &stubResolver{cnames: make(map[string]string)}
@@ -296,11 +296,8 @@ func TestValidateDomain_PermissionDeniedDoesNotValidate(t *testing.T) {
 	assert.Error(t, err, "the domain must still be unservable")
 }
 
-// Validation runs asynchronously, so it can finish after the domain was
-// deleted and then write a stale row back. gorm's Save falls back to an insert
-// when an update affects no rows, which would resurrect the domain as
-// validated; UpdateCustomDomain avoids that by selecting explicit columns.
-// This pins that behaviour, since dropping the Select would reintroduce it.
+// A validation finishing after deletion must reject the stale write, without
+// restoring the registration or reporting successful validation.
 func TestUpdateCustomDomain_DoesNotResurrectDeletedDomain(t *testing.T) {
 	ctx := context.Background()
 	env := setupDomainTest(t)
@@ -315,11 +312,9 @@ func TestUpdateCustomDomain_DoesNotResurrectDeletedDomain(t *testing.T) {
 	require.Nil(t, storedDomain(t, env.store, accountA, "racy.example.com"), "the domain should be gone")
 
 	// What an in-flight validation would write once its CNAME check succeeded.
-	// The write has to succeed for the assertion below to mean anything: a
-	// rejected write would leave the domain absent for the wrong reason.
 	stale.Validated = true
 	_, err = env.store.UpdateCustomDomain(ctx, accountA, stale)
-	require.NoError(t, err, "the update itself must succeed, so absence is not just a failed write")
+	require.Error(t, err, "a deleted registration must reject a late validation")
 
 	assert.Nil(t, storedDomain(t, env.store, accountA, "racy.example.com"),
 		"a late validation write must not recreate a deleted domain")

management/internals/modules/reverseproxy/domain/manager/manager_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -208,6 +209,14 @@ func (s *stubStore) DeleteCustomDomain(context.Context, string, string) error {
 	panic("not used in allow-list tests")
 }
 
+func (s *stubStore) GetExpiredCustomDomains(context.Context, time.Time, domain.ID, int) ([]*domain.Domain, error) {
+	panic("not used in allow-list tests")
+}
+
+func (s *stubStore) DeleteExpiredCustomDomain(context.Context, *domain.Domain, time.Time) (bool, error) {
+	panic("not used in allow-list tests")
+}
+
 // TestGetClusterAllowList_DedicatedGatewayAddressExcluded pins invariant (B)'s
 // chokepoint: a self-addressed settings pin reserves the account's gateway
 // address, so it is dropped from the allow list — which, because the

management/internals/modules/reverseproxy/domain/manager/normalization_test.go

@@ -0,0 +1,83 @@
+package manager
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestCreateDomain_NormalizesName(t *testing.T) {
+	for _, tt := range []struct {
+		name      string
+		input     string
+		canonical string
+	}{
+		{"mixed case", "Apps.Example.COM", "apps.example.com"},
+		{"unicode", "münchen.example.com", "xn--mnchen-3ya.example.com"},
+		{"trailing dot", "apps.example.com.", "apps.example.com"},
+		{"underscore", "My_App.example.com", "my_app.example.com"},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			env := setupDomainTest(t)
+			env.resolver.set("validation."+tt.canonical, testCluster)
+
+			created, err := env.manager.CreateDomain(ctx, accountA, accountAUser, tt.input, testCluster)
+			require.NoError(t, err)
+			assert.Equal(t, tt.canonical, created.Domain, "the response must use the normalized name")
+			assert.True(t, created.Validated, "the CNAME lookup must use the normalized name")
+			stored, err := env.store.GetCustomDomain(ctx, accountA, created.ID)
+			require.NoError(t, err)
+			assert.Equal(t, tt.canonical, stored.Domain, "the database must retain the normalized name")
+
+			_, err = env.manager.CreateDomain(ctx, accountB, accountBUser, tt.canonical, testCluster)
+			require.Error(t, err)
+			sErr, ok := status.FromError(err)
+			require.True(t, ok, "an equivalent name must return a typed conflict")
+			assert.Equal(t, status.AlreadyExists, sErr.Type(), "normalization must precede the availability check")
+		})
+	}
+}
+
+func TestCreateDomain_NormalizedNameCanValidateLater(t *testing.T) {
+	ctx := context.Background()
+	env := setupDomainTest(t)
+	created, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "Apps.Example.COM.", testCluster)
+	require.NoError(t, err)
+	require.False(t, created.Validated, "a missing CNAME must leave the normalized registration pending")
+
+	env.resolver.set("validation.apps.example.com", testCluster)
+	env.manager.ValidateDomain(ctx, accountA, accountAUser, created.ID)
+	stored, err := env.store.GetCustomDomain(ctx, accountA, created.ID)
+	require.NoError(t, err)
+	assert.Equal(t, "apps.example.com", stored.Domain, "retrying validation must retain the normalized name")
+	assert.True(t, stored.Validated, "later validation must look up the normalized name")
+}
+
+func TestCreateDomain_RejectsInvalidName(t *testing.T) {
+	ctx := context.Background()
+	env := setupDomainTest(t)
+	for _, name := range []string{
+		"", ".", "app..example.com", "app.example.com..", "-app.example.com",
+		"app%.example.com", "app!.example.com", "*.example.com", "app example.com",
+		"https://example.com", strings.Repeat("a", 64) + ".example.com",
+	} {
+		t.Run(name, func(t *testing.T) {
+			// A matching DNS response must not make a malformed name acceptable.
+			env.resolver.set("validation."+name, testCluster)
+			_, err := env.manager.CreateDomain(ctx, accountA, accountAUser, name, testCluster)
+			require.Error(t, err)
+			sErr, ok := status.FromError(err)
+			require.True(t, ok, "invalid names must return a typed client error")
+			assert.Equal(t, status.InvalidArgument, sErr.Type(), "malformed names must be rejected before storage")
+		})
+	}
+	stored, err := env.store.ListCustomDomains(ctx, accountA)
+	require.NoError(t, err)
+	assert.Empty(t, stored, "invalid registration attempts must not reserve any names")
+}

management/internals/modules/reverseproxy/proxy/manager.go

@@ -11,7 +11,7 @@ import (
 
 // Manager defines the interface for proxy operations
 type Manager interface {
-	Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *Capabilities) (*Proxy, error)
+	Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress, version string, accountID *string, capabilities *Capabilities) (*Proxy, error)
 	Disconnect(ctx context.Context, proxyID, sessionID string) error
 	Heartbeat(ctx context.Context, p *Proxy) error
 	GetActiveClusterAddresses(ctx context.Context) ([]string, error)

management/internals/modules/reverseproxy/proxy/manager/manager.go

@@ -50,7 +50,7 @@ func NewManager(store store, meter metric.Meter) (*Manager, error) {
 
 // Connect registers a new proxy connection in the database.
 // capabilities may be nil for old proxies that do not report them.
-func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *proxy.Capabilities) (*proxy.Proxy, error) {
+func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress, version string, accountID *string, capabilities *proxy.Capabilities) (*proxy.Proxy, error) {
 	now := time.Now()
 	var caps proxy.Capabilities
 	if capabilities != nil {
@@ -61,6 +61,7 @@ func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddres
 		SessionID:      sessionID,
 		ClusterAddress: clusterAddress,
 		IPAddress:      ipAddress,
+		Version:        truncateVersion(version),
 		AccountID:      accountID,
 		LastSeen:       now,
 		ConnectedAt:    &now,
@@ -78,6 +79,7 @@ func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddres
 		"sessionID":      sessionID,
 		"clusterAddress": clusterAddress,
 		"ipAddress":      ipAddress,
+		"version":        p.Version,
 	}).Info("proxy connected")
 
 	return p, nil
@@ -184,3 +186,13 @@ func (m *Manager) DeleteAccountCluster(ctx context.Context, clusterAddress, acco
 	}
 	return nil
 }
+
+// truncateVersion cuts a proxy-reported version to the column width so an
+// oversized value cannot fail the save and block the connect.
+func truncateVersion(version string) string {
+	runes := []rune(version)
+	if len(runes) <= proxy.MaxVersionLength {
+		return version
+	}
+	return string(runes[:proxy.MaxVersionLength])
+}

management/internals/modules/reverseproxy/proxy/manager/manager_test.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
+	"unicode/utf8"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -124,7 +126,7 @@ func TestConnect_WithAccountID(t *testing.T) {
 	}
 
 	mgr := newTestManager(s)
-	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", &accountID, nil)
+	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", "0.60.0", &accountID, nil)
 	require.NoError(t, err)
 
 	require.NotNil(t, savedProxy)
@@ -132,6 +134,7 @@ func TestConnect_WithAccountID(t *testing.T) {
 	assert.Equal(t, "session-1", savedProxy.SessionID)
 	assert.Equal(t, "cluster.example.com", savedProxy.ClusterAddress)
 	assert.Equal(t, "10.0.0.1", savedProxy.IPAddress)
+	assert.Equal(t, "0.60.0", savedProxy.Version, "reported proxy version should be stored")
 	assert.Equal(t, &accountID, savedProxy.AccountID)
 	assert.Equal(t, proxy.StatusConnected, savedProxy.Status)
 	assert.NotNil(t, savedProxy.ConnectedAt)
@@ -147,7 +150,7 @@ func TestConnect_WithoutAccountID(t *testing.T) {
 	}
 
 	mgr := newTestManager(s)
-	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "eu.proxy.netbird.io", "10.0.0.1", nil, nil)
+	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "eu.proxy.netbird.io", "10.0.0.1", "", nil, nil)
 	require.NoError(t, err)
 
 	require.NotNil(t, savedProxy)
@@ -155,6 +158,29 @@ func TestConnect_WithoutAccountID(t *testing.T) {
 	assert.Equal(t, proxy.StatusConnected, savedProxy.Status)
 }
 
+func TestConnect_TruncatesOversizedVersion(t *testing.T) {
+	var savedProxy *proxy.Proxy
+	s := &mockStore{
+		saveProxyFunc: func(_ context.Context, p *proxy.Proxy) error {
+			savedProxy = p
+			return nil
+		},
+	}
+
+	// Multi-byte runes make sure the cut counts characters, as varchar does,
+	// and never splits a rune into invalid UTF-8.
+	version := strings.Repeat("ü", proxy.MaxVersionLength+10)
+
+	mgr := newTestManager(s)
+	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", version, nil, nil)
+	require.NoError(t, err)
+
+	require.NotNil(t, savedProxy)
+	assert.Equal(t, proxy.MaxVersionLength, utf8.RuneCountInString(savedProxy.Version), "stored version should be cut to the column width")
+	assert.True(t, utf8.ValidString(savedProxy.Version), "stored version should remain valid UTF-8")
+	assert.True(t, strings.HasPrefix(version, savedProxy.Version), "stored version should be a prefix of the reported one")
+}
+
 func TestConnect_StoreError(t *testing.T) {
 	s := &mockStore{
 		saveProxyFunc: func(_ context.Context, _ *proxy.Proxy) error {
@@ -163,7 +189,7 @@ func TestConnect_StoreError(t *testing.T) {
 	}
 
 	mgr := newTestManager(s)
-	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", nil, nil)
+	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", "", nil, nil)
 	assert.Error(t, err)
 }
 

management/internals/modules/reverseproxy/proxy/manager_mock.go

@@ -113,18 +113,18 @@ func (mr *MockManagerMockRecorder) ClusterSupportsPrivate(ctx, clusterAddr any)
 }
 
 // Connect mocks base method.
-func (m *MockManager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *Capabilities) (*Proxy, error) {
+func (m *MockManager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress, version string, accountID *string, capabilities *Capabilities) (*Proxy, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Connect", ctx, proxyID, sessionID, clusterAddress, ipAddress, accountID, capabilities)
+	ret := m.ctrl.Call(m, "Connect", ctx, proxyID, sessionID, clusterAddress, ipAddress, version, accountID, capabilities)
 	ret0, _ := ret[0].(*Proxy)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // Connect indicates an expected call of Connect.
-func (mr *MockManagerMockRecorder) Connect(ctx, proxyID, sessionID, clusterAddress, ipAddress, accountID, capabilities any) *gomock.Call {
+func (mr *MockManagerMockRecorder) Connect(ctx, proxyID, sessionID, clusterAddress, ipAddress, version, accountID, capabilities any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockManager)(nil).Connect), ctx, proxyID, sessionID, clusterAddress, ipAddress, accountID, capabilities)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockManager)(nil).Connect), ctx, proxyID, sessionID, clusterAddress, ipAddress, version, accountID, capabilities)
 }
 
 // CountAccountProxies mocks base method.

management/internals/modules/reverseproxy/proxy/proxy.go

@@ -9,6 +9,9 @@ const (
 	StatusDisconnected = "disconnected"
 )
 
+// MaxVersionLength is the width of the Version column, in characters.
+const MaxVersionLength = 255
+
 // Capabilities describes what a proxy can handle, as reported via gRPC.
 // Nil fields mean the proxy never reported this capability.
 type Capabilities struct {
@@ -31,6 +34,7 @@ type Proxy struct {
 	SessionID      string    `gorm:"type:varchar(36)"`
 	ClusterAddress string    `gorm:"type:varchar(255);not null;index:idx_proxy_cluster_status"`
 	IPAddress      string    `gorm:"type:varchar(45)"`
+	Version        string    `gorm:"type:varchar(255)"`
 	AccountID      *string   `gorm:"type:varchar(255);index:idx_proxy_account_id"`
 	LastSeen       time.Time `gorm:"not null;index:idx_proxy_last_seen"`
 	ConnectedAt    *time.Time

management/internals/modules/reverseproxy/service/manager/domain_validation_test.go

@@ -30,7 +30,7 @@ func withRealDomainManager(t *testing.T, mgr *Manager, testStore store.Store) {
 	proxyMgr, err := proxymanager.NewManager(testStore, noop.NewMeterProvider().Meter(""))
 	require.NoError(t, err)
 
-	_, err = proxyMgr.Connect(ctx, "proxy-1", "session-1", validationTestCluster, "127.0.0.1", nil, nil)
+	_, err = proxyMgr.Connect(ctx, "proxy-1", "session-1", validationTestCluster, "127.0.0.1", "", nil, nil)
 	require.NoError(t, err)
 
 	accountMgr := &mock_server.MockAccountManager{

management/internals/server/boot.go

@@ -31,6 +31,7 @@ import (
 	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 	networkmapdbfactory "github.com/netbirdio/netbird/management/internals/network_map_db/factory"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server/activity"
 	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
@@ -111,7 +112,8 @@ func (s *BaseServer) NetworkMapStore() *networkmapdb.NetworkMapDBStoreImpl {
 			s.Config.StoreConfig.Engine,
 			s.Config.Datadir,
 			s.IntegratedValidator(),
-			s.SettingsManager())
+			s.SettingsManager(),
+		)
 		// networkmap db store supports postgres and sqlite backends only
 		// for other backends a fallback is used, so NotSupportedStoreEngineError
 		// is not a fatal error
@@ -180,24 +182,7 @@ func (s *BaseServer) RateLimiter() *middleware.APIRateLimiter {
 
 func (s *BaseServer) GRPCServer() *grpc.Server {
 	return Create(s, func() *grpc.Server {
-		trustedPeers := s.Config.ReverseProxy.TrustedPeers
-		defaultTrustedPeers := []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")}
-		if len(trustedPeers) == 0 || slices.Equal[[]netip.Prefix](trustedPeers, defaultTrustedPeers) {
-			log.WithContext(context.Background()).Warn("TrustedPeers are configured to default value '0.0.0.0/0', '::/0'. This allows connection IP spoofing.")
-			trustedPeers = defaultTrustedPeers
-		}
-		trustedHTTPProxies := s.Config.ReverseProxy.TrustedHTTPProxies
-		trustedProxiesCount := s.Config.ReverseProxy.TrustedHTTPProxiesCount
-		if len(trustedHTTPProxies) > 0 && trustedProxiesCount > 0 {
-			log.WithContext(context.Background()).Warn("TrustedHTTPProxies and TrustedHTTPProxiesCount both are configured. " +
-				"This is not recommended way to extract X-Forwarded-For. Consider using one of these options.")
-		}
-		realipOpts := []realip.Option{
-			realip.WithTrustedPeers(trustedPeers),
-			realip.WithTrustedProxies(trustedHTTPProxies),
-			realip.WithTrustedProxiesCount(trustedProxiesCount),
-			realip.WithHeaders([]string{realip.XForwardedFor, realip.XRealIp}),
-		}
+		realipOpts := realIPOptions(s.Config.ReverseProxy)
 		proxyUnary, proxyStream, proxyAuthClose := nbgrpc.NewProxyAuthInterceptors(s.Store())
 		s.proxyAuthClose = proxyAuthClose
 		gRPCOpts := []grpc.ServerOption{
@@ -333,7 +318,7 @@ func (s *BaseServer) AccessLogsManager() accesslogs.Manager {
 	})
 }
 
-func loadTLSConfig(certFile string, certKey string) (*tls.Config, error) {
+func loadTLSConfig(certFile, certKey string) (*tls.Config, error) {
 	// Load server's certificate and private key
 	serverCert, err := tls.LoadX509KeyPair(certFile, certKey)
 	if err != nil {
@@ -380,3 +365,37 @@ func streamInterceptor(
 	wrapped.WrappedContext = context.WithValue(ctx, nbContext.RequestIDKey, reqID)
 	return handler(srv, wrapped)
 }
+
+// realIPOptions builds the real-IP middleware options.
+//
+// Empty TrustedPeers trusts all IPv4 and IPv6 sources. Configure TrustedPeers
+// with the reverse proxy address or network.
+//
+// X-Forwarded-For takes precedence over X-Real-IP.
+func realIPOptions(cfg nbconfig.ReverseProxy) []realip.Option {
+	trustedPeers := cfg.TrustedPeers
+	if len(trustedPeers) == 0 {
+		trustedPeers = []netip.Prefix{
+			netip.MustParsePrefix("0.0.0.0/0"),
+			netip.MustParsePrefix("::/0"),
+		}
+	}
+	if idx := slices.IndexFunc(trustedPeers, func(p netip.Prefix) bool { return p.Bits() == 0 }); idx >= 0 {
+		log.WithContext(context.Background()).Warnf("TrustedPeers contains the default route %s, which trusts "+
+			"X-Forwarded-For from every client and allows connection IP spoofing. Set TrustedPeers to the address "+
+			"of your reverse proxy.", trustedPeers[idx])
+	}
+	if cfg.TrustedHTTPProxiesCount > 0 {
+		log.WithContext(context.Background()).Warn(
+			"TrustedHTTPProxiesCount skips X-Forwarded-For entries by position before TrustedHTTPProxies filters by address. " +
+				"An incorrect count may skip the real client IP and produce an incorrect source address.",
+		)
+	}
+
+	return []realip.Option{
+		realip.WithTrustedPeers(trustedPeers),
+		realip.WithTrustedProxies(cfg.TrustedHTTPProxies),
+		realip.WithTrustedProxiesCount(cfg.TrustedHTTPProxiesCount),
+		realip.WithHeaders([]string{realip.XForwardedFor, realip.XRealIp}),
+	}
+}

management/internals/server/realip_test.go

@@ -0,0 +1,179 @@
+package server
+
+import (
+	"context"
+	"io"
+	"net"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/protobuf/types/known/emptypb"
+
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+)
+
+const (
+	realIPProbeMethod       = "/netbird.test.RealIPProbe/Probe"
+	realIPProbeStreamMethod = "/netbird.test.RealIPProbe/ProbeStream"
+)
+
+// realIPProbe records the real IP the middleware derived for each call.
+type realIPProbe struct {
+	got chan string
+}
+
+func (p *realIPProbe) record(ctx context.Context) {
+	addr, _ := realip.FromContext(ctx)
+	p.got <- addr.String()
+}
+
+func (p *realIPProbe) wait(t *testing.T) string {
+	t.Helper()
+
+	select {
+	case got := <-p.got:
+		return got
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for probe")
+		return ""
+	}
+}
+
+func startProbeServer(t *testing.T, cfg nbconfig.ReverseProxy) (*grpc.ClientConn, *realIPProbe) {
+	t.Helper()
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	probe := &realIPProbe{got: make(chan string, 1)}
+	opts := realIPOptions(cfg)
+	srv := grpc.NewServer(
+		grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(opts...)),
+		grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(opts...)),
+	)
+	srv.RegisterService(&grpc.ServiceDesc{
+		ServiceName: "netbird.test.RealIPProbe",
+		HandlerType: (*any)(nil),
+		Methods: []grpc.MethodDesc{{
+			MethodName: "Probe",
+			Handler: func(_ any, ctx context.Context, dec func(any) error, interceptor grpc.UnaryServerInterceptor) (any, error) {
+				req := new(emptypb.Empty)
+				if err := dec(req); err != nil {
+					return nil, err
+				}
+				handler := func(ctx context.Context, _ any) (any, error) {
+					probe.record(ctx)
+					return &emptypb.Empty{}, nil
+				}
+				if interceptor == nil {
+					return handler(ctx, req)
+				}
+				return interceptor(ctx, req, &grpc.UnaryServerInfo{FullMethod: realIPProbeMethod}, handler)
+			},
+		}},
+		Streams: []grpc.StreamDesc{{
+			StreamName:    "ProbeStream",
+			ServerStreams: true,
+			Handler: func(_ any, stream grpc.ServerStream) error {
+				probe.record(stream.Context())
+				return nil
+			},
+		}},
+	}, probe)
+
+	go func() { _ = srv.Serve(listener) }()
+	t.Cleanup(srv.Stop)
+
+	conn, err := grpc.NewClient(listener.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	return conn, probe
+}
+
+func callUnary(t *testing.T, conn *grpc.ClientConn, probe *realIPProbe, kv ...string) string {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	ctx = metadata.AppendToOutgoingContext(ctx, kv...)
+	require.NoError(t, conn.Invoke(ctx, realIPProbeMethod, &emptypb.Empty{}, &emptypb.Empty{}))
+
+	return probe.wait(t)
+}
+
+func callStream(t *testing.T, conn *grpc.ClientConn, probe *realIPProbe, kv ...string) string {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	ctx = metadata.AppendToOutgoingContext(ctx, kv...)
+	desc := &grpc.StreamDesc{StreamName: "ProbeStream", ServerStreams: true}
+	stream, err := conn.NewStream(ctx, desc, realIPProbeStreamMethod)
+	require.NoError(t, err)
+	require.NoError(t, stream.CloseSend())
+	require.ErrorIs(t, stream.RecvMsg(&emptypb.Empty{}), io.EOF)
+
+	return probe.wait(t)
+}
+
+func assertRealIP(t *testing.T, cfg nbconfig.ReverseProxy, want string, kv ...string) {
+	t.Helper()
+
+	conn, probe := startProbeServer(t, cfg)
+	t.Run("unary", func(t *testing.T) {
+		assert.Equal(t, want, callUnary(t, conn, probe, kv...))
+	})
+	t.Run("stream", func(t *testing.T) {
+		assert.Equal(t, want, callStream(t, conn, probe, kv...))
+	})
+}
+
+func TestRealIPDefaultTrustsForwardedHeaders(t *testing.T) {
+	assertRealIP(t, nbconfig.ReverseProxy{}, "203.0.113.44",
+		realip.XForwardedFor, "203.0.113.44",
+		realip.XRealIp, "203.0.113.44",
+	)
+}
+
+func TestRealIPUntrustedPeerIgnoresForwardedHeaders(t *testing.T) {
+	cfg := nbconfig.ReverseProxy{TrustedPeers: []netip.Prefix{netip.MustParsePrefix("10.9.8.7/32")}}
+
+	assertRealIP(t, cfg, "127.0.0.1",
+		realip.XForwardedFor, "203.0.113.44",
+		realip.XRealIp, "203.0.113.44",
+	)
+}
+
+func TestRealIPTrustedPeerHonoursForwardedHeaders(t *testing.T) {
+	cfg := nbconfig.ReverseProxy{TrustedPeers: []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")}}
+
+	assertRealIP(t, cfg, "203.0.113.44",
+		realip.XForwardedFor, "203.0.113.44",
+		realip.XRealIp, "203.0.113.44",
+	)
+}
+
+func TestRealIPReadsXRealIPWhenProxyCountSkipsForwardedFor(t *testing.T) {
+	cfg := nbconfig.ReverseProxy{
+		TrustedPeers:            []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")},
+		TrustedHTTPProxiesCount: 1,
+	}
+
+	t.Run("no X-Forwarded-For", func(t *testing.T) {
+		assertRealIP(t, cfg, "203.0.113.44", realip.XRealIp, "203.0.113.44")
+	})
+	t.Run("single-entry X-Forwarded-For", func(t *testing.T) {
+		assertRealIP(t, cfg, "198.51.100.7",
+			realip.XForwardedFor, "203.0.113.44",
+			realip.XRealIp, "198.51.100.7",
+		)
+	})
+}

management/internals/server/server.go

@@ -23,6 +23,8 @@ import (
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/metrics"
 	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/shared/lifecycle"
+	"github.com/netbirdio/netbird/shared/profiling"
 	"github.com/netbirdio/netbird/util/wsproxy"
 	wsproxyserver "github.com/netbirdio/netbird/util/wsproxy/server"
 	"github.com/netbirdio/netbird/version"
@@ -36,6 +38,8 @@ const (
 	DefaultSelfHostedDomain = "netbird.selfhosted"
 
 	ContainerKeyBaseServer = "baseServer"
+
+	applicationName = "management"
 )
 
 type Server interface {
@@ -66,7 +70,8 @@ type BaseServer struct {
 	disableLegacyManagementPort bool
 	autoResolveDomains          bool
 
-	proxyAuthClose func()
+	proxyAuthClose    func()
+	domainCleanupStop func()
 
 	// grpcExtensions holds additional gRPC services, interceptors, and shutdown
 	// hooks registered by external modules via RegisterGRPCExtension. Populated
@@ -74,12 +79,15 @@ type BaseServer struct {
 	grpcExtensions []GRPCExtension
 
 	listener    net.Listener
+	tlsConfig   *tls.Config
 	certManager *autocert.Manager
 	update      *version.Update
 
 	errCh  chan error
 	wg     sync.WaitGroup
 	cancel context.CancelFunc
+
+	lifecycle.StopHandlers
 }
 
 // Config holds the configuration parameters for creating a new server
@@ -94,6 +102,7 @@ type Config struct {
 	DisableGeoliteUpdate        bool
 	UserDeleteFromIDPEnabled    bool
 	AutoResolveDomains          bool
+	TLSConfig                   *tls.Config
 }
 
 // NewServer initializes and configures a new Server instance
@@ -110,9 +119,13 @@ func NewServer(cfg *Config) *BaseServer {
 		disableLegacyManagementPort: cfg.DisableLegacyManagementPort,
 		mgmtMetricsPort:             cfg.MgmtMetricsPort,
 		autoResolveDomains:          cfg.AutoResolveDomains,
+		tlsConfig:                   cfg.TLSConfig,
 	}
 	s.container[ContainerKeyBaseServer] = s
 
+	stopProfiling := profiling.Start(applicationName)
+	s.OnStop(stopProfiling)
+
 	return s
 }
 
@@ -122,6 +135,14 @@ func (s *BaseServer) AfterInit(fn func(s *BaseServer)) {
 
 // Start begins listening for HTTP requests on the configured address
 func (s *BaseServer) Start(ctx context.Context) error {
+	if err := s.start(ctx); err != nil {
+		s.RunStopHandlers()
+		return err
+	}
+	return nil
+}
+
+func (s *BaseServer) start(ctx context.Context) error {
 	srvCtx, cancel := context.WithCancel(ctx)
 	s.cancel = cancel
 	s.errCh = make(chan error, 4)
@@ -139,21 +160,9 @@ func (s *BaseServer) Start(ctx context.Context) error {
 	}
 	s.EphemeralManager().LoadInitialPeers(srvCtx)
 
-	var tlsConfig *tls.Config
-	tlsEnabled := false
-	if s.Config.HttpConfig.LetsEncryptDomain != "" {
-		s.certManager, err = encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
-		if err != nil {
-			return fmt.Errorf("failed creating LetsEncrypt cert manager: %v", err)
-		}
-		tlsEnabled = true
-	} else if s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "" {
-		tlsConfig, err = loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
-		if err != nil {
-			log.WithContext(srvCtx).Errorf("cannot load TLS credentials: %v", err)
-			return err
-		}
-		tlsEnabled = true
+	tlsEnabled, err := s.setupTLS(srvCtx)
+	if err != nil {
+		return err
 	}
 
 	installationID, err := getInstallationID(srvCtx, s.Store())
@@ -215,8 +224,8 @@ func (s *BaseServer) Start(ctx context.Context) error {
 			log.WithContext(ctx).Infof("running HTTP server (LetsEncrypt challenge handler): %s", cml.Addr().String())
 			s.serveHTTP(ctx, cml, s.certManager.HTTPHandler(nil))
 		}
-	case tlsConfig != nil:
-		s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), tlsConfig)
+	case s.tlsConfig != nil:
+		s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), s.tlsConfig)
 		if err != nil {
 			return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
 		}
@@ -236,14 +245,60 @@ func (s *BaseServer) Start(ctx context.Context) error {
 	s.update.SetOnUpdateListener(func() {
 		log.WithContext(ctx).Infof("your management version, \"%s\", is outdated, a new management version is available. Learn more here: https://github.com/netbirdio/netbird/releases", version.NetbirdVersion())
 	})
+	s.startDomainCleanup(srvCtx)
 
 	return nil
 }
+func (s *BaseServer) startDomainCleanup(ctx context.Context) {
+	if s.domainCleanupStop != nil {
+		return
+	}
+	mgr := s.ReverseProxyDomainManager()
+	ctx, cancel := context.WithCancel(ctx)
+	done := make(chan struct{})
+	s.domainCleanupStop = func() {
+		cancel()
+		<-done
+	}
+	go func() {
+		defer close(done)
+		mgr.RunValidationCleanup(ctx)
+	}()
+}
+
+// setupTLS resolves the listener's TLS source: an injected config wins over the HttpConfig certificate settings
+func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
+	switch {
+	case s.tlsConfig != nil:
+		return true, nil
+	case s.Config.HttpConfig.LetsEncryptDomain != "":
+		certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
+		if err != nil {
+			return false, fmt.Errorf("failed creating LetsEncrypt cert manager: %v", err)
+		}
+		s.certManager = certManager
+		return true, nil
+	case s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "":
+		tlsConfig, err := loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
+		if err != nil {
+			log.WithContext(ctx).Errorf("cannot load TLS credentials: %v", err)
+			return false, err
+		}
+		s.tlsConfig = tlsConfig
+		return true, nil
+	default:
+		return false, nil
+	}
+}
 
 // Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
 func (s *BaseServer) Stop() error {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
+	defer s.RunStopHandlers()
+	if s.domainCleanupStop != nil {
+		s.domainCleanupStop()
+	}
 
 	s.IntegratedValidator().Stop(ctx)
 	if s.GeoLocationManager() != nil {

management/internals/shared/grpc/peer_update_handler.go

@@ -0,0 +1,135 @@
+package grpc
+
+import (
+	"context"
+	"time"
+
+	"github.com/netbirdio/netbird/encryption"
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func PeerUpdateHandlerFactory(
+	peerKey wgtypes.Key,
+	updates chan *network_map.UpdateMessage,
+	secretsManager SecretsManager,
+	srv proto.ManagementService_SyncServer,
+	cleanupfunc func()) *PeerUpdateHandler {
+	return &PeerUpdateHandler{
+		peerKey:        peerKey,
+		updates:        updates,
+		secretsManager: secretsManager,
+		srv:            srv,
+		encrypter:      encryption.DefaultEncrypter{},
+		debouncer:      NewUpdateDebouncer(1000 * time.Millisecond),
+		cleanupFunc:    cleanupfunc,
+	}
+}
+
+// PeerUpdateHandler sends updates to the connected peer until the updates channel is closed.
+// It implements a backpressure mechanism that sends the first update immediately,
+// then debounces subsequent rapid updates, ensuring only the latest update is sent
+// after a quiet period.
+type PeerUpdateHandler struct {
+	peerKey        wgtypes.Key
+	updates        chan *network_map.UpdateMessage
+	appMetrics     telemetry.AppMetrics
+	secretsManager SecretsManager
+	srv            syncSender
+	encrypter      encryption.Encrypter
+	debouncer      Debouncer
+	cleanupFunc    func()
+}
+
+func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler {
+	pu.appMetrics = appMetrics
+	return pu
+}
+
+//go:generate go tool mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+type syncSender interface {
+	Send(*proto.EncryptedMessage) error
+	Context() context.Context
+}
+
+func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
+	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", pu.peerKey.String())
+
+	defer pu.debouncer.Stop()
+
+	for {
+		select {
+		// condition when there are some updates
+		// todo set the updates channel size to 1
+		case update, open := <-pu.updates:
+			if pu.appMetrics != nil {
+				pu.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(pu.updates) + 1)
+			}
+
+			if !open {
+				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", pu.peerKey.String())
+				pu.cleanupFunc()
+				return nil
+			}
+
+			log.WithContext(ctx).Tracef("received an update for peer %s", pu.peerKey.String())
+			if pu.debouncer.ProcessUpdate(update) {
+				// Send immediately (first update or after quiet period)
+				if err := pu.SendUpdate(ctx, update); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// Timer expired - quiet period reached, send pending updates if any
+		case <-pu.debouncer.TimerChannel():
+			pendingUpdates := pu.debouncer.GetPendingUpdates()
+			if len(pendingUpdates) == 0 {
+				continue
+			}
+			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), pu.peerKey.String())
+			for _, pendingUpdate := range pendingUpdates {
+				if err := pu.SendUpdate(ctx, pendingUpdate); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// condition when client <-> server connection has been terminated
+		case <-pu.srv.Context().Done():
+			// happens when connection drops, e.g. client disconnects
+			log.WithContext(ctx).Debugf("stream of peer %s has been closed", pu.peerKey.String())
+			pu.cleanupFunc()
+			return pu.srv.Context().Err()
+		}
+	}
+}
+
+func (pu *PeerUpdateHandler) SendUpdate(ctx context.Context, update *network_map.UpdateMessage) error {
+	key, err := pu.secretsManager.GetWGKey()
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+
+	encryptedResp, err := pu.encrypter.EncryptMessage(pu.peerKey, key, update.Update)
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+	err = pu.srv.Send(&proto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     encryptedResp,
+	})
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed sending update message")
+	}
+	log.WithContext(ctx).Tracef("sent an update to peer %s", pu.peerKey.String())
+	return nil
+}

management/internals/shared/grpc/peer_update_handler_test.go

@@ -0,0 +1,155 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	pb "github.com/golang/protobuf/proto" //nolint
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+func TestSendPeerUpdates_FirstUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	// mock a first update, should send it right away
+	updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true)
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	pu.updates <- &msg
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_TimerUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	updateDebouncer.EXPECT().GetPendingUpdates().Return([]*network_map.UpdateMessage{&msg})
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	timeCh <- time.Now()
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_ServerContextDone(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx, cancel := context.WithCancel(context.TODO())
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	cancel()
+	wg.Wait()
+}
+
+func mustGenerateKey(t *testing.T) wgtypes.Key {
+	t.Helper()
+	k, err := wgtypes.GenerateKey()
+	assert.NoError(t, err)
+	return k
+}
+
+func mustMarshal(t *testing.T, msg *network_map.UpdateMessage) []byte {
+	t.Helper()
+	r, err := pb.Marshal(msg.Update)
+	assert.NoError(t, err)
+	return r
+}
+
+type testEncrypter struct{}
+
+func (testEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
+	return pb.Marshal(message)
+}
+
+type pbMatcher struct {
+	x pb.Message
+}
+
+func (pbm pbMatcher) Matches(x any) bool {
+	msg, ok := x.(pb.Message)
+	if !ok {
+		return false
+	}
+	return pb.Equal(pbm.x, msg)
+}
+
+func (pbm pbMatcher) String() string {
+	return fmt.Sprintf("is equal to %s (%T)", pbm.x, pbm.x)
+}

management/internals/shared/grpc/proxy.go

@@ -102,7 +102,8 @@ type ProxyServiceServer struct {
 
 	mu sync.RWMutex
 	// Manager for reverse proxy operations
-	serviceManager rpservice.Manager
+	serviceManager   rpservice.Manager
+	credentialLimits credentialVerificationLimiter
 	// agentNetworkSynth produces synthesised reverse-proxy services from
 	// Agent Network state. Optional — when nil the snapshot path only ships
 	// persisted services.
@@ -242,9 +243,10 @@ func (s *ProxyServiceServer) cleanupStaleProxies(ctx context.Context) {
 	}
 }
 
-// Close stops background goroutines.
+// Close stops background goroutines and releases credential verification state.
 func (s *ProxyServiceServer) Close() {
 	s.cancel()
+	s.credentialLimits.close()
 }
 
 // SetServiceManager sets the service manager. Must be called before serving.
@@ -412,6 +414,7 @@ func (s *ProxyServiceServer) SetProxyController(proxyController proxy.Controller
 type proxyConnectParams struct {
 	proxyID      string
 	address      string
+	version      string
 	capabilities *proto.ProxyCapabilities
 }
 
@@ -422,6 +425,7 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 		return err
 	}
 	params.capabilities = req.GetCapabilities()
+	params.version = req.GetVersion()
 
 	conn, proxyRecord, err := s.registerProxyConnection(stream.Context(), params, &proxyConnection{
 		stream: stream,
@@ -455,6 +459,7 @@ func (s *ProxyServiceServer) SyncMappings(stream proto.ProxyService_SyncMappings
 		return err
 	}
 	params.capabilities = init.GetCapabilities()
+	params.version = init.GetVersion()
 
 	conn, proxyRecord, err := s.registerProxyConnection(stream.Context(), params, &proxyConnection{
 		syncStream: stream,
@@ -566,7 +571,7 @@ func (s *ProxyServiceServer) registerProxyConnection(ctx context.Context, params
 		}
 	}
 
-	proxyRecord, err := s.proxyManager.Connect(ctx, params.proxyID, sessionID, params.address, peerInfo, accountID, caps)
+	proxyRecord, err := s.proxyManager.Connect(ctx, params.proxyID, sessionID, params.address, peerInfo, params.version, accountID, caps)
 	if err != nil {
 		cancel()
 		if accountID != nil {
@@ -1223,6 +1228,7 @@ func shallowCloneMapping(m *proto.ProxyMapping) *proto.ProxyMapping {
 	}
 }
 
+// Authenticate verifies service credentials and issues a session token.
 func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
 	if err := enforceAccountScope(ctx, req.GetAccountId()); err != nil {
 		return nil, err
@@ -1234,6 +1240,14 @@ func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.Authen
 		return nil, status.Errorf(codes.FailedPrecondition, "get service from store: %v", err)
 	}
 
+	switch req.GetRequest().(type) {
+	case *proto.AuthenticateRequest_Pin, *proto.AuthenticateRequest_Password:
+		key := credentialVerificationKey{accountID: credentialAccountID(service.AccountID), serviceID: credentialServiceID(service.ID)}
+		if err := s.credentialLimits.allow(key); err != nil {
+			return nil, err
+		}
+	}
+
 	authenticated, userId, method := s.authenticateRequest(ctx, req, service)
 
 	// Non-OIDC schemes (PIN/Password/Header) authenticate against per-service

management/internals/shared/grpc/proxy_connect_version_test.go

@@ -0,0 +1,93 @@
+package grpc
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+const (
+	versionTestProxyID = "proxy-a"
+	versionTestCluster = "cluster.example.com"
+	versionTestVersion = "0.60.0"
+)
+
+// hangupStream cancels its context on the first Send, emulating a proxy that
+// disconnects right after receiving the initial snapshot. The legacy stream
+// carries no proxy-to-management messages, so this is the only way for
+// GetMappingUpdate to return.
+type hangupStream struct {
+	recordingStream
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *hangupStream) Send(m *proto.GetMappingUpdateResponse) error {
+	s.cancel()
+	return s.recordingStream.Send(m)
+}
+
+func (s *hangupStream) Context() context.Context { return s.ctx }
+
+// newVersionTestServer wires a server whose proxy manager only accepts a
+// Connect carrying versionTestVersion, so a dropped or mangled version fails
+// the test as an unexpected call.
+func newVersionTestServer(t *testing.T) *ProxyServiceServer {
+	t.Helper()
+	ctrl := gomock.NewController(t)
+
+	svcMgr := rpservice.NewMockManager(ctrl)
+	svcMgr.EXPECT().GetGlobalServices(gomock.Any()).Return(nil, nil)
+
+	proxyMgr := proxy.NewMockManager(ctrl)
+	proxyMgr.EXPECT().
+		Connect(gomock.Any(), versionTestProxyID, gomock.Any(), versionTestCluster, gomock.Any(), versionTestVersion, gomock.Any(), gomock.Any()).
+		Return(&proxy.Proxy{ID: versionTestProxyID, Version: versionTestVersion}, nil)
+	proxyMgr.EXPECT().Disconnect(gomock.Any(), versionTestProxyID, gomock.Any()).Return(nil)
+
+	s := newSnapshotTestServer(t, 10)
+	s.serviceManager = svcMgr
+	s.proxyManager = proxyMgr
+	return s
+}
+
+func TestSyncMappings_ForwardsProxyVersion(t *testing.T) {
+	s := newVersionTestServer(t)
+
+	// The init carries the version, the ack acknowledges the empty snapshot,
+	// and the exhausted fake stream then ends the RPC.
+	stream := &syncRecordingStream{
+		recvMsgs: []*proto.SyncMappingsRequest{
+			{Msg: &proto.SyncMappingsRequest_Init{Init: &proto.SyncMappingsInit{
+				ProxyId: versionTestProxyID,
+				Address: versionTestCluster,
+				Version: versionTestVersion,
+			}}},
+			ackMsg(),
+		},
+	}
+
+	err := s.SyncMappings(stream)
+	require.ErrorContains(t, err, "no more recv messages")
+}
+
+func TestGetMappingUpdate_ForwardsProxyVersion(t *testing.T) {
+	s := newVersionTestServer(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	stream := &hangupStream{ctx: ctx, cancel: cancel}
+
+	err := s.GetMappingUpdate(&proto.GetMappingUpdateRequest{
+		ProxyId: versionTestProxyID,
+		Address: versionTestCluster,
+		Version: versionTestVersion,
+	}, stream)
+	require.ErrorIs(t, err, context.Canceled)
+}

management/internals/shared/grpc/proxy_credential_limiter.go

@@ -0,0 +1,101 @@
+package grpc
+
+import (
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const (
+	credentialVerificationInterval        = 6 * time.Second
+	credentialVerificationBurst           = 5
+	credentialVerificationMaxServices     = 4096
+	credentialVerificationIdleTimeout     = 15 * time.Minute
+	credentialVerificationCleanupInterval = time.Minute
+)
+
+type credentialAccountID string
+type credentialServiceID string
+
+type credentialVerificationKey struct {
+	accountID credentialAccountID
+	serviceID credentialServiceID
+}
+
+type credentialVerificationBudget struct {
+	limiter  *rate.Limiter
+	lastUsed time.Time
+}
+
+// The zero value is ready to use. Budgets are local to this Management process;
+// proxy replicas reaching this process share a service's verification budget.
+type credentialVerificationLimiter struct {
+	mu          sync.Mutex
+	now         func() time.Time
+	services    map[credentialVerificationKey]*credentialVerificationBudget
+	nextCleanup time.Time
+	closed      bool
+}
+
+func (l *credentialVerificationLimiter) allow(key credentialVerificationKey) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.closed {
+		return status.Error(codes.Unavailable, "credential verification is closed")
+	}
+	now := time.Now()
+	if l.now != nil {
+		now = l.now()
+	}
+	l.cleanup(now)
+	budget := l.services[key]
+	if budget == nil {
+		if len(l.services) >= credentialVerificationMaxServices {
+			return credentialVerificationThrottled(credentialVerificationCleanupInterval)
+		}
+		if l.services == nil {
+			l.services = make(map[credentialVerificationKey]*credentialVerificationBudget)
+		}
+		budget = &credentialVerificationBudget{limiter: rate.NewLimiter(rate.Every(credentialVerificationInterval), credentialVerificationBurst)}
+		l.services[key] = budget
+	}
+	budget.lastUsed = now
+	if budget.limiter.AllowN(now, 1) {
+		return nil
+	}
+	delay := max(time.Nanosecond, time.Duration((1-budget.limiter.TokensAt(now))*float64(credentialVerificationInterval)))
+	return credentialVerificationThrottled(delay)
+}
+
+func (l *credentialVerificationLimiter) cleanup(now time.Time) {
+	if now.Before(l.nextCleanup) {
+		return
+	}
+	l.nextCleanup = now.Add(credentialVerificationCleanupInterval)
+	for key, budget := range l.services {
+		if now.Sub(budget.lastUsed) >= credentialVerificationIdleTimeout {
+			delete(l.services, key)
+		}
+	}
+}
+
+func (l *credentialVerificationLimiter) close() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.closed = true
+	l.services = nil
+}
+
+func credentialVerificationThrottled(delay time.Duration) error {
+	s := status.New(codes.ResourceExhausted, "too many credential verification attempts")
+	withRetry, err := s.WithDetails(&errdetails.RetryInfo{RetryDelay: durationpb.New(delay)})
+	if err != nil {
+		return s.Err()
+	}
+	return withRetry.Err()
+}

management/internals/shared/grpc/proxy_credential_limiter_test.go

@@ -0,0 +1,79 @@
+package grpc
+
+import (
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestCredentialVerificationRefillAndIsolation(t *testing.T) {
+	now := time.Now()
+	l := credentialVerificationLimiter{now: func() time.Time { return now }}
+	key := credentialVerificationKey{accountID: "account", serviceID: "service"}
+	for range credentialVerificationBurst {
+		require.NoError(t, l.allow(key))
+	}
+	err := l.allow(key)
+	require.Equal(t, codes.ResourceExhausted, status.Code(err), "the burst must be bounded")
+	now = now.Add(3 * time.Second)
+	err = l.allow(key)
+	require.Equal(t, codes.ResourceExhausted, status.Code(err), "a partially refilled token must not permit a check")
+	details := status.Convert(err).Details()
+	require.Len(t, details, 1, "throttling must provide RetryInfo")
+	retry, ok := details[0].(*errdetails.RetryInfo)
+	require.True(t, ok, "retry details must use the standard message")
+	assert.Equal(t, 3*time.Second, retry.RetryDelay.AsDuration(), "retry hint must reflect time until the next check")
+	now = now.Add(3 * time.Second)
+	require.NoError(t, l.allow(key))
+	assert.Equal(t, codes.ResourceExhausted, status.Code(l.allow(key)), "only one check must refill every six seconds")
+	require.NoError(t, l.allow(credentialVerificationKey{accountID: "other-account", serviceID: key.serviceID}))
+	require.NoError(t, l.allow(credentialVerificationKey{accountID: key.accountID, serviceID: "other-service"}))
+}
+
+func TestCredentialVerificationCapacityAndExpiry(t *testing.T) {
+	now := time.Now()
+	l := credentialVerificationLimiter{now: func() time.Time { return now }}
+	for i := range credentialVerificationMaxServices {
+		require.NoError(t, l.allow(credentialVerificationKey{accountID: "account", serviceID: credentialServiceID(strconv.Itoa(i))}))
+	}
+	key := credentialVerificationKey{accountID: "account", serviceID: "new-service"}
+	assert.Equal(t, codes.ResourceExhausted, status.Code(l.allow(key)), "capacity exhaustion must deny new checks")
+	now = now.Add(credentialVerificationIdleTimeout)
+	for range credentialVerificationBurst {
+		require.NoError(t, l.allow(key))
+	}
+	assert.Equal(t, codes.ResourceExhausted, status.Code(l.allow(key)), "expiry must retain the normal burst bound")
+}
+
+func TestCredentialVerificationConcurrentChecksAndClose(t *testing.T) {
+	var l credentialVerificationLimiter
+	key := credentialVerificationKey{accountID: "account", serviceID: "service"}
+	var admitted atomic.Int32
+	var wg sync.WaitGroup
+	for range 100 {
+		wg.Go(func() {
+			if err := l.allow(key); err == nil {
+				admitted.Add(1)
+			} else {
+				assert.Equal(t, codes.ResourceExhausted, status.Code(err), "excess checks must be throttled")
+			}
+		})
+	}
+	wg.Wait()
+	assert.EqualValues(t, credentialVerificationBurst, admitted.Load(), "concurrent checks must share the burst")
+	for range 10 {
+		wg.Go(l.close)
+		wg.Go(func() { assert.Error(t, l.allow(key)) })
+	}
+	wg.Wait()
+	assert.Empty(t, l.services, "closing must release retained budgets")
+	assert.Equal(t, codes.Unavailable, status.Code(l.allow(key)), "checks after close must fail closed")
+}

management/internals/shared/grpc/proxy_credentials.md

@@ -0,0 +1,18 @@
+# Reverse proxy credential verification
+
+The `ProxyService.Authenticate` RPC limits PIN and password checks before
+verifying their Argon2 hashes. Both methods share one budget per account and
+service: a burst of five checks, replenishing one check every six seconds
+(ten per minute). Successful and failed checks consume the budget. Account
+scope and service lookup run before the limiter.
+
+Excess checks receive gRPC `ResourceExhausted` with a standard `RetryInfo` delay.
+Updated proxies translate it to HTTP 429 and `Retry-After`. Older proxies show
+an authentication-service error but cannot bypass the Management limit.
+
+Budgets are held in memory per Management process and reset on restart. Proxy
+replicas reaching the same Management process share its budgets. Multiple
+Management processes have independent budgets; this is not a cluster-wide
+limit. At most 4,096 service budgets are retained, with idle entries expiring
+after fifteen minutes. Capacity exhaustion denies new checks until entries
+expire. Closing the server releases the retained state.

management/internals/shared/grpc/proxy_credentials_test.go

@@ -0,0 +1,131 @@
+package grpc_test
+
+import (
+	"context"
+	"net"
+	"net/netip"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	servicemanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service/manager"
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
+	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+func credentialServer(t *testing.T) (*nbgrpc.ProxyServiceServer, context.Context, grpc.UnaryServerInterceptor) {
+	t.Helper()
+	ctx := context.Background()
+	s, err := store.NewStore(ctx, types.SqliteStoreEngine, t.TempDir(), nil, false)
+	require.NoError(t, err)
+	t.Cleanup(func() { assert.NoError(t, s.Close(ctx)) })
+	require.NoError(t, s.SaveAccount(ctx, &types.Account{Id: "account"}))
+	keys, err := sessionkey.GenerateKeyPair()
+	require.NoError(t, err)
+	for _, id := range []string{"service", "other-service"} {
+		svc := &service.Service{
+			ID: id, AccountID: "account", Name: id, Domain: id + ".example.com",
+			Enabled: true, SessionPrivateKey: keys.PrivateKey, SessionPublicKey: keys.PublicKey,
+			Auth: service.AuthConfig{
+				PinAuth:      &service.PINAuthConfig{Enabled: true, Pin: "842716"},
+				PasswordAuth: &service.PasswordAuthConfig{Enabled: true, Password: "test-password"},
+			},
+		}
+		require.NoError(t, svc.Auth.HashSecrets())
+		require.NoError(t, s.CreateService(ctx, svc))
+	}
+	account := "account"
+	token, err := types.CreateNewProxyAccessToken("test proxy", time.Hour, &account, "admin")
+	require.NoError(t, err)
+	require.NoError(t, s.SaveProxyAccessToken(ctx, &token.ProxyAccessToken))
+	ctx = metadata.NewIncomingContext(ctx, metadata.Pairs("authorization", "Bearer "+string(token.PlainToken)))
+	ctx = peer.NewContext(ctx, &peer.Peer{Addr: net.TCPAddrFromAddrPort(netip.MustParseAddrPort("192.0.2.1:443"))})
+	server := nbgrpc.NewProxyServiceServer(nil, nil, nil, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil)
+	t.Cleanup(server.Close)
+	server.SetServiceManager(servicemanager.NewManager(s, nil, nil, nil, nil, nil))
+	interceptor, _, closeInterceptor := nbgrpc.NewProxyAuthInterceptors(s)
+	t.Cleanup(closeInterceptor)
+	return server, ctx, interceptor
+}
+
+func TestAuthenticateCredentialRateLimit(t *testing.T) {
+	server, ctx, interceptor := credentialServer(t)
+	authenticate := func(req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
+		response, err := interceptor(ctx, req, &grpc.UnaryServerInfo{FullMethod: "/management.ProxyService/Authenticate"}, func(ctx context.Context, req any) (any, error) {
+			return server.Authenticate(ctx, req.(*proto.AuthenticateRequest))
+		})
+		if err != nil {
+			return nil, err
+		}
+		return response.(*proto.AuthenticateResponse), nil
+	}
+	for i := range 5 {
+		req := &proto.AuthenticateRequest{AccountId: "account", Id: "service"}
+		if i%2 == 0 {
+			req.Request = &proto.AuthenticateRequest_Pin{Pin: &proto.PinRequest{Pin: "000000"}}
+		} else {
+			req.Request = &proto.AuthenticateRequest_Password{Password: &proto.PasswordRequest{Password: "wrong-password"}}
+		}
+		resp, err := authenticate(req)
+		require.NoError(t, err)
+		assert.False(t, resp.GetSuccess(), "incorrect PINs and passwords must be denied")
+		assert.Empty(t, resp.GetSessionToken(), "incorrect credentials must not issue a token")
+	}
+	req := &proto.AuthenticateRequest{AccountId: "account", Id: "service", Request: &proto.AuthenticateRequest_Pin{Pin: &proto.PinRequest{Pin: "842716"}}}
+	resp, err := authenticate(req)
+	assert.Nil(t, resp, "a throttled verification must not return a session")
+	require.Equal(t, codes.ResourceExhausted, status.Code(err), "PIN and password checks must share a service budget even with a valid proxy token")
+	details := status.Convert(err).Details()
+	require.Len(t, details, 1, "throttled responses must include a retry hint")
+	retry, ok := details[0].(*errdetails.RetryInfo)
+	require.True(t, ok, "the hint must use the standard RetryInfo message")
+	assert.Positive(t, retry.RetryDelay.AsDuration(), "the retry delay must be positive")
+	assert.LessOrEqual(t, retry.RetryDelay.AsDuration(), 6*time.Second, "the service must replenish one verification every six seconds")
+	req.AccountId = "another-account"
+	_, err = authenticate(req)
+	assert.Equal(t, codes.PermissionDenied, status.Code(err), "account scope must still be enforced before throttling")
+	req.AccountId = "account"
+	req.Id = "other-service"
+	resp, err = authenticate(req)
+	require.NoError(t, err)
+	assert.True(t, resp.GetSuccess(), "one service's throttle must not block another service")
+	assert.NotEmpty(t, resp.GetSessionToken(), "valid credentials on another service must issue a session")
+}
+
+func TestAuthenticateCredentialConcurrentLimit(t *testing.T) {
+	server, _, _ := credentialServer(t)
+	req := &proto.AuthenticateRequest{AccountId: "account", Id: "service", Request: &proto.AuthenticateRequest_Pin{Pin: &proto.PinRequest{Pin: "000000"}}}
+	var checked, throttled atomic.Int32
+	var wg sync.WaitGroup
+	for range 20 {
+		wg.Go(func() {
+			resp, err := server.Authenticate(context.Background(), req)
+			switch status.Code(err) {
+			case codes.OK:
+				checked.Add(1)
+				assert.False(t, resp.GetSuccess(), "incorrect credentials must be denied")
+			case codes.ResourceExhausted:
+				throttled.Add(1)
+			default:
+				assert.NoError(t, err)
+			}
+		})
+	}
+	wg.Wait()
+	assert.EqualValues(t, 5, checked.Load(), "only the burst budget may reach concurrent credential verification")
+	assert.EqualValues(t, 15, throttled.Load(), "excess concurrent checks must be throttled")
+}

management/internals/shared/grpc/server.go

@@ -337,7 +337,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 
 	s.syncSem.Add(-1)
 
-	return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv, syncStart)
+	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }).
+		WithMetrics(s.appMetrics).HandleUpdates(ctx)
 }
 
 func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementService_JobServer) (wgtypes.Key, error) {
@@ -404,91 +405,6 @@ func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgt
 	}
 }
 
-// handleUpdates sends updates to the connected peer until the updates channel is closed.
-// It implements a backpressure mechanism that sends the first update immediately,
-// then debounces subsequent rapid updates, ensuring only the latest update is sent
-// after a quiet period.
-func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
-
-	// Create a debouncer for this peer connection
-	debouncer := NewUpdateDebouncer(1000 * time.Millisecond)
-	defer debouncer.Stop()
-
-	for {
-		select {
-		// condition when there are some updates
-		// todo set the updates channel size to 1
-		case update, open := <-updates:
-			if s.appMetrics != nil {
-				s.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(updates) + 1)
-			}
-
-			if !open {
-				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String())
-				s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-				return nil
-			}
-
-			log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String())
-			if debouncer.ProcessUpdate(update) {
-				// Send immediately (first update or after quiet period)
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// Timer expired - quiet period reached, send pending updates if any
-		case <-debouncer.TimerChannel():
-			pendingUpdates := debouncer.GetPendingUpdates()
-			if len(pendingUpdates) == 0 {
-				continue
-			}
-			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String())
-			for _, pendingUpdate := range pendingUpdates {
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// condition when client <-> server connection has been terminated
-		case <-srv.Context().Done():
-			// happens when connection drops, e.g. client disconnects
-			log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
-			s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-			return srv.Context().Err()
-		}
-	}
-}
-
-// sendUpdate encrypts the update message using the peer key and the server's wireguard key,
-// then sends the encrypted message to the connected peer via the sync server.
-func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update)
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-	err = srv.Send(&proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
-		Body:     encryptedResp,
-	})
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed sending update message")
-	}
-	log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String())
-	return nil
-}
-
 // sendJob encrypts the update message using the peer key and the server's wireguard key,
 // then sends the encrypted message to the connected peer via the sync server.
 func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Event, srv proto.ManagementService_JobServer) error {

management/internals/shared/grpc/sync_sender_mock.go

@@ -0,0 +1,70 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./peer_update_handler.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	proto "github.com/netbirdio/netbird/shared/management/proto"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MocksyncSender is a mock of syncSender interface.
+type MocksyncSender struct {
+	ctrl     *gomock.Controller
+	recorder *MocksyncSenderMockRecorder
+	isgomock struct{}
+}
+
+// MocksyncSenderMockRecorder is the mock recorder for MocksyncSender.
+type MocksyncSenderMockRecorder struct {
+	mock *MocksyncSender
+}
+
+// NewMocksyncSender creates a new mock instance.
+func NewMocksyncSender(ctrl *gomock.Controller) *MocksyncSender {
+	mock := &MocksyncSender{ctrl: ctrl}
+	mock.recorder = &MocksyncSenderMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MocksyncSender) EXPECT() *MocksyncSenderMockRecorder {
+	return m.recorder
+}
+
+// Context mocks base method.
+func (m *MocksyncSender) Context() context.Context {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Context")
+	ret0, _ := ret[0].(context.Context)
+	return ret0
+}
+
+// Context indicates an expected call of Context.
+func (mr *MocksyncSenderMockRecorder) Context() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MocksyncSender)(nil).Context))
+}
+
+// Send mocks base method.
+func (m *MocksyncSender) Send(arg0 *proto.EncryptedMessage) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Send", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Send indicates an expected call of Send.
+func (mr *MocksyncSenderMockRecorder) Send(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MocksyncSender)(nil).Send), arg0)
+}

management/internals/shared/grpc/token_mgr.go

@@ -25,6 +25,8 @@ import (
 const defaultDuration = 12 * time.Hour
 
 // SecretsManager used to manage TURN and relay secrets
+//
+//go:generate go tool mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
 type SecretsManager interface {
 	GenerateTurnToken() (*Token, error)
 	GenerateRelayToken() (*Token, error)

management/internals/shared/grpc/token_mgr_mock.go

@@ -0,0 +1,111 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./token_mgr.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "go.uber.org/mock/gomock"
+	wgtypes "golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+// MockSecretsManager is a mock of SecretsManager interface.
+type MockSecretsManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockSecretsManagerMockRecorder
+	isgomock struct{}
+}
+
+// MockSecretsManagerMockRecorder is the mock recorder for MockSecretsManager.
+type MockSecretsManagerMockRecorder struct {
+	mock *MockSecretsManager
+}
+
+// NewMockSecretsManager creates a new mock instance.
+func NewMockSecretsManager(ctrl *gomock.Controller) *MockSecretsManager {
+	mock := &MockSecretsManager{ctrl: ctrl}
+	mock.recorder = &MockSecretsManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockSecretsManager) EXPECT() *MockSecretsManagerMockRecorder {
+	return m.recorder
+}
+
+// CancelRefresh mocks base method.
+func (m *MockSecretsManager) CancelRefresh(peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "CancelRefresh", peerKey)
+}
+
+// CancelRefresh indicates an expected call of CancelRefresh.
+func (mr *MockSecretsManagerMockRecorder) CancelRefresh(peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRefresh", reflect.TypeOf((*MockSecretsManager)(nil).CancelRefresh), peerKey)
+}
+
+// GenerateRelayToken mocks base method.
+func (m *MockSecretsManager) GenerateRelayToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateRelayToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateRelayToken indicates an expected call of GenerateRelayToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateRelayToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateRelayToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateRelayToken))
+}
+
+// GenerateTurnToken mocks base method.
+func (m *MockSecretsManager) GenerateTurnToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateTurnToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateTurnToken indicates an expected call of GenerateTurnToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateTurnToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateTurnToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateTurnToken))
+}
+
+// GetWGKey mocks base method.
+func (m *MockSecretsManager) GetWGKey() (wgtypes.Key, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWGKey")
+	ret0, _ := ret[0].(wgtypes.Key)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWGKey indicates an expected call of GetWGKey.
+func (mr *MockSecretsManagerMockRecorder) GetWGKey() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGKey", reflect.TypeOf((*MockSecretsManager)(nil).GetWGKey))
+}
+
+// SetupRefresh mocks base method.
+func (m *MockSecretsManager) SetupRefresh(ctx context.Context, accountID, peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "SetupRefresh", ctx, accountID, peerKey)
+}
+
+// SetupRefresh indicates an expected call of SetupRefresh.
+func (mr *MockSecretsManagerMockRecorder) SetupRefresh(ctx, accountID, peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetupRefresh", reflect.TypeOf((*MockSecretsManager)(nil).SetupRefresh), ctx, accountID, peerKey)
+}

management/internals/shared/grpc/update_debouncer.go

@@ -6,6 +6,14 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 )
 
+//go:generate go tool mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+type Debouncer interface {
+	Stop()
+	TimerChannel() <-chan time.Time
+	ProcessUpdate(update *network_map.UpdateMessage) bool
+	GetPendingUpdates() []*network_map.UpdateMessage
+}
+
 // UpdateDebouncer implements a backpressure mechanism that:
 // - Sends the first update immediately
 // - Coalesces rapid subsequent network map updates (only latest matters)

management/internals/shared/grpc/update_debouncer_mock.go

@@ -0,0 +1,96 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./update_debouncer.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	reflect "reflect"
+	time "time"
+
+	network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockDebouncer is a mock of Debouncer interface.
+type MockDebouncer struct {
+	ctrl     *gomock.Controller
+	recorder *MockDebouncerMockRecorder
+	isgomock struct{}
+}
+
+// MockDebouncerMockRecorder is the mock recorder for MockDebouncer.
+type MockDebouncerMockRecorder struct {
+	mock *MockDebouncer
+}
+
+// NewMockDebouncer creates a new mock instance.
+func NewMockDebouncer(ctrl *gomock.Controller) *MockDebouncer {
+	mock := &MockDebouncer{ctrl: ctrl}
+	mock.recorder = &MockDebouncerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockDebouncer) EXPECT() *MockDebouncerMockRecorder {
+	return m.recorder
+}
+
+// GetPendingUpdates mocks base method.
+func (m *MockDebouncer) GetPendingUpdates() []*network_map.UpdateMessage {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPendingUpdates")
+	ret0, _ := ret[0].([]*network_map.UpdateMessage)
+	return ret0
+}
+
+// GetPendingUpdates indicates an expected call of GetPendingUpdates.
+func (mr *MockDebouncerMockRecorder) GetPendingUpdates() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPendingUpdates", reflect.TypeOf((*MockDebouncer)(nil).GetPendingUpdates))
+}
+
+// ProcessUpdate mocks base method.
+func (m *MockDebouncer) ProcessUpdate(update *network_map.UpdateMessage) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ProcessUpdate", update)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// ProcessUpdate indicates an expected call of ProcessUpdate.
+func (mr *MockDebouncerMockRecorder) ProcessUpdate(update any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessUpdate", reflect.TypeOf((*MockDebouncer)(nil).ProcessUpdate), update)
+}
+
+// Stop mocks base method.
+func (m *MockDebouncer) Stop() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Stop")
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockDebouncerMockRecorder) Stop() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockDebouncer)(nil).Stop))
+}
+
+// TimerChannel mocks base method.
+func (m *MockDebouncer) TimerChannel() <-chan time.Time {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "TimerChannel")
+	ret0, _ := ret[0].(<-chan time.Time)
+	return ret0
+}
+
+// TimerChannel indicates an expected call of TimerChannel.
+func (mr *MockDebouncerMockRecorder) TimerChannel() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimerChannel", reflect.TypeOf((*MockDebouncer)(nil).TimerChannel))
+}

management/internals/shared/grpc/validate_session_test.go

@@ -570,7 +570,7 @@ func (m *testValidateSessionServiceManager) DeleteAccountCluster(_ context.Conte
 
 type testValidateSessionProxyManager struct{}
 
-func (m *testValidateSessionProxyManager) Connect(_ context.Context, _, _, _, _ string, _ *string, _ *proxy.Capabilities) (*proxy.Proxy, error) {
+func (m *testValidateSessionProxyManager) Connect(_ context.Context, _, _, _, _, _ string, _ *string, _ *proxy.Capabilities) (*proxy.Proxy, error) {
 	return nil, nil
 }
 

management/server/account_test.go

@@ -3516,7 +3516,13 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
 
 	eventStore := &activity.InMemoryEventStore{}
 
-	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	// Everything built here watches this context; cancelling it on cleanup stops
+	// the metrics flushers, caches and controllers instead of leaking them for
+	// the rest of the package run.
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	metrics, err := telemetry.NewDefaultAppMetrics(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -3535,8 +3541,6 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
 		Return(nil).
 		AnyTimes()
 
-	ctx := context.Background()
-
 	cacheStore, err := cache.NewStore(ctx, 100*time.Millisecond, 300*time.Millisecond, 100)
 	if err != nil {
 		return nil, nil, err

management/server/activity/codes.go

@@ -284,9 +284,12 @@ const (
 	// AgentNetworkSettingsDeleted indicates that a user deleted the Agent Network account settings, releasing the endpoint
 	AgentNetworkSettingsDeleted Activity = 142
 
+	// CustomDomainValidationExpired indicates that an unvalidated domain registration expired.
+	CustomDomainValidationExpired Activity = 143
+
 	// AccountDebugBundleUploadURLUpdated indicates that a user changed where the
 	// account's peers upload their debug bundles
-	AccountDebugBundleUploadURLUpdated Activity = 143
+	AccountDebugBundleUploadURLUpdated Activity = 144
 
 	AccountDeleted Activity = 99999
 )
@@ -466,9 +469,10 @@ var activityMap = map[Activity]Code{
 	AccountMetricsPushEnabled:  {"Account metrics push enabled", "account.setting.metrics.push.enable"},
 	AccountMetricsPushDisabled: {"Account metrics push disabled", "account.setting.metrics.push.disable"},
 
-	DomainAdded:     {"Domain added", "domain.add"},
-	DomainDeleted:   {"Domain deleted", "domain.delete"},
-	DomainValidated: {"Domain validated", "domain.validate"},
+	DomainAdded:                   {"Domain added", "domain.add"},
+	DomainDeleted:                 {"Domain deleted", "domain.delete"},
+	DomainValidated:               {"Domain validated", "domain.validate"},
+	CustomDomainValidationExpired: {"Unvalidated domain registration expired", "domain.validation.expire"},
 }
 
 // StringCode returns a string code of the activity

management/server/activity/store/sql_store.go

@@ -165,16 +165,16 @@ func (store *Store) Get(ctx context.Context, accountID string, offset, limit int
 	return store.processResult(ctx, events)
 }
 
-// Save an event in the SQLite events table end encrypt the "email" element in meta map
-func (store *Store) Save(_ context.Context, event *activity.Event) (*activity.Event, error) {
+// Save persists an activity event and encrypts deleted user details using the caller's context.
+func (store *Store) Save(ctx context.Context, event *activity.Event) (*activity.Event, error) {
 	eventCopy := event.Copy()
-	meta, err := store.saveDeletedUserEmailAndNameInEncrypted(eventCopy)
+	meta, err := store.saveDeletedUserEmailAndNameInEncrypted(ctx, eventCopy)
 	if err != nil {
 		return nil, err
 	}
 	eventCopy.Meta = meta
 
-	if err = store.db.Create(eventCopy).Error; err != nil {
+	if err = store.db.WithContext(ctx).Create(eventCopy).Error; err != nil {
 		return nil, err
 	}
 
@@ -183,7 +183,7 @@ func (store *Store) Save(_ context.Context, event *activity.Event) (*activity.Ev
 
 // saveDeletedUserEmailAndNameInEncrypted if the meta contains email and name then store it in encrypted way and delete
 // this item from meta map
-func (store *Store) saveDeletedUserEmailAndNameInEncrypted(event *activity.Event) (map[string]any, error) {
+func (store *Store) saveDeletedUserEmailAndNameInEncrypted(ctx context.Context, event *activity.Event) (map[string]any, error) {
 	email, ok := event.Meta["email"]
 	if !ok {
 		return event.Meta, nil
@@ -211,7 +211,7 @@ func (store *Store) saveDeletedUserEmailAndNameInEncrypted(event *activity.Event
 	}
 	deletedUser.Name = encryptedName
 
-	err = store.db.Clauses(clause.OnConflict{
+	err = store.db.WithContext(ctx).Clauses(clause.OnConflict{
 		Columns:   []clause.Column{{Name: "id"}},
 		DoUpdates: clause.AssignmentColumns([]string{"email", "name"}),
 	}).Create(deletedUser).Error

management/server/activity/store/sql_store_test.go

@@ -7,11 +7,49 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/util/crypt"
 )
 
+func TestSave_CancellationWhileWaitingForConnection(t *testing.T) {
+	t.Setenv(storeEngineEnv, "sqlite")
+	key, err := crypt.GenerateKey()
+	require.NoError(t, err)
+	store, err := NewSqlStore(context.Background(), t.TempDir(), key)
+	require.NoError(t, err)
+	t.Cleanup(func() { assert.NoError(t, store.Close(context.Background())) })
+	db, err := store.db.DB()
+	require.NoError(t, err)
+	conn, err := db.Conn(context.Background())
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	result := make(chan error, 1)
+	go func() {
+		_, err := store.Save(ctx, &activity.Event{
+			Timestamp: time.Now().UTC(), Activity: activity.CustomDomainValidationExpired,
+			AccountID: "account-id", TargetID: "domain-id", InitiatorID: activity.SystemInitiator,
+		})
+		result <- err
+	}()
+	select {
+	case err := <-result:
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+		require.NoError(t, conn.Close())
+	case <-time.After(time.Second):
+		// Release the connection so a regression cannot leave the writer running.
+		require.NoError(t, conn.Close())
+		assert.ErrorIs(t, <-result, context.DeadlineExceeded)
+		t.Error("activity writes must stop waiting when their deadline expires")
+	}
+	events, err := store.Get(context.Background(), "account-id", 0, 10, true)
+	require.NoError(t, err)
+	assert.Empty(t, events, "a timed-out write must not persist after the connection is released")
+}
+
 func TestNewSqlStore(t *testing.T) {
 	dataDir := t.TempDir()
 	key, _ := crypt.GenerateKey()

management/server/agentnetwork_budgetrule_realstack_test.go

@@ -89,6 +89,7 @@ func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *t
 
 	// Bootstrap is an explicit settings create; providers have no settings
 	// side effects anymore.
+	seedPrivateProxyCluster(t, am.Store, clusterAddr)
 	before, err := mgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
 	require.NoError(t, err, "CreateSettings must bootstrap the row")
 	require.Equal(t, clusterAddr, before.ProxyAddress, "proxy address pinned at bootstrap")

management/server/agentnetwork_realstack_test.go

@@ -13,6 +13,7 @@ import (
 	networkmap "github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
 	agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/store"
@@ -92,6 +93,7 @@ func TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers(t *testing.T) {
 	// UpdateAccountPeers, which is the path under test.
 	agentMgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
 
+	seedPrivateProxyCluster(t, am.Store, clusterAddr)
 	_, err = agentMgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
 	require.NoError(t, err, "CreateSettings must bootstrap the endpoint")
 	// The bootstrap itself reconciles and queues updates on both channels;
@@ -222,3 +224,22 @@ func synthZoneRData(sync *nbproto.SyncResponse, clusterAddr, fqdn string) string
 	}
 	return ""
 }
+
+// seedPrivateProxyCluster registers a connected proxy with private capabilities in a
+// netbird client for clusterAddr, matching what a real deployment looks like
+// when the account bootstraps: the agent-network gateway service is always
+// private, so its cluster has to be one that can serve private services.
+func seedPrivateProxyCluster(t *testing.T, st store.Store, clusterAddr string) {
+	t.Helper()
+	private := true
+	now := time.Now().UTC()
+	require.NoError(t, st.SaveProxy(context.Background(), &rpproxy.Proxy{
+		ID:             "agent-net-proxy-" + clusterAddr,
+		SessionID:      "agent-net-session",
+		ClusterAddress: clusterAddr,
+		LastSeen:       now,
+		ConnectedAt:    &now,
+		Status:         rpproxy.StatusConnected,
+		Capabilities:   rpproxy.Capabilities{Private: &private},
+	}), "seeding the proxy cluster must succeed")
+}

management/server/event.go

@@ -61,23 +61,34 @@ func (am *DefaultAccountManager) GetEvents(ctx context.Context, accountID, userI
 	return filtered, nil
 }
 
+// StoreEvent records an activity, waiting for expiration events before cleanup can stop.
 func (am *DefaultAccountManager) StoreEvent(ctx context.Context, initiatorID, targetID, accountID string, activityID activity.ActivityDescriber, meta map[string]any) {
-	if isEnabled() {
-		go func() {
-			_, err := am.eventStore.Save(ctx, &activity.Event{
-				Timestamp:   time.Now().UTC(),
-				Activity:    activityID.(activity.Activity),
-				InitiatorID: initiatorID,
-				TargetID:    targetID,
-				AccountID:   accountID,
-				Meta:        meta,
-			})
-			if err != nil {
-				// todo add metric
-				log.WithContext(ctx).Errorf("received an error while storing an activity event, error: %s", err)
-			}
-		}()
+	if !isEnabled() {
+		return
 	}
+	eventStore := am.eventStore
+	save := func(ctx context.Context) {
+		_, err := eventStore.Save(ctx, &activity.Event{
+			Timestamp:   time.Now().UTC(),
+			Activity:    activityID.(activity.Activity),
+			InitiatorID: initiatorID,
+			TargetID:    targetID,
+			AccountID:   accountID,
+			Meta:        meta,
+		})
+		if err != nil {
+			log.WithContext(ctx).Errorf("received an error while storing an activity event, error: %s", err)
+		}
+	}
+	if activityID == activity.CustomDomainValidationExpired {
+		// The domain is already deleted; shutdown must allow its audit write to finish.
+		ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+		defer cancel()
+		save(ctx)
+		return
+	}
+	// Request cancellation must not discard the audit record of a completed operation.
+	go save(context.WithoutCancel(ctx))
 }
 
 type eventUserInfo struct {

management/server/event_test.go

@@ -6,10 +6,52 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/management/server/activity"
+	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
+	"github.com/netbirdio/netbird/util/crypt"
 )
 
+func TestStoreEvent_CanceledContext(t *testing.T) {
+	t.Setenv("NB_EVENT_ACTIVITY_LOG_ENABLED", "true")
+	t.Setenv("NB_ACTIVITY_EVENT_STORE_ENGINE", "sqlite")
+	for _, code := range []activity.Activity{activity.CustomDomainValidationExpired, activity.DomainAdded} {
+		t.Run(code.StringCode(), func(t *testing.T) {
+			dir := t.TempDir()
+			key, err := crypt.GenerateKey()
+			require.NoError(t, err)
+			eventStore, err := activitystore.NewSqlStore(context.Background(), dir, key)
+			require.NoError(t, err)
+			t.Cleanup(func() { assert.NoError(t, eventStore.Close(context.Background())) })
+			manager := &DefaultAccountManager{eventStore: eventStore}
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+
+			// The operation already succeeded when shutdown or the request cancels its context.
+			manager.StoreEvent(ctx, activity.SystemInitiator, "domain-id", "account-id",
+				code, map[string]any{"domain": "expired.example.com"})
+			if code != activity.CustomDomainValidationExpired {
+				require.Eventually(t, func() bool {
+					events, err := eventStore.Get(context.Background(), "account-id", 0, 10, true)
+					return err == nil && len(events) == 1
+				}, time.Second, time.Millisecond, "asynchronous events must survive request cancellation")
+			}
+			require.NoError(t, eventStore.Close(context.Background()))
+
+			reopened, err := activitystore.NewSqlStore(context.Background(), dir, key)
+			require.NoError(t, err)
+			t.Cleanup(func() { assert.NoError(t, reopened.Close(context.Background())) })
+			events, err := reopened.Get(context.Background(), "account-id", 0, 10, true)
+			require.NoError(t, err)
+			require.Len(t, events, 1, "the event must be persisted before shutdown closes the store")
+			assert.Equal(t, code, events[0].Activity, "persist the requested activity")
+			assert.Equal(t, "domain-id", events[0].TargetID, "retain the registration ID")
+			assert.Equal(t, "expired.example.com", events[0].Meta["domain"], "retain the domain name")
+		})
+	}
+}
+
 func generateAndStoreEvents(t *testing.T, manager *DefaultAccountManager, typ activity.Activity, initiatorID, targetID,
 	accountID string, count int) {
 	t.Helper()

management/server/group.go

@@ -101,10 +101,8 @@ func (am *DefaultAccountManager) CreateGroup(ctx context.Context, accountID, use
 			return status.Errorf(status.Internal, "failed to create group: %v", err)
 		}
 
-		for _, peerID := range newGroup.Peers {
-			if err := transaction.AddPeerToGroup(ctx, accountID, peerID, newGroup.ID); err != nil {
-				return status.Errorf(status.Internal, "failed to add peer %s to group %s: %v", peerID, newGroup.ID, err)
-			}
+		if err = syncGroupMembership(ctx, transaction, accountID, newGroup.ID, newGroup.Peers, nil); err != nil {
+			return err
 		}
 
 		snap, err = affectedpeers.Load(ctx, transaction, accountID, change)
@@ -200,6 +198,9 @@ func (am *DefaultAccountManager) UpdateGroup(ctx context.Context, accountID, use
 
 // syncGroupMembership applies the peer membership delta for a group within a transaction.
 func syncGroupMembership(ctx context.Context, transaction store.Store, accountID, groupID string, peersToAdd, peersToRemove []string) error {
+	if err := validateGroupPeers(ctx, transaction, accountID, peersToAdd); err != nil {
+		return err
+	}
 	for _, peerID := range peersToAdd {
 		if err := transaction.AddPeerToGroup(ctx, accountID, peerID, groupID); err != nil {
 			return status.Errorf(status.Internal, "failed to add peer %s to group %s: %v", peerID, groupID, err)
@@ -213,6 +214,25 @@ func syncGroupMembership(ctx context.Context, transaction store.Store, accountID
 	return nil
 }
 
+func validateGroupPeers(ctx context.Context, transaction store.Store, accountID string, peerIDs []string) error {
+	if len(peerIDs) == 0 {
+		return nil
+	}
+
+	peers, err := transaction.GetPeersByIDs(ctx, store.LockingStrengthNone, accountID, peerIDs)
+	if err != nil {
+		return err
+	}
+
+	for _, peerID := range peerIDs {
+		if _, ok := peers[peerID]; !ok {
+			return status.Errorf(status.InvalidArgument, "peer with ID %s not found", peerID)
+		}
+	}
+
+	return nil
+}
+
 // CreateGroups adds new groups to the account.
 // Note: This function does not acquire the global lock.
 // It is the caller's responsibility to ensure proper locking is in place before invoking this method.
@@ -540,7 +560,7 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr
 	change := affectedpeers.Change{OutputPeerIDs: []string{peerID}, LinkGroups: []string{groupID}}
 
 	err := am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
-		if err := transaction.AddPeerToGroup(ctx, accountID, peerID, groupID); err != nil {
+		if err := syncGroupMembership(ctx, transaction, accountID, groupID, []string{peerID}, nil); err != nil {
 			return err
 		}
 
@@ -754,6 +774,14 @@ func validateDeleteGroup(ctx context.Context, transaction store.Store, group *ty
 		return &GroupLinkError{"agent network policy", linkedPolicy.Name}
 	}
 
+	isLinked, linkedRule, err := isGroupLinkedToAgentNetworkBudgetRule(ctx, transaction, group.AccountID, group.ID)
+	if err != nil {
+		return status.Errorf(status.Internal, "failed to check agent network budget rules")
+	}
+	if isLinked {
+		return &GroupLinkError{"agent network budget rule", linkedRule.Name}
+	}
+
 	return checkGroupLinkedToSettings(ctx, transaction, group)
 }
 
@@ -925,6 +953,26 @@ func isGroupLinkedToAgentNetworkPolicy(ctx context.Context, transaction store.St
 	return false, nil
 }
 
+// isGroupLinkedToAgentNetworkBudgetRule checks if a group is a target of any
+// account-level agent network budget rule.
+func isGroupLinkedToAgentNetworkBudgetRule(ctx context.Context, transaction store.Store, accountID string, groupID string) (bool, *agentNetworkTypes.AccountBudgetRule, error) {
+	rules, err := transaction.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("error retrieving agent network budget rules while checking group linkage: %v", err)
+		return false, nil, err
+	}
+
+	for _, rule := range rules {
+		if rule == nil {
+			continue
+		}
+		if slices.Contains(rule.TargetGroups, groupID) {
+			return true, rule, nil
+		}
+	}
+	return false, nil, nil
+}
+
 // areGroupChangesAffectPeers checks if any changes to the specified groups will affect peers.
 // It fetches each collection once and checks all groupIDs against them in memory.
 func areGroupChangesAffectPeers(ctx context.Context, transaction store.Store, accountID string, groupIDs []string) (bool, error) {

management/server/group_test.go

@@ -11,10 +11,10 @@ import (
 	"testing"
 	"time"
 
-	"go.uber.org/mock/gomock"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 	"golang.org/x/exp/maps"
 
 	nbdns "github.com/netbirdio/netbird/dns"
@@ -132,6 +132,11 @@ func TestDefaultAccountManager_DeleteGroup(t *testing.T) {
 			"grp-for-agent-network-policy",
 			"agent network policy",
 		},
+		{
+			"agent network budget rule",
+			"grp-for-agent-network-budget-rule",
+			"agent network budget rule",
+		},
 		{
 			"reverse proxy private service access group",
 			"grp-for-rp-private",
@@ -152,6 +157,16 @@ func TestDefaultAccountManager_DeleteGroup(t *testing.T) {
 				return
 			}
 
+			group, getErr := am.GetGroup(context.Background(), account.Id, testCase.groupID, groupAdminUserID)
+			if getErr != nil {
+				t.Errorf("group %s should still exist after failed deletion: %s", testCase.groupID, getErr)
+				return
+			}
+			if group == nil {
+				t.Errorf("group %s was deleted despite the failed deletion", testCase.groupID)
+				return
+			}
+
 			var sErr *status.Error
 			if errors.As(err, &sErr) {
 				if sErr.Message != testCase.expectedReason {
@@ -240,6 +255,12 @@ func TestDefaultAccountManager_DeleteGroups(t *testing.T) {
 			groupIDs:        []string{"grp-for-agent-network-policy"},
 			expectedReasons: []string{"agent network policy"},
 		},
+		{
+			name:               "agent network budget rule",
+			groupIDs:           []string{"grp-for-agent-network-budget-rule"},
+			expectedReasons:    []string{"agent network budget rule"},
+			expectedNotDeleted: []string{"grp-for-agent-network-budget-rule"},
+		},
 		{
 			name:               "reverse proxy services",
 			groupIDs:           []string{"grp-for-rp-private", "grp-for-rp-bearer"},
@@ -501,6 +522,14 @@ func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *t
 		Peers:     make([]string, 0),
 	}
 
+	groupForAgentNetworkBudgetRule := &types.Group{
+		ID:        "grp-for-agent-network-budget-rule",
+		AccountID: "account-id",
+		Name:      "Group for agent network budget rules",
+		Issued:    types.GroupIssuedAPI,
+		Peers:     make([]string, 0),
+	}
+
 	groupForRPPrivate := &types.Group{
 		ID:        "grp-for-rp-private",
 		AccountID: "account-id",
@@ -573,6 +602,7 @@ func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *t
 	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForUsers)
 	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForIntegration)
 	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForAgentNetworkPolicy)
+	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForAgentNetworkBudgetRule)
 	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForRPPrivate)
 	_ = am.CreateGroup(context.Background(), accountID, groupAdminUserID, groupForRPBearer)
 
@@ -587,6 +617,20 @@ func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *t
 		return nil, nil, err
 	}
 
+	budgetRuleDecoy := agentNetworkTypes.NewAccountBudgetRule(accountID)
+	budgetRuleDecoy.Name = "Unrelated agent network budget rule"
+	budgetRuleDecoy.TargetGroups = []string{"unrelated-group"}
+	if err := am.Store.SaveAgentNetworkBudgetRule(context.Background(), budgetRuleDecoy); err != nil {
+		return nil, nil, err
+	}
+
+	budgetRule := agentNetworkTypes.NewAccountBudgetRule(accountID)
+	budgetRule.Name = "Example agent network budget rule"
+	budgetRule.TargetGroups = []string{groupForAgentNetworkBudgetRule.ID}
+	if err := am.Store.SaveAgentNetworkBudgetRule(context.Background(), budgetRule); err != nil {
+		return nil, nil, err
+	}
+
 	// The decoy services are created first so the linkage check has to scan
 	// past services that do not reference the groups under test.
 	rpServices := []*rpservice.Service{
@@ -1236,3 +1280,82 @@ func Test_IncrementNetworkSerial(t *testing.T) {
 
 	assert.Equal(t, totalPeers, int(account.Network.Serial), "Expected %d serial increases in account %s, got %d", totalPeers, accountID, account.Network.Serial)
 }
+
+func TestDefaultAccountManager_GroupPeersMustBelongToAccount(t *testing.T) {
+	manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+
+	otherAccount, err := createAccount(manager, "other_account", "other_user", "")
+	require.NoError(t, err)
+
+	foreignPeer := &peer2.Peer{
+		ID:        "foreign-peer",
+		AccountID: otherAccount.Id,
+		Key:       "foreign-key",
+		DNSLabel:  "foreign-peer",
+		IP:        uint32ToIP(1),
+	}
+	require.NoError(t, manager.Store.AddPeerToAccount(context.Background(), foreignPeer))
+
+	assertRejected := func(t *testing.T, err error) {
+		t.Helper()
+		require.Error(t, err)
+		s, ok := status.FromError(err)
+		require.True(t, ok, "expected status error, got %v", err)
+		assert.Equal(t, status.InvalidArgument, s.Type(), "peer outside the account should be rejected as invalid argument")
+	}
+
+	t.Run("create rejects foreign peer", func(t *testing.T) {
+		err := manager.CreateGroup(context.Background(), account.Id, userID, &types.Group{
+			Name:   "foreign",
+			Issued: types.GroupIssuedAPI,
+			Peers:  []string{peer1.ID, foreignPeer.ID},
+		})
+		assertRejected(t, err)
+
+		_, err = manager.Store.GetGroupByName(context.Background(), store.LockingStrengthNone, account.Id, "foreign")
+		assert.Error(t, err, "rejected create must not persist the group")
+	})
+
+	t.Run("update rejects foreign and unknown peers", func(t *testing.T) {
+		group := &types.Group{ID: "own", Name: "own", Issued: types.GroupIssuedAPI, Peers: []string{peer1.ID}}
+		require.NoError(t, manager.CreateGroup(context.Background(), account.Id, userID, group))
+
+		group.Peers = []string{peer1.ID, foreignPeer.ID}
+		assertRejected(t, manager.UpdateGroup(context.Background(), account.Id, userID, group))
+
+		group.Peers = []string{peer1.ID, "does-not-exist"}
+		assertRejected(t, manager.UpdateGroup(context.Background(), account.Id, userID, group))
+
+		stored, err := manager.Store.GetGroupByID(context.Background(), store.LockingStrengthNone, account.Id, group.ID)
+		require.NoError(t, err)
+		assert.Equal(t, []string{peer1.ID}, stored.Peers, "rejected updates must not change membership")
+	})
+
+	t.Run("update tolerates and drops pre-existing dangling members", func(t *testing.T) {
+		group := &types.Group{ID: "polluted", Name: "polluted", Issued: types.GroupIssuedAPI, Peers: []string{peer1.ID}}
+		require.NoError(t, manager.CreateGroup(context.Background(), account.Id, userID, group))
+		require.NoError(t, manager.Store.AddPeerToGroup(context.Background(), account.Id, foreignPeer.ID, group.ID))
+
+		group.Peers = []string{peer1.ID, foreignPeer.ID}
+		assert.NoError(t, manager.UpdateGroup(context.Background(), account.Id, userID, group), "keeping an existing member must not be rejected")
+
+		group.Peers = []string{peer1.ID}
+		require.NoError(t, manager.UpdateGroup(context.Background(), account.Id, userID, group))
+
+		stored, err := manager.Store.GetGroupByID(context.Background(), store.LockingStrengthNone, account.Id, group.ID)
+		require.NoError(t, err)
+		assert.Equal(t, []string{peer1.ID}, stored.Peers, "dangling member should be removed once omitted")
+	})
+
+	t.Run("direct add rejects foreign and unknown peers", func(t *testing.T) {
+		group := &types.Group{ID: "direct", Name: "direct", Issued: types.GroupIssuedAPI, Peers: []string{peer1.ID}}
+		require.NoError(t, manager.CreateGroup(context.Background(), account.Id, userID, group))
+
+		assertRejected(t, manager.GroupAddPeer(context.Background(), account.Id, group.ID, foreignPeer.ID))
+		assertRejected(t, manager.GroupAddPeer(context.Background(), account.Id, group.ID, "does-not-exist"))
+
+		stored, err := manager.Store.GetGroupByID(context.Background(), store.LockingStrengthNone, account.Id, group.ID)
+		require.NoError(t, err)
+		assert.Equal(t, []string{peer1.ID}, stored.Peers, "rejected direct adds must not change membership")
+	})
+}

management/server/http/handlers/groups/groups_handler.go

@@ -148,13 +148,10 @@ func (h *handler) updateGroup(w http.ResponseWriter, r *http.Request) {
 		peers = *req.Peers
 	}
 
-	resources := make([]types.Resource, 0)
-	if req.Resources != nil {
-		for _, res := range *req.Resources {
-			resource := types.Resource{}
-			resource.FromAPIRequest(&res)
-			resources = append(resources, resource)
-		}
+	resources, err := resourcesFromAPIRequest(req.Resources)
+	if err != nil {
+		util.WriteError(r.Context(), err, w)
+		return
 	}
 
 	group := types.Group{
@@ -210,13 +207,10 @@ func (h *handler) createGroup(w http.ResponseWriter, r *http.Request) {
 		peers = *req.Peers
 	}
 
-	resources := make([]types.Resource, 0)
-	if req.Resources != nil {
-		for _, res := range *req.Resources {
-			resource := types.Resource{}
-			resource.FromAPIRequest(&res)
-			resources = append(resources, resource)
-		}
+	resources, err := resourcesFromAPIRequest(req.Resources)
+	if err != nil {
+		util.WriteError(r.Context(), err, w)
+		return
 	}
 
 	group := types.Group{
@@ -335,11 +329,30 @@ func toGroupResponse(peers []*nbpeer.Peer, group *types.Group) *api.Group {
 	gr.PeersCount = len(gr.Peers)
 
 	for _, res := range group.Resources {
-		resResp := res.ToAPIResponse()
-		gr.Resources = append(gr.Resources, *resResp)
+		if resResp := res.ToAPIResponse(); resResp != nil {
+			gr.Resources = append(gr.Resources, *resResp)
+		}
 	}
 
 	gr.ResourcesCount = len(gr.Resources)
 
 	return &gr
 }
+
+func resourcesFromAPIRequest(req *[]api.Resource) ([]types.Resource, error) {
+	resources := make([]types.Resource, 0)
+	if req == nil {
+		return resources, nil
+	}
+
+	for _, res := range *req {
+		if res.Id == "" || !types.ResourceType(res.Type).Valid() {
+			return nil, status.Errorf(status.InvalidArgument, "resource id shouldn't be empty and type must be one of: peer, domain, host, subnet")
+		}
+		resource := types.Resource{}
+		resource.FromAPIRequest(&res)
+		resources = append(resources, resource)
+	}
+
+	return resources, nil
+}

management/server/http/handlers/groups/groups_handler_test.go

@@ -8,8 +8,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/netip"
 	"net/http/httptest"
+	"net/netip"
 	"strings"
 	"testing"
 
@@ -208,6 +208,33 @@ func TestWriteGroup(t *testing.T) {
 			expectedStatus: http.StatusUnprocessableEntity,
 			expectedBody:   false,
 		},
+		{
+			name:        "Write Group POST Empty Resource",
+			requestType: http.MethodPost,
+			requestPath: "/api/groups",
+			requestBody: bytes.NewBuffer(
+				[]byte(`{"name":"With Resource","resources":[{}]}`)),
+			expectedStatus: http.StatusUnprocessableEntity,
+			expectedBody:   false,
+		},
+		{
+			name:        "Write Group PUT Empty Resource",
+			requestType: http.MethodPut,
+			requestPath: "/api/groups/id-existed",
+			requestBody: bytes.NewBuffer(
+				[]byte(`{"name":"With Resource","resources":[{"id":"","type":"host"}]}`)),
+			expectedStatus: http.StatusUnprocessableEntity,
+			expectedBody:   false,
+		},
+		{
+			name:        "Write Group POST Unknown Resource Type",
+			requestType: http.MethodPost,
+			requestPath: "/api/groups",
+			requestBody: bytes.NewBuffer(
+				[]byte(`{"name":"With Resource","resources":[{"id":"res-1","type":"banana"}]}`)),
+			expectedStatus: http.StatusUnprocessableEntity,
+			expectedBody:   false,
+		},
 		{
 			name:        "Write Group PUT OK",
 			requestType: http.MethodPut,
@@ -376,6 +403,20 @@ func TestGetAllGroups(t *testing.T) {
 	}
 }
 
+func TestToGroupResponseSkipsEmptyResource(t *testing.T) {
+	group := &types.Group{
+		ID:        "id-resources",
+		Name:      "Resources",
+		Issued:    types.GroupIssuedAPI,
+		Resources: []types.Resource{{}, {ID: "res-1", Type: types.ResourceTypeHost}},
+	}
+
+	got := toGroupResponse(nil, group)
+
+	assert.Equal(t, 1, got.ResourcesCount)
+	assert.Equal(t, []api.Resource{{Id: "res-1", Type: api.ResourceType(types.ResourceTypeHost)}}, got.Resources)
+}
+
 func TestDeleteGroup(t *testing.T) {
 	tt := []struct {
 		name           string

management/server/http/middleware/rate_limiter.go

@@ -14,12 +14,14 @@ import (
 	"golang.org/x/time/rate"
 
 	"github.com/netbirdio/netbird/shared/management/http/util"
+	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 const (
-	RateLimitingEnabledEnv = "NB_API_RATE_LIMITING_ENABLED"
-	RateLimitingBurstEnv   = "NB_API_RATE_LIMITING_BURST"
-	RateLimitingRPMEnv     = "NB_API_RATE_LIMITING_RPM"
+	RateLimitingEnabledEnv        = "NB_API_RATE_LIMITING_ENABLED"
+	RateLimitingBurstEnv          = "NB_API_RATE_LIMITING_BURST"
+	RateLimitingRPMEnv            = "NB_API_RATE_LIMITING_RPM"
+	RateLimitingTrustedProxiesEnv = "NB_API_RATE_LIMITING_TRUSTED_PROXIES"
 
 	defaultAPIRPM   = 6
 	defaultAPIBurst = 500
@@ -35,6 +37,9 @@ type RateLimiterConfig struct {
 	CleanupInterval time.Duration
 	// LimiterTTL defines how long a limiter should be kept after last use (age threshold for removal)
 	LimiterTTL time.Duration
+	// TrustedProxies lists the upstream proxies whose forwarding headers may be
+	// believed. Empty means requests are keyed by their direct peer address.
+	TrustedProxies *trustedproxy.List
 }
 
 // DefaultRateLimiterConfig returns a default configuration
@@ -76,11 +81,18 @@ func RateLimiterConfigFromEnv() (cfg *RateLimiterConfig, enabled bool) {
 		burst = defaultAPIBurst
 	}
 
+	trusted, err := trustedproxy.Parse(os.Getenv(RateLimitingTrustedProxiesEnv))
+	if err != nil {
+		log.Warnf("parsing %s env var: %v, trusting no proxies", RateLimitingTrustedProxiesEnv, err)
+		trusted = nil
+	}
+
 	return &RateLimiterConfig{
 		RequestsPerMinute: float64(rpm),
 		Burst:             burst,
 		CleanupInterval:   6 * time.Hour,
 		LimiterTTL:        24 * time.Hour,
+		TrustedProxies:    trusted,
 	}, os.Getenv(RateLimitingEnabledEnv) == "true"
 }
 
@@ -250,7 +262,7 @@ func (rl *APIRateLimiter) Middleware(next http.Handler) http.Handler {
 			next.ServeHTTP(w, r)
 			return
 		}
-		clientIP := getClientIP(r)
+		clientIP := getClientIP(r, rl.config.TrustedProxies)
 		if !rl.Allow(clientIP) {
 			util.WriteErrorResponse("rate limit exceeded, please try again later", http.StatusTooManyRequests, w)
 			return
@@ -259,8 +271,15 @@ func (rl *APIRateLimiter) Middleware(next http.Handler) http.Handler {
 	})
 }
 
-// getClientIP extracts the client IP address from the request.
-func getClientIP(r *http.Request) string {
+// getClientIP extracts the client IP address from the request. Forwarding headers
+// are used only when the request arrives from a trusted proxy.
+func getClientIP(r *http.Request, trusted *trustedproxy.List) string {
+	if !trusted.Empty() {
+		if addr := trusted.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For")); addr.IsValid() {
+			return addr.String()
+		}
+	}
+
 	ip, _, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
 		return r.RemoteAddr

management/server/http/middleware/rate_limiter_test.go

@@ -9,6 +9,9 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 func TestAPIRateLimiter_Allow(t *testing.T) {
@@ -134,7 +137,7 @@ func TestGetClientIP(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			req := httptest.NewRequest(http.MethodGet, "/test", nil)
 			req.RemoteAddr = tc.remoteAddr
-			assert.Equal(t, tc.expected, getClientIP(req))
+			assert.Equal(t, tc.expected, getClientIP(req, nil))
 		})
 	}
 }
@@ -327,3 +330,47 @@ func TestRateLimiterConfigFromEnv(t *testing.T) {
 	assert.Equal(t, float64(defaultAPIRPM), cfg.RequestsPerMinute, "non-positive rpm must fall back to default")
 	assert.Equal(t, defaultAPIBurst, cfg.Burst, "non-positive burst must fall back to default")
 }
+
+func TestGetClientIP_TrustedProxies(t *testing.T) {
+	trusted, err := trustedproxy.Parse("10.0.0.0/8")
+	require.NoError(t, err)
+
+	tests := []struct {
+		name       string
+		list       *trustedproxy.List
+		remoteAddr string
+		xff        string
+		expected   string
+	}{
+		{
+			name:       "no trusted proxies ignores the header",
+			remoteAddr: "10.0.0.1:5555",
+			xff:        "1.1.1.1, 2.2.2.2",
+			expected:   "10.0.0.1",
+		},
+		{
+			name:       "behind a trusted proxy uses the right-most untrusted hop",
+			list:       trusted,
+			remoteAddr: "10.0.0.1:5555",
+			xff:        "1.1.1.1, 2.2.2.2",
+			expected:   "2.2.2.2",
+		},
+		{
+			name:       "a caller reaching us directly cannot forge the header",
+			list:       trusted,
+			remoteAddr: "203.0.113.5:5555",
+			xff:        "1.1.1.1",
+			expected:   "203.0.113.5",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "/test", nil)
+			req.RemoteAddr = tc.remoteAddr
+			req.Header.Set("X-Forwarded-For", tc.xff)
+
+			assert.Equal(t, tc.expected, getClientIP(req, tc.list))
+		})
+	}
+}

management/server/http/testing/integration/events_handler_integration_test.go

@@ -3,6 +3,7 @@
 package integration
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -10,6 +11,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools"
 	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools/channel"
@@ -34,7 +36,7 @@ func Test_Events_GetAll(t *testing.T) {
 
 	for _, user := range users {
 		t.Run(user.name+" - Get all events", func(t *testing.T) {
-			apiHandler, _, _ := channel.BuildApiBlackBoxWithDBState(t, "../testdata/events.sql", nil, false)
+			apiHandler, accountManager, _ := channel.BuildApiBlackBoxWithDBState(t, "../testdata/events.sql", nil, false)
 
 			// First, perform a mutation to generate an event (create a group as admin)
 			groupBody, err := json.Marshal(&api.GroupRequest{Name: "eventTestGroup"})
@@ -44,7 +46,14 @@ func Test_Events_GetAll(t *testing.T) {
 			createReq := testing_tools.BuildRequest(t, groupBody, http.MethodPost, "/api/groups", testing_tools.TestAdminId)
 			createRecorder := httptest.NewRecorder()
 			apiHandler.ServeHTTP(createRecorder, createReq)
-			assert.Equal(t, http.StatusOK, createRecorder.Code, "Failed to create group to generate event")
+			require.Equal(t, http.StatusOK, createRecorder.Code, "Failed to create group to generate event")
+
+			// Group creation returns before its asynchronous audit write finishes.
+			require.EventuallyWithT(t, func(c *assert.CollectT) {
+				events, err := accountManager.GetEvents(context.Background(), testing_tools.TestAccountId, testing_tools.TestAdminId)
+				assert.NoError(c, err)
+				assert.NotEmpty(c, events, "wait for the group creation event before checking permissions")
+			}, time.Second, 10*time.Millisecond)
 
 			// Now query events
 			req := testing_tools.BuildRequest(t, []byte{}, http.MethodGet, "/api/events", user.userId)

management/server/identity_provider.go

@@ -23,6 +23,9 @@ import (
 	"github.com/netbirdio/netbird/shared/management/status"
 )
 
+// maxDiscoveryDocumentSize caps the discovery document read at 1 MiB. Providers serve a few kilobytes.
+const maxDiscoveryDocumentSize = 1 << 20
+
 // oidcProviderJSON represents the OpenID Connect discovery document
 type oidcProviderJSON struct {
 	Issuer string `json:"issuer"`
@@ -35,6 +38,10 @@ func validateOIDCIssuer(ctx context.Context, issuer string) error {
 
 	httpClient := &http.Client{
 		Timeout: 10 * time.Second,
+		// An issuer that redirects its own discovery document is misconfigured.
+		CheckRedirect: func(*http.Request, []*http.Request) error {
+			return http.ErrUseLastResponse
+		},
 	}
 
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, wellKnown, nil)
@@ -48,22 +55,22 @@ func validateOIDCIssuer(ctx context.Context, issuer string) error {
 	}
 	defer resp.Body.Close()
 
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return fmt.Errorf("%w: unable to read response body: %v", types.ErrIdentityProviderIssuerUnreachable, err)
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("%w: %s", types.ErrIdentityProviderIssuerUnreachable, resp.Status)
 	}
 
-	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("%w: %s: %s", types.ErrIdentityProviderIssuerUnreachable, resp.Status, body)
+	body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryDocumentSize+1))
+	if err != nil || len(body) > maxDiscoveryDocumentSize {
+		return fmt.Errorf("%w: failed to decode provider discovery object", types.ErrIdentityProviderIssuerUnreachable)
 	}
 
 	var p oidcProviderJSON
 	if err := json.Unmarshal(body, &p); err != nil {
-		return fmt.Errorf("%w: failed to decode provider discovery object: %v", types.ErrIdentityProviderIssuerUnreachable, err)
+		return fmt.Errorf("%w: failed to decode provider discovery object", types.ErrIdentityProviderIssuerUnreachable)
 	}
 
 	if p.Issuer != issuer {
-		return fmt.Errorf("%w: expected %q got %q", types.ErrIdentityProviderIssuerMismatch, issuer, p.Issuer)
+		return fmt.Errorf("%w: %q", types.ErrIdentityProviderIssuerMismatch, issuer)
 	}
 
 	return nil
@@ -151,15 +158,15 @@ func (am *DefaultAccountManager) CreateIdentityProvider(ctx context.Context, acc
 		return nil, status.NewPermissionDeniedError()
 	}
 
-	if err := validateIdentityProviderConfig(ctx, idpConfig); err != nil {
-		return nil, err
-	}
-
 	embeddedManager, ok := am.idpManager.(*idp.EmbeddedIdPManager)
 	if !ok {
 		return nil, status.Errorf(status.Internal, "identity provider management requires embedded IdP")
 	}
 
+	if err := validateIdentityProviderConfig(ctx, idpConfig); err != nil {
+		return nil, err
+	}
+
 	// Generate ID if not provided
 	if idpConfig.ID == "" {
 		idpConfig.ID = generateIdentityProviderID(idpConfig.Type)
@@ -188,15 +195,15 @@ func (am *DefaultAccountManager) UpdateIdentityProvider(ctx context.Context, acc
 		return nil, status.NewPermissionDeniedError()
 	}
 
-	if err := validateIdentityProviderConfig(ctx, idpConfig); err != nil {
-		return nil, err
-	}
-
 	embeddedManager, ok := am.idpManager.(*idp.EmbeddedIdPManager)
 	if !ok {
 		return nil, status.Errorf(status.Internal, "identity provider management requires embedded IdP")
 	}
 
+	if err := validateIdentityProviderConfig(ctx, idpConfig); err != nil {
+		return nil, err
+	}
+
 	idpConfig.ID = idpID
 	idpConfig.AccountID = accountID
 

management/server/identity_provider_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -121,7 +122,7 @@ func createManagerWithEmbeddedIdPModeAndSetup(
 }
 
 func TestDefaultAccountManager_CreateIdentityProvider_Validation(t *testing.T) {
-	manager, _, err := createManager(t)
+	manager, _, err := createManagerWithEmbeddedIdP(t)
 	require.NoError(t, err)
 
 	userID := "testingUser"
@@ -233,7 +234,7 @@ func TestUpdateUserAuthWithSingleModeKeepsConfiguredDomain(t *testing.T) {
 }
 
 func TestDefaultAccountManager_UpdateIdentityProvider_Validation(t *testing.T) {
-	manager, _, err := createManager(t)
+	manager, _, err := createManagerWithEmbeddedIdP(t)
 	require.NoError(t, err)
 
 	userID := "testingUser"
@@ -355,3 +356,45 @@ func TestValidateOIDCIssuer_TrailingSlash(t *testing.T) {
 	require.Error(t, err)
 	assert.True(t, errors.Is(err, types.ErrIdentityProviderIssuerMismatch))
 }
+
+func TestValidateOIDCIssuer_DoesNotFollowRedirects(t *testing.T) {
+	var reached bool
+	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		reached = true
+		w.WriteHeader(http.StatusForbidden)
+	}))
+	t.Cleanup(target.Close)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, target.URL+"/redirect-target", http.StatusFound)
+	}))
+	t.Cleanup(srv.Close)
+
+	err := validateOIDCIssuer(context.Background(), srv.URL)
+	require.Error(t, err)
+	assert.False(t, reached, "Redirects are not followed")
+}
+
+func TestValidateOIDCIssuer_BoundsResponseSize(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"issuer":"` + strings.Repeat("a", maxDiscoveryDocumentSize) + `"}`))
+	}))
+	t.Cleanup(srv.Close)
+
+	err := validateOIDCIssuer(context.Background(), srv.URL)
+	require.ErrorIs(t, err, types.ErrIdentityProviderIssuerUnreachable)
+	assert.NotErrorIs(t, err, types.ErrIdentityProviderIssuerMismatch)
+}
+
+func TestValidateOIDCIssuer_RejectsTrailingContent(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"issuer":"http://` + r.Host + `"} {"issuer":"second"}`))
+	}))
+	t.Cleanup(srv.Close)
+
+	err := validateOIDCIssuer(context.Background(), srv.URL)
+	require.ErrorIs(t, err, types.ErrIdentityProviderIssuerUnreachable,
+		"Content after the first object is not a valid discovery document")
+}

management/server/idp/auth0.go

@@ -132,13 +132,7 @@ type ConnectionOptions struct {
 
 // NewAuth0Manager creates a new instance of the Auth0Manager
 func NewAuth0Manager(config Auth0ClientConfig, appMetrics telemetry.AppMetrics) (*Auth0Manager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/authentik.go

@@ -49,13 +49,7 @@ type AuthentikCredentials struct {
 
 // NewAuthentikManager creates a new instance of the AuthentikManager.
 func NewAuthentikManager(config AuthentikClientConfig, appMetrics telemetry.AppMetrics) (*AuthentikManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/azure.go

@@ -54,13 +54,7 @@ type azureProfile map[string]any
 
 // NewAzureManager creates a new instance of the AzureManager.
 func NewAzureManager(config AzureClientConfig, appMetrics telemetry.AppMetrics) (*AzureManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/dex.go

@@ -4,10 +4,8 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
-	"net/http"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/dexidp/dex/api/v2"
 	log "github.com/sirupsen/logrus"
@@ -44,13 +42,7 @@ func NewDexManager(config DexClientConfig, appMetrics telemetry.AppMetrics) (*De
 		return nil, fmt.Errorf("dex IdP configuration is incomplete, GRPCAddr is missing")
 	}
 
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   10 * time.Second,
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 	helper := JsonParser{}
 
 	return &DexManager{

management/server/idp/google_workspace.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
-	"net/http"
 
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/oauth2/google"
@@ -44,13 +43,7 @@ func (gc *GoogleWorkspaceCredentials) Authenticate(_ context.Context) (JWTToken,
 
 // NewGoogleWorkspaceManager creates a new instance of the GoogleWorkspaceManager.
 func NewGoogleWorkspaceManager(ctx context.Context, config GoogleWorkspaceClientConfig, appMetrics telemetry.AppMetrics) (*GoogleWorkspaceManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/jumpcloud.go

@@ -58,13 +58,7 @@ type JumpCloudCredentials struct {
 
 // NewJumpCloudManager creates a new instance of the JumpCloudManager.
 func NewJumpCloudManager(config JumpCloudClientConfig, appMetrics telemetry.AppMetrics) (*JumpCloudManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/keycloak.go

@@ -59,13 +59,7 @@ type keycloakProfile struct {
 
 // NewKeycloakManager creates a new instance of the KeycloakManager.
 func NewKeycloakManager(config KeycloakClientConfig, appMetrics telemetry.AppMetrics) (*KeycloakManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/okta.go

@@ -40,13 +40,7 @@ type OktaCredentials struct {
 
 // NewOktaManager creates a new instance of the OktaManager.
 func NewOktaManager(config OktaClientConfig, appMetrics telemetry.AppMetrics) (*OktaManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 	config.Issuer = baseURL(config.Issuer)

management/server/idp/pocketid.go

@@ -83,13 +83,7 @@ type pocketIdUserGroupDto struct {
 }
 
 func NewPocketIdManager(config PocketIdClientConfig, appMetrics telemetry.AppMetrics) (*PocketIdManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/idp/util.go

@@ -2,6 +2,8 @@ package idp
 
 import (
 	"encoding/json"
+	"errors"
+	"net/http"
 	"net/url"
 	"os"
 	"strings"
@@ -81,6 +83,23 @@ const (
 	defaultTimeout = 10 * time.Second
 )
 
+// errRedirectRefused is returned instead of http.ErrUseLastResponse so the
+// client closes the redirect response rather than handing it back unread.
+var errRedirectRefused = errors.New("redirect refused")
+
+func newHTTPClient() *http.Client {
+	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
+	httpTransport.MaxIdleConns = 5
+
+	return &http.Client{
+		Timeout:   idpTimeout(),
+		Transport: httpTransport,
+		CheckRedirect: func(*http.Request, []*http.Request) error {
+			return errRedirectRefused
+		},
+	}
+}
+
 // idpTimeout returns a timeout value for the IDP
 func idpTimeout() time.Duration {
 	timeoutStr, ok := os.LookupEnv(idpTimeoutEnv)

management/server/idp/zitadel.go

@@ -160,13 +160,7 @@ func verifyJWTConfig(config ZitadelClientConfig) error {
 
 // NewZitadelManager creates a new instance of the ZitadelManager.
 func NewZitadelManager(config ZitadelClientConfig, appMetrics telemetry.AppMetrics) (*ZitadelManager, error) {
-	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
-	httpTransport.MaxIdleConns = 5
-
-	httpClient := &http.Client{
-		Timeout:   idpTimeout(),
-		Transport: httpTransport,
-	}
+	httpClient := newHTTPClient()
 
 	helper := JsonParser{}
 

management/server/migration/migration_custom_domain.go

@@ -0,0 +1,22 @@
+package migration
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+)
+
+// MigrateCustomDomainValidationExpiry gives existing pending registrations a validation window.
+func MigrateCustomDomainValidationExpiry(ctx context.Context, db *gorm.DB) error {
+	result := db.WithContext(ctx).Model(&domain.Domain{}).
+		Where("validated = ? AND validation_expires_at IS NULL", false).
+		Update("validation_expires_at", time.Now().UTC().Add(domain.ValidationTTL))
+	if result.Error != nil {
+		return fmt.Errorf("backfill custom domain validation expiry: %w", result.Error)
+	}
+	return nil
+}

management/server/migration/migration_custom_domain_test.go

@@ -0,0 +1,44 @@
+package migration_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	"github.com/netbirdio/netbird/management/server/migration"
+)
+
+func TestMigrateCustomDomainValidationExpiry(t *testing.T) {
+	db := setupDatabase(t)
+	require.NoError(t, db.AutoMigrate(&domain.Domain{}))
+	t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&domain.Domain{})) })
+	ctx := context.Background()
+	existingDeadline := time.Now().UTC().Add(time.Hour).Truncate(time.Second)
+	rows := []domain.Domain{
+		{ID: "legacy", Domain: "legacy.example.com"},
+		{ID: "validated", Domain: "validated.example.com", Validated: true},
+		{ID: "pending", Domain: "pending.example.com", ValidationExpiresAt: &existingDeadline},
+	}
+	require.NoError(t, db.Create(&rows).Error)
+	before := time.Now().UTC()
+	require.NoError(t, migration.MigrateCustomDomainValidationExpiry(ctx, db))
+	after := time.Now().UTC()
+	var migrated domain.Domain
+	require.NoError(t, db.First(&migrated, "id = ?", "legacy").Error)
+	require.NotNil(t, migrated.ValidationExpiresAt)
+	assert.WithinRange(t, *migrated.ValidationExpiresAt, before.Truncate(time.Millisecond).Add(48*time.Hour), after.Add(48*time.Hour+time.Millisecond), "legacy pending registrations get a full window")
+	deadline := *migrated.ValidationExpiresAt
+	require.NoError(t, migration.MigrateCustomDomainValidationExpiry(ctx, db))
+	require.NoError(t, db.First(&migrated, "id = ?", "legacy").Error)
+	assert.Equal(t, deadline, *migrated.ValidationExpiresAt, "repeated migration must not extend the deadline")
+	var validated, pending domain.Domain
+	require.NoError(t, db.First(&validated, "id = ?", "validated").Error)
+	require.NoError(t, db.First(&pending, "id = ?", "pending").Error)
+	assert.Nil(t, validated.ValidationExpiresAt, "validated domains do not acquire an expiry")
+	require.NotNil(t, pending.ValidationExpiresAt)
+	assert.WithinDuration(t, existingDeadline, *pending.ValidationExpiresAt, 0, "existing deadlines must be preserved")
+}

management/server/networks/resources/types/resource.go

@@ -32,7 +32,7 @@ type NetworkResource struct {
 	ID          string `gorm:"primaryKey"`
 	NetworkID   string `gorm:"index"`
 	AccountID   string `gorm:"index"`
-	PublicID    string `json:"-"`
+	PublicID    string `json:"-" gorm:"index"`
 	Name        string
 	Description string
 	Type        NetworkResourceType

management/server/store/sql_store_access_log.go

@@ -0,0 +1,149 @@
+package store
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// CreateAccessLog creates a new access log entry in the database
+func (s *SqlStore) CreateAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error {
+	result := s.db.Create(logEntry)
+	if result.Error != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"service_id": logEntry.ServiceID,
+			"method":     logEntry.Method,
+			"host":       logEntry.Host,
+			"path":       logEntry.Path,
+		}).Errorf("failed to create access log entry in store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to create access log entry in store")
+	}
+	return nil
+}
+
+// GetAccountAccessLogs retrieves access logs for a given account with pagination and filtering
+func (s *SqlStore) GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) {
+	var logs []*accesslogs.AccessLogEntry
+	var totalCount int64
+
+	baseQuery := s.db.
+		Model(&accesslogs.AccessLogEntry{}).
+		Where(accountIDCondition, accountID)
+
+	baseQuery = s.applyAccessLogFilters(baseQuery, filter)
+
+	if err := baseQuery.Count(&totalCount).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to count access logs: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to count access logs")
+	}
+
+	query := s.db.
+		Where(accountIDCondition, accountID)
+
+	query = s.applyAccessLogFilters(query, filter)
+
+	sortColumns := filter.GetSortColumn()
+	sortOrder := strings.ToUpper(filter.GetSortOrder())
+
+	var orderClauses []string
+	for _, col := range strings.Split(sortColumns, ",") {
+		col = strings.TrimSpace(col)
+		if col != "" {
+			orderClauses = append(orderClauses, col+" "+sortOrder)
+		}
+	}
+	orderClause := strings.Join(orderClauses, ", ")
+
+	query = query.
+		Order(orderClause).
+		Limit(filter.GetLimit()).
+		Offset(filter.GetOffset())
+
+	if lockStrength != LockingStrengthNone {
+		query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	result := query.Find(&logs)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get access logs from store: %v", result.Error)
+		return nil, 0, status.Errorf(status.Internal, "failed to get access logs from store")
+	}
+
+	return logs, totalCount, nil
+}
+
+// DeleteOldAccessLogs deletes all access logs older than the specified time
+func (s *SqlStore) DeleteOldAccessLogs(ctx context.Context, olderThan time.Time) (int64, error) {
+	result := s.db.
+		Where("timestamp < ?", olderThan).
+		Delete(&accesslogs.AccessLogEntry{})
+
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
+		return 0, status.Errorf(status.Internal, "failed to delete old access logs")
+	}
+
+	return result.RowsAffected, nil
+}
+
+// applyAccessLogFilters applies filter conditions to the query
+func (s *SqlStore) applyAccessLogFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB {
+	if filter.Search != nil {
+		searchPattern := "%" + *filter.Search + "%"
+		query = query.Where(
+			"id LIKE ? OR location_connection_ip LIKE ? OR host LIKE ? OR path LIKE ? OR CONCAT(host, path) LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
+			searchPattern, searchPattern, searchPattern, searchPattern, searchPattern, searchPattern, searchPattern,
+		)
+	}
+
+	if filter.SourceIP != nil {
+		query = query.Where("location_connection_ip = ?", *filter.SourceIP)
+	}
+
+	if filter.Host != nil {
+		query = query.Where("host = ?", *filter.Host)
+	}
+
+	if filter.Path != nil {
+		// Support LIKE pattern for path filtering
+		query = query.Where("path LIKE ?", "%"+*filter.Path+"%")
+	}
+
+	if filter.UserID != nil {
+		query = query.Where("user_id = ?", *filter.UserID)
+	}
+
+	if filter.Method != nil {
+		query = query.Where("method = ?", *filter.Method)
+	}
+
+	if filter.Status != nil {
+		switch *filter.Status {
+		case "success":
+			query = query.Where("status_code >= ? AND status_code < ?", 200, 400)
+		case "failed":
+			query = query.Where("status_code < ? OR status_code >= ?", 200, 400)
+		}
+	}
+
+	if filter.StatusCode != nil {
+		query = query.Where("status_code = ?", *filter.StatusCode)
+	}
+
+	if filter.StartDate != nil {
+		query = query.Where("timestamp >= ?", *filter.StartDate)
+	}
+
+	if filter.EndDate != nil {
+		query = query.Where("timestamp <= ?", *filter.EndDate)
+	}
+
+	return query
+}

management/server/store/sql_store_account_onboarding.go

@@ -0,0 +1,70 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// GetAccountOnboarding retrieves the onboarding information for a specific account.
+func (s *SqlStore) GetAccountOnboarding(ctx context.Context, accountID string) (*types.AccountOnboarding, error) {
+	var accountOnboarding types.AccountOnboarding
+	result := s.db.Model(&accountOnboarding).Take(&accountOnboarding, accountIDCondition, accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAccountOnboardingNotFoundError(accountID)
+		}
+		log.WithContext(ctx).Errorf("error when getting account onboarding %s from the store: %s", accountID, result.Error)
+		return nil, status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return &accountOnboarding, nil
+}
+
+// SaveAccountOnboarding updates the onboarding information for a specific account.
+func (s *SqlStore) SaveAccountOnboarding(ctx context.Context, onboarding *types.AccountOnboarding) error {
+	result := s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(onboarding)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("error when saving account onboarding %s in the store: %s", onboarding.AccountID, result.Error)
+		return status.Errorf(status.Internal, "error when saving account onboarding %s in the store: %s", onboarding.AccountID, result.Error)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) getAccountOnboarding(ctx context.Context, accountID string, account *types.Account) error {
+	const query = `SELECT account_id, onboarding_flow_pending, signup_form_pending, created_at, updated_at FROM account_onboardings WHERE account_id = $1`
+	var onboardingFlowPending, signupFormPending sql.NullBool
+	var createdAt, updatedAt sql.NullTime
+	err := s.pool.QueryRow(ctx, query, accountID).Scan(
+		&account.Onboarding.AccountID,
+		&onboardingFlowPending,
+		&signupFormPending,
+		&createdAt,
+		&updatedAt,
+	)
+	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
+		return err
+	}
+	if createdAt.Valid {
+		account.Onboarding.CreatedAt = createdAt.Time
+	}
+	if updatedAt.Valid {
+		account.Onboarding.UpdatedAt = updatedAt.Time
+	}
+	if onboardingFlowPending.Valid {
+		account.Onboarding.OnboardingFlowPending = onboardingFlowPending.Bool
+	}
+	if signupFormPending.Valid {
+		account.Onboarding.SignupFormPending = signupFormPending.Bool
+	}
+	return nil
+}

management/server/store/sql_store_account_onboarding_test.go

@@ -0,0 +1,68 @@
+package store
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+func TestSqlStore_GetAccountOnboarding(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "9439-34653001fc3b-bf1c8084-ba50-4ce7"
+	a, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	t.Logf("Onboarding: %+v", a.Onboarding)
+	err = store.SaveAccount(context.Background(), a)
+	require.NoError(t, err)
+	onboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
+	require.NoError(t, err)
+	require.NotNil(t, onboarding)
+	require.Equal(t, accountID, onboarding.AccountID)
+	require.Equal(t, time.Date(2024, time.October, 2, 14, 1, 38, 210000000, time.UTC), onboarding.CreatedAt.UTC())
+}
+
+func TestSqlStore_SaveAccountOnboarding(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+	t.Run("New onboarding should be saved correctly", func(t *testing.T) {
+		accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+		onboarding := &types.AccountOnboarding{
+			AccountID:             accountID,
+			SignupFormPending:     true,
+			OnboardingFlowPending: true,
+		}
+
+		err = store.SaveAccountOnboarding(context.Background(), onboarding)
+		require.NoError(t, err)
+
+		savedOnboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
+		require.NoError(t, err)
+		require.Equal(t, onboarding.SignupFormPending, savedOnboarding.SignupFormPending)
+		require.Equal(t, onboarding.OnboardingFlowPending, savedOnboarding.OnboardingFlowPending)
+	})
+
+	t.Run("Existing onboarding should be updated correctly", func(t *testing.T) {
+		accountID := "9439-34653001fc3b-bf1c8084-ba50-4ce7"
+		onboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
+		require.NoError(t, err)
+
+		onboarding.OnboardingFlowPending = !onboarding.OnboardingFlowPending
+		onboarding.SignupFormPending = !onboarding.SignupFormPending
+
+		err = store.SaveAccountOnboarding(context.Background(), onboarding)
+		require.NoError(t, err)
+
+		savedOnboarding, err := store.GetAccountOnboarding(context.Background(), accountID)
+		require.NoError(t, err)
+		require.Equal(t, onboarding.SignupFormPending, savedOnboarding.SignupFormPending)
+		require.Equal(t, onboarding.OnboardingFlowPending, savedOnboarding.OnboardingFlowPending)
+	})
+}

management/server/store/sql_store_account_test.go

@@ -0,0 +1,965 @@
+package store
+
+import (
+	"context"
+	"encoding/binary"
+	"fmt"
+	"net"
+	"net/netip"
+	"os"
+	"reflect"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+	proxydomain "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
+	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
+	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+	nbroute "github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/shared/management/status"
+	"github.com/netbirdio/netbird/shared/testing_helpers"
+)
+
+func Test_SaveAccount_Large(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		runLargeTest(t, store)
+	})
+}
+
+func runLargeTest(t *testing.T, store Store) {
+	t.Helper()
+
+	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+	groupALL, err := account.GetGroupAll()
+	if err != nil {
+		t.Fatal(err)
+	}
+	setupKey, _ := types.GenerateDefaultSetupKey()
+	account.SetupKeys[setupKey.Key] = setupKey
+	const numPerAccount = 6000
+	for n := 0; n < numPerAccount; n++ {
+		netIP := sequentialIPv4(n)
+		peerID := fmt.Sprintf("%s-peer-%d", account.Id, n)
+		addr, _ := netip.AddrFromSlice(netIP)
+
+		peer := &nbpeer.Peer{
+			ID:         peerID,
+			Key:        peerID,
+			IP:         addr.Unmap(),
+			Name:       peerID,
+			DNSLabel:   peerID,
+			UserID:     "testuser",
+			Status:     &nbpeer.PeerStatus{Connected: false, LastSeen: time.Now()},
+			SSHEnabled: false,
+		}
+		account.Peers[peerID] = peer
+		group, _ := account.GetGroupAll()
+		group.Peers = append(group.Peers, peerID)
+		user := &types.User{
+			Id:        fmt.Sprintf("%s-user-%d", account.Id, n),
+			AccountID: account.Id,
+		}
+		account.Users[user.Id] = user
+		route := &nbroute.Route{
+			ID:          nbroute.ID(fmt.Sprintf("network-id-%d", n)),
+			Description: "base route",
+			NetID:       nbroute.NetID(fmt.Sprintf("network-id-%d", n)),
+			Network:     netip.MustParsePrefix(netIP.String() + "/24"),
+			NetworkType: nbroute.IPv4Network,
+			Metric:      9999,
+			Masquerade:  false,
+			Enabled:     true,
+			Groups:      []string{groupALL.ID},
+		}
+		account.Routes[route.ID] = route
+
+		group = &types.Group{
+			ID:        fmt.Sprintf("group-id-%d", n),
+			AccountID: account.Id,
+			Name:      fmt.Sprintf("group-id-%d", n),
+			Issued:    "api",
+			Peers:     nil,
+		}
+		account.Groups[group.ID] = group
+
+		nameserver := &nbdns.NameServerGroup{
+			ID:                   fmt.Sprintf("nameserver-id-%d", n),
+			AccountID:            account.Id,
+			Name:                 fmt.Sprintf("nameserver-id-%d", n),
+			Description:          "",
+			NameServers:          []nbdns.NameServer{{IP: netip.MustParseAddr(netIP.String()), NSType: nbdns.UDPNameServerType}},
+			Groups:               []string{group.ID},
+			Primary:              false,
+			Domains:              nil,
+			Enabled:              false,
+			SearchDomainsEnabled: false,
+		}
+		account.NameServerGroups[nameserver.ID] = nameserver
+
+		setupKey, _ := types.GenerateDefaultSetupKey()
+		_, exists := account.SetupKeys[setupKey.Key]
+		if exists {
+			t.Errorf("setup key already exists")
+		}
+		account.SetupKeys[setupKey.Key] = setupKey
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	if len(store.GetAllAccounts(context.Background())) != 1 {
+		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
+	}
+
+	a, err := store.GetAccount(context.Background(), account.Id)
+	if a == nil {
+		t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
+	}
+
+	if a != nil && len(a.Policies) != 1 {
+		t.Errorf("expecting Account to have one policy stored after SaveAccount(), got %d", len(a.Policies))
+	}
+
+	if a != nil && len(a.Policies[0].Rules) != 1 {
+		t.Errorf("expecting Account to have one policy rule stored after SaveAccount(), got %d", len(a.Policies[0].Rules))
+		return
+	}
+
+	if a != nil && len(a.Peers) != numPerAccount {
+		t.Errorf("expecting Account to have %d peers stored after SaveAccount(), got %d",
+			numPerAccount, len(a.Peers))
+		return
+	}
+
+	if a != nil && len(a.Users) != numPerAccount+1 {
+		t.Errorf("expecting Account to have %d users stored after SaveAccount(), got %d",
+			numPerAccount+1, len(a.Users))
+		return
+	}
+
+	if a != nil && len(a.Routes) != numPerAccount {
+		t.Errorf("expecting Account to have %d routes stored after SaveAccount(), got %d",
+			numPerAccount, len(a.Routes))
+		return
+	}
+
+	if a != nil && len(a.NameServerGroups) != numPerAccount {
+		t.Errorf("expecting Account to have %d NameServerGroups stored after SaveAccount(), got %d",
+			numPerAccount, len(a.NameServerGroups))
+		return
+	}
+
+	if a != nil && len(a.NameServerGroups) != numPerAccount {
+		t.Errorf("expecting Account to have %d NameServerGroups stored after SaveAccount(), got %d",
+			numPerAccount, len(a.NameServerGroups))
+		return
+	}
+
+	if a != nil && len(a.SetupKeys) != numPerAccount+1 {
+		t.Errorf("expecting Account to have %d SetupKeys stored after SaveAccount(), got %d",
+			numPerAccount+1, len(a.SetupKeys))
+		return
+	}
+}
+
+// sequentialIPv4 returns a unique IPv4 address for the given index, avoiding
+// the random collisions that would otherwise violate the unique (account_id, ip)
+// index when generating a large number of peers.
+func sequentialIPv4(n int) net.IP {
+	b := make([]byte, 4)
+	binary.BigEndian.PutUint32(b, 0x0A000000+uint32(n))
+	return net.IP(b)
+}
+
+func Test_SaveAccount(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+		setupKey, _ := types.GenerateDefaultSetupKey()
+		account.SetupKeys[setupKey.Key] = setupKey
+		account.Peers["testpeer"] = &nbpeer.Peer{
+			Key:    "peerkey",
+			IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+			IPv6:   netip.MustParseAddr("fd00::1"),
+			Meta:   nbpeer.PeerSystemMeta{},
+			Name:   "peer name",
+			Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+		}
+
+		err := store.SaveAccount(context.Background(), account)
+		require.NoError(t, err)
+
+		account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
+		setupKey, _ = types.GenerateDefaultSetupKey()
+		account2.SetupKeys[setupKey.Key] = setupKey
+		account2.Peers["testpeer2"] = &nbpeer.Peer{
+			Key:    "peerkey2",
+			IP:     netip.AddrFrom4([4]byte{127, 0, 0, 2}),
+			IPv6:   netip.MustParseAddr("fd00::2"),
+			Meta:   nbpeer.PeerSystemMeta{},
+			Name:   "peer name 2",
+			Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+		}
+
+		err = store.SaveAccount(context.Background(), account2)
+		require.NoError(t, err)
+
+		if len(store.GetAllAccounts(context.Background())) != 2 {
+			t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
+		}
+
+		a, err := store.GetAccount(context.Background(), account.Id)
+		if a == nil {
+			t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
+		}
+
+		if a != nil && len(a.Policies) != 1 {
+			t.Errorf("expecting Account to have one policy stored after SaveAccount(), got %d", len(a.Policies))
+		}
+
+		if a != nil && len(a.Policies[0].Rules) != 1 {
+			t.Errorf("expecting Account to have one policy rule stored after SaveAccount(), got %d", len(a.Policies[0].Rules))
+			return
+		}
+
+		if a, err := store.GetAccountByPeerPubKey(context.Background(), "peerkey"); a == nil {
+			t.Errorf("expecting PeerKeyID2AccountID index updated after SaveAccount(): %v", err)
+		}
+
+		if a, err := store.GetAccountByUser(context.Background(), "testuser"); a == nil {
+			t.Errorf("expecting UserID2AccountID index updated after SaveAccount(): %v", err)
+		}
+
+		if a, err := store.GetAccountByPeerID(context.Background(), "testpeer"); a == nil {
+			t.Errorf("expecting PeerID2AccountID index updated after SaveAccount(): %v", err)
+		}
+
+		if a, err := store.GetAccountBySetupKey(context.Background(), setupKey.Key); a == nil {
+			t.Errorf("expecting SetupKeyID2AccountID index updated after SaveAccount(): %v", err)
+		}
+	})
+}
+
+func Test_AccountSettings_SaveAndRetrieve(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	populateFields := testing_helpers.NewPopulateFields().WithCustomFieldSetter(
+		reflect.PointerTo(reflect.TypeOf(types.ExtraSettings{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
+			es := types.ExtraSettings{}
+			reflectedEs := reflect.ValueOf(&es).Elem()
+			n, err := this.PopulateAll(reflectedEs)
+			if err != nil {
+				return n, err
+			}
+			field.Set(reflectedEs.Addr())
+			return n, nil
+		}).WithCustomFieldSetter(
+		reflect.PointerTo(reflect.TypeOf(types.DashboardFeatures{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
+			t := true
+			df := types.DashboardFeatures{AgentNetwork: &t}
+			reflectedDf := reflect.ValueOf(&df).Elem()
+			field.Set(reflectedDf.Addr())
+			return 1, nil
+		}).WithSkippedTag("gorm", "-")
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+		setupKey, _ := types.GenerateDefaultSetupKey()
+		account.SetupKeys[setupKey.Key] = setupKey
+
+		settings := types.Settings{}
+		numOfExportedFields, err := populateFields.PopulateAll(reflect.ValueOf(&settings).Elem())
+		assert.NoError(t, err)
+		assert.Equal(t, 28, numOfExportedFields)
+		account.Settings = &settings
+
+		err = store.SaveAccount(context.Background(), account)
+		assert.NoError(t, err)
+
+		accountFromDb, err := store.GetAccount(context.Background(), account.Id)
+		assert.NoError(t, err)
+		assert.NotNil(t, accountFromDb)
+		assert.NotNil(t, accountFromDb.Settings)
+
+		assert.True(t, reflect.DeepEqual(&settings, accountFromDb.Settings), "created settings and settings retrieved from the db should match")
+	})
+}
+
+func TestSqlite_DeleteAccount(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	t.Cleanup(cleanUp)
+	assert.NoError(t, err)
+
+	testUserID := "testuser"
+	user := types.NewAdminUser(testUserID)
+	user.PATs = map[string]*types.PersonalAccessToken{"testtoken": {
+		ID:   "testtoken",
+		Name: "test token",
+	}}
+
+	account := newAccountWithId(context.Background(), "account_id", testUserID, "")
+	setupKey, _ := types.GenerateDefaultSetupKey()
+	account.SetupKeys[setupKey.Key] = setupKey
+	account.Peers["testpeer"] = &nbpeer.Peer{
+		Key:    "peerkey",
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
+		Meta:   nbpeer.PeerSystemMeta{},
+		Name:   "peer name",
+		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+	account.Users[testUserID] = user
+	account.Networks = []*networkTypes.Network{
+		{
+			ID:          "network_id",
+			AccountID:   account.Id,
+			Name:        "network name",
+			Description: "network description",
+		},
+	}
+	account.NetworkRouters = []*routerTypes.NetworkRouter{
+		{
+			ID:         "router_id",
+			NetworkID:  account.Networks[0].ID,
+			AccountID:  account.Id,
+			PeerGroups: []string{"group_id"},
+			Masquerade: true,
+			Metric:     1,
+		},
+	}
+	account.NetworkResources = []*resourceTypes.NetworkResource{
+		{
+			ID:          "resource_id",
+			NetworkID:   account.Networks[0].ID,
+			AccountID:   account.Id,
+			Name:        "Name",
+			Description: "Description",
+			Type:        "Domain",
+			Address:     "example.com",
+		},
+	}
+
+	account.Services = []*rpservice.Service{
+		{
+			ID:        "service_id",
+			AccountID: account.Id,
+			Name:      "test service",
+			Domain:    "svc.example.com",
+			Enabled:   true,
+			Targets: []*rpservice.Target{
+				{
+					AccountID: account.Id,
+					ServiceID: "service_id",
+					Host:      "localhost",
+					Port:      8080,
+					Protocol:  "http",
+					Enabled:   true,
+				},
+			},
+		},
+	}
+
+	account.Domains = []*proxydomain.Domain{
+		{
+			ID:        "domain_id",
+			Domain:    "custom.example.com",
+			AccountID: account.Id,
+			Validated: true,
+		},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	if len(store.GetAllAccounts(context.Background())) != 1 {
+		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
+	}
+
+	o, err := store.GetAccountOnboarding(context.Background(), account.Id)
+	require.NoError(t, err)
+	require.Equal(t, o.AccountID, account.Id)
+
+	err = store.DeleteAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	_, err = store.GetAccountOnboarding(context.Background(), account.Id)
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting onboarding")
+
+	if len(store.GetAllAccounts(context.Background())) != 0 {
+		t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
+	}
+
+	_, err = store.GetAccountByPeerPubKey(context.Background(), "peerkey")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer public key")
+
+	_, err = store.GetAccountByUser(context.Background(), "testuser")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by user")
+
+	_, err = store.GetAccountByPeerID(context.Background(), "testpeer")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer id")
+
+	_, err = store.GetAccountBySetupKey(context.Background(), setupKey.Key)
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by setup key")
+
+	_, err = store.GetAccount(context.Background(), account.Id)
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by id")
+
+	for _, policy := range account.Policies {
+		var rules []*types.PolicyRule
+		err = store.(*SqlStore).db.Model(&types.PolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for policy rules")
+		require.Len(t, rules, 0, "expecting no policy rules to be found after removing DeleteAccount")
+
+	}
+
+	for _, accountUser := range account.Users {
+		var pats []*types.PersonalAccessToken
+		err = store.(*SqlStore).db.Model(&types.PersonalAccessToken{}).Find(&pats, "user_id = ?", accountUser.Id).Error
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for personal access token")
+		require.Len(t, pats, 0, "expecting no personal access token to be found after removing DeleteAccount")
+
+	}
+
+	for _, network := range account.Networks {
+		routers, err := store.GetNetworkRoutersByNetID(context.Background(), LockingStrengthNone, account.Id, network.ID)
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for network routers")
+		require.Len(t, routers, 0, "expecting no network routers to be found after DeleteAccount")
+
+		resources, err := store.GetNetworkResourcesByNetID(context.Background(), LockingStrengthNone, account.Id, network.ID)
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for network resources")
+		require.Len(t, resources, 0, "expecting no network resources to be found after DeleteAccount")
+	}
+
+	domains, err := store.ListCustomDomains(context.Background(), account.Id)
+	require.NoError(t, err, "expecting no error after DeleteAccount when searching for custom domains")
+	require.Len(t, domains, 0, "expecting no custom domains to be found after DeleteAccount")
+
+	var services []*rpservice.Service
+	err = store.(*SqlStore).db.Model(&rpservice.Service{}).Find(&services, "account_id = ?", account.Id).Error
+	require.NoError(t, err, "expecting no error after DeleteAccount when searching for services")
+	require.Len(t, services, 0, "expecting no services to be found after DeleteAccount")
+
+	var targets []*rpservice.Target
+	err = store.(*SqlStore).db.Model(&rpservice.Target{}).Find(&targets, "account_id = ?", account.Id).Error
+	require.NoError(t, err, "expecting no error after DeleteAccount when searching for service targets")
+	require.Len(t, targets, 0, "expecting no service targets to be found after DeleteAccount")
+}
+
+func Test_GetAccount(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		id := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		account, err := store.GetAccount(context.Background(), id)
+		require.NoError(t, err)
+		require.Equal(t, id, account.Id, "account id should match")
+		require.Equal(t, false, account.Onboarding.OnboardingFlowPending)
+
+		id = "9439-34653001fc3b-bf1c8084-ba50-4ce7"
+
+		account, err = store.GetAccount(context.Background(), id)
+		require.NoError(t, err)
+		require.Equal(t, id, account.Id, "account id should match")
+		require.Equal(t, true, account.Onboarding.OnboardingFlowPending)
+
+		_, err = store.GetAccount(context.Background(), "non-existing-account")
+		assert.Error(t, err)
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok)
+		require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+
+	})
+}
+
+func Test_TestGetAccountByPrivateDomain(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		existingDomain := "test.com"
+
+		account, err := store.GetAccountByPrivateDomain(context.Background(), existingDomain)
+		require.NoError(t, err, "should found account")
+		require.Equal(t, existingDomain, account.Domain, "domains should match")
+
+		_, err = store.GetAccountByPrivateDomain(context.Background(), "missing-domain.com")
+		require.Error(t, err, "should return error on domain lookup")
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok)
+		require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+	})
+}
+
+func TestPostgresql_SaveAccount(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	t.Cleanup(cleanUp)
+	assert.NoError(t, err)
+
+	account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+	setupKey, _ := types.GenerateDefaultSetupKey()
+	account.SetupKeys[setupKey.Key] = setupKey
+	account.Peers["testpeer"] = &nbpeer.Peer{
+		Key:    "peerkey",
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
+		Meta:   nbpeer.PeerSystemMeta{},
+		Name:   "peer name",
+		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	account2 := newAccountWithId(context.Background(), "account_id2", "testuser2", "")
+	setupKey, _ = types.GenerateDefaultSetupKey()
+	account2.SetupKeys[setupKey.Key] = setupKey
+	account2.Peers["testpeer2"] = &nbpeer.Peer{
+		Key:    "peerkey2",
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 2}),
+		IPv6:   netip.MustParseAddr("fd00::2"),
+		Meta:   nbpeer.PeerSystemMeta{},
+		Name:   "peer name 2",
+		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+
+	err = store.SaveAccount(context.Background(), account2)
+	require.NoError(t, err)
+
+	if len(store.GetAllAccounts(context.Background())) != 2 {
+		t.Errorf("expecting 2 Accounts to be stored after SaveAccount()")
+	}
+
+	a, err := store.GetAccount(context.Background(), account.Id)
+	if a == nil {
+		t.Errorf("expecting Account to be stored after SaveAccount(): %v", err)
+	}
+
+	if a != nil && len(a.Policies) != 1 {
+		t.Errorf("expecting Account to have one policy stored after SaveAccount(), got %d", len(a.Policies))
+	}
+
+	if a != nil && len(a.Policies[0].Rules) != 1 {
+		t.Errorf("expecting Account to have one policy rule stored after SaveAccount(), got %d", len(a.Policies[0].Rules))
+		return
+	}
+
+	if a, err := store.GetAccountByPeerPubKey(context.Background(), "peerkey"); a == nil {
+		t.Errorf("expecting PeerKeyID2AccountID index updated after SaveAccount(): %v", err)
+	}
+
+	if a, err := store.GetAccountByUser(context.Background(), "testuser"); a == nil {
+		t.Errorf("expecting UserID2AccountID index updated after SaveAccount(): %v", err)
+	}
+
+	if a, err := store.GetAccountByPeerID(context.Background(), "testpeer"); a == nil {
+		t.Errorf("expecting PeerID2AccountID index updated after SaveAccount(): %v", err)
+	}
+
+	if a, err := store.GetAccountBySetupKey(context.Background(), setupKey.Key); a == nil {
+		t.Errorf("expecting SetupKeyID2AccountID index updated after SaveAccount(): %v", err)
+	}
+}
+
+func TestPostgresql_DeleteAccount(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	t.Cleanup(cleanUp)
+	assert.NoError(t, err)
+
+	testUserID := "testuser"
+	user := types.NewAdminUser(testUserID)
+	user.PATs = map[string]*types.PersonalAccessToken{"testtoken": {
+		ID:   "testtoken",
+		Name: "test token",
+	}}
+
+	account := newAccountWithId(context.Background(), "account_id", testUserID, "")
+	setupKey, _ := types.GenerateDefaultSetupKey()
+	account.SetupKeys[setupKey.Key] = setupKey
+	account.Peers["testpeer"] = &nbpeer.Peer{
+		Key:    "peerkey",
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
+		Meta:   nbpeer.PeerSystemMeta{},
+		Name:   "peer name",
+		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+	account.Users[testUserID] = user
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	if len(store.GetAllAccounts(context.Background())) != 1 {
+		t.Errorf("expecting 1 Accounts to be stored after SaveAccount()")
+	}
+
+	err = store.DeleteAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	if len(store.GetAllAccounts(context.Background())) != 0 {
+		t.Errorf("expecting 0 Accounts to be stored after DeleteAccount()")
+	}
+
+	_, err = store.GetAccountByPeerPubKey(context.Background(), "peerkey")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer public key")
+
+	_, err = store.GetAccountByUser(context.Background(), "testuser")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by user")
+
+	_, err = store.GetAccountByPeerID(context.Background(), "testpeer")
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by peer id")
+
+	_, err = store.GetAccountBySetupKey(context.Background(), setupKey.Key)
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by setup key")
+
+	_, err = store.GetAccount(context.Background(), account.Id)
+	require.Error(t, err, "expecting error after removing DeleteAccount when getting account by id")
+
+	for _, policy := range account.Policies {
+		var rules []*types.PolicyRule
+		err = store.(*SqlStore).db.Model(&types.PolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for policy rules")
+		require.Len(t, rules, 0, "expecting no policy rules to be found after removing DeleteAccount")
+
+	}
+
+	for _, accountUser := range account.Users {
+		var pats []*types.PersonalAccessToken
+		err = store.(*SqlStore).db.Model(&types.PersonalAccessToken{}).Find(&pats, "user_id = ?", accountUser.Id).Error
+		require.NoError(t, err, "expecting no error after removing DeleteAccount when searching for personal access token")
+		require.Len(t, pats, 0, "expecting no personal access token to be found after removing DeleteAccount")
+
+	}
+
+}
+
+func TestPostgresql_TestGetAccountByPrivateDomain(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	assert.NoError(t, err)
+
+	existingDomain := "test.com"
+
+	account, err := store.GetAccountByPrivateDomain(context.Background(), existingDomain)
+	require.NoError(t, err, "should found account")
+	require.Equal(t, existingDomain, account.Domain, "domains should match")
+
+	_, err = store.GetAccountByPrivateDomain(context.Background(), "missing-domain.com")
+	require.Error(t, err, "should return error on domain lookup")
+}
+
+func TestSqlite_GetAccountNetwork(t *testing.T) {
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	_, err = store.GetAccount(context.Background(), existingAccountID)
+	require.NoError(t, err)
+
+	network, err := store.GetAccountNetwork(context.Background(), LockingStrengthNone, existingAccountID)
+	require.NoError(t, err)
+	ip := net.IP{100, 64, 0, 0}.To16()
+	assert.Equal(t, ip, network.Net.IP)
+	assert.Equal(t, net.IPMask{255, 255, 0, 0}, network.Net.Mask)
+	assert.Equal(t, "", network.Dns)
+	assert.Equal(t, "af1c8024-ha40-4ce2-9418-34653101fc3c", network.Identifier)
+	assert.Equal(t, uint64(0), network.Serial)
+}
+
+func TestSqlStore_SaveAccountPersistsAgentNetworkOnly(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.False(t, account.Settings.AgentNetworkOnly, "setting should default to false")
+
+	account.Settings.AgentNetworkOnly = true
+	require.NoError(t, store.SaveAccount(context.Background(), account))
+
+	reloaded, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.True(t, reloaded.Settings.AgentNetworkOnly, "setting should survive a save/load round-trip")
+
+	reloaded.Settings.AgentNetworkOnly = false
+	require.NoError(t, store.SaveAccount(context.Background(), reloaded))
+
+	disabled, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.False(t, disabled.Settings.AgentNetworkOnly, "disabling should persist")
+}
+
+func TestSqlStore_SaveAccountPersistsDashboardFeatures(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.Nil(t, account.Settings.DashboardFeatures, "dashboard features should default to unset")
+
+	agentNetwork := true
+	account.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &agentNetwork}
+	require.NoError(t, store.SaveAccount(context.Background(), account))
+
+	reloaded, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.NotNil(t, reloaded.Settings.DashboardFeatures, "dashboard features should survive a save/load round-trip")
+	require.NotNil(t, reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should be set")
+	require.True(t, *reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should persist as true")
+
+	disabled := false
+	reloaded.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &disabled}
+	require.NoError(t, store.SaveAccount(context.Background(), reloaded))
+
+	reloadedDisabled, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.NotNil(t, reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "agent network flag should remain set")
+	require.False(t, *reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "explicit false should persist")
+}
+
+func TestSqlStore_UpdateAccountDomainAttributes(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	if err != nil {
+		t.Fatal(err)
+	}
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	t.Run("Should update attributes with public domain", func(t *testing.T) {
+		require.NoError(t, err)
+		domain := "example.com"
+		category := "public"
+		IsDomainPrimaryAccount := false
+		err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
+		require.NoError(t, err)
+		account, err := store.GetAccount(context.Background(), accountID)
+		require.NoError(t, err)
+		require.Equal(t, domain, account.Domain)
+		require.Equal(t, category, account.DomainCategory)
+		require.Equal(t, IsDomainPrimaryAccount, account.IsDomainPrimaryAccount)
+	})
+
+	t.Run("Should update attributes with private domain", func(t *testing.T) {
+		require.NoError(t, err)
+		domain := "test.com"
+		category := "private"
+		IsDomainPrimaryAccount := true
+		err = store.UpdateAccountDomainAttributes(context.Background(), accountID, domain, category, IsDomainPrimaryAccount)
+		require.NoError(t, err)
+		account, err := store.GetAccount(context.Background(), accountID)
+		require.NoError(t, err)
+		require.Equal(t, domain, account.Domain)
+		require.Equal(t, category, account.DomainCategory)
+		require.Equal(t, IsDomainPrimaryAccount, account.IsDomainPrimaryAccount)
+	})
+
+	t.Run("Should fail when account does not exist", func(t *testing.T) {
+		require.NoError(t, err)
+		domain := "test.com"
+		category := "private"
+		IsDomainPrimaryAccount := true
+		err = store.UpdateAccountDomainAttributes(context.Background(), "non-existing-account-id", domain, category, IsDomainPrimaryAccount)
+		require.Error(t, err)
+	})
+
+}
+
+func TestSqlStore_GetDNSSettings(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name        string
+		accountID   string
+		expectError bool
+	}{
+		{
+			name:        "retrieve existing account dns settings",
+			accountID:   "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectError: false,
+		},
+		{
+			name:        "retrieve non-existing account dns settings",
+			accountID:   "non-existing",
+			expectError: true,
+		},
+		{
+			name:        "retrieve dns settings with empty account ID",
+			accountID:   "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, tt.accountID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, dnsSettings)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, dnsSettings)
+			}
+		})
+	}
+}
+
+func TestSqlStore_SaveDNSSettings(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	dnsSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+
+	dnsSettings.DisabledManagementGroups = []string{"groupA", "groupB"}
+	err = store.SaveDNSSettings(context.Background(), accountID, dnsSettings)
+	require.NoError(t, err)
+
+	saveDNSSettings, err := store.GetAccountDNSSettings(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	require.Equal(t, saveDNSSettings, dnsSettings)
+}
+
+func TestSqlStore_GetAccountCreatedBy(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name        string
+		accountID   string
+		expectError bool
+		createdBy   string
+	}{
+		{
+			name:        "existing account ID",
+			accountID:   "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectError: false,
+			createdBy:   "edafee4e-63fb-11ec-90d6-0242ac120003",
+		},
+		{
+			name:        "non-existing account ID",
+			accountID:   "nonexistent",
+			expectError: true,
+		},
+		{
+			name:        "empty account ID",
+			accountID:   "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			createdBy, err := store.GetAccountCreatedBy(context.Background(), LockingStrengthNone, tt.accountID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Empty(t, createdBy)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, createdBy)
+				require.Equal(t, tt.createdBy, createdBy)
+			}
+		})
+	}
+
+}
+
+func TestSqlStore_GetAccountMeta(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	accountMeta, err := store.GetAccountMeta(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	require.NotNil(t, accountMeta)
+	require.Equal(t, accountID, accountMeta.AccountID)
+	require.Equal(t, "edafee4e-63fb-11ec-90d6-0242ac120003", accountMeta.CreatedBy)
+	require.Equal(t, "test.com", accountMeta.Domain)
+	require.Equal(t, "private", accountMeta.DomainCategory)
+	require.Equal(t, time.Date(2024, time.October, 2, 14, 1, 38, 210000000, time.UTC), accountMeta.CreatedAt.UTC())
+}
+
+func TestSqlStore_GetAnyAccountID(t *testing.T) {
+	t.Run("should return account ID when accounts exist", func(t *testing.T) {
+		store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+		t.Cleanup(cleanup)
+		require.NoError(t, err)
+
+		accountID, err := store.GetAnyAccountID(context.Background())
+		require.NoError(t, err)
+		assert.Equal(t, "bf1c8084-ba50-4ce7-9439-34653001fc3b", accountID)
+	})
+
+	t.Run("should return error when no accounts exist", func(t *testing.T) {
+		store, cleanup, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+		t.Cleanup(cleanup)
+		require.NoError(t, err)
+
+		accountID, err := store.GetAnyAccountID(context.Background())
+		require.Error(t, err)
+		sErr, ok := status.FromError(err)
+		assert.True(t, ok)
+		assert.Equal(t, sErr.Type(), status.NotFound)
+		assert.Empty(t, accountID)
+	})
+}

management/server/store/sql_store_agent_network_access_log.go

@@ -0,0 +1,268 @@
+package store
+
+import (
+	"context"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// CreateAgentNetworkAccessLog persists a flattened agent-network access-log
+// entry together with its authorising-group child rows in a single
+// transaction.
+func (s *SqlStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		// Idempotent on the log id / (log_id, group_id) so a proxy resend of the
+		// same entry can't fail the request.
+		if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(entry).Error; err != nil {
+			return err
+		}
+		if len(groups) > 0 {
+			if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"account_id": entry.AccountID,
+			"service_id": entry.ServiceID,
+			"model":      entry.Model,
+		}).Errorf("failed to create agent-network access log entry in store: %v", err)
+		return status.Errorf(status.Internal, "failed to create agent-network access log entry in store")
+	}
+	return nil
+}
+
+// DeleteOldAgentNetworkAccessLogs deletes an account's access-log rows (and
+// their authorising-group child rows) older than the cutoff. Usage records are
+// untouched — they are the long-term aggregate. Returns the number of log rows
+// deleted.
+func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
+	var deleted int64
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		// Remove group child rows for the soon-to-be-deleted logs first.
+		if err := tx.Exec(
+			"DELETE FROM agent_network_access_log_group WHERE account_id = ? AND log_id IN (SELECT id FROM agent_network_access_log WHERE account_id = ? AND timestamp < ?)",
+			accountID, accountID, olderThan,
+		).Error; err != nil {
+			return err
+		}
+		res := tx.Where("account_id = ? AND timestamp < ?", accountID, olderThan).
+			Delete(&agentNetworkTypes.AgentNetworkAccessLog{})
+		if res.Error != nil {
+			return res.Error
+		}
+		deleted = res.RowsAffected
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to delete old agent-network access logs for account %s: %v", accountID, err)
+		return 0, status.Errorf(status.Internal, "failed to delete old agent-network access logs")
+	}
+	return deleted, nil
+}
+
+// GetAgentNetworkAccessLogs retrieves flattened agent-network access logs for
+// an account with server-side pagination, filtering and sorting. Authorising
+// group ids are hydrated from the group child table for the returned page.
+func (s *SqlStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
+	var logs []*agentNetworkTypes.AgentNetworkAccessLog
+	var totalCount int64
+
+	countQuery := s.applyAgentNetworkAccessLogFilters(
+		s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
+		filter,
+	)
+	if err := countQuery.Count(&totalCount).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to count agent-network access logs: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access logs")
+	}
+
+	query := s.applyAgentNetworkAccessLogFilters(
+		s.db.Where(accountIDCondition, accountID),
+		filter,
+	).
+		Order(filter.GetSortColumn() + " " + filter.GetSortOrder()).
+		Limit(filter.GetLimit()).
+		Offset(filter.GetOffset())
+
+	if lockStrength != LockingStrengthNone {
+		query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	if err := query.Find(&logs).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get agent-network access logs from store: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access logs from store")
+	}
+
+	if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, logs); err != nil {
+		return nil, 0, err
+	}
+
+	return logs, totalCount, nil
+}
+
+// applyAgentNetworkAccessLogFilters applies the filter conditions to a query.
+func (s *SqlStore) applyAgentNetworkAccessLogFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
+	if filter.Search != nil {
+		p := "%" + *filter.Search + "%"
+		query = query.Where(
+			"id LIKE ? OR host LIKE ? OR path LIKE ? OR model LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
+			p, p, p, p, p, p,
+		)
+	}
+	if filter.UserID != nil {
+		query = query.Where("user_id = ?", *filter.UserID)
+	}
+	if filter.SessionID != nil {
+		query = query.Where("session_id = ?", *filter.SessionID)
+	}
+	if filter.Decision != nil {
+		query = query.Where("decision = ?", *filter.Decision)
+	}
+	if filter.PathPrefix != nil {
+		query = query.Where("path LIKE ?", *filter.PathPrefix+"%")
+	}
+	if len(filter.ProviderIDs) > 0 {
+		query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
+	}
+	if len(filter.Models) > 0 {
+		query = query.Where("model IN ?", filter.Models)
+	}
+	if len(filter.GroupIDs) > 0 {
+		query = query.Where(
+			"id IN (SELECT log_id FROM agent_network_access_log_group WHERE group_id IN ?)",
+			filter.GroupIDs,
+		)
+	}
+	if filter.StartDate != nil {
+		query = query.Where("timestamp >= ?", *filter.StartDate)
+	}
+	if filter.EndDate != nil {
+		query = query.Where("timestamp <= ?", *filter.EndDate)
+	}
+	return query
+}
+
+// hydrateAgentNetworkAccessLogGroups loads the authorising group ids for the
+// given page of entries and assigns them onto each entry's GroupIDs field.
+func (s *SqlStore) hydrateAgentNetworkAccessLogGroups(ctx context.Context, accountID string, logs []*agentNetworkTypes.AgentNetworkAccessLog) error {
+	if len(logs) == 0 {
+		return nil
+	}
+
+	ids := make([]string, 0, len(logs))
+	for _, l := range logs {
+		ids = append(ids, l.ID)
+	}
+
+	var rows []agentNetworkTypes.AgentNetworkAccessLogGroup
+	if err := s.db.
+		Where(accountIDCondition, accountID).
+		Where("log_id IN ?", ids).
+		Find(&rows).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to hydrate agent-network access log groups: %v", err)
+		return status.Errorf(status.Internal, "failed to hydrate agent-network access log groups")
+	}
+
+	byLog := make(map[string][]string, len(logs))
+	for _, r := range rows {
+		byLog[r.LogID] = append(byLog[r.LogID], r.GroupID)
+	}
+	for _, l := range logs {
+		l.GroupIDs = byLog[l.ID]
+	}
+	return nil
+}
+
+// agentNetworkSessionKeyExpr is the SQL group key for session-grouped access
+// logs: the row's session id, or — when the client sent none — the row id, so
+// session-less requests each form their own singleton group. COALESCE/NULLIF
+// are standard SQL, so this stays portable across SQLite and Postgres.
+const agentNetworkSessionKeyExpr = "COALESCE(NULLIF(session_id, ''), id)"
+
+// GetAgentNetworkAccessLogSessions retrieves agent-network access logs grouped
+// by session, with server-side pagination, filtering and sorting at the session
+// level. It paginates over the distinct session keys (ordered by the requested
+// session-level aggregate), fetches every entry for the page's sessions, and
+// folds them into per-session summaries. The returned count is the number of
+// matching sessions. Filters apply to the entries, so a session's summary
+// reflects only its filter-matching requests.
+func (s *SqlStore) GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error) {
+	// Count distinct sessions via a grouped subquery — portable and avoids
+	// relying on COUNT(DISTINCT <expr>) quoting quirks.
+	sessionsSubquery := s.applyAgentNetworkAccessLogFilters(
+		s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
+		filter,
+	).
+		Select(agentNetworkSessionKeyExpr + " AS session_key").
+		Group(agentNetworkSessionKeyExpr)
+
+	var totalCount int64
+	if err := s.db.Table("(?) AS sessions", sessionsSubquery).Count(&totalCount).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to count agent-network access-log sessions: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access-log sessions")
+	}
+
+	// The page of session keys, ordered by the session-level aggregate. The
+	// session-key tiebreaker keeps pagination deterministic when the primary
+	// aggregate ties.
+	type sessionKeyRow struct {
+		SessionKey string
+	}
+	var keyRows []sessionKeyRow
+	keyQuery := s.applyAgentNetworkAccessLogFilters(
+		s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
+		filter,
+	).
+		Select(agentNetworkSessionKeyExpr + " AS session_key").
+		Group(agentNetworkSessionKeyExpr).
+		Order(filter.GetSessionSortExpr() + " " + filter.GetSortOrder()).
+		Order("session_key ASC").
+		Limit(filter.GetLimit()).
+		Offset(filter.GetOffset())
+	if err := keyQuery.Scan(&keyRows).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to list agent-network access-log session keys: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to list agent-network access-log session keys")
+	}
+	if len(keyRows) == 0 {
+		return nil, totalCount, nil
+	}
+
+	keys := make([]string, 0, len(keyRows))
+	for _, r := range keyRows {
+		keys = append(keys, r.SessionKey)
+	}
+
+	// All entries for the page's sessions, contiguous per session and oldest
+	// first within each — the fold relies on that ordering.
+	var entries []*agentNetworkTypes.AgentNetworkAccessLog
+	entriesQuery := s.applyAgentNetworkAccessLogFilters(
+		s.db.Where(accountIDCondition, accountID),
+		filter,
+	).
+		Where(agentNetworkSessionKeyExpr+" IN ?", keys).
+		Order(agentNetworkSessionKeyExpr + ", timestamp ASC")
+
+	if lockStrength != LockingStrengthNone {
+		entriesQuery = entriesQuery.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	if err := entriesQuery.Find(&entries).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get agent-network access-log session entries: %v", err)
+		return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access-log session entries")
+	}
+
+	if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, entries); err != nil {
+		return nil, 0, err
+	}
+
+	return agentNetworkTypes.FoldAccessLogSessions(keys, entries), totalCount, nil
+}

management/server/store/sql_store_agent_network_usage.go

@@ -0,0 +1,91 @@
+package store
+
+import (
+	"context"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// CreateAgentNetworkUsage persists a stripped agent-network usage record
+// together with its authorising-group child rows in a single transaction.
+func (s *SqlStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		// Idempotent on the usage id / (usage_id, group_id) so a proxy resend of
+		// the same entry can't fail the request.
+		if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(usage).Error; err != nil {
+			return err
+		}
+		if len(groups) > 0 {
+			if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"account_id": usage.AccountID,
+			"model":      usage.Model,
+		}).Errorf("failed to create agent-network usage record in store: %v", err)
+		return status.Errorf(status.Internal, "failed to create agent-network usage record in store")
+	}
+	return nil
+}
+
+// GetAgentNetworkUsageRows returns the stripped usage rows for an account that
+// match the filter (date / user / group / provider / model). Aggregation into
+// time buckets happens in the manager so granularities stay engine-portable.
+func (s *SqlStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
+	var rows []*agentNetworkTypes.AgentNetworkUsage
+
+	query := s.applyAgentNetworkUsageFilters(
+		s.db.Where(accountIDCondition, accountID),
+		filter,
+	).Order("timestamp ASC")
+
+	if lockStrength != LockingStrengthNone {
+		query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	if err := query.Find(&rows).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get agent-network usage rows from store: %v", err)
+		return nil, status.Errorf(status.Internal, "failed to get agent-network usage rows from store")
+	}
+	return rows, nil
+}
+
+// applyAgentNetworkUsageFilters applies the shared access-log filter's
+// date/user/group/provider/model conditions to a usage-table query. Pagination,
+// sort and free-text search are ignored — the overview is an aggregate.
+func (s *SqlStore) applyAgentNetworkUsageFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
+	if filter.UserID != nil {
+		query = query.Where("user_id = ?", *filter.UserID)
+	}
+	if filter.SessionID != nil {
+		query = query.Where("session_id = ?", *filter.SessionID)
+	}
+	if len(filter.ProviderIDs) > 0 {
+		query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
+	}
+	if len(filter.Models) > 0 {
+		query = query.Where("model IN ?", filter.Models)
+	}
+	if len(filter.GroupIDs) > 0 {
+		query = query.Where(
+			"id IN (SELECT usage_id FROM agent_network_request_usage_group WHERE group_id IN ?)",
+			filter.GroupIDs,
+		)
+	}
+	if filter.StartDate != nil {
+		query = query.Where("timestamp >= ?", *filter.StartDate)
+	}
+	if filter.EndDate != nil {
+		query = query.Where("timestamp <= ?", *filter.EndDate)
+	}
+	return query
+}

management/server/store/sql_store_agentnetwork.go

@@ -315,6 +315,36 @@ func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength
 	return settings, nil
 }
 
+// HasGatewayClusterPinnedByOtherAccount reports whether another account has a
+// labeled agent network gateway pinned beneath host, making host its cluster.
+// A self-addressed endpoint on the very same hostname is not counted: that
+// collision is the domain unique index's to refuse, as a conflict. Case-folded,
+// since a settings row written before hostnames were normalised may carry
+// capitals; one row per account, so the scan is cheap.
+func (s *SqlStore) HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
+	return s.countGatewayRowsByOtherAccount(ctx, "LOWER(proxy_address) = LOWER(?) AND LOWER(domain) <> LOWER(proxy_address)", host, accountID)
+}
+
+// HasGatewayEndpointByOtherAccount reports whether host is another account's
+// agent network endpoint hostname (domain). Case-folded for the same reason as
+// HasGatewayClusterPinnedByOtherAccount.
+func (s *SqlStore) HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
+	return s.countGatewayRowsByOtherAccount(ctx, "LOWER(domain) = LOWER(?)", host, accountID)
+}
+
+func (s *SqlStore) countGatewayRowsByOtherAccount(ctx context.Context, predicate, host, accountID string) (bool, error) {
+	var count int64
+	result := s.db.
+		Model(&agentNetworkTypes.Settings{}).
+		Where(predicate+" AND account_id != ?", host, accountID).
+		Count(&count)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to check agent network gateway claims at host: %v", result.Error)
+		return false, status.Errorf(status.Internal, "check agent network gateway claims at host")
+	}
+	return count > 0, nil
+}
+
 // GetAgentNetworkSettingsByProxyAddress returns every Settings row whose
 // gateway is served by the proxy declaring the given cluster address. Used by
 // cluster-scoped synthesis to find the accounts a shared proxy serves.

management/server/store/sql_store_custom_domain.go

@@ -0,0 +1,146 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// GetCustomDomainsCounts returns the total and validated custom domain counts.
+func (s *SqlStore) GetCustomDomainsCounts(ctx context.Context) (int64, int64, error) {
+	var total, validated int64
+	if err := s.db.Model(&domain.Domain{}).Count(&total).Error; err != nil {
+		return 0, 0, err
+	}
+	if err := s.db.Model(&domain.Domain{}).Where("validated = ?", true).Count(&validated).Error; err != nil {
+		return 0, 0, err
+	}
+	return total, validated, nil
+}
+
+func (s *SqlStore) GetCustomDomain(ctx context.Context, accountID string, domainID string) (*domain.Domain, error) {
+	tx := s.db
+
+	customDomain := &domain.Domain{}
+	result := tx.Take(&customDomain, accountAndIDQueryCondition, accountID, domainID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get custom domain from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
+	}
+
+	return customDomain, nil
+}
+
+func (s *SqlStore) ListFreeDomains(ctx context.Context, accountID string) ([]string, error) {
+	return nil, nil
+}
+
+func (s *SqlStore) ListCustomDomains(ctx context.Context, accountID string) ([]*domain.Domain, error) {
+	tx := s.db
+
+	var domains []*domain.Domain
+	result := tx.Find(&domains, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get reverse proxy custom domains from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get reverse proxy custom domains from store")
+	}
+
+	return domains, nil
+}
+
+// GetCustomDomainByName returns the custom domain row holding the given name,
+// regardless of which account owns it.
+func (s *SqlStore) GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error) {
+	customDomain := &domain.Domain{}
+	result := s.db.Take(customDomain, "domain = ?", domainName)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainName)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get custom domain by name from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
+	}
+
+	return customDomain, nil
+}
+
+func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error) {
+	newDomain := &domain.Domain{
+		ID:            xid.New().String(), // Generate our own ID because gorm doesn't always configure the database to handle this for us.
+		Domain:        domainName,
+		AccountID:     accountID,
+		TargetCluster: targetCluster,
+		Type:          domain.TypeCustom,
+		Validated:     validated,
+	}
+	if !validated {
+		expiresAt := time.Now().UTC().Add(domain.ValidationTTL)
+		newDomain.ValidationExpiresAt = &expiresAt
+	}
+	result := s.db.Create(newDomain)
+	if result.Error != nil {
+		// The unique index is the last guard when two requests clear the
+		// manager's availability check at the same time. The one that loses the
+		// insert is a conflict, not an internal failure.
+		var count int64
+		if err := s.db.Model(&domain.Domain{}).Where("domain = ?", domainName).Count(&count).Error; err == nil && count > 0 {
+			// The insert error is logged even on this path: the name being taken
+			// is what the caller has to act on, but if the insert also failed for
+			// an unrelated reason the operator still needs to see it.
+			log.WithContext(ctx).Warnf("create reverse proxy custom domain %s rejected, name already registered: %v", domainName, result.Error)
+			return nil, status.Errorf(status.AlreadyExists, "domain %s is already registered", domainName)
+		}
+
+		log.WithContext(ctx).Errorf("failed to create reverse proxy custom domain to store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to create reverse proxy custom domain to store")
+	}
+
+	return newDomain, nil
+}
+
+// UpdateCustomDomain completes validation only while the original registration is pending.
+func (s *SqlStore) UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error) {
+	if !d.Validated {
+		return nil, status.Errorf(status.InvalidArgument, "custom domain update must complete validation")
+	}
+	result := s.db.WithContext(ctx).Model(&domain.Domain{}).
+		Where(accountAndIDQueryCondition, accountID, d.ID).
+		Where("domain = ? AND target_cluster = ?", d.Domain, d.TargetCluster).
+		Where("validated = ? AND validation_expires_at > ?", false, time.Now().UTC()).
+		Update("validated", true)
+	if result.Error != nil {
+		return nil, fmt.Errorf("validate custom domain in store: %w", result.Error)
+	}
+	if result.RowsAffected == 0 {
+		return nil, status.Errorf(status.PreconditionFailed, "custom domain registration is no longer pending validation")
+	}
+
+	return d, nil
+}
+
+func (s *SqlStore) DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error {
+	result := s.db.Delete(domain.Domain{}, accountAndIDQueryCondition, accountID, domainID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete reverse proxy custom domain from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete reverse proxy custom domain from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.Errorf(status.NotFound, "reverse proxy custom domain %s not found", domainID)
+	}
+
+	return nil
+}

management/server/store/sql_store_dns_record.go

@@ -0,0 +1,109 @@
+package store
+
+import (
+	"context"
+	"errors"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/netbirdio/netbird/management/internals/modules/zones/records"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) CreateDNSRecord(ctx context.Context, record *records.Record) error {
+	result := s.db.Create(record)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to create dns record to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to create dns record to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) UpdateDNSRecord(ctx context.Context, record *records.Record) error {
+	result := s.db.Select("*").Save(record)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to update dns record to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to update dns record to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteDNSRecord(ctx context.Context, accountID, zoneID, recordID string) error {
+	result := s.db.Delete(&records.Record{}, "account_id = ? AND zone_id = ? AND id = ?", accountID, zoneID, recordID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete dns record from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete dns record from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewDNSRecordNotFoundError(recordID)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) GetDNSRecordByID(ctx context.Context, lockStrength LockingStrength, accountID, zoneID, recordID string) (*records.Record, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var record *records.Record
+	result := tx.Where("account_id = ? AND zone_id = ? AND id = ?", accountID, zoneID, recordID).Take(&record)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewDNSRecordNotFoundError(recordID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get dns record from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get dns record from store")
+	}
+
+	return record, nil
+}
+
+func (s *SqlStore) GetZoneDNSRecords(ctx context.Context, lockStrength LockingStrength, accountID, zoneID string) ([]*records.Record, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var recordsList []*records.Record
+	result := tx.Where("account_id = ? AND zone_id = ?", accountID, zoneID).Find(&recordsList)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get zone dns records from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get zone dns records from store")
+	}
+
+	return recordsList, nil
+}
+
+func (s *SqlStore) GetZoneDNSRecordsByName(ctx context.Context, lockStrength LockingStrength, accountID, zoneID, name string) ([]*records.Record, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var recordsList []*records.Record
+	result := tx.Where("account_id = ? AND zone_id = ? AND name = ?", accountID, zoneID, name).Find(&recordsList)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get zone dns records by name from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get zone dns records by name from store")
+	}
+
+	return recordsList, nil
+}
+
+func (s *SqlStore) DeleteZoneDNSRecords(ctx context.Context, accountID, zoneID string) error {
+	result := s.db.Delete(&records.Record{}, "account_id = ? AND zone_id = ?", accountID, zoneID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete zone dns records from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete zone dns records from store")
+	}
+
+	return nil
+}

management/server/store/sql_store_dns_record_test.go

@@ -0,0 +1,260 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/zones"
+	"github.com/netbirdio/netbird/management/internals/modules/zones/records"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlStore_CreateDNSRecord(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+
+	err = store.CreateDNSRecord(context.Background(), record)
+	require.NoError(t, err)
+
+	savedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
+	require.NoError(t, err)
+	require.NotNil(t, savedRecord)
+	assert.Equal(t, record.ID, savedRecord.ID)
+	assert.Equal(t, record.Name, savedRecord.Name)
+	assert.Equal(t, record.Type, savedRecord.Type)
+	assert.Equal(t, record.Content, savedRecord.Content)
+	assert.Equal(t, record.TTL, savedRecord.TTL)
+	assert.Equal(t, zone.ID, savedRecord.ZoneID)
+}
+
+func TestSqlStore_GetDNSRecordByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), record)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name        string
+		accountID   string
+		zoneID      string
+		recordID    string
+		expectError bool
+	}{
+		{
+			name:        "retrieve existing record",
+			accountID:   accountID,
+			zoneID:      zone.ID,
+			recordID:    record.ID,
+			expectError: false,
+		},
+		{
+			name:        "retrieve non-existing record",
+			accountID:   accountID,
+			zoneID:      zone.ID,
+			recordID:    "non-existing",
+			expectError: true,
+		},
+		{
+			name:        "retrieve with empty record ID",
+			accountID:   accountID,
+			zoneID:      zone.ID,
+			recordID:    "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			savedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, tt.accountID, tt.zoneID, tt.recordID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, savedRecord)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, savedRecord)
+				assert.Equal(t, tt.recordID, savedRecord.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_GetZoneDNSRecords(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	recordA := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), recordA)
+	require.NoError(t, err)
+
+	recordAAAA := records.NewRecord(accountID, zone.ID, "ipv6.example.com", records.RecordTypeAAAA, "2001:db8::1", 300)
+	err = store.CreateDNSRecord(context.Background(), recordAAAA)
+	require.NoError(t, err)
+
+	recordCNAME := records.NewRecord(accountID, zone.ID, "alias.example.com", records.RecordTypeCNAME, "www.example.com", 300)
+	err = store.CreateDNSRecord(context.Background(), recordCNAME)
+	require.NoError(t, err)
+
+	allRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
+	require.NoError(t, err)
+	require.NotNil(t, allRecords)
+	assert.Equal(t, 3, len(allRecords))
+
+	recordIDs := make(map[string]bool)
+	for _, r := range allRecords {
+		recordIDs[r.ID] = true
+	}
+	assert.True(t, recordIDs[recordA.ID])
+	assert.True(t, recordIDs[recordAAAA.ID])
+	assert.True(t, recordIDs[recordCNAME.ID])
+}
+
+func TestSqlStore_GetZoneDNSRecordsByName(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record1 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), record1)
+	require.NoError(t, err)
+
+	record2 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeAAAA, "2001:db8::1", 300)
+	err = store.CreateDNSRecord(context.Background(), record2)
+	require.NoError(t, err)
+
+	record3 := records.NewRecord(accountID, zone.ID, "mail.example.com", records.RecordTypeA, "192.168.1.2", 600)
+	err = store.CreateDNSRecord(context.Background(), record3)
+	require.NoError(t, err)
+
+	recordsByName, err := store.GetZoneDNSRecordsByName(context.Background(), LockingStrengthNone, accountID, zone.ID, "www.example.com")
+	require.NoError(t, err)
+	require.NotNil(t, recordsByName)
+	assert.Equal(t, 2, len(recordsByName))
+
+	for _, r := range recordsByName {
+		assert.Equal(t, "www.example.com", r.Name)
+	}
+}
+
+func TestSqlStore_UpdateDNSRecord(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), record)
+	require.NoError(t, err)
+
+	record.Name = "api.example.com"
+	record.Content = "192.168.1.100"
+	record.TTL = 600
+
+	err = store.UpdateDNSRecord(context.Background(), record)
+	require.NoError(t, err)
+
+	updatedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
+	require.NoError(t, err)
+	require.NotNil(t, updatedRecord)
+	assert.Equal(t, "api.example.com", updatedRecord.Name)
+	assert.Equal(t, "192.168.1.100", updatedRecord.Content)
+	assert.Equal(t, 600, updatedRecord.TTL)
+}
+
+func TestSqlStore_DeleteDNSRecord(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), record)
+	require.NoError(t, err)
+
+	err = store.DeleteDNSRecord(context.Background(), accountID, zone.ID, record.ID)
+	require.NoError(t, err)
+
+	deletedRecord, err := store.GetDNSRecordByID(context.Background(), LockingStrengthNone, accountID, zone.ID, record.ID)
+	require.Error(t, err)
+	require.Nil(t, deletedRecord)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, sErr.Type(), status.NotFound)
+}
+
+func TestSqlStore_DeleteZoneDNSRecords(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	zone := zones.NewZone(accountID, "Test Zone", "example.com", true, false, []string{"group1"})
+	err = store.CreateZone(context.Background(), zone)
+	require.NoError(t, err)
+
+	record1 := records.NewRecord(accountID, zone.ID, "www.example.com", records.RecordTypeA, "192.168.1.1", 300)
+	err = store.CreateDNSRecord(context.Background(), record1)
+	require.NoError(t, err)
+
+	record2 := records.NewRecord(accountID, zone.ID, "mail.example.com", records.RecordTypeA, "192.168.1.2", 600)
+	err = store.CreateDNSRecord(context.Background(), record2)
+	require.NoError(t, err)
+
+	allRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
+	require.NoError(t, err)
+	assert.Equal(t, 2, len(allRecords))
+
+	err = store.DeleteZoneDNSRecords(context.Background(), accountID, zone.ID)
+	require.NoError(t, err)
+
+	remainingRecords, err := store.GetZoneDNSRecords(context.Background(), LockingStrengthNone, accountID, zone.ID)
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(remainingRecords))
+}

management/server/store/sql_store_domain_expiration.go

@@ -0,0 +1,60 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// GetExpiredCustomDomains lists pending registrations in stable batches across accounts.
+func (s *SqlStore) GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error) {
+	var domains []*domain.Domain
+	result := s.db.WithContext(ctx).
+		Where("validated = ? AND validation_expires_at <= ? AND id > ?", false, now, string(afterID)).
+		Order("id").Limit(limit).Find(&domains)
+	if result.Error != nil {
+		return nil, fmt.Errorf("list expired custom domains: %w", result.Error)
+	}
+	return domains, nil
+}
+
+// DeleteExpiredCustomDomain deletes an expired registration only if no service uses its namespace.
+func (s *SqlStore) DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error) {
+	db := s.db.WithContext(ctx)
+	services := customDomainServices(db, d)
+	result := db.Where(accountAndIDQueryCondition, d.AccountID, d.ID).
+		Where("domain = ? AND validated = ? AND validation_expires_at <= ?", d.Domain, false, now).
+		Where("NOT EXISTS (?)", services.Select("1")).Delete(&domain.Domain{})
+	if result.Error != nil {
+		return false, fmt.Errorf("delete expired custom domain: %w", result.Error)
+	}
+	if result.RowsAffected > 0 {
+		return true, nil
+	}
+	var count int64
+	if err := customDomainServices(db, d).Count(&count).Error; err != nil {
+		return false, fmt.Errorf("check expired custom domain services: %w", err)
+	}
+	if count > 0 {
+		return false, status.Errorf(status.PreconditionFailed, "expired custom domain still has dependent services")
+	}
+	return false, nil
+}
+
+func customDomainServices(db *gorm.DB, d *domain.Domain) *gorm.DB {
+	name := strings.ToLower(strings.TrimSuffix(d.Domain, "."))
+	// Shared domain validation permits underscores, and older rows may contain
+	// other LIKE metacharacters.
+	escaped := strings.NewReplacer("!", "!!", "%", "!%", "_", "!_").Replace(name)
+	return db.Model(&rpservice.Service{}).Where(
+		"LOWER(domain) IN ? OR LOWER(domain) LIKE ? ESCAPE '!' OR LOWER(domain) LIKE ? ESCAPE '!'",
+		[]string{name, name + "."}, "%."+escaped, "%."+escaped+".",
+	)
+}

management/server/store/sql_store_domain_expiration_test.go

@@ -0,0 +1,78 @@
+package store
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+)
+
+func TestDeleteExpiredCustomDomain_ServiceDependencies(t *testing.T) {
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		now := time.Now().UTC()
+		db := store.(*SqlStore).db
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
+		for _, tt := range []struct {
+			name        string
+			domainName  string
+			serviceHost string
+			protected   bool
+		}{
+			{"exact", "example.com", "example.com", true},
+			{"subdomain", "example.com", "deep.app.example.com", true},
+			{"case", "example.com", "APP.EXAMPLE.COM.", true},
+			{"suffix-boundary", "example.com", "notexample.com", false},
+			{"literal underscore", "a_b.example.com", "app.a_b.example.com", true},
+			{"underscore wildcard", "a_b.example.com", "app.axb.example.com", false},
+			{"legacy percent wildcard", "a%b.example.com", "app.axxb.example.com", false},
+			{"legacy escape character", "a!b.example.com", "app.ab.example.com", false},
+		} {
+			t.Run(tt.name, func(t *testing.T) {
+				d, err := store.CreateCustomDomain(ctx, "owner", tt.domainName, "cluster", false)
+				require.NoError(t, err)
+				require.NoError(t, db.Model(d).Update("validation_expires_at", now.Add(-time.Hour)).Error)
+				svc := &rpservice.Service{ID: "legacy", AccountID: "owner", Domain: tt.serviceHost}
+				require.NoError(t, store.CreateService(ctx, svc))
+				deleted, err := store.DeleteExpiredCustomDomain(ctx, d, now)
+				if tt.protected {
+					require.Error(t, err)
+					assert.False(t, deleted, "service namespaces must remain reserved")
+				} else {
+					require.NoError(t, err)
+					assert.True(t, deleted, "a hostname outside the namespace must not prevent cleanup")
+				}
+				require.NoError(t, db.Delete(svc).Error)
+				require.NoError(t, db.Delete(d).Error)
+			})
+		}
+	})
+}
+
+func TestDeleteExpiredCustomDomain_RechecksValidation(t *testing.T) {
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
+		d, err := store.CreateCustomDomain(ctx, "owner", "validated.example.com", "cluster", false)
+		require.NoError(t, err)
+		d, err = store.GetCustomDomain(ctx, "owner", d.ID)
+		require.NoError(t, err)
+		stale := d.Copy()
+		d.Validated = true
+		_, err = store.UpdateCustomDomain(ctx, "owner", d)
+		require.NoError(t, err)
+		deleted, err := store.DeleteExpiredCustomDomain(ctx, stale, time.Now().Add(domain.ValidationTTL))
+		require.NoError(t, err)
+		assert.False(t, deleted, "a stale cleanup candidate must not delete a validated registration")
+		stored, err := store.GetCustomDomain(ctx, "owner", d.ID)
+		require.NoError(t, err)
+		assert.True(t, stored.Validated, "the validated registration must remain usable")
+		require.NotNil(t, stored.ValidationExpiresAt)
+		assert.Equal(t, stale.ValidationExpiresAt, stored.ValidationExpiresAt, "validation must preserve the original deadline")
+	})
+}

management/server/store/sql_store_get_account_test.go

@@ -64,7 +64,7 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
 
 	_, err := store.CreateCustomDomain(ctx, accountID, "example.com", "eu.proxy.netbird.io", true)
 	require.NoError(t, err, "creating the first custom domain must succeed")
-	_, err = store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
+	pending, err := store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
 	require.NoError(t, err, "creating the second custom domain must succeed")
 
 	account, err := store.GetAccount(ctx, accountID)
@@ -75,6 +75,10 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
 	for _, d := range account.Domains {
 		require.NotNil(t, d)
 		byDomain[d.Domain] = d.TargetCluster
+		if d.ID == pending.ID {
+			require.NotNil(t, d.ValidationExpiresAt)
+			assert.WithinDuration(t, *pending.ValidationExpiresAt, *d.ValidationExpiresAt, time.Millisecond, "both account loaders must preserve the validation deadline")
+		}
 	}
 	assert.Equal(t, "eu.proxy.netbird.io", byDomain["example.com"], "custom domain must carry its target cluster")
 	assert.Equal(t, "us.proxy.netbird.io", byDomain["apps.acme.io"], "custom domain must carry its target cluster")

management/server/store/sql_store_group.go

@@ -0,0 +1,356 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// CreateGroups creates the given list of groups to the database.
+// groupUpsertColumns is the explicit allowlist of columns that get updated when
+// CreateGroups / UpdateGroups hit a PK conflict. public_id is intentionally
+// omitted so a caller passing an entity with the zero value (e.g. an HTTP
+// handler-built struct) cannot reset the persisted public_id during an upsert.
+// Keep this in sync with the Group schema in management/server/types/group.go.
+func groupUpsertColumns() clause.Set {
+	return clause.AssignmentColumns([]string{
+		"account_id",
+		"name",
+		"issued",
+		"integration_ref_id",
+		"integration_ref_integration_type",
+		"resources",
+	})
+}
+
+func (s *SqlStore) CreateGroups(ctx context.Context, accountID string, groups []*types.Group) error {
+	if len(groups) == 0 {
+		return nil
+	}
+
+	return s.db.Transaction(func(tx *gorm.DB) error {
+		result := tx.
+			Clauses(
+				clause.OnConflict{
+					Columns:   []clause.Column{{Name: "id"}},
+					Where:     clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
+					DoUpdates: groupUpsertColumns(),
+				},
+			).
+			Omit(clause.Associations).
+			Create(&groups)
+		if result.Error != nil {
+			log.WithContext(ctx).Errorf("failed to save groups to store: %v", result.Error)
+			return status.Errorf(status.Internal, "failed to save groups to store")
+		}
+
+		return nil
+	})
+}
+
+// UpdateGroups updates the given list of groups to the database.
+func (s *SqlStore) UpdateGroups(ctx context.Context, accountID string, groups []*types.Group) error {
+	if len(groups) == 0 {
+		return nil
+	}
+
+	return s.db.Transaction(func(tx *gorm.DB) error {
+		result := tx.
+			Clauses(
+				clause.OnConflict{
+					Columns:   []clause.Column{{Name: "id"}},
+					Where:     clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
+					DoUpdates: groupUpsertColumns(),
+				},
+			).
+			Omit(clause.Associations).
+			Create(&groups)
+		if result.Error != nil {
+			log.WithContext(ctx).Errorf("failed to save groups to store: %v", result.Error)
+			return status.Errorf(status.Internal, "failed to save groups to store")
+		}
+
+		return nil
+	})
+}
+
+func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.Group, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var groups []*types.Group
+	result := tx.Preload(clause.Associations).Find(&groups, accountIDCondition, accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
+		}
+		log.WithContext(ctx).Errorf("failed to get account groups from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get account groups from the store")
+	}
+
+	for _, g := range groups {
+		g.LoadGroupPeers()
+	}
+
+	return groups, nil
+}
+
+func (s *SqlStore) GetResourceGroups(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) ([]*types.Group, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var groups []*types.Group
+
+	likePattern := `%"ID":"` + resourceID + `"%`
+
+	result := tx.
+		Preload(clause.Associations).
+		Where("resources LIKE ?", likePattern).
+		Find(&groups)
+
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, nil
+		}
+		return nil, result.Error
+	}
+
+	for _, g := range groups {
+		g.LoadGroupPeers()
+	}
+
+	return groups, nil
+}
+
+func (s *SqlStore) getGroups(ctx context.Context, accountID string) ([]*types.Group, error) {
+	const query = `SELECT id, account_id, public_id, name, issued, resources, integration_ref_id, integration_ref_integration_type FROM groups WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+	groups, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*types.Group, error) {
+		var g types.Group
+		var resources []byte
+		var refID sql.NullInt64
+		var refType sql.NullString
+		err := row.Scan(&g.ID, &g.AccountID, &g.PublicID, &g.Name, &g.Issued, &resources, &refID, &refType)
+		if err == nil {
+			if refID.Valid {
+				g.IntegrationReference.ID = int(refID.Int64)
+			}
+			if refType.Valid {
+				g.IntegrationReference.IntegrationType = refType.String
+			}
+			if resources != nil {
+				_ = json.Unmarshal(resources, &g.Resources)
+			} else {
+				g.Resources = []types.Resource{}
+			}
+			g.GroupPeers = []types.GroupPeer{}
+			g.Peers = []string{}
+		}
+		return &g, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return groups, nil
+}
+
+// AddResourceToGroup adds a resource to a group. Method always needs to run n a transaction
+func (s *SqlStore) AddResourceToGroup(ctx context.Context, accountId string, groupID string, resource *types.Resource) error {
+	var group types.Group
+	result := s.db.Where(accountAndIDQueryCondition, accountId, groupID).Take(&group)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return status.NewGroupNotFoundError(groupID)
+		}
+
+		return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
+	}
+
+	for _, res := range group.Resources {
+		if res.ID == resource.ID {
+			return nil
+		}
+	}
+
+	group.Resources = append(group.Resources, *resource)
+
+	if err := s.db.Save(&group).Error; err != nil {
+		return status.Errorf(status.Internal, "issue updating group: %s", err)
+	}
+
+	return nil
+}
+
+// RemoveResourceFromGroup removes a resource from a group. Method always needs to run in a transaction
+func (s *SqlStore) RemoveResourceFromGroup(ctx context.Context, accountId string, groupID string, resourceID string) error {
+	var group types.Group
+	result := s.db.Where(accountAndIDQueryCondition, accountId, groupID).Take(&group)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return status.NewGroupNotFoundError(groupID)
+		}
+
+		return status.Errorf(status.Internal, "issue finding group: %s", result.Error)
+	}
+
+	for i, res := range group.Resources {
+		if res.ID == resourceID {
+			group.Resources = append(group.Resources[:i], group.Resources[i+1:]...)
+			break
+		}
+	}
+
+	if err := s.db.Save(&group).Error; err != nil {
+		return status.Errorf(status.Internal, "issue updating group: %s", err)
+	}
+
+	return nil
+}
+
+// GetGroupByID retrieves a group by ID and account ID.
+func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrength, accountID, groupID string) (*types.Group, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var group *types.Group
+	result := tx.Preload(clause.Associations).Take(&group, accountAndIDQueryCondition, accountID, groupID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.NewGroupNotFoundError(groupID)
+		}
+		log.WithContext(ctx).Errorf("failed to get group from store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get group from store")
+	}
+
+	group.LoadGroupPeers()
+
+	return group, nil
+}
+
+// GetGroupByName retrieves a group by name and account ID.
+func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStrength, accountID, groupName string) (*types.Group, error) {
+	tx := s.db
+
+	var group types.Group
+
+	// TODO: This fix is accepted for now, but if we need to handle this more frequently
+	// we may need to reconsider changing the types.
+	query := tx.Preload(clause.Associations)
+
+	result := query.
+		Model(&types.Group{}).
+		Joins("LEFT JOIN group_peers ON group_peers.group_id = groups.id").
+		Where("groups.account_id = ? AND groups.name = ?", accountID, groupName).
+		Group("groups.id").
+		Order("COUNT(group_peers.peer_id) DESC").
+		Limit(1).
+		First(&group)
+	if err := result.Error; err != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewGroupNotFoundError(groupName)
+		}
+		log.WithContext(ctx).Errorf("failed to get group by name from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get group by name from store")
+	}
+
+	group.LoadGroupPeers()
+
+	return &group, nil
+}
+
+// GetGroupsByIDs retrieves groups by their IDs and account ID.
+func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*types.Group, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var groups []*types.Group
+	result := tx.Preload(clause.Associations).Find(&groups, accountAndIDsQueryCondition, accountID, groupIDs)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get groups by ID's from store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get groups by ID's from store")
+	}
+
+	groupsMap := make(map[string]*types.Group)
+	for _, group := range groups {
+		group.LoadGroupPeers()
+		groupsMap[group.ID] = group
+	}
+
+	return groupsMap, nil
+}
+
+// CreateGroup creates a group in the store.
+func (s *SqlStore) CreateGroup(ctx context.Context, group *types.Group) error {
+	if group == nil {
+		return status.Errorf(status.InvalidArgument, "group is nil")
+	}
+
+	if err := s.db.Omit(clause.Associations).Create(group).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to save group to store: %v", err)
+		return status.Errorf(status.Internal, "failed to save group to store")
+	}
+
+	return nil
+}
+
+// UpdateGroup updates a group in the store.
+func (s *SqlStore) UpdateGroup(ctx context.Context, group *types.Group) error {
+	if group == nil {
+		return status.Errorf(status.InvalidArgument, "group is nil")
+	}
+
+	if err := s.db.Omit(clause.Associations, "public_id").Save(group).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to save group to store: %v", err)
+		return status.Errorf(status.Internal, "failed to save group to store")
+	}
+
+	return nil
+}
+
+// DeleteGroup deletes a group from the database.
+func (s *SqlStore) DeleteGroup(ctx context.Context, accountID, groupID string) error {
+	result := s.db.Select(clause.Associations).
+		Delete(&types.Group{}, accountAndIDQueryCondition, accountID, groupID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to delete group from store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to delete group from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewGroupNotFoundError(groupID)
+	}
+
+	return nil
+}
+
+// DeleteGroups deletes groups from the database.
+func (s *SqlStore) DeleteGroups(ctx context.Context, accountID string, groupIDs []string) error {
+	result := s.db.Select(clause.Associations).
+		Delete(&types.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 nil
+}

management/server/store/sql_store_group_peer.go

@@ -0,0 +1,229 @@
+package store
+
+import (
+	"context"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) getGroupPeers(ctx context.Context, groupIDs []string) ([]types.GroupPeer, error) {
+	if len(groupIDs) == 0 {
+		return nil, nil
+	}
+	const query = `SELECT account_id, group_id, peer_id FROM group_peers WHERE group_id = ANY($1)`
+	rows, err := s.pool.Query(ctx, query, groupIDs)
+	if err != nil {
+		return nil, err
+	}
+	groupPeers, err := pgx.CollectRows(rows, pgx.RowToStructByName[types.GroupPeer])
+	if err != nil {
+		return nil, err
+	}
+	return groupPeers, nil
+}
+
+// AddPeerToAllGroup adds a peer to the 'All' group. Method always needs to run in a transaction
+func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error {
+	var groupID string
+	_ = s.db.Model(types.Group{}).
+		Select("id").
+		Where("account_id = ? AND name = ?", accountID, "All").
+		Limit(1).
+		Scan(&groupID)
+
+	if groupID == "" {
+		return status.Errorf(status.NotFound, "group 'All' not found for account %s", accountID)
+	}
+
+	err := s.db.Clauses(clause.OnConflict{
+		Columns:   []clause.Column{{Name: "group_id"}, {Name: "peer_id"}},
+		DoNothing: true,
+	}).Create(&types.GroupPeer{
+		AccountID: accountID,
+		GroupID:   groupID,
+		PeerID:    peerID,
+	}).Error
+	if err != nil {
+		return status.Errorf(status.Internal, "error adding peer to group 'All': %v", err)
+	}
+
+	return nil
+}
+
+// AddPeerToGroup adds a peer to a group
+func (s *SqlStore) AddPeerToGroup(ctx context.Context, accountID, peerID, groupID string) error {
+	peer := &types.GroupPeer{
+		AccountID: accountID,
+		GroupID:   groupID,
+		PeerID:    peerID,
+	}
+
+	err := s.db.Clauses(clause.OnConflict{
+		Columns:   []clause.Column{{Name: "group_id"}, {Name: "peer_id"}},
+		DoNothing: true,
+	}).Create(peer).Error
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to add peer %s to group %s for account %s: %v", peerID, groupID, accountID, err)
+		return status.Errorf(status.Internal, "failed to add peer to group")
+	}
+
+	return nil
+}
+
+// RemovePeerFromGroup removes a peer from a group
+func (s *SqlStore) RemovePeerFromGroup(ctx context.Context, peerID string, groupID string) error {
+	err := s.db.
+		Delete(&types.GroupPeer{}, "group_id = ? AND peer_id = ?", groupID, peerID).Error
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to remove peer %s from group %s: %v", peerID, groupID, err)
+		return status.Errorf(status.Internal, "failed to remove peer from group")
+	}
+
+	return nil
+}
+
+// RemovePeerFromAllGroups removes a peer from all groups
+func (s *SqlStore) RemovePeerFromAllGroups(ctx context.Context, peerID string) error {
+	err := s.db.
+		Delete(&types.GroupPeer{}, "peer_id = ?", peerID).Error
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to remove peer %s from all groups: %v", peerID, err)
+		return status.Errorf(status.Internal, "failed to remove peer from all groups")
+	}
+
+	return nil
+}
+
+// GetPeerGroups retrieves all groups assigned to a specific peer in a given account.
+func (s *SqlStore) GetPeerGroups(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]*types.Group, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var groups []*types.Group
+	query := tx.
+		Joins("JOIN group_peers ON group_peers.group_id = groups.id").
+		Where("groups.account_id = ? AND group_peers.peer_id = ?", accountId, peerId).
+		Preload(clause.Associations).
+		Find(&groups)
+
+	if query.Error != nil {
+		return nil, query.Error
+	}
+
+	for _, group := range groups {
+		group.LoadGroupPeers()
+	}
+
+	return groups, nil
+}
+
+// GetPeerGroupIDs retrieves all group IDs assigned to a specific peer in a given account.
+func (s *SqlStore) GetPeerGroupIDs(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var groupIDs []string
+	query := tx.
+		Model(&types.GroupPeer{}).
+		Where("account_id = ? AND peer_id = ?", accountId, peerId).
+		Pluck("group_id", &groupIDs)
+
+	if query.Error != nil {
+		if errors.Is(query.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "no groups found for peer %s in account %s", peerId, accountId)
+		}
+		log.WithContext(ctx).Errorf("failed to get group IDs for peer %s in account %s: %v", peerId, accountId, query.Error)
+		return nil, status.Errorf(status.Internal, "failed to get group IDs for peer from store")
+	}
+
+	return groupIDs, nil
+}
+
+func (s *SqlStore) GetAccountGroupPeers(ctx context.Context, lockStrength LockingStrength, accountID string) (map[string]map[string]struct{}, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peers []types.GroupPeer
+	result := tx.Find(&peers, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get account group peers from store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get account group peers from store")
+	}
+
+	groupPeers := make(map[string]map[string]struct{})
+	for _, peer := range peers {
+		if _, exists := groupPeers[peer.GroupID]; !exists {
+			groupPeers[peer.GroupID] = make(map[string]struct{})
+		}
+		groupPeers[peer.GroupID][peer.PeerID] = struct{}{}
+	}
+
+	return groupPeers, nil
+}
+
+func (s *SqlStore) GetPeersByGroupIDs(ctx context.Context, accountID string, groupIDs []string) ([]*nbpeer.Peer, error) {
+	if len(groupIDs) == 0 {
+		return []*nbpeer.Peer{}, nil
+	}
+
+	var peers []*nbpeer.Peer
+	peerIDsSubquery := s.db.Model(&types.GroupPeer{}).
+		Select("DISTINCT peer_id").
+		Where("account_id = ? AND group_id IN ?", accountID, groupIDs)
+
+	result := s.db.Where("account_id = ? AND id IN (?)", accountID, peerIDsSubquery).Find(&peers)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get peers by group IDs: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get peers by group IDs")
+	}
+
+	return peers, nil
+}
+
+func (s *SqlStore) GetPeerIDsByGroups(ctx context.Context, accountID string, groupIDs []string) ([]string, error) {
+	if len(groupIDs) == 0 {
+		return nil, nil
+	}
+
+	var peerIDs []string
+	result := s.db.Model(&types.GroupPeer{}).
+		Select("DISTINCT peer_id").
+		Where("account_id = ? AND group_id IN ?", accountID, groupIDs).
+		Pluck("peer_id", &peerIDs)
+	if result.Error != nil {
+		return nil, status.Errorf(status.Internal, "failed to get peer IDs by groups: %s", result.Error)
+	}
+
+	return peerIDs, nil
+}
+
+func (s *SqlStore) GetGroupIDsByPeerIDs(ctx context.Context, accountID string, peerIDs []string) ([]string, error) {
+	if len(peerIDs) == 0 {
+		return nil, nil
+	}
+
+	var groupIDs []string
+	result := s.db.Model(&types.GroupPeer{}).
+		Select("DISTINCT group_id").
+		Where("account_id = ? AND peer_id IN ?", accountID, peerIDs).
+		Pluck("group_id", &groupIDs)
+	if result.Error != nil {
+		return nil, status.Errorf(status.Internal, "failed to get group IDs by peers: %s", result.Error)
+	}
+
+	return groupIDs, nil
+}

management/server/store/sql_store_group_peer_test.go

@@ -0,0 +1,210 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+func TestSqlStore_AddPeerToGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	peerID := "cfefqs706sqkneg59g4g"
+	groupID := "cfefqs706sqkneg59g4h"
+
+	group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err, "failed to get group")
+	require.Len(t, group.Peers, 0, "group should have 0 peers")
+
+	err = store.AddPeerToGroup(context.Background(), accountID, peerID, groupID)
+	require.NoError(t, err, "failed to add peer to group")
+
+	group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err, "failed to get group")
+	require.Len(t, group.Peers, 1, "group should have 1 peers")
+	require.Contains(t, group.Peers, peerID)
+}
+
+func TestSqlStore_AddPeerToAllGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	groupID := "cfefqs706sqkneg59g3g"
+
+	peer := &nbpeer.Peer{
+		ID:        "peer1",
+		AccountID: accountID,
+		DNSLabel:  "peer1.domain.test",
+	}
+
+	group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err, "failed to get group")
+	require.Len(t, group.Peers, 2, "group should have 2 peers")
+	require.NotContains(t, group.Peers, peer.ID)
+
+	err = store.AddPeerToAccount(context.Background(), peer)
+	require.NoError(t, err, "failed to add peer to account")
+
+	err = store.AddPeerToAllGroup(context.Background(), accountID, peer.ID)
+	require.NoError(t, err, "failed to add peer to all group")
+
+	group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err, "failed to get group")
+	require.Len(t, group.Peers, 3, "group should have  peers")
+	require.Contains(t, group.Peers, peer.ID)
+}
+
+func TestSqlStore_GetPeerGroups(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	peerID := "cfefqs706sqkneg59g4g"
+
+	groups, err := store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
+	require.NoError(t, err)
+	assert.Len(t, groups, 1)
+	assert.Equal(t, groups[0].Name, "All")
+
+	err = store.AddPeerToGroup(context.Background(), accountID, peerID, "cfefqs706sqkneg59g4h")
+	require.NoError(t, err)
+
+	groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
+	require.NoError(t, err)
+	assert.Len(t, groups, 2)
+
+	foreignPeerID := "foreign-peer"
+	err = store.AddPeerToGroup(context.Background(), accountID, foreignPeerID, "cfefqs706sqkneg59g4h")
+	require.NoError(t, err)
+
+	groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, "other-account", foreignPeerID)
+	require.NoError(t, err)
+	assert.Empty(t, groups, "groups of another account must not be returned")
+}
+
+func TestSqlStore_GetPeersByGroupIDs(t *testing.T) {
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	group1ID := "test-group-1"
+	group2ID := "test-group-2"
+	emptyGroupID := "empty-group"
+
+	peer1 := "cfefqs706sqkneg59g4g"
+	peer2 := "cfeg6sf06sqkneg59g50"
+
+	tests := []struct {
+		name          string
+		groupIDs      []string
+		expectedPeers []string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve peers from single group with multiple peers",
+			groupIDs:      []string{group1ID},
+			expectedPeers: []string{peer1, peer2},
+			expectedCount: 2,
+		},
+		{
+			name:          "retrieve peers from single group with one peer",
+			groupIDs:      []string{group2ID},
+			expectedPeers: []string{peer1},
+			expectedCount: 1,
+		},
+		{
+			name:          "retrieve peers from multiple groups (with overlap)",
+			groupIDs:      []string{group1ID, group2ID},
+			expectedPeers: []string{peer1, peer2}, // should deduplicate
+			expectedCount: 2,
+		},
+		{
+			name:          "retrieve peers from existing 'All' group",
+			groupIDs:      []string{"cfefqs706sqkneg59g3g"}, // All group from test data
+			expectedPeers: []string{peer1, peer2},
+			expectedCount: 2,
+		},
+		{
+			name:          "retrieve peers from empty group",
+			groupIDs:      []string{emptyGroupID},
+			expectedPeers: []string{},
+			expectedCount: 0,
+		},
+		{
+			name:          "retrieve peers from non-existing group",
+			groupIDs:      []string{"non-existing-group"},
+			expectedPeers: []string{},
+			expectedCount: 0,
+		},
+		{
+			name:          "empty group IDs list",
+			groupIDs:      []string{},
+			expectedPeers: []string{},
+			expectedCount: 0,
+		},
+		{
+			name:          "mix of existing and non-existing groups",
+			groupIDs:      []string{group1ID, "non-existing-group"},
+			expectedPeers: []string{peer1, peer2},
+			expectedCount: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+			t.Cleanup(cleanup)
+			require.NoError(t, err)
+
+			ctx := context.Background()
+
+			groups := []*types.Group{
+				{
+					ID:        group1ID,
+					AccountID: accountID,
+				},
+				{
+					ID:        group2ID,
+					AccountID: accountID,
+				},
+			}
+			require.NoError(t, store.CreateGroups(ctx, accountID, groups))
+
+			otherAccount := newAccountWithId(ctx, "other-account", "other-user", "")
+			require.NoError(t, store.SaveAccount(ctx, otherAccount))
+			foreignPeer := &nbpeer.Peer{ID: "foreign-peer", AccountID: otherAccount.Id}
+			require.NoError(t, store.AddPeerToAccount(ctx, foreignPeer))
+
+			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group1ID))
+			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer2, group1ID))
+			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group2ID))
+			require.NoError(t, store.AddPeerToGroup(ctx, accountID, foreignPeer.ID, group1ID))
+
+			peers, err := store.GetPeersByGroupIDs(ctx, accountID, tt.groupIDs)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+
+			if tt.expectedCount > 0 {
+				actualPeerIDs := make([]string, len(peers))
+				for i, peer := range peers {
+					actualPeerIDs[i] = peer.ID
+				}
+				assert.ElementsMatch(t, tt.expectedPeers, actualPeerIDs)
+
+				// Verify all returned peers belong to the correct account
+				for _, peer := range peers {
+					assert.Equal(t, accountID, peer.AccountID)
+				}
+			}
+		})
+	}
+}

management/server/store/sql_store_group_test.go

@@ -0,0 +1,289 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlite_GetGroupByName(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	if err != nil {
+		t.Fatal(err)
+	}
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	group, err := store.GetGroupByName(context.Background(), LockingStrengthNone, accountID, "All")
+	require.NoError(t, err)
+	require.True(t, group.IsGroupAll())
+}
+
+func TestSqlStore_GetGroupsByIDs(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	tests := []struct {
+		name          string
+		groupIDs      []string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve existing groups by existing IDs",
+			groupIDs:      []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
+			expectedCount: 2,
+		},
+		{
+			name:          "empty group IDs list",
+			groupIDs:      []string{},
+			expectedCount: 0,
+		},
+		{
+			name:          "non-existing group IDs",
+			groupIDs:      []string{"nonexistent1", "nonexistent2"},
+			expectedCount: 0,
+		},
+		{
+			name:          "mixed existing and non-existing group IDs",
+			groupIDs:      []string{"cfefqs706sqkneg59g4g", "nonexistent"},
+			expectedCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			groups, err := store.GetGroupsByIDs(context.Background(), LockingStrengthNone, accountID, tt.groupIDs)
+			require.NoError(t, err)
+			require.Len(t, groups, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_CreateGroup(t *testing.T) {
+	if os.Getenv("CI") == "true" {
+		t.Log("Skipping MySQL test on CI")
+	}
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.MysqlStoreEngine))
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	group := &types.Group{
+		ID:         "group-id",
+		AccountID:  accountID,
+		Issued:     "api",
+		Peers:      []string{},
+		Resources:  []types.Resource{},
+		GroupPeers: []types.GroupPeer{},
+	}
+	err = store.CreateGroup(context.Background(), group)
+	require.NoError(t, err)
+
+	savedGroup, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, "group-id")
+	require.NoError(t, err)
+	require.Equal(t, savedGroup, group)
+}
+
+func TestSqlStore_CreateUpdateGroups(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	groups := []*types.Group{
+		{
+			ID:         "group-1",
+			AccountID:  accountID,
+			Issued:     "api",
+			Peers:      []string{},
+			Resources:  []types.Resource{},
+			GroupPeers: []types.GroupPeer{},
+		},
+		{
+			ID:         "group-2",
+			AccountID:  accountID,
+			Issued:     "integration",
+			Peers:      []string{},
+			Resources:  []types.Resource{},
+			GroupPeers: []types.GroupPeer{},
+		},
+	}
+	err = store.CreateGroups(context.Background(), accountID, groups)
+	require.NoError(t, err)
+
+	groups[1].Peers = []string{}
+	err = store.UpdateGroups(context.Background(), accountID, groups)
+	require.NoError(t, err)
+
+	group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groups[1].ID)
+	require.NoError(t, err)
+	require.Equal(t, groups[1], group)
+}
+
+func TestSqlStore_DeleteGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	tests := []struct {
+		name        string
+		groupID     string
+		expectError bool
+	}{
+		{
+			name:        "delete existing group",
+			groupID:     "cfefqs706sqkneg59g4g",
+			expectError: false,
+		},
+		{
+			name:        "delete non-existing group",
+			groupID:     "non-existing-group-id",
+			expectError: true,
+		},
+		{
+			name:        "delete with empty group ID",
+			groupID:     "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := store.DeleteGroup(context.Background(), accountID, tt.groupID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+			} else {
+				require.NoError(t, err)
+
+				group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, tt.groupID)
+				require.Error(t, err)
+				require.Nil(t, group)
+			}
+		})
+	}
+}
+
+func TestSqlStore_DeleteGroups(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	tests := []struct {
+		name        string
+		groupIDs    []string
+		expectError bool
+	}{
+		{
+			name:        "delete multiple existing groups",
+			groupIDs:    []string{"cfefqs706sqkneg59g4g", "cfefqs706sqkneg59g3g"},
+			expectError: false,
+		},
+		{
+			name:        "delete non-existing groups",
+			groupIDs:    []string{"non-existing-id-1", "non-existing-id-2"},
+			expectError: false,
+		},
+		{
+			name:        "delete with empty group IDs list",
+			groupIDs:    []string{},
+			expectError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := store.DeleteGroups(context.Background(), accountID, tt.groupIDs)
+			if tt.expectError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+
+				for _, groupID := range tt.groupIDs {
+					group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+					require.Error(t, err)
+					require.Nil(t, group)
+				}
+			}
+		})
+	}
+}
+
+func TestSqlStore_AddAndRemoveResourceFromGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	resourceId := "ctc4nci7qv9061u6ilfg"
+	groupID := "cs1tnh0hhcjnqoiuebeg"
+
+	res := &types.Resource{
+		ID:   resourceId,
+		Type: "host",
+	}
+	err = store.AddResourceToGroup(context.Background(), accountID, groupID, res)
+	require.NoError(t, err)
+
+	group, err := store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err)
+	require.Contains(t, group.Resources, *res)
+
+	groups, err := store.GetResourceGroups(context.Background(), LockingStrengthNone, accountID, resourceId)
+	require.NoError(t, err)
+	require.Len(t, groups, 1)
+
+	err = store.RemoveResourceFromGroup(context.Background(), accountID, groupID, res.ID)
+	require.NoError(t, err)
+
+	group, err = store.GetGroupByID(context.Background(), LockingStrengthNone, accountID, groupID)
+	require.NoError(t, err)
+	require.NotContains(t, group.Resources, *res)
+}
+
+func TestSqlStore_SaveGroups_LargeBatch(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	accountGroups, err := store.GetAccountGroups(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	require.Len(t, accountGroups, 3)
+
+	groupsToSave := make([]*types.Group, 0)
+
+	for i := 1; i <= 8000; i++ {
+		groupsToSave = append(groupsToSave, &types.Group{
+			ID:        fmt.Sprintf("%d", i),
+			AccountID: accountID,
+			Name:      fmt.Sprintf("group-%d", i),
+		})
+	}
+
+	err = store.CreateGroups(context.Background(), accountID, groupsToSave)
+	require.NoError(t, err)
+
+	accountGroups, err = store.GetAccountGroups(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	require.Equal(t, 8003, len(accountGroups))
+}

management/server/store/sql_store_installation.go

@@ -0,0 +1,29 @@
+package store
+
+import (
+	"context"
+
+	"gorm.io/gorm/clause"
+)
+
+type installation struct {
+	ID                  uint `gorm:"primaryKey"`
+	InstallationIDValue string
+}
+
+func (s *SqlStore) SaveInstallationID(_ context.Context, ID string) error {
+	installation := installation{InstallationIDValue: ID}
+	installation.ID = uint(s.installationPK)
+
+	return s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(&installation).Error
+}
+
+func (s *SqlStore) GetInstallationID() string {
+	var installation installation
+
+	if result := s.db.Take(&installation, idQueryCondition, s.installationPK); result.Error != nil {
+		return ""
+	}
+
+	return installation.InstallationIDValue
+}

management/server/store/sql_store_job.go

@@ -0,0 +1,103 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// SaveJob persists a job in DB
+func (s *SqlStore) CreatePeerJob(ctx context.Context, job *types.Job) error {
+	result := s.db.Create(job)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to create job in store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to create job in store")
+	}
+	return nil
+}
+
+func (s *SqlStore) CompletePeerJob(ctx context.Context, job *types.Job) error {
+	result := s.db.
+		Model(&types.Job{}).
+		Where(idQueryCondition, job.ID).
+		Updates(job)
+
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to update job in store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to update job in store")
+	}
+	return nil
+}
+
+// job was pending for too long and has been cancelled
+func (s *SqlStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
+	now := time.Now().UTC()
+	result := s.db.
+		Model(&types.Job{}).
+		Where(accountAndPeerIDQueryCondition+" AND id = ?"+" AND status = ?", accountID, peerID, jobID, types.JobStatusPending).
+		Updates(types.Job{
+			Status:       types.JobStatusFailed,
+			FailedReason: reason,
+			CompletedAt:  &now,
+		})
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to mark pending jobs as Failed job in store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to mark pending job as Failed in store")
+	}
+	return nil
+}
+
+// job was pending for too long and has been cancelled
+func (s *SqlStore) MarkAllPendingJobsAsFailed(ctx context.Context, accountID, peerID, reason string) error {
+	now := time.Now().UTC()
+	result := s.db.
+		Model(&types.Job{}).
+		Where(accountAndPeerIDQueryCondition+" AND status = ?", accountID, peerID, types.JobStatusPending).
+		Updates(types.Job{
+			Status:       types.JobStatusFailed,
+			FailedReason: reason,
+			CompletedAt:  &now,
+		})
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to mark pending jobs as Failed job in store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to mark pending job as Failed in store")
+	}
+	return nil
+}
+
+// GetJobByID fetches job by ID
+func (s *SqlStore) GetPeerJobByID(ctx context.Context, accountID, jobID string) (*types.Job, error) {
+	var job types.Job
+	err := s.db.
+		Where(accountAndIDQueryCondition, accountID, jobID).
+		First(&job).Error
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, status.Errorf(status.NotFound, "job %s not found", jobID)
+	}
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to fetch job from store: %s", err)
+		return nil, err
+	}
+	return &job, nil
+}
+
+// get all jobs
+func (s *SqlStore) GetPeerJobs(ctx context.Context, accountID, peerID string) ([]*types.Job, error) {
+	var jobs []*types.Job
+	err := s.db.
+		Where(accountAndPeerIDQueryCondition, accountID, peerID).
+		Order("created_at DESC").
+		Find(&jobs).Error
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to fetch jobs from store: %s", err)
+		return nil, err
+	}
+
+	return jobs, nil
+}

management/server/store/sql_store_name_server_group.go

@@ -0,0 +1,124 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) getNameServerGroups(ctx context.Context, accountID string) ([]nbdns.NameServerGroup, error) {
+	const query = `SELECT id, account_id, public_id, name, description, name_servers, groups, "primary", domains, enabled, search_domains_enabled FROM name_server_groups WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+	nsgs, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nbdns.NameServerGroup, error) {
+		var n nbdns.NameServerGroup
+		var ns, groups, domains []byte
+		var primary, enabled, searchDomainsEnabled sql.NullBool
+		err := row.Scan(&n.ID, &n.AccountID, &n.PublicID, &n.Name, &n.Description, &ns, &groups, &primary, &domains, &enabled, &searchDomainsEnabled)
+		if err == nil {
+			if primary.Valid {
+				n.Primary = primary.Bool
+			}
+			if enabled.Valid {
+				n.Enabled = enabled.Bool
+			}
+			if searchDomainsEnabled.Valid {
+				n.SearchDomainsEnabled = searchDomainsEnabled.Bool
+			}
+			if ns != nil {
+				_ = json.Unmarshal(ns, &n.NameServers)
+			} else {
+				n.NameServers = []nbdns.NameServer{}
+			}
+			if groups != nil {
+				_ = json.Unmarshal(groups, &n.Groups)
+			} else {
+				n.Groups = []string{}
+			}
+			if domains != nil {
+				_ = json.Unmarshal(domains, &n.Domains)
+			} else {
+				n.Domains = []string{}
+			}
+		}
+		return n, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return nsgs, nil
+}
+
+// GetAccountNameServerGroups retrieves name server groups for an account.
+func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var nsGroups []*nbdns.NameServerGroup
+	result := tx.Find(&nsGroups, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get name server groups from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get name server groups from store")
+	}
+
+	return nsGroups, nil
+}
+
+// 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) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var nsGroup *nbdns.NameServerGroup
+	result := tx.
+		Take(&nsGroup, accountAndIDQueryCondition, accountID, nsGroupID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.NewNameServerGroupNotFoundError(nsGroupID)
+		}
+		log.WithContext(ctx).Errorf("failed to get name server group from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get name server group from store")
+	}
+
+	return nsGroup, nil
+}
+
+// SaveNameServerGroup saves a name server group to the database.
+func (s *SqlStore) SaveNameServerGroup(ctx context.Context, nameServerGroup *nbdns.NameServerGroup) error {
+	result := s.db.Save(nameServerGroup)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to save name server group to the store: %s", err)
+		return status.Errorf(status.Internal, "failed to save name server group to store")
+	}
+	return nil
+}
+
+// DeleteNameServerGroup deletes a name server group from the database.
+func (s *SqlStore) DeleteNameServerGroup(ctx context.Context, accountID, nsGroupID string) error {
+	result := s.db.Delete(&nbdns.NameServerGroup{}, accountAndIDQueryCondition, accountID, nsGroupID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to delete name server group from the store: %s", err)
+		return status.Errorf(status.Internal, "failed to delete name server group from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewNameServerGroupNotFoundError(nsGroupID)
+	}
+
+	return nil
+}

management/server/store/sql_store_name_server_group_test.go

@@ -0,0 +1,143 @@
+package store
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlStore_GetAccountNameServerGroups(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name          string
+		accountID     string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve name server groups by existing account ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectedCount: 1,
+		},
+		{
+			name:          "non-existing account ID",
+			accountID:     "nonexistent",
+			expectedCount: 0,
+		},
+		{
+			name:          "empty account ID",
+			accountID:     "",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetAccountNameServerGroups(context.Background(), LockingStrengthNone, tt.accountID)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+		})
+	}
+
+}
+
+func TestSqlStore_GetNameServerByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name        string
+		nsGroupID   string
+		expectError bool
+	}{
+		{
+			name:        "retrieve existing nameserver group",
+			nsGroupID:   "csqdelq7qv97ncu7d9t0",
+			expectError: false,
+		},
+		{
+			name:        "retrieve non-existing nameserver group",
+			nsGroupID:   "non-existing",
+			expectError: true,
+		},
+		{
+			name:        "retrieve with empty nameserver group ID",
+			nsGroupID:   "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, tt.nsGroupID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, nsGroup)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, nsGroup)
+				require.Equal(t, tt.nsGroupID, nsGroup.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_SaveNameServerGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	nsGroup := &nbdns.NameServerGroup{
+		ID:        "ns-group-id",
+		AccountID: accountID,
+		Name:      "NS Group",
+		NameServers: []nbdns.NameServer{
+			{
+				IP:     netip.MustParseAddr("8.8.8.8"),
+				NSType: 1,
+				Port:   53,
+			},
+		},
+		Groups:               []string{"groupA"},
+		Primary:              true,
+		Enabled:              true,
+		SearchDomainsEnabled: false,
+	}
+
+	err = store.SaveNameServerGroup(context.Background(), nsGroup)
+	require.NoError(t, err)
+
+	saveNSGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, nsGroup.ID)
+	require.NoError(t, err)
+	require.Equal(t, saveNSGroup, nsGroup)
+}
+
+func TestSqlStore_DeleteNameServerGroup(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	nsGroupID := "csqdelq7qv97ncu7d9t0"
+
+	err = store.DeleteNameServerGroup(context.Background(), accountID, nsGroupID)
+	require.NoError(t, err)
+
+	nsGroup, err := store.GetNameServerGroupByID(context.Background(), LockingStrengthNone, accountID, nsGroupID)
+	require.Error(t, err)
+	require.Nil(t, nsGroup)
+}

management/server/store/sql_store_network.go

@@ -0,0 +1,91 @@
+package store
+
+import (
+	"context"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) getNetworks(ctx context.Context, accountID string) ([]*networkTypes.Network, error) {
+	const query = `SELECT id, account_id, public_id, name, description FROM networks WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+	networks, err := pgx.CollectRows(rows, pgx.RowToStructByName[networkTypes.Network])
+	if err != nil {
+		return nil, err
+	}
+	result := make([]*networkTypes.Network, len(networks))
+	for i := range networks {
+		result[i] = &networks[i]
+	}
+	return result, nil
+}
+
+func (s *SqlStore) GetAccountNetworks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*networkTypes.Network, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var networks []*networkTypes.Network
+	result := tx.Find(&networks, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get networks from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get networks from store")
+	}
+
+	return networks, nil
+}
+
+func (s *SqlStore) GetNetworkByID(ctx context.Context, lockStrength LockingStrength, accountID, networkID string) (*networkTypes.Network, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var network *networkTypes.Network
+	result := tx.Take(&network, accountAndIDQueryCondition, accountID, networkID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkNotFoundError(networkID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get network from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network from store")
+	}
+
+	return network, nil
+}
+
+func (s *SqlStore) SaveNetwork(ctx context.Context, network *networkTypes.Network) error {
+	result := s.db.Save(network)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save network to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save network to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteNetwork(ctx context.Context, accountID, networkID string) error {
+	result := s.db.Delete(&networkTypes.Network{}, accountAndIDQueryCondition, accountID, networkID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete network from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete network from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewNetworkNotFoundError(networkID)
+	}
+
+	return nil
+}

management/server/store/sql_store_network_resource.go

@@ -0,0 +1,167 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) getNetworkResources(ctx context.Context, accountID string) ([]*resourceTypes.NetworkResource, error) {
+	const query = `SELECT id, network_id, account_id, public_id, name, description, type, domain, prefix, enabled FROM network_resources WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+	resources, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (resourceTypes.NetworkResource, error) {
+		var r resourceTypes.NetworkResource
+		var prefix []byte
+		var enabled sql.NullBool
+		err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Name, &r.Description, &r.Type, &r.Domain, &prefix, &enabled)
+		if err == nil {
+			if enabled.Valid {
+				r.Enabled = enabled.Bool
+			}
+			if prefix != nil {
+				_ = json.Unmarshal(prefix, &r.Prefix)
+			}
+		}
+		return r, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	result := make([]*resourceTypes.NetworkResource, len(resources))
+	for i := range resources {
+		result[i] = &resources[i]
+	}
+	return result, nil
+}
+
+func (s *SqlStore) GetNetworkResourcesByNetID(ctx context.Context, lockStrength LockingStrength, accountID, networkID string) ([]*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources []*resourceTypes.NetworkResource
+	result := tx.
+		Find(&netResources, "account_id = ? AND network_id = ?", accountID, networkID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get network resources from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resources from store")
+	}
+
+	return netResources, nil
+}
+
+func (s *SqlStore) GetNetworkResourcesByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources []*resourceTypes.NetworkResource
+	result := tx.
+		Find(&netResources, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get network resources from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resources from store")
+	}
+
+	return netResources, nil
+}
+
+func (s *SqlStore) GetNetworkResourceByID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources *resourceTypes.NetworkResource
+	result := tx.
+		Take(&netResources, accountAndIDQueryCondition, accountID, resourceID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkResourceNotFoundError(resourceID)
+		}
+		log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resource from store")
+	}
+
+	return netResources, nil
+}
+
+// GetNetworkResourceByIDOrPublicID retrieves a network resource by either its ID or its
+// PublicID. See GetPolicyByIDOrPublicID for why peer-reported references need both.
+func (s *SqlStore) GetNetworkResourceByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources *resourceTypes.NetworkResource
+	result := tx.
+		Take(&netResources, accountAndAnyIDQueryCondition, accountID, resourceID, resourceID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkResourceNotFoundError(resourceID)
+		}
+		log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resource from store")
+	}
+
+	return netResources, nil
+}
+
+func (s *SqlStore) GetNetworkResourceByName(ctx context.Context, lockStrength LockingStrength, accountID, resourceName string) (*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources *resourceTypes.NetworkResource
+	result := tx.
+		Take(&netResources, "account_id = ? AND name = ?", accountID, resourceName)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkResourceNotFoundError(resourceName)
+		}
+		log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resource from store")
+	}
+
+	return netResources, nil
+}
+
+func (s *SqlStore) SaveNetworkResource(ctx context.Context, resource *resourceTypes.NetworkResource) error {
+	result := s.db.Save(resource)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save network resource to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save network resource to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteNetworkResource(ctx context.Context, accountID, resourceID string) error {
+	result := s.db.Delete(&resourceTypes.NetworkResource{}, accountAndIDQueryCondition, accountID, resourceID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete network resource from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete network resource from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewNetworkResourceNotFoundError(resourceID)
+	}
+
+	return nil
+}

management/server/store/sql_store_network_resource_test.go

@@ -0,0 +1,161 @@
+package store
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlStore_GetNetworkResourcesByNetID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	tests := []struct {
+		name          string
+		networkID     string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve resources by existing network ID",
+			networkID:     "ct286bi7qv930dsrrug0",
+			expectedCount: 1,
+		},
+		{
+			name:          "retrieve resources by non-existing network ID",
+			networkID:     "non-existent",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			netResources, err := store.GetNetworkResourcesByNetID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
+			require.NoError(t, err)
+			require.Len(t, netResources, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_GetNetworkResourceByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name          string
+		netResourceID string
+		expectError   bool
+	}{
+		{
+			name:          "retrieve existing network resource ID",
+			netResourceID: "ctc4nci7qv9061u6ilfg",
+			expectError:   false,
+		},
+		{
+			name:          "retrieve non-existing network resource ID",
+			netResourceID: "non-existing",
+			expectError:   true,
+		},
+		{
+			name:          "retrieve network with empty resource ID",
+			netResourceID: "",
+			expectError:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			netResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, tt.netResourceID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, netResource)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, netResource)
+				require.Equal(t, tt.netResourceID, netResource.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_GetNetworkResourceByIDOrPublicID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	netResourceID := "ctc4nci7qv9061u6ilfg"
+
+	netResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, netResourceID)
+	require.NoError(t, err)
+	require.NotEmpty(t, netResource.PublicID)
+
+	for _, id := range []string{netResourceID, netResource.PublicID} {
+		netResource, err := store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		require.Equal(t, netResourceID, netResource.ID)
+	}
+
+	netResource, err = store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, sErr.Type(), status.NotFound)
+	require.Nil(t, netResource)
+}
+
+func TestSqlStore_SaveNetworkResource(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	networkID := "ct286bi7qv930dsrrug0"
+
+	netResource, err := resourceTypes.NewNetworkResource(accountID, networkID, "resource-name", "", "example.com", []string{}, true)
+	require.NoError(t, err)
+
+	err = store.SaveNetworkResource(context.Background(), netResource)
+	require.NoError(t, err)
+
+	savedNetResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, netResource.ID)
+	require.NoError(t, err)
+	require.Equal(t, netResource.ID, savedNetResource.ID)
+	require.Equal(t, netResource.Name, savedNetResource.Name)
+	require.Equal(t, netResource.NetworkID, savedNetResource.NetworkID)
+	require.Equal(t, netResource.Type, resourceTypes.NetworkResourceType("domain"))
+	require.Equal(t, netResource.Domain, "example.com")
+	require.Equal(t, netResource.AccountID, savedNetResource.AccountID)
+	require.Equal(t, netResource.Prefix, netip.Prefix{})
+}
+
+func TestSqlStore_DeleteNetworkResource(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	netResourceID := "ctc4nci7qv9061u6ilfg"
+
+	err = store.DeleteNetworkResource(context.Background(), accountID, netResourceID)
+	require.NoError(t, err)
+
+	netResource, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, netResourceID)
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, sErr.Type())
+	require.Nil(t, netResource)
+}

management/server/store/sql_store_network_router.go

@@ -0,0 +1,208 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
+	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) getNetworkRouters(ctx context.Context, accountID string) ([]*routerTypes.NetworkRouter, error) {
+	const query = `SELECT id, network_id, account_id, public_id, peer, peer_groups, masquerade, metric, enabled FROM network_routers WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+	routers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (routerTypes.NetworkRouter, error) {
+		var r routerTypes.NetworkRouter
+		var peerGroups []byte
+		var masquerade, enabled sql.NullBool
+		var metric sql.NullInt64
+		err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Peer, &peerGroups, &masquerade, &metric, &enabled)
+		if err == nil {
+			if masquerade.Valid {
+				r.Masquerade = masquerade.Bool
+			}
+			if enabled.Valid {
+				r.Enabled = enabled.Bool
+			}
+			if metric.Valid {
+				r.Metric = int(metric.Int64)
+			}
+			if peerGroups != nil {
+				_ = json.Unmarshal(peerGroups, &r.PeerGroups)
+			}
+		}
+		return r, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	result := make([]*routerTypes.NetworkRouter, len(routers))
+	for i := range routers {
+		result[i] = &routers[i]
+	}
+	return result, nil
+}
+
+func (s *SqlStore) GetNetworkRoutersByNetID(ctx context.Context, lockStrength LockingStrength, accountID, netID string) ([]*routerTypes.NetworkRouter, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netRouters []*routerTypes.NetworkRouter
+	result := tx.
+		Find(&netRouters, "account_id = ? AND network_id = ?", accountID, netID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get network routers from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network routers from store")
+	}
+
+	return netRouters, nil
+}
+
+func (s *SqlStore) GetNetworkRoutersByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*routerTypes.NetworkRouter, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netRouters []*routerTypes.NetworkRouter
+	result := tx.
+		Find(&netRouters, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get network routers from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network routers from store")
+	}
+
+	return netRouters, nil
+}
+
+func (s *SqlStore) GetNetworkRouterByID(ctx context.Context, lockStrength LockingStrength, accountID, routerID string) (*routerTypes.NetworkRouter, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netRouter *routerTypes.NetworkRouter
+	result := tx.
+		Take(&netRouter, accountAndIDQueryCondition, accountID, routerID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkRouterNotFoundError(routerID)
+		}
+		log.WithContext(ctx).Errorf("failed to get network router from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network router from store")
+	}
+
+	return netRouter, nil
+}
+
+func (s *SqlStore) CreateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error {
+	if err := s.db.Create(router).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to create network router in store: %v", err)
+		return status.Errorf(status.Internal, "failed to create network router in store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) UpdateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error {
+	result := s.db.
+		Select("*").
+		Where(accountAndIDQueryCondition, router.AccountID, router.ID).
+		Updates(router)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to update network router in store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to update network router in store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewNetworkRouterNotFoundError(router.ID)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteNetworkRouter(ctx context.Context, accountID, routerID string) error {
+	result := s.db.Delete(&routerTypes.NetworkRouter{}, accountAndIDQueryCondition, accountID, routerID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete network router from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete network router from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewNetworkRouterNotFoundError(routerID)
+	}
+
+	return nil
+}
+
+// GetRoutingPeerNetworks returns the distinct network names where the peer is assigned as a routing peer
+// in an enabled network router, either directly or via peer groups.
+func (s *SqlStore) GetRoutingPeerNetworks(_ context.Context, accountID, peerID string) ([]string, error) {
+	var routers []*routerTypes.NetworkRouter
+	if err := s.db.Select("peer, peer_groups, network_id").Where("account_id = ? AND enabled = true", accountID).Find(&routers).Error; err != nil {
+		return nil, status.Errorf(status.Internal, "failed to get enabled routers: %v", err)
+	}
+
+	if len(routers) == 0 {
+		return nil, nil
+	}
+
+	var groupPeers []types.GroupPeer
+	if err := s.db.Select("group_id").Where("account_id = ? AND peer_id = ?", accountID, peerID).Find(&groupPeers).Error; err != nil {
+		return nil, status.Errorf(status.Internal, "failed to get peer group memberships: %v", err)
+	}
+
+	groupSet := make(map[string]struct{}, len(groupPeers))
+	for _, gp := range groupPeers {
+		groupSet[gp.GroupID] = struct{}{}
+	}
+
+	networkIDs := make(map[string]struct{})
+	for _, r := range routers {
+		if r.Peer == peerID {
+			networkIDs[r.NetworkID] = struct{}{}
+		} else if r.Peer == "" {
+			for _, pg := range r.PeerGroups {
+				if _, ok := groupSet[pg]; ok {
+					networkIDs[r.NetworkID] = struct{}{}
+					break
+				}
+			}
+		}
+	}
+
+	if len(networkIDs) == 0 {
+		return nil, nil
+	}
+
+	ids := make([]string, 0, len(networkIDs))
+	for id := range networkIDs {
+		ids = append(ids, id)
+	}
+
+	var networks []*networkTypes.Network
+	if err := s.db.Select("name").Where("account_id = ? AND id IN ?", accountID, ids).Find(&networks).Error; err != nil {
+		return nil, status.Errorf(status.Internal, "failed to get networks: %v", err)
+	}
+
+	names := make([]string, 0, len(networks))
+	for _, n := range networks {
+		names = append(names, n.Name)
+	}
+
+	return names, nil
+}

management/server/store/sql_store_network_router_test.go

@@ -0,0 +1,161 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlStore_GetNetworkRoutersByNetID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	tests := []struct {
+		name          string
+		networkID     string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve routers by existing network ID",
+			networkID:     "ct286bi7qv930dsrrug0",
+			expectedCount: 1,
+		},
+		{
+			name:          "retrieve routers by non-existing network ID",
+			networkID:     "non-existent",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			routers, err := store.GetNetworkRoutersByNetID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
+			require.NoError(t, err)
+			require.Len(t, routers, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_GetNetworkRouterByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name            string
+		networkRouterID string
+		expectError     bool
+	}{
+		{
+			name:            "retrieve existing network router ID",
+			networkRouterID: "ctc20ji7qv9ck2sebc80",
+			expectError:     false,
+		},
+		{
+			name:            "retrieve non-existing network router ID",
+			networkRouterID: "non-existing",
+			expectError:     true,
+		},
+		{
+			name:            "retrieve network with empty router ID",
+			networkRouterID: "",
+			expectError:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			networkRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, tt.networkRouterID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, networkRouter)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, networkRouter)
+				require.Equal(t, tt.networkRouterID, networkRouter.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_CreateNetworkRouter(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	networkID := "ct286bi7qv930dsrrug0"
+
+	netRouter, err := routerTypes.NewNetworkRouter(accountID, networkID, "", []string{"net-router-grp"}, true, 0, true)
+	require.NoError(t, err)
+
+	err = store.CreateNetworkRouter(context.Background(), netRouter)
+	require.NoError(t, err)
+
+	savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, netRouter.ID)
+	require.NoError(t, err)
+	require.Equal(t, netRouter, savedNetRouter)
+}
+
+func TestSqlStore_UpdateNetworkRouter(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	networkID := "ct286bi7qv930dsrrug0"
+	routerID := "ctc20ji7qv9ck2sebc80"
+
+	netRouter := &routerTypes.NetworkRouter{
+		ID:         routerID,
+		AccountID:  accountID,
+		NetworkID:  networkID,
+		Peer:       "",
+		PeerGroups: []string{"net-router-grp"},
+		Masquerade: true,
+		Metric:     42,
+		Enabled:    true,
+	}
+
+	err = store.UpdateNetworkRouter(context.Background(), netRouter)
+	require.NoError(t, err)
+
+	savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, routerID)
+	require.NoError(t, err)
+	require.Equal(t, netRouter, savedNetRouter)
+
+	// Updating a router under a different account must not match any row.
+	netRouter.AccountID = "non-existent-account"
+	err = store.UpdateNetworkRouter(context.Background(), netRouter)
+	require.Error(t, err)
+}
+
+func TestSqlStore_DeleteNetworkRouter(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	netRouterID := "ctc20ji7qv9ck2sebc80"
+
+	err = store.DeleteNetworkRouter(context.Background(), accountID, netRouterID)
+	require.NoError(t, err)
+
+	netRouter, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, netRouterID)
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, sErr.Type())
+	require.Nil(t, netRouter)
+}

management/server/store/sql_store_network_test.go

@@ -0,0 +1,128 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func TestSqlStore_GetAccountNetworks(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name          string
+		accountID     string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve networks by existing account ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectedCount: 1,
+		},
+
+		{
+			name:          "retrieve networks by non-existing account ID",
+			accountID:     "non-existent",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			networks, err := store.GetAccountNetworks(context.Background(), LockingStrengthNone, tt.accountID)
+			require.NoError(t, err)
+			require.Len(t, networks, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_GetNetworkByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name        string
+		networkID   string
+		expectError bool
+	}{
+		{
+			name:        "retrieve existing network ID",
+			networkID:   "ct286bi7qv930dsrrug0",
+			expectError: false,
+		},
+		{
+			name:        "retrieve non-existing network ID",
+			networkID:   "non-existing",
+			expectError: true,
+		},
+		{
+			name:        "retrieve network with empty ID",
+			networkID:   "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			network, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, tt.networkID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, network)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, network)
+				require.Equal(t, tt.networkID, network.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_SaveNetwork(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	network := &networkTypes.Network{
+		ID:        "net-id",
+		AccountID: accountID,
+		Name:      "net",
+	}
+
+	err = store.SaveNetwork(context.Background(), network)
+	require.NoError(t, err)
+
+	savedNet, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, network.ID)
+	require.NoError(t, err)
+	require.Equal(t, network, savedNet)
+}
+
+func TestSqlStore_DeleteNetwork(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	networkID := "ct286bi7qv930dsrrug0"
+
+	err = store.DeleteNetwork(context.Background(), accountID, networkID)
+	require.NoError(t, err)
+
+	network, err := store.GetNetworkByID(context.Background(), LockingStrengthNone, accountID, networkID)
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, sErr.Type())
+	require.Nil(t, network)
+}

management/server/store/sql_store_peer.go

@@ -0,0 +1,781 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net"
+	"net/netip"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) SavePeer(ctx context.Context, 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.transaction(func(tx *gorm.DB) error {
+		// check if peer exists before saving
+		var peerID string
+		result := tx.Model(&nbpeer.Peer{}).Select("id").Take(&peerID, accountAndIDQueryCondition, accountID, peer.ID)
+		if result.Error != nil {
+			if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+				return status.Errorf(status.NotFound, peerNotFoundFMT, peer.ID)
+			}
+			return result.Error
+		}
+
+		if peerID == "" {
+			return status.Errorf(status.NotFound, peerNotFoundFMT, peer.ID)
+		}
+
+		result = tx.Model(&nbpeer.Peer{}).Where(accountAndIDQueryCondition, accountID, peer.ID).Save(peerCopy)
+		if result.Error != nil {
+			return status.Errorf(status.Internal, "failed to save peer to store: %v", result.Error)
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *SqlStore) SavePeerStatus(ctx context.Context, accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
+	var peerCopy nbpeer.Peer
+	peerCopy.Status = &peerStatus
+
+	fieldsToUpdate := []string{
+		"peer_status_last_seen", "peer_status_session_started_at",
+		"peer_status_connected", "peer_status_login_expired",
+		"peer_status_requires_approval",
+	}
+	result := s.db.Model(&nbpeer.Peer{}).
+		Select(fieldsToUpdate).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Updates(&peerCopy)
+	if result.Error != nil {
+		return status.Errorf(status.Internal, "failed to save peer status to store: %v", result.Error)
+	}
+
+	if result.RowsAffected == 0 {
+		return status.Errorf(status.NotFound, peerNotFoundFMT, peerID)
+	}
+
+	return nil
+}
+
+// MarkPeerConnectedIfNewerSession is an atomic optimistic-locked update.
+// The peer is marked connected with the given session token only when
+// the stored SessionStartedAt is strictly smaller than the incoming
+// one — equivalently, when no newer stream has already taken ownership.
+// The sentinel zero (set on peer creation or after a disconnect) counts
+// as the smallest possible token. This is the write half of the
+// fencing protocol described on PeerStatus.SessionStartedAt.
+//
+// The post-write side effects in the caller — geo lookup,
+// schedulePeerLoginExpiration, checkAndSchedulePeerInactivityExpiration,
+// OnPeersUpdated — all run AFTER this method returns and are deliberately
+// outside the database write so they cannot extend the row-lock window.
+//
+// LastSeen is set to the database's clock (CURRENT_TIMESTAMP) at the
+// moment the row is written. The caller never supplies LastSeen because
+// the value would otherwise drift under lock contention — a Go-side
+// time.Now() taken before the write can land minutes later than the
+// actual UPDATE under load, which previously caused real ordering bugs.
+func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error) {
+	result := s.db.WithContext(ctx).
+		Model(&nbpeer.Peer{}).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Where("peer_status_session_started_at < ?", newSessionStartedAt).
+		Updates(map[string]any{
+			"peer_status_connected":          true,
+			"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
+			"peer_status_session_started_at": newSessionStartedAt,
+			"peer_status_login_expired":      false,
+		})
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
+	}
+	return result.RowsAffected > 0, nil
+}
+
+// MarkPeerDisconnectedIfSameSession is an atomic optimistic-locked update.
+// The peer is marked disconnected only when the stored SessionStartedAt
+// matches the incoming token — meaning the stream that owns the current
+// session is the one ending. If a newer stream has already replaced the
+// session, the update is skipped. LastSeen is set to CURRENT_TIMESTAMP at
+// write time; see MarkPeerConnectedIfNewerSession for the rationale.
+//
+// A zero sessionStartedAt is rejected at the call site; the underlying
+// WHERE on equality would otherwise match every never-connected peer.
+func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
+	if sessionStartedAt == 0 {
+		return false, nil
+	}
+	result := s.db.WithContext(ctx).
+		Model(&nbpeer.Peer{}).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Where("peer_status_session_started_at = ?", sessionStartedAt).
+		Updates(map[string]any{
+			"peer_status_connected":          false,
+			"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
+			"peer_status_session_started_at": int64(0),
+		})
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "mark peer disconnected: %v", result.Error)
+	}
+	return result.RowsAffected > 0, nil
+}
+
+// ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
+func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
+	result := s.db.Model(&nbpeer.Peer{}).
+		Where("account_id = ? AND peer_status_requires_approval = ?", accountID, true).
+		Update("peer_status_requires_approval", false)
+	if result.Error != nil {
+		return 0, status.Errorf(status.Internal, "failed to approve pending account peers: %v", result.Error)
+	}
+
+	return int(result.RowsAffected), nil
+}
+
+// RefreshPeerLastSeen updates only peer_status_last_seen. Every other status
+// column is left untouched: peer_status_connected and
+// peer_status_session_started_at belong to the sync stream that owns the
+// session, and a blind write here would corrupt the fencing
+// MarkPeerConnectedIfNewerSession relies on.
+//
+// LastSeen comes from the database clock for the same reason it does there: a
+// Go-side timestamp is taken before the write and can land after a connect that
+// used CURRENT_TIMESTAMP, dragging the column backwards.
+//
+// staleBefore carries the caller's throttle into the same statement, so
+// concurrent requests for one peer collapse into a single write instead of
+// each racing on its own stale read. The column is nullable — Status is an
+// embedded pointer, so a peer stored without one leaves it NULL — and NULL
+// loses every comparison, hence the explicit branch for a peer never seen.
+func (s *SqlStore) RefreshPeerLastSeen(ctx context.Context, accountID, peerID string, staleBefore time.Time) (bool, error) {
+	result := s.db.WithContext(ctx).
+		Model(&nbpeer.Peer{}).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Where("(peer_status_last_seen IS NULL OR peer_status_last_seen < ?)", staleBefore).
+		Update("peer_status_last_seen", gorm.Expr("CURRENT_TIMESTAMP"))
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "refresh peer last seen: %v", result.Error)
+	}
+
+	return result.RowsAffected > 0, nil
+}
+
+func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Peer, error) {
+	const query = `SELECT id, account_id, key, ip, name, dns_label, user_id, ssh_key, ssh_enabled, login_expiration_enabled,
+	inactivity_expiration_enabled, last_login, created_at, ephemeral, extra_dns_labels, allow_extra_dns_labels, meta_hostname,
+	meta_go_os, meta_kernel, meta_core, meta_platform, meta_os, meta_os_version, meta_wt_version, meta_ui_version,
+	meta_kernel_version, meta_network_addresses, meta_system_serial_number, meta_system_product_name, meta_system_manufacturer,
+	meta_environment, meta_flags, meta_files, meta_capabilities, peer_status_last_seen, peer_status_session_started_at,
+	peer_status_connected, peer_status_login_expired, peer_status_requires_approval, location_connection_ip,
+	location_country_code, location_city_name, location_geo_name_id, proxy_meta_embedded, proxy_meta_cluster, ipv6, meta_sync_message_version
+	FROM peers WHERE account_id = $1`
+	rows, err := s.pool.Query(ctx, query, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	peers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nbpeer.Peer, error) {
+		var p nbpeer.Peer
+		p.Status = &nbpeer.PeerStatus{}
+		var (
+			lastLogin, createdAt                                                                            sql.NullTime
+			sshEnabled, loginExpirationEnabled, inactivityExpirationEnabled, ephemeral, allowExtraDNSLabels sql.NullBool
+			peerStatusLastSeen                                                                              sql.NullTime
+			peerStatusSessionStartedAt                                                                      sql.NullInt64
+			peerStatusConnected, peerStatusLoginExpired, peerStatusRequiresApproval, proxyEmbedded          sql.NullBool
+			ip, extraDNS, netAddr, env, flags, files, capabilities, connIP, ipv6                            []byte
+			metaHostname, metaGoOS, metaKernel, metaCore, metaPlatform                                      sql.NullString
+			metaOS, metaOSVersion, metaWtVersion, metaUIVersion, metaKernelVersion                          sql.NullString
+			metaSystemSerialNumber, metaSystemProductName, metaSystemManufacturer                           sql.NullString
+			locationCountryCode, locationCityName, proxyCluster                                             sql.NullString
+			locationGeoNameID                                                                               sql.NullInt64
+			metaSyncMessageVersion                                                                          sql.NullInt32
+		)
+
+		err := row.Scan(&p.ID, &p.AccountID, &p.Key, &ip, &p.Name, &p.DNSLabel, &p.UserID, &p.SSHKey, &sshEnabled,
+			&loginExpirationEnabled, &inactivityExpirationEnabled, &lastLogin, &createdAt, &ephemeral, &extraDNS,
+			&allowExtraDNSLabels, &metaHostname, &metaGoOS, &metaKernel, &metaCore, &metaPlatform,
+			&metaOS, &metaOSVersion, &metaWtVersion, &metaUIVersion, &metaKernelVersion, &netAddr,
+			&metaSystemSerialNumber, &metaSystemProductName, &metaSystemManufacturer, &env, &flags, &files, &capabilities,
+			&peerStatusLastSeen, &peerStatusSessionStartedAt, &peerStatusConnected, &peerStatusLoginExpired,
+			&peerStatusRequiresApproval, &connIP, &locationCountryCode, &locationCityName, &locationGeoNameID,
+			&proxyEmbedded, &proxyCluster, &ipv6, &metaSyncMessageVersion)
+
+		if err == nil {
+			if lastLogin.Valid {
+				p.LastLogin = &lastLogin.Time
+			}
+			if createdAt.Valid {
+				p.CreatedAt = createdAt.Time
+			}
+			if sshEnabled.Valid {
+				p.SSHEnabled = sshEnabled.Bool
+			}
+			if loginExpirationEnabled.Valid {
+				p.LoginExpirationEnabled = loginExpirationEnabled.Bool
+			}
+			if inactivityExpirationEnabled.Valid {
+				p.InactivityExpirationEnabled = inactivityExpirationEnabled.Bool
+			}
+			if ephemeral.Valid {
+				p.Ephemeral = ephemeral.Bool
+			}
+			if allowExtraDNSLabels.Valid {
+				p.AllowExtraDNSLabels = allowExtraDNSLabels.Bool
+			}
+			if peerStatusLastSeen.Valid {
+				p.Status.LastSeen = peerStatusLastSeen.Time
+			}
+			if peerStatusSessionStartedAt.Valid {
+				p.Status.SessionStartedAt = peerStatusSessionStartedAt.Int64
+			}
+			if peerStatusConnected.Valid {
+				p.Status.Connected = peerStatusConnected.Bool
+			}
+			if peerStatusLoginExpired.Valid {
+				p.Status.LoginExpired = peerStatusLoginExpired.Bool
+			}
+			if peerStatusRequiresApproval.Valid {
+				p.Status.RequiresApproval = peerStatusRequiresApproval.Bool
+			}
+			if metaHostname.Valid {
+				p.Meta.Hostname = metaHostname.String
+			}
+			if metaGoOS.Valid {
+				p.Meta.GoOS = metaGoOS.String
+			}
+			if metaKernel.Valid {
+				p.Meta.Kernel = metaKernel.String
+			}
+			if metaCore.Valid {
+				p.Meta.Core = metaCore.String
+			}
+			if metaPlatform.Valid {
+				p.Meta.Platform = metaPlatform.String
+			}
+			if metaOS.Valid {
+				p.Meta.OS = metaOS.String
+			}
+			if metaOSVersion.Valid {
+				p.Meta.OSVersion = metaOSVersion.String
+			}
+			if metaWtVersion.Valid {
+				p.Meta.WtVersion = metaWtVersion.String
+			}
+			if metaUIVersion.Valid {
+				p.Meta.UIVersion = metaUIVersion.String
+			}
+			if metaKernelVersion.Valid {
+				p.Meta.KernelVersion = metaKernelVersion.String
+			}
+			if metaSystemSerialNumber.Valid {
+				p.Meta.SystemSerialNumber = metaSystemSerialNumber.String
+			}
+			if metaSystemProductName.Valid {
+				p.Meta.SystemProductName = metaSystemProductName.String
+			}
+			if metaSystemManufacturer.Valid {
+				p.Meta.SystemManufacturer = metaSystemManufacturer.String
+			}
+			if locationCountryCode.Valid {
+				p.Location.CountryCode = locationCountryCode.String
+			}
+			if locationCityName.Valid {
+				p.Location.CityName = locationCityName.String
+			}
+			if locationGeoNameID.Valid {
+				p.Location.GeoNameID = uint(locationGeoNameID.Int64)
+			}
+			if proxyEmbedded.Valid {
+				p.ProxyMeta.Embedded = proxyEmbedded.Bool
+			}
+			if proxyCluster.Valid {
+				p.ProxyMeta.Cluster = proxyCluster.String
+			}
+			if ip != nil {
+				_ = json.Unmarshal(ip, &p.IP)
+			}
+			if ipv6 != nil {
+				_ = json.Unmarshal(ipv6, &p.IPv6)
+			}
+			if extraDNS != nil {
+				_ = json.Unmarshal(extraDNS, &p.ExtraDNSLabels)
+			}
+			if netAddr != nil {
+				_ = json.Unmarshal(netAddr, &p.Meta.NetworkAddresses)
+			}
+			if env != nil {
+				_ = json.Unmarshal(env, &p.Meta.Environment)
+			}
+			if flags != nil {
+				_ = json.Unmarshal(flags, &p.Meta.Flags)
+			}
+			if files != nil {
+				_ = json.Unmarshal(files, &p.Meta.Files)
+			}
+			if capabilities != nil {
+				_ = json.Unmarshal(capabilities, &p.Meta.Capabilities)
+			}
+			if connIP != nil {
+				_ = json.Unmarshal(connIP, &p.Location.ConnectionIP)
+			}
+			if metaSyncMessageVersion.Valid {
+				p.Meta.SyncMessageVersion = int(metaSyncMessageVersion.Int32)
+			}
+		}
+		return p, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return peers, nil
+}
+
+func (s *SqlStore) GetAccountByPeerID(ctx context.Context, peerID string) (*types.Account, error) {
+	var peer nbpeer.Peer
+	result := s.db.Select("account_id").Take(&peer, idQueryCondition, peerID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
+		}
+		return nil, status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	if peer.AccountID == "" {
+		return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
+	}
+
+	return s.GetAccount(ctx, peer.AccountID)
+}
+
+func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (*types.Account, error) {
+	var peer nbpeer.Peer
+	result := s.db.Select("account_id").Take(&peer, GetKeyQueryCondition(s), peerKey)
+
+	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.NewGetAccountFromStoreError(result.Error)
+	}
+
+	if peer.AccountID == "" {
+		return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
+	}
+
+	return s.GetAccount(ctx, peer.AccountID)
+}
+
+func (s *SqlStore) GetAccountIDByPeerPubKey(ctx context.Context, peerKey string) (string, error) {
+	var peer nbpeer.Peer
+	var accountID string
+	result := s.db.Model(&peer).Select("account_id").Where(GetKeyQueryCondition(s), peerKey).Take(&accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
+		}
+		return "", status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return accountID, nil
+}
+
+func (s *SqlStore) GetAccountIDByPeerID(ctx context.Context, lockStrength LockingStrength, peerID string) (string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var accountID string
+	result := tx.Model(&nbpeer.Peer{}).
+		Select("account_id").Where(idQueryCondition, peerID).Take(&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) GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountID string) ([]netip.Addr, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var ipJSONStrings []string
+
+	result := tx.Model(&nbpeer.Peer{}).
+		Where("account_id = ?", accountID).
+		Pluck("ip", &ipJSONStrings)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "no peers found for the account")
+		}
+		return nil, status.Errorf(status.Internal, "issue getting IPs from store: %s", result.Error)
+	}
+
+	ips := make([]netip.Addr, len(ipJSONStrings))
+	for i, ipJSON := range ipJSONStrings {
+		var ip netip.Addr
+		if err := json.Unmarshal([]byte(ipJSON), &ip); err != nil {
+			return nil, status.Errorf(status.Internal, "issue parsing IP JSON from store")
+		}
+		ips[i] = ip.Unmap()
+	}
+
+	return ips, nil
+}
+
+func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string, dnsLabel string) ([]string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var labels []string
+	result := tx.Model(&nbpeer.Peer{}).
+		Where("account_id = ? AND dns_label LIKE ?", accountID, dnsLabel+"%").
+		Pluck("dns_label", &labels)
+
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "no peers found for the account")
+		}
+		log.WithContext(ctx).Errorf("error when getting dns labels from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "issue getting dns labels from store: %s", result.Error)
+	}
+
+	return labels, nil
+}
+
+func (s *SqlStore) GetPeerByPeerPubKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (*nbpeer.Peer, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peer nbpeer.Peer
+	result := tx.Take(&peer, GetKeyQueryCondition(s), peerKey)
+
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewPeerNotFoundError(peerKey)
+		}
+		return nil, status.Errorf(status.Internal, "issue getting peer from store: %s", result.Error)
+	}
+
+	return &peer, nil
+}
+
+// GetAccountPeers retrieves peers for an account.
+func (s *SqlStore) GetAccountPeers(ctx context.Context, lockStrength LockingStrength, accountID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error) {
+	var peers []*nbpeer.Peer
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+	query := tx.Where(accountIDCondition, accountID)
+
+	if nameFilter != "" {
+		query = query.Where("name LIKE ?", "%"+nameFilter+"%")
+	}
+	if ipFilter != "" {
+		query = query.Where("ip LIKE ? OR ipv6 LIKE ?", "%"+ipFilter+"%", "%"+ipFilter+"%")
+	}
+
+	if err := query.Find(&peers).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) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peers []*nbpeer.Peer
+
+	// Exclude peers added via setup keys, as they are not user-specific and have an empty user_id.
+	if userID == "" {
+		return peers, nil
+	}
+
+	result := tx.
+		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 {
+	if err := s.db.Create(peer).Error; err != nil {
+		return status.Errorf(status.Internal, "issue adding peer to account: %s", err)
+	}
+
+	return nil
+}
+
+// GetPeerByID retrieves a peer by its ID and account ID.
+func (s *SqlStore) GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID, peerID string) (*nbpeer.Peer, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peer *nbpeer.Peer
+	result := tx.
+		Take(&peer, accountAndIDQueryCondition, accountID, peerID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewPeerNotFoundError(peerID)
+		}
+		return nil, status.Errorf(status.Internal, "failed to get peer from store")
+	}
+
+	return peer, nil
+}
+
+// GetPeersByIDs retrieves peers by their IDs and account ID.
+func (s *SqlStore) GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peers []*nbpeer.Peer
+	result := tx.Find(&peers, accountAndIDsQueryCondition, accountID, peerIDs)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get peers by ID's from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get peers by ID's from the store")
+	}
+
+	peersMap := make(map[string]*nbpeer.Peer)
+	for _, peer := range peers {
+		peersMap[peer.ID] = peer
+	}
+
+	return peersMap, nil
+}
+
+// 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) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peers []*nbpeer.Peer
+	result := tx.
+		Where("login_expiration_enabled = ? AND peer_status_login_expired != ? AND user_id IS NOT NULL AND user_id != ''", true, 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) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peers []*nbpeer.Peer
+	result := tx.
+		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) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var allEphemeralPeers, batchPeers []*nbpeer.Peer
+	result := tx.
+		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, accountID string, peerID string) error {
+	result := s.db.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) GetPeerByIP(ctx context.Context, lockStrength LockingStrength, accountID string, ip net.IP) (*nbpeer.Peer, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	column := "ip"
+	if ip.To4() == nil {
+		column = "ipv6"
+	}
+	jsonValue := fmt.Sprintf(`"%s"`, ip.String())
+
+	var peer nbpeer.Peer
+	result := tx.
+		Take(&peer, fmt.Sprintf("account_id = ? AND %s = ?", column), accountID, jsonValue)
+	if result.Error != nil {
+		// A tunnel-IP miss is an expected outcome (e.g. the proxy's
+		// ValidateTunnelPeer probing an address that isn't in the
+		// account roster); surface it as NotFound so callers can tell
+		// it apart from a real store failure.
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "peer with ip %s not found", ip.String())
+		}
+		return nil, status.Errorf(status.Internal, "failed to get peer from store")
+	}
+
+	return &peer, nil
+}
+
+func (s *SqlStore) GetPeerIdByLabel(ctx context.Context, lockStrength LockingStrength, accountID string, hostname string) (string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peerID string
+	result := tx.Model(&nbpeer.Peer{}).
+		Select("id").
+		// Where(" = ?", hostname).
+		Where("account_id = ? AND dns_label = ?", accountID, hostname).
+		Limit(1).
+		Scan(&peerID)
+
+	if peerID == "" {
+		return "", gorm.ErrRecordNotFound
+	}
+
+	return peerID, result.Error
+}
+
+// GetEmbeddedProxyPeerIDsByCluster returns peer IDs of all embedded proxy peers
+// in the account, grouped by their ProxyCluster. The map is nil when no embedded
+// proxy peers exist.
+func (s *SqlStore) GetEmbeddedProxyPeerIDsByCluster(ctx context.Context, accountID string) (map[string][]string, error) {
+	type row struct {
+		ID      string
+		Cluster string
+	}
+	var rows []row
+	result := s.db.Model(&nbpeer.Peer{}).
+		Select("id, proxy_meta_cluster AS cluster").
+		Where("account_id = ? AND proxy_meta_embedded = ?", accountID, true).
+		Scan(&rows)
+	if result.Error != nil {
+		return nil, status.Errorf(status.Internal, "failed to get embedded proxy peers: %s", result.Error)
+	}
+
+	out := make(map[string][]string, len(rows))
+	for _, r := range rows {
+		out[r.Cluster] = append(out[r.Cluster], r.ID)
+	}
+	return out, nil
+}
+
+func (s *SqlStore) GetUserIDByPeerKey(ctx context.Context, lockStrength LockingStrength, peerKey string) (string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var userID string
+	result := tx.Model(&nbpeer.Peer{}).
+		Select("user_id").
+		Take(&userID, GetKeyQueryCondition(s), peerKey)
+
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "peer not found: index lookup failed")
+		}
+		return "", status.Errorf(status.Internal, "failed to get user ID by peer key")
+	}
+
+	return userID, nil
+}
+
+func (s *SqlStore) GetPeerIDByKey(ctx context.Context, lockStrength LockingStrength, key string) (string, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peerID string
+	result := tx.Model(&nbpeer.Peer{}).
+		Select("id").
+		Where(GetKeyQueryCondition(s), key).
+		Limit(1).
+		Scan(&peerID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get peer ID by key: %s", result.Error)
+		return "", status.Errorf(status.Internal, "failed to get peer ID by key")
+	}
+
+	return peerID, nil
+}

management/server/store/sql_store_peer_test.go

@@ -0,0 +1,901 @@
+package store
+
+import (
+	"context"
+	"encoding/binary"
+	"fmt"
+	"net"
+	"net/netip"
+	"reflect"
+	"sort"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/management/server/util"
+	"github.com/netbirdio/netbird/shared/management/status"
+	"github.com/netbirdio/netbird/shared/testing_helpers"
+)
+
+// TestSqlStore_GetPeerByIP_NotFound pins the not-found semantics the
+// proxy's ValidateTunnelPeer relies on: a tunnel-IP that isn't in the
+// account roster must surface as a NotFound error (not a generic
+// Internal) so callers can distinguish an expected miss from a real
+// store failure. A known IP still resolves.
+func TestSqlStore_GetPeerByIP_NotFound(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		const accountID = "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		peer, err := store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("192.168.0.0"))
+		require.NoError(t, err, "known tunnel IP must resolve")
+		require.NotNil(t, peer)
+
+		_, err = store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("100.65.0.99"))
+		require.Error(t, err, "unknown tunnel IP must error")
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok, "error must be a status error")
+		require.Equal(t, status.NotFound, parsedErr.Type(), "tunnel-IP miss must be NotFound, not Internal")
+	})
+}
+
+func TestSqlStore_SavePeer(t *testing.T) {
+	populateFields := testing_helpers.NewPopulateFields()
+
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
+		require.NoError(t, err)
+
+		metadata := nbpeer.PeerSystemMeta{}
+		reflectedMetadata := reflect.ValueOf(&metadata).Elem()
+
+		numOfFields, err := populateFields.PopulateAll(reflectedMetadata)
+		assert.NoError(t, err)
+		assert.Equal(t, 33, numOfFields)
+
+		// save status of non-existing peer
+		peer := &nbpeer.Peer{
+			Key:       "peerkey",
+			ID:        "testpeer",
+			IP:        netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1"),
+			Meta:      metadata, //nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
+			Name:      "peer name",
+			Status:    &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+			CreatedAt: time.Now().UTC(),
+		}
+		ctx := context.Background()
+		err = store.SavePeer(ctx, account.Id, peer)
+		assert.Error(t, err)
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok)
+		require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+
+		// save new status of existing peer
+		account.Peers[peer.ID] = peer
+
+		err = store.SaveAccount(context.Background(), account)
+		require.NoError(t, err)
+
+		updatedPeer := peer.Copy()
+		updatedPeer.Status.Connected = false
+		updatedPeer.Meta.Hostname = "updatedpeer"
+
+		err = store.SavePeer(ctx, account.Id, updatedPeer)
+		require.NoError(t, err)
+
+		account, err = store.GetAccount(context.Background(), account.Id)
+		require.NoError(t, err)
+
+		actual := account.Peers[peer.ID]
+		assert.Equal(t, updatedPeer.Meta, actual.Meta)
+		assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
+		assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
+		assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
+		assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+	})
+}
+
+func TestSqlStore_SavePeerStatus(t *testing.T) {
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	assert.NoError(t, err)
+
+	account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
+	require.NoError(t, err)
+
+	// save status of non-existing peer
+	newStatus := nbpeer.PeerStatus{Connected: false, LastSeen: time.Now().UTC()}
+	err = store.SavePeerStatus(context.Background(), account.Id, "non-existing-peer", newStatus)
+	assert.Error(t, err)
+	parsedErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+
+	// save new status of existing peer
+	account.Peers["testpeer"] = &nbpeer.Peer{
+		Key:    "peerkey",
+		ID:     "testpeer",
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
+		Meta:   nbpeer.PeerSystemMeta{},
+		Name:   "peer name",
+		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	err = store.SavePeerStatus(context.Background(), account.Id, "testpeer", newStatus)
+	require.NoError(t, err)
+
+	account, err = store.GetAccount(context.Background(), account.Id)
+	require.NoError(t, err)
+
+	actual := account.Peers["testpeer"].Status
+	assert.Equal(t, newStatus.Connected, actual.Connected)
+	assert.Equal(t, newStatus.LoginExpired, actual.LoginExpired)
+	assert.Equal(t, newStatus.RequiresApproval, actual.RequiresApproval)
+	assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+
+	newStatus.Connected = true
+
+	err = store.SavePeerStatus(context.Background(), account.Id, "testpeer", newStatus)
+	require.NoError(t, err)
+
+	account, err = store.GetAccount(context.Background(), account.Id)
+	require.NoError(t, err)
+
+	actual = account.Peers["testpeer"].Status
+	assert.Equal(t, newStatus.Connected, actual.Connected)
+	assert.Equal(t, newStatus.LoginExpired, actual.LoginExpired)
+	assert.Equal(t, newStatus.RequiresApproval, actual.RequiresApproval)
+	assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+}
+
+func TestSqlite_GetTakenIPs(t *testing.T) {
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	defer cleanup()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	_, err = store.GetAccount(context.Background(), existingAccountID)
+	require.NoError(t, err)
+
+	takenIPs, err := store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
+	require.NoError(t, err)
+	assert.Equal(t, []netip.Addr{}, takenIPs)
+
+	peer1 := &nbpeer.Peer{
+		ID:        "peer1",
+		AccountID: existingAccountID,
+		Key:       "key1",
+		DNSLabel:  "peer1",
+		IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
+	}
+	err = store.AddPeerToAccount(context.Background(), peer1)
+	require.NoError(t, err)
+
+	takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
+	require.NoError(t, err)
+	ip1 := netip.AddrFrom4([4]byte{1, 1, 1, 1})
+	assert.Equal(t, []netip.Addr{ip1}, takenIPs)
+
+	peer2 := &nbpeer.Peer{
+		ID:        "peer1second",
+		AccountID: existingAccountID,
+		Key:       "key2",
+		DNSLabel:  "peer1-1",
+		IP:        netip.AddrFrom4([4]byte{2, 2, 2, 2}),
+		IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
+	}
+	err = store.AddPeerToAccount(context.Background(), peer2)
+	require.NoError(t, err)
+
+	takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
+	require.NoError(t, err)
+	ip2 := netip.AddrFrom4([4]byte{2, 2, 2, 2})
+	assert.Equal(t, []netip.Addr{ip1, ip2}, takenIPs)
+}
+
+func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+		peerHostname := "peer1"
+
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
+
+		labels, err := store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
+		require.NoError(t, err)
+		assert.Equal(t, []string{}, labels)
+
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			Key:       "key1",
+			DNSLabel:  "peer1",
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
+		}
+		err = store.AddPeerToAccount(context.Background(), peer1)
+		require.NoError(t, err)
+
+		labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
+		require.NoError(t, err)
+		assert.Equal(t, []string{"peer1"}, labels)
+
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			Key:       "key2",
+			DNSLabel:  "peer1-1",
+			IP:        netip.AddrFrom4([4]byte{2, 2, 2, 2}),
+			IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
+		}
+		err = store.AddPeerToAccount(context.Background(), peer2)
+		require.NoError(t, err)
+
+		labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthNone, existingAccountID, peerHostname)
+		require.NoError(t, err)
+
+		expected := []string{"peer1", "peer1-1"}
+		sort.Strings(expected)
+		sort.Strings(labels)
+		assert.Equal(t, expected, labels)
+	})
+}
+
+func Test_AddPeerWithSameDnsLabel(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
+
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			Key:       "key1",
+			DNSLabel:  "peer1.domain.test",
+		}
+		err = store.AddPeerToAccount(context.Background(), peer1)
+		require.NoError(t, err)
+
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			Key:       "key2",
+			DNSLabel:  "peer1.domain.test",
+		}
+		err = store.AddPeerToAccount(context.Background(), peer2)
+		require.Error(t, err)
+	})
+}
+
+func Test_AddPeerWithSameIP(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
+
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			Key:       "key1",
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
+		}
+		err = store.AddPeerToAccount(context.Background(), peer1)
+		require.NoError(t, err)
+
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			Key:       "key2",
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
+		}
+		err = store.AddPeerToAccount(context.Background(), peer2)
+		require.Error(t, err)
+	})
+}
+
+func TestSqlStore_GetPeerByID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name        string
+		peerID      string
+		expectError bool
+	}{
+		{
+			name:        "retrieve existing peer",
+			peerID:      "cfefqs706sqkneg59g4g",
+			expectError: false,
+		},
+		{
+			name:        "retrieve non-existing peer",
+			peerID:      "non-existing",
+			expectError: true,
+		},
+		{
+			name:        "retrieve with empty peer ID",
+			peerID:      "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, tt.peerID)
+			if tt.expectError {
+				require.Error(t, err)
+				sErr, ok := status.FromError(err)
+				require.True(t, ok)
+				require.Equal(t, sErr.Type(), status.NotFound)
+				require.Nil(t, peer)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, peer)
+				require.Equal(t, tt.peerID, peer.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_GetPeersByIDs(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	tests := []struct {
+		name          string
+		peerIDs       []string
+		expectedCount int
+	}{
+		{
+			name:          "retrieve existing peers by existing IDs",
+			peerIDs:       []string{"cfefqs706sqkneg59g4g", "cfeg6sf06sqkneg59g50"},
+			expectedCount: 2,
+		},
+		{
+			name:          "empty peer IDs list",
+			peerIDs:       []string{},
+			expectedCount: 0,
+		},
+		{
+			name:          "non-existing peer IDs",
+			peerIDs:       []string{"nonexistent1", "nonexistent2"},
+			expectedCount: 0,
+		},
+		{
+			name:          "mixed existing and non-existing peer IDs",
+			peerIDs:       []string{"cfeg6sf06sqkneg59g50", "nonexistent"},
+			expectedCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetPeersByIDs(context.Background(), LockingStrengthNone, accountID, tt.peerIDs)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_AddPeerToAccount(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_policy_migrate.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	peer := &nbpeer.Peer{
+		ID:        "peer1",
+		AccountID: accountID,
+		Key:       "key",
+		IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
+		Meta: nbpeer.PeerSystemMeta{
+			Hostname:  "hostname",
+			GoOS:      "linux",
+			Kernel:    "Linux",
+			Core:      "21.04",
+			Platform:  "x86_64",
+			OS:        "Ubuntu",
+			WtVersion: "development",
+			UIVersion: "development",
+		},
+		Name:     "peer.test",
+		DNSLabel: "peer",
+		Status: &nbpeer.PeerStatus{
+			LastSeen:         time.Now().UTC(),
+			Connected:        true,
+			LoginExpired:     false,
+			RequiresApproval: false,
+		},
+		SSHKey:                      "ssh-key",
+		SSHEnabled:                  false,
+		LoginExpirationEnabled:      true,
+		InactivityExpirationEnabled: false,
+		LastLogin:                   util.ToPtr(time.Now().UTC()),
+		CreatedAt:                   time.Now().UTC(),
+		Ephemeral:                   true,
+	}
+	err = store.AddPeerToAccount(context.Background(), peer)
+	require.NoError(t, err, "failed to add peer to account")
+
+	storedPeer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, peer.ID)
+	require.NoError(t, err, "failed to get peer")
+
+	assert.Equal(t, peer.ID, storedPeer.ID)
+	assert.Equal(t, peer.AccountID, storedPeer.AccountID)
+	assert.Equal(t, peer.Key, storedPeer.Key)
+	assert.Equal(t, peer.IP.String(), storedPeer.IP.String())
+	assert.Equal(t, peer.Meta, storedPeer.Meta)
+	assert.Equal(t, peer.Name, storedPeer.Name)
+	assert.Equal(t, peer.DNSLabel, storedPeer.DNSLabel)
+	assert.Equal(t, peer.SSHKey, storedPeer.SSHKey)
+	assert.Equal(t, peer.SSHEnabled, storedPeer.SSHEnabled)
+	assert.Equal(t, peer.LoginExpirationEnabled, storedPeer.LoginExpirationEnabled)
+	assert.Equal(t, peer.InactivityExpirationEnabled, storedPeer.InactivityExpirationEnabled)
+	assert.WithinDurationf(t, peer.GetLastLogin(), storedPeer.GetLastLogin().UTC(), time.Millisecond, "LastLogin should be equal")
+	assert.WithinDurationf(t, peer.CreatedAt, storedPeer.CreatedAt.UTC(), time.Millisecond, "CreatedAt should be equal")
+	assert.Equal(t, peer.Ephemeral, storedPeer.Ephemeral)
+	assert.Equal(t, peer.Status.Connected, storedPeer.Status.Connected)
+	assert.Equal(t, peer.Status.LoginExpired, storedPeer.Status.LoginExpired)
+	assert.Equal(t, peer.Status.RequiresApproval, storedPeer.Status.RequiresApproval)
+	assert.WithinDurationf(t, peer.Status.LastSeen, storedPeer.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+}
+
+func TestSqlStore_GetAccountPeers(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name          string
+		accountID     string
+		nameFilter    string
+		ipFilter      string
+		expectedCount int
+	}{
+		{
+			name:          "should retrieve peers for an existing account ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectedCount: 5,
+		},
+		{
+			name:          "should return no peers for a non-existing account ID",
+			accountID:     "nonexistent",
+			expectedCount: 0,
+		},
+		{
+			name:          "should return no peers for an empty account ID",
+			accountID:     "",
+			expectedCount: 0,
+		},
+		{
+			name:          "should filter peers by name",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			nameFilter:    "expiredhost",
+			expectedCount: 1,
+		},
+		{
+			name:          "should filter peers by partial name",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			nameFilter:    "host",
+			expectedCount: 4,
+		},
+		{
+			name:          "should filter peers by ip",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			ipFilter:      "100.64.39.54",
+			expectedCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetAccountPeers(context.Background(), LockingStrengthNone, tt.accountID, tt.nameFilter, tt.ipFilter)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+		})
+	}
+
+}
+
+func TestSqlStore_GetAccountPeersWithExpiration(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name            string
+		accountID       string
+		expectedCount   int
+		expectedPeerIDs []string
+	}{
+		{
+			name:            "should retrieve only non-expired peers with expiration enabled",
+			accountID:       "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectedCount:   1,
+			expectedPeerIDs: []string{"notexpired01"},
+		},
+		{
+			name:          "should return no peers with expiration for a non-existing account ID",
+			accountID:     "nonexistent",
+			expectedCount: 0,
+		},
+		{
+			name:          "should return no peers with expiration for a empty account ID",
+			accountID:     "",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetAccountPeersWithExpiration(context.Background(), LockingStrengthNone, tt.accountID)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+			for i, peer := range peers {
+				assert.Equal(t, tt.expectedPeerIDs[i], peer.ID)
+			}
+		})
+	}
+}
+
+func TestSqlStore_GetAccountPeersWithExpiration_ExcludesAlreadyExpired(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	peers, err := store.GetAccountPeersWithExpiration(context.Background(), LockingStrengthNone, accountID)
+	require.NoError(t, err)
+
+	// Verify the already-expired peer (cg05lnblo1hkg2j514p0) is not returned
+	for _, peer := range peers {
+		assert.NotEqual(t, "cg05lnblo1hkg2j514p0", peer.ID, "already expired peer should not be returned")
+		assert.False(t, peer.Status.LoginExpired, "returned peers should not have LoginExpired set")
+	}
+}
+
+func TestSqlStore_GetAccountPeersWithInactivity(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name          string
+		accountID     string
+		expectedCount int
+	}{
+		{
+			name:          "should retrieve peers with inactivity for an existing account ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			expectedCount: 1,
+		},
+		{
+			name:          "should return no peers with inactivity for a non-existing account ID",
+			accountID:     "nonexistent",
+			expectedCount: 0,
+		},
+		{
+			name:          "should return no peers with inactivity for an empty account ID",
+			accountID:     "",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetAccountPeersWithInactivity(context.Background(), LockingStrengthNone, tt.accountID)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_GetAllEphemeralPeers(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/storev1.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	peers, err := store.GetAllEphemeralPeers(context.Background(), LockingStrengthNone)
+	require.NoError(t, err)
+	require.Len(t, peers, 1)
+	require.True(t, peers[0].Ephemeral)
+}
+
+func TestSqlStore_GetUserPeers(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name          string
+		accountID     string
+		userID        string
+		expectedCount int
+	}{
+		{
+			name:          "should retrieve peers for existing account ID and user ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			userID:        "f4f6d672-63fb-11ec-90d6-0242ac120003",
+			expectedCount: 1,
+		},
+		{
+			name:          "should return no peers for non-existing account ID with existing user ID",
+			accountID:     "nonexistent",
+			userID:        "f4f6d672-63fb-11ec-90d6-0242ac120003",
+			expectedCount: 0,
+		},
+		{
+			name:          "should return no peers for non-existing user ID with existing account ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			userID:        "nonexistent_user",
+			expectedCount: 0,
+		},
+		{
+			name:          "should retrieve peers for another valid account ID and user ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			userID:        "edafee4e-63fb-11ec-90d6-0242ac120003",
+			expectedCount: 3,
+		},
+		{
+			name:          "should return no peers for existing account ID with empty user ID",
+			accountID:     "bf1c8084-ba50-4ce7-9439-34653001fc3b",
+			userID:        "",
+			expectedCount: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			peers, err := store.GetUserPeers(context.Background(), LockingStrengthNone, tt.accountID, tt.userID)
+			require.NoError(t, err)
+			require.Len(t, peers, tt.expectedCount)
+		})
+	}
+}
+
+func TestSqlStore_DeletePeer(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	peerID := "csrnkiq7qv9d8aitqd50"
+
+	err = store.DeletePeer(context.Background(), accountID, peerID)
+	require.NoError(t, err)
+
+	peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, accountID, peerID)
+	require.Error(t, err)
+	require.Nil(t, peer)
+}
+
+func BenchmarkGetAccountPeers(b *testing.B) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store_with_expired_peers.sql", b.TempDir())
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.Cleanup(cleanup)
+
+	numberOfPeers := 1000
+	numberOfGroups := 200
+	numberOfPeersPerGroup := 500
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+	peers := make([]*nbpeer.Peer, 0, numberOfPeers)
+	for i := 0; i < numberOfPeers; i++ {
+		peer := &nbpeer.Peer{
+			ID:        fmt.Sprintf("peer-%d", i),
+			AccountID: accountID,
+			Key:       fmt.Sprintf("key-%d", i),
+			DNSLabel:  fmt.Sprintf("peer%d.example.com", i),
+			IP:        intToIPv4(uint32(i)),
+		}
+		err = store.AddPeerToAccount(context.Background(), peer)
+		if err != nil {
+			b.Fatalf("Failed to add peer: %v", err)
+		}
+		peers = append(peers, peer)
+	}
+
+	for i := 0; i < numberOfGroups; i++ {
+		groupID := fmt.Sprintf("group-%d", i)
+		group := &types.Group{
+			ID:        groupID,
+			AccountID: accountID,
+		}
+		err = store.CreateGroup(context.Background(), group)
+		if err != nil {
+			b.Fatalf("Failed to create group: %v", err)
+		}
+		for j := 0; j < numberOfPeersPerGroup; j++ {
+			peerIndex := (i*numberOfPeersPerGroup + j) % numberOfPeers
+			err = store.AddPeerToGroup(context.Background(), accountID, peers[peerIndex].ID, groupID)
+			if err != nil {
+				b.Fatalf("Failed to add peer to group: %v", err)
+			}
+		}
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peers[i%numberOfPeers].ID)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func intToIPv4(n uint32) netip.Addr {
+	var b [4]byte
+	binary.BigEndian.PutUint32(b[:], n)
+	return netip.AddrFrom4(b)
+}
+
+func TestSqlStore_GetUserIDByPeerKey(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	userID := "test-user-123"
+	peerKey := "peer-key-abc"
+
+	peer := &nbpeer.Peer{
+		ID:        "test-peer-1",
+		Key:       peerKey,
+		AccountID: existingAccountID,
+		UserID:    userID,
+		IP:        netip.AddrFrom4([4]byte{10, 0, 0, 1}),
+		IPv6:      netip.MustParseAddr("fd00::a00:1"),
+		DNSLabel:  "test-peer-1",
+	}
+
+	err = store.AddPeerToAccount(context.Background(), peer)
+	require.NoError(t, err)
+
+	retrievedUserID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, peerKey)
+	require.NoError(t, err)
+	assert.Equal(t, userID, retrievedUserID)
+}
+
+func TestSqlStore_GetUserIDByPeerKey_NotFound(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	nonExistentPeerKey := "non-existent-peer-key"
+
+	userID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, nonExistentPeerKey)
+	require.Error(t, err)
+	assert.Equal(t, "", userID)
+}
+
+func TestSqlStore_GetUserIDByPeerKey_NoUserID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	peerKey := "peer-key-abc"
+
+	peer := &nbpeer.Peer{
+		ID:        "test-peer-1",
+		Key:       peerKey,
+		AccountID: existingAccountID,
+		UserID:    "",
+		IP:        netip.AddrFrom4([4]byte{10, 0, 0, 1}),
+		IPv6:      netip.MustParseAddr("fd00::a00:1"),
+		DNSLabel:  "test-peer-1",
+	}
+
+	err = store.AddPeerToAccount(context.Background(), peer)
+	require.NoError(t, err)
+
+	retrievedUserID, err := store.GetUserIDByPeerKey(context.Background(), LockingStrengthNone, peerKey)
+	require.NoError(t, err)
+	assert.Equal(t, "", retrievedUserID)
+}
+
+func TestSqlStore_ApproveAccountPeers(t *testing.T) {
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		accountID := "test-account"
+		ctx := context.Background()
+
+		account := newAccountWithId(ctx, accountID, "testuser", "example.com")
+		err := store.SaveAccount(ctx, account)
+		require.NoError(t, err)
+
+		peers := []*nbpeer.Peer{
+			{
+				ID:        "peer1",
+				AccountID: accountID,
+				DNSLabel:  "peer1.netbird.cloud",
+				Key:       "peer1-key",
+				IP:        netip.MustParseAddr("100.64.0.1"),
+				IPv6:      netip.MustParseAddr("fd00::1"),
+				Status: &nbpeer.PeerStatus{
+					RequiresApproval: true,
+					LastSeen:         time.Now().UTC(),
+				},
+			},
+			{
+				ID:        "peer2",
+				AccountID: accountID,
+				DNSLabel:  "peer2.netbird.cloud",
+				Key:       "peer2-key",
+				IP:        netip.MustParseAddr("100.64.0.2"),
+				IPv6:      netip.MustParseAddr("fd00::2"),
+				Status: &nbpeer.PeerStatus{
+					RequiresApproval: true,
+					LastSeen:         time.Now().UTC(),
+				},
+			},
+			{
+				ID:        "peer3",
+				AccountID: accountID,
+				DNSLabel:  "peer3.netbird.cloud",
+				Key:       "peer3-key",
+				IP:        netip.MustParseAddr("100.64.0.3"),
+				IPv6:      netip.MustParseAddr("fd00::3"),
+				Status: &nbpeer.PeerStatus{
+					RequiresApproval: false,
+					LastSeen:         time.Now().UTC(),
+				},
+			},
+		}
+
+		for _, peer := range peers {
+			err = store.AddPeerToAccount(ctx, peer)
+			require.NoError(t, err)
+		}
+
+		t.Run("approve all pending peers", func(t *testing.T) {
+			count, err := store.ApproveAccountPeers(ctx, accountID)
+			require.NoError(t, err)
+			assert.Equal(t, 2, count)
+
+			allPeers, err := store.GetAccountPeers(ctx, LockingStrengthNone, accountID, "", "")
+			require.NoError(t, err)
+
+			for _, peer := range allPeers {
+				assert.False(t, peer.Status.RequiresApproval, "peer %s should not require approval", peer.ID)
+			}
+		})
+
+		t.Run("no peers to approve", func(t *testing.T) {
+			count, err := store.ApproveAccountPeers(ctx, accountID)
+			require.NoError(t, err)
+			assert.Equal(t, 0, count)
+		})
+
+		t.Run("non-existent account", func(t *testing.T) {
+			count, err := store.ApproveAccountPeers(ctx, "non-existent")
+			require.NoError(t, err)
+			assert.Equal(t, 0, count)
+		})
+	})
+}

management/server/store/sql_store_personal_access_token.go

@@ -0,0 +1,178 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/management/server/util"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// 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) GetTokenIDByHashedToken(ctx context.Context, hashedToken string) (string, error) {
+	var token types.PersonalAccessToken
+	result := s.db.Take(&token, "hashed_token = ?", hashedToken)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
+		}
+		log.WithContext(ctx).Errorf("error when getting token from the store: %s", result.Error)
+		return "", status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	return token.ID, nil
+}
+
+func (s *SqlStore) getPersonalAccessTokens(ctx context.Context, userIDs []string) ([]types.PersonalAccessToken, error) {
+	if len(userIDs) == 0 {
+		return nil, nil
+	}
+	const query = `SELECT id, user_id, name, hashed_token, expiration_date, created_by, created_at, last_used FROM personal_access_tokens WHERE user_id = ANY($1)`
+	rows, err := s.pool.Query(ctx, query, userIDs)
+	if err != nil {
+		return nil, err
+	}
+	pats, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (types.PersonalAccessToken, error) {
+		var pat types.PersonalAccessToken
+		var expirationDate, lastUsed, createdAt sql.NullTime
+		err := row.Scan(&pat.ID, &pat.UserID, &pat.Name, &pat.HashedToken, &expirationDate, &pat.CreatedBy, &createdAt, &lastUsed)
+		if err == nil {
+			if expirationDate.Valid {
+				pat.ExpirationDate = &expirationDate.Time
+			}
+			if createdAt.Valid {
+				pat.CreatedAt = createdAt.Time
+			}
+			if lastUsed.Valid {
+				pat.LastUsed = &lastUsed.Time
+			}
+		}
+		return pat, err
+	})
+	if err != nil {
+		return nil, err
+	}
+	return pats, nil
+}
+
+// GetPATByHashedToken returns a PersonalAccessToken by its hashed token.
+func (s *SqlStore) GetPATByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken string) (*types.PersonalAccessToken, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var pat types.PersonalAccessToken
+	result := tx.Take(&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) (*types.PersonalAccessToken, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var pat types.PersonalAccessToken
+	result := tx.
+		Take(&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) ([]*types.PersonalAccessToken, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var pats []*types.PersonalAccessToken
+	result := tx.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, patID string) error {
+	patCopy := types.PersonalAccessToken{
+		LastUsed: util.ToPtr(time.Now().UTC()),
+	}
+
+	fieldsToUpdate := []string{"last_used"}
+	result := s.db.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, pat *types.PersonalAccessToken) error {
+	result := s.db.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, userID, patID string) error {
+	result := s.db.Delete(&types.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
+}