From 262fa8a01ed991fcbc9ab6fb2e65c8c4e0fcf7d9 Mon Sep 17 00:00:00 2001
From: Claire Wang <clairew@users.noreply.github.com>
Date: Wed, 3 Apr 2024 10:55:28 -0400
Subject: [PATCH] ipn/ipnlocal: populate peers' capabilities (#11365)

Populates capabilties field of peers in ipn status.
Updates tailscale/corp#17516

Signed-off-by: Claire Wang <claire@tailscale.com>
---
 ipn/ipnlocal/local.go      |  8 ++++
 ipn/ipnlocal/local_test.go | 94 ++++++++++++++++++++++++++++++++++++++
 ipn/ipnstate/ipnstate.go   |  6 +++
 3 files changed, 108 insertions(+)

diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index b954afffc..a83ebdce3 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -894,6 +894,14 @@ func (b *LocalBackend) populatePeerStatusLocked(sb *ipnstate.StatusBuilder) {
 			ExitNode:        p.StableID() != "" && p.StableID() == exitNodeID,
 			SSH_HostKeys:    p.Hostinfo().SSH_HostKeys().AsSlice(),
 			Location:        p.Hostinfo().Location(),
+			Capabilities:    p.Capabilities().AsSlice(),
+		}
+		if cm := p.CapMap(); cm.Len() > 0 {
+			ps.CapMap = make(tailcfg.NodeCapMap, cm.Len())
+			cm.Range(func(k tailcfg.NodeCapability, v views.Slice[tailcfg.RawMessage]) bool {
+				ps.CapMap[k] = v.AsSlice()
+				return true
+			})
 		}
 		peerStatusFromNode(ps, p)
 
diff --git a/ipn/ipnlocal/local_test.go b/ipn/ipnlocal/local_test.go
index 2713fff7d..0119bc7ac 100644
--- a/ipn/ipnlocal/local_test.go
+++ b/ipn/ipnlocal/local_test.go
@@ -736,6 +736,100 @@ func TestStatusWithoutPeers(t *testing.T) {
 	}
 }
 
+func TestStatusPeerCapabilities(t *testing.T) {
+	tests := []struct {
+		name                     string
+		peers                    []tailcfg.NodeView
+		expectedPeerCapabilities map[tailcfg.StableNodeID][]tailcfg.NodeCapability
+		expectedPeerCapMap       map[tailcfg.StableNodeID]tailcfg.NodeCapMap
+	}{
+		{
+			name: "peers-with-capabilities",
+			peers: []tailcfg.NodeView{
+				(&tailcfg.Node{
+					ID:              1,
+					StableID:        "foo",
+					IsWireGuardOnly: true,
+					Hostinfo:        (&tailcfg.Hostinfo{}).View(),
+					Capabilities:    []tailcfg.NodeCapability{tailcfg.CapabilitySSH},
+					CapMap: (tailcfg.NodeCapMap)(map[tailcfg.NodeCapability][]tailcfg.RawMessage{
+						tailcfg.CapabilitySSH: nil,
+					}),
+				}).View(),
+				(&tailcfg.Node{
+					ID:           2,
+					StableID:     "bar",
+					Hostinfo:     (&tailcfg.Hostinfo{}).View(),
+					Capabilities: []tailcfg.NodeCapability{tailcfg.CapabilityAdmin},
+					CapMap: (tailcfg.NodeCapMap)(map[tailcfg.NodeCapability][]tailcfg.RawMessage{
+						tailcfg.CapabilityAdmin: {`{"test": "true}`},
+					}),
+				}).View(),
+			},
+			expectedPeerCapabilities: map[tailcfg.StableNodeID][]tailcfg.NodeCapability{
+				tailcfg.StableNodeID("foo"): {tailcfg.CapabilitySSH},
+				tailcfg.StableNodeID("bar"): {tailcfg.CapabilityAdmin},
+			},
+			expectedPeerCapMap: map[tailcfg.StableNodeID]tailcfg.NodeCapMap{
+				tailcfg.StableNodeID("foo"): (tailcfg.NodeCapMap)(map[tailcfg.NodeCapability][]tailcfg.RawMessage{
+					tailcfg.CapabilitySSH: nil,
+				}),
+				tailcfg.StableNodeID("bar"): (tailcfg.NodeCapMap)(map[tailcfg.NodeCapability][]tailcfg.RawMessage{
+					tailcfg.CapabilityAdmin: {`{"test": "true}`},
+				}),
+			},
+		},
+		{
+			name: "peers-without-capabilities",
+			peers: []tailcfg.NodeView{
+				(&tailcfg.Node{
+					ID:              1,
+					StableID:        "foo",
+					IsWireGuardOnly: true,
+					Hostinfo:        (&tailcfg.Hostinfo{}).View(),
+				}).View(),
+				(&tailcfg.Node{
+					ID:       2,
+					StableID: "bar",
+					Hostinfo: (&tailcfg.Hostinfo{}).View(),
+				}).View(),
+			},
+		},
+	}
+	b := newTestLocalBackend(t)
+
+	var cc *mockControl
+	b.SetControlClientGetterForTesting(func(opts controlclient.Options) (controlclient.Client, error) {
+		cc = newClient(t, opts)
+
+		t.Logf("ccGen: new mockControl.")
+		cc.called("New")
+		return cc, nil
+	})
+	b.Start(ipn.Options{})
+	b.Login(nil)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b.setNetMapLocked(&netmap.NetworkMap{
+				SelfNode: (&tailcfg.Node{
+					MachineAuthorized: true,
+					Addresses:         ipps("100.101.101.101"),
+				}).View(),
+				Peers: tt.peers,
+			})
+			got := b.Status()
+			for _, peer := range got.Peer {
+				if !reflect.DeepEqual(peer.Capabilities, tt.expectedPeerCapabilities[peer.ID]) {
+					t.Errorf("peer capabilities: expected %v got %v", tt.expectedPeerCapabilities, peer.Capabilities)
+				}
+				if !reflect.DeepEqual(peer.CapMap, tt.expectedPeerCapMap[peer.ID]) {
+					t.Errorf("peer capmap: expected %v got %v", tt.expectedPeerCapMap, peer.CapMap)
+				}
+			}
+		})
+	}
+}
+
 // legacyBackend was the interface between Tailscale frontends
 // (e.g. cmd/tailscale, iOS/MacOS/Windows GUIs) and the tailscale
 // backend (e.g. cmd/tailscaled) running on the same machine.
diff --git a/ipn/ipnstate/ipnstate.go b/ipn/ipnstate/ipnstate.go
index c9ad0d0da..d86bc1d26 100644
--- a/ipn/ipnstate/ipnstate.go
+++ b/ipn/ipnstate/ipnstate.go
@@ -496,6 +496,12 @@ func (sb *StatusBuilder) AddPeer(peer key.NodePublic, st *PeerStatus) {
 	if t := st.KeyExpiry; t != nil {
 		e.KeyExpiry = ptr.To(*t)
 	}
+	if v := st.CapMap; v != nil {
+		e.CapMap = v
+	}
+	if v := st.Capabilities; v != nil {
+		e.Capabilities = v
+	}
 	e.Location = st.Location
 }