From 524b8b97188e4817b0b616ff97241d5604cf041f Mon Sep 17 00:00:00 2001 From: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:03:40 +0200 Subject: [PATCH] [management] prewarm a posture check cache on network map generation (#7093) --- .../network_map/controller/controller.go | 2 + management/server/account_test.go | 1 + management/server/types/account.go | 3 + management/server/types/account_components.go | 115 ++++++++++++++++-- .../types/account_posture_validation_test.go | 72 +++++++++++ .../server/types/networkmap_benchmark_test.go | 37 ++++++ 6 files changed, 223 insertions(+), 7 deletions(-) create mode 100644 management/server/types/account_posture_validation_test.go diff --git a/management/internals/controllers/network_map/controller/controller.go b/management/internals/controllers/network_map/controller/controller.go index 5785004db..356dc9f67 100644 --- a/management/internals/controllers/network_map/controller/controller.go +++ b/management/internals/controllers/network_map/controller/controller.go @@ -176,6 +176,7 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin semaphore := make(chan struct{}, 10) c.injectAllProxyPolicies(ctx, account) + account.PrecomputePostureValidation(ctx) dnsCache := &cache.DNSConfigCache{} dnsDomain := c.GetDNSDomain(account.Settings) peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain) @@ -357,6 +358,7 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s // network map that omitted the synth DNS zone, and the agent kept // resolving against the stale or absent record. c.injectAllProxyPolicies(ctx, account) + account.PrecomputePostureValidation(ctx) dnsCache := &cache.DNSConfigCache{} dnsDomain := c.GetDNSDomain(account.Settings) peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain) diff --git a/management/server/account_test.go b/management/server/account_test.go index 3c0bb25da..73126a496 100644 --- a/management/server/account_test.go +++ b/management/server/account_test.go @@ -1757,6 +1757,7 @@ func TestAccount_Copy(t *testing.T) { AccountID: "account1", }, }, + PostureValidation: map[string]map[string]bool{"1": {"1": true}}, } err := hasNilField(account) if err != nil { diff --git a/management/server/types/account.go b/management/server/types/account.go index 474825281..4616fe26b 100644 --- a/management/server/types/account.go +++ b/management/server/types/account.go @@ -91,6 +91,8 @@ type Account struct { Onboarding AccountOnboarding `gorm:"foreignKey:AccountID;references:id;constraint:OnDelete:CASCADE"` ReverseProxyFreeDomainNonce string + + PostureValidation map[string]map[string]bool `gorm:"-"` } // this class is used by gorm only @@ -874,6 +876,7 @@ func (a *Account) Copy() *Account { Services: services, Onboarding: a.Onboarding, Domains: domains, + PostureValidation: a.PostureValidation, } } diff --git a/management/server/types/account_components.go b/management/server/types/account_components.go index 6fc904c0b..624a778fe 100644 --- a/management/server/types/account_components.go +++ b/management/server/types/account_components.go @@ -10,6 +10,8 @@ import ( nbdns "github.com/netbirdio/netbird/dns" "github.com/netbirdio/netbird/management/internals/modules/zones" routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types" + nbpeer "github.com/netbirdio/netbird/management/server/peer" + "github.com/netbirdio/netbird/management/server/posture" "github.com/netbirdio/netbird/management/server/telemetry" "github.com/netbirdio/netbird/route" ) @@ -506,8 +508,8 @@ func (a *Account) getPeersGroupsPoliciesRoutes( func (a *Account) getPeersFromGroups(ctx context.Context, groups []string, peerID string, sourcePostureChecksIDs []string, validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) ([]string, bool) { peerInGroups := false - filteredPeerIDs := make([]string, 0, len(groups)) - seenPeerIds := make(map[string]struct{}, len(groups)) + var filteredPeerIDs []string + var seenPeerIds map[string]struct{} for _, gid := range groups { group := a.GetGroup(gid) @@ -547,6 +549,17 @@ func (a *Account) getPeersFromGroups(ctx context.Context, groups []string, peerI return filteredPeerIDs, peerInGroups } + if seenPeerIds == nil { + totalGroupPeers := 0 + for _, g := range groups { + if grp := a.GetGroup(g); grp != nil { + totalGroupPeers += len(grp.Peers) + } + } + filteredPeerIDs = make([]string, 0, totalGroupPeers) + seenPeerIds = make(map[string]struct{}, totalGroupPeers) + } + for _, pid := range group.Peers { if _, seen := seenPeerIds[pid]; seen { continue @@ -589,21 +602,109 @@ func (a *Account) validatePostureChecksOnPeerGetFailed(ctx context.Context, sour } for _, postureChecksID := range sourcePostureChecksID { + if valid, cached := a.cachedPostureCheckResult(postureChecksID, peerID); cached { + if !valid { + return false, postureChecksID + } + continue + } + postureChecks := a.GetPostureChecks(postureChecksID) if postureChecks == nil { continue } - for _, check := range postureChecks.GetChecks() { - isValid, _ := check.Check(ctx, *peer) - if !isValid { - return false, postureChecksID - } + if !peerPassesPostureChecks(ctx, postureChecks.GetChecks(), peer) { + return false, postureChecksID } } return true, "" } +// PrecomputePostureValidation evaluates every posture check referenced by an enabled +// policy once against the peers of that policy's source groups and stores the results, +// so the per-peer network map calculations that follow look them up instead of +// re-evaluating checks for every peer pair. It must be called before the account is +// shared across goroutines; lookups not covered by the precomputed results fall back +// to direct evaluation. +func (a *Account) PrecomputePostureValidation(ctx context.Context) { + if len(a.PostureChecks) == 0 { + a.PostureValidation = nil + return + } + + checkPeerIDs := make(map[string]map[string]struct{}) + for _, policy := range a.Policies { + if !policy.Enabled || len(policy.SourcePostureChecks) == 0 { + continue + } + + peerIDs := a.getUniquePeerIDsFromGroupsIDs(ctx, policy.SourceGroups()) + for _, rule := range policy.Rules { + if rule.SourceResource.Type == ResourceTypePeer && rule.SourceResource.ID != "" { + peerIDs = append(peerIDs, rule.SourceResource.ID) + } + } + + for _, postureChecksID := range policy.SourcePostureChecks { + set := checkPeerIDs[postureChecksID] + if set == nil { + set = make(map[string]struct{}, len(peerIDs)) + checkPeerIDs[postureChecksID] = set + } + for _, pid := range peerIDs { + set[pid] = struct{}{} + } + } + } + + results := make(map[string]map[string]bool, len(checkPeerIDs)) + for postureChecksID, peerIDs := range checkPeerIDs { + results[postureChecksID] = a.evaluatePostureChecksForPeers(ctx, postureChecksID, peerIDs) + } + a.PostureValidation = results +} + +func (a *Account) evaluatePostureChecksForPeers(ctx context.Context, postureChecksID string, peerIDs map[string]struct{}) map[string]bool { + postureChecks := a.GetPostureChecks(postureChecksID) + if postureChecks == nil { + return nil + } + + checks := postureChecks.GetChecks() + results := make(map[string]bool, len(peerIDs)) + for peerID := range peerIDs { + peer, ok := a.Peers[peerID] + if !ok || peer == nil { + continue + } + results[peerID] = peerPassesPostureChecks(ctx, checks, peer) + } + return results +} + +func (a *Account) cachedPostureCheckResult(postureChecksID, peerID string) (bool, bool) { + results, ok := a.PostureValidation[postureChecksID] + if !ok { + return false, false + } + if results == nil { + return true, true + } + valid, found := results[peerID] + return valid, found +} + +func peerPassesPostureChecks(ctx context.Context, checks []posture.Check, peer *nbpeer.Peer) bool { + for _, check := range checks { + isValid, _ := check.Check(ctx, *peer) + if !isValid { + return false + } + } + return true +} + func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChecksIDs []string, validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) []string { var dest []string for _, peerID := range inputPeers { diff --git a/management/server/types/account_posture_validation_test.go b/management/server/types/account_posture_validation_test.go new file mode 100644 index 000000000..69a744bea --- /dev/null +++ b/management/server/types/account_posture_validation_test.go @@ -0,0 +1,72 @@ +package types_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + nbdns "github.com/netbirdio/netbird/dns" + "github.com/netbirdio/netbird/management/server/posture" +) + +func TestPrecomputePostureValidation_MatchesDirectEvaluation(t *testing.T) { + account, validatedPeers := scalableTestAccount(60, 5) + + account.PostureChecks = append(account.PostureChecks, &posture.Checks{ + ID: "posture-check-strict", Name: "Strict version", + Checks: posture.ChecksDefinition{ + NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.50.0"}, + }, + }) + account.Policies[0].SourcePostureChecks = []string{"posture-check-ver", "posture-check-unknown"} + account.Policies[1].SourcePostureChecks = []string{"posture-check-strict"} + account.Policies[2].SourcePostureChecks = []string{"posture-check-ver"} + account.Policies[2].Enabled = false + + ctx := context.Background() + resourcePolicies := account.GetResourcePoliciesMap() + routers := account.GetResourceRoutersMap() + + type result struct { + peers map[string]struct{} + postureFailedPeers map[string]map[string]struct{} + } + snapshot := func() map[string]result { + results := make(map[string]result, len(account.Peers)) + for peerID := range account.Peers { + components := account.GetPeerNetworkMapComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil) + require.NotNil(t, components) + peerSet := make(map[string]struct{}, len(components.Peers)) + for id := range components.Peers { + peerSet[id] = struct{}{} + } + results[peerID] = result{peers: peerSet, postureFailedPeers: components.PostureFailedPeers} + } + return results + } + + direct := snapshot() + account.PrecomputePostureValidation(ctx) + memoized := snapshot() + + require.Equal(t, len(direct), len(memoized)) + for peerID, want := range direct { + got := memoized[peerID] + assert.Equal(t, want.peers, got.peers, "visible peers changed for %s", peerID) + assert.Equal(t, want.postureFailedPeers, got.postureFailedPeers, "posture failed peers changed for %s", peerID) + } +} + +func TestPrecomputePostureValidation_NoPostureChecks(t *testing.T) { + account, validatedPeers := scalableTestAccount(10, 2) + account.PostureChecks = nil + + ctx := context.Background() + account.PrecomputePostureValidation(ctx) + + components := account.GetPeerNetworkMapComponents(ctx, "peer-0", nbdns.CustomZone{}, nil, validatedPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil) + require.NotNil(t, components) + assert.NotEmpty(t, components.Peers) +} diff --git a/management/server/types/networkmap_benchmark_test.go b/management/server/types/networkmap_benchmark_test.go index 38272e7b0..118ef9baf 100644 --- a/management/server/types/networkmap_benchmark_test.go +++ b/management/server/types/networkmap_benchmark_test.go @@ -86,6 +86,43 @@ func BenchmarkNetworkMapGeneration_AllPeers(b *testing.B) { b.ReportAllocs() b.ResetTimer() for range b.N { + account.PrecomputePostureValidation(ctx) + for _, peerID := range peerIDs { + _ = account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs) + } + } + }) + } +} + +// BenchmarkNetworkMapGeneration_AllPeersPostureChecks benchmarks the UpdateAccountPeers +// hot path with a posture check attached to the account-wide policy, so posture +// validation runs for every source peer of every target peer's map. +func BenchmarkNetworkMapGeneration_AllPeersPostureChecks(b *testing.B) { + skipCIBenchmark(b) + scales := []benchmarkScale{ + {"500peers_20groups", 500, 20}, + {"1000peers_50groups", 1000, 50}, + } + + for _, scale := range scales { + account, validatedPeers := scalableTestAccount(scale.peers, scale.groups) + account.Policies[0].SourcePostureChecks = []string{"posture-check-ver"} + ctx := context.Background() + + peerIDs := make([]string, 0, len(account.Peers)) + for peerID := range account.Peers { + peerIDs = append(peerIDs, peerID) + } + + b.Run("components/"+scale.name, func(b *testing.B) { + resourcePolicies := account.GetResourcePoliciesMap() + routers := account.GetResourceRoutersMap() + groupIDToUserIDs := account.GetActiveGroupUsers() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + account.PrecomputePostureValidation(ctx) for _, peerID := range peerIDs { _ = account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs) }