From 03fd19958a42c184f5d471385e4b656ad08a96b6 Mon Sep 17 00:00:00 2001
From: Johan Siebens <johan.siebens@gmail.com>
Date: Tue, 24 May 2022 20:35:05 +0200
Subject: [PATCH] fix: display tagged devices in user profile

---
 internal/domain/machine.go        |  4 ++++
 internal/handlers/poll_net_map.go | 20 +++++++++++---------
 internal/mapping/mapping.go       | 25 ++++++++++++++++++-------
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/internal/domain/machine.go b/internal/domain/machine.go
index a9212ce..10e51ae 100644
--- a/internal/domain/machine.go
+++ b/internal/domain/machine.go
@@ -61,6 +61,10 @@ func (m *Machine) HasTag(tag string) bool {
 	return false
 }
 
+func (m *Machine) HasTags() bool {
+	return len(m.Tags) != 0
+}
+
 func (m *Machine) IsAllowedIP(i netaddr.IP) bool {
 	if m.HasIP(i) {
 		return true
diff --git a/internal/handlers/poll_net_map.go b/internal/handlers/poll_net_map.go
index 19505e7..bb49ec8 100644
--- a/internal/handlers/poll_net_map.go
+++ b/internal/handlers/poll_net_map.go
@@ -219,7 +219,7 @@ func (h *PollNetMapHandler) createKeepAliveResponse(binder bind.Binder, request
 func (h *PollNetMapHandler) createMapResponse(m *domain.Machine, binder bind.Binder, request *tailcfg.MapRequest, delta bool, prevSyncedPeerIDs map[uint64]bool) ([]byte, map[uint64]bool, error) {
 	ctx := context.TODO()
 
-	node, err := mapping.ToNode(m, true)
+	node, user, err := mapping.ToNode(m, true)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -229,11 +229,7 @@ func (h *PollNetMapHandler) createMapResponse(m *domain.Machine, binder bind.Bin
 		return nil, nil, err
 	}
 
-	users, err := h.repository.ListUsers(ctx, m.TailnetID)
-	if err != nil {
-		return nil, nil, err
-	}
-
+	var users = []tailcfg.UserProfile{*user}
 	var changedPeers []*tailcfg.Node
 	var removedPeers []tailcfg.NodeID
 
@@ -243,19 +239,25 @@ func (h *PollNetMapHandler) createMapResponse(m *domain.Machine, binder bind.Bin
 	}
 
 	syncedPeerIDs := map[uint64]bool{}
+	syncedUserIDs := map[tailcfg.UserID]bool{}
 
 	for _, peer := range candidatePeers {
 		if peer.IsExpired() {
 			continue
 		}
 		if domain.IsValidPeer(policies, m, &peer) || domain.IsValidPeer(policies, &peer, m) {
-			n, err := mapping.ToNode(&peer, h.brokers(peer.TailnetID).IsConnected(peer.ID))
+			n, u, err := mapping.ToNode(&peer, h.brokers(peer.TailnetID).IsConnected(peer.ID))
 			if err != nil {
 				return nil, nil, err
 			}
 			changedPeers = append(changedPeers, n)
 			syncedPeerIDs[peer.ID] = true
 			delete(prevSyncedPeerIDs, peer.ID)
+
+			if _, ok := syncedUserIDs[u.ID]; !ok {
+				users = append(users, *u)
+				syncedUserIDs[u.ID] = true
+			}
 		}
 	}
 
@@ -287,7 +289,7 @@ func (h *PollNetMapHandler) createMapResponse(m *domain.Machine, binder bind.Bin
 			DERPMap:      derpMap,
 			Domain:       dnsname.SanitizeHostname(m.Tailnet.Name),
 			Peers:        changedPeers,
-			UserProfiles: mapping.ToUserProfiles(users),
+			UserProfiles: users,
 			ControlTime:  &controlTime,
 		}
 	} else {
@@ -297,7 +299,7 @@ func (h *PollNetMapHandler) createMapResponse(m *domain.Machine, binder bind.Bin
 			DERPMap:      derpMap,
 			PeersChanged: changedPeers,
 			PeersRemoved: removedPeers,
-			UserProfiles: mapping.ToUserProfiles(users),
+			UserProfiles: users,
 			ControlTime:  &controlTime,
 		}
 	}
diff --git a/internal/mapping/mapping.go b/internal/mapping/mapping.go
index 088dbd6..04c6fe6 100644
--- a/internal/mapping/mapping.go
+++ b/internal/mapping/mapping.go
@@ -72,22 +72,22 @@ func ToDNSConfig(tailnet *domain.Tailnet, c *domain.DNSConfig) *tailcfg.DNSConfi
 	return config
 }
 
-func ToNode(m *domain.Machine, connected bool) (*tailcfg.Node, error) {
+func ToNode(m *domain.Machine, connected bool) (*tailcfg.Node, *tailcfg.UserProfile, error) {
 	nKey, err := util.ParseNodePublicKey(m.NodeKey)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	mKey, err := util.ParseMachinePublicKey(m.MachineKey)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	var discoKey key.DiscoPublic
 	if m.DiscoKey != "" {
 		dKey, err := util.ParseDiscoPublicKey(m.DiscoKey)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		discoKey = *dKey
 	}
@@ -101,7 +101,7 @@ func ToNode(m *domain.Machine, connected bool) (*tailcfg.Node, error) {
 	if !m.IPv4.IsZero() {
 		ipv4, err := m.IPv4.Prefix(32)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		addrs = append(addrs, ipv4)
 		allowedIPs = append(allowedIPs, ipv4)
@@ -110,7 +110,7 @@ func ToNode(m *domain.Machine, connected bool) (*tailcfg.Node, error) {
 	if !m.IPv6.IsZero() {
 		ipv6, err := m.IPv6.Prefix(128)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		addrs = append(addrs, ipv6)
 		allowedIPs = append(allowedIPs, ipv6)
@@ -169,7 +169,18 @@ func ToNode(m *domain.Machine, connected bool) (*tailcfg.Node, error) {
 		n.LastSeen = &l
 	}
 
-	return &n, nil
+	var user = ToUserProfile(m.User)
+
+	if m.HasTags() {
+		n.User = tailcfg.UserID(m.ID)
+		user = tailcfg.UserProfile{
+			ID:          tailcfg.UserID(m.ID),
+			LoginName:   "tagged-devices",
+			DisplayName: "Tagged Devices",
+		}
+	}
+
+	return &n, &user, nil
 }
 
 func ToUserProfile(u domain.User) tailcfg.UserProfile {