Support context during connect

Signed-off-by: Luca Bigliardi <shammash@google.com>

client/connection.go

@@ -2,7 +2,9 @@ package client
 
 import (
 	"bufio"
+	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -36,7 +38,7 @@ type Conn struct {
 
 	// I/O stuff to server
 	dialer      *net.Dialer
-	proxyDialer proxy.Dialer
+	proxyDialer proxy.ContextDialer
 	sock        net.Conn
 	io          *bufio.ReadWriter
 	in          chan *Line
@@ -304,11 +306,16 @@ func (conn *Conn) initialise() {
 // Config.Server to host, Config.Pass to pass if one is provided, and then
 // calls Connect.
 func (conn *Conn) ConnectTo(host string, pass ...string) error {
+	return conn.ConnectToContext(context.Background(), host, pass...)
+}
+
+// ConnectToContext works like ConnectTo but uses the provided context.
+func (conn *Conn) ConnectToContext(ctx context.Context, host string, pass ...string) error {
 	conn.cfg.Server = host
 	if len(pass) > 0 {
 		conn.cfg.Pass = pass[0]
 	}
-	return conn.Connect()
+	return conn.ConnectContext(ctx)
 }
 
 // Connect connects the IRC client to the server configured in Config.Server.
@@ -323,10 +330,15 @@ func (conn *Conn) ConnectTo(host string, pass ...string) error {
 // handler for the CONNECTED event is used to perform any initial client work
 // like joining channels and sending messages.
 func (conn *Conn) Connect() error {
+	return conn.ConnectContext(context.Background())
+}
+
+// ConnectContext works like Connect but uses the provided context.
+func (conn *Conn) ConnectContext(ctx context.Context) error {
 	// We don't want to hold conn.mu while firing the REGISTER event,
 	// and it's much easier and less error prone to defer the unlock,
 	// so the connect mechanics have been delegated to internalConnect.
-	err := conn.internalConnect()
+	err := conn.internalConnect(ctx)
 	if err == nil {
 		conn.dispatch(&Line{Cmd: REGISTER, Time: time.Now()})
 	}
@@ -334,7 +346,7 @@ func (conn *Conn) Connect() error {
 }
 
 // internalConnect handles the work of actually connecting to the server.
-func (conn *Conn) internalConnect() error {
+func (conn *Conn) internalConnect(ctx context.Context) error {
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 	conn.initialise()
@@ -359,20 +371,24 @@ func (conn *Conn) internalConnect() error {
 		if err != nil {
 			return err
 		}
-		conn.proxyDialer, err = proxy.FromURL(proxyURL, conn.dialer)
+		proxyDialer, err := proxy.FromURL(proxyURL, conn.dialer)
 		if err != nil {
 			return err
 		}
-
+		contextProxyDialer, ok := proxyDialer.(proxy.ContextDialer)
+		if !ok {
+			return errors.New("Dialer for proxy does not support context")
+		}
+		conn.proxyDialer = contextProxyDialer
 		logging.Info("irc.Connect(): Connecting to %s.", conn.cfg.Server)
-		if s, err := conn.proxyDialer.Dial("tcp", conn.cfg.Server); err == nil {
+		if s, err := conn.proxyDialer.DialContext(ctx, "tcp", conn.cfg.Server); err == nil {
 			conn.sock = s
 		} else {
 			return err
 		}
 	} else {
 		logging.Info("irc.Connect(): Connecting to %s.", conn.cfg.Server)
-		if s, err := conn.dialer.Dial("tcp", conn.cfg.Server); err == nil {
+		if s, err := conn.dialer.DialContext(ctx, "tcp", conn.cfg.Server); err == nil {
 			conn.sock = s
 		} else {
 			return err

go.mod

@@ -1,8 +1,8 @@
 module github.com/fluffle/goirc
 
 require (
-	github.com/golang/mock v1.1.1
+	github.com/golang/mock v1.5.0
-	golang.org/x/net v0.0.0-20180926154720-4dfa2610cdf3
+	golang.org/x/net v0.0.0-20210119194325-5f4716e94777
 )
 
 go 1.13

go.sum

@@ -1,4 +1,11 @@
 github.com/golang/mock v1.1.1 h1:G5FRp8JnTd7RQH5kemVNlMeyXQAztQ3mOWV95KxsXH8=
 github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
+github.com/golang/mock v1.5.0 h1:jlYHihg//f7RRwuPfptm04yp4s7O6Kw8EZiVYIGcH0g=
+github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
 golang.org/x/net v0.0.0-20180926154720-4dfa2610cdf3 h1:dgd4x4kJt7G4k4m93AYLzM8Ni6h2qLTfh9n9vXJT3/0=
 golang.org/x/net v0.0.0-20180926154720-4dfa2610cdf3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20210119194325-5f4716e94777 h1:003p0dJM77cxMSyCPFphvZf/Y5/NXf5fzg6ufd1/Oew=
+golang.org/x/net v0.0.0-20210119194325-5f4716e94777/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=