ssh: add (*Client).DialContext method

ssh/tcpip.go

@@ -5,6 +5,7 @@
 package ssh
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -332,6 +333,40 @@ func (l *tcpListener) Addr() net.Addr {
 	return l.laddr
 }
 
+// DialContext initiates a connection to the addr from the remote host.
+//
+// The provided Context must be non-nil. If the context expires before the
+// connection is complete, an error is returned. Once successfully connected,
+// any expiration of the context will not affect the connection.
+//
+// See func Dial for additional information.
+func (c *Client) DialContext(ctx context.Context, n, addr string) (net.Conn, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	type connErr struct {
+		conn net.Conn
+		err  error
+	}
+	ch := make(chan connErr)
+	go func() {
+		conn, err := c.Dial(n, addr)
+		select {
+		case ch <- connErr{conn, err}:
+		case <-ctx.Done():
+			if conn != nil {
+				conn.Close()
+			}
+		}
+	}()
+	select {
+	case res := <-ch:
+		return res.conn, res.err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
 // Dial initiates a connection to the addr from the remote host.
 // The resulting connection has a zero LocalAddr() and RemoteAddr().
 func (c *Client) Dial(n, addr string) (net.Conn, error) {

ssh/tcpip_test.go

@@ -5,7 +5,10 @@
 package ssh
 
 import (
+	"context"
+	"net"
 	"testing"
+	"time"
 )
 
 func TestAutoPortListenBroken(t *testing.T) {
@@ -18,3 +21,33 @@ func TestAutoPortListenBroken(t *testing.T) {
 		t.Errorf("version %q marked as broken", works)
 	}
 }
+
+func TestClientImplementsDialContext(t *testing.T) {
+	type ContextDialer interface {
+		DialContext(context.Context, string, string) (net.Conn, error)
+	}
+	// Belt and suspenders assertion, since package net does not
+	// declare a ContextDialer type.
+	var _ ContextDialer = &net.Dialer{}
+	var _ ContextDialer = &Client{}
+}
+
+func TestClientDialContextWithCancel(t *testing.T) {
+	c := &Client{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err := c.DialContext(ctx, "tcp", "localhost:1000")
+	if err != context.Canceled {
+		t.Errorf("DialContext: got nil error, expected %v", context.Canceled)
+	}
+}
+
+func TestClientDialContextWithDeadline(t *testing.T) {
+	c := &Client{}
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now())
+	defer cancel()
+	_, err := c.DialContext(ctx, "tcp", "localhost:1000")
+	if err != context.DeadlineExceeded {
+		t.Errorf("DialContext: got nil error, expected %v", context.DeadlineExceeded)
+	}
+}

ssh/test/dial_unix_test.go

@@ -10,6 +10,7 @@ package test
 // direct-tcpip and direct-streamlocal functional tests
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
@@ -47,7 +48,11 @@ func testDial(t *testing.T, n, listenAddr string, x dialTester) {
 		}
 	}()
 
-	conn, err := sshConn.Dial(n, l.Addr().String())
+	ctx, cancel := context.WithCancel(context.Background())
+	conn, err := sshConn.DialContext(ctx, n, l.Addr().String())
+	// Canceling the context after dial should have no effect
+	// on the opened connection.
+	cancel()
 	if err != nil {
 		t.Fatalf("Dial: %v", err)
 	}