[v14] Allow tsh ssh to use predicate expressions/search keywords from proxy templates

lib/client/api.go

@@ -1305,32 +1305,31 @@ func (tc *TeleportClient) getTargetNodes(ctx context.Context, clt client.GetReso
 		return []string{options.HostAddress}, nil
 	}
 
-	// use the target node that was explicitly provided if valid
-	if len(tc.Labels) == 0 {
-		// detect the common error when users use host:port address format
-		_, port, err := net.SplitHostPort(tc.Host)
-		// client has used host:port notation
-		if err == nil {
-			return nil, trace.BadParameter("please use ssh subcommand with '--port=%v' flag instead of semicolon", port)
+	// Query for nodes if labels, fuzzy search, or predicate expressions were provided.
+	if len(tc.Labels) > 0 || len(tc.SearchKeywords) > 0 || tc.PredicateExpression != "" {
+		nodes, err := client.GetAllResources[types.Server](ctx, clt, tc.ResourceFilter(types.KindNode))
+		if err != nil {
+			return nil, trace.Wrap(err)
 		}
 
-		addr := net.JoinHostPort(tc.Host, strconv.Itoa(tc.HostPort))
-		return []string{addr}, nil
-	}
+		retval := make([]string, 0, len(nodes))
+		for _, resource := range nodes {
+			// always dial nodes by UUID
+			retval = append(retval, fmt.Sprintf("%s:0", resource.GetName()))
+		}
 
-	// find the nodes matching the labels that were provided
-	nodes, err := client.GetAllResources[types.Server](ctx, clt, tc.ResourceFilter(types.KindNode))
-	if err != nil {
-		return nil, trace.Wrap(err)
+		return retval, nil
 	}
 
-	retval := make([]string, 0, len(nodes))
-	for _, resource := range nodes {
-		// always dial nodes by UUID
-		retval = append(retval, fmt.Sprintf("%s:0", resource.GetName()))
+	// detect the common error when users use host:port address format
+	_, port, err := net.SplitHostPort(tc.Host)
+	// client has used host:port notation
+	if err == nil {
+		return nil, trace.BadParameter("please use ssh subcommand with '--port=%v' flag instead of semicolon", port)
 	}
 
-	return retval, nil
+	addr := net.JoinHostPort(tc.Host, strconv.Itoa(tc.HostPort))
+	return []string{addr}, nil
 }
 
 // ReissueUserCerts issues new user certs based on params and stores them in

lib/client/api_test.go

@@ -33,6 +33,8 @@ import (
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 
+	apiclient "github.com/gravitational/teleport/api/client"
+	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/utils/grpc/interceptors"
@@ -1244,3 +1246,83 @@ func TestIsErrorResolvableWithRelogin(t *testing.T) {
 		})
 	}
 }
+
+type fakeResourceClient struct {
+	apiclient.GetResourcesClient
+
+	nodes []*types.ServerV2
+}
+
+func (f fakeResourceClient) GetResources(ctx context.Context, req *proto.ListResourcesRequest) (*proto.ListResourcesResponse, error) {
+	out := make([]*proto.PaginatedResource, 0, len(f.nodes))
+	for _, n := range f.nodes {
+		out = append(out, &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: n}})
+	}
+
+	return &proto.ListResourcesResponse{Resources: out}, nil
+}
+
+func TestGetTargetNodes(t *testing.T) {
+	tests := []struct {
+		name      string
+		options   SSHOptions
+		labels    map[string]string
+		search    []string
+		predicate string
+		host      string
+		port      int
+		clt       fakeResourceClient
+		expected  []string
+	}{
+		{
+			name: "options override",
+			options: SSHOptions{
+				HostAddress: "test:1234",
+			},
+			expected: []string{"test:1234"},
+		},
+		{
+			name:     "explicit target",
+			host:     "test",
+			port:     1234,
+			expected: []string{"test:1234"},
+		},
+		{
+			name:     "labels",
+			labels:   map[string]string{"foo": "bar"},
+			expected: []string{"abcd:0"},
+			clt:      fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "labels"}}}},
+		},
+		{
+			name:     "search",
+			search:   []string{"foo", "bar"},
+			expected: []string{"abcd:0"},
+			clt:      fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "search"}}}},
+		},
+		{
+			name:      "predicate",
+			predicate: `resource.spec.hostname == "test"`,
+			expected:  []string{"abcd:0"},
+			clt:       fakeResourceClient{nodes: []*types.ServerV2{{Metadata: types.Metadata{Name: "abcd"}, Spec: types.ServerSpecV2{Hostname: "predicate"}}}},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			clt := TeleportClient{
+				Config: Config{
+					Tracer:              tracing.NoopTracer(""),
+					Labels:              test.labels,
+					SearchKeywords:      test.search,
+					PredicateExpression: test.predicate,
+					Host:                test.host,
+					HostPort:            test.port,
+				},
+			}
+
+			match, err := clt.getTargetNodes(context.Background(), test.clt, test.options)
+			require.NoError(t, err)
+			require.EqualValues(t, test.expected, match)
+		})
+	}
+}