policy/v2: use Alias and SSHUser types in SSHPolicyTest

hscontrol/policy/v2/sshtest.go

@@ -230,8 +230,13 @@ func runSSHPolicyTest(
 	nodes views.Slice[types.NodeView],
 	cache map[types.NodeID]*tailcfg.SSHPolicy,
 ) SSHPolicyTestResult {
+	srcLabel := ""
+	if test.Src != nil {
+		srcLabel = test.Src.String()
+	}
+
 	res := SSHPolicyTestResult{
-		Src:    test.Src,
+		Src:    srcLabel,
 		Passed: true,
 	}
 
@@ -239,7 +244,7 @@ func runSSHPolicyTest(
 	if err != nil {
 		res.Passed = false
 		res.Errors = append(res.Errors,
-			fmt.Sprintf("failed to resolve source %q: %v", test.Src, err))
+			fmt.Sprintf("failed to resolve source %q: %v", srcLabel, err))
 
 		return res
 	}
@@ -247,7 +252,7 @@ func runSSHPolicyTest(
 	if len(srcAddrs) == 0 {
 		res.Passed = false
 		res.Errors = append(res.Errors,
-			fmt.Sprintf("source %q resolved to no IP addresses", test.Src))
+			fmt.Sprintf("source %q resolved to no IP addresses", srcLabel))
 
 		return res
 	}
@@ -290,7 +295,7 @@ func runSSHPolicyTest(
 	for _, user := range test.Accept {
 		evaluateAssertion(
 			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user,
+			srcAddrs, dstNodes, user.String(),
 			assertAccept, &res,
 		)
 	}
@@ -298,7 +303,7 @@ func runSSHPolicyTest(
 	for _, user := range test.Deny {
 		evaluateAssertion(
 			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user,
+			srcAddrs, dstNodes, user.String(),
 			assertDeny, &res,
 		)
 	}
@@ -306,7 +311,7 @@ func runSSHPolicyTest(
 	for _, user := range test.Check {
 		evaluateAssertion(
 			pol, users, nodes, cache,
-			srcAddrs, dstNodes, user,
+			srcAddrs, dstNodes, user.String(),
 			assertCheck, &res,
 		)
 	}
@@ -438,23 +443,22 @@ func appendUserDst(m map[string][]string, user, dst string) map[string][]string
 	return m
 }
 
-// resolveSSHTestSource resolves the src alias into a list of
+// resolveSSHTestSource resolves the typed src alias into a list of
 // netip.Addr (one per principal address the SSH compiler would emit
 // for the same source). For user-shaped sources, srcUserID returns the
 // resolved user's ID so autogroup:self destinations can scope to the
 // same user. Returns ID 0 when the source is a tag, host, or IP.
 func resolveSSHTestSource(
-	src string,
+	src Alias,
 	pol *Policy,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
 ) ([]netip.Addr, uint, error) {
-	alias, err := parseAlias(src)
-	if err != nil {
-		return nil, 0, fmt.Errorf("invalid alias: %w", err)
+	if src == nil {
+		return nil, 0, nil
 	}
 
-	addrs, err := alias.Resolve(pol, users, nodes)
+	addrs, err := src.Resolve(pol, users, nodes)
 	if err != nil {
 		return nil, 0, fmt.Errorf("resolving: %w", err)
 	}
@@ -470,7 +474,7 @@ func resolveSSHTestSource(
 
 	var userID uint
 
-	u, ok := alias.(*Username)
+	u, ok := src.(*Username)
 	if ok {
 		resolved, rErr := u.resolveUser(users)
 		if rErr == nil {
@@ -490,7 +494,7 @@ func resolveSSHTestSource(
 // node's IPs via InIPSet (the same primitive the SSH compiler uses to
 // decide whether a node is a destination of a given rule).
 func resolveSSHTestDestNodes(
-	dsts []string,
+	dsts SSHTestDestinations,
 	pol *Policy,
 	users []types.User,
 	nodes views.Slice[types.NodeView],
@@ -503,12 +507,8 @@ func resolveSSHTestDestNodes(
 		emptyDsts []string
 	)
 
-	for _, dst := range dsts {
-		alias, err := parseAlias(dst)
-		if err != nil {
-			return nil, nil, fmt.Errorf("invalid destination %q: %w", dst, err)
-		}
-
+	for _, alias := range dsts {
+		dstLabel := alias.String()
 		matched := false
 
 		if ag, ok := alias.(*AutoGroup); ok && ag.Is(AutoGroupSelf) {
@@ -519,7 +519,7 @@ func resolveSSHTestDestNodes(
 			// (matches SaaS, which treats a no-node dst as a
 			// failing assertion).
 			if srcUserID == 0 {
-				emptyDsts = append(emptyDsts, dst)
+				emptyDsts = append(emptyDsts, dstLabel)
 
 				continue
 			}
@@ -548,7 +548,7 @@ func resolveSSHTestDestNodes(
 			}
 
 			if !matched {
-				emptyDsts = append(emptyDsts, dst)
+				emptyDsts = append(emptyDsts, dstLabel)
 			}
 
 			continue
@@ -556,11 +556,11 @@ func resolveSSHTestDestNodes(
 
 		ips, err := alias.Resolve(pol, users, nodes)
 		if err != nil {
-			return nil, nil, fmt.Errorf("resolving destination %q: %w", dst, err)
+			return nil, nil, fmt.Errorf("resolving destination %q: %w", dstLabel, err)
 		}
 
 		if ips == nil || ips.Empty() {
-			emptyDsts = append(emptyDsts, dst)
+			emptyDsts = append(emptyDsts, dstLabel)
 
 			continue
 		}
@@ -570,7 +570,7 @@ func resolveSSHTestDestNodes(
 		// the resolved prefixes.
 		set, err := prefixesToIPSet(ips.Prefixes())
 		if err != nil {
-			return nil, nil, fmt.Errorf("building IPSet for %q: %w", dst, err)
+			return nil, nil, fmt.Errorf("building IPSet for %q: %w", dstLabel, err)
 		}
 
 		for _, n := range nodes.All() {
@@ -589,7 +589,7 @@ func resolveSSHTestDestNodes(
 		}
 
 		if !matched {
-			emptyDsts = append(emptyDsts, dst)
+			emptyDsts = append(emptyDsts, dstLabel)
 		}
 	}
 

hscontrol/policy/v2/test.go

@@ -7,6 +7,7 @@ import (
 	"slices"
 	"strings"
 
+	"github.com/go-json-experiment/json"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"tailscale.com/tailcfg"
@@ -67,27 +68,92 @@ type PolicyTest struct {
 type SSHPolicyTest struct {
 	// Src is a single source alias (user, group, tag, host, or IP). Same
 	// shape as PolicyTest.Src — Tailscale only supports one src per entry.
-	Src string `json:"src"`
+	Src Alias `json:"src"`
 
 	// Dst lists destination host aliases the test exercises. Tags, hosts,
 	// and the SSH-compatible autogroups are valid; ports, CIDR ranges, and
 	// autogroup:internet are rejected at parse time.
-	Dst []string `json:"dst"`
+	Dst SSHTestDestinations `json:"dst"`
 
 	// Accept lists SSH login users that must be allowed by an action:accept
 	// or action:check rule when Src connects to each entry in Dst.
-	Accept []string `json:"accept,omitempty"`
+	Accept []SSHUser `json:"accept,omitempty"`
 
 	// Deny lists SSH login users that must NOT be allowed by any rule when
 	// Src connects to each entry in Dst.
-	Deny []string `json:"deny,omitempty"`
+	Deny []SSHUser `json:"deny,omitempty"`
 
 	// Check lists SSH login users that must reach every dst via an
 	// action:check rule specifically (the HoldAndDelegate signal on the
 	// compiled SSH policy). An action:accept rule alone does not satisfy
 	// a check assertion — SaaS keeps the two categories distinct so
 	// policy authors can pin sensitive logins to check rules.
-	Check []string `json:"check,omitempty"`
+	Check []SSHUser `json:"check,omitempty"`
+}
+
+// SSHTestDestinations is the list of destination aliases an sshTests entry
+// targets. Unmarshalling reuses the same alias parser the rest of the
+// policy engine drives so each element lands as a typed Alias; the parse-
+// time shape rules in validateSSHTestDestination continue to enforce the
+// SSH-specific restrictions (no :port, no CIDR, no autogroup:internet,
+// known tag).
+type SSHTestDestinations []Alias
+
+// UnmarshalJSON walks the JSON array, dispatching each element through
+// AliasEnc so trimming and prefix detection match the rest of the parser.
+func (d *SSHTestDestinations) UnmarshalJSON(b []byte) error {
+	var aliases []AliasEnc
+
+	err := json.Unmarshal(b, &aliases, policyJSONOpts...)
+	if err != nil {
+		return err
+	}
+
+	*d = make([]Alias, len(aliases))
+	for i, a := range aliases {
+		(*d)[i] = a.Alias
+	}
+
+	return nil
+}
+
+// UnmarshalJSON drives the typed shape of SSHPolicyTest. The wire format
+// is unchanged: src is a JSON string parsed through parseAlias; dst is an
+// array of strings handled by SSHTestDestinations; accept/deny/check are
+// arrays of strings handled per element by SSHUser.UnmarshalJSON. An
+// empty src string lands as a nil Alias so the empty-src case stays a
+// validation-time error with the SaaS-aligned ErrSSHTestEmptySrc body
+// rather than a raw parser failure.
+func (t *SSHPolicyTest) UnmarshalJSON(b []byte) error {
+	var raw struct {
+		Src    string              `json:"src"`
+		Dst    SSHTestDestinations `json:"dst"`
+		Accept []SSHUser           `json:"accept,omitempty"`
+		Deny   []SSHUser           `json:"deny,omitempty"`
+		Check  []SSHUser           `json:"check,omitempty"`
+	}
+
+	err := json.Unmarshal(b, &raw, policyJSONOpts...)
+	if err != nil {
+		return err
+	}
+
+	trimmedSrc := strings.TrimSpace(raw.Src)
+	if trimmedSrc != "" {
+		alias, parseErr := parseAlias(trimmedSrc)
+		if parseErr != nil {
+			return parseErr
+		}
+
+		t.Src = alias
+	}
+
+	t.Dst = raw.Dst
+	t.Accept = raw.Accept
+	t.Deny = raw.Deny
+	t.Check = raw.Check
+
+	return nil
 }
 
 // PolicyTestResult is the outcome of a single PolicyTest.

hscontrol/policy/v2/types.go

@@ -843,6 +843,12 @@ type Alias interface {
 	Validate() error
 	UnmarshalJSON(b []byte) error
 
+	// String renders the alias back to its policy-file form. Implementations
+	// are expected to return a value that round-trips through parseAlias for
+	// any alias the parser accepted, so callers can use it as a stable
+	// identity in rendered errors and logs.
+	String() string
+
 	// Resolve resolves the Alias to an IPSet. The IPSet will contain all the IP
 	// addresses that the Alias represents within Headscale. It is the product
 	// of the Alias and the Policy, Users and Nodes.
@@ -3286,7 +3292,7 @@ func validateSSHTests(pol *Policy, tests []SSHPolicyTest) error {
 	var errs []error
 
 	for i, t := range tests {
-		if t.Src == "" {
+		if t.Src == nil {
 			errs = append(errs, fmt.Errorf("sshTest %d: %w", i, ErrSSHTestEmptySrc))
 		}
 
@@ -3317,11 +3323,8 @@ func validateSSHTests(pol *Policy, tests []SSHPolicyTest) error {
 // (only valid in ACL destinations, not SSH ones). Tag entries must
 // reference a tag that exists in tagOwners; bare hosts must resolve to a
 // single-address prefix.
-func validateSSHTestDestination(pol *Policy, dst string) error {
-	alias, err := parseAlias(dst)
-	if err != nil {
-		return fmt.Errorf("%w %q", ErrSSHTestDstDisallowedElement, dst)
-	}
+func validateSSHTestDestination(pol *Policy, alias Alias) error {
+	dst := alias.String()
 
 	switch a := alias.(type) {
 	case *AutoGroup:
@@ -3333,12 +3336,12 @@ func validateSSHTestDestination(pol *Policy, dst string) error {
 		}
 
 	case *Prefix:
-		// A CIDR literal in dst is rejected. A bare IP parses as a Prefix
-		// with no slash in the input string — distinguish on the raw text
-		// the same way validateTestDestination does.
-		if strings.Contains(dst, "/") {
-			return fmt.Errorf("%w %q", ErrSSHTestDstDisallowedElement, dst)
-		}
+		// IP / CIDR dsts are rejected: SSH dsts name hosts, tags, or
+		// SSH-compatible autogroups, not raw addresses. Bare IPs and
+		// `/BitLen` literals parse to the same Prefix value so this
+		// branch covers both with one disallowed-element answer.
+		_ = a
+		return fmt.Errorf("%w %q", ErrSSHTestDstDisallowedElement, dst)
 
 	case *Tag:
 		// A tag must be declared in tagOwners. The `tag:server:22` shape

hscontrol/policy/v2/types_test.go

@@ -6003,9 +6003,12 @@ func TestUnmarshalPolicySSHTests(t *testing.T) {
 				t.Helper()
 				require.Len(t, pol.SSHTests, 1)
 				got := pol.SSHTests[0]
-				require.Equal(t, "thor@example.org", got.Src)
-				require.Equal(t, []string{"tag:server"}, got.Dst)
-				require.Equal(t, []string{"root"}, got.Accept)
+				require.IsType(t, (*Username)(nil), got.Src)
+				require.Equal(t, "thor@example.org", got.Src.String())
+				require.Len(t, got.Dst, 1)
+				require.IsType(t, (*Tag)(nil), got.Dst[0])
+				require.Equal(t, "tag:server", got.Dst[0].String())
+				require.Equal(t, []SSHUser{"root"}, got.Accept)
 				require.Empty(t, got.Deny)
 				require.Empty(t, got.Check)
 			},
@@ -6030,9 +6033,9 @@ func TestUnmarshalPolicySSHTests(t *testing.T) {
 				t.Helper()
 				require.Len(t, pol.SSHTests, 1)
 				got := pol.SSHTests[0]
-				require.Equal(t, []string{"root"}, got.Accept)
-				require.Equal(t, []string{"nobody"}, got.Deny)
-				require.Equal(t, []string{"alice"}, got.Check)
+				require.Equal(t, []SSHUser{"root"}, got.Accept)
+				require.Equal(t, []SSHUser{"nobody"}, got.Deny)
+				require.Equal(t, []SSHUser{"alice"}, got.Check) //nolint:goconst
 			},
 		},
 		{

integration/cli_policy_test.go

@@ -200,9 +200,9 @@ func TestSSHTestsRejectFailingPolicy(t *testing.T) {
 		},
 		SSHTests: []policyv2.SSHPolicyTest{
 			{
-				Src:    user1,
-				Dst:    []string{"autogroup:member"},
-				Accept: []string{"root"},
+				Src:    usernamep(user1),
+				Dst:    policyv2.SSHTestDestinations{new(policyv2.AutoGroupMember)},
+				Accept: []policyv2.SSHUser{policyv2.SSHUser("root")},
 			},
 		},
 	}
@@ -215,9 +215,9 @@ func TestSSHTestsRejectFailingPolicy(t *testing.T) {
 	badPolicy := goodPolicy
 	badPolicy.SSHTests = []policyv2.SSHPolicyTest{
 		{
-			Src:    user2,
-			Dst:    []string{"autogroup:member"},
-			Accept: []string{"root"},
+			Src:    usernamep(user2),
+			Dst:    policyv2.SSHTestDestinations{new(policyv2.AutoGroupMember)},
+			Accept: []policyv2.SSHUser{policyv2.SSHUser("root")},
 		},
 	}