From 64560bbf6e16202d7299f9875e80ae1a2c6ce7f0 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Wed, 18 Dec 2024 15:10:16 +0100
Subject: [PATCH] expand principles to correct login name

and if fix an issue where nodeip principles might not expand to all
relevant IPs instead of taking the first in a prefix.

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/policy/acls.go | 90 ++++++++++++++++++++++++++++++----------
 1 file changed, 67 insertions(+), 23 deletions(-)

diff --git a/hscontrol/policy/acls.go b/hscontrol/policy/acls.go
index 88461d27..3d7a6f4a 100644
--- a/hscontrol/policy/acls.go
+++ b/hscontrol/policy/acls.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 	"net/netip"
 	"os"
 	"slices"
@@ -361,37 +362,67 @@ func (pol *ACLPolicy) CompileSSHPolicy(
 			)
 		}
 
-		principals := make([]*tailcfg.SSHPrincipal, 0, len(sshACL.Sources))
-		for innerIndex, rawSrc := range sshACL.Sources {
-			if isWildcard(rawSrc) {
-				principals = append(principals, &tailcfg.SSHPrincipal{
+		var principals []*tailcfg.SSHPrincipal
+		for innerIndex, srcToken := range sshACL.Sources {
+			if isWildcard(srcToken) {
+				principals = []*tailcfg.SSHPrincipal{{
 					Any: true,
-				})
-			} else if isGroup(rawSrc) {
-				users, err := pol.expandUsersFromGroup(rawSrc)
+				}}
+				break
+			}
+
+			// If the token is a group, expand the users and validate
+			// them. Then use the .Username() to get the login name
+			// that corresponds with the User info in the netmap.
+			if isGroup(srcToken) {
+				usersFromGroup, err := pol.expandUsersFromGroup(srcToken)
 				if err != nil {
 					return nil, fmt.Errorf("parsing SSH policy, expanding user from group, index: %d->%d: %w", index, innerIndex, err)
 				}
 
-				for _, user := range users {
+				for _, userStr := range usersFromGroup {
+					user, err := findUserFromTokenOrErr(users, userStr)
+					if err != nil {
+						log.Trace().Err(err).Msg("user not found")
+						continue
+					}
+
 					principals = append(principals, &tailcfg.SSHPrincipal{
-						UserLogin: user,
-					})
-				}
-			} else {
-				expandedSrcs, err := pol.ExpandAlias(
-					peers,
-					users,
-					rawSrc,
-				)
-				if err != nil {
-					return nil, fmt.Errorf("parsing SSH policy, expanding alias, index: %d->%d: %w", index, innerIndex, err)
-				}
-				for _, expandedSrc := range expandedSrcs.Prefixes() {
-					principals = append(principals, &tailcfg.SSHPrincipal{
-						NodeIP: expandedSrc.Addr().String(),
+						UserLogin: user.Username(),
 					})
 				}
+
+				continue
+			}
+
+			// Try to check if the token is a user, if it is, then we
+			// can use the .Username() to get the login name that
+			// corresponds with the User info in the netmap.
+			// TODO(kradalby): This is a bit of a hack, and it should go
+			// away with the new policy where users can be reliably determined.
+			if user, err := findUserFromTokenOrErr(users, srcToken); err == nil {
+				principals = append(principals, &tailcfg.SSHPrincipal{
+					UserLogin: user.Username(),
+				})
+				continue
+			}
+
+			// This is kind of then non-ideal scenario where we dont really know
+			// what to do with the token, so we expand it to IP addresses of nodes.
+			// The pro here is that we have a pretty good lockdown on the mapping
+			// between users and node, but it can explode if a user owns many nodes.
+			ips, err := pol.ExpandAlias(
+				peers,
+				users,
+				srcToken,
+			)
+			if err != nil {
+				return nil, fmt.Errorf("parsing SSH policy, expanding alias, index: %d->%d: %w", index, innerIndex, err)
+			}
+			for addr := range ipSetAll(ips) {
+				principals = append(principals, &tailcfg.SSHPrincipal{
+					NodeIP: addr.String(),
+				})
 			}
 		}
 
@@ -411,6 +442,19 @@ func (pol *ACLPolicy) CompileSSHPolicy(
 	}, nil
 }
 
+// ipSetAll returns a function that iterates over all the IPs in the IPSet.
+func ipSetAll(ipSet *netipx.IPSet) iter.Seq[netip.Addr] {
+	return func(yield func(netip.Addr) bool) {
+		for _, rng := range ipSet.Ranges() {
+			for ip := rng.From(); ip.Compare(rng.To()) <= 0; ip = ip.Next() {
+				if !yield(ip) {
+					return
+				}
+			}
+		}
+	}
+}
+
 func sshCheckAction(duration string) (*tailcfg.SSHAction, error) {
 	sessionLength, err := time.ParseDuration(duration)
 	if err != nil {