[client] Return the context error when the SSH handshake fails with it (#7426)

* [client] Return the context error when the SSH handshake fails on a context deadline

The handshake mapped the context deadline onto the socket but returned the
raw socket error. Which error surfaces depends on a race between the x/crypto
ssh readLoop and kexLoop goroutines: the kexLoop write fails with i/o timeout
and closes the conn, and the readLoop then reports use of closed network
connection. Callers checking errors.Is(err, context.DeadlineExceeded) never
matched, and TestSSHClient_ContextCancellation flaked on the FreeBSD job.

Handshake now wraps the context error when the context is done or its
deadline has passed. The deadline comparison is needed because the socket
deadline and the context timer fire independently, so ctx.Err() can still be
nil when the deadline-triggered socket error arrives.

* [client] Close the silent test server conn without racing t.Cleanup

The accept goroutine registered the conn close via t.Cleanup, which can run
after the test's cleanup list has already been drained, leaving the accepted
connection open. The goroutine now holds the conn until a cleanup-closed
channel signals the end of the test and closes it on the way out.

* [client] Bind the SSH handshake to the context instead of a socket deadline

Mapping only the context deadline onto the socket left context cancellation
unobserved: an in-flight handshake kept running until the deadline, and the
error classification had to guess whether a raw socket error was caused by
the deadline. Closing the conn from context.AfterFunc covers both deadline
and cancellation, and ctx.Err() is already set by the time the close-induced
error surfaces, so the time-based DeadlineExceeded attribution is no longer
needed. The stop() result guards the window between a successful handshake
and the AfterFunc firing so a closed conn is never handed back as a client.

client/ssh/handshake.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"net"
-	"time"
 
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
@@ -13,26 +12,23 @@ import (
 
 // Handshake runs the SSH client handshake on an already dialed conn and
 // returns the resulting client. Dialing bounds only the TCP establishment;
-// without a deadline on the socket a peer that accepts and then goes silent
-// blocks the handshake forever, so the context deadline is applied to conn
-// for the duration of the handshake. conn is closed on any error.
+// a peer that accepts and then goes silent would block the handshake forever,
+// so conn is closed as soon as ctx is done, which unblocks the handshake and
+// surfaces the context error. conn is closed on any error.
 func Handshake(ctx context.Context, conn net.Conn, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
-	if deadline, ok := ctx.Deadline(); ok {
-		if err := conn.SetDeadline(deadline); err != nil {
-			closeHandshake(conn, "conn after deadline error")
-			return nil, fmt.Errorf("set handshake deadline: %w", err)
-		}
-	}
+	stop := context.AfterFunc(ctx, func() { closeHandshake(conn, "conn on context done") })
 
 	sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
 	if err != nil {
-		closeHandshake(conn, "conn after handshake error")
-		return nil, fmt.Errorf("ssh handshake: %w", err)
+		if stop() {
+			closeHandshake(conn, "conn after handshake error")
+		}
+		return nil, handshakeError(ctx, err)
 	}
 
-	if err := conn.SetDeadline(time.Time{}); err != nil {
-		closeHandshake(sshConn, "ssh conn after deadline clear error")
-		return nil, fmt.Errorf("clear handshake deadline: %w", err)
+	if !stop() {
+		closeHandshake(sshConn, "ssh conn after context done")
+		return nil, fmt.Errorf("ssh handshake: %w", ctx.Err())
 	}
 
 	return ssh.NewClient(sshConn, chans, reqs), nil
@@ -43,3 +39,10 @@ func closeHandshake(c io.Closer, label string) {
 		log.Debugf("ssh: close %s: %v", label, err)
 	}
 }
+
+func handshakeError(ctx context.Context, err error) error {
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		return fmt.Errorf("ssh handshake: %w: %w", ctxErr, err)
+	}
+	return fmt.Errorf("ssh handshake: %w", err)
+}

client/ssh/handshake_test.go

@@ -0,0 +1,90 @@
+package ssh
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
+)
+
+func TestHandshake_ContextDeadlineWrapped(t *testing.T) {
+	conn := dialSilentServer(t)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	_, err := Handshake(ctx, conn, conn.RemoteAddr().String(), testClientConfig())
+	require.Error(t, err)
+	require.True(t, errors.Is(err, context.DeadlineExceeded), "expected context.DeadlineExceeded, got: %v", err)
+}
+
+func TestHandshake_ContextCancelUnblocks(t *testing.T) {
+	conn := dialSilentServer(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	time.AfterFunc(50*time.Millisecond, cancel)
+
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := Handshake(ctx, conn, conn.RemoteAddr().String(), testClientConfig())
+		errCh <- err
+	}()
+
+	select {
+	case err := <-errCh:
+		require.Error(t, err)
+		require.True(t, errors.Is(err, context.Canceled), "expected context.Canceled, got: %v", err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("handshake did not return after context cancellation")
+	}
+}
+
+func TestHandshake_NonContextErrorNotWrapped(t *testing.T) {
+	conn := dialSilentServer(t)
+	require.NoError(t, conn.Close())
+
+	_, err := Handshake(context.Background(), conn, conn.RemoteAddr().String(), testClientConfig())
+	require.Error(t, err)
+	require.False(t, errors.Is(err, context.Canceled))
+	require.False(t, errors.Is(err, context.DeadlineExceeded))
+}
+
+func testClientConfig() *ssh.ClientConfig {
+	return &ssh.ClientConfig{
+		User:            "test",
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+}
+
+// dialSilentServer returns a client conn to a server that accepts and never
+// sends anything, so the SSH handshake blocks until the context is done.
+func dialSilentServer(t *testing.T) net.Conn {
+	t.Helper()
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = listener.Close() })
+
+	done := make(chan struct{})
+	t.Cleanup(func() { close(done) })
+
+	go func() {
+		c, err := listener.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = c.Close() }()
+		<-done
+	}()
+
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	return conn
+}