From 221de0174506d4670a24c285dc07470b61d94b75 Mon Sep 17 00:00:00 2001
From: Claire Wang <clairew@users.noreply.github.com>
Date: Tue, 19 Mar 2024 18:56:06 -0400
Subject: [PATCH] control/controlclient: fix sending peer capmap changes
 (#11457)

Instead of just checking if a peer capmap is nil, compare the previous
state peer capmap with the new peer capmap.
Updates tailscale/corp#17516

Signed-off-by: Claire Wang <claire@tailscale.com>
---
 control/controlclient/map.go      | 17 +++++++++++++++--
 control/controlclient/map_test.go | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/control/controlclient/map.go b/control/controlclient/map.go
index d44310028..e03e57d7a 100644
--- a/control/controlclient/map.go
+++ b/control/controlclient/map.go
@@ -664,9 +664,22 @@ func peerChangeDiff(was tailcfg.NodeView, n *tailcfg.Node) (_ *tailcfg.PeerChang
 				pc().Cap = n.Cap
 			}
 		case "CapMap":
-			if n.CapMap != nil {
-				pc().CapMap = n.CapMap
+			if len(n.CapMap) != was.CapMap().Len() {
+				if n.CapMap == nil {
+					pc().CapMap = make(tailcfg.NodeCapMap)
+				} else {
+					pc().CapMap = n.CapMap
+				}
+				break
 			}
+			was.CapMap().Range(func(k tailcfg.NodeCapability, v views.Slice[tailcfg.RawMessage]) bool {
+				nv, ok := n.CapMap[k]
+				if !ok || !views.SliceEqual(v, views.SliceOf(nv)) {
+					pc().CapMap = n.CapMap
+					return false
+				}
+				return true
+			})
 		case "Tags":
 			if !views.SliceEqual(was.Tags(), views.SliceOf(n.Tags)) {
 				return nil, false
diff --git a/control/controlclient/map_test.go b/control/controlclient/map_test.go
index bdbfaf6e4..817e730ca 100644
--- a/control/controlclient/map_test.go
+++ b/control/controlclient/map_test.go
@@ -848,7 +848,36 @@ func TestPeerChangeDiff(t *testing.T) {
 			a:    &tailcfg.Node{ID: 1, SelfNodeV6MasqAddrForThisPeer: ptr.To(netip.MustParseAddr("2001::3456"))},
 			b:    &tailcfg.Node{ID: 1, SelfNodeV6MasqAddrForThisPeer: ptr.To(netip.MustParseAddr("2001::3006"))},
 			want: nil,
-		}}
+		},
+		{
+			name: "patch-capmap-add-value-to-existing-key",
+			a:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			b:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: []tailcfg.RawMessage{"true"}}},
+			want: &tailcfg.PeerChange{NodeID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: []tailcfg.RawMessage{"true"}}},
+		},
+		{
+			name: "patch-capmap-add-new-key",
+			a:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			b:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil, tailcfg.CapabilityDebug: nil}},
+			want: &tailcfg.PeerChange{NodeID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil, tailcfg.CapabilityDebug: nil}},
+		}, {
+			name: "patch-capmap-remove-key",
+			a:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			b:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{}},
+			want: &tailcfg.PeerChange{NodeID: 1, CapMap: tailcfg.NodeCapMap{}},
+		}, {
+			name: "patch-capmap-add-key-to-empty-map",
+			a:    &tailcfg.Node{ID: 1},
+			b:    &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			want: &tailcfg.PeerChange{NodeID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+		},
+		{
+			name:      "patch-capmap-no-change",
+			a:         &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			b:         &tailcfg.Node{ID: 1, CapMap: tailcfg.NodeCapMap{tailcfg.CapabilityAdmin: nil}},
+			wantEqual: true,
+		},
+	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			pc, ok := peerChangeDiff(tt.a.View(), tt.b)