From 7414c276f61dd0505525e6ae75f99591dc3391b7 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Mon, 19 Oct 2026 09:26:59 +0000 Subject: [PATCH] mapper: limit capability copies to requested peers Small peer updates should not copy capabilities for the entire tailnet. --- hscontrol/mapper/builder.go | 10 +- hscontrol/mapper/mapper_test.go | 347 +++++++++++++++++++++++++++++ hscontrol/policy/pm.go | 22 +- hscontrol/policy/v2/policy.go | 45 +++- hscontrol/policy/v2/policy_test.go | 345 ++++++++++++++++++++++++++++ hscontrol/state/state.go | 6 + 6 files changed, 751 insertions(+), 24 deletions(-) diff --git a/hscontrol/mapper/builder.go b/hscontrol/mapper/builder.go index e3f794d58..2f3e14b23 100644 --- a/hscontrol/mapper/builder.go +++ b/hscontrol/mapper/builder.go @@ -262,10 +262,8 @@ func (b *MapResponseBuilder) buildTailPeers(peers views.Slice[types.NodeView]) ( return nil, err } - // Snapshot the per-node policy CapMap once per peer-list build - // instead of locking the policy manager per peer. The per-call - // path used to take pm.mu N times for an N-peer response. - allCapMaps := b.mapper.state.NodeCapMaps() + // Snapshot these peers' policy caps under one policy-manager lock. + peerCapMaps := b.mapper.state.NodeCapMapsForPeers(peers) // Build tail nodes with per-peer via-aware route function. tailPeers := make([]*tailcfg.Node, 0, peers.Len()) @@ -278,7 +276,7 @@ func (b *MapResponseBuilder) buildTailPeers(peers views.Slice[types.NodeView]) ( // filtering side-effect inside TailNode survives. tn, err := peer.TailNode(b.capVer, func(_ types.NodeID) []netip.Prefix { return b.mapper.state.RoutesForPeer(node, peer, matchers) - }, b.mapper.cfg, allCapMaps[peer.ID()]) + }, b.mapper.cfg, peerCapMaps[peer.ID()]) if err != nil { // One peer with invalid data (e.g. an empty or over-long // GivenName that fails GetFQDN) must not blank out the map for @@ -304,7 +302,7 @@ func (b *MapResponseBuilder) buildTailPeers(peers views.Slice[types.NodeView]) ( // the peer satisfies the cap's emission condition; every other // cap stays off the peer view, leaving CapMap empty for most // peers. [policyv2.PeerCapMap] encodes those conditions. - tn.CapMap = policyv2.PeerCapMap(peer, allCapMaps[peer.ID()]) + tn.CapMap = policyv2.PeerCapMap(peer, peerCapMaps[peer.ID()]) tailPeers = append(tailPeers, tn) } diff --git a/hscontrol/mapper/mapper_test.go b/hscontrol/mapper/mapper_test.go index 37f4bf4e7..60e1548a1 100644 --- a/hscontrol/mapper/mapper_test.go +++ b/hscontrol/mapper/mapper_test.go @@ -3,6 +3,7 @@ package mapper import ( stdcmp "cmp" "database/sql" + "encoding/json" "errors" "fmt" "maps" @@ -17,14 +18,18 @@ import ( "github.com/google/go-cmp/cmp/cmpopts" "github.com/juanfont/headscale/hscontrol/db" "github.com/juanfont/headscale/hscontrol/dns" + policyv2 "github.com/juanfont/headscale/hscontrol/policy/v2" "github.com/juanfont/headscale/hscontrol/state" "github.com/juanfont/headscale/hscontrol/types" "github.com/juanfont/headscale/hscontrol/types/change" + "github.com/juanfont/headscale/hscontrol/util/zlog/zf" + "github.com/rs/zerolog" "github.com/rs/zerolog/log" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" "tailscale.com/tailcfg" + "tailscale.com/tailcfg/nodecap" "tailscale.com/types/dnstype" "tailscale.com/types/views" ) @@ -1545,3 +1550,345 @@ func BenchmarkGenerateUserProfiles(b *testing.B) { } } } + +// Frozen from d00f135bf before selected capability snapshots; only names and cmp alias differ. +func (b *MapResponseBuilder) buildTailPeersCapMapsReference(peers views.Slice[types.NodeView]) ([]*tailcfg.Node, error) { + node, ok := b.mapper.state.GetNodeByID(b.nodeID) + if !ok { + return nil, ErrNodeNotFoundMapper + } + + // [State.RoutesForPeer] needs the unreduced matchers: every rule where the + // node is source or destination, not only the [State.FilterForNode] rules + // where it is destination. + matchers, err := b.mapper.state.MatchersForNode(node) + if err != nil { + return nil, err + } + + // Snapshot the per-node policy CapMap once per peer-list build + // instead of locking the policy manager per peer. The per-call + // path used to take pm.mu N times for an N-peer response. + allCapMaps := b.mapper.state.NodeCapMaps() + + // Build tail nodes with per-peer via-aware route function. + tailPeers := make([]*tailcfg.Node, 0, peers.Len()) + + for _, peer := range peers.All() { + // Pass the peer's policy CapMap as selfPolicyCaps so per-peer + // address-shape rules (today: disable-ipv4) apply consistently + // in the viewer's netmap. The CapMap merge into tn.CapMap is + // overwritten by the PeerCapMap call below; only the address + // filtering side-effect inside TailNode survives. + tn, err := peer.TailNode(b.capVer, func(_ types.NodeID) []netip.Prefix { + return b.mapper.state.RoutesForPeer(node, peer, matchers) + }, b.mapper.cfg, allCapMaps[peer.ID()]) + if err != nil { + // One peer with invalid data (e.g. an empty or over-long + // GivenName that fails GetFQDN) must not blank out the map for + // every node that can see it. Drop the offending peer, log it + // with the identity an operator needs to fix it, and keep + // building from the remaining valid peers. + log.Warn().Caller(). + Err(err). + Uint64(zf.NodeID, peer.ID().Uint64()). + Str(zf.NodeHostname, peer.Hostname()). + Uint64("map.viewer.node.id", b.nodeID.Uint64()). + Msgf("dropping peer %d from map response: invalid node data; fix with `headscale nodes rename %d `", peer.ID(), peer.ID()) + + continue + } + + // [tailcfg.Node.CapMap] on a peer carries the small set of + // caps the Tailscale client reads from the peer view rather + // than the self view (suggest-exit-node, dns-subdomain-resolve + // — see [tailcfg.NodeAttrSuggestExitNode] and + // [tailcfg.NodeAttrDNSSubdomainResolve]). + // The Tailscale-hosted control plane stamps these only when + // the peer satisfies the cap's emission condition; every other + // cap stays off the peer view, leaving CapMap empty for most + // peers. [policyv2.PeerCapMap] encodes those conditions. + tn.CapMap = policyv2.PeerCapMap(peer, allCapMaps[peer.ID()]) + + tailPeers = append(tailPeers, tn) + } + + // Peers is always returned sorted by Node.ID. + slices.SortStableFunc(tailPeers, func(a, b *tailcfg.Node) int { + return stdcmp.Compare(a.ID, b.ID) + }) + + return tailPeers, nil +} + +func (b *MapResponseBuilder) withPeersCapMapsReference(peers views.Slice[types.NodeView]) *MapResponseBuilder { + tailPeers, err := b.buildTailPeersCapMapsReference(peers) + if err != nil { + b.addError(err) + return b + } + + b.resp.Peers = tailPeers + + return b +} + +func (m *mapper) fullMapResponseCapMapsReference( + nodeID types.NodeID, + capVer tailcfg.CapabilityVersion, +) (*tailcfg.MapResponse, error) { + peers := m.state.ListPeers(nodeID) + + return m.NewMapResponseBuilder(nodeID). + WithDebugType(fullResponseDebug). + WithCapabilityVersion(capVer). + WithSelfNode(). + WithDERPMap(). + WithDomain(). + WithCollectServicesDisabled(). + WithDebugConfig(). + WithSSHPolicy(). + WithDNSConfig(). + WithUserProfiles(peers). + WithPacketFilters(). + withPeersCapMapsReference(peers). + Build() +} + +// A star policy gives the viewer all peers without retaining a dense N² peer map. +func capMapsMapperFixture(tb testing.TB, n int, attrs string) (*mapper, []types.NodeView) { + tb.Helper() + + p4 := netip.MustParsePrefix("100.64.0.0/10") + p6 := netip.MustParsePrefix("fd7a:115c:a1e0::/48") + cfg := &types.Config{ + Database: types.DatabaseConfig{ + Type: types.DatabaseSqlite, + Sqlite: types.SqliteConfig{Path: tb.TempDir() + "/caps.db"}, + }, + PrefixV4: &p4, PrefixV6: &p6, + IPAllocation: types.IPAllocationStrategySequential, + BaseDomain: "headscale.test", + Policy: types.PolicyConfig{Mode: types.PolicyModeDB}, + Tuning: types.Tuning{ + NodeStoreBatchSize: state.TestBatchSize, + NodeStoreBatchTimeout: state.TestBatchTimeout, + }, + } + database, err := db.NewHeadscaleDatabase(cfg) + require.NoError(tb, err) + + user := database.CreateUserForTest("capuser") + nodes := make(types.Nodes, n) + ip4, ip6 := netip.MustParseAddr("100.64.0.1"), netip.MustParseAddr("fd7a:115c:a1e0::1") + + for i := range nodes { + name := fmt.Sprintf("capnode%d", i+1) + + nodes[i] = &types.Node{ + ID: types.NodeID(i + 1), Hostname: name, GivenName: name, + UserID: &user.ID, User: user, + IPv4: new(ip4), IPv6: new(ip6), + Hostinfo: &tailcfg.Hostinfo{}, + } + if i%8 == 1 { + routes := []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")} + nodes[i].ApprovedRoutes = routes + nodes[i].Hostinfo.RoutableIPs = routes + } + + ip4, ip6 = ip4.Next(), ip6.Next() + } + + require.NoError(tb, database.DB.Omit("User").CreateInBatches(nodes, 100).Error) + _, err = database.SetPolicy(`{ + "acls":[{"action":"accept","src":["100.64.0.1"],"dst":["*:*"]}], + "nodeAttrs":` + attrs + `}`) + require.NoError(tb, err) + require.NoError(tb, database.Close()) + + s, err := state.NewState(cfg) + require.NoError(tb, err) + tb.Cleanup(func() { require.NoError(tb, s.Close()) }) + s.SetDERPMap(&tailcfg.DERPMap{Regions: map[tailcfg.DERPRegionID]*tailcfg.DERPRegion{}}) + stored := s.ListNodes().AsSlice() + slices.SortFunc(stored, func(a, b types.NodeView) int { return stdcmp.Compare(a.ID(), b.ID()) }) + require.Len(tb, stored, n) + + return &mapper{state: s, cfg: cfg}, stored +} + +func TestBuildTailPeersCapMapsReference(t *testing.T) { + for _, attrs := range []struct { + name, policy string + }{ + {"empty", `[]`}, + {"partial", `[{"target":["100.64.0.2"],"attr":["disable-ipv4","suggest-exit-node"]}]`}, + {"all", `[{"target":["*"],"attr":["disable-ipv4","suggest-exit-node"]}]`}, + } { + t.Run(attrs.name, func(t *testing.T) { + m, nodes := capMapsMapperFixture(t, 5, attrs.policy) + duplicate := nodes[1].AsStruct() + duplicate.GivenName = "same-id-different-name" + missing := nodes[2].AsStruct() + missing.ID = 999 + invalid := nodes[3].AsStruct() + + invalid.GivenName = "" + for _, selection := range []struct { + name string + peers []types.NodeView + }{ + {"nil", nil}, + {"empty", []types.NodeView{}}, + {"one", nodes[1:2]}, + {"full", nodes[1:]}, + {"missing", []types.NodeView{missing.View()}}, + {"duplicates", []types.NodeView{nodes[3], duplicate.View(), nodes[1], nodes[2], nodes[1]}}, + {"invalid-name", []types.NodeView{invalid.View(), nodes[1]}}, + } { + t.Run(selection.name, func(t *testing.T) { + peers := views.SliceOf(selection.peers) + before, err := json.Marshal(m.state.NodeCapMaps()) + require.NoError(t, err) + + builder := m.NewMapResponseBuilder(nodes[0].ID()).WithCapabilityVersion(tailcfg.CurrentCapabilityVersion) + want, wantErr := builder.buildTailPeersCapMapsReference(peers) + got, gotErr := builder.buildTailPeers(peers) + require.Equal(t, wantErr, gotErr) + require.Equal(t, want, got, "nil/empty, ordering, duplicates, routes and peer caps must match") + require.NotNil(t, got) + + for _, peer := range got { + if peer.ID == nodes[1].ID().NodeID() { + require.Contains(t, peer.CapMap, nodecap.SuggestExitNode) + + if attrs.name != "empty" { + require.Equal(t, []netip.Prefix{netip.PrefixFrom(nodes[1].IPv6().Get(), 128)}, peer.Addresses) + require.NotContains(t, peer.CapMap, nodecap.Cap("disable-ipv4"), "self caps must not leak into the peer view") + } + } + + if peer.CapMap != nil { + peer.CapMap["caller-only"] = nil + delete(peer.CapMap, nodecap.SuggestExitNode) + } + } + + after, err := json.Marshal(m.state.NodeCapMaps()) + require.NoError(t, err) + require.Equal(t, before, after, "rendering must not mutate manager-owned inner maps") + + again, err := builder.buildTailPeers(peers) + require.NoError(t, err) + require.Equal(t, want, again, "output mutations must not affect later builds") + }) + } + + builder := m.NewMapResponseBuilder(999) + want, wantErr := builder.buildTailPeersCapMapsReference(views.SliceOf(nodes)) + got, gotErr := builder.buildTailPeers(views.SliceOf(nodes)) + require.Equal(t, want, got) + require.Equal(t, wantErr, gotErr) + require.ErrorIs(t, gotErr, ErrNodeNotFoundMapper) + + for _, pol := range []string{ + "", + `{"nodeAttrs":[{"target":["*"],"attr":["randomize-client-port"]}]}`, + `{"randomizeClientPort":true}`, + } { + if pol != "" { + _, err := m.state.SetPolicy([]byte(pol)) + require.NoError(t, err) + } + + want, err := m.fullMapResponseCapMapsReference(nodes[0].ID(), tailcfg.CurrentCapabilityVersion) + require.NoError(t, err) + got, err := m.fullMapResponse(nodes[0].ID(), tailcfg.CurrentCapabilityVersion) + require.NoError(t, err) + + want.ControlTime, got.ControlTime = nil, nil + require.Equal(t, want, got) + } + }) + } +} + +var ( + capMapsBuilderBenchmarkSink []*tailcfg.Node + capMapsResponseBenchmarkSink *tailcfg.MapResponse +) + +func BenchmarkBuildTailPeersCapMaps(b *testing.B) { + previousLevel := zerolog.GlobalLevel() + + zerolog.SetGlobalLevel(zerolog.InfoLevel) + b.Cleanup(func() { zerolog.SetGlobalLevel(previousLevel) }) + + for _, n := range []int{100, 1000, 10000} { + for _, attrs := range []struct{ name, policy string }{ + {"empty", `[]`}, + {"partial", `[{"target":["100.64.0.2"],"attr":["disable-ipv4"]}]`}, + {"all", `[{"target":["*"],"attr":["disable-ipv4"]}]`}, + } { + b.Run(fmt.Sprintf("N%d/%s", n, attrs.name), func(b *testing.B) { + m, nodes := capMapsMapperFixture(b, n, attrs.policy) + builder := m.NewMapResponseBuilder(nodes[0].ID()).WithCapabilityVersion(tailcfg.CurrentCapabilityVersion) + + allPeers := views.SliceOf(nodes[1:]) + for _, k := range []int{0, 1, 8, allPeers.Len()} { + peers := allPeers.Slice(0, k) + b.Run(fmt.Sprintf("K%d/reference", k), func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + var err error + + capMapsBuilderBenchmarkSink, err = builder.buildTailPeersCapMapsReference(peers) + if err != nil { + b.Fatal(err) + } + } + }) + b.Run(fmt.Sprintf("K%d/selected", k), func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + var err error + + capMapsBuilderBenchmarkSink, err = builder.buildTailPeers(peers) + if err != nil { + b.Fatal(err) + } + } + }) + } + + b.Run("fullresponse/reference", func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + var err error + + capMapsResponseBenchmarkSink, err = m.fullMapResponseCapMapsReference(nodes[0].ID(), tailcfg.CurrentCapabilityVersion) + if err != nil { + b.Fatal(err) + } + } + }) + b.Run("fullresponse/selected", func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + var err error + + capMapsResponseBenchmarkSink, err = m.fullMapResponse(nodes[0].ID(), tailcfg.CurrentCapabilityVersion) + if err != nil { + b.Fatal(err) + } + } + }) + }) + } + } +} diff --git a/hscontrol/policy/pm.go b/hscontrol/policy/pm.go index ba8a61051..9ad30f525 100644 --- a/hscontrol/policy/pm.go +++ b/hscontrol/policy/pm.go @@ -69,19 +69,23 @@ type PolicyManager interface { // both fields are empty and the caller falls back to existing behavior. ViaRoutesForPeer(viewer, peer types.NodeView) types.ViaRouteResult - // NodeCapMap returns the policy-derived CapMap for the given node, - // or nil when no nodeAttrs entry targets it. The returned map is - // owned by the manager; treat it as read-only and copy before - // merging into a [tailcfg.Node]. It describes the node's own - // capabilities, not a per-viewer view. + // NodeCapMap returns compiled policy capabilities for the given node, + // or nil when none target it. The map is a defensive clone; its values + // are shared and read-only. It describes the node's own capabilities, + // not a per-viewer view. NodeCapMap(id types.NodeID) tailcfg.NodeCapMap - // NodeCapMaps returns a snapshot of the per-node policy CapMap so - // callers can amortise lock acquisitions over a peer loop. The - // outer map is a fresh container; the inner [tailcfg.NodeCapMap] - // values are shared with the manager and read-only. + // NodeCapMaps snapshots every node's policy CapMap. The outer map is + // caller-owned; inner maps are shared with the manager and read-only. + // NodeCapMapsForPeers limits copying when fewer peers are requested. NodeCapMaps() map[types.NodeID]tailcfg.NodeCapMap + // NodeCapMapsForPeers snapshots requested peer caps under one lock. It may + // include unselected entries when copying the whole compiled table is less work. + // The outer map is caller-owned; inner maps are shared and read-only. + // Empty selections or no compiled policy capabilities return nil; missing IDs read nil. + NodeCapMapsForPeers(peers views.Slice[types.NodeView]) map[types.NodeID]tailcfg.NodeCapMap + // NodesWithChangedCapMap returns the IDs of nodes whose nodeAttrs // CapMap shifted during recent updateLocked calls. The buffer // drains on read; callers consume it once per update cycle to diff --git a/hscontrol/policy/v2/policy.go b/hscontrol/policy/v2/policy.go index fce05656e..d5ff21043 100644 --- a/hscontrol/policy/v2/policy.go +++ b/hscontrol/policy/v2/policy.go @@ -2068,10 +2068,9 @@ func (pm *PolicyManager) refreshNodeAttrsLocked(newMap map[types.NodeID]tailcfg. } } -// NodeCapMap returns the policy-derived CapMap for the given node, or -// nil when the node has no nodeAttrs entries that target it. The -// returned map is a defensive clone — caller mutations cannot reach -// the manager-owned cache. +// NodeCapMap returns compiled policy capabilities for the given node, or nil +// when none target it. The map is a defensive clone; its values are shared +// and read-only. func (pm *PolicyManager) NodeCapMap(id types.NodeID) tailcfg.NodeCapMap { if pm == nil { return nil @@ -2091,11 +2090,9 @@ func (pm *PolicyManager) NodeCapMap(id types.NodeID) tailcfg.NodeCapMap { return out } -// NodeCapMaps returns a snapshot of the per-node policy CapMap. The -// mapper calls this once per request to amortise lock acquisitions -// over a peer-loop instead of taking the lock per peer. The returned -// map is a fresh container; the inner [tailcfg.NodeCapMap] values are -// shared with the manager and must be treated as read-only. +// NodeCapMaps snapshots every node's policy CapMap. The outer map is +// caller-owned; inner maps are shared with the manager and read-only. +// NodeCapMapsForPeers limits copying when fewer peers are requested. func (pm *PolicyManager) NodeCapMaps() map[types.NodeID]tailcfg.NodeCapMap { if pm == nil { return nil @@ -2110,6 +2107,36 @@ func (pm *PolicyManager) NodeCapMaps() map[types.NodeID]tailcfg.NodeCapMap { return out } +// NodeCapMapsForPeers snapshots requested peer caps under one lock. It may +// include unselected entries when copying the whole compiled table is less work. +// The outer map is caller-owned; inner maps are shared and read-only. +// Empty selections or no compiled policy capabilities return nil; missing IDs read nil. +func (pm *PolicyManager) NodeCapMapsForPeers(peers views.Slice[types.NodeView]) map[types.NodeID]tailcfg.NodeCapMap { + if pm == nil || peers.Len() == 0 { + return nil + } + + pm.mu.RLock() + defer pm.mu.RUnlock() + + if len(pm.nodeAttrsMap) == 0 { + return nil + } + + out := make(map[types.NodeID]tailcfg.NodeCapMap, min(peers.Len(), len(pm.nodeAttrsMap))) + if peers.Len() >= len(pm.nodeAttrsMap) { + maps.Copy(out, pm.nodeAttrsMap) + } else { + for _, peer := range peers.All() { + if caps, ok := pm.nodeAttrsMap[peer.ID()]; ok { + out[peer.ID()] = caps + } + } + } + + return out +} + // NodesWithChangedCapMap returns the IDs of nodes whose nodeAttrs // CapMap shifted across one or more [PolicyManager.updateLocked] calls // since the last drain. The buffer drains on return. diff --git a/hscontrol/policy/v2/policy_test.go b/hscontrol/policy/v2/policy_test.go index 5e28831b1..18eca4d92 100644 --- a/hscontrol/policy/v2/policy_test.go +++ b/hscontrol/policy/v2/policy_test.go @@ -2,6 +2,7 @@ package v2 import ( "fmt" + "maps" "net/netip" "slices" "testing" @@ -11,10 +12,13 @@ import ( "github.com/juanfont/headscale/hscontrol/policy/matcher" "github.com/juanfont/headscale/hscontrol/types" "github.com/puzpuzpuz/xsync/v4" + "github.com/rs/zerolog" "github.com/stretchr/testify/require" "pgregory.net/rapid" "tailscale.com/net/tsaddr" "tailscale.com/tailcfg" + "tailscale.com/tailcfg/nodecap" + "tailscale.com/types/views" ) func node(name, ipv4, ipv6 string, user types.User) *types.Node { @@ -3830,3 +3834,344 @@ func TestFailedSetNodesKeepsCompiledPolicy(t *testing.T) { require.Equal(t, wantSSH, gotSSH, "node %s SSH after the retry", n.Hostname) } } + +// nodeCapMapsReference is the complete bulk accessor from d00f135bf. +func (pm *PolicyManager) nodeCapMapsReference() map[types.NodeID]tailcfg.NodeCapMap { + if pm == nil { + return nil + } + + pm.mu.RLock() + defer pm.mu.RUnlock() + + out := make(map[types.NodeID]tailcfg.NodeCapMap, len(pm.nodeAttrsMap)) + maps.Copy(out, pm.nodeAttrsMap) + + return out +} + +func TestNodeCapMapsForPeersReference(t *testing.T) { + var nilManager *PolicyManager + require.Nil(t, nilManager.NodeCapMaps()) + require.Nil(t, nilManager.nodeCapMapsReference()) + require.Nil(t, nilManager.NodeCapMapsForPeers(views.SliceOf([]types.NodeView{(&types.Node{ID: 1}).View()}))) + + for _, attrs := range []struct { + name string + caps map[types.NodeID]tailcfg.NodeCapMap + }{ + {"nil", nil}, + {"empty", map[types.NodeID]tailcfg.NodeCapMap{}}, + {"single", map[types.NodeID]tailcfg.NodeCapMap{3: {"disable-ipv4": nil}}}, + {"partial", map[types.NodeID]tailcfg.NodeCapMap{ + 1: nil, + 2: {}, + 3: {"disable-ipv4": nil, "value": {`{"a":1}`}}, + 5: {"randomize-client-port": {}}, + }}, + } { + t.Run(attrs.name, func(t *testing.T) { + pm := &PolicyManager{nodeAttrsMap: attrs.caps} + all := pm.nodeCapMapsReference() + require.NotNil(t, all) + require.Equal(t, all, pm.NodeCapMaps()) + + for _, selection := range []struct { + name string + ids []types.NodeID + }{ + {"nil", nil}, + {"empty", []types.NodeID{}}, + {"missing", []types.NodeID{99}}, + {"one", []types.NodeID{3}}, + {"below threshold", []types.NodeID{5, 3, 3}}, + {"threshold missing", []types.NodeID{99, 99, 99, 99}}, + {"duplicates", []types.NodeID{3, 5, 3, 99, 1, 2}}, + {"full", []types.NodeID{5, 4, 3, 2, 1}}, + } { + t.Run(selection.name, func(t *testing.T) { + var peers []types.NodeView + for _, id := range selection.ids { + peers = append(peers, (&types.Node{ID: id}).View()) + } + + if selection.ids != nil && len(peers) == 0 { + peers = []types.NodeView{} + } + + before := slices.Clone(peers) + + var want map[types.NodeID]tailcfg.NodeCapMap + if len(peers) > 0 && len(all) > 0 { + want = make(map[types.NodeID]tailcfg.NodeCapMap) + + if len(peers) >= len(all) { + maps.Copy(want, all) + } else { + for _, peer := range peers { + if caps, ok := all[peer.ID()]; ok { + want[peer.ID()] = caps + } + } + } + } + + got := pm.NodeCapMapsForPeers(views.SliceOf(peers)) + require.Equal(t, want, got) + require.Equal(t, before, peers) + + for _, peer := range peers { + require.Equal(t, all[peer.ID()], got[peer.ID()], "requested lookups must match the bulk snapshot") + } + + if got != nil { + clear(got) + got[99] = tailcfg.NodeCapMap{"caller-only": nil} + } + + require.Equal(t, all, pm.NodeCapMaps(), "outer mutations must not reach the manager") + }) + } + + wantAll := pm.nodeCapMapsReference() + + clear(all) + require.Equal(t, wantAll, pm.NodeCapMaps(), "bulk container remains independently owned") + }) + } +} + +func TestNodeCapMapsForPeersSnapshots(t *testing.T) { + user := types.User{ID: 1, Name: "capuser"} + nodes := types.Nodes{ + node("one", "100.64.0.1", "fd7a:115c:a1e0::1", user), + node("two", "100.64.0.2", "fd7a:115c:a1e0::2", user), + } + nodes[0].ID, nodes[1].ID = 1, 2 + pm, err := NewPolicyManager([]byte(`{"nodeAttrs":[{"target":["*"],"attr":["disable-ipv4"]}]}`), []types.User{user}, nodes.ViewSlice()) + require.NoError(t, err) + + old := pm.NodeCapMapsForPeers(nodes.ViewSlice()) + all := pm.nodeCapMapsReference() + + require.Len(t, old, 2) + require.Equal(t, all, old) + // Probe the documented shared-inner contract, then restore before any policy update. + old[1]["alias-probe"] = []tailcfg.RawMessage{`{"shared":true}`} + + require.Contains(t, pm.nodeAttrsMap[1], nodecap.Cap("alias-probe")) + require.Same(t, &old[1]["alias-probe"][0], &all[1]["alias-probe"][0]) + delete(old[1], "alias-probe") + + owned := pm.NodeCapMap(1) + delete(owned, "disable-ipv4") + require.Contains(t, old[1], nodecap.Cap("disable-ipv4"), "single-node accessor still clones its inner map") + + _, err = pm.SetPolicy([]byte(`{"nodeAttrs":[{"target":["100.64.0.2"],"attr":["randomize-client-port"]}]}`)) + require.NoError(t, err) + + current := pm.NodeCapMapsForPeers(nodes.ViewSlice()) + require.Equal(t, pm.nodeCapMapsReference(), current) + require.NotContains(t, current, types.NodeID(1)) + require.Contains(t, current[2], nodecap.Cap("randomize-client-port")) + + for _, caps := range old { + require.Equal(t, tailcfg.NodeCapMap{"disable-ipv4": nil}, caps, "old snapshot survives policy replacement") + } + + _, err = pm.SetPolicy([]byte(`{}`)) + require.NoError(t, err) + require.Nil(t, pm.NodeCapMapsForPeers(nodes.ViewSlice())) + require.Contains(t, current[2], nodecap.Cap("randomize-client-port")) + require.NotNil(t, pm.NodeCapMaps(), "bulk empty result must remain non-nil") +} + +func TestNodeCapMapsForPeersGlobalRandomizeClientPort(t *testing.T) { + users, nodes := benchNodes(2) + pm, err := NewPolicyManager([]byte(`{"randomizeClientPort":true}`), users, nodes.ViewSlice()) + require.NoError(t, err) + + all := pm.nodeCapMapsReference() + require.Len(t, all, 2) + + peers := nodes.ViewSlice().Slice(0, 1) + got := pm.NodeCapMapsForPeers(peers) + require.Equal(t, map[types.NodeID]tailcfg.NodeCapMap{ + nodes[0].ID: {"randomize-client-port": nil}, + }, got) + require.Equal(t, all, pm.NodeCapMapsForPeers(nodes.ViewSlice())) + require.Nil(t, pm.NodeCapMapsForPeers(views.SliceOf([]types.NodeView(nil)))) +} + +func TestNodeCapMapsForPeersConcurrentPolicy(t *testing.T) { + for _, k := range []int{3, 8} { + t.Run(fmt.Sprintf("K%d", k), func(t *testing.T) { + users, nodes := benchNodes(8) + policies := []string{ + `{"nodeAttrs":[{"target":["*"],"attr":["disable-ipv4"]}]}`, + `{"nodeAttrs":[{"target":["*"],"attr":["randomize-client-port"]}]}`, + } + pm, err := NewPolicyManager([]byte(policies[0]), users, nodes.ViewSlice()) + require.NoError(t, err) + + peers := nodes.ViewSlice().Slice(0, k) + generations := []tailcfg.NodeCapMap{ + {"disable-ipv4": nil}, + {"randomize-client-port": nil}, + } + checkSnapshot := func(generation int) { + t.Helper() + + snapshot := pm.NodeCapMapsForPeers(peers) + if len(snapshot) != peers.Len() { + t.Errorf("snapshot has %d entries, want %d", len(snapshot), peers.Len()) + } + + first := snapshot[nodes[0].ID] + if !maps.EqualFunc(generations[0], first, slices.Equal) && + !maps.EqualFunc(generations[1], first, slices.Equal) { + t.Errorf("snapshot has unexpected capabilities: %v", first) + } + + for _, caps := range snapshot { + if !maps.EqualFunc(first, caps, slices.Equal) { + t.Error("snapshot mixed policy generations") + } + + if generation >= 0 && !maps.EqualFunc(generations[generation], caps, slices.Equal) { + t.Errorf("snapshot did not observe published generation %d: %v", generation, caps) + } + } + } + checkSnapshot(0) + + published := make(chan int) + acknowledged := make(chan struct{}) + writerErrors := make(chan error, 1) + + go func() { + defer close(published) + + var err error + defer func() { writerErrors <- err }() + + for i := range 100 { + generation := (i + 1) % 2 + + _, err = pm.SetPolicy([]byte(policies[generation])) + if err != nil { + return + } + + published <- generation + + <-acknowledged + } + }() + + var observed [2]int + + for active := true; active; { + checkSnapshot(-1) + + select { + case generation, ok := <-published: + if !ok { + active = false + continue + } + + checkSnapshot(generation) + observed[generation]++ + + acknowledged <- struct{}{} + default: + } + } + + require.NoError(t, <-writerErrors) + checkSnapshot(0) + require.Equal(t, [2]int{50, 50}, observed) + }) + } +} + +var nodeCapMapsBenchmarkSink map[types.NodeID]tailcfg.NodeCapMap + +// Each call constructs a fresh outer snapshot; there is no retained selector index. +// Direct method calls keep interface adapters out of this isolated measurement. +func BenchmarkNodeCapMapsSelection(b *testing.B) { + previousLevel := zerolog.GlobalLevel() + + zerolog.SetGlobalLevel(zerolog.InfoLevel) + b.Cleanup(func() { zerolog.SetGlobalLevel(previousLevel) }) + + for _, n := range []int{100, 1000, 10000} { + _, nodes := benchNodes(n) + allPeers := nodes.ViewSlice() + + for _, stride := range []int{0, 10, 1, n} { + b.Run(fmt.Sprintf("N%d/attrStride%d", n, stride), func(b *testing.B) { + pm := &PolicyManager{nodeAttrsMap: make(map[types.NodeID]tailcfg.NodeCapMap)} + + if stride > 0 { + for i := 0; i < n; i += stride { + pm.nodeAttrsMap[nodes[i].ID] = tailcfg.NodeCapMap{"disable-ipv4": {}} + } + } + + for _, k := range []int{0, 1, 8, n} { + peers := allPeers.Slice(0, k) + b.Run(fmt.Sprintf("K%d/reference", k), func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + nodeCapMapsBenchmarkSink = pm.nodeCapMapsReference() + } + }) + b.Run(fmt.Sprintf("K%d/selected", k), func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + nodeCapMapsBenchmarkSink = pm.NodeCapMapsForPeers(peers) + } + }) + } + }) + } + } +} + +func BenchmarkNodeCapMapsConstruction(b *testing.B) { + previousLevel := zerolog.GlobalLevel() + + zerolog.SetGlobalLevel(zerolog.InfoLevel) + b.Cleanup(func() { zerolog.SetGlobalLevel(previousLevel) }) + + for _, n := range []int{100, 1000, 10000} { + users, nodes := benchNodes(n) + peers := nodes.ViewSlice().Slice(0, 1) + + for _, attrs := range []string{`[]`, `[{"target":["*"],"attr":["disable-ipv4"]}]`} { + pol := []byte(`{"acls":[],"nodeAttrs":` + attrs + `}`) + for _, selected := range []bool{false, true} { + b.Run(fmt.Sprintf("N%d/attrs%t/selected%t", n, attrs != `[]`, selected), func(b *testing.B) { + b.ReportAllocs() + + for b.Loop() { + pm, err := NewPolicyManager(pol, users, nodes.ViewSlice()) + if err != nil { + b.Fatal(err) + } + + if selected { + nodeCapMapsBenchmarkSink = pm.NodeCapMapsForPeers(peers) + } else { + nodeCapMapsBenchmarkSink = pm.nodeCapMapsReference() + } + } + }) + } + } + } +} diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index e73238486..d465375f7 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -1317,6 +1317,12 @@ func (s *State) NodeCapMaps() map[types.NodeID]tailcfg.NodeCapMap { return s.polMan.NodeCapMaps() } +// NodeCapMapsForPeers returns one policy snapshot containing the requested peer caps. +// It may include other entries; inner maps are shared and read-only. +func (s *State) NodeCapMapsForPeers(peers views.Slice[types.NodeView]) map[types.NodeID]tailcfg.NodeCapMap { + return s.polMan.NodeCapMapsForPeers(peers) +} + // NodeCanHaveTag checks if a node is allowed to have a specific tag. func (s *State) NodeCanHaveTag(node types.NodeView, tag string) bool { return s.polMan.NodeCanHaveTag(node, tag)