Make Ctrl-C work when an SSH connection stalls

internal/command/ssh/ssh_terminal.go

@@ -96,7 +96,7 @@ func SSHConnect(p *SSHParams, addr string) error {
 		defer endSpin()
 	}
 
-	if err := sshClient.Connect(context.Background()); err != nil {
+	if err := sshClient.Connect(p.Ctx); err != nil {
 		return errors.Wrap(err, "error connecting to SSH server")
 	}
 	defer sshClient.Close()
@@ -115,7 +115,7 @@ func SSHConnect(p *SSHParams, addr string) error {
 		TermEnv:  "xterm",
 	}
 
-	if err := sshClient.Shell(context.Background(), sessIO, p.Cmd, ssh.SessionTarget{}); err != nil {
+	if err := sshClient.Shell(p.Ctx, sessIO, p.Cmd, ssh.SessionTarget{}); err != nil {
 		return errors.Wrap(err, "ssh shell")
 	}
 

ssh/client.go

@@ -75,7 +75,9 @@ func (c *Client) Connect(ctx context.Context) error {
 		HostKeyAlgorithms: []string{ssh.KeyAlgoED25519},
 	}
 
-	respCh := make(chan connResp)
+	// Buffered so the handshake goroutine can always send and exit, even once
+	// nobody is left to read the result.
+	respCh := make(chan connResp, 1)
 
 	// ssh.NewClientConn doesn't take a context, so we need to handle cancelation on our end
 	go func() {
@@ -91,19 +93,22 @@ func (c *Client) Connect(ctx context.Context) error {
 		respCh <- connResp{nil, conn, client}
 	}()
 
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case resp := <-respCh:
-			if resp.err != nil {
-				return resp.err
-			}
-			c.conn = resp.conn
-			c.Client = resp.client
-
-			return nil
+	select {
+	case <-ctx.Done():
+		// Closing the socket is what unblocks the handshake above, which owns
+		// tcpConn from here on and has no other way to learn we gave up.
+		tcpConn.Close()
+
+		return ctx.Err()
+	case resp := <-respCh:
+		if resp.err != nil {
+			// ssh.NewClientConn closes tcpConn itself when the handshake fails.
+			return resp.err
 		}
+		c.conn = resp.conn
+		c.Client = resp.client
+
+		return nil
 	}
 }
 

ssh/client_test.go

@@ -1,6 +1,15 @@
 package ssh
 
-import "testing"
+import (
+	"context"
+	"crypto/ed25519"
+	"crypto/rand"
+	"errors"
+	"net"
+	"testing"
+
+	"golang.org/x/crypto/ssh"
+)
 
 func TestSessionTargetValidate(t *testing.T) {
 	for _, tc := range []struct {
@@ -40,3 +49,66 @@ func TestSessionTargetValidate(t *testing.T) {
 		})
 	}
 }
+
+func TestConnectStopsWhenTheContextIsCanceled(t *testing.T) {
+	// The far end never speaks SSH, so the handshake blocks on the version
+	// exchange until something closes the socket under it.
+	local, remote := net.Pipe()
+	defer remote.Close()
+
+	certificate, privateKey := testCredentials(t)
+
+	client := &Client{
+		Addr:        "192.0.2.1:22",
+		User:        "root",
+		Certificate: certificate,
+		PrivateKey:  privateKey,
+		Dial: func(context.Context, string, string) (net.Conn, error) {
+			return local, nil
+		},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	if err := client.Connect(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+
+	// A socket left open here strands the handshake goroutine on it.
+	if _, err := remote.Read(make([]byte, 1)); err == nil {
+		t.Fatal("expected the socket to be closed, it still reads")
+	}
+}
+
+// testCredentials returns a self-signed user certificate and its private key,
+// which is as much as Connect parses before it dials.
+func testCredentials(t *testing.T) (certificate, privateKey string) {
+	t.Helper()
+
+	pub, priv, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	signer, err := ssh.NewSignerFromKey(priv)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sshPub, err := ssh.NewPublicKey(pub)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cert := &ssh.Certificate{
+		Key:         sshPub,
+		CertType:    ssh.UserCert,
+		ValidBefore: ssh.CertTimeInfinity,
+	}
+	if err := cert.SignCert(rand.Reader, signer); err != nil {
+		t.Fatal(err)
+	}
+
+	return string(ssh.MarshalAuthorizedKey(cert)), string(MarshalED25519PrivateKey(priv, "test"))
+}