From 8f90e678c5e1ac6cd2353317b28acfb8bc869454 Mon Sep 17 00:00:00 2001
From: rosstimothy <39066650+rosstimothy@users.noreply.github.com>
Date: Mon, 12 Sep 2022 09:31:11 -0400
Subject: [PATCH] Dial by UUID for label based ssh (#16257)

Switches label based ssh to always use the node UUID instead of its
address. This prevents issues that occurred when dialing by address
if the public_addr isn't set correctly.

Closes #3214
---
 lib/client/api.go    | 42 ++++++++++++++++++------------------------
 tool/tsh/tsh_test.go | 10 ++++------
 2 files changed, 22 insertions(+), 30 deletions(-)

diff --git a/lib/client/api.go b/lib/client/api.go
index 1996aa90506ae..67be97f8a098c 100644
--- a/lib/client/api.go
+++ b/lib/client/api.go
@@ -1635,37 +1635,31 @@ func (tc *TeleportClient) getTargetNodes(ctx context.Context, proxy *ProxyClient
 	)
 	defer span.End()
 
-	var (
-		err    error
-		nodes  []types.Server
-		retval = make([]string, 0)
-	)
-	if tc.Labels != nil && len(tc.Labels) > 0 {
-		nodes, err = proxy.FindNodesByFilters(ctx, *tc.DefaultResourceFilter())
-		if err != nil {
-			return nil, trace.Wrap(err)
-		}
-		for i := 0; i < len(nodes); i++ {
-			addr := nodes[i].GetAddr()
-			if addr == "" {
-				// address is empty, try dialing by UUID instead.
-				addr = fmt.Sprintf("%s:0", nodes[i].GetName())
-			}
-			retval = append(retval, addr)
-		}
-	}
-	if len(nodes) == 0 {
+	// 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)
+			return nil, trace.BadParameter("please use ssh subcommand with '--port=%v' flag instead of semicolon", port)
 		}
+
 		addr := net.JoinHostPort(tc.Host, strconv.Itoa(tc.HostPort))
-		retval = append(retval, addr)
+		return []string{addr}, nil
 	}
+
+	// find the nodes matching the labels that were provided
+	nodes, err := proxy.FindNodesByFilters(ctx, *tc.DefaultResourceFilter())
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	retval := make([]string, 0, len(nodes))
+	for i := 0; i < len(nodes); i++ {
+		// always dial nodes by UUID
+		retval = append(retval, fmt.Sprintf("%s:0", nodes[i].GetName()))
+	}
+
 	return retval, nil
 }
 
diff --git a/tool/tsh/tsh_test.go b/tool/tsh/tsh_test.go
index be35e361564e6..c378a19ca5e22 100644
--- a/tool/tsh/tsh_test.go
+++ b/tool/tsh/tsh_test.go
@@ -1048,12 +1048,10 @@ func TestSSHOnMultipleNodes(t *testing.T) {
 					RequireSessionMFA: true,
 				},
 			},
-			setup:        registerPasswordlessDeviceWithWebauthnSolver,
-			hostLabels:   "env=dev",
-			errAssertion: require.Error,
-			// label is missing however tsh will try to connect over to agent "env=dev:3022"
-			// and fail, however that's why it require sign
-			deviceSignCount: 1,
+			setup:           registerPasswordlessDeviceWithWebauthnSolver,
+			hostLabels:      "env=dev",
+			errAssertion:    require.Error,
+			deviceSignCount: 0,
 			stdoutAssertion: require.Empty,
 		},
 	}