From 0bc0cb81318f4f6e40428d5b175cdad6c1595c7e Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 24 Jun 2026 05:31:27 +0000
Subject: [PATCH] tstest/natlab/vmtest: retry SSHExec on transient SSH failures

Add a retry loop with BatchMode=yes to absorb the race window
between Env.Start() returning (when tta reports the tailscale
backend as Running) and cloud-init finishing the user/SSH-key
setup. In CI, the second VM's tta agent has been observed
connecting only a few hundred milliseconds before the test SSHes
in, which is inside the window where /root/.ssh/authorized_keys
hasn't fully landed yet. SSH key auth then fails and ssh(1) falls
back to interactive password prompts (3x), wasting time and
producing a confusing "Permission denied (publickey,password)"
error.

BatchMode=yes makes the client fail fast on auth failure instead
of prompting, and the retry loop handles SSH transport-level
errors (exit code 255) for up to 30 seconds with 500ms backoff.
Remote command non-zero exits still pass through unchanged.

Fixes #20228

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
Change-Id: I17f7422e9e27bf7b995f505c0184cbb2b230ed81
---
 tstest/natlab/vmtest/vmtest.go | 44 +++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/tstest/natlab/vmtest/vmtest.go b/tstest/natlab/vmtest/vmtest.go
index fdf4d4245..7d385f09e 100644
--- a/tstest/natlab/vmtest/vmtest.go
+++ b/tstest/natlab/vmtest/vmtest.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -1302,21 +1303,42 @@ func (e *Env) AddRoute(n *Node, prefix, via string) {
 // SSHExec runs a command on a cloud VM via its debug SSH NIC.
 // Only works for cloud VMs that have the debug NIC and SSH key configured.
 // Returns stdout and any error.
+//
+// SSH transport-level errors (exit code 255: connection refused, auth
+// failure, etc.) are retried for up to ~30s to absorb the race window
+// between Env.Start() returning (when tta reports the tailscale backend
+// as Running) and cloud-init finishing the user/SSH-key setup. The remote
+// command's own non-zero exit codes are returned to the caller without
+// retry.
 func (e *Env) SSHExec(n *Node, cmd string) (string, error) {
 	if n.sshPort == 0 {
 		return "", fmt.Errorf("node %s has no SSH debug port", n.name)
 	}
-	sshCmd := exec.Command("ssh",
-		"-o", "StrictHostKeyChecking=no",
-		"-o", "UserKnownHostsFile=/dev/null",
-		"-o", "ConnectTimeout=5",
-		"-o", "LogLevel=ERROR",
-		"-i", "/tmp/vmtest_key",
-		"-p", fmt.Sprintf("%d", n.sshPort),
-		"root@127.0.0.1",
-		cmd)
-	out, err := sshCmd.CombinedOutput()
-	return string(out), err
+	deadline := time.Now().Add(30 * time.Second)
+	for {
+		sshCmd := exec.Command("ssh",
+			"-o", "StrictHostKeyChecking=no",
+			"-o", "UserKnownHostsFile=/dev/null",
+			"-o", "BatchMode=yes",
+			"-o", "ConnectTimeout=5",
+			"-o", "LogLevel=ERROR",
+			"-i", "/tmp/vmtest_key",
+			"-p", fmt.Sprintf("%d", n.sshPort),
+			"root@127.0.0.1",
+			cmd)
+		out, err := sshCmd.CombinedOutput()
+		if err == nil {
+			return string(out), nil
+		}
+		var exitErr *exec.ExitError
+		if !errors.As(err, &exitErr) || exitErr.ExitCode() != 255 {
+			return string(out), err
+		}
+		if time.Now().After(deadline) {
+			return string(out), err
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
 }
 
 // DumpStatus logs the tailscale status of a node, including its peers and their