diff --git a/e2e/agentnetwork/account_delete_test.go b/e2e/agentnetwork/account_delete_test.go new file mode 100644 index 000000000..a8b52766f --- /dev/null +++ b/e2e/agentnetwork/account_delete_test.go @@ -0,0 +1,292 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// accountDeleteModel is a made-up model id the provider enumerates and prices, +// so the chat routes to the mock upstream and is metered deterministically. +const accountDeleteModel = "e2e-account-delete-model" + +// agentNetworkConfigTables are deleted with the account, in its transaction. +var agentNetworkConfigTables = []string{ + "agent_network_settings", + "agent_network_providers", + "agent_network_policies", + "agent_network_guardrails", + "agent_network_budget_rules", +} + +// TestAccountDelete_RemovesAgentNetworkState deletes an account that has a full +// Agent Network setup and has served traffic, and checks what that leaves +// behind, end to end: +// +// - the proxy stops running the account's gateway, instead of keeping its +// mappings and provider API keys in memory until it next resyncs; +// - the configuration rows go with the account, while access logs and usage +// records stay for retention; +// - the account's consumption counters are swept once the cleanup runs; +// - the gateway domain is free for another account to claim. +// +// It runs on a dedicated server, since deleting the shared account would take +// every other test down with it. +func TestAccountDelete_RemovesAgentNetworkState(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + fresh, err := harnessStartFresh(ctx, t) + require.NoError(t, err, "start dedicated combined server") + + accounts, err := fresh.API().Accounts.List(ctx) + require.NoError(t, err, "list accounts") + require.Len(t, accounts, 1, "a fresh server has exactly the bootstrapped account") + accountID := accounts[0].Id + + cluster := harness.AgentNetworkCluster + settings, err := fresh.CreateSettings(ctx, api.AgentNetworkSettingsCreateRequest{ProxyAddress: &cluster}) + require.NoError(t, err, "bootstrap agent-network endpoint") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned at bootstrap") + + env := provisionAccountDeleteEnv(t, ctx, fresh, settings.Endpoint) + chatThrough(t, ctx, env) + + // Preconditions: the proxy runs the account's gateway, and the request left + // the traffic-driven rows the rest of the test expects to outlive the delete. + requireEventually(t, ctx, 60*time.Second, "proxy should run a client for the account", func() bool { + return proxyRunsAccount(t, ctx, env.proxy, accountID) + }) + requireEventually(t, ctx, accessLogIngestWindow, "the request should leave consumption, usage and access-log rows", func() bool { + counts := accountRowCounts(t, fresh, accountID, + "agent_network_consumption", "agent_network_request_usage", "agent_network_access_log") + return counts["agent_network_consumption"] > 0 && + counts["agent_network_request_usage"] > 0 && + counts["agent_network_access_log"] > 0 + }) + + require.NoError(t, fresh.API().Accounts.Delete(ctx, accountID), "delete account") + + // The proxy is told to drop the gateway. A proxy that only learns on its + // next resync keeps serving the deleted account with its provider API keys. + if !eventually(ctx, 60*time.Second, func() bool { return proxyDroppedAccount(t, ctx, env.proxy, accountID) }) { + t.Errorf("proxy still runs a client for deleted account %s\n=== proxy logs ===\n%s", + accountID, env.proxy.Logs(context.Background())) + } + + counts := accountRowCounts(t, fresh, accountID, append(slices.Clone(agentNetworkConfigTables), + "agent_network_request_usage", "agent_network_access_log")...) + for _, table := range agentNetworkConfigTables { + assert.Zero(t, counts[table], "%s rows should be deleted with the account", table) + } + assert.NotZero(t, counts["agent_network_request_usage"], "usage records should be kept") + assert.NotZero(t, counts["agent_network_access_log"], "access logs should be left for retention") + + // The cleanup's first pass runs at startup, and whether instance setup is + // open again is only re-evaluated then. + require.NoError(t, fresh.Restart(ctx), "restart combined server") + requireEventually(t, ctx, 60*time.Second, "the cleanup should sweep the deleted account's consumption counters", func() bool { + return accountRowCounts(t, fresh, accountID, "agent_network_consumption")["agent_network_consumption"] == 0 + }) + + // A new account can claim the deleted account's gateway domain: its + // settings row no longer holds the global unique index. + _, err = fresh.Bootstrap(ctx) + require.NoError(t, err, "bootstrap a second account once the first is gone") + claimed, err := fresh.CreateSettings(ctx, api.AgentNetworkSettingsCreateRequest{Endpoint: &settings.Endpoint}) + require.NoError(t, err, "a new account should be able to claim the deleted account's gateway domain") + assert.Equal(t, settings.Endpoint, claimed.Endpoint, "the new account should hold the released domain") +} + +// accountDeleteEnv is a connected gateway for one account: a proxy running the +// debug endpoint, a client peer, and the resolved endpoint. +type accountDeleteEnv struct { + endpoint string + proxyIP string + client *harness.Client + proxy *harness.Proxy +} + +// provisionAccountDeleteEnv gives the server's account one of every Agent +// Network configuration row (provider, guardrail, policy, budget rule; the +// settings row is the caller's) and brings up a proxy and a client. The policy +// and budget rule switch on usage metering, so a request records consumption. +func provisionAccountDeleteEnv(t *testing.T, ctx context.Context, srv *harness.Combined, endpoint string) accountDeleteEnv { + t.Helper() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock vLLM upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-account-delete"}) + require.NoError(t, err, "create group") + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-account-delete-client", + Type: "reusable", + ExpiresIn: 86400, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + + apiKey := "sk-account-delete-e2e" + models := []api.AgentNetworkProviderModel{{Id: accountDeleteModel, InputPer1k: 0.01, OutputPer1k: 0.02}} + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "account-delete", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &apiKey, + Enabled: ptr(true), + Models: &models, + }) + require.NoError(t, err, "create provider") + + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-account-delete" + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{accountDeleteModel} + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + + limits := api.AgentNetworkPolicyLimits{ + TokenLimit: api.AgentNetworkPolicyTokenLimit{ + Enabled: true, + GroupCap: 10_000_000, + UserCap: 10_000_000, + WindowSeconds: 60, + }, + } + _, err = srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-account-delete", + Enabled: ptr(true), + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{guard.Id}, + Limits: &limits, + }) + require.NoError(t, err, "create policy") + + _, err = srv.CreateBudgetRule(ctx, api.AgentNetworkBudgetRuleRequest{ + Name: "e2e-account-delete", + Limits: limits, + TargetGroups: &[]string{grp.Id}, + }) + require.NoError(t, err, "create budget rule") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-account-delete-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{"NB_PROXY_DEBUG_ENDPOINT": "true"}) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + proxyIP, err := cl.ResolveProxyIP(ctx, endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + return accountDeleteEnv{endpoint: endpoint, proxyIP: proxyIP, client: cl, proxy: px} +} + +// chatThrough drives one chat through the gateway, retrying to absorb +// first-call tunnel and DNS jitter. +func chatThrough(t *testing.T, ctx context.Context, env accountDeleteEnv) { + t.Helper() + var code int + var body string + ok := eventually(ctx, 90*time.Second, func() bool { + c, b, err := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, + accountDeleteModel, "Reply with exactly: pong", "e2e-session-account-delete") + code, body = c, b + return err == nil && c == 200 + }) + require.True(t, ok, "chat must return 200, last got %d: %s\n=== proxy logs ===\n%s", + code, body, env.proxy.Logs(context.Background())) +} + +// proxyRunsAccount reports whether a lookup succeeded and shows the proxy +// running a client for the account. +func proxyRunsAccount(t *testing.T, ctx context.Context, px *harness.Proxy, accountID string) bool { + t.Helper() + runs, ok := lookupProxyAccount(t, ctx, px, accountID) + return ok && runs +} + +// proxyDroppedAccount reports whether a lookup succeeded and shows the proxy +// no longer running a client for the account. A failed lookup confirms +// nothing, so it keeps the caller polling rather than passing the check. +func proxyDroppedAccount(t *testing.T, ctx context.Context, px *harness.Proxy, accountID string) bool { + t.Helper() + runs, ok := lookupProxyAccount(t, ctx, px, accountID) + return ok && !runs +} + +// lookupProxyAccount asks the proxy whether it runs a client for the account; +// ok is false when the lookup itself failed. +func lookupProxyAccount(t *testing.T, ctx context.Context, px *harness.Proxy, accountID string) (runs, ok bool) { + t.Helper() + clients, err := px.DebugClients(ctx) + if err != nil { + t.Logf("proxy debug clients: %v", err) + return false, false + } + return slices.ContainsFunc(clients, func(c harness.ProxyDebugClient) bool { return c.AccountID == accountID }), true +} + +// accountRowCounts counts the account's rows in each table, read from a +// snapshot of the management store. +func accountRowCounts(t *testing.T, srv *harness.Combined, accountID string, tables ...string) map[string]int64 { + t.Helper() + dbPath, err := srv.SnapshotStoreDB(t.TempDir()) + require.NoError(t, err, "snapshot management sqlite store") + db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{}) + require.NoError(t, err, "open store snapshot") + sqlDB, err := db.DB() + require.NoError(t, err) + defer func() { _ = sqlDB.Close() }() + + counts := make(map[string]int64, len(tables)) + for _, table := range tables { + var n int64 + require.NoError(t, db.Table(table).Where("account_id = ?", accountID).Count(&n).Error, "count %s rows", table) + counts[table] = n + } + return counts +} + +// eventually polls cond every two seconds until it holds or timeout passes. +func eventually(ctx context.Context, timeout time.Duration, cond func() bool) bool { + deadline := time.Now().Add(timeout) + for { + if cond() { + return true + } + if time.Now().After(deadline) || !waitBeforeRetry(ctx, 2*time.Second) { + return false + } + } +} + +// requireEventually fails the test now if cond does not hold within timeout. +func requireEventually(t *testing.T, ctx context.Context, timeout time.Duration, msg string, cond func() bool) { + t.Helper() + require.True(t, eventually(ctx, timeout, cond), msg) +} diff --git a/e2e/harness/agentnetwork.go b/e2e/harness/agentnetwork.go index e51f2dd7a..8f688836d 100644 --- a/e2e/harness/agentnetwork.go +++ b/e2e/harness/agentnetwork.go @@ -135,6 +135,11 @@ func (c *Combined) DeleteGuardrail(ctx context.Context, id string) error { return anDelete(ctx, c, "/api/agent-network/guardrails/"+id) } +// CreateBudgetRule creates an account-level agent-network budget rule. +func (c *Combined) CreateBudgetRule(ctx context.Context, req api.AgentNetworkBudgetRuleRequest) (api.AgentNetworkBudgetRule, error) { + return anRequest[api.AgentNetworkBudgetRule](ctx, c, http.MethodPost, "/api/agent-network/budget-rules", req) +} + // CreateSettings bootstraps the account's agent-network settings row, // assigning the immutable endpoint. Exactly one of req.ProxyAddress (labeled // endpoint beneath that cluster) and req.Endpoint (self-addressed dedicated diff --git a/e2e/harness/combined.go b/e2e/harness/combined.go index e03f9f256..ea451e3ad 100644 --- a/e2e/harness/combined.go +++ b/e2e/harness/combined.go @@ -305,6 +305,33 @@ func (c *Combined) SnapshotStoreDB(dstDir string) (string, error) { return dst, nil } +// Restart stops and starts the combined container, keeping its bind-mounted +// data dir, and waits for the API again. The host port can change across a +// restart, so BaseURL and the authenticated client are refreshed. Work that +// management only does at startup (such as the agent-network cleanup's first +// pass, or re-evaluating whether instance setup is required) runs again. +func (c *Combined) Restart(ctx context.Context) error { + if err := c.container.Stop(ctx, nil); err != nil { + return fmt.Errorf("stop combined container: %w", err) + } + if err := c.container.Start(ctx); err != nil { + return fmt.Errorf("start combined container: %w", err) + } + host, err := c.container.Host(ctx) + if err != nil { + return fmt.Errorf("container host: %w", err) + } + mapped, err := c.container.MappedPort(ctx, nat.Port(combinedHTTPPort)) + if err != nil { + return fmt.Errorf("mapped port: %w", err) + } + c.BaseURL = fmt.Sprintf("http://%s:%s", host, mapped.Port()) + if c.PAT != "" { + c.api = rest.New(c.BaseURL, c.PAT) + } + return nil +} + // Logs returns the combined server container logs, for diagnostics. func (c *Combined) Logs(ctx context.Context) string { return containerLogs(ctx, c.container) diff --git a/e2e/harness/proxy.go b/e2e/harness/proxy.go index 3d709b439..ee458908b 100644 --- a/e2e/harness/proxy.go +++ b/e2e/harness/proxy.go @@ -3,13 +3,17 @@ package harness import ( + "bytes" "context" + "encoding/json" "fmt" + "io" "os" "time" "github.com/docker/docker/api/types/container" "github.com/testcontainers/testcontainers-go" + tcexec "github.com/testcontainers/testcontainers-go/exec" "github.com/testcontainers/testcontainers-go/wait" ) @@ -114,6 +118,41 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverride return &Proxy{container: ctr, workDir: workDir}, nil } +// ProxyDebugClient is one per-account embedded client the proxy runs, as the +// proxy's debug endpoint reports it. +type ProxyDebugClient struct { + AccountID string `json:"account_id"` + ServiceCount int `json:"service_count"` + ServiceKeys []string `json:"service_keys"` +} + +// DebugClients lists the per-account clients the proxy is running, through +// the proxy's own debug CLI inside the container. The proxy must be started +// with NB_PROXY_DEBUG_ENDPOINT=true. +func (p *Proxy) DebugClients(ctx context.Context) ([]ProxyDebugClient, error) { + code, reader, err := p.container.Exec(ctx, + []string{"/usr/bin/netbird-proxy", "debug", "clients", "--json"}, tcexec.Multiplexed()) + if err != nil { + return nil, fmt.Errorf("exec debug clients: %w", err) + } + out, _ := io.ReadAll(reader) + if code != 0 { + return nil, fmt.Errorf("debug clients exited %d: %s", code, string(out)) + } + // stderr is multiplexed in; the JSON document starts at the first brace. + start := bytes.IndexByte(out, '{') + if start < 0 { + return nil, fmt.Errorf("no JSON in debug clients output: %s", string(out)) + } + var resp struct { + Clients []ProxyDebugClient `json:"clients"` + } + if err := json.NewDecoder(bytes.NewReader(out[start:])).Decode(&resp); err != nil { + return nil, fmt.Errorf("decode debug clients output: %w", err) + } + return resp.Clients, nil +} + // Logs returns the proxy container logs, for diagnostics on failure. func (p *Proxy) Logs(ctx context.Context) string { return containerLogs(ctx, p.container) diff --git a/management/internals/modules/agentnetwork/accesslog_cleanup_realstore_test.go b/management/internals/modules/agentnetwork/accesslog_cleanup_realstore_test.go new file mode 100644 index 000000000..0e294a560 --- /dev/null +++ b/management/internals/modules/agentnetwork/accesslog_cleanup_realstore_test.go @@ -0,0 +1,76 @@ +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + nbtypes "github.com/netbirdio/netbird/management/server/types" +) + +// TestCleanupAccessLogs_RealStore_DeletedAccount covers a deleted account's access logs. +// The sweep is driven by settings rows, which go with the account, so without a fallback +// those logs would never expire. They get the default retention instead. A live account +// can delete its own settings row, so "no settings" must not be mistaken for "deleted": +// that account's logs are left alone, as are those of an account that keeps logs forever. +func TestCleanupAccessLogs_RealStore_DeletedAccount(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + const ( + deletedAccountID = "acc-deleted" + keepAccountID = "acc-keep-forever" + noSettingsAccountID = "acc-live-no-settings" + ) + old := time.Now().UTC().AddDate(0, 0, -(types.DefaultAccessLogRetentionDays + 10)) + recent := time.Now().UTC().AddDate(0, 0, -1) + + require.NoError(t, s.SaveAccount(ctx, &nbtypes.Account{Id: keepAccountID})) + require.NoError(t, s.SaveAccount(ctx, &nbtypes.Account{Id: noSettingsAccountID})) + + keepSettings := types.DefaultSettings(keepAccountID) + keepSettings.Domain = "keep.gw.example.com" + keepSettings.AccessLogRetentionDays = 0 + require.NoError(t, s.SaveAgentNetworkSettings(ctx, keepSettings)) + + mkLog := func(id, accountID string, ts time.Time) { + t.Helper() + entry := &types.AgentNetworkAccessLog{ + ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts, StatusCode: 200, Model: "gpt-4o", + } + groups := []types.AgentNetworkAccessLogGroup{{LogID: id, GroupID: "grp-eng", AccountID: accountID}} + require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, entry, groups)) + } + mkLog("deleted-old", deletedAccountID, old) + mkLog("deleted-recent", deletedAccountID, recent) + mkLog("keep-old", keepAccountID, old) + mkLog("no-settings-old", noSettingsAccountID, old) + + m := &managerImpl{store: s} + m.cleanupAccessLogsOnce(ctx) + + logIDs := func(accountID string) []string { + t.Helper() + logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, accountID, + types.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50}) + require.NoError(t, err) + ids := make([]string, 0, len(logs)) + for _, l := range logs { + ids = append(ids, l.ID) + } + return ids + } + assert.Equal(t, []string{"deleted-recent"}, logIDs(deletedAccountID), + "a deleted account should have logs past the default retention swept") + assert.Equal(t, []string{"keep-old"}, logIDs(keepAccountID), + "an account with retention disabled should keep its old logs") + assert.Equal(t, []string{"no-settings-old"}, logIDs(noSettingsAccountID), + "a live account without a settings row should keep its old logs") +} diff --git a/management/internals/modules/agentnetwork/manager.go b/management/internals/modules/agentnetwork/manager.go index d1a5ebd7b..71ff53214 100644 --- a/management/internals/modules/agentnetwork/manager.go +++ b/management/internals/modules/agentnetwork/manager.go @@ -80,6 +80,9 @@ type Manager interface { ListAccessLogSessions(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLogSession, int64, error) GetUsageOverview(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter, granularity types.UsageGranularity) ([]*types.AgentNetworkUsageBucket, error) StartAccessLogCleanup(ctx context.Context, cleanupIntervalHours int) + // RemoveAccountGateway drops the account's gateway mappings from the + // proxies. It runs as an account deletion hook. + RemoveAccountGateway(ctx context.Context, accountID string) error RecordConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds, tokensIn, tokensOut int64, costUSD float64) error RecordAccountBudgetUsage(ctx context.Context, accountID, userID string, groupIDs []string, tokensIn, tokensOut int64, costUSD float64) error RecordUsage(ctx context.Context, in RecordUsageInput) error @@ -1350,8 +1353,8 @@ func (m *managerImpl) scopeFilterToCaller(ctx context.Context, accountID, userID // StartAccessLogCleanup launches a background sweep that periodically deletes // each account's agent-network access-log rows older than that account's -// AccessLogRetentionDays. Usage records are never swept. A non-positive -// interval defaults to 24h. +// AccessLogRetentionDays, and the consumption counters of deleted accounts. +// Usage records are never swept. A non-positive interval defaults to 24h. func (m *managerImpl) StartAccessLogCleanup(ctx context.Context, cleanupIntervalHours int) { if cleanupIntervalHours <= 0 { cleanupIntervalHours = 24 @@ -1362,21 +1365,40 @@ func (m *managerImpl) StartAccessLogCleanup(ctx context.Context, cleanupInterval ticker := time.NewTicker(interval) defer ticker.Stop() - m.cleanupAccessLogsOnce(ctx) // run once on startup + m.cleanupOnce(ctx) // run once on startup for { select { case <-ctx.Done(): return case <-ticker.C: - m.cleanupAccessLogsOnce(ctx) + m.cleanupOnce(ctx) } } }() } +func (m *managerImpl) cleanupOnce(ctx context.Context) { + m.cleanupAccessLogsOnce(ctx) + m.cleanupDeletedAccountConsumption(ctx) +} + +// cleanupDeletedAccountConsumption deletes the consumption counters of accounts +// that no longer exist. Best-effort: a failure is logged and retried next sweep. +func (m *managerImpl) cleanupDeletedAccountConsumption(ctx context.Context) { + deleted, err := m.store.DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx) + if err != nil { + log.WithContext(ctx).Warnf("agent-network consumption cleanup: %v", err) + return + } + if deleted > 0 { + log.WithContext(ctx).Infof("agent-network consumption cleanup: deleted %d counters of deleted accounts", deleted) + } +} + // cleanupAccessLogsOnce sweeps every account's expired access-log rows against -// its configured retention. Best-effort: a per-account failure is logged and -// the sweep continues. +// its configured retention. Deleted accounts, whose settings rows went with +// them, get the default retention. Best-effort: a per-account failure is +// logged and the sweep continues. func (m *managerImpl) cleanupAccessLogsOnce(ctx context.Context) { settings, err := m.store.GetAllAgentNetworkSettings(ctx, store.LockingStrengthNone) if err != nil { @@ -1384,18 +1406,31 @@ func (m *managerImpl) cleanupAccessLogsOnce(ctx context.Context) { return } for _, s := range settings { - if s.AccessLogRetentionDays <= 0 { - continue // keep indefinitely - } - cutoff := time.Now().UTC().AddDate(0, 0, -s.AccessLogRetentionDays) - deleted, err := m.store.DeleteOldAgentNetworkAccessLogs(ctx, s.AccountID, cutoff) - if err != nil { - log.WithContext(ctx).Warnf("agent-network access-log cleanup for account %s: %v", s.AccountID, err) - continue - } - if deleted > 0 { - log.WithContext(ctx).Infof("agent-network access-log cleanup: deleted %d rows for account %s (retention %d days)", deleted, s.AccountID, s.AccessLogRetentionDays) - } + m.cleanupAccountAccessLogs(ctx, s.AccountID, s.AccessLogRetentionDays) + } + + deleted, err := m.store.GetDeletedAccountIDsWithAgentNetworkAccessLogs(ctx) + if err != nil { + log.WithContext(ctx).Errorf("agent-network access-log cleanup: list deleted accounts: %v", err) + return + } + for _, accountID := range deleted { + m.cleanupAccountAccessLogs(ctx, accountID, types.DefaultAccessLogRetentionDays) + } +} + +func (m *managerImpl) cleanupAccountAccessLogs(ctx context.Context, accountID string, retentionDays int) { + if retentionDays <= 0 { + return // keep indefinitely + } + cutoff := time.Now().UTC().AddDate(0, 0, -retentionDays) + deleted, err := m.store.DeleteOldAgentNetworkAccessLogs(ctx, accountID, cutoff) + if err != nil { + log.WithContext(ctx).Warnf("agent-network access-log cleanup for account %s: %v", accountID, err) + return + } + if deleted > 0 { + log.WithContext(ctx).Infof("agent-network access-log cleanup: deleted %d rows for account %s (retention %d days)", deleted, accountID, retentionDays) } } @@ -1545,6 +1580,8 @@ func (*mockManager) GetUsageOverview(_ context.Context, _, _ string, _ types.Age func (*mockManager) StartAccessLogCleanup(_ context.Context, _ int) {} +func (*mockManager) RemoveAccountGateway(_ context.Context, _ string) error { return nil } + func (*mockManager) RecordConsumption(_ context.Context, _ string, _ types.ConsumptionDimension, _ string, _, _, _ int64, _ float64) error { return nil } diff --git a/management/internals/modules/agentnetwork/reconcile.go b/management/internals/modules/agentnetwork/reconcile.go index 69e684014..20d0bb42a 100644 --- a/management/internals/modules/agentnetwork/reconcile.go +++ b/management/internals/modules/agentnetwork/reconcile.go @@ -2,8 +2,10 @@ package agentnetwork import ( "context" + "fmt" log "github.com/sirupsen/logrus" + goproto "google.golang.org/protobuf/proto" rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" "github.com/netbirdio/netbird/management/server/types" @@ -81,18 +83,66 @@ func (m *managerImpl) reconcile(ctx context.Context, accountID string) { } m.reconcileMu.Unlock() - for _, entry := range creates { - entry.mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED - m.proxyController.SendServiceUpdateToCluster(ctx, accountID, entry.mapping, entry.cluster) + m.sendMappings(ctx, accountID, creates, proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED) + m.sendMappings(ctx, accountID, updates, proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED) + m.sendMappings(ctx, accountID, deletes, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED) +} + +// sendMappings sends each entry as updateType. It sends a copy: the entries' +// mappings are shared with reconcileCache, which another reconcile or +// RemoveAccountGateway may be reading, so they are never written. +func (m *managerImpl) sendMappings(ctx context.Context, accountID string, entries []syntheticMapping, updateType proto.ProxyMappingUpdateType) { + for _, entry := range entries { + update := goproto.Clone(entry.mapping).(*proto.ProxyMapping) + update.Type = updateType + m.proxyController.SendServiceUpdateToCluster(ctx, accountID, update, entry.cluster) } - for _, entry := range updates { - entry.mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED - m.proxyController.SendServiceUpdateToCluster(ctx, accountID, entry.mapping, entry.cluster) +} + +// RemoveAccountGateway tells the proxies to drop every mapping of the account's +// gateway, so a deleted account's proxy config, provider API keys included, does +// not linger in proxy memory until the next resync. It is an account deletion +// hook: it runs before the account's data is removed, the last point at which +// the mappings can be synthesised from the store. The cache alone would miss +// them, since it is per instance and empty after a restart. If the deletion +// then fails, the gateway stays down until the account's next change reconciles +// it back. +func (m *managerImpl) RemoveAccountGateway(ctx context.Context, accountID string) error { + if m.proxyController == nil { + return nil } - for _, entry := range deletes { - entry.mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED - m.proxyController.SendServiceUpdateToCluster(ctx, accountID, entry.mapping, entry.cluster) + + services, err := SynthesizeServices(ctx, m.store, accountID) + if err != nil { + return fmt.Errorf("synthesise agent network services: %w", err) } + oidcCfg := m.proxyController.GetOIDCValidationConfig() + removed := make(map[string]syntheticMapping, len(services)) + for _, svc := range services { + if svc == nil || svc.ID == "" { + continue + } + removed[svc.ID] = syntheticMapping{ + mapping: svc.ToProtoMapping(rpservice.Delete, "", oidcCfg), + cluster: svc.ProxyCluster, + } + } + + m.reconcileMu.Lock() + for id, entry := range m.reconcileCache[accountID] { + if _, ok := removed[id]; !ok { + removed[id] = entry + } + } + delete(m.reconcileCache, accountID) + m.reconcileMu.Unlock() + + entries := make([]syntheticMapping, 0, len(removed)) + for _, entry := range removed { + entries = append(entries, entry) + } + m.sendMappings(ctx, accountID, entries, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED) + return nil } // diffMappings classifies the previous→current transition for a single diff --git a/management/internals/modules/agentnetwork/reconcile_test.go b/management/internals/modules/agentnetwork/reconcile_test.go index ab3b08481..2cfea9828 100644 --- a/management/internals/modules/agentnetwork/reconcile_test.go +++ b/management/internals/modules/agentnetwork/reconcile_test.go @@ -2,6 +2,8 @@ package agentnetwork import ( "context" + "sync" + "sync/atomic" "testing" "go.uber.org/mock/gomock" @@ -12,6 +14,7 @@ import ( "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" "github.com/netbirdio/netbird/management/server/store" "github.com/netbirdio/netbird/shared/management/proto" + "github.com/netbirdio/netbird/shared/management/status" ) func newReconcileMgr(t *testing.T, ctrl *gomock.Controller) (*managerImpl, *store.MockStore, *proxy.MockController) { @@ -287,3 +290,154 @@ func TestDiffMappings_RemovedServiceIsDeletedOnItsOwnCluster(t *testing.T) { assert.Equal(t, "brave-otter.gateway.example.com", deletes[0].cluster) } } + +// TestRemoveAccountGateway_EmitsRemovedFromStore — account deletion runs on an +// instance that may never have reconciled the account, so its cache is empty. +// The mappings are synthesised from the store, still intact before the delete, +// and each is sent as REMOVED to the cluster that serves it. +func TestRemoveAccountGateway_EmitsRemovedFromStore(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + provider := newReconcileTestProvider() + policy := newReconcileTestPolicy(provider.ID, "grp-eng") + + expectReconcileSynthInputs(mockStore, ctx, []*types.Provider{provider}, []*types.Policy{policy}, []*types.Guardrail{}) + mockProxy.EXPECT().GetOIDCValidationConfig().Return(proxy.OIDCValidationConfig{}) + + var sent []*proto.ProxyMapping + mockProxy.EXPECT(). + SendServiceUpdateToCluster(ctx, "acct-1", gomock.Any(), "eu.proxy.netbird.io"). + Do(func(_ context.Context, _ string, m *proto.ProxyMapping, _ string) { + sent = append(sent, m) + }) + + require.NoError(t, mgr.RemoveAccountGateway(ctx, "acct-1")) + + require.Len(t, sent, 1, "the account's one gateway mapping must be removed") + assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED, sent[0].Type, "the update must be a removal") + assert.Equal(t, "agent-net-svc-acct-1", sent[0].Id, "the removal must name the account's gateway service") +} + +// TestRemoveAccountGateway_AlsoRemovesCachedMappings — a mapping this instance +// last sent but the store no longer synthesises (here, one on another cluster) +// is removed too, and the account's cache entry is cleared. +func TestRemoveAccountGateway_AlsoRemovesCachedMappings(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + mgr.reconcileCache["acct-1"] = map[string]syntheticMapping{ + "stale-svc": {mapping: &proto.ProxyMapping{Id: "stale-svc"}, cluster: "us.proxy.netbird.io"}, + } + + // Settings but no providers: the store synthesises nothing. + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1"). + Return(newReconcileTestSettings(), nil) + mockStore.EXPECT(). + GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1"). + Return([]*types.Provider{}, nil) + mockProxy.EXPECT().GetOIDCValidationConfig().Return(proxy.OIDCValidationConfig{}) + + var sent []*proto.ProxyMapping + mockProxy.EXPECT(). + SendServiceUpdateToCluster(ctx, "acct-1", gomock.Any(), "us.proxy.netbird.io"). + Do(func(_ context.Context, _ string, m *proto.ProxyMapping, _ string) { + sent = append(sent, m) + }) + + require.NoError(t, mgr.RemoveAccountGateway(ctx, "acct-1")) + + require.Len(t, sent, 1, "the cached mapping must be removed from its own cluster") + assert.Equal(t, "stale-svc", sent[0].Id) + assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED, sent[0].Type) + mgr.reconcileMu.Lock() + _, present := mgr.reconcileCache["acct-1"] + mgr.reconcileMu.Unlock() + assert.False(t, present, "the deleted account's cache entry must be cleared") +} + +// TestRemoveAccountGateway_SynthFailureAbortsDeletion — if the mappings cannot +// be read, nothing is sent and the error is returned, which as an account +// deletion hook keeps the account rather than leaving its gateway running. +func TestRemoveAccountGateway_SynthFailureAbortsDeletion(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, _ := newReconcileMgr(t, ctrl) + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1"). + Return(nil, status.Errorf(status.Internal, "store unavailable")) + + assert.Error(t, mgr.RemoveAccountGateway(ctx, "acct-1"), "a failed synthesis must fail the hook") +} + +func TestRemoveAccountGateway_NilProxyController_NoOp(t *testing.T) { + mgr := &managerImpl{reconcileCache: make(map[string]map[string]syntheticMapping)} + // Must not panic and must not query the store. + assert.NoError(t, mgr.RemoveAccountGateway(context.Background(), "acct-1")) +} + +// TestReconcile_ConcurrentWithGatewayChanges — while an account's gateway +// flaps (its policy is removed and re-added between reads), concurrent +// reconciles and RemoveAccountGateway share the cached mappings: one caches a +// mapping and sends it, another finds it gone and sends its removal. Run under +// -race: neither path may write a cached mapping, only copies of it. +func TestReconcile_ConcurrentWithGatewayChanges(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + // gomock serialises every call on the controller's mutex, which would give + // the race detector the ordering the code under test lacks. The sends go + // through a fake that takes no lock. + mgr.proxyController = unsyncedSender{MockController: mockProxy} + provider := newReconcileTestProvider() + policy := newReconcileTestPolicy(provider.ID, "grp-eng") + + var reads atomic.Int64 + mockStore.EXPECT().GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1").Return(newReconcileTestSettings(), nil).AnyTimes() + mockStore.EXPECT().GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Provider{provider}, nil).AnyTimes() + mockStore.EXPECT().GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, "acct-1"). + DoAndReturn(func(context.Context, store.LockingStrength, string) ([]*types.Policy, error) { + if reads.Add(1)%2 == 0 { + return []*types.Policy{}, nil + } + return []*types.Policy{policy}, nil + }).AnyTimes() + mockStore.EXPECT().GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Guardrail{}, nil).AnyTimes() + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func(remove bool) { + defer wg.Done() + for j := 0; j < 50; j++ { + if remove && j%10 == 0 { + _ = mgr.RemoveAccountGateway(ctx, "acct-1") + continue + } + mgr.reconcile(ctx, "acct-1") + } + }(i == 0) + } + wg.Wait() +} + +// unsyncedSender answers the calls reconcile makes on every pass without any +// locking, so concurrent callers are not ordered by the fake itself. +type unsyncedSender struct { + *proxy.MockController +} + +func (unsyncedSender) GetOIDCValidationConfig() proxy.OIDCValidationConfig { + return proxy.OIDCValidationConfig{} +} + +func (unsyncedSender) SendServiceUpdateToCluster(context.Context, string, *proto.ProxyMapping, string) {} diff --git a/management/internals/server/modules.go b/management/internals/server/modules.go index 6b1365f3b..4840e40ad 100644 --- a/management/internals/server/modules.go +++ b/management/internals/server/modules.go @@ -103,6 +103,7 @@ func (s *BaseServer) AccountManager() account.Manager { s.AfterInit(func(s *BaseServer) { accountManager.SetServiceManager(s.ServiceManager()) + accountManager.AddAccountDeletionHook(s.AgentNetworkManager().RemoveAccountGateway) }) return accountManager diff --git a/management/server/account.go b/management/server/account.go index 6ccf673f5..340bcc84b 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -112,6 +112,9 @@ type DefaultAccountManager struct { permissionsManager permissions.Manager disableDefaultPolicy bool + + deletionHooksMu sync.RWMutex + deletionHooks []account.DeletionHook } var _ account.Manager = (*DefaultAccountManager)(nil) @@ -120,6 +123,32 @@ func (am *DefaultAccountManager) SetServiceManager(serviceManager service.Manage am.serviceManager = serviceManager } +// AddAccountDeletionHook registers hook to run on every account deletion. Hooks run in +// registration order, and the first one to fail stops the rest and aborts the deletion. +// It panics on a nil hook: dropping one silently would skip that hook's cleanup on every +// deletion, so the wiring bug surfaces at startup instead. +func (am *DefaultAccountManager) AddAccountDeletionHook(hook account.DeletionHook) { + if hook == nil { + panic("nil account deletion hook") + } + am.deletionHooksMu.Lock() + defer am.deletionHooksMu.Unlock() + am.deletionHooks = append(am.deletionHooks, hook) +} + +func (am *DefaultAccountManager) runAccountDeletionHooks(ctx context.Context, accountID string) error { + am.deletionHooksMu.RLock() + hooks := slices.Clone(am.deletionHooks) + am.deletionHooksMu.RUnlock() + + for _, hook := range hooks { + if err := hook(ctx, accountID); err != nil { + return fmt.Errorf("account deletion hook: %w", err) + } + } + return nil +} + func isUniqueConstraintError(err error) bool { switch { case strings.Contains(err.Error(), "(SQLSTATE 23505)"), @@ -889,6 +918,10 @@ func (am *DefaultAccountManager) DeleteAccount(ctx context.Context, accountID, u return status.Errorf(status.Internal, "failed to build user infos for account %s: %v", accountID, err) } + if err = am.runAccountDeletionHooks(ctx, accountID); err != nil { + return err + } + if err = am.deleteAccountUsers(ctx, accountID, userID, account.Users, userInfosMap); err != nil { return err } diff --git a/management/server/account/deletion_hook.go b/management/server/account/deletion_hook.go new file mode 100644 index 000000000..17c817444 --- /dev/null +++ b/management/server/account/deletion_hook.go @@ -0,0 +1,14 @@ +package account + +import "context" + +// DeletionHook runs when an account is deleted, after the caller's permission to delete +// it has been checked and before any of its users or data are removed. It lets code that +// keeps per-account state outside the store tear that state down while the account still +// exists. +// +// A hook that returns an error aborts the deletion and the account is kept. The caller +// sees the error, so a hook that wants a specific response returns a status error. A +// retried deletion runs every hook again, and a later step can still fail after the hooks +// succeed, so a hook must be idempotent and must tolerate the account surviving it. +type DeletionHook func(ctx context.Context, accountID string) error diff --git a/management/server/account_test.go b/management/server/account_test.go index c63782ca8..8c735b28e 100644 --- a/management/server/account_test.go +++ b/management/server/account_test.go @@ -958,6 +958,101 @@ func TestAccountManager_DeleteAccount(t *testing.T) { assert.Len(t, pats, 0) } +func TestAccountManager_DeleteAccount_RunsDeletionHooks(t *testing.T) { + manager, _, err := createManager(t) + require.NoError(t, err) + + ownerID := "account_creator" + account, err := createAccount(manager, "test_account", ownerID, "") + require.NoError(t, err) + + // Each hook records its call and checks the account is still in the store, which is + // the point of running before deletion: a hook must be able to read what it cleans up. + var calls []string + hook := func(name string) nbAccount.DeletionHook { + return func(ctx context.Context, accountID string) error { + calls = append(calls, name+":"+accountID) + _, err := manager.Store.GetAccount(ctx, accountID) + assert.NoError(t, err, "account should still exist while hook %s runs", name) + return nil + } + } + manager.AddAccountDeletionHook(hook("first")) + manager.AddAccountDeletionHook(hook("second")) + + require.NoError(t, manager.DeleteAccount(context.Background(), account.Id, ownerID)) + + assert.Equal(t, []string{"first:" + account.Id, "second:" + account.Id}, calls, + "hooks should run once each, in registration order, with the deleted account's ID") + _, err = manager.Store.GetAccount(context.Background(), account.Id) + assert.Error(t, err, "account should be deleted after the hooks succeed") +} + +func TestAccountManager_DeleteAccount_DeletionHookErrorAbortsDeletion(t *testing.T) { + manager, _, err := createManager(t) + require.NoError(t, err) + + ownerID := "account_creator" + account, err := createAccount(manager, "test_account", ownerID, "") + require.NoError(t, err) + + manager.AddAccountDeletionHook(func(context.Context, string) error { + return status.Errorf(status.PreconditionFailed, "teardown refused") + }) + secondCalled := false + manager.AddAccountDeletionHook(func(context.Context, string) error { + secondCalled = true + return nil + }) + + err = manager.DeleteAccount(context.Background(), account.Id, ownerID) + require.Error(t, err) + + // The hook's status type has to survive the wrapping, since the HTTP layer maps it + // to the response code. + sErr, ok := status.FromError(err) + require.True(t, ok, "error should carry the hook's status error, got %v", err) + assert.Equal(t, status.PreconditionFailed, sErr.Type(), "status type should be the hook's") + assert.False(t, secondCalled, "hooks after a failing one should not run") + + _, err = manager.Store.GetAccount(context.Background(), account.Id) + assert.NoError(t, err, "account should survive a failing hook") + _, err = manager.Store.GetUserByUserID(context.Background(), store.LockingStrengthNone, ownerID) + assert.NoError(t, err, "account owner should survive a failing hook") +} + +func TestAccountManager_AddAccountDeletionHook_RejectsNil(t *testing.T) { + manager, _, err := createManager(t) + require.NoError(t, err) + + assert.PanicsWithValue(t, "nil account deletion hook", func() { + manager.AddAccountDeletionHook(nil) + }, "registering a nil hook should panic instead of breaking a later deletion") +} + +func TestAccountManager_DeleteAccount_DeletionHooksSkippedWithoutPermission(t *testing.T) { + manager, _, err := createManager(t) + require.NoError(t, err) + + ownerID := "account_creator" + account, err := createAccount(manager, "test_account", ownerID, "") + require.NoError(t, err) + + adminID := "regular_admin" + account.Users[adminID] = types.NewAdminUser(adminID) + require.NoError(t, manager.Store.SaveAccount(context.Background(), account)) + + called := false + manager.AddAccountDeletionHook(func(context.Context, string) error { + called = true + return nil + }) + + err = manager.DeleteAccount(context.Background(), account.Id, adminID) + require.Error(t, err, "only the owner may delete the account") + assert.False(t, called, "hooks should not run for a caller who may not delete the account") +} + func BenchmarkTest_GetAccountWithclaims(b *testing.B) { claims := auth.UserAuth{ Domain: "example.com", diff --git a/management/server/store/sql_store_account.go b/management/server/store/sql_store_account.go index 8729c1f5a..43fe09861 100644 --- a/management/server/store/sql_store_account.go +++ b/management/server/store/sql_store_account.go @@ -18,6 +18,7 @@ import ( "gorm.io/gorm/clause" nbdns "github.com/netbirdio/netbird/dns" + agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" nbpeer "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/types" "github.com/netbirdio/netbird/route" @@ -162,6 +163,10 @@ func (s *SqlStore) DeleteAccount(ctx context.Context, account *types.Account) er return result.Error } + if err := deleteAgentNetworkAccountConfig(tx, account.Id); err != nil { + return err + } + result = tx.Select(clause.Associations).Delete(account) if result.Error != nil { return result.Error @@ -179,6 +184,29 @@ func (s *SqlStore) DeleteAccount(ctx context.Context, account *types.Account) er return err } +// deleteAgentNetworkAccountConfig removes the account's agent network configuration. These +// tables are not account associations, so deleting the account does not reach them. The +// settings row holds the account's globally unique gateway domain and the provider rows +// hold its upstream API keys. Tables that grow with traffic are left out: consumption +// counters and access logs are swept in the background, and usage records are kept. +func deleteAgentNetworkAccountConfig(tx *gorm.DB, accountID string) error { + // Dependents first: policies point at providers and guardrails, and settings + // go last, as DeleteSettings refuses while providers exist. + models := []any{ + &agentNetworkTypes.Policy{}, + &agentNetworkTypes.Provider{}, + &agentNetworkTypes.Guardrail{}, + &agentNetworkTypes.AccountBudgetRule{}, + &agentNetworkTypes.Settings{}, + } + for _, model := range models { + if err := tx.Delete(model, "account_id = ?", accountID).Error; err != nil { + return fmt.Errorf("delete %T rows: %w", model, err) + } + } + return nil +} + func (s *SqlStore) UpdateAccountDomainAttributes(ctx context.Context, accountID string, domain string, category string, isPrimaryDomain bool) error { accountCopy := types.Account{ Domain: domain, diff --git a/management/server/store/sql_store_account_test.go b/management/server/store/sql_store_account_test.go index 4c3b3f5fd..ec5976982 100644 --- a/management/server/store/sql_store_account_test.go +++ b/management/server/store/sql_store_account_test.go @@ -16,6 +16,7 @@ import ( "github.com/stretchr/testify/require" nbdns "github.com/netbirdio/netbird/dns" + agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" 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" @@ -401,6 +402,25 @@ func TestSqlite_DeleteAccount(t *testing.T) { require.NoError(t, err) require.Equal(t, o.AccountID, account.Id) + err = store.CreateAgentNetworkSettings(context.Background(), &agentNetworkTypes.Settings{ + AccountID: account.Id, + Domain: "gw.example.com", + ProxyAddress: "gw.example.com", + }) + require.NoError(t, err) + + agentNetworkConfig := []any{ + &agentNetworkTypes.Provider{ID: "an_provider", AccountID: account.Id, APIKey: "sk-test"}, + &agentNetworkTypes.Policy{ID: "an_policy", AccountID: account.Id}, + &agentNetworkTypes.Guardrail{ID: "an_guardrail", AccountID: account.Id}, + &agentNetworkTypes.AccountBudgetRule{ID: "an_budget_rule", AccountID: account.Id}, + } + for _, row := range agentNetworkConfig { + require.NoError(t, store.(*SqlStore).db.Create(row).Error, "creating %T", row) + } + otherProvider := &agentNetworkTypes.Provider{ID: "other_provider", AccountID: "other_account"} + require.NoError(t, store.(*SqlStore).db.Create(otherProvider).Error) + err = store.DeleteAccount(context.Background(), account) require.NoError(t, err) @@ -465,6 +485,32 @@ func TestSqlite_DeleteAccount(t *testing.T) { 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") + + _, err = store.GetAgentNetworkSettings(context.Background(), LockingStrengthNone, account.Id) + require.Error(t, err, "expecting agent network settings to be deleted with the account") + sErr, ok := status.FromError(err) + require.True(t, ok, "expecting a status error when getting agent network settings, got %v", err) + require.Equal(t, status.NotFound, sErr.Type(), "expecting agent network settings to be deleted with the account") + + // The domain is globally unique, so a leftover row would keep it from another account. + err = store.CreateAgentNetworkSettings(context.Background(), &agentNetworkTypes.Settings{ + AccountID: "other_account", + Domain: "gw.example.com", + ProxyAddress: "gw.example.com", + }) + require.NoError(t, err, "expecting the deleted account's gateway domain to be free for another account") + + for _, row := range agentNetworkConfig { + var count int64 + err = store.(*SqlStore).db.Model(row).Where("account_id = ?", account.Id).Count(&count).Error + require.NoError(t, err, "counting %T rows after DeleteAccount", row) + assert.Zero(t, count, "expecting no %T rows to be found after DeleteAccount", row) + } + + var otherProviders int64 + err = store.(*SqlStore).db.Model(&agentNetworkTypes.Provider{}).Where("account_id = ?", "other_account").Count(&otherProviders).Error + require.NoError(t, err) + assert.Equal(t, int64(1), otherProviders, "expecting another account's agent network provider to survive DeleteAccount") } func Test_GetAccount(t *testing.T) { diff --git a/management/server/store/sql_store_agent_network_access_log.go b/management/server/store/sql_store_agent_network_access_log.go index a1ae0150d..479dfa45e 100644 --- a/management/server/store/sql_store_agent_network_access_log.go +++ b/management/server/store/sql_store_agent_network_access_log.go @@ -69,6 +69,23 @@ func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountI return deleted, nil } +// GetDeletedAccountIDsWithAgentNetworkAccessLogs returns the IDs of accounts that no +// longer exist but still have access-log rows. The retention sweep is driven by settings +// rows, which are deleted with the account, so it uses this to find logs it would +// otherwise never expire. +func (s *SqlStore) GetDeletedAccountIDsWithAgentNetworkAccessLogs(ctx context.Context) ([]string, error) { + var accountIDs []string + err := s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}). + Distinct("account_id"). + Where("NOT EXISTS (SELECT 1 FROM accounts WHERE accounts.id = agent_network_access_log.account_id)"). + Pluck("account_id", &accountIDs).Error + if err != nil { + log.WithContext(ctx).Errorf("failed to get deleted accounts with agent-network access logs: %v", err) + return nil, status.Errorf(status.Internal, "failed to get deleted accounts with agent-network access logs") + } + return accountIDs, 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. diff --git a/management/server/store/sql_store_agentnetwork.go b/management/server/store/sql_store_agentnetwork.go index 4fe77d994..bdb1f97d6 100644 --- a/management/server/store/sql_store_agentnetwork.go +++ b/management/server/store/sql_store_agentnetwork.go @@ -663,6 +663,21 @@ func (s *SqlStore) IncrementAgentNetworkConsumptionBatch( return nil } +// DeleteAgentNetworkConsumptionOfDeletedAccounts deletes every consumption counter whose +// account no longer exists and returns the number of rows deleted. Counters grow with +// traffic, so they are swept in the background instead of in the account-deletion +// transaction, and the sweep also catches counters a proxy writes after the deletion. +func (s *SqlStore) DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx context.Context) (int64, error) { + res := s.db. + Where("NOT EXISTS (SELECT 1 FROM accounts WHERE accounts.id = agent_network_consumption.account_id)"). + Delete(&agentNetworkTypes.Consumption{}) + if res.Error != nil { + log.WithContext(ctx).Errorf("failed to delete agent-network consumption of deleted accounts: %v", res.Error) + return 0, status.Errorf(status.Internal, "failed to delete agent-network consumption of deleted accounts") + } + return res.RowsAffected, nil +} + // ListAgentNetworkConsumption returns every consumption row recorded // for the account, ordered by window_start descending. Backs the // dashboard's basic counter view. diff --git a/management/server/store/sql_store_agentnetwork_accesslog_test.go b/management/server/store/sql_store_agentnetwork_accesslog_test.go index 8ba79a062..b8a3560e0 100644 --- a/management/server/store/sql_store_agentnetwork_accesslog_test.go +++ b/management/server/store/sql_store_agentnetwork_accesslog_test.go @@ -9,6 +9,7 @@ import ( "github.com/stretchr/testify/require" agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/types" ) // TestAgentNetworkUsage_RealStore_RoundTrip drives CreateAgentNetworkUsage and @@ -300,3 +301,37 @@ func TestDeleteOldAgentNetworkAccessLogs(t *testing.T) { require.NoError(t, err) require.Len(t, usage, 1, "usage record for the deleted log must survive") } + +// TestDeleteAgentNetworkConsumptionOfDeletedAccounts verifies that the sweep removes the +// consumption counters of accounts that no longer exist and leaves live accounts' counters, +// including those of a live account without a settings row. +func TestDeleteAgentNetworkConsumptionOfDeletedAccounts(t *testing.T) { + runTestForAllEngines(t, "", func(t *testing.T, s Store) { + ctx := context.Background() + const ( + liveAccountID = "acc-anet-consumption-live" + deletedAccountID = "acc-anet-consumption-deleted" + ) + require.NoError(t, s.SaveAccount(ctx, &types.Account{Id: liveAccountID})) + + windowStart := time.Now().UTC().Truncate(time.Hour) + for _, accountID := range []string{liveAccountID, deletedAccountID} { + for _, dimID := range []string{"user-1", "user-2"} { + require.NoError(t, s.IncrementAgentNetworkConsumption(ctx, accountID, + agentNetworkTypes.DimensionUser, dimID, 3600, windowStart, 10, 5, 0.01)) + } + } + + deleted, err := s.DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx) + require.NoError(t, err) + assert.Equal(t, int64(2), deleted, "both of the deleted account's counters should be removed") + + rows, err := s.ListAgentNetworkConsumption(ctx, LockingStrengthNone, deletedAccountID) + require.NoError(t, err) + assert.Empty(t, rows, "the deleted account should have no consumption counters left") + + rows, err = s.ListAgentNetworkConsumption(ctx, LockingStrengthNone, liveAccountID) + require.NoError(t, err) + assert.Len(t, rows, 2, "the live account's consumption counters should survive") + }) +} diff --git a/management/server/store/store.go b/management/server/store/store.go index 6d810643c..465f84413 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -320,6 +320,7 @@ type Store interface { GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) + GetDeletedAccountIDsWithAgentNetworkAccessLogs(ctx context.Context) ([]string, error) GetServiceTargetByTargetID(ctx context.Context, lockStrength LockingStrength, accountID string, targetID string) (*rpservice.Target, error) GetTargetsByServiceID(ctx context.Context, lockStrength LockingStrength, accountID string, serviceID string) ([]*rpservice.Target, error) DeleteTarget(ctx context.Context, accountID string, serviceID string, targetID uint) error @@ -388,6 +389,7 @@ type Store interface { GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error) GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error) + DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx context.Context) (int64, error) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go index 4dd7f9286..cd9e7334d 100644 --- a/management/server/store/store_mock.go +++ b/management/server/store/store_mock.go @@ -457,6 +457,21 @@ func (mr *MockStoreMockRecorder) DeleteAgentNetworkBudgetRule(ctx, accountID, ru return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkBudgetRule), ctx, accountID, ruleID) } +// DeleteAgentNetworkConsumptionOfDeletedAccounts mocks base method. +func (m *MockStore) DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx context.Context) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAgentNetworkConsumptionOfDeletedAccounts", ctx) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteAgentNetworkConsumptionOfDeletedAccounts indicates an expected call of DeleteAgentNetworkConsumptionOfDeletedAccounts. +func (mr *MockStoreMockRecorder) DeleteAgentNetworkConsumptionOfDeletedAccounts(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkConsumptionOfDeletedAccounts", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkConsumptionOfDeletedAccounts), ctx) +} + // DeleteAgentNetworkGuardrail mocks base method. func (m *MockStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error { m.ctrl.T.Helper() @@ -1972,6 +1987,21 @@ func (mr *MockStoreMockRecorder) GetDNSRecordByID(ctx, lockStrength, accountID, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDNSRecordByID", reflect.TypeOf((*MockStore)(nil).GetDNSRecordByID), ctx, lockStrength, accountID, zoneID, recordID) } +// GetDeletedAccountIDsWithAgentNetworkAccessLogs mocks base method. +func (m *MockStore) GetDeletedAccountIDsWithAgentNetworkAccessLogs(ctx context.Context) ([]string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetDeletedAccountIDsWithAgentNetworkAccessLogs", ctx) + ret0, _ := ret[0].([]string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetDeletedAccountIDsWithAgentNetworkAccessLogs indicates an expected call of GetDeletedAccountIDsWithAgentNetworkAccessLogs. +func (mr *MockStoreMockRecorder) GetDeletedAccountIDsWithAgentNetworkAccessLogs(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeletedAccountIDsWithAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).GetDeletedAccountIDsWithAgentNetworkAccessLogs), ctx) +} + // GetEmbeddedProxyPeerIDsByCluster mocks base method. func (m *MockStore) GetEmbeddedProxyPeerIDsByCluster(ctx context.Context, accountID string) (map[string][]string, error) { m.ctrl.T.Helper()