From a5d2cdfcedf4f0d567421c728c3eee328c8fd61d Mon Sep 17 00:00:00 2001
From: Claire Wang <claire@tailscale.com>
Date: Tue, 19 Mar 2024 11:55:31 -0400
Subject: [PATCH] control/controlclient: fix sending peer capmap changes
 Instead of just checking if a peer capmap is nil, also check if the previous
 state's capmap was nil. Updates tailscale/corp#17516 Co-authored-by: Adrian
 Dewhurst <adrian@tailscale.com>

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

diff --git a/control/controlclient/map.go b/control/controlclient/map.go
index d44310028002a..c460ce0d7a317 100644
--- a/control/controlclient/map.go
+++ b/control/controlclient/map.go
@@ -665,7 +665,18 @@ func peerChangeDiff(was tailcfg.NodeView, n *tailcfg.Node) (_ *tailcfg.PeerChang
 			}
 		case "CapMap":
 			if n.CapMap != nil {
-				pc().CapMap = n.CapMap
+				for k, v := range n.CapMap {
+					wasCap, ok := was.CapMap().GetOk(k)
+					if !ok || !views.SliceEqual(wasCap, views.SliceOf(v)) {
+						pc().CapMap = n.CapMap
+						break
+					}
+				}
+				if was.CapMap().IsNil() || len(n.CapMap) != was.CapMap().Len() {
+					pc().CapMap = n.CapMap
+				}
+			} else if !was.CapMap().IsNil() {
+				pc().CapMap = make(tailcfg.NodeCapMap)
 			}
 		case "Tags":
 			if !views.SliceEqual(was.Tags(), views.SliceOf(n.Tags)) {
diff --git a/control/controlclient/map_test.go b/control/controlclient/map_test.go
index bdbfaf6e4aed9..817e730caab2f 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)